Only concatenate after all batches are done

samfundev · samfundev · commit d788fb49bf1f · 2023-06-24T15:51:46.000-04:00
diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py
@@ -405,6 +405,7 @@ def eval(self, tokens: Sequence[int]):
         """
         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 @@ def eval(self, tokens: Sequence[int]):
             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,
