#!/bin/bash

#
# Read arguments
#
POSITIONAL=()
LM="MLM"
RELOAD_MODEL=""
while [[ $# -gt 0 ]]
do
key="$1"
case $key in
  --src)
    SRC="$2"; shift 2;;
  --tgt)
    TGT="$2"; shift 2;;
  --emb_dim)
    EMB_DIM="$2"; shift 2;;
  --n_layers)
    N_LAYERS="$2"; shift 2;;
  --n_heads)
    N_HEADS="$2"; shift 2;;
  --batch_size)
    BATCH_SIZE="$2"; shift 2;;
  --type)
    TYPE="$2"; shift 2;;
  --data)
    DATA_OVERWRITE="$2"; shift 2;;
  --exp)
    EXP="$2"; shift 2;;
  --reload_model)
    RELOAD_MODEL="--reload_model $2"; shift 2;;
  --tlm)
    LM="TLM"; shift 1;;
  *)
  POSITIONAL+=("$1")
  shift
  ;;
esac
done
set -- "${POSITIONAL[@]}"

if [ -z "$DATA_OVERWRITE" ]; then
	DATA_PATH=data/${TYPE}/mlm/${SRC}-${TGT}/processed
else
	DATA_PATH=$DATA_OVERWRITE
fi
MLM_STEPS="${SRC},${TGT}"
if [ "$LM" = "TLM" ]; then
  MLM_STEPS += ",${SRC}-${TGT}"
fi

PARAMS=${SRC}_${TGT}_e${EMB_DIM}_l${N_LAYERS}_h${N_HEADS}_b${BATCH_SIZE}
MODELS=models/$TYPE/mlm
MODELS_TGT=$MODELS/mlm/${SRC}-${TGT}
mkdir -p $MODELS_TGT
python3 XLM/train.py  \
    --exp_name mlm_${PARAMS} \
    --exp_id $EXP \
    --dump_path $MODELS_TGT $RELOAD_MODEL \
    --data_path $DATA_PATH \
    --lgs ${SRC}-${TGT} \
    --mlm_steps $MLM_STEPS \
    --emb_dim $EMB_DIM \
    --n_layers $N_LAYERS \
    --n_heads $N_HEADS \
    --dropout 0.1 \
    --attention_dropout 0.1 \
    --gelu_activation true \
    --batch_size $BATCH_SIZE \
    --optimizer adam,lr=0.0001,weight_decay=0.01 \
    --epoch_size $(wc -l < $DATA_PATH/train.$SRC) \
    --max_epoch 1000 \
    --validation_metrics _valid_mlm_ppl \
    --stopping_criterion _valid_mlm_ppl,10

