Initial Commit

.gitignore

@@ -0,0 +1 @@
+.idea

README.md

@@ -0,0 +1,188 @@
+```bash
+conda install pytorch torchvision torchaudio pytorch-cuda=11.6 -c pytorch -c nvidia
+```
+
+# Contrastive Learning for Many-to-many Multilingual Neural Machine Translation(mCOLT/mRASP2), ACL2021
+The code for training mCOLT/mRASP2, a multilingual neural machine translation training method, implemented based on [fairseq](https://github.com/pytorch/fairseq). 
+
+**mRASP2**: [paper](https://arxiv.org/abs/2105.09501) [blog](https://medium.com/@panxiao1994/mrasp2-multilingual-nmt-advances-via-contrastive-learning-ac8c4c35d63)
+
+**mRASP**: [paper](https://www.aclweb.org/anthology/2020.emnlp-main.210.pdf),
+[code](https://github.com/linzehui/mRASP)
+
+---
+## News
+We have released two versions, this version is the original one. In this implementation:
+- You should first merge all data, by pre-pending language token before each sentence to indicate the language.
+- AA/RAS muse be done off-line (before binarize), check [this toolkit](https://github.com/linzehui/mRASP/blob/master/preprocess).
+
+**New implementation**: https://github.com/PANXiao1994/mRASP2/tree/new_impl
+
+* Acknowledgement: This work is supported by [Bytedance](https://bytedance.com). We thank [Chengqi](https://github.com/zhaocq-nlp) for uploading all files and checkpoints.
+
+## Introduction
+
+mRASP2/mCOLT, representing multilingual Contrastive Learning for Transformer, is a multilingual neural machine translation model that supports complete many-to-many multilingual machine translation. It employs both parallel corpora and multilingual corpora in a unified training framework. For detailed information please refer to the paper.  
+
+![img.png](docs/img.png)
+
+## Pre-requisite
+```bash
+pip install -r requirements.txt
+# install fairseq
+git clone https://github.com/pytorch/fairseq
+cd fairseq
+pip install --editable ./
+```
+
+## Training Data and Checkpoints
+We release our preprocessed training data and checkpoints in the following.
+### Dataset
+
+We merge 32 English-centric language pairs, resulting in 64 directed translation pairs in total. The original 32 language pairs corpus contains about 197M pairs of sentences. We get about 262M pairs of sentences after applying RAS, since we keep both the original sentences and the substituted sentences. We release both the original dataset and dataset after applying RAS.
+
+| Dataset | #Pair |
+| --- | --- |
+| [32-lang-pairs-TRAIN](https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/acl2021/mrasp2/bin_parallel/download.sh) | 197603294 |
+| [32-lang-pairs-RAS-TRAIN](https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/acl2021/mrasp2/bin_parallel_ras/download.sh) | 262662792 |
+| [mono-split-a](https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/acl2021/mrasp2/bin_mono_split_a/download.sh) | - |
+| [mono-split-b](https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/acl2021/mrasp2/bin_mono_split_b/download.sh) | - |
+| [mono-split-c](https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/acl2021/mrasp2/bin_mono_split_c/download.sh) | - |
+| [mono-split-d](https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/acl2021/mrasp2/bin_mono_split_d/download.sh) | - |
+| [mono-split-e](https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/acl2021/mrasp2/bin_mono_split_e/download.sh) | - |
+| [mono-split-de-fr-en](https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/acl2021/mrasp2/bin_mono_de_fr_en/download.sh) | - |
+| [mono-split-nl-pl-pt](https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/acl2021/mrasp2/bin_mono_nl_pl_pt/download.sh) | - |
+| [32-lang-pairs-DEV-en-centric](https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/acl2021/mrasp2/bin_dev_en_centric/download.sh) | - |
+| [32-lang-pairs-DEV-many-to-many](https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/acl2021/mrasp2/bin_dev_m2m/download.sh) | - |
+| [Vocab](https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/acl2021/mrasp2/bpe_vocab) | - |
+| [BPE Code](https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/emnlp2020/mrasp/pretrain/dataset/codes.bpe.32000) | - |
+
+
+### Checkpoints & Results
+* **Please note that the provided checkpoint is sightly different from that in the paper.** In the following sections, we report the results of the provided checkpoints.
+
+#### English-centric Directions
+We report **tokenized BLEU** in the following table. Please click the model links to download. It is in pytorch format. (check eval.sh for details)
+
+|Models  | [6e6d-no-mono](https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/acl2021/mrasp2/6e6d_no_mono.pt) | [12e12d-no-mono](https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/acl2021/mrasp2/12e12d_no_mono.pt) | [12e12d](https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/acl2021/mrasp2/12e12d_last.pt) |
+| --- | --- | --- | --- |
+| en2cs/wmt16 | 21.0 | 22.3 | 23.8 |
+| cs2en/wmt16 | 29.6 | 32.4 | 33.2 |
+| en2fr/wmt14 | 42.0 | 43.3 | 43.4 |
+| fr2en/wmt14 | 37.8 | 39.3 | 39.5 |
+| en2de/wmt14 | 27.4 | 29.2 | 29.5 |
+| de2en/wmt14 | 32.2 | 34.9 | 35.2 |
+| en2zh/wmt17 | 33.0 | 34.9 | 34.1 |
+| zh2en/wmt17 | 22.4 | 24.0 | 24.4 |
+| en2ro/wmt16 | 26.6 | 28.1 | 28.7 |
+| ro2en/wmt16 | 36.8 | 39.0 | 39.1 |
+| en2tr/wmt16 | 18.6 | 20.3 | 21.2 |
+| tr2en/wmt16 | 22.2 | 25.5 | 26.1 |
+| en2ru/wmt19 | 17.4 | 18.5 | 19.2 |
+| ru2en/wmt19 | 22.0 | 23.2 | 23.6 |
+| en2fi/wmt17 | 20.2 | 22.1 | 22.9 |
+| fi2en/wmt17 | 26.1 | 29.5 | 29.7 |
+| en2es/wmt13 | 32.8 | 34.1 | 34.6 |
+| es2en/wmt13 | 32.8 | 34.6 | 34.7 |
+| en2it/wmt09 | 28.9 | 30.0 | 30.8 |
+| it2en/wmt09 | 31.4 | 32.7 | 32.8 |
+
+#### Unsupervised Directions
+We report **tokenized BLEU** in the following table. (check eval.sh for details)
+
+| | 12e12d |
+| --- | --- |
+| en2pl/wmt20 | 6.2 |
+| pl2en/wmt20 | 13.5 |
+| en2nl/iwslt14 | 8.8 |
+| nl2en/iwslt14 | 27.1 |
+| en2pt/opus100 | 18.9 |
+| pt2en/opus100 | 29.2 |
+
+#### Zero-shot Directions
+* row: source language
+* column: target language
+We report **[sacreBLEU](https://github.com/mozilla/sacreBLEU)** in the following table.
+
+| 12e12d  | ar | zh | nl | fr | de | ru |
+| --- | --- | --- | --- | --- | --- | --- |
+| ar | - | 32.5 | 3.2 | 22.8 | 11.2 | 16.7 |
+| zh | 6.5 | - | 1.9 | 32.9 | 7.6 | 23.7 |
+| nl | 1.7 | 8.2 | - | 7.5 | 10.2 | 2.9 |
+| fr | 6.2 | 42.3 | 7.5 | - | 18.9 | 24.4 |
+| de | 4.9 | 21.6 | 9.2 | 24.7 | - | 14.4 |
+| ru | 7.1 | 40.6 | 4.5 | 29.9 | 13.5 | - |
+
+## Training
+```bash
+export NUM_GPU=4 && bash train_w_mono.sh ${model_config}
+```
+* We give example of `${model_config}` in `${PROJECT_REPO}/examples/configs/parallel_mono_12e12d_contrastive.yml`
+
+## Inference
+* You must pre-pend the corresponding language token to the source side before binarize the test data.
+```bash
+fairseq-generate ${test_path} \
+    --user-dir ${repo_dir}/mcolt \
+    -s ${src} \
+    -t ${tgt} \
+    --skip-invalid-size-inputs-valid-test \
+    --path ${ckpts} \
+    --max-tokens ${batch_size} \
+    --task translation_w_langtok \
+    ${options} \
+    --lang-prefix-tok "LANG_TOK_"`echo "${tgt} " | tr '[a-z]' '[A-Z]'` \
+    --max-source-positions ${max_source_positions} \
+    --max-target-positions ${max_target_positions} \
+    --nbest 1 | grep -E '[S|H|P|T]-[0-9]+' > ${final_res_file}
+
+python fairseq/fairseq_cli/preprocess.py --dataset-impl raw --srcdict ckpt/bpe_vocab --tgtdict ckpt/bpe_vocab --testpref test/input -s zh -t en
+
+python fairseq/fairseq_cli/interactive.py /mnt/data2/siqiouyang/demo/mRASP2/data-bin \
+  --user-dir mcolt \
+  -s en \
+  -t zh \
+  --skip-invalid-size-inputs-valid-test \
+  --path ckpt/12e12d_last.pt \
+  --max-tokens 1024 \
+  --task translation_w_langtok \
+  --lang-prefix-tok "LANG_TOK_"`echo "zh " | tr '[a-z]' '[A-Z]'` \
+  --max-source-positions 1024 \
+  --max-target-positions 1024 \
+  --nbest 1 \
+  --bpe subword_nmt \
+  --bpe-codes ckpt/codes.bpe.32000 \
+  --post-process --tokenizer moses \
+  --input ./test/input.en | grep -E '[D]-[0-9]+' > test/output.zh.no_bpe.moses
+  
+python3 ${repo_dir}/scripts/utils.py ${res_file} ${ref_file} || exit 1;
+```
+
+## Synonym dictionaries
+We use the bilingual synonym dictionaries provised by [MUSE](https://github.com/facebookresearch/MUSE).
+
+We generate multilingual synonym dictionaries using [this script](https://github.com/linzehui/mRASP/blob/master/preprocess/tools/ras/multi_way_word_graph.py), and apply
+RAS using [this script](https://github.com/linzehui/mRASP/blob/master/preprocess/tools/ras/random_alignment_substitution_w_multi.sh).
+
+| Description | File | Size |
+| --- | --- | --- |
+| dep=1 | [synonym_dict_raw_dep1](https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/acl2021/mrasp2/synonym_dict_raw_dep1) | 138.0 M |
+| dep=2 | [synonym_dict_raw_dep2](https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/acl2021/mrasp2/synonym_dict_raw_dep2) | 1.6 G |
+| dep=3 | [synonym_dict_raw_dep3](https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/acl2021/mrasp2/synonym_dict_raw_dep3) | 2.2 G |
+
+## Contact
+Please contact me via e-mail `panxiao94@163.com` or via [wechat/zhihu](https://fork-ball-95c.notion.site/mRASP2-4e9b3450d5aa4137ae1a2c46d5f3c1fa) or join [the slack group](https://mrasp2.slack.com/join/shared_invite/zt-10k9710mb-MbDHzDboXfls2Omd8cuWqA)!
+
+## Citation
+Please cite as:
+```
+@inproceedings{mrasp2,
+  title = {Contrastive Learning for Many-to-many Multilingual Neural Machine Translation},
+  author= {Xiao Pan and
+           Mingxuan Wang and
+           Liwei Wu and
+           Lei Li},
+  booktitle = {Proceedings of ACL 2021},
+  year = {2021},
+}
+```

data-bin/preprocess.log

@@ -0,0 +1,4 @@
+Namespace(no_progress_bar=False, log_interval=100, log_format=None, log_file=None, aim_repo=None, aim_run_hash=None, tensorboard_logdir=None, wandb_project=None, azureml_logging=False, seed=1, cpu=False, tpu=False, bf16=False, memory_efficient_bf16=False, fp16=False, memory_efficient_fp16=False, fp16_no_flatten_grads=False, fp16_init_scale=128, fp16_scale_window=None, fp16_scale_tolerance=0.0, on_cpu_convert_precision=False, min_loss_scale=0.0001, threshold_loss_scale=None, amp=False, amp_batch_retries=2, amp_init_scale=128, amp_scale_window=None, user_dir=None, empty_cache_freq=0, all_gather_list_size=16384, model_parallel_size=1, quantization_config_path=None, profile=False, reset_logging=False, suppress_crashes=False, use_plasma_view=False, plasma_path='/tmp/plasma', criterion='cross_entropy', tokenizer=None, bpe=None, optimizer=None, lr_scheduler='fixed', scoring='bleu', task='translation', source_lang='en', target_lang='zh', trainpref=None, validpref=None, testpref='test/input', align_suffix=None, destdir='data-bin', thresholdtgt=0, thresholdsrc=0, tgtdict='ckpt/bpe_vocab', srcdict='ckpt/bpe_vocab', nwordstgt=-1, nwordssrc=-1, alignfile=None, dataset_impl='raw', joined_dictionary=False, only_source=False, padding_factor=8, workers=1, dict_only=False)
+Wrote preprocessed data to data-bin
+Namespace(no_progress_bar=False, log_interval=100, log_format=None, log_file=None, aim_repo=None, aim_run_hash=None, tensorboard_logdir=None, wandb_project=None, azureml_logging=False, seed=1, cpu=False, tpu=False, bf16=False, memory_efficient_bf16=False, fp16=False, memory_efficient_fp16=False, fp16_no_flatten_grads=False, fp16_init_scale=128, fp16_scale_window=None, fp16_scale_tolerance=0.0, on_cpu_convert_precision=False, min_loss_scale=0.0001, threshold_loss_scale=None, amp=False, amp_batch_retries=2, amp_init_scale=128, amp_scale_window=None, user_dir=None, empty_cache_freq=0, all_gather_list_size=16384, model_parallel_size=1, quantization_config_path=None, profile=False, reset_logging=False, suppress_crashes=False, use_plasma_view=False, plasma_path='/tmp/plasma', criterion='cross_entropy', tokenizer=None, bpe=None, optimizer=None, lr_scheduler='fixed', scoring='bleu', task='translation', source_lang='zh', target_lang='en', trainpref=None, validpref=None, testpref='test/input', align_suffix=None, destdir='data-bin', thresholdtgt=0, thresholdsrc=0, tgtdict='ckpt/bpe_vocab', srcdict='ckpt/bpe_vocab', nwordstgt=-1, nwordssrc=-1, alignfile=None, dataset_impl='raw', joined_dictionary=False, only_source=False, padding_factor=8, workers=1, dict_only=False)
+Wrote preprocessed data to data-bin

data-bin/test.en-zh.en

@@ -0,0 +1 @@
+LANG_TOK_EN Hello my friend!

data-bin/test.en-zh.zh

@@ -0,0 +1 @@
+LANG_TOK_ZH

data-bin/test.zh-en.en

@@ -0,0 +1 @@
+LANG_TOK_EN Hello my friend!

data-bin/test.zh-en.zh

@@ -0,0 +1 @@
+LANG_TOK_ZH 你好！

eval.sh

@@ -0,0 +1,166 @@
+#!/usr/bin/env bash
+
+# repo_dir: root directory of the project
+repo_dir="$( cd "$( dirname "$0" )" && pwd )"
+echo "==== Working directory: ====" >&2
+echo "${repo_dir}" >&2
+echo "============================" >&2
+
+
+test_config=$1
+source ${repo_dir}/scripts/load_config.sh ${test_config} ${repo_dir}
+model_dir=$2
+choice=$3  # all|best|last
+
+model_dir=${repo_dir}/model
+data_dir=${repo_dir}/data
+res_path=${model_dir}/results
+
+mkdir -p ${model_dir} ${data_dir} ${res_path}
+
+testset_name=data_testset_1_name
+testset_path=data_testset_1_path
+testset_ref=data_testset_1_ref
+testset_direc=data_testset_1_direction
+i=1
+testsets=""
+while [[ ! -z ${!testset_path} && ! -z ${!testset_direc} ]]; do
+    dataname=${!testset_name}
+    mkdir -p ${data_dir}/${!testset_direc}/${dataname} ${data_dir}/ref/${!testset_direc}/${dataname}
+    cp ${!testset_path}/* ${data_dir}/${!testset_direc}/${dataname}/
+    cp ${!testset_ref}/* ${data_dir}/ref/${!testset_direc}/${dataname}/
+    if [[ $testsets == "" ]]; then
+        testsets=${!testset_direc}/${dataname}
+    else
+        testsets=${testsets}:${!testset_direc}/${dataname}
+    fi
+    i=$((i+1))
+    testset_name=testset_${i}_name
+    testset_path=testset_${i}_path
+    testset_ref=testset_${i}_ref
+    testset_direc=testset_${i}_direction
+done
+
+IFS=':' read -r -a testset_list <<< ${testsets}
+
+
+bleu () {
+    src=$1
+    tgt=$2
+    res_file=$3
+    ref_file=$4
+    if [[ -f ${res_file} ]]; then
+        f_dirname=`dirname ${res_file}`
+        python3 ${repo_dir}/scripts/utils.py ${res_file} ${ref_file} || exit 1;
+        input_file="${f_dirname}/hypo.out.nobpe"
+        output_file="${f_dirname}/hypo.out.nobpe.final"
+        # form command
+        cmd="cat ${input_file}"
+        lang_token="LANG_TOK_"`echo "${tgt} " | tr '[a-z]' '[A-Z]'`
+        if [[ $tgt == "fr" ]]; then
+            cmd=$cmd" | sed -Ee 's/\"([^\"]*)\"/« \1 »/g'"
+        elif [[ $tgt == "zh" ]]; then
+            tokenizer="zh"
+        elif [[ $tgt == "ja" ]]; then
+            tokenizer="ja-mecab"
+        fi
+        [[ -z $tokenizer ]] && tokenizer="none"
+        cmd=$cmd" | sed -e s'|${lang_token} ||g' > ${output_file}"
+        eval $cmd || { echo "$cmd FAILED !"; exit 1; }
+        cat ${output_file} | sacrebleu -l ${src}-${tgt} -tok $tokenizer --short "${f_dirname}/ref.out" | awk '{print $3}'
+    else
+        echo "${res_file} not exist!" >&2 && exit 1;
+    fi
+}
+
+# monitor
+# ${ckptname}/${direction}/${testname}/orig.txt
+(inotifywait -r -m -e close_write ${res_path} |
+while read path action file; do
+    if [[ "$file" =~ .*txt$ ]]; then
+        tmp_str="${path%/*}"
+        testname="${tmp_str##*/}"
+        tmp_str="${tmp_str%/*}"
+        direction="${tmp_str##*/}"
+        tmp_str="${tmp_str%/*}"
+        ckptname="${tmp_str##*/}"
+        src_lang="${direction%2*}"
+        tgt_lang="${direction##*2}"
+        res_file=$path$file
+        ref_file=${data_dir}/ref/${direction}/${testname}/dev.${tgt_lang}
+        bleuscore=`bleu ${src_lang} ${tgt_lang} ${res_file} ${ref_file}`
+        bleu_str="$(date "+%Y-%m-%d %H:%M:%S")\t${ckptname}\t${direction}/${testname}\t$bleuscore"
+        echo -e ${bleu_str}  # to stdout
+        echo -e ${bleu_str} >> ${model_dir}/summary.log
+    fi
+done) &
+
+
+if [[ ${choice} == "all" ]]; then
+    filelist=`ls -la ${model_dir} | sort -k6,7 -r | awk '{print $NF}' | grep .pt$ | tr '\n' ' '`
+elif [[ ${choice} == "best" ]]; then
+    filelist="${model_dir}/checkpoint_best.pt"
+elif [[ ${choice} == "last" ]]; then
+    filelist="${model_dir}/checkpoint_last.pt"
+else
+    echo "invalid choice!" && exit 2;
+fi
+
+N=${NUM_GPU}
+#export CUDA_VISIBLE_DEVICES=$(seq -s ',' 0 $(($N - 1)) )
+
+
+infer_test () {
+    test_path=$1
+    ckpts=$2
+    gpu=$3
+    final_res_file=$4
+    src=$5
+    tgt=$6
+    gpu_cmd="CUDA_VISIBLE_DEVICES=$gpu "
+    lang_token="LANG_TOK_"`echo "${tgt} " | tr '[a-z]' '[A-Z]'`
+    [[ -z ${max_source_positions} ]] && max_source_positions=1024
+    [[ -z ${max_target_positions} ]] && max_target_positions=1024
+    command=${gpu_cmd}"fairseq-generate ${test_path} \
+    --user-dir ${repo_dir}/mcolt \
+    -s ${src} \
+    -t ${tgt} \
+    --skip-invalid-size-inputs-valid-test \
+    --path ${ckpts} \
+    --max-tokens 1024 \
+    --task translation_w_langtok \
+    ${options} \
+    --lang-prefix-tok ${lang_token} \
+    --max-source-positions ${max_source_positions} \
+    --max-target-positions ${max_target_positions} \
+    --nbest 1 | grep -E '[S|H|P|T]-[0-9]+' > ${final_res_file}
+    "
+    echo "$command"
+}
+
+export -f infer_test
+i=0
+(for ckpt in ${filelist}
+do
+    for testset in "${testset_list[@]}"
+    do
+        ckptbase=`basename $ckpt`
+        ckptname="${ckptbase%.*}"
+        direction="${testset%/*}"
+        testname="${testset##*/}"
+        src_lang="${direction%2*}"
+        tgt_lang="${direction##*2}"
+
+        ((i=i%N)); ((i++==0)) && wait
+        test_path=${data_dir}/${testset}
+
+        echo "-----> "${ckptname}" | "${direction}/$testname" <-----" >&2
+        if [[ ! -d ${res_path}/${ckptname}/${direction}/${testname} ]]; then
+            mkdir -p ${res_path}/${ckptname}/${direction}/${testname}
+        fi
+        final_res_file="${res_path}/${ckptname}/${direction}/${testname}/orig.txt"
+        command=`infer_test ${test_path} ${model_dir}/${ckptname}.pt $((i-1)) ${final_res_file} ${src_lang} ${tgt_lang}`
+        echo "${command}"
+        eval $command &
+    done
+done)

examples/configs/eval_benchmarks.yml

@@ -0,0 +1,80 @@
+data_testset_1:
+  direction: en2de
+  name: wmt14
+  path: data/binarized/en_de/en2de/wmt14
+  ref: data/dev/en2de/wmt14
+data_testset_10:
+  direction: ru2en
+  name: newstest2019
+  path: data/binarized/en_ru/ru2en/newstest2019
+  ref: data/dev/ru2en/newstest2019
+data_testset_11:
+  direction: en2fi
+  name: newstest2017
+  path: data/binarized/en_fi/en2fi/newstest2017
+  ref: data/dev/en2fi/newstest2017
+data_testset_12:
+  direction: fi2en
+  name: newstest2017
+  path: data/binarized/en_fi/fi2en/newstest2017
+  ref: data/dev/fi2en/newstest2017
+data_testset_13:
+  direction: en2cs
+  name: newstest2016
+  path: data/binarized/en_cs/en2cs/newstest2016
+  ref: data/dev/en2cs/newstest2016
+data_testset_14:
+  direction: cs2en
+  name: newstest2016
+  path: data/binarized/en_cs/cs2en/newstest2016
+  ref: data/dev/cs2en/newstest2016
+data_testset_15:
+  direction: en2et
+  name: newstest2018
+  path: data/binarized/en_et/en2et/newstest2018
+  ref: data/dev/en2et/newstest2018
+data_testset_16:
+  direction: et2en
+  name: newstest2018
+  path: data/binarized/en_et/et2en/newstest2018
+  ref: data/dev/et2en/newstest2018
+data_testset_2:
+  direction: de2en
+  name: wmt14
+  path: data/binarized/en_de/de2en/wmt14
+  ref: data/dev/de2en/wmt14
+data_testset_3:
+  direction: en2fr
+  name: newstest2014
+  path: data/binarized/en_fr/en2fr/newstest2014
+  ref: data/dev/en2fr/newstest2014
+data_testset_4:
+  direction: fr2en
+  name: newstest2014
+  path: data/binarized/en_fr/fr2en/newstest2014
+  ref: data/dev/fr2en/newstest2014
+data_testset_5:
+  direction: en2ro
+  name: wmt16
+  path: data/binarized/en_ro/en_ro/wmt16
+  ref: data/dev/en_ro/wmt16
+data_testset_6:
+  direction: ro2en
+  name: wmt16
+  path: data/binarized/en_ro/en_ro/wmt16
+  ref: data/dev/en_ro/wmt16
+data_testset_7:
+  direction: en2zh
+  name: wmt17
+  path: data/binarized/en_zh/en2zh/wmt17
+  ref: data/dev/en2zh/wmt17
+data_testset_8:
+  direction: zh2en
+  name: wmt17
+  path: data/binarized/en_zh/zh2en/wmt17
+  ref: data/dev/zh2en/wmt17
+data_testset_9:
+  direction: en2ru
+  name: newstest2019
+  path: data/binarized/en_ru/en2ru/newstest2019
+  ref: data/dev/en2ru/newstest2019

examples/configs/parallel_mono_12e12d_contrastive.yml

@@ -0,0 +1,44 @@
+model_dir: model/pretrain/lab/multilingual/l2r/multi_bpe32k/parallel_mono_contrastive_1/transformer_big_t2t_12e12d
+data_1: data/multilingual/bin/merged_deduped_ras
+data_mono_1: data/multilingual/bin/mono_only/splitaa
+data_mono_2: data/multilingual/bin/mono_only/splitab
+data_mono_3: data/multilingual/bin/mono_only/splitac
+data_mono_4: data/multilingual/bin/mono_only/splitad
+data_mono_5: data/multilingual/bin/mono_only/splitae
+data_mono_6: data/multilingual/bin/mono_only/mono_de_fr_en
+data_mono_7: data/multilingual/bin/mono_only/mono_nl_pl_pt
+source_lang: src
+target_lang: trg
+task: translation_w_mono
+parallel_ratio: 0.2
+mono_ratio: 0.07
+arch: transformer_big_t2t_12e12d
+share_all_embeddings: true
+encoder_learned_pos: true
+decoder_learned_pos: true
+max_source_positions: 1024
+max_target_positions: 1024
+dropout: 0.1
+criterion: label_smoothed_cross_entropy_with_contrastive
+contrastive_lambda: 1.0
+temperature: 0.1
+lr: 0.0003
+clip_norm: 10.0
+optimizer: adam
+adam_eps: 1e-06
+weight_decay: 0.01
+warmup_updates: 10000
+label_smoothing: 0.1
+lr_scheduler: polynomial_decay
+min_lr: -1
+max_tokens: 1536
+update_freq: 30
+max_update: 5000000
+no_scale_embedding: true
+layernorm_embedding: true
+save_interval_updates: 2000
+skip_invalid_size_inputs_valid_test: true
+log_interval: 500
+num_workers: 1
+fp16: true
+seed: 33122

mcolt/__init__.py

@@ -0,0 +1,4 @@
+from .arches import *
+from . criterions import *
+from .data import *
+from .tasks import *

mcolt/arches/__init__.py

@@ -0,0 +1 @@
+from .transformer import *

mcolt/arches/transformer.py

@@ -0,0 +1,380 @@
+from fairseq.models import register_model_architecture
+
+
+@register_model_architecture('transformer', 'transformer_bigger')
+def transformer_bigger(args):
+    args.attention_dropout = getattr(args, 'attention_dropout', 0.3)
+    args.activation_dropout = getattr(args, 'activation_dropout', 0.3)
+    args.dropout = getattr(args, 'dropout', 0.1)
+    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 15000)
+    args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 15000)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+    
+    
+@register_model_architecture('transformer', 'transformer_bigger_16384')
+def transformer_bigger_16384(args):
+    args.attention_dropout = getattr(args, 'attention_dropout', 0.1)
+    args.activation_dropout = getattr(args, 'activation_dropout', 0.1)
+    args.dropout = getattr(args, 'dropout', 0.1)
+    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 16384)
+    args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 16384)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'transformer_bigger_no_share')
+def transformer_bigger_no_share(args):
+    args.attention_dropout = getattr(args, 'attention_dropout', 0.3)
+    args.activation_dropout = getattr(args, 'activation_dropout', 0.3)
+    args.dropout = getattr(args, 'dropout', 0.1)
+    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 15000)
+    args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 15000)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', False)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'transformer_deeper')
+def transformer_deeper(args):
+    args.encoder_layers = getattr(args, 'encoder_layers', 15)
+    args.dense = False
+    args.bottleneck_component = getattr(args, 'bottleneck_component', 'mean_pool')
+    args.attention_dropout = getattr(args, 'attention_dropout', 0.1)
+    args.activation_dropout = getattr(args, 'activation_dropout', 0.1)
+    args.dropout = getattr(args, 'dropout', 0.1)
+    # args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 15000)
+    # args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 15000)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'transformer_deeper_no_share')
+def transformer_deeper_no_share(args):
+    args.encoder_layers = getattr(args, 'encoder_layers', 15)
+    args.dense = False
+    args.bottleneck_component = getattr(args, 'bottleneck_component', 'mean_pool')
+    args.attention_dropout = getattr(args, 'attention_dropout', 0.1)
+    args.activation_dropout = getattr(args, 'activation_dropout', 0.1)
+    args.dropout = getattr(args, 'dropout', 0.1)
+    # args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 15000)
+    # args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 15000)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', False)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'transformer_deeper_dense')
+def transformer_deeper_no_share(args):
+    args.encoder_layers = getattr(args, 'encoder_layers', 15)
+    args.dense = True
+    args.bottleneck_component = 'mean_pool'
+    args.attention_dropout = getattr(args, 'attention_dropout', 0.1)
+    args.activation_dropout = getattr(args, 'activation_dropout', 0.1)
+    args.dropout = getattr(args, 'dropout', 0.1)
+    # args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 15000)
+    # args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 15000)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'transformer_deeper_dense_no_share')
+def transformer_deeper_no_share(args):
+    args.encoder_layers = getattr(args, 'encoder_layers', 15)
+    args.dense = True
+    args.bottleneck_component = 'mean_pool'
+    args.attention_dropout = getattr(args, 'attention_dropout', 0.1)
+    args.activation_dropout = getattr(args, 'activation_dropout', 0.1)
+    args.dropout = getattr(args, 'dropout', 0.1)
+    # args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 15000)
+    # args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 15000)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', False)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'transformer_big')
+def transformer_big(args):
+    args.dropout = getattr(args, 'dropout', 0.1)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'transformer_big_emb512')
+def transformer_big_emb512(args):
+    args.dropout = getattr(args, 'dropout', 0.1)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 512)
+    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 512)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'transformer_big_no_share')
+def transformer_big_no_share(args):
+    args.dropout = getattr(args, 'dropout', 0.1)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', False)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'transformer_big_16e4d')
+def transformer_big_16e4d(args):
+    args.dropout = getattr(args, 'dropout', 0.2)
+    args.encoder_layers = getattr(args, 'encoder_layers', 16)
+    args.decoder_layers = getattr(args, 'decoder_layers', 4)
+    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 1024)
+    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 1024)
+    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 4096)
+    args.encoder_attention_heads = getattr(args, 'encoder_attention_heads', 16)
+    args.decoder_attention_heads = getattr(args, 'decoder_attention_heads', 16)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'transformer_big_16e6d')
+def transformer_big_16e6d(args):
+    args.dropout = getattr(args, 'dropout', 0.2)
+    args.encoder_layers = getattr(args, 'encoder_layers', 16)
+    args.decoder_layers = getattr(args, 'decoder_layers', 6)
+    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 1024)
+    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 1024)
+    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 4096)
+    args.encoder_attention_heads = getattr(args, 'encoder_attention_heads', 16)
+    args.decoder_attention_heads = getattr(args, 'decoder_attention_heads', 16)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'transformer_base')
+def transformer_bigger(args):
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    from fairseq.models.transformer import transformer_wmt_en_de
+    transformer_wmt_en_de(args)
+
+
+@register_model_architecture('transformer', 'transformer_mid_50e6d')
+def transformer_mid_50e6d(args):
+    args.dropout = getattr(args, 'dropout', 0.1)
+    args.encoder_layers = getattr(args, 'encoder_layers', 50)
+    args.decoder_layers = getattr(args, 'decoder_layers', 6)
+    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 768)
+    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 768)
+    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 3072)
+    args.encoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 3072)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'transformer_big_t2t_12e12d')
+def transformer_big_t2t_12e12d(args):
+    args.dropout = getattr(args, 'dropout', 0.1)
+    args.encoder_layers = getattr(args, 'encoder_layers', 12)
+    args.decoder_layers = getattr(args, 'decoder_layers', 12)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'mix_transformer_mid_50e6d')
+def mix_transformer_mid_50e6d(args):
+    args.mix_prepost_norm  = getattr(args, "mix_prepost_norm", True)
+    args.dropout = getattr(args, 'dropout', 0.1)
+    args.encoder_layers = getattr(args, 'encoder_layers', 50)
+    args.decoder_layers = getattr(args, 'decoder_layers', 6)
+    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 768)
+    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 768)
+    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 3072)
+    args.encoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 3072)
+    args.encoder_normalize_before = getattr(args, "encoder_normalize_before", False)
+    args.decoder_normalize_before = getattr(args, "decoder_normalize_before", False)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    args.mix_type = getattr(args, "mix_type", "learnable")
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 're_zero_transformer_mid_50e6d')
+def re_zero_transformer_mid_50e6d(args):
+    args.dropout = getattr(args, 'dropout', 0.1)
+    args.encoder_layers = getattr(args, 'encoder_layers', 50)
+    args.decoder_layers = getattr(args, 'decoder_layers', 6)
+    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 768)
+    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 768)
+    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 3072)
+    args.encoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 3072)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    args.re_zero = getattr(args, "re_zero", True)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+    
+
+@register_model_architecture('transformer', 'transformer_mid_50e3d_ed3072')
+def transformer_mid_50e3d_ed3072(args):
+    args.dropout = getattr(args, 'dropout', 0.1)
+    args.encoder_layers = getattr(args, 'encoder_layers', 50)
+    args.decoder_layers = getattr(args, 'decoder_layers', 3)
+    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 768)
+    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 768)
+    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 3072)
+    args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 3072)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'mix_transformer_mid_50e6d_3000fix_10000decay')
+def mix_transformer_mid_50e6d_3000fix_10000decay(args):
+    args.mix_prepost_norm  = getattr(args, "mix_prepost_norm", True)
+    args.mix_type = getattr(args, "mix_type", "step_moving")
+    args.pre_steps = getattr(args, "pre_steps", 3000)
+    args.change_steps = getattr(args, "change_steps", 10000)
+
+    args.dropout = getattr(args, 'dropout', 0.1)
+    args.encoder_layers = getattr(args, 'encoder_layers', 50)
+    args.decoder_layers = getattr(args, 'decoder_layers', 6)
+    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 768)
+    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 768)
+    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 3072)
+    args.encoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 3072)
+    args.encoder_normalize_before = getattr(args, "encoder_normalize_before", False)
+    args.decoder_normalize_before = getattr(args, "decoder_normalize_before", False)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'mix_transformer_mid_50e6d_7000fix_7000decay')
+def mix_transformer_mid_50e6d_3000fix_10000decay(args):
+    args.mix_prepost_norm = getattr(args, "mix_prepost_norm", True)
+    args.mix_type = getattr(args, "mix_type", "step_moving")
+    args.pre_steps = getattr(args, "pre_steps", 7000)
+    args.change_steps = getattr(args, "change_steps", 7000)
+
+    args.dropout = getattr(args, 'dropout', 0.1)
+    args.encoder_layers = getattr(args, 'encoder_layers', 50)
+    args.decoder_layers = getattr(args, 'decoder_layers', 6)
+    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 768)
+    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 768)
+    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 3072)
+    args.encoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 3072)
+    args.encoder_normalize_before = getattr(args, "encoder_normalize_before", False)
+    args.decoder_normalize_before = getattr(args, "decoder_normalize_before", False)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'transformer_mid_75e6d')
+def transformer_mid_75e6d(args):
+    args.dropout = getattr(args, 'dropout', 0.1)
+    args.encoder_layers = getattr(args, 'encoder_layers', 75)
+    args.decoder_layers = getattr(args, 'decoder_layers', 6)
+    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 768)
+    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 768)
+    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 3072)
+    args.encoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 3072)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'transformer_mid_25e6d')
+def transformer_mid_25e6d(args):
+    args.dropout = getattr(args, 'dropout', 0.1)
+    args.encoder_layers = getattr(args, 'encoder_layers', 25)
+    args.decoder_layers = getattr(args, 'decoder_layers', 6)
+    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 768)
+    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 768)
+    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 3072)
+    args.encoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 3072)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+    
+
+@register_model_architecture('transformer', 'transformer_mid_25e6d_ed3072')
+def transformer_mid_25e6d_ed3072(args):
+    args.dropout = getattr(args, 'dropout', 0.1)
+    args.encoder_layers = getattr(args, 'encoder_layers', 25)
+    args.decoder_layers = getattr(args, 'decoder_layers', 6)
+    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 768)
+    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 768)
+    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 3072)
+    args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 3072)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'transformer_mid_25e6d_e3072_d4096')
+def transformer_mid_25e6d_e3072_d4096(args):
+    args.dropout = getattr(args, 'dropout', 0.1)
+    args.encoder_layers = getattr(args, 'encoder_layers', 25)
+    args.decoder_layers = getattr(args, 'decoder_layers', 6)
+    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 768)
+    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 768)
+    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 3072)
+    args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 4096)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    # args.share_all_embeddings = getattr(args, 'share_all_embeddings', True)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+# def transformer_fixed_multihead(args):
+#     args.head_dim = getattr(args, 'head_dim', 128)
+#     from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+#     transformer_wmt_en_de_big_t2t(args)
+
+@register_model_architecture('transformer', 'transformer_fixed_multihead_base')
+def transformer_fixed_multihead_base(args):
+    args.head_dim = getattr(args, 'head_dim', 128)
+    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 512)
+    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 512)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'transformer_fixed_multihead_embed_1024_nhead_16_hdim_128')
+def transformer_fixed_multihead_embed_1024_nhead_16_hdim_128(args):
+    args.head_dim = getattr(args, 'head_dim', 128)
+    args.dropout = getattr(args, 'dropout', 0.1)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'transformer_fixed_multihead_embed_1024_nhead_16_hdim_256')
+def transformer_fixed_multihead_embed_1024_nhead_16_hdim_128(args):
+    args.head_dim = getattr(args, 'head_dim', 256)
+    args.dropout = getattr(args, 'dropout', 0.1)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'transformer_fh_16x128_layer_12')
+def transformer_fh_16x128_layer_12(args):
+    args.head_dim = getattr(args, 'head_dim', 128)
+    args.dropout = getattr(args, 'dropout', 0.1)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)
+
+
+@register_model_architecture('transformer', 'transformer_fh_16x256_layer_12')
+def transformer_fh_16x256_layer_12(args):
+    args.head_dim = getattr(args, 'head_dim', 256)
+    args.dropout = getattr(args, 'dropout', 0.1)
+    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
+    from fairseq.models.transformer import transformer_wmt_en_de_big_t2t
+    transformer_wmt_en_de_big_t2t(args)

