#!/bin/bash
# ==============================================================================
# Universal lmms-eval Evaluation Script
# Supports: GQA, MMMU, TextVQA, POPE, MathVista, MMBench, and more
# NO GPT Judge Required!
# ==============================================================================

set -eo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# ==============================================================================
# Benchmark Presets (lmms-eval task names)
# ==============================================================================
declare -A BENCHMARK_TASKS=(
    ["gqa"]="gqa"
    ["mmmu"]="mmmu_val"
    ["textvqa"]="textvqa_val"
    ["pope"]="pope"
    ["mathvista"]="mathvista_testmini"
    ["mmbench"]="mmbench_en_dev"
    ["mme"]="mme"
    ["ai2d"]="ai2d"
    ["scienceqa"]="scienceqa_img"
    ["realworldqa"]="realworldqa"
    ["ocrbench"]="ocrbench"
    ["docvqa"]="docvqa_val"
    ["chartqa"]="chartqa"
    ["vqav2"]="vqav2_val"
    ["seed"]="seedbench"
    ["mmstar"]="mmstar"
    ["mmvet"]="mmvet"
)

# ==============================================================================
# Parse Arguments
# ==============================================================================
show_help() {
    echo "Usage: $0 <benchmark> [benchmark2 ...] [options]"
    echo ""
    echo "lmms-eval: NO GPT Judge Required!"
    echo ""
    echo "Benchmarks:"
    echo "  gqa        - GQA (VQA, exact match)"
    echo "  mmmu       - MMMU Validation (MCQ)"
    echo "  textvqa    - TextVQA Validation (VQA)"
    echo "  pope       - POPE (Y/N classification)"
    echo "  mathvista  - MathVista TestMini (math reasoning)"
    echo "  mmbench    - MMBench English Dev (MCQ)"
    echo "  mme        - MME (perception + reasoning)"
    echo "  ai2d       - AI2D (diagram understanding)"
    echo "  scienceqa  - ScienceQA IMG (science questions)"
    echo "  ocrbench   - OCRBench (OCR capability)"
    echo "  docvqa     - DocVQA Validation (document QA)"
    echo "  chartqa    - ChartQA (chart understanding)"
    echo "  vqav2      - VQAv2 Validation"
    echo "  seed       - SEEDBench"
    echo "  mmstar     - MMStar"
    echo "  mmvet      - MM-Vet"
    echo ""
    echo "Options (set via environment variables):"
    echo "  MODEL_PATH       - Path to model checkpoint or HF model ID"
    echo "  MODEL_TYPE       - lmms-eval model type (default: qwen2_vl)"
    echo "  BATCH_SIZE       - Batch size (default: 1)"
    echo "  OUTPUT_DIR       - Output directory (default: ./eval_outputs/lmms_eval)"
    echo "  LIMIT            - Limit number of samples (optional, for debugging)"
    echo "  NUM_GPUS         - Number of GPUs for tensor parallel (default: 1)"
    echo ""
    echo "Examples:"
    echo "  $0 gqa"
    echo "  $0 mmmu pope textvqa"
    echo "  MODEL_PATH=Qwen/Qwen2.5-VL-7B-Instruct $0 mathvista"
    echo "  LIMIT=100 $0 gqa  # Only run 100 samples"
    exit 0
}

if [ $# -eq 0 ] || [ "$1" == "-h" ] || [ "$1" == "--help" ]; then
    show_help
fi

# Collect benchmarks from arguments
BENCHMARKS="$@"

# ==============================================================================
# Configuration
# ==============================================================================
PYTHON_BIN=${PYTHON_BIN:-python}
MODEL_PATH=${MODEL_PATH:-Qwen/Qwen2.5-VL-3B-Instruct}
MODEL_TYPE=${MODEL_TYPE:-qwen2_vl}
BATCH_SIZE=${BATCH_SIZE:-1}
OUTPUT_DIR=${OUTPUT_DIR:-$(pwd)/eval_outputs/lmms_eval}
LIMIT=${LIMIT:-}
NUM_GPUS=${NUM_GPUS:-1}
LOG_SAMPLES=${LOG_SAMPLES:-true}

mkdir -p "${OUTPUT_DIR}"

# ==============================================================================
# Build Task List
# ==============================================================================
TASK_LIST=""

for benchmark in ${BENCHMARKS}; do
    benchmark_lower=$(echo "$benchmark" | tr '[:upper:]' '[:lower:]')
    
    if [ -z "${BENCHMARK_TASKS[$benchmark_lower]}" ]; then
        echo "Warning: Unknown benchmark '$benchmark', treating as task name"
        TASK_LIST="${TASK_LIST},${benchmark}"
    else
        TASK_LIST="${TASK_LIST},${BENCHMARK_TASKS[$benchmark_lower]}"
    fi
done

# Remove leading comma
TASK_LIST=$(echo "$TASK_LIST" | sed 's/^,//')

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

# ==============================================================================
# Print Configuration
# ==============================================================================
echo "=============================================="
echo "lmms-eval Evaluation Configuration"
echo "=============================================="
echo "BENCHMARKS: ${BENCHMARKS}"
echo "TASKS: ${TASK_LIST}"
echo "MODEL_PATH: ${MODEL_PATH}"
echo "MODEL_TYPE: ${MODEL_TYPE}"
echo "BATCH_SIZE: ${BATCH_SIZE}"
echo "OUTPUT_DIR: ${OUTPUT_DIR}"
echo "NUM_GPUS: ${NUM_GPUS}"
if [ -n "$LIMIT" ]; then
    echo "LIMIT: ${LIMIT}"
fi
echo ""
echo "Advantage: NO GPT/LLM Judge Required!"
echo "=============================================="

# ==============================================================================
# Build Command
# ==============================================================================
CMD="${PYTHON_BIN} -m lmms_eval"
CMD="${CMD} --model ${MODEL_TYPE}"
CMD="${CMD} --model_args pretrained=${MODEL_PATH}"
CMD="${CMD} --tasks ${TASK_LIST}"
CMD="${CMD} --batch_size ${BATCH_SIZE}"
CMD="${CMD} --output_path ${OUTPUT_DIR}"

if [ "$LOG_SAMPLES" == "true" ]; then
    CMD="${CMD} --log_samples"
fi

if [ -n "$LIMIT" ]; then
    CMD="${CMD} --limit ${LIMIT}"
fi

# Add tensor parallel if using multiple GPUs
if [ "$NUM_GPUS" -gt 1 ]; then
    CMD="${CMD} --model_args pretrained=${MODEL_PATH},device_map=auto"
fi

echo "Command: ${CMD}"
echo ""

# ==============================================================================
# Run Evaluation
# ==============================================================================
eval ${CMD}

echo "=============================================="
echo "Done! Results are under: ${OUTPUT_DIR}"
echo "=============================================="

