#!/bin/bash
#SBATCH -A nvr_elm_llm                                                                #account
#SBATCH -p batch_block1,batch_block3,batch_block4                                                         #partition
#SBATCH -t 02: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
#SBATCH --exclusive

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

# 默认配置
task=${1:-gsm8k}
model_path=${2:-output_models/finetune_qwen_1.5B_nemotron_shizhe_no_reasoning_pad_bd_size_128_8nodes_seq_length_2k/checkpoint-15000}
batch_size=${3:-32}
use_block_cache=${4:-True}
block_cache_refresh_interval=${5:-32}
small_block_size=${6:-32}
dynamic_threshold=${7:-False}
factor=${8:-1}
threshold=${9:-0.9}

# 构建动态output路径，包含关键参数信息
output_path="eval_results/${task}_batch${batch_size}_block_cache_${use_block_cache}_dynamic_thresh_${dynamic_threshold}_thresh${threshold}"

# 创建必要的目录
mkdir -p ${output_path}
mkdir -p workdir
mkdir -p workdir/slurm_log/eval

echo "开始执行batch evaluation..."
echo "任务: ${task}"
echo "模型路径: ${model_path}"
echo "批次大小: ${batch_size}"
echo "输出路径: ${output_path}"
echo "使用块缓存: ${use_block_cache}"
echo "块缓存刷新间隔: ${block_cache_refresh_interval}"
echo "小块大小: ${small_block_size}"
echo "动态阈值: ${dynamic_threshold}"
echo "因子: ${factor}"
echo "阈值: ${threshold}"

# 构建评估命令
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}"',block_cache_refresh_interval='"${block_cache_refresh_interval}"',small_block_size='"${small_block_size}"',dynamic_threshold='"${dynamic_threshold}"',factor='"${factor}"',threshold='"${threshold}"' \
    --output_path '"${output_path}"' \
    --log_samples
'

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

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

echo "评估完成！结果保存在: ${output_path}"

# 使用说明
echo ""
echo "使用方法:"
echo "  sbatch $0 [task] [model_path] [batch_size] [use_block_cache] [block_cache_refresh_interval] [small_block_size] [dynamic_threshold] [factor] [threshold]"
echo ""
echo "示例:"
echo "  sbatch $0 gsm8k /path/to/model 64 True 64 64 True 1 0.8"
echo "  sbatch $0 mbpp /path/to/model 16 False 32 32 False 1 0.9"
