fix: Fix speculative decoding

abetlen · Sep 26, 2024 · 9992c50 · 9992c50
1 parent 65222bc
commit 9992c50
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py
@@ -807,8 +807,10 @@ def sample(
                 grammar=grammar,
             )
 
+        ridx = idx - self.n_tokens if idx is not None else -1
+
         assert self.ctx is not None
-        token = self._sampler.sample(self._ctx, -1)
+        token = self._sampler.sample(self._ctx, ridx)
         if tmp_sampler:
             self._sampler = None
         return token