mcolt/criterions/__init__.py

@@ -0,0 +1 @@
+from .label_smoothed_cross_entropy_with_contrastive import *

mcolt/criterions/label_smoothed_cross_entropy_with_contrastive.py

@@ -0,0 +1,123 @@
+import math
+
+from fairseq.criterions import register_criterion
+from fairseq.criterions.label_smoothed_cross_entropy import LabelSmoothedCrossEntropyCriterion
+from fairseq import metrics, utils
+
+from collections import deque
+
+import torch
+import torch.nn as nn
+
+
+@register_criterion("label_smoothed_cross_entropy_with_contrastive")
+class LabelSmoothedCrossEntropyCriterionWithContrastive(
+    LabelSmoothedCrossEntropyCriterion
+):
+    def __init__(self, task, sentence_avg, label_smoothing, ignore_prefix_size=0, report_accuracy=False,
+                 contrastive_lambda=0.0,
+                 temperature=1.0):
+        super().__init__(task, sentence_avg, label_smoothing, ignore_prefix_size, report_accuracy)
+        self.contrastive_lambda = contrastive_lambda
+        self.temperature = temperature
+    
+    @staticmethod
+    def add_args(parser):
+        LabelSmoothedCrossEntropyCriterion.add_args(parser)
+        parser.add_argument("--contrastive-lambda", type=float,
+                            default=0.0,
+                            help="The contrastive loss weight")
+        parser.add_argument("--temperature", type=float,
+                            default=1.0,)
+    
+    def swap_sample(self, sample):
+        target = sample["target"]
+        prev_output_tokens = sample["net_input"]["prev_output_tokens"]
+        src_tokens = torch.cat((prev_output_tokens[:, :1], sample["net_input"]['src_tokens']), dim=-1)
+        return {
+            "net_input": {
+                "src_tokens": target.contiguous(),
+                "src_lengths": (target != self.padding_idx).int().sum(dim=1),
+                "prev_output_tokens": src_tokens[:, :-1].contiguous()
+            },
+            'nsentences': sample['nsentences'],
+            'ntokens': utils.item((src_tokens[:, 1:] != self.padding_idx).int().sum().data),
+            "target": src_tokens[:, 1:].contiguous(),
+            "id": sample["id"],
+        }
+    
+    def forward(self, model, sample, reduce=True):
+        net_output = model(**sample["net_input"])
+        loss, nll_loss = self.compute_loss(model, net_output, sample, reduce=reduce)
+        encoder_out = model.encoder.forward(sample["net_input"]["src_tokens"], sample["net_input"]["src_lengths"]).encoder_out
+        reverse_sample = self.swap_sample(sample)
+        reversed_encoder_out = model.encoder.forward(reverse_sample["net_input"]["src_tokens"], reverse_sample["net_input"]["src_lengths"]).encoder_out
+        contrastive_loss = self.get_contrastive_loss(
+            encoder_out,
+            reversed_encoder_out,
+            sample,
+            reverse_sample,
+        )
+        sample_size = (
+            sample["target"].size(0) if self.sentence_avg else sample["ntokens"]
+        )
+        nsentences = sample["target"].size(0)
+        ntokens = sample["ntokens"]
+        all_loss = loss + contrastive_loss * self.contrastive_lambda * ntokens / nsentences
+        logging_output = {
+            "loss": loss.data,
+            "nll_loss": nll_loss.data,
+            "ntokens": ntokens,
+            "nsentences": nsentences,
+            "sample_size": sample_size,
+        }
+        if isinstance(contrastive_loss, int):
+            logging_output["contrastive_loss"] = 0
+        else:
+            logging_output["contrastive_loss"] = utils.item(contrastive_loss.data)
+        
+        return all_loss, sample_size, logging_output
+    
+    def similarity_function(self, ):
+        return nn.CosineSimilarity(dim=-1)
+    
+    def get_contrastive_loss(self, encoder_out1, encoder_out2, sample1, sample2):
+        
+        def _sentence_embedding(encoder_out, sample):
+            encoder_output = encoder_out.transpose(0, 1)
+            src_tokens = sample["net_input"]["src_tokens"]
+            mask = (src_tokens != self.padding_idx)
+            encoder_embedding = (encoder_output * mask.unsqueeze(-1)).sum(dim=1) / mask.float().sum(dim=1).unsqueeze(-1)  # [batch, hidden_size]
+            return encoder_embedding
+        
+        encoder_embedding1 = _sentence_embedding(encoder_out1, sample1)  # [batch, hidden_size]
+        encoder_embedding2 = _sentence_embedding(encoder_out2, sample2)  # [batch, hidden_size]
+        
+        batch_size = encoder_embedding2.shape[0]
+        feature_dim = encoder_embedding2.shape[1]
+        anchor_feature = encoder_embedding1
+        contrast_feature = encoder_embedding2
+        
+        similarity_function = self.similarity_function()
+        anchor_dot_contrast = similarity_function(anchor_feature.expand((batch_size, batch_size, feature_dim)),
+                                                  torch.transpose(contrast_feature.expand((batch_size, batch_size, feature_dim)), 0, 1))
+        
+        loss = -nn.LogSoftmax(0)(torch.div(anchor_dot_contrast, self.temperature)).diag().sum()
+        
+        return loss
+    
+    @classmethod
+    def reduce_metrics(cls, logging_outputs) -> None:
+        super().reduce_metrics(logging_outputs)
+        nsentences = utils.item(
+            sum(log.get("nsentences", 0) for log in logging_outputs)
+        )
+        contrastive_loss = utils.item(
+            sum(log.get("contrastive_loss", 0) for log in logging_outputs)
+        )
+        metrics.log_scalar(
+            "contrastive_loss",
+            contrastive_loss / nsentences / math.log(2),
+            nsentences,
+            round=3,
+        )

