Merge branch 'main' into riedgar-ms/llama-cpp-ub-remove

guidance-ai · Apr 23, 2024 · 5d090a7 · 5d090a7
2 parents abd5ace + 0d06746
commit 5d090a7
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/guidance/models/llama_cpp/_llama_cpp.py b/guidance/models/llama_cpp/_llama_cpp.py
@@ -198,7 +198,8 @@ def get_logits(self, token_ids, forced_bytes, current_temp):
 
         # get the logits
         logits = llama_cpp.llama_get_logits(self.model_obj.ctx)
-        logits = logits[(n_tokens - 1) * self._n_vocab : n_tokens * self._n_vocab]
+        if llama_cpp.__version__ < "0.2.58":
+            logits = logits[(n_tokens - 1) * self._n_vocab : n_tokens * self._n_vocab]
         logits = np.ctypeslib.as_array(logits, shape=(self._n_vocab,)).copy()
 
         self._cached_logits = logits