#!/bin/bash
# Fast_dLLM Evaluation Script using lmms-eval with custom plugin
# 
# Usage:
#   MODEL_PATH=output_models/xxx/checkpoint-2000 bash eval_fast_dllm.sh mmmu
#   MODEL_PATH=output_models/xxx/checkpoint-2000 bash eval_fast_dllm.sh gqa mmmu mathvista
#
# This script uses lmms-eval with the fast_dllm model type which supports:
# - AutoModelForCausalLM with trust_remote_code=True
# - Custom checkpoints that don't have processor config (uses processor_path)

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)"

# Model configuration
MODEL_PATH="${MODEL_PATH:-}"
PROCESSOR_PATH="${PROCESSOR_PATH:-Qwen/Qwen2.5-VL-3B-Instruct}"  # Default processor for finetuned checkpoints
BATCH_SIZE="${BATCH_SIZE:-1}"
NUM_GPUS="${NUM_GPUS:-1}"

# Output configuration
OUTPUT_DIR="${OUTPUT_DIR:-${PROJECT_ROOT}/eval_outputs/lmms_eval}"

# Task mapping (same as eval_lmms.sh)
declare -A TASK_MAP=(
    ["gqa"]="gqa"
    ["mmmu"]="mmmu_val"
    ["textvqa"]="textvqa_val"
    ["pope"]="pope"
    ["mathvista"]="mathvista_testmini"
    ["mmbench"]="mmbench_en_dev"
    ["ai2d"]="ai2d"
    ["chartqa"]="chartqa"
    ["docvqa"]="docvqa_val"
    ["ocrbench"]="ocrbench"
    ["realworldqa"]="realworldqa"
    ["scienceqa"]="scienceqa_img"
)

# Parse arguments
if [ $# -eq 0 ]; then
    echo "Usage: MODEL_PATH=<path> bash $0 <benchmark1> [benchmark2] ..."
    echo ""
    echo "Available benchmarks:"
    for key in "${!TASK_MAP[@]}"; do
        echo "  $key       -> ${TASK_MAP[$key]}"
    done
    echo ""
    echo "Examples:"
    echo "  MODEL_PATH=output_models/ar_sft/checkpoint-2000 bash $0 mmmu"
    echo "  MODEL_PATH=output_models/mdm_sft/checkpoint-2000 bash $0 gqa mmmu mathvista"
    exit 1
fi

if [ -z "$MODEL_PATH" ]; then
    echo "Error: MODEL_PATH is required"
    echo "Usage: MODEL_PATH=<path> bash $0 <benchmark1> [benchmark2] ..."
    exit 1
fi

# Build task list
TASKS=""
for arg in "$@"; do
    if [ -n "${TASK_MAP[$arg]}" ]; then
        if [ -n "$TASKS" ]; then
            TASKS="${TASKS},${TASK_MAP[$arg]}"
        else
            TASKS="${TASK_MAP[$arg]}"
        fi
    else
        echo "Warning: Unknown benchmark '$arg', skipping..."
    fi
done

if [ -z "$TASKS" ]; then
    echo "Error: No valid tasks specified"
    exit 1
fi

echo "=============================================="
echo "Fast_dLLM Evaluation (lmms-eval plugin)"
echo "=============================================="
echo "MODEL_PATH: $MODEL_PATH"
echo "PROCESSOR_PATH: $PROCESSOR_PATH"
echo "TASKS: $TASKS"
echo "BATCH_SIZE: $BATCH_SIZE"
echo "NUM_GPUS: $NUM_GPUS"
echo "OUTPUT_DIR: $OUTPUT_DIR"
echo "=============================================="

# Set up plugin path
cd "$PROJECT_ROOT"
export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH}"
export LMMS_EVAL_PLUGINS="chengyue_script.dvlm.evaluation.lmms_plugin"

# Build model_args
MODEL_ARGS="pretrained=${MODEL_PATH},processor_path=${PROCESSOR_PATH}"

# Run evaluation
if [ "$NUM_GPUS" -gt 1 ]; then
    echo "Running with $NUM_GPUS GPUs using accelerate..."
    accelerate launch --num_processes=${NUM_GPUS} -m lmms_eval \
        --model fast_dllm \
        --model_args "$MODEL_ARGS" \
        --tasks "$TASKS" \
        --batch_size "$BATCH_SIZE" \
        --output_path "$OUTPUT_DIR" \
        --log_samples
else
    echo "Running on single GPU..."
    python -m lmms_eval \
        --model fast_dllm \
        --model_args "$MODEL_ARGS" \
        --tasks "$TASKS" \
        --batch_size "$BATCH_SIZE" \
        --output_path "$OUTPUT_DIR" \
        --log_samples
fi

echo ""
echo "Evaluation complete! Results saved to: $OUTPUT_DIR"


