Fix last_n_tokens_size

This commit is contained in:
Andrei Betlen 2023-05-17 01:42:51 -04:00
parent 7e55244540
commit f11e2a781c

View file

@ -295,7 +295,9 @@ class Llama:
assert self.ctx is not None assert self.ctx is not None
assert len(self.eval_logits) > 0 assert len(self.eval_logits) > 0
n_vocab = int(llama_cpp.llama_n_vocab(self.ctx)) 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 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] logits = self.eval_logits[-1]
data = (llama_cpp.llama_token_data * n_vocab)( data = (llama_cpp.llama_token_data * n_vocab)(
*[ *[