| |
| |
| |
| |
| """ |
| TODO (huxu): fairseq wrapper class for all dataset you defined: mostly MMDataset. |
| """ |
|
|
| from collections import OrderedDict |
|
|
| from torch.utils.data import Dataset |
| from torch.utils.data.dataloader import default_collate |
| from fairseq.data import FairseqDataset, data_utils |
|
|
|
|
| class FairseqMMDataset(FairseqDataset): |
| """ |
| A wrapper class for MMDataset for fairseq. |
| """ |
|
|
| def __init__(self, mmdataset): |
| if not isinstance(mmdataset, Dataset): |
| raise TypeError("mmdataset must be of type `torch.utils.data.dataset`.") |
| self.mmdataset = mmdataset |
|
|
| def set_epoch(self, epoch, **unused): |
| super().set_epoch(epoch) |
| self.epoch = epoch |
|
|
| def __getitem__(self, idx): |
| with data_utils.numpy_seed(43211, self.epoch, idx): |
| return self.mmdataset[idx] |
|
|
| def __len__(self): |
| return len(self.mmdataset) |
|
|
| def collater(self, samples): |
| if hasattr(self.mmdataset, "collator"): |
| return self.mmdataset.collator(samples) |
| if len(samples) == 0: |
| return {} |
| if isinstance(samples[0], dict): |
| batch = OrderedDict() |
| for key in samples[0]: |
| if samples[0][key] is not None: |
| batch[key] = default_collate([sample[key] for sample in samples]) |
| return batch |
| else: |
| return default_collate(samples) |
|
|
| def size(self, index): |
| """dummy implementation: we don't use --max-tokens""" |
| return 1 |
|
|
| def num_tokens(self, index): |
| """dummy implementation: we don't use --max-tokens""" |
| return 1 |
|
|
