From f11e2a781c6f3a6de03d67a52fa529e6c147f1b3 Mon Sep 17 00:00:00 2001 From: Andrei Betlen Date: Wed, 17 May 2023 01:42:51 -0400 Subject: [PATCH] Fix last_n_tokens_size --- llama_cpp/llama.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py index 4c8ba39..44363a8 100644 --- a/llama_cpp/llama.py +++ b/llama_cpp/llama.py @@ -295,7 +295,9 @@ class Llama: assert self.ctx is not None assert len(self.eval_logits) > 0 n_vocab = int(llama_cpp.llama_n_vocab(self.ctx)) + n_ctx = int(llama_cpp.llama_n_ctx(self.ctx)) top_k = llama_cpp.c_int(n_vocab) if top_k.value <= 0 else top_k + last_n_tokens_size = llama_cpp.c_int(n_ctx) if last_n_tokens_size.value < 0 else last_n_tokens_size logits = self.eval_logits[-1] data = (llama_cpp.llama_token_data * n_vocab)( *[