#!/bin/bash

GPUID=1
echo "Run on GPU $GPUID"

# data
PROJECT_ROOT=$(dirname "$(readlink -f "$0")")/..
DATA_ROOT=./dataset/snips/
SHOT=20



# model
MODEL_TYPE=bert
# params
LR=5e-5
CONTROLER_LR=5e-5
WEIGHT_DECAY=5e-6
EPOCH=70
SEED=0

ADAM_EPS=1e-8
ADAM_BETA1=0.9
ADAM_BETA2=0.98
WARMUP=20

USE_CONTROLER=weight
MT_ALPHA1=0.995
MT_ALPHA2=0.99
TRAIN_BATCH=16
UN_TRAIN_BATCH=32
EVAL_BATCH=32
export MAX_LENGTH=64
export BERT_MODEL=bert-base-uncased
MODEL_NAME=bert-base-uncased


# self-training parameters
REINIT=1
BEGIN_STEP=2000
LABEL_MODE=hard
PERIOD=3000
HP_LABEL=5.0

python3 ./src/preprocess.py $DATA_ROOT/train_whole_${SHOT}.txt $BERT_MODEL $MAX_LENGTH > $DATA_ROOT/train.txt
python3 ./src/preprocess.py $DATA_ROOT/unlabeled_train_whole_${SHOT}.txt $BERT_MODEL $MAX_LENGTH > $DATA_ROOT/unlabeled_train.txt
python3 ./src/preprocess.py $DATA_ROOT/valid_whole.txt $BERT_MODEL $MAX_LENGTH > $DATA_ROOT/dev.txt
python3 ./src/preprocess.py $DATA_ROOT/test_whole.txt $BERT_MODEL $MAX_LENGTH > $DATA_ROOT/test.txt
cat $DATA_ROOT/train.txt $DATA_ROOT/unlabeled_train.txt $DATA_ROOT/dev.txt $DATA_ROOT/test.txt | cut -d " " -f 2 | grep -v "^$"| sort | uniq > $DATA_ROOT/labels.txt

# output
OUTPUT=$PROJECT_ROOT/outputs/snips/${SHOT}/${MODEL_TYPE}_reinit${REINIT}_begin${BEGIN_STEP}_period${PERIOD}_${LABEL_MODE}_hp${HP_LABEL}_${EPOCH}_${LR}_${USE_CONTROLER}/

[ -e $OUTPUT/script  ] || mkdir -p $OUTPUT/script
cp -f $(readlink -f "$0") $OUTPUT/script
rsync -ruzC --exclude-from=$PROJECT_ROOT/.gitignore --exclude 'dataset' --exclude 'pretrained_model' --exclude 'outputs' $PROJECT_ROOT/ $OUTPUT/src

CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=$GPUID python3 ./src/MetaST.py --data_dir $DATA_ROOT \
  --model_type $MODEL_TYPE --model_name_or_path $MODEL_NAME \
  --learning_rate $LR \
  --weight_decay $WEIGHT_DECAY \
  --adam_epsilon $ADAM_EPS \
  --adam_beta1 $ADAM_BETA1 \
  --adam_beta2 $ADAM_BETA2 \
  --num_train_epochs $EPOCH \
  --warmup_steps $WARMUP \
  --per_gpu_train_batch_size $TRAIN_BATCH \
  --per_pgu_labeled_batch_size 4 \
  --per_gpu_unlabeled_train_batch_size $UN_TRAIN_BATCH \
  --per_gpu_eval_batch_size $EVAL_BATCH \
  --logging_steps 200 \
  --save_steps 1000 \
  --do_train \
  --do_eval \
  --do_predict \
  --overwrite_cache \
  --evaluate_during_training \
  --output_dir $OUTPUT \
  --cache_dir $PROJECT_ROOT/pretrained_model \
  --mt_alpha1 $MT_ALPHA1 \
  --mt_alpha2 $MT_ALPHA2  \
  --seed $SEED \
  --max_seq_length $MAX_LENGTH \
  --overwrite_output_dir \
  --self_training_reinit $REINIT --self_training_begin_step $BEGIN_STEP --self_training_period $PERIOD \
  --self_training_label_mode $LABEL_MODE --self_training_hp_label $HP_LABEL \
  --controler_step_size $CONTROLER_LR \
  --use_controler \
  --controler_sampling_steps -1 \
  --per_gpu_meta_batch_size 32 \
  --labeled_beta 1 \
  --labeled_beta_decay 0.9\
  --use_labeled_loss \
  --max_steps 20000 \
  --sampling_strategy loss_decay \
  --use_token_weight 1 \
  --smoothness_type max \



