#!/bin/bash
#SBATCH -A nvr_elm_llm                                                                #account
#SBATCH -p batch_block1,batch_block3,batch_block4                                                         #partition
#SBATCH -t 04:00:00                                                                   #wall time limit, hr:min:sec
#SBATCH -N 1                                                                           #number of nodes
#SBATCH -J batch_eval_qwen                                                             #job name
#SBATCH --gpus-per-node 8
#SBATCH -e workdir/error.log
#SBATCH -o workdir/out.log

# Set the environment variables first before running the command.
export HF_ALLOW_CODE_EVAL=1
export HF_DATASETS_TRUST_REMOTE_CODE=true

# 模型路径 - 7B AR模型
model_path=../evalplus/Qwen2.5-7B-Instruct
log_path=eval_results/qwen2_5_7B_ar_batch_eval.log

# 任务配置
task=gsm8k

# 创建必要的目录
mkdir -p workdir
mkdir -p workdir/slurm_log/qwen2_5_7B_ar_batch_eval
mkdir -p eval_results

echo "开始执行batch evaluation..."
echo "任务: ${task}"
echo "模型路径: ${model_path}"

# 定义要遍历的参数
batch_sizes=(1 2 4 8 16 32 64)

# 计数器
experiment_count=0
total_experiments=${#batch_sizes[@]}

echo "总共需要运行 ${total_experiments} 个实验"

# 遍历所有batch size
for batch_size in "${batch_sizes[@]}"; do
    experiment_count=$((experiment_count + 1))
    
    echo ""
    echo "=========================================="
    echo "实验 ${experiment_count}/${total_experiments}"
    echo "任务: ${task}"
    echo "批次大小: ${batch_size}"
    echo "=========================================="

    # 构建评估命令
    eval_cmd='
    export HF_ALLOW_CODE_EVAL=1;
    export HF_DATASETS_TRUST_REMOTE_CODE=true;
    accelerate launch -m lm_eval --confirm_run_unsafe_code --model hf \
        --model_args pretrained='"${model_path}"' \
        --tasks '"${task}"' \
        --trust_remote_code \
        --apply_chat_template \
        --fewshot_as_multiturn \
        --batch_size '"${batch_size}"'
    '

    echo "执行评估命令..."
    echo $eval_cmd

    # 使用srun执行评估命令，日志路径包含关键参数信息
    srun -o workdir/slurm_log/qwen2_5_7B_ar_batch_eval/slurm_%A_%j_${task}_batch_size_${batch_size}.log -e workdir/slurm_log/qwen2_5_7B_ar_batch_eval/error_%A_%j_${task}_batch_size_${batch_size}.log \
        bash -c "${eval_cmd}"

    echo "实验 ${experiment_count} 完成！"
done

echo ""
echo "所有评估完成！总共运行了 ${total_experiments} 个实验"
