mirror of
https://github.com/facebookresearch/fairseq.git
synced 2024-09-11 17:25:31 +03:00
f766c9a0d5
Summary: This generalizes BacktranslationDataset to allow us to use any SequenceGenerator class. For example, if we want to use this model in PyTorch Translate, we can pass the following to BacktraanslationDataset init: (1) a PyTorch Translate SequenceGenerator class as generator_class and (2) the appropriate args for initializing that class as kwargs. Reviewed By: xianxl Differential Revision: D10156552 fbshipit-source-id: 0495d825bf4727da96d0d9a40dc434135ff3486c
67 lines
2.4 KiB
Python
67 lines
2.4 KiB
Python
# 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
|
|
from fairseq import sequence_generator
|
|
|
|
|
|
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)
|
|
|
|
def test_backtranslation_dataset(self):
|
|
"""
|
|
SequenceGenerator kwargs are same as defaults from fairseq/options.py
|
|
"""
|
|
backtranslation_dataset = BacktranslationDataset(
|
|
tgt_dataset=self.tgt_dataset,
|
|
tgt_dict=self.tgt_dict,
|
|
backtranslation_model=self.model,
|
|
max_len_a=0,
|
|
max_len_b=200,
|
|
beam_size=2,
|
|
unk_penalty=0,
|
|
sampling=False,
|
|
generator_class=sequence_generator.SequenceGenerator,
|
|
)
|
|
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]])
|
|
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)
|
|
|
|
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()
|