| 40 | |
| 41 | |
| 42 | class CompositionTestDataset(Dataset): |
| 43 | def __init__(self, file_path, vocab, split=None): |
| 44 | with open(file_path, 'r') as f: |
| 45 | self.data = json.load(f) |
| 46 | if split: |
| 47 | self.data = [d for d in self.data if d.get("type") == split] |
| 48 | self.vocab = vocab |
| 49 | self.token_pattern = re.compile(r'<[^>]+>') |
| 50 | |
| 51 | def tokenize(self, text): |
| 52 | return self.token_pattern.findall(text) |
| 53 | |
| 54 | def tokens_to_ids(self, tokens): |
| 55 | return [self.vocab[token] for token in tokens] |
| 56 | |
| 57 | def __len__(self): |
| 58 | return len(self.data) |
| 59 | |
| 60 | def __getitem__(self, idx): |
| 61 | sample = self.data[idx] |
| 62 | |
| 63 | input_tokens = self.tokenize(sample['input_text']) |
| 64 | target_tokens = self.tokenize(sample['target_text']) |
| 65 | test_type = sample['type'] |
| 66 | |
| 67 | if target_tokens and target_tokens[-1] == "</a>": |
| 68 | target_tokens = target_tokens[:-1] |
| 69 | |
| 70 | input_ids = torch.tensor(self.tokens_to_ids(input_tokens), dtype=torch.long) |
| 71 | target_ids = torch.tensor(self.tokens_to_ids(target_tokens), dtype=torch.long) |
| 72 | |
| 73 | return input_ids, target_ids, test_type |
| 74 | |
| 75 | |
| 76 | def custom_collate(batch): |
no outgoing calls
no test coverage detected