Fix top_k value. Closes abetlen#220

abetlen · abetlen · commit 7e5524454052 · 2023-05-17T01:41:42.000-04:00
diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py
@@ -295,6 +295,7 @@ def _sample(
         assert self.ctx is not None
         assert len(self.eval_logits) > 0
         n_vocab = int(llama_cpp.llama_n_vocab(self.ctx))
+        top_k = llama_cpp.c_int(n_vocab) if top_k.value <= 0 else top_k
         logits = self.eval_logits[-1]
         data = (llama_cpp.llama_token_data * n_vocab)(
             *[
