Format
This commit is contained in:
parent
34c505edf2
commit
4f2b5d0b53
|
@ -445,8 +445,12 @@ class Llama:
|
|||
# Save logits
|
||||
rows = n_tokens if self.params.logits_all else 1
|
||||
cols = self._n_vocab
|
||||
offset = 0 if self.params.logits_all else n_tokens - 1 # NOTE: Only save the last token logits if logits_all is False
|
||||
self.scores[self.n_tokens + offset: self.n_tokens + n_tokens, :].reshape(-1)[:] = llama_cpp.llama_get_logits(self.ctx)[:rows * cols]
|
||||
offset = (
|
||||
0 if self.params.logits_all else n_tokens - 1
|
||||
) # NOTE: Only save the last token logits if logits_all is False
|
||||
self.scores[self.n_tokens + offset : self.n_tokens + n_tokens, :].reshape(
|
||||
-1
|
||||
)[:] = llama_cpp.llama_get_logits(self.ctx)[: rows * cols]
|
||||
# Update n_tokens
|
||||
self.n_tokens += n_tokens
|
||||
|
||||
|
|
Loading…
Reference in a new issue