From d6be5333e1e28dd07cfec5babd6332c7d1f50788 Mon Sep 17 00:00:00 2001 From: Andrew Lapp Date: Tue, 13 Feb 2024 17:26:07 +0000 Subject: [PATCH] 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 Co-authored-by: Andrei --- llama_cpp/llama.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py index 4869a9d..8d726d3 100644 --- a/llama_cpp/llama.py +++ b/llama_cpp/llama.py @@ -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(