def test_rnn_fp16(self):
with contextlib.redirect_stdout(StringIO()):
with tempfile.TemporaryDirectory("test_rnn_fp16") as data_dir:
create_dummy_data(data_dir)
train_translation_model(
data_dir,
[
"--fp16",
"--arch",
"rnn",
"--cell-type",
"lstm",
"--sequence-lstm",
"--reverse-source",
"--encoder-bidirectional",
"--encoder-layers",
"2",
"--encoder-embed-dim",
"8",
"--encoder-hidden-dim",
"16",
"--decoder-layers",
"2",
"--decoder-embed-dim",
"8",
"--decoder-hidden-dim",
"16",
"--decoder-out-embed-dim",
"8",
"--attention-type",
"dot",
],
)
generate_main(data_dir)
def test_semi_supervised_transformer(self):
with contextlib.redirect_stdout(StringIO()):
with tempfile.TemporaryDirectory("test_transformer") as data_dir:
create_dummy_data(data_dir)
train_translation_model(
data_dir,
[
"--arch",
"semi_supervised_transformer",
# semi-supervised task args:
"--task",
"pytorch_translate_semi_supervised",
"--train-mono-source-text-file",
os.path.join(data_dir, "train.in"),
"--train-mono-target-text-file",
os.path.join(data_dir, "train.out"),
# transformer args:
"--encoder-embed-dim",
"4",
"--encoder-ffn-embed-dim",
"4",
"--encoder-attention-heads",
"2",
"--encoder-layers",
"1",
"--decoder-embed-dim",
"4",
"--decoder-ffn-embed-dim",
"4",
"--decoder-attention-heads",
"2",
"--decoder-layers",
"1",
],
)
def test_transformer_multigpu(self):
with contextlib.redirect_stdout(StringIO()):
with tempfile.TemporaryDirectory("test_transformer") as data_dir:
create_dummy_data(data_dir)
train_translation_model(
data_dir,
[
"--arch",
"ptt_transformer",
"--encoder-embed-dim",
"256",
"--encoder-ffn-embed-dim",
"512",
"--encoder-attention-heads",
"4",
"--encoder-layers",
"3",
"--decoder-embed-dim",
"256",
"--decoder-ffn-embed-dim",
"512",
"--decoder-attention-heads",
"4",
"--decoder-layers",
"3",
"--distributed-world-size",
str(torch.cuda.device_count()),
"--local-num-gpus",
str(torch.cuda.device_count()),
],
)
generate_main(data_dir)
def test_transformer(self):
with contextlib.redirect_stdout(StringIO()):
with tempfile.TemporaryDirectory("test_transformer") as data_dir:
create_dummy_data(data_dir)
train_translation_model(
data_dir,
[
"--arch",
"ptt_transformer",
"--encoder-embed-dim",
"8",
"--encoder-ffn-embed-dim",
"16",
"--encoder-attention-heads",
"4",
"--encoder-layers",
"3",
"--decoder-embed-dim",
"8",
"--decoder-ffn-embed-dim",
"16",
"--decoder-attention-heads",
"4",
"--decoder-layers",
"3",
],
)
generate_main(data_dir)
def test_char_rnn(self):
with contextlib.redirect_stdout(StringIO()):
with tempfile.TemporaryDirectory("test_char_rnn") as data_dir:
create_dummy_data(data_dir)
train_translation_model(
data_dir,
[
"--arch",
"char_source",
"--char-embed-dim",
"4",
"--char-rnn-units",
"8",
"--char-rnn-layers",
"1",
"--char-source-max-vocab-size",
"26",
"--cell-type",
"lstm",
"--sequence-lstm",
"--encoder-dropout-in",
"0",
"--encoder-dropout-out",
"0",
"--decoder-dropout-in",
"0.2",
"--decoder-dropout-out",
"0.2",
"--encoder-layers",
"2",
"--encoder-embed-dim",
"8",
"--encoder-hidden-dim",
"16",
"--decoder-layers",
"2",
"--decoder-embed-dim",
"8",
"--decoder-hidden-dim",
"16",
"--decoder-out-embed-dim",
"8",
"--attention-type",
"dot",
],
)
generate_main(
data_dir,
[
"--char-source-vocab-file",
os.path.join(data_dir, "char-dictionary-in.txt"),
],
)
def test_char_aware_hybrid(self):
with contextlib.redirect_stdout(StringIO()):
with tempfile.TemporaryDirectory("test_char_aware") as data_dir:
create_dummy_data(data_dir)
train_translation_model(
data_dir,
[
"--arch",
"char_aware_hybrid",
"--char-embed-dim",
"4",
"--char-cnn-params",
"[(10, 1), (20,2)]",
"--char-cnn-nonlinear-fn",
"relu",
"--char-cnn-num-highway-layers",
"2",
"--char-source-max-vocab-size",
"26",
"--char-target-max-vocab-size",
"26",
"--encoder-embed-dim",
"8",
"--encoder-ffn-embed-dim",
"16",
"--encoder-attention-heads",
"4",
"--encoder-layers",
"3",
"--decoder-embed-dim",
"8",
"--decoder-attention-heads",
"4",
"--decoder-layers",
"2",
"--decoder-lstm-units",
"16",
"--decoder-out-embed-dim",
"8",
],
)
generate_main(
data_dir,
[
"--char-source-vocab-file",
os.path.join(data_dir, "char-src-dictionary-in.txt"),
"--char-target-vocab-file",
os.path.join(data_dir, "char-tgt-dictionary-in.txt"),
],
)
def test_word_prediction(self):
""" Tests a word prediction model, which will use a learned vocab
reduction via the word prediction model. It uses a custom criterion
(word_prediction) on top of label_smoothed_cross_entropy so we pass the
corresponding word_prediction criterion flag in during training.
"""
with contextlib.redirect_stdout(StringIO()):
with tempfile.TemporaryDirectory("test_word_pred") as data_dir:
create_dummy_data(data_dir)
train_translation_model(
data_dir=data_dir,
extra_flags=[
"--arch",
"rnn_word_pred",
"--predictor-hidden-dim",
"32",
"--topk-labels-per-source-token",
"5",
"--cell-type",
"lstm",
"--sequence-lstm",
"--reverse-source",
"--encoder-bidirectional",
"--encoder-layers",
"2",
"--encoder-embed-dim",
"8",
"--encoder-hidden-dim",
"16",
"--decoder-layers",
"2",
"--decoder-embed-dim",
"8",
"--decoder-hidden-dim",
"16",
"--decoder-out-embed-dim",
"8",
"--attention-type",
"dot",
],
criterion=[
"--criterion",
"word_prediction",
"--label-smoothing",
"0.1",
],
)
generate_main(data_dir)
def test_denoising_autoencoder(self):
"""
Tests denoising autoencoder task. Important flags:
`--train-mono-*-text-file`, `--task`, `--arch`, and
`--denoising-target-mono`.
"""
with contextlib.redirect_stdout(StringIO()):
with tempfile.TemporaryDirectory("test_rnn") as data_dir:
create_dummy_data(data_dir)
train_translation_model(
data_dir,
[
"--task",
"pytorch_translate_denoising_autoencoder",
"--train-mono-source-text-file",
os.path.join(data_dir, "train.in"),
"--train-mono-target-text-file",
os.path.join(data_dir, "train.out"),
"--arch",
"semi_supervised_rnn",
"--denoising-target-mono",
"--cell-type",
"lstm",
"--sequence-lstm",
"--reverse-source",
"--encoder-bidirectional",
"--encoder-layers",
"2",
"--encoder-embed-dim",
"8",
"--encoder-hidden-dim",
"16",
"--decoder-layers",
"2",
"--decoder-embed-dim",
"8",
"--decoder-hidden-dim",
"16",
"--decoder-out-embed-dim",
"8",
"--attention-type",
"dot",
],
)
def test_pretrained_masked_lm_for_translation(self):
with contextlib.redirect_stdout(StringIO()):
with tempfile.TemporaryDirectory("test_mlm") as data_dir:
create_dummy_data(data_dir)
train_translation_model(
data_dir,
[
"--arch",
"xlm_base",
data_dir,
# semi-supervised task args:
"--task",
"pytorch_translate_cross_lingual_lm",
# transformer args:
"--encoder-embed-dim",
"4",
"--encoder-ffn-embed-dim",
"4",
"--encoder-attention-heads",
"2",
"--encoder-layers",
"1",
# dict files
"--source-vocab-file",
os.path.join(data_dir, "dictionary-in.txt"),
"--target-vocab-file",
os.path.join(data_dir, "dictionary-out.txt"),
# additoinal ones
"--dataset-impl",
"raw",
"--monolingual-langs",
"in,out",
"--save-only",
"--masked-lm-only",
"--num-segment",
"2",
],
criterion=["--criterion", "legacy_masked_lm_loss"],
)
def test_multilingual_hybrid(self):
"""
Tests multilingual translation task. Important flags:
`--multilingual-*-binary-path`, `--task`, `--arch`,
`--source-vocabulary`, `--target-vocabulary`, `--vocabulary`.
"""
with contextlib.redirect_stdout(StringIO()):
with tempfile.TemporaryDirectory(
"test_multilingual_hybrid") as data_dir:
create_dummy_multilingual_data(data_dir)
train_translation_model(
data_dir,
[
"--task",
"pytorch_translate_multilingual_task",
"--arch",
"multilingual_hybrid_transformer_rnn",
"--encoder-embed-dim",
"8",
"--encoder-ffn-embed-dim",
"16",
"--encoder-attention-heads",
"4",
"--encoder-layers",
"3",
"--decoder-embed-dim",
"8",
"--decoder-attention-heads",
"4",
"--decoder-layers",
"2",
"--decoder-lstm-units",
"16",
"--decoder-out-embed-dim",
"8",
"--lang-pairs",
"xh-en,zu-en",
"--multilingual-train-text-file",
("xh-en:"
f"{os.path.join(data_dir, 'train.xhen.xh')},"
f"{os.path.join(data_dir, 'train.xhen.en')}"),
"--multilingual-eval-text-file",
("xh-en:"
f"{os.path.join(data_dir, 'tune.xhen.xh')},"
f"{os.path.join(data_dir, 'tune.xhen.en')}"),
"--multilingual-train-text-file",
("zu-en:"
f"{os.path.join(data_dir, 'train.zuen.zu')},"
f"{os.path.join(data_dir, 'train.zuen.en')}"),
"--multilingual-eval-text-file",
("zu-en:"
f"{os.path.join(data_dir, 'tune.zuen.zu')},"
f"{os.path.join(data_dir, 'tune.zuen.en')}"),
# set these to empty to satisfy argument validation
"--train-source-text-file",
"",
"--train-target-text-file",
"",
"--eval-source-text-file",
"",
"--eval-target-text-file",
"",
],
# fairseq MultlilingualTranslationTask expects mandatory
# data directory positional argument
set_empty_data_positional_arg=True,
set_lang_args=False,
)
def test_multilingual(self):
with contextlib.redirect_stdout(StringIO()):
with tempfile.TemporaryDirectory("test_multilingual") as data_dir:
create_dummy_multilingual_data(data_dir)
train_translation_model(
data_dir,
[
"--task",
"pytorch_translate_multilingual",
"--arch",
"rnn",
"--cell-type",
"lstm",
"--sequence-lstm",
"--reverse-source",
"--encoder-bidirectional",
"--encoder-layers",
"2",
"--encoder-embed-dim",
"8",
"--encoder-hidden-dim",
"16",
"--decoder-layers",
"2",
"--decoder-embed-dim",
"8",
"--decoder-hidden-dim",
"16",
"--decoder-out-embed-dim",
"8",
"--attention-type",
"dot",
"--multiling-encoder-lang",
"xh",
"--multiling-encoder-lang",
"zu",
"--multiling-encoder-lang",
"en",
"--multiling-decoder-lang",
"xh",
"--multiling-decoder-lang",
"en",
"--multiling-source-lang",
"xh",
"--multiling-target-lang",
"en",
"--multiling-train-source-text-file",
os.path.join(data_dir, "train.xhen.xh"),
"--multiling-train-target-text-file",
os.path.join(data_dir, "train.xhen.en"),
"--multiling-eval-source-text-file",
os.path.join(data_dir, "tune.xhen.xh"),
"--multiling-eval-target-text-file",
os.path.join(data_dir, "tune.xhen.en"),
"--multiling-source-lang",
"zu",
"--multiling-target-lang",
"en",
"--multiling-train-source-text-file",
os.path.join(data_dir, "train.zuen.zu"),
"--multiling-train-target-text-file",
os.path.join(data_dir, "train.zuen.en"),
"--multiling-eval-source-text-file",
os.path.join(data_dir, "tune.zuen.zu"),
"--multiling-eval-target-text-file",
os.path.join(data_dir, "tune.zuen.en"),
"--multiling-source-lang",
"en",
"--multiling-target-lang",
"xh",
"--multiling-train-source-text-file",
os.path.join(data_dir, "train.xhen.en"),
"--multiling-train-target-text-file",
os.path.join(data_dir, "train.xhen.xh"),
"--multiling-eval-source-text-file",
os.path.join(data_dir, "tune.xhen.en"),
"--multiling-eval-target-text-file",
os.path.join(data_dir, "tune.xhen.xh"),
# set these to empty to satisfy argument validation
"--train-source-text-file",
"",
"--train-target-text-file",
"",
"--eval-source-text-file",
"",
"--eval-target-text-file",
"",
],
)
for langpair, src, tgt in [
("xhen", "xh", "en"),
("zuen", "zu", "en"),
("xhen", "en", "xh"),
]:
generate_main(
data_dir,
[
"--task",
"pytorch_translate_multilingual",
"--multiling-source-lang",
src,
"--multiling-target-lang",
tgt,
"--source-vocab-file",
os.path.join(data_dir,
f"dictionary-src-{src}.txt"),
"--target-vocab-file",
os.path.join(data_dir,
f"dictionary-trg-{tgt}.txt"),
"--source-text-file",
os.path.join(data_dir, f"tune.{langpair}.{src}"),
"--target-text-file",
os.path.join(data_dir, f"tune.{langpair}.{tgt}"),
],
)