Only concatenate after all batches are done

This commit is contained in:
samfundev 2023-06-24 15:51:46 -04:00
parent b6f9388436
commit d788fb49bf
No known key found for this signature in database
GPG key ID: 90D590E2A224B85D

View file

@ -405,6 +405,7 @@ class Llama:
""" """
assert self.ctx is not None assert self.ctx is not None
n_ctx = self._n_ctx n_ctx = self._n_ctx
scores = []
for i in range(0, len(tokens), self.n_batch): for i in range(0, len(tokens), self.n_batch):
batch = tokens[i : min(len(tokens), i + self.n_batch)] batch = tokens[i : min(len(tokens), i + self.n_batch)]
n_past = min(n_ctx - len(batch), len(self._input_ids)) 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_view = llama_cpp.llama_get_logits(self.ctx)
logits = [logits_view[i * cols : (i + 1) * cols] for i in range(rows)] logits = [logits_view[i * cols : (i + 1) * cols] for i in range(rows)]
self.eval_logits.extend(logits) self.eval_logits.extend(logits)
self._scores: npt.NDArray[np.single] = np.concatenate( scores.append(np.array(logits, dtype=np.single))
(self._scores, np.array(logits, dtype=np.single)), axis=0 self._scores = np.concatenate(scores)
)
def _sample( def _sample(
self, self,