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

# # 检查参数
# if [ $# -ne 1 ]; then
#     echo "用法: $0 <checkpoint_path>"
#     echo "例如: $0 output_models/finetune_qwen_0.5B_shizhe_1M_subset_bd_size_64_mdm_group"
#     exit 1
# fi

checkpoint_path=output_models/finetune_qwen_1.5B_nemotron_shizhe_no_reasoning_pad_bd_size_128_8nodes_seq_length_2k

# 设置环境变量
export HF_ALLOW_CODE_EVAL=1
export HF_DATASETS_TRUST_REMOTE_CODE=true
export CUDA_HOME=/lustre/fsw/portfolios/nvr/users/chengyuew/anaconda3/envs/lmflow

# 评估参数设置
task=gsm8k
num_fewshot=4
eval_interval=100  # 每隔100步进行一次评估

# 创建输出目录
output_dir="eval_results/$(basename ${checkpoint_path})_${task}"
log_dir="log/checkpoint_eval_${task}"
mkdir -p ${output_dir} ${log_dir}

echo "开始从路径 ${checkpoint_path} 中提取checkpoint进行评估"
echo "评估任务: ${task}"
echo "评估间隔: ${eval_interval} steps"
echo "输出目录: ${output_dir}"

# 检查路径是否存在
if [ ! -d "${checkpoint_path}" ]; then
    echo "错误: 路径 ${checkpoint_path} 不存在"
    exit 1
fi

# 函数：扫描checkpoint目录
scan_checkpoints() {
    local checkpoints_found=($(find "${checkpoint_path}" -name "checkpoint-*" -type d | sort -V))
    echo "${checkpoints_found[@]}"
}

# 函数：获取checkpoint的step数
get_checkpoint_step() {
    local ckpt_path="$1"
    basename "${ckpt_path}" | sed 's/checkpoint-//'
}

# 函数：检查checkpoint是否需要评估
should_evaluate_checkpoint() {
    local step="$1"
    local first_step="$2"
    local eval_interval="$3"
    
    if [ $((step - first_step)) -eq 0 ] || [ $(( (step - first_step) % eval_interval )) -eq 0 ]; then
        return 0  # 需要评估
    else
        return 1  # 不需要评估
    fi
}

# 初始扫描
echo "开始监控checkpoint目录: ${checkpoint_path}"
echo "评估间隔: ${eval_interval} steps"
echo "等待checkpoint出现..."

