fix: Additional fixes for speculative decoding

abetlen · abetlen · commit e975dabf74b3 · 2024-09-26T01:11:02.000-04:00
diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py
@@ -930,7 +930,7 @@ def generate(
 
                 sample_idx += 1
                 if stopping_criteria is not None and stopping_criteria(
-                    self._input_ids, self._scores[-1, :]
+                    self._input_ids[: sample_idx], self._scores[sample_idx - self.n_tokens, :]
                 ):
                     return
                 tokens_or_none = yield token
