| | |
| | |
| | |
| | |
| |
|
| | from collections import OrderedDict |
| |
|
| | import torch |
| | from torch.utils.data.dataloader import default_collate |
| |
|
| | from . import FairseqDataset |
| |
|
| |
|
| | def _flatten(dico, prefix=None): |
| | """Flatten a nested dictionary.""" |
| | new_dico = OrderedDict() |
| | if isinstance(dico, dict): |
| | prefix = prefix + "." if prefix is not None else "" |
| | for k, v in dico.items(): |
| | if v is None: |
| | continue |
| | new_dico.update(_flatten(v, prefix + k)) |
| | elif isinstance(dico, list): |
| | for i, v in enumerate(dico): |
| | new_dico.update(_flatten(v, prefix + ".[" + str(i) + "]")) |
| | else: |
| | new_dico = OrderedDict({prefix: dico}) |
| | return new_dico |
| |
|
| |
|
| | def _unflatten(dico): |
| | """Unflatten a flattened dictionary into a nested dictionary.""" |
| | new_dico = OrderedDict() |
| | for full_k, v in dico.items(): |
| | full_k = full_k.split(".") |
| | node = new_dico |
| | for k in full_k[:-1]: |
| | if k.startswith("[") and k.endswith("]"): |
| | k = int(k[1:-1]) |
| | if k not in node: |
| | node[k] = OrderedDict() |
| | node = node[k] |
| | node[full_k[-1]] = v |
| | return new_dico |
| |
|
| |
|
| | class NestedDictionaryDataset(FairseqDataset): |
| | def __init__(self, defn, sizes=None): |
| | super().__init__() |
| | self.defn = _flatten(defn) |
| | self.sizes = [sizes] if not isinstance(sizes, (list, tuple)) else sizes |
| |
|
| | first = None |
| | for v in self.defn.values(): |
| | if not isinstance( |
| | v, |
| | ( |
| | FairseqDataset, |
| | torch.utils.data.Dataset, |
| | ), |
| | ): |
| | raise ValueError("Expected Dataset but found: {}".format(v.__class__)) |
| | first = first or v |
| | if len(v) > 0: |
| | assert len(v) == len(first), "dataset lengths must match" |
| |
|
| | self._len = len(first) |
| |
|
| | def __getitem__(self, index): |
| | return OrderedDict((k, ds[index]) for k, ds in self.defn.items()) |
| |
|
| | def __len__(self): |
| | return self._len |
| |
|
| | def collater(self, samples): |
| | """Merge a list of samples to form a mini-batch. |
| | |
| | Args: |
| | samples (List[dict]): samples to collate |
| | |
| | Returns: |
| | dict: a mini-batch suitable for forwarding with a Model |
| | """ |
| | if len(samples) == 0: |
| | return {} |
| | sample = OrderedDict() |
| | for k, ds in self.defn.items(): |
| | try: |
| | sample[k] = ds.collater([s[k] for s in samples]) |
| | except NotImplementedError: |
| | sample[k] = default_collate([s[k] for s in samples]) |
| | return _unflatten(sample) |
| |
|
| | def num_tokens(self, index): |
| | """Return the number of tokens in a sample. This value is used to |
| | enforce ``--max-tokens`` during batching.""" |
| | return max(s[index] for s in self.sizes) |
| |
|
| | def size(self, index): |
| | """Return an example's size as a float or tuple. This value is used when |
| | filtering a dataset with ``--max-positions``.""" |
| | if len(self.sizes) == 1: |
| | return self.sizes[0][index] |
| | else: |
| | return (s[index] for s in self.sizes) |
| |
|
| | @property |
| | def supports_prefetch(self): |
| | """Whether this dataset supports prefetching.""" |
| | return any(ds.supports_prefetch for ds in self.defn.values()) |
| |
|
| | def prefetch(self, indices): |
| | """Prefetch the data required for this epoch.""" |
| | for ds in self.defn.values(): |
| | if getattr(ds, "supports_prefetch", False): |
| | ds.prefetch(indices) |
| |
|
| | @property |
| | def can_reuse_epoch_itr_across_epochs(self): |
| | return all(ds.can_reuse_epoch_itr_across_epochs for ds in self.defn.values()) |
| |
|
| | def set_epoch(self, epoch): |
| | super().set_epoch(epoch) |
| | for ds in self.defn.values(): |
| | ds.set_epoch(epoch) |
| |
|