# 等待第一个checkpoint出现
while true; do
    initial_checkpoints=($(scan_checkpoints))
    if [ ${#initial_checkpoints[@]} -gt 0 ]; then
        break
    fi
    echo "等待checkpoint出现... (当前时间: $(date))"
    sleep 30
done

# 获取第一个checkpoint的step数
first_checkpoint=${initial_checkpoints[0]}
first_step=$(get_checkpoint_step "${first_checkpoint}")
echo "发现第一个checkpoint: ${first_checkpoint} (step ${first_step})"

# 记录已评估的checkpoint
evaluated_checkpoints=()

# 创建评估结果汇总文件
summary_file="${output_dir}/evaluation_summary.txt"
echo "Checkpoint Evaluation Summary - ${task}" > ${summary_file}
echo "===============================================" >> ${summary_file}
echo "评估时间: $(date)" >> ${summary_file}
echo "评估任务: ${task}" >> ${summary_file}
echo "Few-shot数量: ${num_fewshot}" >> ${summary_file}
echo "评估间隔: ${eval_interval} steps" >> ${summary_file}
echo "" >> ${summary_file}

# 主循环：持续监控和评估新的checkpoint
echo "开始持续监控checkpoint目录..."
while true; do
    # 重新扫描checkpoint目录
    current_checkpoints=($(scan_checkpoints))
    
    if [ ${#current_checkpoints[@]} -eq 0 ]; then
        echo "没有找到checkpoint，等待中... (当前时间: $(date))"
        sleep 60
        continue
    fi
    
    # 检查是否有新的checkpoint需要评估
    new_checkpoints_found=false
    
    for ckpt in "${current_checkpoints[@]}"; do
        step=$(get_checkpoint_step "${ckpt}")
        
        # 检查是否已经评估过
        if [[ " ${evaluated_checkpoints[@]} " =~ " ${step} " ]]; then
            continue
        fi
        
        # 检查是否需要评估
        if should_evaluate_checkpoint "${step}" "${first_step}" "${eval_interval}"; then
            # 检查是否已经存在评估结果
            ckpt_output_dir="${output_dir}/checkpoint-${step}"
            if [ -d "${ckpt_output_dir}" ] && [ "$(ls -1 ${ckpt_output_dir}/results_*.json 2>/dev/null | wc -l)" -gt 0 ]; then
                echo "跳过 checkpoint-${step} (已存在评估结果)"
                evaluated_checkpoints+=("${step}")
                continue
            fi
            
            echo "发现新的checkpoint需要评估: checkpoint-${step}..."
            new_checkpoints_found=true
            
            # 创建当前checkpoint的评估结果目录
            mkdir -p ${ckpt_output_dir}
            
            # 运行评估
            eval_log="${log_dir}/eval_checkpoint-${step}.log"
            
            srun -o ${eval_log} -e ${eval_log} \
                accelerate launch chengyue_script/eval_mdm_qwen.py \
                --tasks ${task} \
                --num_fewshot ${num_fewshot} \
                --confirm_run_unsafe_code \
                --fewshot_as_multiturn \
                --apply_chat_template \
                --model dfm_qwen \
                --model_args model_path=${ckpt},show_speed=True \
                --output_path ${ckpt_output_dir}/results.json
            
            # 检查评估是否成功
            if [ $? -eq 0 ]; then
                echo "✓ checkpoint-${step} 评估完成"
                
                # 标记为已评估
                evaluated_checkpoints+=("${step}")
                
                # 提取评估结果并添加到汇总文件
                results_files=($(ls -1 ${ckpt_output_dir}/results_*.json 2>/dev/null))
                if [ ${#results_files[@]} -gt 0 ]; then
                    latest_results="${results_files[-1]}"  # 获取最新的结果文件
                    echo "checkpoint-${step}: 评估完成" >> ${summary_file}
                    echo "  结果文件: ${latest_results}" >> ${summary_file}
                    
                    # 尝试提取关键指标
                    if command -v jq &> /dev/null; then
                        strict_accuracy=$(jq -r '.results.gsm8k."exact_match,strict-match"' "${latest_results}" 2>/dev/null)
                        flexible_accuracy=$(jq -r '.results.gsm8k."exact_match,flexible-extract"' "${latest_results}" 2>/dev/null)
                        
                        if [ "$strict_accuracy" != "null" ] && [ "$strict_accuracy" != "" ]; then
                            echo "  GSM8K严格匹配准确率: ${strict_accuracy}" >> ${summary_file}
                        fi
                        if [ "$flexible_accuracy" != "null" ] && [ "$flexible_accuracy" != "" ]; then
                            echo "  GSM8K灵活提取准确率: ${flexible_accuracy}" >> ${summary_file}
                        fi
                    fi
                    echo "" >> ${summary_file}
                fi
            else
                echo "✗ checkpoint-${step} 评估失败，查看日志: ${eval_log}"
                echo "checkpoint-${step}: 评估失败" >> ${summary_file}
                echo "  错误日志: ${eval_log}" >> ${summary_file}
                echo "" >> ${summary_file}
            fi
            
            # 添加延迟避免GPU过热
            sleep 10
        fi
    done
    
    # 如果没有新的checkpoint需要评估，等待一段时间
    if [ "$new_checkpoints_found" = false ]; then
        echo "当前没有新的checkpoint需要评估，等待中... (当前时间: $(date))"
        sleep 60
    fi
    
    # 显示当前状态
    echo "已评估的checkpoint: ${evaluated_checkpoints[*]}"
    echo "当前发现的checkpoint总数: ${#current_checkpoints[@]}"
    echo "----------------------------------------"
done

echo "所有评估完成！"
echo "结果保存在: ${output_dir}"
echo "汇总文件: ${summary_file}"
echo "详细日志: ${log_dir}"

# 显示最终汇总
echo ""
echo "评估汇总:"
cat ${summary_file}
