#!/bin/bash
PROJECT_DIR=${HOME}"/models/t5-political-evaluator"
export PYTHONPATH=${PROJECT_DIR}

FINETUNE_STEPS=5000
EVAL_PREFIX="ul2test/eval_political_"
MODEL_BUCKET_DIR="gs://nb-t5x-us-central2/finetuned/"

CHECKPOINT_LIST=(
  "t5-data/pretrained_models/t5x/mt5_small/checkpoint_1000000"
  "north-t5x/pretrained_models/base/norwegian_NCC_plus_English_t5x_base/checkpoint_1500000"
  "north-t5x/pretrained_models/base/norwegian_NCC_plus_English_pluss100k_lm_t5x_base/checkpoint_1600000"
  "north-t5x/pretrained_models/base/scandinavian3k_t5x_base/checkpoint_3000000"
  "t5-data/pretrained_models/t5x/mt5_base/checkpoint_1000000"
  "nb-t5/t5/t5_1_1_base_from_t5_scandi_unigram/checkpoint_2000000"
  "nb-t5/t5/t5_1_1_base_from_t5_scandi_unigram/checkpoint_3000000"
  "nb-t5/t5/t5_1_1_base_from_t5_scandi_unigram/checkpoint_4000000"
  "t5-data/pretrained_models/t5x/mt5_large/checkpoint_1000000"
  "north-t5x/pretrained_models/large/norwegian_NCC_plus_English_t5x_large/checkpoint_1500000"
  "north-t5x/pretrained_models/large/norwegian_NCC_plus_English_pluss100k_lm_t5x_large/checkpoint_1600000"
  "north-t5x/pretrained_models/large/scandinavian3k_t5x_large/checkpoint_3000000"
  "nb-t5/t5/t5_1_1_large_from_t5_scandi_unigram/checkpoint_2000000"
  "nb-t5/t5/t5_1_1_large_from_t5_scandi_unigram/checkpoint_3000000"
  "nb-t5/t5/t5_1_1_large_from_t5_scandi_unigram/checkpoint_4000000"
)

NAME_LIST=(
  "mt5_small"
  "north_t5_base_NCC"
  "north_t5_base_NCC_lm"
  "north_t5_base_scand3M"
  "mt5_base"
  "sab_base_2"
  "sab_base_3"
  "sab_base_4"
  "mt5_large"
  "north_t5_large_NCC"
  "north_t5_large_NCC_lm"
  "north_t5_large_scand3M"
  "sab_large_2"
  "sab_large_3"
  "sab_large_4"
)

TASK_LIST=("parliament_max300_mt5" "parliament_max300_mt5" "parliament_max300_mt5" "parliament_max300_mt5" "parliament_max300_mt5" "parliament_max300_scand" "parliament_max300_scand" "parliament_max300_scand" "parliament_max300_mt5" "parliament_max300_mt5" "parliament_max300_mt5" "parliament_max300_mt5" "parliament_max300_scand" "parliament_max300_scand" "parliament_max300_scand")

GIN_LIST=(
  "finetune_classification_small_mt5.gin"
  "finetune_classification_base_mt5.gin"
  "finetune_classification_base_mt5.gin"
  "finetune_classification_base_mt5.gin"
  "finetune_classification_base_mt5.gin"
  "finetune_classification_base_scand.gin"
  "finetune_classification_base_scand.gin"
  "finetune_classification_base_scand.gin"
  "finetune_classification_large_mt5.gin"
  "finetune_classification_large_mt5.gin"
  "finetune_classification_large_mt5.gin"
  "finetune_classification_large_mt5.gin"
  "finetune_classification_large_scand.gin"
  "finetune_classification_large_scand.gin"
  "finetune_classification_large_scand.gin"
)

START_LIST=(1000000 1500000 1600000 3000000 1000000 2000000 3000000 4000000 1000000 1500000 1600000 3000000 2000000 3000000 4000000)
EXP_LIST=(100 101 102 103 104 105 106 107 108 109 110 111 112 113 114)




VERSION_LIST=("1" "2" "3" "4" "5") 

index=$(($1 + 0))

