fairseq/data/multilingual/sampled_multi_dataset.py

# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from typing import List
from enum import Enum
from collections import OrderedDict
from collections import defaultdict
from bisect import bisect_right
import hashlib
import logging
import datetime
import time

import numpy as np
import torch

from fairseq import distributed_utils
from fairseq.data import plasma_utils, FairseqDataset


def get_time_gap(s, e):
    return (datetime.datetime.fromtimestamp(e) - datetime.datetime.fromtimestamp(s)).__str__()


logger = logging.getLogger(__name__)


def default_virtual_size_func(datasets, ratios, max_scale_up=1.5):
    sizes = [len(d) for d in datasets]
    if ratios is None:
        return sum(sizes)
    largest_idx = np.argmax(sizes)
    largest_r = ratios[largest_idx]
    largest_s = sizes[largest_idx]
    # set virtual sizes relative to the largest dataset
    virtual_sizes = [(r / largest_r) * largest_s for r in ratios]
    vsize = sum(virtual_sizes)
    max_size = sum(sizes) * max_scale_up
    return int(vsize if vsize < max_size else max_size)


class CollateFormat(Enum):
    single = 1
    ordered_dict = 2


class SampledMultiDataset(FairseqDataset):
    """Samples from multiple sub-datasets according to given sampling ratios.
    Args:
        datasets (
            List[~torch.utils.data.Dataset]
            or OrderedDict[str, ~torch.utils.data.Dataset]
        ): datasets
        sampling_ratios (List[float]): list of probability of each dataset to be sampled
            (default: None, which corresponds to concating all dataset together).
        batch_by_size (bool): whether or not to batch by sequence length
            (default: True).
        seed (int): RNG seed to use (default: 2).
        epoch (int): starting epoch number (default: 1).
        eval_key (str, optional): a key used at evaluation time that causes
            this instance to pass-through batches from *datasets[eval_key]*.
        collate_format (CollateFormat):  collater output format, either CollateFormat.ordered_dict or
            CollateFormat.single (default: CollateFormat.single) where CollateFormat.single configures
            the collater to output batches of data mixed from all sub-datasets,
            and CollateFormat.ordered_dict configures the collater to output a dictionary of batches indexed by keys
            of sub-datasets.
            Note that not all sub-datasets will present in a single batch in both formats.
        virtual_size (int, or callable): the expected virtual size of the dataset (default: default_virtual_size_func).
        split (str): the split of the data, e.g. 'train', 'valid' or 'test'.
        shared_collater (bool): whether or not to all sub-datasets have the same collater.
    """

    def __init__(
            self,
            datasets,
            sampling_ratios=None,
            batch_by_size=False,
            seed=2,
            epoch=1,
            eval_key=None,
            collate_format=CollateFormat.single,
            virtual_size=default_virtual_size_func,
            split='',
            shared_collater=False,
    ):
        super().__init__()
        self.batch_by_size = batch_by_size
        self.shared_collater = shared_collater

        if isinstance(datasets, OrderedDict):
            self.keys = list(datasets.keys())
            datasets = list(datasets.values())
        elif isinstance(datasets, List):
            self.keys = list(range(len(datasets)))
        else:
            raise AssertionError()
        self.datasets = datasets
        self.split = split

        self.eval_key = eval_key
        if self.eval_key is not None:
            self.collate_format = CollateFormat.single
        else:
            self.collate_format = collate_format

        self.seed = seed
        self._cur_epoch = None
        self._cur_indices = None
        self._sizes = None
        self._ordered_indices = None
        self.virtual_size_per_dataset = None
        # caching properties
        self._reset_cached_properties()
        self.setup_sampling(sampling_ratios, virtual_size)
        self.cumulated_sizes = None
        self.virtual_size_per_dataset = None
        self._size_cache = {}
        self.set_epoch(epoch)

    def _clean_if_not_none(self, var_list):
        for v in var_list:
            if v is not None:
                del v

    def _reset_cached_properties(self):
        self._clean_if_not_none([
            self._sizes, self._ordered_indices, self._cur_indices
        ])
        self._sizes = None
        self._ordered_indices = None
        self._cur_indices = None

    def setup_sampling(self, sample_ratios, virtual_size):
        sizes = [len(d) for d in self.datasets]
        if sample_ratios is None:
            # default back to concating datasets
            self.sample_ratios = None
            self.virtual_size = sum(sizes)
        else:
            if not isinstance(sample_ratios, np.ndarray):
                sample_ratios = np.array(sample_ratios)
            self.sample_ratios = plasma_utils.PlasmaArray(sample_ratios)
            virtual_size = default_virtual_size_func if virtual_size is None else virtual_size
            self.virtual_size = (
                virtual_size(self.datasets, self.sample_ratios.array) if callable(virtual_size)
                else virtual_size)

    def adjust_sampling(self, epoch, sampling_ratios, virtual_size):
        if sampling_ratios is not None:
            sampling_ratios = self._sync_sample_ratios(sampling_ratios)
            self.setup_sampling(sampling_ratios, virtual_size)

    def _sync_sample_ratios(self, ratios):
        # in case the ratios are not precisely the same across processes
        # also to ensure every procresses update the ratios in the same pace
        ratios = torch.DoubleTensor(ratios)
        if torch.distributed.is_initialized():
            if torch.cuda.is_available():
                distributed_utils.all_reduce(ratios.cuda())
            else:
                distributed_utils.all_reduce(ratios)
            ret = ratios.cpu()
            ret = ret.numpy()
        return ret

    def random_choice_in_dataset(self, rng, dataset, choice_size):
        if hasattr(dataset, 'random_choice_in_dataset'):
            return dataset.random_choice_in_dataset(rng, choice_size)
        dataset_size = len(dataset)
        return rng.choice(dataset_size, choice_size, replace=(choice_size > dataset_size))

    def get_virtual_indices(self, rng, datasets, sample_ratios, virtual_size):
        def get_counts(sample_ratios):
            counts = np.array([virtual_size * r for r in sample_ratios], dtype=np.int64)
            diff = virtual_size - counts.sum()
            assert diff >= 0
            # due to round-offs, the size might not match the desired sizes
            if diff > 0:
                dataset_indices = rng.choice(len(sample_ratios), size=diff, p=sample_ratios)
                for i in dataset_indices:
                    counts[i] += 1
            return counts

        def get_in_dataset_indices(datasets, sizes, sample_ratios):
            counts = get_counts(sample_ratios)
            # uniformally sample desired counts for each dataset
            # if the desired counts are large, sample with replacement:
            indices = [
                self.random_choice_in_dataset(rng, d, c)
                for c, d in zip(counts, datasets)]
            return indices

        sizes = [len(d) for d in datasets]
        if sample_ratios is None:
            # default back to concating datasets
            in_dataset_indices = [list(range(s)) for s in sizes]
            virtual_sizes_per_dataset = sizes
        else:
            sample_ratios = sample_ratios.array
            ratios = sample_ratios / sample_ratios.sum()
            in_dataset_indices = get_in_dataset_indices(datasets, sizes, ratios)
            virtual_sizes_per_dataset = [len(d) for d in in_dataset_indices]
        virtual_sizes_per_dataset = np.array(virtual_sizes_per_dataset, np.int64)
        cumulative_sizes = np.cumsum(virtual_sizes_per_dataset)
        assert sum(virtual_sizes_per_dataset) == virtual_size
        assert cumulative_sizes[-1] == virtual_size
        if virtual_size < sum(sizes):
            logger.warning(
                f'virtual data size ({virtual_size}) is less than real data size ({sum(sizes)}).'
                ' If virtual size << real data size, there could be data coverage issue.'
            )
        in_dataset_indices = np.hstack(in_dataset_indices)
        return in_dataset_indices, cumulative_sizes, virtual_sizes_per_dataset

    def _get_dataset_and_index(self, index):
        i = bisect_right(self.cumulated_sizes.array, index)
        return i, self._cur_indices.array[index]

    def __getitem__(self, index):
        ds_idx, ds_sample_idx = self._get_dataset_and_index(index)
        ret = (ds_idx, self.datasets[ds_idx][ds_sample_idx])
        return ret

    def num_tokens(self, index):
        ds_idx, ds_sample_idx = self._get_dataset_and_index(index)
        return self.datasets[ds_idx].num_tokens(ds_sample_idx)

    def size(self, index):
        if self._sizes is not None:
            return self._sizes[index]
        ds_idx, ds_sample_idx = self._get_dataset_and_index(index)
        return self.datasets[ds_idx].size(ds_sample_idx)

    def __len__(self):
        return self.virtual_size

    def collater(self, samples, **extra_args):
        """Merge a list of samples to form a mini-batch."""
        if len(samples) == 0:
            return None
        if self.collate_format == 'ordered_dict':
            collect_samples = [[] for _ in range(len(self.datasets))]
            for (i, sample) in samples:
                collect_samples[i].append(sample)
            return OrderedDict([
                (self.keys[i], dataset.collater(collect_samples[i]))
                for i, (key, dataset) in enumerate(zip(self.keys, self.datasets))
                if len(collect_samples[i]) > 0
            ])
        elif self.shared_collater:
            return self.datasets[0].collater(
                [s for _, s in samples]
            )
        else:
            samples_dict = defaultdict(list)
            pad_to_length = defaultdict(int) if 'pad_to_length' not in extra_args else extra_args['pad_to_length']
            for ds_idx, s in samples:
                pad_to_length['source'] = max(pad_to_length['source'], s['source'].size(0))
                if s['target'] is not None:
                    pad_to_length['target'] = max(pad_to_length['target'], s['target'].size(0))
                samples_dict[ds_idx].append(s)
            batches = [
                self.datasets[i].collater(samples_dict[i], pad_to_length=pad_to_length)
                for i in range(len(self.datasets))
                if len(samples_dict[i]) > 0
            ]

            def straight_data(tensors):
                batch = torch.cat(tensors, dim=0)
                return batch

            src_lengths = straight_data([b['net_input']['src_lengths'] for b in batches])
            src_lengths, sort_order = src_lengths.sort(descending=True)

            def straight_order(tensors):
                batch = straight_data(tensors)
                return batch.index_select(0, sort_order)

            batch = {
                'id': straight_order([b['id'] for b in batches]),
                'nsentences': sum(b['nsentences'] for b in batches),
                'ntokens': sum(b['ntokens'] for b in batches),
                'net_input': {
                    'src_tokens': straight_order([b['net_input']['src_tokens'] for b in batches]),
                    'src_lengths': src_lengths,
                },
                'target': straight_order([b['target'] for b in batches]) if batches[0]['target'] is not None else None,
            }
            if 'prev_output_tokens' in batches[0]['net_input']:
                batch['net_input']['prev_output_tokens'] = straight_order(
                    [b['net_input']['prev_output_tokens'] for b in batches])
            if 'src_lang_id' in batches[0]['net_input']:
                batch['net_input']['src_lang_id'] = straight_order([b['net_input']['src_lang_id'] for b in batches])
            if 'tgt_lang_id' in batches[0]:
                batch['tgt_lang_id'] = straight_order([b['tgt_lang_id'] for b in batches])
            return batch

    @property
    def sizes(self):
        if self._sizes is not None:
            return self._sizes
        start_time = time.time()
        size_cache = self._size_cache
        ret = []
        for i in range(len(self)):
            ds_idx, ds_sample_idx = self._get_dataset_and_index(i)
            if (ds_idx, ds_sample_idx) in size_cache:
                ret.append(size_cache[(ds_idx, ds_sample_idx)])
            else:
                s = self.datasets[ds_idx].size(ds_sample_idx)
                size_cache[(ds_idx, ds_sample_idx)] = s
                ret.append(s)
        logger.debug(f'sizes() calling time: {get_time_gap(start_time, time.time())}')
        self._sizes = np.array(ret, np.int64)
        return self._sizes

    def ordered_indices(self):
        if self._ordered_indices is not None:
            return self._ordered_indices

        if self.batch_by_size:
            # No need to do shuffle as the data items are already randomized
            indices = np.arange(len(self))
            sizes = self.sizes
            tgt_sizes = sizes[:, 1] if len(sizes.shape) > 0 and sizes.shape[1] > 1 else None
            src_sizes = sizes[:, 0] if len(sizes.shape) > 0 and sizes.shape[1] > 1 else sizes

            # sort by target length, then source length
            if tgt_sizes is not None:
                indices = indices[
                    np.argsort(tgt_sizes[indices], kind='mergesort')
                ]
            sort_indices = indices[np.argsort(src_sizes[indices], kind='mergesort')]
        else:
            sort_indices = np.arange(len(self))
        self._ordered_indices = sort_indices
        return sort_indices

    def prefetch(self, indices):
        prefetch_indices = [[] for _ in range(len(self.datasets))]
        for i in indices:
            ds_idx, ds_sample_idx = self._get_dataset_and_index(i)
            prefetch_indices[ds_idx].append(ds_sample_idx)
        for i in range(len(prefetch_indices)):
            self.datasets[i].prefetch(prefetch_indices[i])

    def set_epoch(self, epoch):
        super().set_epoch(epoch)
        if epoch == self._cur_epoch:
            # re-enter so return
            return
        for d in self.datasets:
            if hasattr(d, 'set_epoch'):
                d.set_epoch(epoch)
        self._cur_epoch = epoch
        self._establish_virtual_datasets()

    def _establish_virtual_datasets(self):
        if self.sample_ratios is None and self._cur_indices is not None:
            # not a samping dataset, no need to resample if indices are already established
            return
        self._reset_cached_properties()

        start_time = time.time()
        # Generate a weighted sample of indices as a function of the
        # random seed and the current epoch.
        rng = np.random.RandomState(
           [
               int(hashlib.sha1(str(self.__class__.__name__).encode('utf-8')).hexdigest(), 16) % (2 ** 32),
               self.seed % (2 ** 32),  # global seed
               self._cur_epoch,  # epoch index,
           ]
        )
        indices, cumulated_sizes, virtual_size_per_dataset = self.get_virtual_indices(
            rng, self.datasets, self.sample_ratios, self.virtual_size)

        self._clean_if_not_none([
            self.cumulated_sizes, self.virtual_size_per_dataset
        ])
        self._cur_indices = plasma_utils.PlasmaArray(indices)
        self.cumulated_sizes = plasma_utils.PlasmaArray(cumulated_sizes)
        self.virtual_size_per_dataset = plasma_utils.PlasmaArray(virtual_size_per_dataset)

        raw_sizes = [len(d) for d in self.datasets]
        sampled_sizes = self.virtual_size_per_dataset.array
        logger.info(f'[{self.split}] Raw sizes: {str(dict(zip(self.keys, raw_sizes)))}; '
                    f'raw total size: {sum(raw_sizes)}')
        logger.info(f'[{self.split}] Resampled sizes: {str(dict(zip(self.keys, sampled_sizes)))}; '
                    f'resampled total size: {sum(sampled_sizes)}')
        if self.sample_ratios is not None:
            logger.info(f'[{self.split}] Upsampling ratios: {str(dict(zip(self.keys, self.sample_ratios.array)))}')
        else:
            logger.info(f'[{self.split}] A concat dataset')
        logger.debug(f'[{self.split}] virtual dataset established time: {get_time_gap(start_time, time.time())}')