#!/bin/bash
#SBATCH --account project_account
#SBATCH --job-name hi2
#SBATCH --nodes=1
#SBATCH --gpus-per-node=1
#SBATCH --ntasks-per-node=5  # Allocate 30 CPU tasks
#SBATCH --cpus-per-task=4     # Each task gets 2 CPUs (60 CPUs total)
#SBATCH --time=08:20:00
#SBATCH --mail-type=BEGIN,END,FAIL

conda init
conda activate fxt
conda env list 
nvidia-smi
export PYTHONPATH=.

cat $0
echo "--------------------"

export PYTHONPATH=$(pwd)
export HF_HOME="cache"
export WANDB_CACHE_DIR="cache"

C=configs/finetune/xnli_routing_1bp.yml
GPUS=1
config_file=configs/accelerate/gpu_1.yaml
models=(
    # fxt_baseline_1_bp_6_priors_0.1_en_hard_no_binomial/_2025-04-29_22-00-31
    # fxt_baseline_1_bp_6_priors_0.2_en_hard_no_binomial/_2025-04-29_17-06-35
    # fxt_baseline_1_bp_6_priors_0.3_en_hard_no_binomial/_2025-04-29_16-29-47
    # fxt_baseline_1_bp_6_priors_0.3_en_hard_yes_binomial/_2025-04-30_06-11-02
    fxt_baseline_1_bp_6_priors_0.3_en_hard_no_binomial_lambda_1/_2025-05-04_09-31-46
    # fxt_baseline_1_bp_6_priors_0.3_en_hard_no_binomial_lambda_3/_2025-05-03_14-43-31
    # fxt_baseline_1_bp_6_priors_BPE_50k/_2025-05-10_23-51-40    
    # fxt_baseline_1_bp_6_priors_BPE/_2025-05-09_20-40-49
    )

SEEDS=(42)
LRS=(5e-5)
BSZS=(8)
gradient_accumulation_steps=8
LANGS=(ur)

for model in "${models[@]}"; do
    echo "Starting with model ${model}"
    work_dir="model_ckpts/downstream/${model}/"
    echo $work_dir
    for SEED in "${SEEDS[@]}";do
        echo "Starting with seed ${SEED}"

        for LR in "${LRS[@]}";do

            for BSZ in "${BSZS[@]}";do

                for language in "${LANGS[@]}"; do
                    echo 'Finding free port'
                    PORT=$(python -c 'import socket; s=socket.socket(); s.bind(("", 0)); print(s.getsockname()[1]); s.close()')
                    accelerate launch --main_process_port=$PORT --config_file=$config_file --num_processes="$GPUS" src/finetune/train_classification.py  \
                        --config_file "$C" \
                        --pretrained_path "model_ckpts/${model}/" \
                        --work_dir $work_dir \
                        --language $language \
                        --lr  $LR \
                        --batch_size $BSZ \
                        --seed $SEED \
                        --gradient_accumulation_steps $gradient_accumulation_steps \
                        --use_best_model False \

                done
            done
        done
    done
done


#rm -r /fs/ess/project_account/your_username/experiments/fxt-base/model_ckpts/downstream/fxt_baseline_1_bp_6_priors_0.3_en_hard_yes_binomial__2025-04-30_06-11-02_/*/*/*/epoch*