| 52 | return dataset |
| 53 | |
| 54 | def _load(self, split): |
| 55 | dataset = datasets.load_dataset(os.path.join(self.data_dir, self.task_name, f'{self.task_name}.py'), split=split) |
| 56 | if self.return_source_length: |
| 57 | dataset = dataset.map(partial(self.add_original_src_length, tokenizer=self.tokenizer)) |
| 58 | dataset = dataset.map(self.add_prompt) |
| 59 | dataset = dataset.map(partial(self.convert_to_features, tokenizer=self.tokenizer), batched=True) |
| 60 | print(f'Example in {split} set:') |
| 61 | print(dataset[0]) |
| 62 | return dataset |
| 63 | |
| 64 | def add_original_src_length(self, example, tokenizer): |
| 65 | return { |