| 514 | self.batch.logits[n_tokens - 1] = True |
| 515 | |
| 516 | def add_sequence(self, batch: Sequence[int], seq_id: int, logits_all: bool): |
| 517 | n_tokens = len(batch) |
| 518 | n_tokens0 = self.batch.n_tokens |
| 519 | self.batch.n_tokens += n_tokens |
| 520 | for i in range(n_tokens): |
| 521 | j = n_tokens0 + i |
| 522 | self.batch.token[j] = batch[i] |
| 523 | self.batch.pos[j] = i |
| 524 | self.batch.seq_id[j][0] = seq_id |
| 525 | self.batch.n_seq_id[j] = 1 |
| 526 | self.batch.logits[j] = logits_all |
| 527 | self.batch.logits[n_tokens0 + n_tokens - 1] = True |
| 528 | |
| 529 | |
| 530 | class LlamaTokenDataArray: |