Skip to content

Commit 5cae104

Browse files
feat: Improve Llama.eval performance by avoiding list conversion (abetlen#1476)
Co-authored-by: Andrei <[email protected]>
1 parent 087cc0b commit 5cae104

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

llama_cpp/llama.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -562,12 +562,12 @@ def eval(self, tokens: Sequence[int]):
562562
if self.context_params.logits_all:
563563
rows = n_tokens
564564
cols = self._n_vocab
565-
logits = self._ctx.get_logits()[: rows * cols]
565+
logits = np.ctypeslib.as_array(self._ctx.get_logits(), shape=(rows * cols, ))
566566
self.scores[n_past : n_past + n_tokens, :].reshape(-1)[: :] = logits
567567
else:
568568
rows = 1
569569
cols = self._n_vocab
570-
logits = self._ctx.get_logits()[: rows * cols]
570+
logits = np.ctypeslib.as_array(self._ctx.get_logits(), shape=(rows * cols, ))
571571
self.scores[n_past + n_tokens - 1, :].reshape(-1)[: :] = logits
572572
# Update n_tokens
573573
self.n_tokens += n_tokens

0 commit comments

Comments
 (0)