| 777 | |
| 778 | |
| 779 | class CustomizationDataset(torch.utils.data.Dataset): |
| 780 | def __init__(self, args, split, tokenizer): |
| 781 | task, data_dir = args.task.lower(), args.data_dir |
| 782 | self.max_src_length, self.max_tgt_length = args.src_seq_length, args.tgt_seq_length |
| 783 | self.split = split |
| 784 | self.tokenizer = tokenizer |
| 785 | self.mask_pad_token = args.mask_pad_token |
| 786 | self.no_block_position = args.no_block_position |
| 787 | self.task_mask = args.task_mask |
| 788 | if split == "train": |
| 789 | filename = "train" |
| 790 | elif split == "dev": |
| 791 | filename = "val" |
| 792 | elif split == "test": |
| 793 | filename = "test" |
| 794 | else: |
| 795 | raise NotImplementedError(split) |
| 796 | print_rank_0(f"Creating {task}-{split} dataset from {data_dir}") |
| 797 | self.dataset_name = split |
| 798 | source_texts, target_texts = [], [] |
| 799 | with open(os.path.join(data_dir, f"{filename}.source"), |
| 800 | encoding='utf-8') as file: |
| 801 | for line in file: |
| 802 | line = line.strip() |
| 803 | source_texts.append(line) |
| 804 | with open(os.path.join(data_dir, f"{filename}.target"), encoding='utf-8') as file: |
| 805 | for line in file: |
| 806 | line = line.strip() |
| 807 | target_texts.append(line) |
| 808 | self.examples, self.example_list = {}, [] |
| 809 | for idx, (source_text, target_text) in enumerate(zip(source_texts, target_texts)): |
| 810 | if (idx + 1) % 20000 == 0: |
| 811 | print_rank_0(f"Complete {idx + 1} examples") |
| 812 | guid = "%s-%s" % (split, idx) |
| 813 | meta = {"ref": target_text} |
| 814 | example = InputExample(guid=guid, text_a=source_text, text_b=target_text, meta=meta) |
| 815 | self.examples[guid] = example |
| 816 | self.example_list.append(example) |
| 817 | print_rank_0(f"Return {len(self.examples)} {split} examples") |
| 818 | |
| 819 | def __len__(self): |
| 820 | return len(self.example_list) |
| 821 | |
| 822 | def __getitem__(self, idx): |
| 823 | example = self.example_list[idx] |
| 824 | cls_id = self.tokenizer.get_command('ENC').Id |
| 825 | mask_token = 'sMASK' if self.task_mask else 'MASK' |
| 826 | mask_id = self.tokenizer.get_command(mask_token).Id |
| 827 | eos_id = self.tokenizer.get_command('eos').Id |
| 828 | pad_id = self.tokenizer.get_command('pad').Id |
| 829 | sop_id = self.tokenizer.get_command('sop').Id |
| 830 | eop_id = self.tokenizer.get_command('eop').Id |
| 831 | source_text, target_text = example.text_a, example.text_b |
| 832 | source_tokens = self.tokenizer.EncodeAsIds(source_text).tokenization |
| 833 | if len(source_tokens) + 3 > self.max_src_length: |
| 834 | source_tokens = source_tokens[-(self.max_src_length - 3):] |
| 835 | source_tokens = [cls_id] + source_tokens + [mask_id, eos_id] |
| 836 | context_length = len(source_tokens) |
no outgoing calls
no test coverage detected