| 148 | |
| 149 | |
| 150 | def format_dataset(dataset, dataset_format): |
| 151 | if ( |
| 152 | dataset_format == 'alpaca' or dataset_format == 'alpaca-clean' or |
| 153 | (dataset_format is None and args.dataset in ['alpaca', 'alpaca-clean']) |
| 154 | ): |
| 155 | dataset = dataset.map(extract_alpaca_dataset, remove_columns=['instruction']) |
| 156 | elif dataset_format == 'oasst1' or (dataset_format is None and args.dataset == 'oasst1'): |
| 157 | dataset = dataset.map(lambda x: { |
| 158 | 'input': '', |
| 159 | 'output': x['text'], |
| 160 | }) |
| 161 | elif dataset_format == 'pt' or (dataset_format is None and args.dataset in ['c4', 'redpajama']): |
| 162 | block_size = args.pt_context_len |
| 163 | column_names = list(dataset["train"].features) |
| 164 | text_column_name = "text" if "text" in column_names else column_names[0] |
| 165 | |
| 166 | def tokenize_function(examples): |
| 167 | output = tokenizer(examples[text_column_name]) |
| 168 | return output |
| 169 | tokenized_datasets = dataset.map( |
| 170 | tokenize_function, |
| 171 | batched=True, |
| 172 | remove_columns=column_names, |
| 173 | num_proc=args.preprocessing_num_workers, |
| 174 | load_from_cache_file=not args.overwrite_cache, |
| 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, |
| 195 | num_proc=args.preprocessing_num_workers, |
| 196 | load_from_cache_file=not args.overwrite_cache, |
| 197 | desc=f"Grouping texts in chunks of {block_size}", |
| 198 | ) |
| 199 | # Remove unused columns for instruction-tuning |
| 200 | if not dataset_format == 'pt': |
| 201 | dataset = dataset.remove_columns( |
| 202 | [col for col in dataset.column_names['train'] if col not in ['input', 'output']] |
| 203 | ) |
| 204 | return dataset |
| 205 | |
| 206 | # Load dataset. |
| 207 | print(f"loading {args.dataset}") |