#!/bin/bash
PROJECT_DIR=${HOME}"/models/t5-nynorsk-norbench"
export PYTHONPATH=${PROJECT_DIR}
echo "PROJECT_DIR is set to: ${PROJECT_DIR}"


FINETUNE_STEPS=20000
EVAL_PREFIX="norbench/norbench_document_sentiment_bs128"
MODEL_BUCKET_DIR="gs://pere-north-t5x/finetuned/"

CHECKPOINT_LIST=(
  "pere-north-t5x/pretrained_models/base/norwegian_NCC_plus_English_t5x_base/checkpoint_1500000"
  "pere-north-t5x/pretrained_models/base/norwegian_NCC_plus_English_t5x_base/checkpoint_1500000"
  "pere-north-t5x/pretrained_models/base/norwegian_NCC_plus_English_t5x_base/checkpoint_1500000"
  "pere-north-t5x/pretrained_models/base/norwegian_NCC_plus_English_t5x_base/checkpoint_1500000"
)

NAME_LIST=(
  "north_t5_base_NCC_LR_0_0006"
  "north_t5_base_NCC_LR_0_0008"
  "north_t5_base_NCC_LR_0_001"
  "north_t5_base_NCC_LR_0_002"
)

TASK_LIST=("document_sentiment" "document_sentiment" "document_sentiment" "document_sentiment")
LR_LIST=(0.0006 0.0008 0.001 0.002)


GIN_LIST=(
  "norbench_document_sentiment_bs128.gin"
  "norbench_document_sentiment_bs128.gin"
  "norbench_document_sentiment_bs128.gin"
  "norbench_document_sentiment_bs128.gin"
)

START_LIST=(1500000 1500000 1500000 1500000)
EXP_LIST=(1 2 3 4)


op through each experiment in EXP_LIST
for i in "${!EXP_LIST[@]}"; do
    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]}
    LR=${LR_LIST[i]}
    EXP=${EXP_LIST[i]}
    MODEL_DIR="${MODEL_BUCKET_DIR}${EVAL_PREFIX}_exp${EXP}_${NAME_LIST[i]}"

    command="python3 ../../t5x/t5x/train.py --gin_search_paths=\"./\" --gin.TRAIN_STEPS=${TRAIN_STEPS} --gin.LR=${LR}  --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