mcolt/data/__init__.py

@@ -0,0 +1 @@
+from .subsample_language_pair_dataset import SubsampleLanguagePairDataset

mcolt/data/subsample_language_pair_dataset.py

@@ -0,0 +1,124 @@
+from fairseq.data import BaseWrapperDataset, LanguagePairDataset, plasma_utils
+import numpy as np
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class SubsampleLanguagePairDataset(BaseWrapperDataset):
+    """Subsamples a given dataset by a specified ratio. Subsampling is done on the number of examples
+
+    Args:
+        dataset (~torch.utils.data.Dataset): dataset to subsample
+        size_ratio(float): the ratio to subsample to. must be between 0 and 1 (exclusive)
+    """
+    
+    def __init__(self, dataset, size_ratio, weights=None, replace=False, seed=0, epoch=1):
+        super().__init__(dataset)
+        assert size_ratio <= 1
+        self.actual_size = np.ceil(len(dataset) * size_ratio).astype(int)
+        logger.info(
+            "subsampled dataset from {} to {} (ratio={})".format(
+                len(self.dataset), self.actual_size, size_ratio
+            )
+        )
+        self.src_dict = self.dataset.src_dict
+        self.tgt_dict = self.dataset.tgt_dict
+        self.left_pad_source = self.dataset.left_pad_source
+        self.left_pad_target = self.dataset.left_pad_target
+        self.seed = seed
+        self._cur_epoch = None
+        self._cur_indices = None
+        self.replace = replace
+        if weights is None:
+            self.weights = None
+        else:
+            assert len(weights) == len(dataset)
+            weights_arr = np.array(weights, dtype=np.float64)
+            weights_arr /= weights_arr.sum()
+            self.weights = plasma_utils.PlasmaArray(weights_arr)
+        self.set_epoch(epoch)
+    
+    def __getitem__(self, index):
+        index = self._cur_indices.array[index]
+        return self.dataset.__getitem__(index)
+    
+    def __len__(self):
+        return self.actual_size
+    
+    @property
+    def sizes(self):
+        return self.dataset.sizes[self._cur_indices.array]
+
+    @property
+    def src_sizes(self):
+        return self.dataset.src_sizes[self._cur_indices.array]
+
+    @property
+    def tgt_sizes(self):
+        return self.dataset.tgt_sizes[self._cur_indices.array]
+    
+    @property
+    def name(self):
+        return self.dataset.name
+    
+    def num_tokens(self, index):
+        index = self._cur_indices.array[index]
+        return self.dataset.num_tokens(index)
+    
+    def size(self, index):
+        index = self._cur_indices.array[index]
+        return self.dataset.size(index)
+    
+    def ordered_indices(self):
+        if self.shuffle:
+            indices = np.random.permutation(len(self)).astype(np.int64)
+        else:
+            indices = np.arange(len(self), dtype=np.int64)
+        # sort by target length, then source length
+        if self.tgt_sizes is not None:
+            indices = indices[np.argsort(self.tgt_sizes[indices], kind="mergesort")]
+        return indices[np.argsort(self.src_sizes[indices], kind="mergesort")]
+    
+    def prefetch(self, indices):
+        indices = self._cur_indices.array[indices]
+        self.dataset.prefetch(indices)
+    
+    @property
+    def can_reuse_epoch_itr_across_epochs(self):
+        return False
+    
+    def set_epoch(self, epoch):
+        logger.info("SubsampleLanguagePairDataset.set_epoch: {}".format(epoch))
+        super().set_epoch(epoch)
+        
+        if epoch == self._cur_epoch:
+            return
+        
+        self._cur_epoch = epoch
+        
+        # Generate a weighted sample of indices as a function of the
+        # random seed and the current epoch.
+        
+        rng = np.random.RandomState(
+            [
+                42,  # magic number
+                self.seed % (2 ** 32),  # global seed
+                self._cur_epoch,  # epoch index
+            ]
+        )
+        self._cur_indices = plasma_utils.PlasmaArray(
+            rng.choice(
+                len(self.dataset),
+                self.actual_size,
+                replace=self.replace,
+                p=(None if self.weights is None else self.weights.array),
+            )
+        )
+        
+        logger.info(
+            "Dataset is sub-sampled: {} -> {}, first 3 ids are: {}".format(len(self.dataset), self.actual_size,
+                                                                           ",".join(
+                                                                               [str(_i) for _i in
+                                                                                self._cur_indices.array[:3]])))

