(self)
| 51 | self._epoch = 0 |
| 52 | |
| 53 | def __iter__(self): |
| 54 | g = torch.Generator() |
| 55 | g.manual_seed(self._epoch + self.rank) |
| 56 | if self.rng_state is not None: |
| 57 | g.set_state(self.rng_state) |
| 58 | |
| 59 | rand_it = self.randint(0, self.buffer_size, g=g) |
| 60 | if self.states is not None: |
| 61 | self.data.load_state_dict(self.states) |
| 62 | |
| 63 | # max number of tokens allowed in the chunk buffer |
| 64 | n_tokens = self.buffer_size * self.context_len |
| 65 | |
| 66 | while True: |
| 67 | for sample in self.tokenize(self.data): |
| 68 | # keep appending the samples to the token buffer |
| 69 | self.tokens += sample |
| 70 | # if the token buffer is full, start sampling |
| 71 | # NOTE: we first convert the token ids to a tensor of shape [n_chunks, context_len] for efficiency |
| 72 | if len(self.buffer) == 0 and len(self.tokens) >= n_tokens: |
| 73 | self.buffer = torch.tensor(self.tokens[:n_tokens], dtype=self.dtype).view(self.buffer_size, -1) |
| 74 | self.tokens = self.tokens[n_tokens:] |
| 75 | if len(self.buffer) == self.buffer_size: |
| 76 | yield from self.sample(rand_it) |
| 77 | |
| 78 | n_chunks = len(self.tokens) // self.context_len |
| 79 | # handle the left tokens in the buffer |
| 80 | if n_chunks > 0: |
| 81 | n_tokens = n_chunks * self.context_len |
| 82 | indices = torch.randperm(n_chunks, generator=g).tolist() |
| 83 | self.buffer = torch.tensor(self.tokens[:n_tokens], dtype=torch.long).view(n_chunks, -1) |
| 84 | self.tokens = self.tokens[n_tokens:] |
| 85 | for i in indices: |
| 86 | yield {'input_ids': self.buffer[i]} |
| 87 | |
| 88 | def tokenize(self, data, batch_size: int = 64): |
| 89 | texts, states = [], [] |
nothing calls this directly
no test coverage detected