Fix top_k value. Closes #220

2023-05-17 01:41:42 -04:00 · 2023-05-17 01:41:42 -04:00 · 7e55244540
commit 7e55244540
parent e37a808bc0
1 changed files with 1 additions and 0 deletions
--- a/llama_cpp/llama.py
+++ b/llama_cpp/llama.py
@ -295,6 +295,7 @@ class Llama:
        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)(
            *[