mcolt/tasks/__init__.py

@@ -0,0 +1,2 @@
+from .translation_w_mono import *
+from .translation_w_langtok import *

mcolt/tasks/translation_w_langtok.py

@@ -0,0 +1,476 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import itertools
+import json
+import logging
+import os
+import torch
+from argparse import Namespace
+
+import numpy as np
+from fairseq import metrics, options, utils
+from fairseq.data import (
+    AppendTokenDataset,
+    ConcatDataset,
+    LanguagePairDataset,
+    PrependTokenDataset,
+    StripTokenDataset,
+    TruncateDataset,
+    data_utils,
+    encoders,
+    indexed_dataset,
+)
+from fairseq.tasks.translation import TranslationTask
+from fairseq.tasks import register_task, LegacyFairseqTask
+
+EVAL_BLEU_ORDER = 4
+
+logger = logging.getLogger(__name__)
+
+
+def load_langpair_dataset(
+        data_path,
+        split,
+        src,
+        src_dict,
+        tgt,
+        tgt_dict,
+        combine,
+        dataset_impl,
+        upsample_primary,
+        left_pad_source,
+        left_pad_target,
+        max_source_positions,
+        max_target_positions,
+        prepend_bos=False,
+        load_alignments=False,
+        truncate_source=False,
+        append_source_id=False,
+        num_buckets=0,
+        shuffle=True,
+        pad_to_multiple=1,
+):
+    def split_exists(split, src, tgt, lang, data_path):
+        filename = os.path.join(data_path, "{}.{}-{}.{}".format(split, src, tgt, lang))
+        return os.path.exists(filename)
+    
+    src_datasets = []
+    tgt_datasets = []
+    
+    for k in itertools.count():
+        split_k = split + (str(k) if k > 0 else "")
+        
+        # infer langcode
+        if split_exists(split_k, src, tgt, src, data_path):
+            prefix = os.path.join(data_path, "{}.{}-{}.".format(split_k, src, tgt))
+        elif split_exists(split_k, tgt, src, src, data_path):
+            prefix = os.path.join(data_path, "{}.{}-{}.".format(split_k, tgt, src))
+        else:
+            if k > 0:
+                break
+            else:
+                raise FileNotFoundError(
+                    "Dataset not found: {} ({})".format(split, data_path)
+                )
+        
+        src_dataset = data_utils.load_indexed_dataset(
+            prefix + src, src_dict, dataset_impl
+        )
+        if truncate_source:
+            src_dataset = AppendTokenDataset(
+                TruncateDataset(
+                    StripTokenDataset(src_dataset, src_dict.eos()),
+                    max_source_positions - 1,
+                ),
+                src_dict.eos(),
+            )
+        src_datasets.append(src_dataset)
+        
+        tgt_dataset = data_utils.load_indexed_dataset(
+            prefix + tgt, tgt_dict, dataset_impl
+        )
+        if tgt_dataset is not None:
+            tgt_datasets.append(tgt_dataset)
+        
+        logger.info(
+            "{} {} {}-{} {} examples".format(
+                data_path, split_k, src, tgt, len(src_datasets[-1])
+            )
+        )
+        
+        if not combine:
+            break
+    
+    assert len(src_datasets) == len(tgt_datasets) or len(tgt_datasets) == 0
+    
+    if len(src_datasets) == 1:
+        src_dataset = src_datasets[0]
+        tgt_dataset = tgt_datasets[0] if len(tgt_datasets) > 0 else None
+    else:
+        sample_ratios = [1] * len(src_datasets)
+        sample_ratios[0] = upsample_primary
+        src_dataset = ConcatDataset(src_datasets, sample_ratios)
+        if len(tgt_datasets) > 0:
+            tgt_dataset = ConcatDataset(tgt_datasets, sample_ratios)
+        else:
+            tgt_dataset = None
+    
+    if prepend_bos:
+        assert hasattr(src_dict, "bos_index") and hasattr(tgt_dict, "bos_index")
+        src_dataset = PrependTokenDataset(src_dataset, src_dict.bos())
+        if tgt_dataset is not None:
+            tgt_dataset = PrependTokenDataset(tgt_dataset, tgt_dict.bos())
+    
+    eos = None
+    if append_source_id:
+        src_dataset = AppendTokenDataset(
+            src_dataset, src_dict.index("[{}]".format(src))
+        )
+        if tgt_dataset is not None:
+            tgt_dataset = AppendTokenDataset(
+                tgt_dataset, tgt_dict.index("[{}]".format(tgt))
+            )
+        eos = tgt_dict.index("[{}]".format(tgt))
+    
+    align_dataset = None
+    if load_alignments:
+        align_path = os.path.join(data_path, "{}.align.{}-{}".format(split, src, tgt))
+        if indexed_dataset.dataset_exists(align_path, impl=dataset_impl):
+            align_dataset = data_utils.load_indexed_dataset(
+                align_path, None, dataset_impl
+            )
+    
+    tgt_dataset_sizes = tgt_dataset.sizes if tgt_dataset is not None else None
+    return LanguagePairDataset(
+        src_dataset,
+        src_dataset.sizes,
+        src_dict,
+        tgt_dataset,
+        tgt_dataset_sizes,
+        tgt_dict,
+        left_pad_source=left_pad_source,
+        left_pad_target=left_pad_target,
+        align_dataset=align_dataset,
+        eos=eos,
+        num_buckets=num_buckets,
+        shuffle=shuffle,
+        pad_to_multiple=pad_to_multiple,
+    )
+
+
+@register_task("translation_w_langtok")
+class TranslationWithLangtokTask(LegacyFairseqTask):
+    """
+    Translate from one (source) language to another (target) language.
+
+    Args:
+        src_dict (~fairseq.data.Dictionary): dictionary for the source language
+        tgt_dict (~fairseq.data.Dictionary): dictionary for the target language
+
+    .. note::
+
+        The translation task is compatible with :mod:`fairseq-train`,
+        :mod:`fairseq-generate` and :mod:`fairseq-interactive`.
+
+    The translation task provides the following additional command-line
+    arguments:
+
+    .. argparse::
+        :ref: fairseq.tasks.translation_parser
+        :prog:
+    """
+    
+    @staticmethod
+    def add_args(parser):
+        """Add task-specific arguments to the parser."""
+        # fmt: off
+        parser.add_argument('data', help='colon separated path to data directories list, \
+                            will be iterated upon during epochs in round-robin manner; \
+                            however, valid and test data are always in the first directory to \
+                            avoid the need for repeating them in all directories')
+        parser.add_argument('-s', '--source-lang', default=None, metavar='SRC',
+                            help='source language')
+        parser.add_argument('-t', '--target-lang', default=None, metavar='TARGET',
+                            help='target language')
+        parser.add_argument('--load-alignments', action='store_true',
+                            help='load the binarized alignments')
+        parser.add_argument('--left-pad-source', default='True', type=str, metavar='BOOL',
+                            help='pad the source on the left')
+        parser.add_argument('--left-pad-target', default='False', type=str, metavar='BOOL',
+                            help='pad the target on the left')
+        parser.add_argument('--max-source-positions', default=1024, type=int, metavar='N',
+                            help='max number of tokens in the source sequence')
+        parser.add_argument('--max-target-positions', default=1024, type=int, metavar='N',
+                            help='max number of tokens in the target sequence')
+        parser.add_argument('--upsample-primary', default=1, type=int,
+                            help='amount to upsample primary dataset')
+        parser.add_argument('--truncate-source', action='store_true', default=False,
+                            help='truncate source to max-source-positions')
+        parser.add_argument('--num-batch-buckets', default=0, type=int, metavar='N',
+                            help='if >0, then bucket source and target lengths into N '
+                                 'buckets and pad accordingly; this is useful on TPUs '
+                                 'to minimize the number of compilations')
+        parser.add_argument('--lang-prefix-tok', default=None, type=str, help="starting token in decoder")
+        
+        # options for reporting BLEU during validation
+        parser.add_argument('--eval-bleu', action='store_true',
+                            help='evaluation with BLEU scores')
+        parser.add_argument('--eval-bleu-detok', type=str, default="space",
+                            help='detokenize before computing BLEU (e.g., "moses"); '
+                                 'required if using --eval-bleu; use "space" to '
+                                 'disable detokenization; see fairseq.data.encoders '
+                                 'for other options')
+        parser.add_argument('--eval-bleu-detok-args', type=str, metavar='JSON',
+                            help='args for building the tokenizer, if needed')
+        parser.add_argument('--eval-tokenized-bleu', action='store_true', default=False,
+                            help='compute tokenized BLEU instead of sacrebleu')
+        parser.add_argument('--eval-bleu-remove-bpe', nargs='?', const='@@ ', default=None,
+                            help='remove BPE before computing BLEU')
+        parser.add_argument('--eval-bleu-args', type=str, metavar='JSON',
+                            help='generation args for BLUE scoring, '
+                                 'e.g., \'{"beam": 4, "lenpen": 0.6}\'')
+        parser.add_argument('--eval-bleu-print-samples', action='store_true',
+                            help='print sample generations during validation')
+        # fmt: on
+    
+    def __init__(self, args, src_dict, tgt_dict):
+        super().__init__(args)
+        self.src_dict = src_dict
+        self.tgt_dict = tgt_dict
+    
+    @classmethod
+    def setup_task(cls, args, **kwargs):
+        """Setup the task (e.g., load dictionaries).
+
+        Args:
+            args (argparse.Namespace): parsed command-line arguments
+        """
+        args.left_pad_source = utils.eval_bool(args.left_pad_source)
+        args.left_pad_target = utils.eval_bool(args.left_pad_target)
+        
+        paths = utils.split_paths(args.data)
+        assert len(paths) > 0
+        # find language pair automatically
+        if args.source_lang is None or args.target_lang is None:
+            args.source_lang, args.target_lang = data_utils.infer_language_pair(
+                paths[0]
+            )
+        if args.source_lang is None or args.target_lang is None:
+            raise Exception(
+                "Could not infer language pair, please provide it explicitly"
+            )
+        
+        # load dictionaries
+        src_dict = cls.load_dictionary(
+            os.path.join(paths[0], "dict.{}.txt".format(args.source_lang))
+        )
+        tgt_dict = cls.load_dictionary(
+            os.path.join(paths[0], "dict.{}.txt".format(args.target_lang))
+        )
+        assert src_dict.pad() == tgt_dict.pad()
+        assert src_dict.eos() == tgt_dict.eos()
+        assert src_dict.unk() == tgt_dict.unk()
+        logger.info("[{}] dictionary: {} types".format(args.source_lang, len(src_dict)))
+        logger.info("[{}] dictionary: {} types".format(args.target_lang, len(tgt_dict)))
+        
+        return cls(args, src_dict, tgt_dict)
+    
+    def load_dataset(self, split, epoch=1, combine=False, **kwargs):
+        """Load a given dataset split.
+
+        Args:
+            split (str): name of the split (e.g., train, valid, test)
+        """
+        paths = utils.split_paths(self.args.data)
+        assert len(paths) > 0
+        if split != getattr(self.args, "train_subset", None):
+            # if not training data set, use the first shard for valid and test
+            paths = paths[:1]
+        data_path = paths[(epoch - 1) % len(paths)]
+        
+        # infer langcode
+        src, tgt = self.args.source_lang, self.args.target_lang
+        
+        self.datasets[split] = load_langpair_dataset(
+            data_path,
+            split,
+            src,
+            self.src_dict,
+            tgt,
+            self.tgt_dict,
+            combine=combine,
+            dataset_impl=self.args.dataset_impl,
+            upsample_primary=self.args.upsample_primary,
+            left_pad_source=self.args.left_pad_source,
+            left_pad_target=self.args.left_pad_target,
+            max_source_positions=self.args.max_source_positions,
+            max_target_positions=self.args.max_target_positions,
+            load_alignments=self.args.load_alignments,
+            truncate_source=self.args.truncate_source,
+            num_buckets=self.args.num_batch_buckets,
+            shuffle=(split != "test"),
+            pad_to_multiple=self.args.required_seq_len_multiple,
+        )
+    
+    def build_dataset_for_inference(self, src_tokens, src_lengths, constraints=None):
+        return LanguagePairDataset(
+            src_tokens,
+            src_lengths,
+            self.source_dictionary,
+            tgt_dict=self.target_dictionary,
+            constraints=constraints,
+        )
+    
+    def build_model(self, args):
+        model = super().build_model(args)
+        if getattr(args, "eval_bleu", False):
+            assert getattr(args, "eval_bleu_detok", None) is not None, (
+                "--eval-bleu-detok is required if using --eval-bleu; "
+                "try --eval-bleu-detok=moses (or --eval-bleu-detok=space "
+                "to disable detokenization, e.g., when using sentencepiece)"
+            )
+            detok_args = json.loads(getattr(args, "eval_bleu_detok_args", "{}") or "{}")
+            self.tokenizer = encoders.build_tokenizer(
+                Namespace(
+                    tokenizer=getattr(args, "eval_bleu_detok", None), **detok_args
+                )
+            )
+            
+            gen_args = json.loads(getattr(args, "eval_bleu_args", "{}") or "{}")
+            self.sequence_generator = self.build_generator(
+                [model], Namespace(**gen_args)
+            )
+        return model
+    
+    def valid_step(self, sample, model, criterion):
+        loss, sample_size, logging_output = super().valid_step(sample, model, criterion)
+        if self.args.eval_bleu:
+            bleu = self._inference_with_bleu(self.sequence_generator, sample, model)
+            logging_output["_bleu_sys_len"] = bleu.sys_len
+            logging_output["_bleu_ref_len"] = bleu.ref_len
+            # we split counts into separate entries so that they can be
+            # summed efficiently across workers using fast-stat-sync
+            assert len(bleu.counts) == EVAL_BLEU_ORDER
+            for i in range(EVAL_BLEU_ORDER):
+                logging_output["_bleu_counts_" + str(i)] = bleu.counts[i]
+                logging_output["_bleu_totals_" + str(i)] = bleu.totals[i]
+        return loss, sample_size, logging_output
+    
+    def inference_step(
+            self, generator, models, sample, prefix_tokens=None, constraints=None
+    ):
+        if self.args.lang_prefix_tok is None:
+            prefix_tokens = None
+        else:
+            prefix_tokens = self.target_dictionary.index(self.args.lang_prefix_tok)
+            assert prefix_tokens != self.target_dictionary.unk_index
+        with torch.no_grad():
+            net_input = sample["net_input"]
+            if "src_tokens" in net_input:
+                src_tokens = net_input["src_tokens"]
+            elif "source" in net_input:
+                src_tokens = net_input["source"]
+            else:
+                raise Exception("expected src_tokens or source in net input")
+            
+            # bsz: total number of sentences in beam
+            # Note that src_tokens may have more than 2 dimenions (i.e. audio features)
+            bsz, _ = src_tokens.size()[:2]
+            if prefix_tokens is not None:
+                if isinstance(prefix_tokens, int):
+                    prefix_tokens = torch.LongTensor([prefix_tokens]).unsqueeze(1)  # 1,1
+                    prefix_tokens = prefix_tokens.expand(bsz, -1)
+                    prefix_tokens = prefix_tokens.to(src_tokens.device)
+            return generator.generate(models, sample, prefix_tokens=prefix_tokens)
+    
+    def reduce_metrics(self, logging_outputs, criterion):
+        super().reduce_metrics(logging_outputs, criterion)
+        if self.args.eval_bleu:
+            
+            def sum_logs(key):
+                return sum(log.get(key, 0) for log in logging_outputs)
+            
+            counts, totals = [], []
+            for i in range(EVAL_BLEU_ORDER):
+                counts.append(sum_logs("_bleu_counts_" + str(i)))
+                totals.append(sum_logs("_bleu_totals_" + str(i)))
+            
+            if max(totals) > 0:
+                # log counts as numpy arrays -- log_scalar will sum them correctly
+                metrics.log_scalar("_bleu_counts", np.array(counts))
+                metrics.log_scalar("_bleu_totals", np.array(totals))
+                metrics.log_scalar("_bleu_sys_len", sum_logs("_bleu_sys_len"))
+                metrics.log_scalar("_bleu_ref_len", sum_logs("_bleu_ref_len"))
+                
+                def compute_bleu(meters):
+                    import inspect
+                    import sacrebleu
+                    
+                    fn_sig = inspect.getfullargspec(sacrebleu.compute_bleu)[0]
+                    if "smooth_method" in fn_sig:
+                        smooth = {"smooth_method": "exp"}
+                    else:
+                        smooth = {"smooth": "exp"}
+                    bleu = sacrebleu.compute_bleu(
+                        correct=meters["_bleu_counts"].sum,
+                        total=meters["_bleu_totals"].sum,
+                        sys_len=meters["_bleu_sys_len"].sum,
+                        ref_len=meters["_bleu_ref_len"].sum,
+                        **smooth
+                    )
+                    return round(bleu.score, 2)
+                
+                metrics.log_derived("bleu", compute_bleu)
+    
+    def max_positions(self):
+        """Return the max sentence length allowed by the task."""
+        return (self.args.max_source_positions, self.args.max_target_positions)
+    
+    @property
+    def source_dictionary(self):
+        """Return the source :class:`~fairseq.data.Dictionary`."""
+        return self.src_dict
+    
+    @property
+    def target_dictionary(self):
+        """Return the target :class:`~fairseq.data.Dictionary`."""
+        return self.tgt_dict
+    
+    def _inference_with_bleu(self, generator, sample, model):
+        import sacrebleu
+        
+        def decode(toks, escape_unk=False):
+            s = self.tgt_dict.string(
+                toks.int().cpu(),
+                self.args.eval_bleu_remove_bpe,
+                # The default unknown string in fairseq is `<unk>`, but
+                # this is tokenized by sacrebleu as `< unk >`, inflating
+                # BLEU scores. Instead, we use a somewhat more verbose
+                # alternative that is unlikely to appear in the real
+                # reference, but doesn't get split into multiple tokens.
+                unk_string=("UNKNOWNTOKENINREF" if escape_unk else "UNKNOWNTOKENINHYP"),
+            )
+            if self.tokenizer:
+                s = self.tokenizer.decode(s)
+            return s
+        
+        gen_out = self.inference_step(generator, [model], sample, prefix_tokens=None)
+        hyps, refs = [], []
+        for i in range(len(gen_out)):
+            hyps.append(decode(gen_out[i][0]["tokens"]))
+            refs.append(
+                decode(
+                    utils.strip_pad(sample["target"][i], self.tgt_dict.pad()),
+                    escape_unk=True,  # don't count <unk> as matches to the hypo
+                )
+            )
+        if self.args.eval_bleu_print_samples:
+            logger.info("example hypothesis: " + hyps[0])
+            logger.info("example reference: " + refs[0])
+        if self.args.eval_tokenized_bleu:
+            return sacrebleu.corpus_bleu(hyps, [refs], tokenize="none")
+        else:
+            return sacrebleu.corpus_bleu(hyps, [refs])

