Clean up FairseqTask so that it's easier to extend/add new tasks (2e507d3c) · Commits · Simon Will / fairseq

eval_lm.py

+4 −2

Original line number	Diff line number	Diff line
		@@ -59,11 +59,13 @@ def main(parsed_args):

		assert len(models) > 0

		itr = data.EpochBatchIterator(
		itr = task.get_batch_iterator(
		dataset=task.dataset(args.gen_subset),
		max_tokens=args.max_tokens or 36000,
		max_sentences=args.max_sentences,
		max_positions=models[0].max_positions(),
		max_positions=utils.resolve_max_positions(*[
		model.max_positions() for model in models
		]),
		num_shards=args.num_shards,
		shard_id=args.shard_id,
		ignore_invalid_inputs=True,

fairseq/data/data_utils.py

+125 −66

Original line number	Diff line number	Diff line
		@@ -12,8 +12,6 @@ import os
		import numpy as np
		import torch

		from . import FairseqDataset


		def infer_language_pair(path):
		"""Infer language pair from filename: <split>.<lang1>-<lang2>.(...).idx"""
		@@ -99,42 +97,35 @@ def collate_tokens(values, pad_idx, eos_idx, left_pad, move_eos_to_beginning=Fal


		class EpochBatchIterator(object):
		"""Iterate over a FairseqDataset and yield batches bucketed by size.
		"""A multi-epoch iterator over a :class:`~torch.utils.data.Dataset`.

		Compared to :class:`~torch.utils.data.DataLoader`, this iterator:

		Batches may contain sequences of different lengths. This iterator can be
		reused across multiple epochs with the next_epoch_itr() method.
		- can be reused across multiple epochs with the :func:`next_epoch_itr`
		method (optionally shuffled between epochs)
		- can be serialized/deserialized with the :func:`state_dict` and
		:func:`load_state_dict` methods
		- supports sharding with the ``num_shards`` and ``shard_id`` arguments

		Args:
		dataset: a FairseqDataset
		max_tokens: max number of tokens in each batch
		max_sentences: max number of sentences in each batch
		max_positions: max sentence length supported by the model
		ignore_invalid_inputs: don't raise Exception for sentences that are too long
		required_batch_size_multiple: require batch size to be a multiple of N
		seed: seed for random number generator for reproducibility
		num_shards: shard the data iterator into N shards
		shard_id: which shard of the data iterator to return
		dataset (Dataset): dataset from which to load the data
		batch_sampler (Sampler): an iterator over batches of indices
		seed (int, optional): seed for random number generator for
		reproducibility. Default: ``1``
		num_shards (int, optional): shard the data iterator into N
		shards. Default: ``1``
		shard_id (int, optional): which shard of the data iterator to
		return. Default: ``0``
		"""

		def __init__(
		self, dataset, max_tokens=None, max_sentences=None, max_positions=None,
		ignore_invalid_inputs=False, required_batch_size_multiple=1, seed=1,
		num_shards=1, shard_id=0,
		):
		assert isinstance(dataset, FairseqDataset)
		def __init__(self, dataset, batch_sampler, seed=1, num_shards=1, shard_id=0):
		assert isinstance(dataset, torch.utils.data.Dataset)
		self.dataset = dataset
		self.max_tokens = max_tokens if max_tokens is not None else float('Inf')
		self.max_sentences = max_sentences if max_sentences is not None else float('Inf')
		self.max_positions = max_positions
		self.ignore_invalid_inputs = ignore_invalid_inputs
		self.bsz_mult = required_batch_size_multiple
		self.frozen_batches = tuple(batch_sampler)
		self.seed = seed
		self.num_shards = num_shards
		self.shard_id = shard_id

		with numpy_seed(self.seed):
		self.frozen_batches = tuple(self._batch_generator())

		self.epoch = 0
		self._cur_epoch_itr = None
		self._next_epoch_itr = None
		@@ -143,7 +134,13 @@ class EpochBatchIterator(object):
		return len(self.frozen_batches)

		def next_epoch_itr(self, shuffle=True):
		"""Shuffle batches and return a new iterator over the dataset."""
		"""
		Return a new iterator over the dataset.

		Args:
		shuffle (bool, optional): shuffle batches before returning the
		iterator. Default: ``True``
		"""
		if self._next_epoch_itr is not None:
		self._cur_epoch_itr = self._next_epoch_itr
		self._next_epoch_itr = None
		@@ -153,10 +150,12 @@ class EpochBatchIterator(object):
		return self._cur_epoch_itr

		def end_of_epoch(self):
		"""Returns whether the most recent epoch iterator has been exhausted"""
		return not self._cur_epoch_itr.has_next()

		@property
		def iterations_in_epoch(self):
		"""The number of consumed batches in the current epoch."""
		if self._cur_epoch_itr is not None:
		return self._cur_epoch_itr.count
		elif self._next_epoch_itr is not None:
		@@ -193,38 +192,119 @@ class EpochBatchIterator(object):
		batch_sampler=ShardedIterator(batches, self.num_shards, self.shard_id, fill_value=[]),
		))

		def _batch_generator(self):

		@contextlib.contextmanager
		def numpy_seed(seed):
		"""Context manager which seeds the NumPy PRNG with the specified seed and
		restores the state afterward"""
		if seed is None:
		yield
		return
		state = np.random.get_state()
		np.random.seed(seed)
		try:
		yield
		finally:
		np.random.set_state(state)


		def collect_filtered(function, iterable, filtered):
		"""
		Similar to :func:`filter` but collects filtered elements in ``filtered``.

		Args:
		function (callable): function that returns ``False`` for elements that
		should be filtered
		iterable (iterable): iterable to filter
		filtered (list): list to store filtered elements
		"""
		for el in iterable:
		if function(el):
		yield el
		else:
		filtered.append(el)


		def filter_by_size(indices, size_fn, max_positions, raise_exception=False):
		"""
		Filter indices based on their size.

		Args:
		indices (List[int]): ordered list of dataset indices
		size_fn (callable): function that returns the size of a given index
		max_positions (tuple): filter elements larger than this size.
		Comparisons are done component-wise.
		raise_exception (bool, optional): if ``True``, raise an exception
		if any elements are filtered. Default: ``False``
		"""
		def check_size(idx):
		if isinstance(max_positions, float) or isinstance(max_positions, int):
		return size_fn(idx) < max_positions
		else:
		return all(a <= b for a, b in zip(size_fn(idx), max_positions))

		ignored = []
		itr = collect_filtered(check_size, indices, ignored)
		for idx in itr:
		if len(ignored) > 0 and raise_exception:
		raise Exception((
		'Size of sample #{} is invalid (={}) since max_positions={}, '
		'skip this example with --skip-invalid-size-inputs-valid-test'
		).format(idx, self.size(idx), max_positions))
		yield idx

		if len(ignored) > 0:
		print((
		'\| WARNING: {} samples have invalid sizes and will be skipped, '
		'max_positions={}, first few sample ids={}'
		).format(len(ignored), max_positions, ignored[:10]))


		def batch_by_size(
		indices, num_tokens_fn, max_tokens=None, max_sentences=None,
		required_batch_size_multiple=1,
		):
		"""
		Yield mini-batches of indices bucketed by size. Batches may contain
		sequences of different lengths.

		Args:
		indices (List[int]): ordered list of dataset indices
		num_tokens_fn (callable): function that returns the number of tokens at
		a given index
		max_tokens (int, optional): max number of tokens in each batch.
		Default: ``None``
		max_sentences (int, optional): max number of sentences in each
		batch. Default: ``None``
		required_batch_size_multiple (int, optional): require batch size to
		be a multiple of N. Default: ``1``
		"""
		max_tokens = max_tokens if max_tokens is not None else float('Inf')
		max_sentences = max_sentences if max_sentences is not None else float('Inf')
		bsz_mult = required_batch_size_multiple

		batch = []

		def is_batch_full(num_tokens):
		if len(batch) == 0:
		return False
		if len(batch) == self.max_sentences:
		if len(batch) == max_sentences:
		return True
		if num_tokens > self.max_tokens:
		if num_tokens > max_tokens:
		return True
		return False

		sample_len = 0
		sample_lens = []
		ignored = []
		for idx in self.dataset.ordered_indices():
		if not self.dataset.valid_size(idx, self.max_positions):
		if self.ignore_invalid_inputs:
		ignored.append(idx)
		continue
		raise Exception((
		'Size of sample #{} is invalid, max_positions={}, skip this '
		'example with --skip-invalid-size-inputs-valid-test'
		).format(idx, self.max_positions))

		sample_lens.append(self.dataset.num_tokens(idx))
		for idx in indices:
		sample_lens.append(num_tokens_fn(idx))
		sample_len = max(sample_len, sample_lens[-1])
		num_tokens = (len(batch) + 1) * sample_len
		if is_batch_full(num_tokens):
		mod_len = max(
		self.bsz_mult * (len(batch) // self.bsz_mult),
		len(batch) % self.bsz_mult,
		bsz_mult * (len(batch) // bsz_mult),
		len(batch) % bsz_mult,
		)
		yield batch[:mod_len]
		batch = batch[mod_len:]
		@@ -235,24 +315,3 @@ class EpochBatchIterator(object):

		if len(batch) > 0:
		yield batch

		if len(ignored) > 0:
		print((
		'\| WARNING: {} samples have invalid sizes and will be skipped, '
		'max_positions={}, first few sample ids={}'
		).format(len(ignored), self.max_positions, ignored[:10]))


		@contextlib.contextmanager
		def numpy_seed(seed):
		"""Context manager which seeds the NumPy PRNG with the specified seed and
		restores the state afterward"""
		if seed is None:
		yield
		return
		state = np.random.get_state()
		np.random.seed(seed)
		try:
		yield
		finally:
		np.random.set_state(state)

fairseq/data/fairseq_dataset.py

+18 −6

Original line number	Diff line number	Diff line
		@@ -7,6 +7,8 @@

		import torch.utils.data

		from fairseq.data import data_utils


		class FairseqDataset(torch.utils.data.Dataset):
		"""A dataset that provides helpers for batching."""
		@@ -18,7 +20,14 @@ class FairseqDataset(torch.utils.data.Dataset):
		raise NotImplementedError

		def collater(self, samples):
		"""Merge a list of samples to form a mini-batch."""
		"""Merge a list of samples to form a mini-batch.

		Args:
		samples (List[int]): sample indices to collate

		Returns:
		dict: a mini-batch suitable for forwarding with a Model
		"""
		raise NotImplementedError

		def get_dummy_batch(self, num_tokens, max_positions):
		@@ -26,13 +35,16 @@ class FairseqDataset(torch.utils.data.Dataset):
		raise NotImplementedError

		def num_tokens(self, index):
		"""Return an example's length (number of tokens), used for batching."""
		"""Return the number of tokens in a sample. This value is used to
		enforce ``--max-tokens`` during batching."""
		raise NotImplementedError

		def ordered_indices(self):
		"""Ordered indices for batching."""
		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``."""
		raise NotImplementedError

		def valid_size(self, index, max_positions):
		"""Check if an example's size is valid according to max_positions."""
		def ordered_indices(self):
		"""Return an ordered list of indices. Batches will be constructed based
		on this order."""
		raise NotImplementedError

fairseq/data/language_pair_dataset.py

+63 −21

Original line number	Diff line number	Diff line
		@@ -8,6 +8,8 @@
		import numpy as np
		import torch

		from fairseq import utils

		from . import data_utils, FairseqDataset


		@@ -59,7 +61,27 @@ def collate(samples, pad_idx, eos_idx, left_pad_source=True, left_pad_target=Fal


		class LanguagePairDataset(FairseqDataset):
		"""A pair of torch.utils.data.Datasets."""
		"""
		A pair of torch.utils.data.Datasets.

		Args:
		src (torch.utils.data.Dataset): source dataset to wrap
		src_sizes (List[int]): source sentence lengths
		src_dict (fairseq.data.Dictionary): source vocabulary
		tgt (torch.utils.data.Dataset, optional): target dataset to wrap
		tgt_sizes (List[int], optional): target sentence lengths
		tgt_dict (fairseq.data.Dictionary, optional): target vocabulary
		left_pad_source (bool, optional): pad source tensors on the left side.
		Default: ``True``
		left_pad_target (bool, optional): pad target tensors on the left side.
		Default: ``False``
		max_source_positions (int, optional): max number of tokens in the source
		sentence. Default: ``1024``
		max_target_positions (int, optional): max number of tokens in the target
		sentence. Default: ``1024``
		shuffle (bool, optional): shuffle dataset elements before batching.
		Default: ``True``
		"""

		def __init__(
		self, src, src_sizes, src_dict,
		@@ -95,15 +117,43 @@ class LanguagePairDataset(FairseqDataset):
		return len(self.src)

		def collater(self, samples):
		"""Merge a list of samples to form a mini-batch."""
		"""Merge a list of samples to form a mini-batch.

		Returned mini-batches contain the following keys:
		- `id` (torch.LongTensor): example IDs in the original input order
		- `ntokens` (int): total number of tokens in the batch
		- `net_input` (dict): the input to the Model, containing keys:
		- `src_tokens` (torch.LongTensor): a padded 2D Tensor of tokens in
		the source sentence of shape `(bsz, src_len)`. Padding will appear
		on the left if ``left_pad_source`` is True.
		- `src_lengths` (torch.LongTensor): 1D Tensor of the unpadded lengths
		of each source sentence of shape `(bsz)`
		- `prev_output_tokens` (torch.LongTensor): a padded 2D Tensor of
		tokens in the target sentence, shifted right by one position for
		input feeding/teacher forcing, of shape `(bsz, tgt_len)`. Padding
		will appear on the left if ``left_pad_target`` is True.
		- `target` (torch.LongTensor): a padded 2D Tensor of tokens in the
		target sentence of shape `(bsz, tgt_len)`. Padding will appear on the
		left if ``left_pad_target`` is True.

		Args:
		samples (List[dict]): samples to collate

		Returns:
		dict: a mini-batch suitable for forwarding with a Model
		"""
		return collate(
		samples, pad_idx=self.src_dict.pad(), eos_idx=self.src_dict.eos(),
		left_pad_source=self.left_pad_source, left_pad_target=self.left_pad_target,
		)

		def get_dummy_batch(self, num_tokens, max_positions, src_len=128, tgt_len=128):
		max_source_positions, max_target_positions = self._get_max_positions(max_positions)
		src_len, tgt_len = min(src_len, max_source_positions), min(tgt_len, max_target_positions)
		"""Return a dummy batch with a given number of tokens."""
		src_len, tgt_len = utils.resolve_max_positions(
		(src_len, tgt_len),
		max_positions,
		(self.max_source_positions, self.max_target_positions),
		)
		bsz = num_tokens // max(src_len, tgt_len)
		return self.collater([
		{
		@@ -115,11 +165,18 @@ class LanguagePairDataset(FairseqDataset):
		])

		def num_tokens(self, index):
		"""Return an example's length (number of tokens), used for batching."""
		"""Return the number of tokens in a sample. This value is used to
		enforce ``--max-tokens`` during batching."""
		return max(self.src_sizes[index], self.tgt_sizes[index] if self.tgt_sizes is not None else 0)

		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``."""
		return (self.src_sizes[index], self.tgt_sizes[index] if self.tgt_sizes is not None else 0)

		def ordered_indices(self):
		"""Ordered indices for batching."""
		"""Return an ordered list of indices. Batches will be constructed based
		on this order."""
		if self.shuffle:
		indices = np.random.permutation(len(self))
		else:
		@@ -127,18 +184,3 @@ class LanguagePairDataset(FairseqDataset):
		if self.tgt_sizes is not None:
		indices = indices[np.argsort(self.tgt_sizes[indices], kind='mergesort')]
		return indices[np.argsort(self.src_sizes[indices], kind='mergesort')]

		def valid_size(self, index, max_positions):
		"""Check if an example's size is valid according to max_positions."""
		max_source_positions, max_target_positions = self._get_max_positions(max_positions)
		return (
		self.src_sizes[index] <= max_source_positions
		and (self.tgt_sizes is None or self.tgt_sizes[index] <= max_target_positions)
		)

		def _get_max_positions(self, max_positions):
		if max_positions is None:
		return self.max_source_positions, self.max_target_positions
		assert len(max_positions) == 2
		max_src_pos, max_tgt_pos = max_positions
		return min(self.max_source_positions, max_src_pos), min(self.max_target_positions, max_tgt_pos)

fairseq/data/monolingual_dataset.py

+40 −10

Original line number	Diff line number	Diff line
		@@ -31,7 +31,16 @@ def collate(samples, pad_idx, eos_idx):


		class MonolingualDataset(FairseqDataset):
		"""A wrapper around torch.utils.data.Dataset for monolingual data."""
		"""
		A wrapper around torch.utils.data.Dataset for monolingual data.

		Args:
		dataset (torch.utils.data.Dataset): dataset to wrap
		sizes (List[int]): sentence lengths
		vocab (fairseq.data.Dictionary): vocabulary
		shuffle (bool, optional): shuffle the elements before batching.
		Default: ``True``
		"""

		def __init__(self, dataset, sizes, vocab, shuffle):
		self.dataset = dataset
		@@ -47,11 +56,30 @@ class MonolingualDataset(FairseqDataset):
		return len(self.dataset)

		def collater(self, samples):
		"""Merge a list of samples to form a mini-batch."""
		"""Merge a list of samples to form a mini-batch.

		Returned mini-batches contain the following keys:
		- `id` (torch.LongTensor): example IDs in the original input order
		- `ntokens` (int): total number of tokens in the batch
		- `net_input` (dict): the input to the Model, containing keys:
		- `src_tokens` (torch.LongTensor): a padded 2D Tensor of tokens in
		the source sentence of shape `(bsz, src_len)`. Padding will appear
		on the right.
		- `target` (torch.LongTensor): a padded 2D Tensor of tokens in the
		target sentence of shape `(bsz, tgt_len)`. Padding will appear on the
		right.

		Args:
		samples (List[dict]): samples to collate

		Returns:
		dict: a mini-batch suitable for forwarding with a Model
		"""
		return collate(samples, self.vocab.pad(), self.vocab.eos())

		def get_dummy_batch(self, num_tokens, max_positions, tgt_len=128):
		assert isinstance(max_positions, float) or isinstance(max_positions, int)
		"""Return a dummy batch with a given number of tokens."""
		if isinstance(max_positions, float) or isinstance(max_positions, int):
		tgt_len = min(tgt_len, max_positions)
		bsz = num_tokens // tgt_len
		target = self.vocab.dummy_sentence(tgt_len + 1)
		@@ -62,19 +90,21 @@ class MonolingualDataset(FairseqDataset):
		])

		def num_tokens(self, index):
		"""Return an example's length (number of tokens), used for batching."""
		"""Return the number of tokens in a sample. This value is used to
		enforce ``--max-tokens`` during batching."""
		return self.sizes[index]

		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``."""
		return self.sizes[index]

		def ordered_indices(self):
		"""Ordered indices for batching."""
		"""Return an ordered list of indices. Batches will be constructed based
		on this order."""
		if self.shuffle:
		order = [np.random.permutation(len(self))]
		else:
		order = [np.arange(len(self))]
		order.append(np.flip(self.sizes, 0))
		return np.lexsort(order)

		def valid_size(self, index, max_positions):
		"""Check if an example's size is valid according to max_positions."""
		assert isinstance(max_positions, float) or isinstance(max_positions, int)
		return self.sizes[index] <= max_positions