| 73 | |
| 74 | # Input should be a dataset dict object , with text column name being 'text' |
| 75 | def tokenize_and_group_text(raw_datasets, tokenizer, block_size, stride, padding_index, num_process): |
| 76 | |
| 77 | def tokenize_function(examples): |
| 78 | output = tokenizer(examples["text"]) |
| 79 | return output |
| 80 | |
| 81 | def group_texts(examples): |
| 82 | # Concatenate all texts. |
| 83 | concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()} |
| 84 | total_length = len(concatenated_examples[list(examples.keys())[0]]) |
| 85 | if total_length >= block_size: |
| 86 | total_length = (total_length // block_size) * block_size |
| 87 | |
| 88 | input_ids = [] |
| 89 | attention_mask = [] |
| 90 | labels = [] |
| 91 | # We implement a sliding window, so all tokens have a non-zero context in their prediction. |
| 92 | # We then mask the duplicate tokens' labels, to not count any token twice in the loss. |
| 93 | for i in tqdm(range(0, total_length, stride), total=total_length): |
| 94 | begin_loc = max(i + stride - block_size, 0) |
| 95 | end_loc = min(i + stride, total_length) |
| 96 | trg_len = end_loc - i |
| 97 | cur_input_ids = concatenated_examples['input_ids'][begin_loc:end_loc] |
| 98 | cur_labels = list(cur_input_ids) |
| 99 | cur_labels[:-trg_len] = [padding_index] * (len(cur_labels) - trg_len) |
| 100 | |
| 101 | if len(cur_input_ids) < block_size: |
| 102 | padding_size = block_size - len(cur_input_ids) |
| 103 | pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id |
| 104 | cur_input_ids += [pad_token_id] * padding_size |
| 105 | cur_labels += [padding_index] * padding_size |
| 106 | input_ids.append(cur_input_ids) |
| 107 | attention_mask.append([1] * len(cur_labels)) |
| 108 | labels.append(cur_labels) |
| 109 | |
| 110 | result = {'input_ids': input_ids, 'labels': labels, 'attention_mask': attention_mask} |
| 111 | return result |
| 112 | |
| 113 | tokenized_datasets = raw_datasets.map( |
| 114 | tokenize_function, |
| 115 | batched=True, |
| 116 | num_proc=num_process, |
| 117 | remove_columns=raw_datasets["train"].column_names, |
| 118 | load_from_cache_file=True, |
| 119 | desc="Running tokenizer on dataset", |
| 120 | ) |
| 121 | |
| 122 | # calculate total tokens |
| 123 | for split in tokenized_datasets: |
| 124 | total_tokens = 0 |
| 125 | for example in tokenized_datasets[split]: |
| 126 | total_tokens += len(example['input_ids']) |
| 127 | logger.info(f"Total tokens in dataset split {split}: {total_tokens}") |
| 128 | |
| 129 | lm_datasets = tokenized_datasets.map( |
| 130 | group_texts, |
| 131 | batched=True, |
| 132 | num_proc=num_process, |