fix sample_idx off-by-one error

Andrew Lapp · Andrew Lapp · commit a4677ad2548c · 2024-02-12T14:38:22.000-06:00
diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py
@@ -658,6 +658,7 @@ def generate(
         while True:
             self.eval(tokens)
             while sample_idx < self.n_tokens:
+                next_sample_idx = sample_idx + 1
                 token = self.sample(
                     top_k=top_k,
                     top_p=top_p,
@@ -674,7 +675,7 @@ def generate(
                     logits_processor=logits_processor,
                     grammar=grammar,
                     penalize_nl=penalize_nl,
-                    idx=sample_idx,
+                    idx=next_sample_idx,
                 )
 
                 sample_idx += 1
