wehub-resource-sync
255 lines
9.4 KiB
Python
255 lines
9.4 KiB
Python
# 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.
|
|
|
|
import logging
|
|
import time
|
|
from collections import OrderedDict
|
|
from typing import Dict, List, Optional
|
|
|
|
import numpy as np
|
|
from fairseq.data import data_utils
|
|
|
|
from . import FairseqDataset
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class MultiCorpusDataset(FairseqDataset):
|
|
"""
|
|
Stores multiple instances of FairseqDataset together.
|
|
Unless batch_sample=True, requires each instance
|
|
to be the same dataset, as the collate method needs to work on batches with
|
|
samples from each dataset.
|
|
|
|
Allows specifying a distribution over the datasets to use. Note that unlike
|
|
MultiCorpusSampledDataset, this distribution allows sampling for each item,
|
|
rather than on a batch level. Note that datasets with sampling probabilty
|
|
of 0 will be skipped.
|
|
|
|
Each time ordered_indices() is called, a new sample is generated with
|
|
the specified distribution.
|
|
|
|
Args:
|
|
datasets: a OrderedDict of FairseqDataset instances.
|
|
distribution: a List containing the probability of getting an utterance from
|
|
corresponding dataset
|
|
seed: random seed for sampling the datsets
|
|
sort_indices: if true, will sort the ordered indices by size
|
|
batch_sample: if true, will ensure each batch is from a single dataset
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
datasets: Dict[str, FairseqDataset],
|
|
distribution: List[float],
|
|
seed: int,
|
|
sort_indices: bool = False,
|
|
batch_sample: bool = False,
|
|
distributed_rank: Optional[int] = None,
|
|
):
|
|
super().__init__()
|
|
assert isinstance(datasets, OrderedDict)
|
|
assert len(datasets) == len(distribution)
|
|
assert sum(distribution) == 1
|
|
self.datasets = datasets
|
|
self.distribution = distribution
|
|
self.seed = seed
|
|
self.sort_indices = sort_indices
|
|
self.batch_sample = batch_sample
|
|
self.distributed_rank = distributed_rank
|
|
|
|
# Avoid repeated conversions to list later
|
|
self.dataset_list = list(datasets.values())
|
|
self.total_num_instances = 0
|
|
|
|
first_dataset = self.dataset_list[0]
|
|
|
|
self.num_instances_per_dataset = []
|
|
self.dataset_offsets = []
|
|
for i, dataset in enumerate(self.dataset_list):
|
|
assert isinstance(dataset, FairseqDataset)
|
|
assert type(dataset) is type(first_dataset)
|
|
self.num_instances_per_dataset.append(
|
|
0 if self.distribution[i] == 0 else len(dataset)
|
|
)
|
|
self.dataset_offsets.append(self.total_num_instances)
|
|
self.total_num_instances += self.num_instances_per_dataset[i]
|
|
|
|
def ordered_indices(self):
|
|
start = time.time()
|
|
with data_utils.numpy_seed(self.seed, self.epoch):
|
|
logger.info(f"sampling new dataset with seed {self.seed} epoch {self.epoch}")
|
|
sampled_indices = []
|
|
num_selected_instances = 0
|
|
|
|
# For each dataset i, sample self.distribution[i] * self.total_num_instances
|
|
for i, key in enumerate(self.datasets):
|
|
if self.distribution[i] == 0:
|
|
# skip dataset if sampling probability is 0
|
|
continue
|
|
|
|
if i < len(self.datasets) - 1:
|
|
num_instances = int(self.distribution[i] * self.total_num_instances)
|
|
high = self.dataset_offsets[i + 1]
|
|
else:
|
|
num_instances = self.total_num_instances - num_selected_instances
|
|
high = self.total_num_instances
|
|
|
|
logger.info(f"sampling {num_instances} from {key} dataset")
|
|
num_selected_instances += num_instances
|
|
|
|
# First, add k copies of the dataset where k = num_instances // len(dataset).
|
|
# This ensures an equal distribution of the data points as much as possible.
|
|
# For the remaining entries randomly sample them
|
|
dataset_size = len(self.datasets[key])
|
|
num_copies = num_instances // dataset_size
|
|
dataset_indices = (
|
|
np.random.permutation(high - self.dataset_offsets[i])
|
|
+ self.dataset_offsets[i]
|
|
)[: num_instances - num_copies * dataset_size]
|
|
if num_copies > 0:
|
|
sampled_indices += list(
|
|
np.concatenate(
|
|
(
|
|
np.repeat(
|
|
np.arange(self.dataset_offsets[i], high), num_copies
|
|
),
|
|
dataset_indices,
|
|
)
|
|
)
|
|
)
|
|
else:
|
|
sampled_indices += list(dataset_indices)
|
|
|
|
assert (
|
|
len(sampled_indices) == self.total_num_instances
|
|
), f"{len(sampled_indices)} vs {self.total_num_instances}"
|
|
|
|
np.random.shuffle(sampled_indices)
|
|
if self.sort_indices:
|
|
sampled_indices.sort(key=lambda i: self.num_tokens(i))
|
|
|
|
logger.info(
|
|
"multi_corpus_dataset ordered_indices took {}s".format(
|
|
time.time() - start
|
|
)
|
|
)
|
|
return np.array(sampled_indices, dtype=np.int64)
|
|
|
|
def _map_index(self, index: int):
|
|
"""
|
|
If dataset A has length N and dataset B has length M
|
|
then index 1 maps to index 1 of dataset A, and index N + 1
|
|
maps to index 1 of B.
|
|
"""
|
|
counter = 0
|
|
for num_instances, key in zip(self.num_instances_per_dataset, self.datasets):
|
|
if index < counter + num_instances:
|
|
return index - counter, key
|
|
counter += num_instances
|
|
raise ValueError(
|
|
"Invalid index: {}, max: {}".format(index, self.total_num_instances)
|
|
)
|
|
|
|
def __len__(self):
|
|
"""
|
|
Length of this dataset is the sum of individual datasets
|
|
"""
|
|
return self.total_num_instances
|
|
|
|
def __getitem__(self, index):
|
|
new_index, key = self._map_index(index)
|
|
try:
|
|
item = self.datasets[key][new_index]
|
|
item["full_id"] = index
|
|
return item
|
|
except Exception as e:
|
|
e.args = (f"Error from {key} dataset", *e.args)
|
|
raise
|
|
|
|
def collater(self, samples):
|
|
"""
|
|
If we are doing batch sampling, then pick the right collater to use.
|
|
|
|
Otherwise we assume all collaters are the same.
|
|
"""
|
|
if len(samples) == 0:
|
|
return None
|
|
if "full_id" in samples[0]:
|
|
_, key = self._map_index(samples[0]["full_id"])
|
|
try:
|
|
batch = self.datasets[key].collater(samples)
|
|
except Exception:
|
|
print(f"Collating failed for key {key}", flush=True)
|
|
raise
|
|
return batch
|
|
else:
|
|
# Subclasses may override __getitem__ to not specify full_id
|
|
return list(self.datasets.values())[0].collater(samples)
|
|
|
|
def num_tokens(self, index: int):
|
|
index, key = self._map_index(index)
|
|
return self.datasets[key].num_tokens(index)
|
|
|
|
def size(self, index: int):
|
|
index, key = self._map_index(index)
|
|
return self.datasets[key].size(index)
|
|
|
|
@property
|
|
def can_reuse_epoch_itr_across_epochs(self):
|
|
return False
|
|
|
|
def set_epoch(self, epoch, **unused):
|
|
super().set_epoch(epoch)
|
|
logger.info(f"setting epoch of multi_corpus_dataset to {epoch}")
|
|
self.epoch = epoch
|
|
|
|
@property
|
|
def supports_prefetch(self):
|
|
return False
|
|
|
|
@property
|
|
def supports_fetch_outside_dataloader(self):
|
|
return all(
|
|
self.datasets[key].supports_fetch_outside_dataloader
|
|
for key in self.datasets
|
|
)
|
|
|
|
def batch_by_size(
|
|
self,
|
|
indices,
|
|
max_tokens=None,
|
|
max_sentences=None,
|
|
required_batch_size_multiple=1,
|
|
):
|
|
if not self.batch_sample:
|
|
return super().batch_by_size(
|
|
indices, max_tokens, max_sentences, required_batch_size_multiple
|
|
)
|
|
|
|
dataset_indices = {key: [] for key in self.datasets}
|
|
for i in indices:
|
|
_, key = self._map_index(i)
|
|
dataset_indices[key].append(i)
|
|
|
|
batches = []
|
|
for key in dataset_indices:
|
|
cur_batches = super().batch_by_size(
|
|
np.array(dataset_indices[key], dtype=np.int64),
|
|
max_tokens,
|
|
max_sentences,
|
|
required_batch_size_multiple,
|
|
)
|
|
logger.info(f"Created {len(cur_batches)} batches for dataset {key}")
|
|
batches += cur_batches
|
|
|
|
# If this dataset is used in a distributed training setup,
|
|
# then shuffle such that the order is seeded by the distributed rank
|
|
# as well
|
|
if self.distributed_rank is not None:
|
|
with data_utils.numpy_seed(self.seed, self.epoch, self.distributed_rank):
|
|
np.random.shuffle(batches)
|
|
return batches
|