MODEL_ID="Qwen/Qwen2.5-7B"
MODEL_NAME="qwen25_7b"

warmup=0.1
lr=5e-5
wd=0.01
bs=128
accum=1
total_bs=$(($bs * $accum))
scheduler=cosine
epochs=3
maxgradnorm=1.0

mixes=(
    'mix-clubs-v1-10sample-textonly-people-v1-10sample-qa5' 'clubs-v1'
    'mix-clubs-v1-10sample-textonly-planets-v1-10sample-qa5' 'clubs-v1'
    'mix-clubs-v1-10sample-textonly-software-v1-10sample-qa5' 'clubs-v1'
    'mix-people-v1-10sample-textonly-clubs-v1-10sample-qa5' 'people-v1'
    'mix-people-v1-10sample-textonly-planets-v1-10sample-qa5' 'people-v1'
    'mix-people-v1-10sample-textonly-software-v1-10sample-qa5' 'people-v1'
    'mix-planets-v1-10sample-textonly-clubs-v1-10sample-qa5' 'planets-v1'
    'mix-planets-v1-10sample-textonly-people-v1-10sample-qa5' 'planets-v1'
    'mix-planets-v1-10sample-textonly-software-v1-10sample-qa5' 'planets-v1'
    'mix-software-v1-10sample-textonly-clubs-v1-10sample-qa5' 'software-v1'
    'mix-software-v1-10sample-textonly-people-v1-10sample-qa5' 'software-v1'
    'mix-software-v1-10sample-textonly-planets-v1-10sample-qa5' 'software-v1'
)

for i in 1 2 3; do
    for samples in 10; do
        for ((j=0; j<${#mixes[@]}; j+=2)); do
            mix="${mixes[j]}"
            task="${mixes[j+1]}"
            OUTPUT_PATH="outputs/models-v1/test1-${MODEL_NAME}-warmup_${warmup}-lr_${lr}-wd_${wd}-bs_${total_bs}-${scheduler}-epochs_${epochs}-${mix}-gradnorm_${maxgradnorm}-${i}/"

            echo ${OUTPUT_PATH}
            if [ ! -d "${OUTPUT_PATH}" ]; then
                python scripts/train.py \
                    --model_name "${MODEL_ID}" \
                    --dataset_name NAME/iclr-dev \
                    --dataset_config "${mix}" \
                    --train_batch_size ${bs} \
                    --gradient_accumulation_steps ${accum} \
                    --num_train_epochs ${epochs} \
                    --learning_rate ${lr} \
                    --weight_decay ${wd} \
                    --lr_scheduler_type ${scheduler} \
                    --logging_steps 0.01 \
                    --warmup_ratio ${warmup} \
                    --seed ${i} \
                    --output_dir "${OUTPUT_PATH}" \
                    --gpu_flops 990e12 \
                    --max_grad_norm ${maxgradnorm}
            fi

            for epoch in $(seq 1 "${epochs}"); do
                CKPT_PATH="${OUTPUT_PATH}/checkpoint-${epoch}"
                LMEVAL_RESULTS="${CKPT_PATH}/eval"
                if [ ! -d ${LMEVAL_RESULTS} ]; then
                    echo ${CKPT_PATH}
                    lm_eval \
                        --model sglang \
                        --model_args "pretrained=${CKPT_PATH},mem_fraction_static=0.8" \
                        --include_path tasks/ \
                        --tasks "fake-${task}" \
                        --batch_size auto \
                        --output_path "${LMEVAL_RESULTS}"
                fi

                # rm -f ${CKPT_PATH}/model*.safetensors
                # rm -f ${CKPT_PATH}/*.{json,txt,jinja}
            done
        done
    done
done
