MASTER_PORT=$((RANDOM % 50001 + 10000))

forget_losses=(
    IDK+AP
)

# You can specify any forget task from 1 to 10
# The standard TOFU benchmark is task 1
task_list=(1)

# Pass to python script
export TASK_LIST=$(IFS=,; echo "${task_list[*]}")

learning_rates=(
    1e-5
)

mask=false
dual_optim=true
use_LoRA=false
save_root_base=/path/to/closer-look-LLM-unlearning/results/tofu_phi1-5_final/ours

forget_coeff=0.1
regularization_coeff=1.0

save_checkpoint=true

num_epochs=5

save_steps=steps_per_epoch

forget_lr=(7.98e-06)  # Fixed forget learning rate

declare -A eval_steps_map
eval_steps_map[forget01]="1 2 3 4 5 6"
eval_steps_map[forget05]="6 12 18 24 30 31"
eval_steps_map[forget10]="12 24 36 48 60 62"

for flr in ${forget_lr[@]}; do
    save_root="${save_root_base}/lr_${flr}"
    for split in forget01 forget05 forget10; do
        eval_steps=(${eval_steps_map[$split]})  # Assign correct eval steps based on forget level
        for forget_loss in ${forget_losses[@]}; do
            for lr in ${learning_rates[@]}; do
                for task_id in ${task_list[@]}; do
                    COMMON="use_LoRA=$use_LoRA forget_coeff=$forget_coeff regularization_coeff=$regularization_coeff lr=$lr forget_lr=$flr split=$split forget_loss=$forget_loss num_epochs=$num_epochs \
                        mask=$mask fix_ref_model=$fix_ref_model save_root=$save_root save_checkpoint=$save_checkpoint dual_optim=$dual_optim"
                    CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node=2 --master_port=$MASTER_PORT \
                            forget.py \
                            --config-name=phi1-5_tofu.yaml \
                            task_id=$task_id \
                            save_steps=$save_steps \
                            $COMMON
                    for step in ${eval_steps[@]}; do
                        CUDA_VISIBLE_DEVICES=0 torchrun --nproc_per_node=1 --master_port=$MASTER_PORT \
                                eval.py \
                                --config-name=phi1-5_tofu.yaml \
                                task_id=$task_id \
                                eval_unlearn_step=$step \
                                $COMMON
                    done
                done
            done
        done
    done
done