mcolt/tasks/translation_w_mono.py

@@ -0,0 +1,214 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+from fairseq import options, utils
+from fairseq.data import (
+    ConcatDataset,
+    data_utils,
+    LanguagePairDataset)
+
+from ..data import SubsampleLanguagePairDataset
+
+import logging
+from fairseq.tasks import register_task
+from fairseq.tasks.translation import TranslationTask, load_langpair_dataset
+
+logger = logging.getLogger(__name__)
+
+
+def concat_language_pair_dataset(*language_pair_datasets, up_sample_ratio=None,
+                                 all_dataset_upsample_ratio=None):
+    logger.info("To cancat the language pairs")
+    dataset_number = len(language_pair_datasets)
+    if dataset_number == 1:
+        return language_pair_datasets[0]
+    elif dataset_number < 1:
+        raise ValueError("concat_language_pair_dataset needs at least on dataset")
+    # for dataset in language_pair_datasets:
+    #     assert isinstance(dataset, LanguagePairDataset), "concat_language_pair_dataset can only concat language pair" \
+    #                                                      "dataset"
+    
+    src_list = [language_pair_datasets[0].src]
+    tgt_list = [language_pair_datasets[0].tgt]
+    src_dict = language_pair_datasets[0].src_dict
+    tgt_dict = language_pair_datasets[0].tgt_dict
+    left_pad_source = language_pair_datasets[0].left_pad_source
+    left_pad_target = language_pair_datasets[0].left_pad_target
+    
+    logger.info("To construct the source dataset list and the target dataset list")
+    for dataset in language_pair_datasets[1:]:
+        assert dataset.src_dict == src_dict
+        assert dataset.tgt_dict == tgt_dict
+        assert dataset.left_pad_source == left_pad_source
+        assert dataset.left_pad_target == left_pad_target
+        src_list.append(dataset.src)
+        tgt_list.append(dataset.tgt)
+    logger.info("Have constructed the source dataset list and the target dataset list")
+    
+    if all_dataset_upsample_ratio is None:
+        sample_ratio = [1] * len(src_list)
+        sample_ratio[0] = up_sample_ratio
+    else:
+        sample_ratio = [int(t) for t in all_dataset_upsample_ratio.strip().split(",")]
+        assert len(sample_ratio) == len(src_list)
+    src_dataset = ConcatDataset(src_list, sample_ratios=sample_ratio)
+    tgt_dataset = ConcatDataset(tgt_list, sample_ratios=sample_ratio)
+    res = LanguagePairDataset(
+        src_dataset, src_dataset.sizes, src_dict,
+        tgt_dataset, tgt_dataset.sizes, tgt_dict,
+        left_pad_source=left_pad_source,
+        left_pad_target=left_pad_target,
+    )
+    logger.info("Have created the concat language pair dataset")
+    return res
+
+
+@register_task('translation_w_mono')
+class TranslationWithMonoTask(TranslationTask):
+    """
+    Translate from one (source) language to another (target) language.
+
+    Args:
+        src_dict (~fairseq.data.Dictionary): dictionary for the source language
+        tgt_dict (~fairseq.data.Dictionary): dictionary for the target language
+
+    .. note::
+
+        The translation task is compatible with :mod:`fairseq-train`,
+        :mod:`fairseq-generate` and :mod:`fairseq-interactive`.
+
+    The translation task provides the following additional command-line
+    arguments:
+
+    .. argparse::
+        :ref: fairseq.tasks.translation_parser
+        :prog:
+    """
+    
+    @staticmethod
+    def add_args(parser):
+        """Add task-specific arguments to the parser."""
+        # fmt: off
+        TranslationTask.add_args(parser)
+        parser.add_argument('--mono-data', default=None, help='monolingual data, split by :')
+        parser.add_argument('--mono-one-split-each-epoch', action='store_true', default=False, help='use on split of monolingual data at each epoch')
+        parser.add_argument('--parallel-ratio', default=1.0, type=float, help='subsample ratio of parallel data')
+        parser.add_argument('--mono-ratio', default=1.0, type=float, help='subsample ratio of mono data')
+    
+    def __init__(self, args, src_dict, tgt_dict):
+        super().__init__(args, src_dict, tgt_dict)
+        self.src_dict = src_dict
+        self.tgt_dict = tgt_dict
+        self.update_number = 0
+    
+    @classmethod
+    def setup_task(cls, args, **kwargs):
+        """Setup the task (e.g., load dictionaries).
+
+        Args:
+            args (argparse.Namespace): parsed command-line arguments
+        """
+        args.left_pad_source = options.eval_bool(args.left_pad_source)
+        args.left_pad_target = options.eval_bool(args.left_pad_target)
+        if getattr(args, 'raw_text', False):
+            utils.deprecation_warning('--raw-text is deprecated, please use --dataset-impl=raw')
+            args.dataset_impl = 'raw'
+        elif getattr(args, 'lazy_load', False):
+            utils.deprecation_warning('--lazy-load is deprecated, please use --dataset-impl=lazy')
+            args.dataset_impl = 'lazy'
+        
+        paths = utils.split_paths(args.data)
+        assert len(paths) > 0
+        # find language pair automatically
+        if args.source_lang is None or args.target_lang is None:
+            args.source_lang, args.target_lang = data_utils.infer_language_pair(paths[0])
+        if args.source_lang is None or args.target_lang is None:
+            raise Exception('Could not infer language pair, please provide it explicitly')
+        
+        # load dictionaries
+        src_dict = cls.load_dictionary(os.path.join(paths[0], 'dict.{}.txt'.format(args.source_lang)))
+        tgt_dict = cls.load_dictionary(os.path.join(paths[0], 'dict.{}.txt'.format(args.target_lang)))
+        assert src_dict.pad() == tgt_dict.pad()
+        assert src_dict.eos() == tgt_dict.eos()
+        assert src_dict.unk() == tgt_dict.unk()
+        logger.info('| [{}] dictionary: {} types'.format(args.source_lang, len(src_dict)))
+        logger.info('| [{}] dictionary: {} types'.format(args.target_lang, len(tgt_dict)))
+        
+        return cls(args, src_dict, tgt_dict)
+    
+    def load_dataset(self, split, epoch=0, combine=False, **kwargs):
+        """Load a given dataset split.
+
+        Args:
+            split (str): name of the split (e.g., train, valid, test)
+        """
+        logger.info("To load the dataset {}".format(split))
+        paths = utils.split_paths(self.args.data)
+        assert len(paths) > 0
+        if split != getattr(self.args, "train_subset", None):
+            # if not training data set, use the first shard for valid and test
+            paths = paths[:1]
+        data_path = paths[(epoch - 1) % len(paths)]
+        
+        mono_paths = utils.split_paths(self.args.mono_data)
+        
+        # infer langcode
+        src, tgt = self.args.source_lang, self.args.target_lang
+        
+        parallel_data = load_langpair_dataset(
+            data_path, split, src, self.src_dict, tgt, self.tgt_dict,
+            combine=combine, dataset_impl=self.args.dataset_impl,
+            upsample_primary=self.args.upsample_primary,
+            left_pad_source=self.args.left_pad_source,
+            left_pad_target=self.args.left_pad_target,
+            max_source_positions=self.args.max_source_positions,
+            max_target_positions=self.args.max_target_positions,
+            load_alignments=self.args.load_alignments,
+            num_buckets=self.args.num_batch_buckets,
+            shuffle=(split != "test"),
+            pad_to_multiple=self.args.required_seq_len_multiple,
+        )
+        if split == "train":
+            parallel_data = SubsampleLanguagePairDataset(parallel_data, size_ratio=self.args.parallel_ratio,
+                                                         seed=self.args.seed,
+                                                         epoch=epoch)
+            if self.args.mono_one_split_each_epoch:
+                mono_path = mono_paths[(epoch - 1) % len(mono_paths)]  # each at one epoch
+                mono_data = load_langpair_dataset(
+                    mono_path, split, src, self.src_dict, tgt, self.tgt_dict,
+                    combine=combine, dataset_impl=self.args.dataset_impl,
+                    upsample_primary=self.args.upsample_primary,
+                    left_pad_source=self.args.left_pad_source,
+                    left_pad_target=self.args.left_pad_target,
+                    max_source_positions=self.args.max_source_positions,
+                    shuffle=(split != "test"),
+                    max_target_positions=self.args.max_target_positions,
+                )
+                mono_data = SubsampleLanguagePairDataset(mono_data, size_ratio=self.args.mono_ratio,
+                                                         seed=self.args.seed,
+                                                         epoch=epoch)
+                all_dataset = [parallel_data, mono_data]
+            else:
+                mono_datas = []
+                for mono_path in mono_paths:
+                    mono_data = load_langpair_dataset(
+                        mono_path, split, src, self.src_dict, tgt, self.tgt_dict,
+                        combine=combine, dataset_impl=self.args.dataset_impl,
+                        upsample_primary=self.args.upsample_primary,
+                        left_pad_source=self.args.left_pad_source,
+                        left_pad_target=self.args.left_pad_target,
+                        max_source_positions=self.args.max_source_positions,
+                        shuffle=(split != "test"),
+                        max_target_positions=self.args.max_target_positions,
+                    )
+                    mono_data = SubsampleLanguagePairDataset(mono_data, size_ratio=self.args.mono_ratio,
+                                                             seed=self.args.seed,
+                                                             epoch=epoch)
+                    mono_datas.append(mono_data)
+                all_dataset = [parallel_data] + mono_datas
+            self.datasets[split] = ConcatDataset(all_dataset)
+        else:
+            self.datasets[split] = parallel_data

