#!/bin/bash

source /opt/rh/devtoolset-10/enable

export TRANSFORMERS_CACHE='/data/private_models/xx_models/huggingface'
function count_gpus() {
    perl -e '$slots=split /,/, $ENV{"SLURM_STEP_GPUS"};
    $node = qx[scontrol show hostnames $ENV{"SLURM_JOB_NODELIST"}];
    chomp($node);
    $gpus = qx[ssh $node nvidia-smi --list-gpus | wc -l];
    chomp($gpus);
    print "$gpus"'
}
export NGPUS=$(count_gpus)
echo "Total GPUs=$NGPUS"

export LR=5e-6

export ADV_ATTACK=$1
export NUM_ADV_STEPS=$2
export MODEL=$3
export TRAIN_TASK=$4
export TEST_TASK=$5
export OUT_FILE=$6
export NUM_OPTIM_TOKENS=$7
export EPOCHS=$8
export BATCH_SIZE=$9
export MASTER_PORT=$((29000 + RANDOM % 1000))

echo "MODEL=$MODEL"
echo "TRAIN_TASK=$TRAIN_TASK"
echo "TEST_TASK=$TEST_TASK"
echo "Attack $ADV_ATTACK with $NUM_ADV_STEPS steps"
echo "OUT_FILE=$OUT_FILE"
echo "DO_TRAIN=True"
echo "DO_EVAL=True"
echo "NUM_OPTIM_TOKENS=$NUM_OPTIM_TOKENS"
echo "EPOCHS=$EPOCHS"
echo "BATCH_SIZE=$BATCH_SIZE"
echo "Master port=$MASTER_PORT"

accelerate launch --config_file ../configs/default_accelerate_config.yaml --num_processes $NGPUS --main_process_port $MASTER_PORT \
    ../tune_adv_accelerate.py -v \
    --model $MODEL \
    --learning_rate $LR \
    --do_train \
    --do_eval \
    --dropout 0 \
    --ngpus $NGPUS \
    --nepochs $EPOCHS \
    --weight_decay 0 \
    --warmup_ratio 0.03 \
    --max_length 512 \
    --batch_size $BATCH_SIZE \
    --train_tasks $TRAIN_TASK \
    --test_tasks $TEST_TASK \
    --num_adv_steps $NUM_ADV_STEPS \
    --adv_attack $ADV_ATTACK \
    --num_optim_tokens $NUM_OPTIM_TOKENS \
    --save \
    --out_file $OUT_FILE \

