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

* 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 <andrew@rew.la>
Co-authored-by: Andrei <abetlen@gmail.com>
This commit is contained in:
Andrew Lapp 2024-02-13 17:26:07 +00:00 committed by GitHub
parent f7cdf78788
commit d6be5333e1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -557,7 +557,7 @@ class Llama:
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(