#!/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

# 任务配置
task=gsm8k
# task=mbpp
# task=humaneval
# task=minerva_math

# 模型路径
model_path=models/Fast_dLLM_1.5B

# 评估参数
block_cache_refresh_interval=32
dynamic_threshold=False
factor=1
threshold=0.9

# 创建必要的目录
mkdir -p workdir
mkdir -p workdir/slurm_log/small_block_experiment_threshold_0.9

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

# 定义要遍历的参数
small_block_sizes=(4 8 16)
batch_sizes=(1 32)
use_block_cache_options=(True False)

# 计数器
experiment_count=0
total_experiments=$((${#small_block_sizes[@]} * ${#batch_sizes[@]} * ${#use_block_cache_options[@]}))

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

# 嵌套for循环遍历所有参数组合
for small_block_size in "${small_block_sizes[@]}"; do
    for batch_size in "${batch_sizes[@]}"; do
        for use_block_cache in "${use_block_cache_options[@]}"; do
            experiment_count=$((experiment_count + 1))
            
            echo ""
            echo "=========================================="
            echo "实验 ${experiment_count}/${total_experiments}"
            echo "小块大小: ${small_block_size}"
            echo "批次大小: ${batch_size}"
            echo "使用块缓存: ${use_block_cache}"
            echo "动态阈值: ${dynamic_threshold}"
            echo "阈值: ${threshold}"
            echo "=========================================="

            # 构建评估命令
            eval_cmd='
            export HF_ALLOW_CODE_EVAL=1;
            export HF_DATASETS_TRUST_REMOTE_CODE=true;
            accelerate launch chengyue_script/eval_mdm_qwen_batch.py \
                --tasks '"${task}"' \
                --batch_size '"${batch_size}"' \
                --confirm_run_unsafe_code \
                --model dfm_qwen \
                --fewshot_as_multiturn \
                --apply_chat_template \
                --model_args model_path='"${model_path}"',show_speed=True,use_block_cache='"${use_block_cache}"',small_block_size='"${small_block_size}"',threshold='"${threshold}"',bd_size=32,max_new_tokens=512
            '

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

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

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

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