if [ $index -lt 1 ] || [ $index -gt ${#CHECKPOINT_LIST[@]} ] || [ $# -ne 1 ]; then
	echo "Error: You need to provide the number of the checkpoints below as a parameter."
	for i in "${!CHECKPOINT_LIST[@]}"; do
	        echo "$((i+1)). ${CHECKPOINT_LIST[i]}"
	done	
	
	  exit 1
fi

for v in "${VERSION_LIST[@]}"; do
	i=($index-1)
    INITIAL_CHECKPOINT_PATH="gs://${CHECKPOINT_LIST[i]}"
    
    TRAIN_STEPS=$((START_LIST[i]+FINETUNE_STEPS))
    GIN_FILE=${GIN_LIST[i]}
    MIXTURE_OR_TASK_NAME=${TASK_LIST[i]}
    MODEL_DIR="${MODEL_BUCKET_DIR}${EVAL_PREFIX}exp${EXP_LIST[i]}_${NAME_LIST[i]}_v${v}"

    command="python3 ../../t5x/t5x/train.py --gin_search_paths=\"./\" --gin.TRAIN_STEPS=${TRAIN_STEPS} --gin_file=${GIN_FILE} --gin.INITIAL_CHECKPOINT_PATH=\\\"${INITIAL_CHECKPOINT_PATH}\\\" --gin.MIXTURE_OR_TASK_NAME=\\\"${MIXTURE_OR_TASK_NAME}\\\" --gin.MODEL_DIR=\\\"${MODEL_DIR}\\\""
    echo "${command}"
    # Uncomment the next line to run the command:
    eval "${command}"
done

#python3 ../../t5x/t5x/train.py --gin_search_paths="./" --gin.TRAIN_STEPS=${TRAIN_STEPS} --gin_file="finetune_translate_base.gin" --gin.INITIAL_CHECKPOINT_PATH=${INITIAL_CHECKPOINT_PATH}  --gin.MIXTURE_OR_TASK_NAME=\"translate_long\" --gin.MODEL_DIR=\"gs://nb-t5x-us-central2/finetuned/nynorsk_NCC_base_v1\" &&
#python3 ../../t5x/t5x/train.py --gin_search_paths="./" --gin.TRAIN_STEPS=${TRAIN_STEPS} --gin_file="finetune_translate_base.gin" --gin.INITIAL_CHECKPOINT_PATH=${INITIAL_CHECKPOINT_PATH}  --gin.MIXTURE_OR_TASK_NAME=\"translate_long\" --gin.MODEL_DIR=\"gs://nb-t5x-us-central2/finetuned/nynorsk_NCC_base_v2\" &&
#python3 ../../t5x/t5x/train.py --gin_search_paths="./" --gin.TRAIN_STEPS=${TRAIN_STEPS} --gin_file="finetune_translate_base.gin" --gin.INITIAL_CHECKPOINT_PATH=${INITIAL_CHECKPOINT_PATH}  --gin.MIXTURE_OR_TASK_NAME=\"translate_long\" --gin.MODEL_DIR=\"gs://nb-t5x-us-central2/finetuned/nynorsk_NCC_base_v3\" &&
#python3 ../../t5x/t5x/train.py --gin_search_paths="./" --gin.TRAIN_STEPS=${TRAIN_STEPS} --gin_file="finetune_translate_base.gin" --gin.INITIAL_CHECKPOINT_PATH=${INITIAL_CHECKPOINT_PATH}  --gin.MIXTURE_OR_TASK_NAME=\"translate_long\" --gin.MODEL_DIR=\"gs://nb-t5x-us-central2/finetuned/nynorsk_NCC_base_v4\" &&
#python3 ../../t5x/t5x/train.py --gin_search_paths="./" --gin.TRAIN_STEPS=${TRAIN_STEPS} --gin_file="finetune_translate_base.gin" --gin.INITIAL_CHECKPOINT_PATH=${INITIAL_CHECKPOINT_PATH}  --gin.MIXTURE_OR_TASK_NAME=\"translate_long\" --gin.MODEL_DIR=\"gs://nb-t5x-us-central2/finetuned/nynorsk_NCC_base_v5\"