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(

Original file line number	Diff line number	Diff line change
`@@ -557,7 +557,7 @@ def sample(`
`557`	`557`	`logits[:] = (`
`558`	`558`	`logits_processor(self._input_ids, logits)`
`559`	`559`	`if idx is None`
`560`		`- else logits_processor(self._input_ids[:idx], logits)`
	`560`	`+ else logits_processor(self._input_ids[:idx + 1], logits)`
`561`	`561`	`)`
`562`	`562`
`563`	`563`	`sampling_params = _LlamaSamplingParams(`