2018-09-25 21:02:34 +03:00
|
|
|
# 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 unittest
|
|
|
|
|
|
|
|
import tests.utils as test_utils
|
|
|
|
import torch
|
|
|
|
from fairseq.data.backtranslation_dataset import BacktranslationDataset
|
2018-10-03 04:20:01 +03:00
|
|
|
from fairseq import sequence_generator
|
2018-09-25 21:02:34 +03:00
|
|
|
|
|
|
|
|
|
|
|
class TestBacktranslationDataset(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
|
|
self.tgt_dict, self.w1, self.w2, self.src_tokens, self.src_lengths, self.model = (
|
|
|
|
test_utils.sequence_generator_setup()
|
|
|
|
)
|
|
|
|
|
|
|
|
dummy_src_samples = self.src_tokens
|
|
|
|
|
|
|
|
self.tgt_dataset = test_utils.TestDataset(data=dummy_src_samples)
|
|
|
|
|
2018-10-04 04:18:03 +03:00
|
|
|
def _backtranslation_dataset_helper(self, remove_eos_at_src):
|
2018-10-03 04:20:01 +03:00
|
|
|
"""
|
|
|
|
SequenceGenerator kwargs are same as defaults from fairseq/options.py
|
|
|
|
"""
|
2018-09-25 21:02:34 +03:00
|
|
|
backtranslation_dataset = BacktranslationDataset(
|
|
|
|
tgt_dataset=self.tgt_dataset,
|
|
|
|
tgt_dict=self.tgt_dict,
|
|
|
|
backtranslation_model=self.model,
|
Explicitly list out generation args for backtranslation dataset
Summary:
Using argparse Namespace hides the actual args that are expected and makes code harder to read.
Note the difference in style for the args list
def __init__(
self,
tgt_dataset,
tgt_dict,
backtranslation_model,
unkpen,
sampling,
beam,
max_len_a,
max_len_b,
):
instead of
def __init__(
self, tgt_dataset, tgt_dict, backtranslation_model, unkpen, sampling,
beam, max_len_a, max_len_b,
):
Reviewed By: dpacgopinath
Differential Revision: D10152331
fbshipit-source-id: 6539ccba09d48acf23759996b7e32fb329b3e3f6
2018-10-03 01:10:21 +03:00
|
|
|
max_len_a=0,
|
|
|
|
max_len_b=200,
|
2018-10-03 04:20:01 +03:00
|
|
|
beam_size=2,
|
|
|
|
unk_penalty=0,
|
|
|
|
sampling=False,
|
2018-10-04 04:18:03 +03:00
|
|
|
remove_eos_at_src=remove_eos_at_src,
|
2018-10-03 04:20:01 +03:00
|
|
|
generator_class=sequence_generator.SequenceGenerator,
|
2018-09-25 21:02:34 +03:00
|
|
|
)
|
|
|
|
dataloader = torch.utils.data.DataLoader(
|
|
|
|
backtranslation_dataset,
|
|
|
|
batch_size=2,
|
|
|
|
collate_fn=backtranslation_dataset.collater,
|
|
|
|
)
|
|
|
|
backtranslation_batch_result = next(iter(dataloader))
|
|
|
|
|
|
|
|
eos, pad, w1, w2 = self.tgt_dict.eos(), self.tgt_dict.pad(), self.w1, self.w2
|
|
|
|
|
|
|
|
# Note that we sort by src_lengths and add left padding, so actually
|
|
|
|
# ids will look like: [1, 0]
|
|
|
|
expected_src = torch.LongTensor([[w1, w2, w1, eos], [pad, pad, w1, eos]])
|
2018-10-04 04:18:03 +03:00
|
|
|
if remove_eos_at_src:
|
|
|
|
expected_src = expected_src[:, :-1]
|
2018-09-25 21:02:34 +03:00
|
|
|
expected_tgt = torch.LongTensor([[w1, w2, eos], [w1, w2, eos]])
|
|
|
|
generated_src = backtranslation_batch_result["net_input"]["src_tokens"]
|
|
|
|
tgt_tokens = backtranslation_batch_result["target"]
|
|
|
|
|
|
|
|
self.assertTensorEqual(expected_src, generated_src)
|
|
|
|
self.assertTensorEqual(expected_tgt, tgt_tokens)
|
|
|
|
|
2018-10-04 04:18:03 +03:00
|
|
|
def test_backtranslation_dataset_no_eos_at_src(self):
|
|
|
|
self._backtranslation_dataset_helper(remove_eos_at_src=True)
|
|
|
|
|
|
|
|
def test_backtranslation_dataset_with_eos_at_src(self):
|
|
|
|
self._backtranslation_dataset_helper(remove_eos_at_src=False)
|
|
|
|
|
2018-09-25 21:02:34 +03:00
|
|
|
def assertTensorEqual(self, t1, t2):
|
|
|
|
self.assertEqual(t1.size(), t2.size(), "size mismatch")
|
|
|
|
self.assertEqual(t1.ne(t2).long().sum(), 0)
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
unittest.main()
|