| 5 | from seq2struct.utils import registry |
| 6 | |
| 7 | class ZippedDataset(torch.utils.data.Dataset): |
| 8 | def __init__(self, *components): |
| 9 | assert len(components) >= 1 |
| 10 | lengths = [len(c) for c in components] |
| 11 | assert all( |
| 12 | lengths[0] == other for other in lengths[1:]), "Lengths don't match: {}".format(lengths) |
| 13 | self.components = components |
| 14 | |
| 15 | def __getitem__(self, idx): |
| 16 | return tuple(c[idx] for c in self.components) |
| 17 | |
| 18 | def __len__(self): |
| 19 | return len(self.components[0]) |
| 20 | |
| 21 | |
| 22 | @registry.register('model', 'EncDec') |