Iterable dataset that returns constant length chunks of tokens from stream of text files. Args: tokenizer (Tokenizer): The processor used for proccessing the data. dataset (dataset.Dataset): Dataset with text files. infinite (bool): If True the iterat
| 128 | |
| 129 | |
| 130 | class ConstantLengthDataset(IterableDataset): |
| 131 | """ |
| 132 | Iterable dataset that returns constant length chunks of tokens from stream of text files. |
| 133 | Args: |
| 134 | tokenizer (Tokenizer): The processor used for proccessing the data. |
| 135 | dataset (dataset.Dataset): Dataset with text files. |
| 136 | infinite (bool): If True the iterator is reset after dataset reaches end else stops. |
| 137 | seq_length (int): Length of token sequences to return. |
| 138 | num_of_sequences (int): Number of token sequences to keep in buffer. |
| 139 | chars_per_token (int): Number of characters per token used to estimate number of tokens in text buffer. |
| 140 | """ |
| 141 | |
| 142 | def __init__( |
| 143 | self, |
| 144 | tokenizer, |
| 145 | dataset, |
| 146 | infinite=False, |
| 147 | seq_length=1024, |
| 148 | num_of_sequences=1024, |
| 149 | chars_per_token=3.6, |
| 150 | input_column_name="prompt", |
| 151 | output_column_name="completion" |
| 152 | ): |
| 153 | self.tokenizer = tokenizer |
| 154 | self.concat_token_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else args.eos_token_id |
| 155 | self.dataset = dataset |
| 156 | self.seq_length = seq_length |
| 157 | self.infinite = infinite |
| 158 | self.current_size = 0 |
| 159 | self.max_buffer_size = seq_length * chars_per_token * num_of_sequences |
| 160 | self.input_column_name = input_column_name |
| 161 | self.output_column_name = output_column_name |
| 162 | |
| 163 | def __iter__(self): |
| 164 | iterator = iter(self.dataset) |
| 165 | more_examples = True |
| 166 | while more_examples: |
| 167 | buffer, buffer_len = [], 0 |
| 168 | while True: |
| 169 | if buffer_len >= self.max_buffer_size: |
| 170 | break |
| 171 | try: |
| 172 | buffer.append(prepare_sample_text(next(iterator), self.input_column_name, self.output_column_name)) |
| 173 | buffer_len += len(buffer[-1]) |
| 174 | except StopIteration: |
| 175 | if self.infinite: |
| 176 | iterator = iter(self.dataset) |
| 177 | else: |
| 178 | more_examples = False |
| 179 | break |
| 180 | tokenized_inputs = self.tokenizer(buffer, truncation=False)["input_ids"] |
| 181 | all_token_ids = [] |
| 182 | for tokenized_input in tokenized_inputs: |
| 183 | all_token_ids.extend(tokenized_input + [self.concat_token_id]) |
| 184 | for i in range(0, len(all_token_ids), self.seq_length): |
| 185 | input_ids = all_token_ids[i : i + self.seq_length] |
| 186 | if len(input_ids) == self.seq_length: |
| 187 | self.current_size += 1 |