#!/bin/bash
#
# Read arguments
#
POSITIONAL=()
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"; LM="$3"; shift 3;;
  --model_path)
    SRC_MODEL="$2"; shift 2;;
  --exp)
    EXP="$2"; shift 2;;
  --data)
    DATA_OVERWRITE="$2"; shift 2;;
  *)
  POSITIONAL+=("$1")
  shift
  ;;
esac
done
set -- "${POSITIONAL[@]}"
if [ -z "$DATA_OVERWRITE" ]; then
  DATA_PATH=data/${TYPE}/tune/${SRC}-${TGT}/processed
else
  DATA_PATH=$DATA_OVERWRITE
fi
PARAMS=${SRC}_${TGT}_e${EMB_DIM}_l${N_LAYERS}_h${N_HEADS}_b${BATCH_SIZE}
MODELS=models/${TYPE}/${LM}/
if [ -z "$SRC_MODEL" ]; then
  MODELS_SRC=$MODELS/mlm/${SRC}-${TGT}/mlm_${PARAMS}/${EXP}/checkpoint.pth
else
  MODELS_SRC=$SRC_MODEL
fi
MODELS_TGT=$MODELS/tuned/${SRC}-${TGT}
echo $MODELS_SRC
echo $MODELS_TGT
echo $DATA_PATH
mkdir -p $MODELS_TGT
python3 XLM/train.py \
    --exp_name tuned_${PARAMS} \
    --exp_id ${EXP} \
    --dump_path $MODELS_TGT \
    --reload_model ${MODELS_SRC},${MODELS_SRC} \
    --data_path $DATA_PATH \
    --lgs "${SRC}-${TGT}" \
    --mt_steps "$SRC-$TGT" \
    --encoder_only false \
    --emb_dim $EMB_DIM \
    --n_layers $N_LAYERS \
    --n_heads $N_HEADS \
    --dropout 0.1 \
    --attention_dropout 0.1 \
    --epoch_size $(wc -l < $DATA_PATH/train.$SRC) \
    --max_epoch 1000 \
    --gelu_activation true \
    --batch_size $BATCH_SIZE \
    --optimizer adam,lr=0.0001,weight_decay=0.01 \
    --eval_bleu true \
    --eval_chrf true \
    --stopping_criterion "valid_${SRC}-${TGT}_mt_chrf,10"

