From d788fb49bf1ff2f41e651ded4e5b788f2185caad Mon Sep 17 00:00:00 2001 From: samfundev Date: Sat, 24 Jun 2023 15:51:46 -0400 Subject: [PATCH] Only concatenate after all batches are done --- llama_cpp/llama.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py index a0b2030..d367601 100644 --- a/llama_cpp/llama.py +++ b/llama_cpp/llama.py @@ -405,6 +405,7 @@ class Llama: """ assert self.ctx is not None n_ctx = self._n_ctx + scores = [] for i in range(0, len(tokens), self.n_batch): batch = tokens[i : min(len(tokens), i + self.n_batch)] n_past = min(n_ctx - len(batch), len(self._input_ids)) @@ -430,9 +431,8 @@ class Llama: logits_view = llama_cpp.llama_get_logits(self.ctx) logits = [logits_view[i * cols : (i + 1) * cols] for i in range(rows)] self.eval_logits.extend(logits) - self._scores: npt.NDArray[np.single] = np.concatenate( - (self._scores, np.array(logits, dtype=np.single)), axis=0 - ) + scores.append(np.array(logits, dtype=np.single)) + self._scores = np.concatenate(scores) def _sample( self,