TASK='iwslt14'
MODE='test'
WORLD_SIZE=$2

SRC_LANG=$4   # en
TRG_LANG=$5   # de

SAVE_DIR='./results/'
DATA_DIR='./data/'

TEST_SRC_FILE='bpe.test.en'
TEST_TRG_FILE='bpe.test.de'
TEST_TOKEN_REF_SRC_FILE='test.en'
TEST_TOKEN_REF_TRG_FILE='test.de'
TEST_REF_SRC_FILE='ref.en'
TEST_REF_TRG_FILE='ref.de'
TEST_BATCH_SIZE=128 #128 #16 #256 #100 for Ti GPU, 60 for non-Ti GPU | 60 is standard | 1 is for measure time
TEST_TOKEN_SIZE=9000

TEST_MAX_LENGTH=150 #250

MODEL=$3

DIM_MODEL=512
DIM_WEMB=$DIM_MODEL # should be the same
DIM_FF=1024
N_LAYERS=6
TM_N_HEAD=4

echo "5 Beam width, IWSLT14 Testset"

# Testing configuration
BEAM_WIDTH=5
TEST_SUBDIR='iwslt14_en-de_150max_indnegatt_admin_nmt_radam_0.0005lr_0.1ls_0os_3600Tk_1sort_1US_50pat_0.3drop_6layer_4head_512dim_norm_res_1neghead_separam_1.0poslambda_1.0neglambda_full_0negkey_'
TEST_MODEL_FILE=$SRC_LANG'2'$TRG_LANG'.ensemble_model.best.pth'
TEST_MODEL_ARGS='model.args.pkl'

SACREBLEU_TOKENIZER='13a' # '13a' is default
SACREBLEU_LOWERCASE=0 # 0 is default

TRAINABLE_METRICS=0


CUDA_VISIBLE_DEVICES=$1 torchrun --rdzv_backend=c10d --rdzv_endpoint=localhost:0 --nnodes=1\
        --nproc_per_node=$WORLD_SIZE nmt_run.py \
        --world_size=$WORLD_SIZE \
        --translation_task=$TASK --mode=$MODE --src_lang=$SRC_LANG --trg_lang=$TRG_LANG\
        --save_dir=$SAVE_DIR --data_dir=$DATA_DIR --model=$MODEL \
        --dim_wemb=$DIM_WEMB --dim_model=$DIM_MODEL \
        --tm_n_layers=$N_LAYERS --tm_dim_ff=$DIM_FF \
        --test_batch_size=$TEST_BATCH_SIZE --test_token_size=$TEST_TOKEN_SIZE \
        --test_max_length=$TEST_MAX_LENGTH --tm_n_head=$TM_N_HEAD\
        --test_subdir=$TEST_SUBDIR --test_src_file=$TEST_SRC_FILE --test_trg_file=$TEST_TRG_FILE\
        --test_ref_src_file=$TEST_REF_SRC_FILE --test_ref_trg_file=$TEST_REF_TRG_FILE \
        --test_token_ref_src_file=$TEST_TOKEN_REF_SRC_FILE \
        --test_token_ref_trg_file=$TEST_TOKEN_REF_TRG_FILE \
        --beam_width=$BEAM_WIDTH --test_model_file=$TEST_MODEL_FILE\
        --test_model_args=$TEST_MODEL_ARGS \
        --sacrebleu_tokenizer=$SACREBLEU_TOKENIZER --sacrebleu_lowercase=$SACREBLEU_LOWERCASE \
        --trainable_metrics=$TRAINABLE_METRICS
