diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py
index 034d7c0..6dc113a 100644
--- a/llama_cpp/llama.py
+++ b/llama_cpp/llama.py
@@ -491,7 +491,7 @@ class _LlamaContext:
     def sample_repetition_penalties(
         self,
         candidates: "_LlamaTokenDataArray",
-        last_tokens_data: llama_cpp.Array[llama_cpp.llama_token],
+        last_tokens_data: "llama_cpp.Array[llama_cpp.llama_token]",
         penalty_last_n: int,
         penalty_repeat: float,
         penalty_freq: float,