requirements.txt

@@ -0,0 +1,5 @@
+subword-nmt
+sacrebleu
+sacremoses
+kytea
+six

scripts/load_config.sh

@@ -0,0 +1,48 @@
+#!/usr/bin/env bash
+
+
+function parse_yaml {
+   local prefix=$2
+   local s='[[:space:]]*' w='[a-zA-Z0-9_]*' fs=$(echo @|tr @ '\034')
+   sed -ne "s|^\($s\):|\1|" \
+        -e "s|^\($s\)\($w\)$s:$s[\"']\(.*\)[\"']$s\$|\1$fs\2$fs\3|p" \
+        -e "s|^\($s\)\($w\)$s:$s\(.*\)$s\$|\1$fs\2$fs\3|p" $1 |
+   awk -F$fs '{
+      indent = length($1)/2;
+      vname[indent] = $2;
+      for (i in vname) {if (i > indent) {delete vname[i]}}
+      if (length($3) > 0) {
+         vn=""; for (i=0; i<indent; i++) {vn=(vn)(vname[i])("_")}
+         printf("%s%s%s=\"%s\"\n", "'$prefix'",vn, $2, $3);
+      }
+   }'
+}
+main_config_yml=$1
+local_root=$2
+if [[ ${main_config_yml} == "hdfs://"* ]]; then
+    config_filename=`basename ${main_config_yml}`
+    echo 'download config from ${main_config_yml}...'
+    local_config="${local_root}/config" && mkdir -p ${local_config}
+    hadoop fs -get ${main_config_yml} ${local_config}/
+    echo 'finish download config from ${main_config_yml}...'
+    main_config_yml=${local_config}/${config_filename}
+fi
+
+compgen -A variable > ~/.env-vars
+eval $(parse_yaml ${main_config_yml})
+
+# set option flags
+options=""
+for var in `compgen -A variable | grep -Fxvf  ~/.env-vars`
+do
+    if [[ ${var} == "model_"* || ${var} == "data_"* || ${var} == "options" ]]; then
+        continue
+    fi
+    if [[ ${!var} == "true" ]]; then
+        varname=`echo ${var} | sed 's/\_/\-/g'`
+        options=${options}" --${varname}"
+    else
+        varname=`echo ${var} | sed 's/\_/\-/g'`
+        options=${options}" --${varname} ${!var}"
+    fi
+done

