| 45 | |
| 46 | |
| 47 | class CompositionTestDataset(Dataset): |
| 48 | def __init__(self, data_dir, file, vocab, all_splits=None): |
| 49 | self.data = [] |
| 50 | with open(os.path.join(data_dir, file), 'r') as f: |
| 51 | data = json.load(f) |
| 52 | if all_splits: |
| 53 | for sample in data: |
| 54 | if sample['type'] in all_splits: |
| 55 | self.data.append(sample) |
| 56 | else: |
| 57 | self.data = data |
| 58 | self.vocab = vocab |
| 59 | self.token_pattern = re.compile(r'<[^>]+>') |
| 60 | |
| 61 | def tokenize(self, text): |
| 62 | return self.token_pattern.findall(text) |
| 63 | |
| 64 | def tokens_to_ids(self, tokens): |
| 65 | return [self.vocab[token] for token in tokens] |
| 66 | |
| 67 | def __len__(self): |
| 68 | return len(self.data) |
| 69 | |
| 70 | def __getitem__(self, idx): |
| 71 | sample = self.data[idx] |
| 72 | |
| 73 | input_tokens = self.tokenize(sample['input_text']) |
| 74 | target_tokens = self.tokenize(sample['target_text']) |
| 75 | test_type = sample['type'] |
| 76 | |
| 77 | # input_tokens.append('<sep>') |
| 78 | |
| 79 | if target_tokens and target_tokens[-1] == "</a>": |
| 80 | target_tokens = target_tokens[:-1] |
| 81 | |
| 82 | input_ids = torch.tensor(self.tokens_to_ids(input_tokens), dtype=torch.long) |
| 83 | target_ids = torch.tensor(self.tokens_to_ids(target_tokens), dtype=torch.long) |
| 84 | |
| 85 | return input_ids, target_ids, test_type |
| 86 | |
| 87 | |
| 88 | def custom_collate(batch, max_len=25): |
no outgoing calls
no test coverage detected