fix: sample idx off-by-one error for logit_processors (#1179)

lapp0 · Andrew Lapp · abetlen · web-flow · commit d6be5333e1e2 · 2024-02-13T12:26:07.000-05:00
* fix sample_idx off-by-one error

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

---------

Co-authored-by: Andrew Lapp &lt;andrew@rew.la&gt;
Co-authored-by: Andrei &lt;abetlen@gmail.com&gt;
diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py
@@ -557,7 +557,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(
