DEVICE=0

SPLITS=( forget01 forget05 forget10 )
NUM_SAMPLES=100
for SEED in 0 1 2
do
    for RANK in 4 8 16 32
    do
        for METHOD in grad_ascent grad_diff idk npo_grad_diff grad_diff_hinge
        do
            for split_idx in "${!SPLITS[@]}"
            do
                SPLIT=${SPLITS[$split_idx]}
                if [ "${SPLIT}" = "forget01" ]
                then
                    STEPS=( 1 2 3 4 5 )
                elif [ "${SPLIT}" = "forget05" ]
                then
                    STEPS=( 6 12 18 24 30 )
                elif [ "${SPLIT}" = "forget10" ]
                then
                    STEPS=( 60 48 36 24 12 )
                else
                    echo "Incorrect split ${SPLIT}"
                fi
    
                for step_idx in "${!STEPS[@]}"
                do
                    CKPT=${STEPS[$step_idx]}
                    CUDA_VISIBLE_DEVICES=${DEVICE} python run_save_generations.py \
                        --split ${SPLIT} \
                        --num_samples ${NUM_SAMPLES} \
                        --base_model_path ./llm_weights/ft_epoch5_lr1e-05_llama3.1-8b_full_wd0.01_seed${SEED}/checkpoint-625/ \
                        --adapter_path ./llm_weights/ft_epoch5_lr1e-05_llama3.1-8b_full_wd0.01_seed${SEED}/checkpoint-625/${METHOD}_target-all_r-${RANK}_0.0001_${SPLIT}_5_seed${SEED}/checkpoint-${CKPT}
                done
            done
        done
    done
done