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"],
)