scripts/utils.py

@@ -0,0 +1,116 @@
+import re
+import os
+import sys
+from tqdm import tqdm
+
+
+def remove_bpe(line, bpe_symbol="@@ "):
+    line = line.replace("\n", '')
+    line = (line + ' ').replace(bpe_symbol, '').rstrip()
+    return line + ("\n")
+
+
+def remove_bpe_fn(i=sys.stdin, o=sys.stdout, bpe="@@ "):
+    lines = tqdm(i)
+    lines = map(lambda x: remove_bpe(x, bpe), lines)
+    # _write_lines(lines, f=o)
+    for line in lines:
+        o.write(line)
+
+
+def reprocess(fle):
+    # takes in a file of generate.py translation generate_output
+    # returns a source dict and hypothesis dict, where keys are the ID num (as a string)
+    # and values and the corresponding source and translation. There may be several translations
+    # per source, so the values for hypothesis_dict are lists.
+    # parses output of generate.py
+    
+    with open(fle, 'r') as f:
+        txt = f.read()
+    
+    """reprocess generate.py output"""
+    p = re.compile(r"[STHP][-]\d+\s*")
+    hp = re.compile(r"(\s*[-]?\d+[.]?\d+(e[+-]?\d+)?\s*)|(\s*(-inf)\s*)")
+    source_dict = {}
+    hypothesis_dict = {}
+    score_dict = {}
+    target_dict = {}
+    pos_score_dict = {}
+    lines = txt.split("\n")
+    
+    for line in lines:
+        line += "\n"
+        prefix = re.search(p, line)
+        if prefix is not None:
+            assert len(prefix.group()) > 2, "prefix id not found"
+            _, j = prefix.span()
+            id_num = prefix.group()[2:]
+            id_num = int(id_num)
+            line_type = prefix.group()[0]
+            if line_type == "H":
+                h_txt = line[j:]
+                hypo = re.search(hp, h_txt)
+                assert hypo is not None, ("regular expression failed to find the hypothesis scoring")
+                _, i = hypo.span()
+                score = hypo.group()
+                hypo_str = h_txt[i:]
+                # if r2l:  # todo: reverse score as well
+                #     hypo_str = " ".join(reversed(hypo_str.strip().split(" "))) + "\n"
+                if id_num in hypothesis_dict:
+                    hypothesis_dict[id_num].append(hypo_str)
+                    score_dict[id_num].append(float(score))
+                else:
+                    hypothesis_dict[id_num] = [hypo_str]
+                    score_dict[id_num] = [float(score)]
+            
+            elif line_type == "S":
+                source_dict[id_num] = (line[j:])
+            elif line_type == "T":
+                # target_dict[id_num] = (line[j:])
+                continue
+            elif line_type == "P":
+                pos_scores = (line[j:]).split()
+                pos_scores = [float(x) for x in pos_scores]
+                if id_num in pos_score_dict:
+                    pos_score_dict[id_num].append(pos_scores)
+                else:
+                    pos_score_dict[id_num] = [pos_scores]
+    
+    return source_dict, hypothesis_dict, score_dict, target_dict, pos_score_dict
+
+
+def get_hypo_and_ref(fle, hyp_file, ref_input, ref_file, rank=0):
+    with open(ref_input, 'r') as f:
+        refs = f.readlines()
+    _, hypo_dict, _, _, _ = reprocess(fle)
+    assert rank < len(hypo_dict[0])
+    maxkey = max(hypo_dict, key=int)
+    f_hyp = open(hyp_file, "w")
+    f_ref = open(ref_file, "w")
+    for idx in range(maxkey + 1):
+        if idx not in hypo_dict:
+            continue
+        f_hyp.write(hypo_dict[idx][rank])
+        f_ref.write(refs[idx])
+    f_hyp.close()
+    f_ref.close()
+
+
+def recover_bpe(hyp_file):
+    f_hyp = open(hyp_file, "r")
+    f_hyp_out = open(hyp_file + ".nobpe", "w")
+    for _s in ["hyp"]:
+        f = eval("f_{}".format(_s))
+        fout = eval("f_{}_out".format(_s))
+        remove_bpe_fn(i=f, o=fout)
+    f_hyp.close()
+    f_hyp_out.close()
+
+
+if __name__ == "__main__":
+    filename = sys.argv[1]
+    ref_in = sys.argv[2]
+    hypo_file = os.path.join(os.path.dirname(filename), "hypo.out")
+    ref_out = os.path.join(os.path.dirname(filename), "ref.out")
+    get_hypo_and_ref(filename, hypo_file, ref_in, ref_out)
+    recover_bpe(hypo_file)

