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

for i in 1 2 3; do
    for domain in 'planets-v1' 'clubs-v1' 'software-v1' 'people-v1'; do
        for samples in 10; do
            for variation in 'textonly' 'qa5' 'qaonly5'; do
                OUTPUT_PATH="outputs/models/indomain-${MODEL_NAME}-warmup_${warmup}-lr_${lr}-wd_${wd}-bs_${total_bs}-${scheduler}-epochs_${epochs}-${domain}-samples_${samples}-${variation}-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 "${domain}-${samples}sample-${variation}" \
                        --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}" \
                            --include_path tasks/ \
                            --tasks "fake-${domain}" \
                            --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
done
