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