Refactor BacktranslationDataset to be more reusable (#354) (3c19878f) · Commits · Simon Will / fairseq

fairseq/data/init.py

+2 −2

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

		from .dictionary import Dictionary, TruncatedDictionary
		from .fairseq_dataset import FairseqDataset
		from .append_eos_dataset import AppendEosDataset
		from .backtranslation_dataset import BacktranslationDataset
		from .concat_dataset import ConcatDataset
		from .indexed_dataset import IndexedDataset, IndexedCachedDataset, IndexedInMemoryDataset, IndexedRawTextDataset
		@@ -15,6 +14,7 @@ from .language_pair_dataset import LanguagePairDataset
		from .monolingual_dataset import MonolingualDataset
		from .round_robin_zip_datasets import RoundRobinZipDatasets
		from .token_block_dataset import TokenBlockDataset
		from .transform_eos_dataset import TransformEosDataset

		from .iterators import (
		CountingIterator,
		@@ -24,7 +24,6 @@ from .iterators import (
		)

		__all__ = [
		'AppendEosDataset',
		'BacktranslationDataset',
		'ConcatDataset',
		'CountingIterator',
		@@ -41,4 +40,5 @@ __all__ = [
		'RoundRobinZipDatasets',
		'ShardedIterator',
		'TokenBlockDataset',
		'TransformEosDataset',
		]

fairseq/data/append_eos_dataset.py

deleted100644 → 0

+0 −23

Original line number	Diff line number	Diff line
		# Copyright (c) 2017-present, Facebook, Inc.
		# All rights reserved.
		#
		# This source code is licensed under the license found in the LICENSE file in
		# the root directory of this source tree. An additional grant of patent rights
		# can be found in the PATENTS file in the same directory.

		import torch


		class AppendEosDataset(torch.utils.data.Dataset):
		"""A dataset wrapper that appends EOS to each item."""

		def __init__(self, dataset, eos):
		self.dataset = dataset
		self.eos = eos

		def __getitem__(self, index):
		item = torch.cat([self.dataset[index], torch.LongTensor([self.eos])])
		return item

		def __len__(self):
		return len(self.dataset)

fairseq/data/backtranslation_dataset.py

+106 −126

Original line number	Diff line number	Diff line
		@@ -7,144 +7,138 @@

		import torch

		from fairseq import sequence_generator
		from fairseq import utils

		from . import FairseqDataset, language_pair_dataset
		from . import FairseqDataset


		def backtranslate_samples(samples, collate_fn, generate_fn, cuda=True):
		"""Backtranslate a list of samples.

		Given an input (samples) of the form:

		[{'id': 1, 'source': 'hallo welt'}]

		this will return:

		[{'id': 1, 'source': 'hello world', 'target': 'hallo welt'}]

		Args:
		samples (List[dict]): samples to backtranslate. Individual samples are
		expected to have a 'source' key, which will become the 'target'
		after backtranslation.
		collate_fn (callable): function to collate samples into a mini-batch
		generate_fn (callable): function to generate backtranslations
		cuda (bool): use GPU for generation (default: ``True``)

		Returns:
		List[dict]: an updated list of samples with a backtranslated source
		"""
		collated_samples = collate_fn(samples)
		s = utils.move_to_cuda(collated_samples) if cuda else collated_samples
		generated_sources = generate_fn(s['net_input'])

		def update_sample(sample, generated_source):
		sample['target'] = sample['source'] # the original source becomes the target
		sample['source'] = generated_source
		return sample

		# Go through each tgt sentence in batch and its corresponding best
		# generated hypothesis and create a backtranslation data pair
		# {id: id, source: generated backtranslation, target: original tgt}
		return [
		update_sample(
		sample=input_sample,
		generated_source=hypos[0]['tokens'].cpu(), # highest scoring hypo is first
		)
		for input_sample, hypos in zip(samples, generated_sources)
		]


		class BacktranslationDataset(FairseqDataset):
		def __init__(
		self,
		tgt_dataset,
		tgt_dict,
		backtranslation_model,
		backtranslation_fn,
		max_len_a,
		max_len_b,
		remove_eos_at_src=False,
		generator_class=sequence_generator.SequenceGenerator,
		output_collater=None,
		cuda=True,
		**kwargs
		):
		"""
		Sets up a backtranslation dataset which takes a tgt batch, generates
		a src using a tgt-src backtranslation_model, and returns the
		corresponding {generated src, input tgt} batch
		a src using a tgt-src backtranslation function (backtranslation_fn),
		and returns the corresponding `{generated src, input tgt}` batch.

		Args:
		tgt_dataset: dataset which will be used to build self.tgt_dataset --
		a LanguagePairDataset with tgt dataset as the source dataset and
		None as the target dataset. Should NOT have padding so that
		src_lengths are accurately calculated by language_pair_dataset
		collate function.
		We use language_pair_dataset here to encapsulate the tgt_dataset
		so we can re-use the LanguagePairDataset collater to format the
		batches in the structure that SequenceGenerator expects.
		Note: tgt_dataset samples should not have EOS at end if
		the tgt-src model expects an input without EOS. This dataset
		does not enforce this, you should enforce that in preprocessing.
		tgt_dict: tgt dictionary (typically a joint src/tgt BPE dictionary)
		backtranslation_model: tgt-src model to use in the SequenceGenerator
		to generate backtranslations from tgt batches
		max_len_a, max_len_b: args passed into generate() function of
		the backtranslation SequenceGenerator
		remove_eos_at_src: whether we should remove EOS from the source
		dialect text generated by the backtranslation model.
		generator_class: which SequenceGenerator class to use for
		backtranslation. Output of generate() should be the same format
		as fairseq's SequenceGenerator
		tgt_dataset (~fairseq.data.FairseqDataset): the dataset to be
		backtranslated. Only the source side of this dataset will be
		used. After backtranslation, the source sentences in this
		dataset will be returned as the targets.
		backtranslation_fn (callable): function to call to generate
		backtranslations. This is typically the `generate` method of a
		:class:`~fairseq.sequence_generator.SequenceGenerator` object.
		max_len_a, max_len_b (int, int): will be used to compute
		`maxlen = max_len_a * src_len + max_len_b`, which will be
		passed into backtranslation_fn.
		output_collater (callable, optional): function to call on the
		backtranslated samples to create the final batch (default:
		``tgt_dataset.collater``)
		cuda: use GPU for generation
		kwargs: generation args to init the backtranslation
		SequenceGenerator
		"""
		self.tgt_dataset = language_pair_dataset.LanguagePairDataset(
		src=tgt_dataset,
		src_sizes=tgt_dataset.sizes,
		src_dict=tgt_dict,
		tgt=None,
		tgt_sizes=None,
		tgt_dict=None,
		)

		self.tgt_dataset = tgt_dataset
		self.backtranslation_fn = backtranslation_fn
		self.max_len_a = max_len_a
		self.max_len_b = max_len_b
		self.remove_eos_at_src = remove_eos_at_src

		self.backtranslation_generator = generator_class(
		models=[backtranslation_model],
		tgt_dict=tgt_dict,
		**kwargs,
		)

		self.output_collater = output_collater if output_collater is not None \
		else tgt_dataset.collater
		self.cuda = cuda if torch.cuda.is_available() else False
		if self.cuda:
		self.backtranslation_generator.cuda()

		def __getitem__(self, index):
		"""
		Returns a single sample. Multiple samples are fed to the collater to
		create a backtranslation batch. Note you should always use collate_fn
		BacktranslationDataset.collater() below if given the option to
		specify which collate_fn to use (e.g. in a dataloader which uses this
		BacktranslationDataset -- see corresponding unittest for an example).
		Returns a single sample from tgt_dataset. Note that backtranslation is
		not applied in this step; use :func:`collater` instead to backtranslate
		a batch of samples.
		"""
		return self.tgt_dataset[index]

		def __len__(self):
		"""
		The length of the backtranslation dataset is the length of tgt.
		"""
		return len(self.tgt_dataset)

		def collater(self, samples):
		"""
		Using the samples from the tgt dataset, load a collated tgt sample to
		feed to the backtranslation model. Then take the generated translation
		with best score as the source and the orignal net input as the target.
		"""
		collated_tgt_only_sample = self.tgt_dataset.collater(samples=samples)
		backtranslation_hypos = self._generate_hypotheses(
		sample=collated_tgt_only_sample
		)
		"""Merge and backtranslate a list of samples to form a mini-batch.

		# Go through each tgt sentence in batch and its corresponding best
		# generated hypothesis and create a backtranslation data pair
		# {id: id, source: generated backtranslation, target: original tgt}
		generated_samples = []
		for input_sample, hypos in zip(samples, backtranslation_hypos):
		original_tgt = input_sample["source"].cpu()
		generated_source = hypos[0]["tokens"].cpu() # first hypo is best hypo

		# Append EOS to the tgt sentence if it does not have an EOS
		# This is the case if the samples in monolingual tgt_dataset don't
		# have an EOS appended to the end of each sentence.
		eos = self.tgt_dataset.src_dict.eos()
		if original_tgt[-1] != eos:
		original_tgt = torch.cat([original_tgt, torch.LongTensor([eos])])

		# The generated source dialect backtranslation will have an EOS.
		# If we want our parallel data source to not have an EOS, we will
		# have to remove it.
		if self.remove_eos_at_src:
		assert generated_source[-1] == eos, (
		"Expected generated backtranslation to have eos (id: "
		"{eos}) at end, but instead found token id "
		"{generated_source[-1]} at end."
		).format(eos=eos, generated_source=generated_source)
		generated_source = generated_source[:-1]

		generated_samples.append(
		{
		"id": input_sample["id"],
		"source": generated_source,
		"target": original_tgt,
		}
		)
		Using the samples from tgt_dataset, load a collated target sample to
		feed to the backtranslation model. Then take the backtranslation with
		the best score as the source and the original input as the target.

		Note: we expect tgt_dataset to provide a function `collater()` that
		will collate samples into the format expected by backtranslation_fn.
		After backtranslation, we will feed the new list of samples (i.e., the
		`(backtranslated source, original source)` pairs) to output_collater
		and return the result.

		return language_pair_dataset.collate(
		samples=generated_samples,
		pad_idx=self.tgt_dataset.src_dict.pad(),
		eos_idx=self.tgt_dataset.src_dict.eos(),
		Args:
		samples (List[dict]): samples to backtranslate and collate

		Returns:
		dict: a mini-batch with keys coming from output_collater
		"""
		samples = backtranslate_samples(
		samples=samples,
		collate_fn=self.tgt_dataset.collater,
		generate_fn=(
		lambda net_input: self.backtranslation_fn(
		net_input,
		maxlen=int(
		self.max_len_a * net_input['src_tokens'].size(1) + self.max_len_b
		),
		)
		),
		cuda=self.cuda,
		)
		return self.output_collater(samples)

		def get_dummy_batch(self, num_tokens, max_positions):
		"""Just use the tgt dataset get_dummy_batch"""
		@@ -162,27 +156,13 @@ class BacktranslationDataset(FairseqDataset):
		"""Just use the tgt dataset size"""
		return self.tgt_dataset.valid_size(index, max_positions)

		def _generate_hypotheses(self, sample):
		"""
		Generates hypotheses from a LanguagePairDataset collated / batched
		sample. Note in this case, sample["target"] is None, and
		sample["net_input"]["src_tokens"] is really in tgt language.
		"""
		s = utils.move_to_cuda(sample) if self.cuda else sample
		input = s["net_input"]
		srclen = input["src_tokens"].size(1)
		hypos = self.backtranslation_generator.generate(
		input,
		maxlen=int(
		self.max_len_a * srclen + self.max_len_b
		),
		)
		return hypos

		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``.
		Here, we return src dataset size as tgt dataset size as an approximation.
		We do not know src size until we backtranslate and generate src sentences.
		"""Return an example's size as a float or tuple. This value is used
		when filtering a dataset with ``--max-positions``.

		Note: we use tgt_dataset to approximate the length of the source
		sentence, since we do not know the actual length until after
		backtranslation.
		"""
		return (self.tgt_dataset.size(index)[0], self.tgt_dataset.size(index)[0])
		tgt_size = self.tgt_dataset.size(index)[0]
		return (tgt_size, tgt_size)

fairseq/data/concat_dataset.py

+7 −0

Original line number	Diff line number	Diff line
		# Copyright (c) 2017-present, Facebook, Inc.
		# All rights reserved.
		#
		# This source code is licensed under the license found in the LICENSE file in
		# the root directory of this source tree. An additional grant of patent rights
		# can be found in the PATENTS file in the same directory.

		import bisect

		import numpy as np

fairseq/data/language_pair_dataset.py

+1 −1

Original line number	Diff line number	Diff line
		@@ -55,13 +55,13 @@ def collate(

		batch = {
		'id': id,
		'nsentences': len(samples),
		'ntokens': ntokens,
		'net_input': {
		'src_tokens': src_tokens,
		'src_lengths': src_lengths,
		},
		'target': target,
		'nsentences': samples[0]['source'].size(0),
		}
		if prev_output_tokens is not None:
		batch['net_input']['prev_output_tokens'] = prev_output_tokens