| 175 | desc="Running tokenizer on dataset", |
| 176 | ) |
| 177 | def group_texts(examples): |
| 178 | # Concatenate all texts. |
| 179 | concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()} |
| 180 | total_length = len(concatenated_examples[list(examples.keys())[0]]) |
| 181 | # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can |
| 182 | # customize this part to your needs. |
| 183 | if total_length >= block_size: |
| 184 | total_length = (total_length // block_size) * block_size |
| 185 | # Split by chunks of max_len. |
| 186 | result = { |
| 187 | k: [t[i : i + block_size] for i in range(0, total_length, block_size)] |
| 188 | for k, t in concatenated_examples.items() |
| 189 | } |
| 190 | result["labels"] = result["input_ids"].copy() |
| 191 | return result |
| 192 | dataset = tokenized_datasets.map( |
| 193 | group_texts, |
| 194 | batched=True, |