# Copyright (c) Microsoft Corporation. 
# Licensed under the MIT license.

# Set pretrained model name
MODEL_NAME=microsoft/amos

# GLUE task name, from ['MNLI', 'QQP', 'QNLI', 'SST-2', 'CoLA', 'RTE', 'MRPC', 'STS-B']
TASK=$1

# Path to GLUE dataset 'path/to/glue_data'
GLUE_PATH=$2

# Output path for results and fine-tuned model
OUT_PATH=$3

export DATASET_PATH=$GLUE_PATH/$TASK

export TASK_NAME=$(echo "$TASK" | tr '[:upper:]' '[:lower:]')

# Set max sequence length
export MAX_LEN=512

# Set path to cache train & dev features (tokenized, only use for this tokenizer!)
export TRAIN_CACHE=${DATASET_PATH}/$TASK_NAME.amos_cased.$MAX_LEN.cache
export DEV_CACHE=${DATASET_PATH}/$TASK_NAME.amos_cased.$MAX_LEN.cache

# Setting the hyperparameters for the run.
export BSZ=$4
export LR=$5
export EPOCH=$6
export WM=$7
export SEED=$8

# Set path to save the finetuned model and result score
export OUTPUT_PATH=$OUT_PATH/$TASK-$BSZ-$LR-$EPOCH-$WM-$SEED

mkdir -p $OUTPUT_PATH
touch $OUTPUT_PATH/train.log

python run_glue.py \
    --model_type amos --model_name_or_path $MODEL_NAME --task_name $TASK_NAME \
    --data_dir $DATASET_PATH --cached_train_file $TRAIN_CACHE --cached_dev_file $DEV_CACHE \
    --config_name $MODEL_NAME \
    --do_train --evaluate_during_training --logging_steps 1000 --output_dir $OUTPUT_PATH --max_grad_norm 0 --gradient_accumulation_steps 1 \
    --max_seq_length $MAX_LEN --per_gpu_train_batch_size $BSZ --learning_rate $LR \
    --num_train_epochs $EPOCH --weight_decay 0.01 --warmup_ratio $WM \
    --adam_epsilon 1e-6 --adam_betas "0.9,0.98" \
    --dropout_prob 0.1 --cls_dropout_prob 0.1 \
    --seed $SEED \
    --overwrite_output_dir |& tee $OUTPUT_PATH/train.log

# Add the following for fp16 training
# --fp16_init_loss_scale 128.0 --fp16 --fp16_opt_level O2 
