self._scores is indexed differently, only modify the index within self._input_ids

Andrew Lapp · Andrew Lapp · commit 493d6a70f8bb · 2024-02-13T10:52:54.000-06:00
diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py
@@ -553,7 +553,7 @@ def sample(
             logits[:] = (
                 logits_processor(self._input_ids, logits)
                 if idx is None
-                else logits_processor(self._input_ids[:idx], logits)
+                else logits_processor(self._input_ids[:idx + 1], logits)
             )
 
         sampling_params = _LlamaSamplingParams(
@@ -658,7 +658,6 @@ 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,
@@ -675,7 +674,7 @@ def generate(
                     logits_processor=logits_processor,
                     grammar=grammar,
                     penalize_nl=penalize_nl,
-                    idx=next_sample_idx,
+                    idx=sample_idx,
                 )
 
                 sample_idx += 1
