| 9 | |
| 10 | |
| 11 | class CompositionDataset(Dataset): |
| 12 | def __init__(self, file_path, vocab): |
| 13 | with open(file_path, 'r') as f: |
| 14 | self.data = json.load(f) |
| 15 | self.vocab = vocab |
| 16 | self.token_pattern = re.compile(r'<[^>]+>') |
| 17 | |
| 18 | def tokenize(self, text): |
| 19 | return self.token_pattern.findall(text) |
| 20 | |
| 21 | def tokens_to_ids(self, tokens): |
| 22 | return [self.vocab[token] for token in tokens] |
| 23 | |
| 24 | def __len__(self): |
| 25 | return len(self.data) |
| 26 | |
| 27 | def __getitem__(self, idx): |
| 28 | sample = self.data[idx] |
| 29 | |
| 30 | input_tokens = self.tokenize(sample['input_text']) |
| 31 | target_tokens = self.tokenize(sample['target_text']) |
| 32 | |
| 33 | if target_tokens and target_tokens[-1] == "</a>": |
| 34 | target_tokens = target_tokens[:-1] |
| 35 | |
| 36 | input_ids = torch.tensor(self.tokens_to_ids(input_tokens), dtype=torch.long) |
| 37 | target_ids = torch.tensor(self.tokens_to_ids(target_tokens), dtype=torch.long) |
| 38 | |
| 39 | return input_ids, target_ids |
| 40 | |
| 41 | |
| 42 | class CompositionTestDataset(Dataset): |
no outgoing calls
no test coverage detected