| 86 | yield {'input_ids': self.buffer[i]} |
| 87 | |
| 88 | def tokenize(self, data, batch_size: int = 64): |
| 89 | texts, states = [], [] |
| 90 | for sample in data: |
| 91 | texts.append(sample['text']) |
| 92 | states.append(self.data.state_dict()) |
| 93 | if len(texts) == batch_size: |
| 94 | for s, tokenized in zip(states, self.tokenizer(texts, return_attention_mask=False)['input_ids']): |
| 95 | self.states = s |
| 96 | yield tokenized |
| 97 | texts, states = [], [] |
| 98 | if len(texts) > 0: |
| 99 | for s, tokenized in zip(states, self.tokenizer(texts, return_attention_mask=False)['input_ids']): |
| 100 | self.states = s |
| 101 | yield tokenized |
| 102 | |
| 103 | def sample(self, indices): |
| 104 | n_tokens = (len(self.tokens) // self.context_len) * self.context_len |