test/input.en

@@ -0,0 +1 @@
+LANG_TOK_EN Hello my friend!

test/input.zh

@@ -0,0 +1 @@
+LANG_TOK_ZH 你好！

test/output.en.no_bpe

@@ -0,0 +1 @@
+D-0	-0.34370458126068115	LANG_TOK_EN ANG_TOK_ZH Hello !

test/output.en.no_bpe.moses

@@ -0,0 +1 @@
+D-0	-1.3185505867004395	LANG_TOK_EN Hello!

test/output.zh

@@ -0,0 +1,3 @@
+S-0	L@@ AN@@ G@@ _@@ T@@ OK@@ _@@ EN H@@ ello my fri@@ end@@ !
+H-0	-0.6148621439933777	LANG_TOK_ZH 你@@ 好 ， 我 的 朋@@ 友 ！
+P-0	-2.1448 -1.4575 -0.0638 -0.8495 -0.6207 -0.1953 -0.2082 -0.0769 -0.3801 -0.1517

test/output.zh.no_bpe

@@ -0,0 +1 @@
+D-0	-0.6148621439933777	LANG_TOK_ZH 你好 ， 我 的 朋友 ！

test/output.zh.no_bpe.moses

@@ -0,0 +1 @@
+D-0	-0.7665940523147583	LANG_TOK_ZH 你好 ， 我的朋友 ！

train_w_mono.sh

@@ -0,0 +1,56 @@
+#!/usr/bin/env bash
+
+# repo_dir: root directory of the project
+repo_dir="$( cd "$( dirname "$0" )" && pwd )"
+echo "==== Working directory: ====" >&2
+echo "${repo_dir}" >&2
+echo "============================" >&2
+
+main_config=$1
+source ${repo_dir}/scripts/load_config.sh ${main_config} ${repo_dir}
+
+model_dir=${repo_dir}/model
+data_dir=${repo_dir}/data
+
+mkdir -p ${model_dir} ${data_dir}/mono
+
+
+# parallel data
+data_var=data_1
+i=1
+data=""
+while [[ ! -z ${!data_var} ]]; do
+    if [[ $data == "" ]]; then
+        data=${!data_var}
+    else
+        data=$data:${!data_var}
+    fi
+    i=$((i+1))
+    data_var=data_$i
+done
+
+# mono data
+mono_data_var=data_mono_1
+y=1
+mono_data=""
+while [[ ! -z ${!mono_data_var} ]]; do
+    if [[ ${mono_data} == "" ]]; then
+        mono_data=${!mono_data_var}
+    else
+        mono_data=${mono_data}:${!mono_data_var}
+    fi
+    y=$((y+1))
+    mono_data_var=data_mono_$y
+done
+
+
+command="CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} fairseq-train ${data} \
+  --user-dir ${repo_dir}/mcolt \
+  --save-dir ${model_dir} \
+  --mono-data ${mono_data} \
+  ${options} \
+  --ddp-backend no_c10d 1>&2"
+
+echo $command
+eval $command
+