BSZ=16
SPLITS=( forget01 forget05 forget10 )
RETAINS=( retain99 retain95 retain90 )
CKPTS=( 618 593 562 )
DEVICES=( 0 0 0 0 0 )
DEVICE=0

# Measure likelihoods of retain model on retain model samples
for SEED in 0 1 2
do
    for split_idx in "${!SPLITS[@]}"
    do
        SPLIT=${SPLITS[$split_idx]}
        RETAIN=${RETAINS[$split_idx]}
        CKPT=${CKPTS[$split_idx]}
        CUDA_VISIBLE_DEVICES=${DEVICE} python run_compute_likelihoods.py \
            --split ${SPLIT} \
            --gen_model ./llm_weights/ft_epoch5_lr1e-05_llama3.1-8b_${RETAIN}_wd0.01_seed${SEED}/checkpoint-${CKPT} \
            --test_model ./llm_weights/ft_epoch5_lr1e-05_llama3.1-8b_${RETAIN}_wd0.01_seed${SEED}/checkpoint-${CKPT} \
            --batch_size ${BSZ} \
            --tag self
    done
done

for METHOD in grad_ascent grad_diff npo_grad_diff idk grad_diff_hinge
do
    for SEED in 0 1 2
    do
        for RANK in 4 8 16 32
        do
            for split_idx in "${!SPLITS[@]}"
            do
                SPLIT=${SPLITS[$split_idx]}
                RETAIN=${RETAINS[$split_idx]}
                CKPT=${CKPTS[$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=( 12 24 36 48 60 )
                else
                    echo "Incorrect split ${SPLIT}"
                fi

                # measure likelihood of unlearned model on retain model samples
                for step_idx in "${!STEPS[@]}"
                do
                    STEP=${STEPS[$step_idx]}
                    DEVICE=${DEVICES[$step_idx]}
                    CUDA_VISIBLE_DEVICES=${DEVICE} python run_compute_likelihoods.py \
                        --split ${SPLIT} \
                        --gen_model ./llm_weights/ft_epoch5_lr1e-05_llama3.1-8b_${RETAIN}_wd0.01_seed${SEED}/checkpoint-${CKPT} \
                        --test_model ./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-${STEP} \
                        --batch_size ${BSZ} \
                        --tag ${RETAIN}
                done

                # measure likelihood of unlearned model on unlearned model samples
                for step_idx in "${!STEPS[@]}"
                do
                    STEP=${STEPS[$step_idx]}
                    DEVICE=${DEVICES[$step_idx]}
                    CUDA_VISIBLE_DEVICES=${DEVICE} python run_compute_likelihoods.py \
                        --split ${SPLIT} \
                        --gen_model ./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-${STEP} \
                        --test_model ./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-${STEP} \
                        --batch_size ${BSZ} \
                        --tag self
                done

                # measure likelihood of retain model on unlearned model samples
                for step_idx in "${!STEPS[@]}"
                do
                    STEP=${STEPS[$step_idx]}
                    DEVICE=${DEVICES[$step_idx]}
                    CUDA_VISIBLE_DEVICES=${DEVICE} python run_compute_likelihoods.py \
                        --split ${SPLIT} \
                        --gen_model ./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-${STEP} \
                        --test_model ./llm_weights/ft_epoch5_lr1e-05_llama3.1-8b_${RETAIN}_wd0.01_seed${SEED}/checkpoint-${CKPT} \
                        --batch_size ${BSZ} \
                        --tag ${METHOD}_r${RANK}_step${STEP}
                done
            done
        done
    done
done
