#!/bin/bash
# Prompt: HumanEval/MBPP eval wrapper for sbatch
# ==============================================================================
# HumanEval / MBPP Evaluation Script
# Tests text-only code generation ability of dVLM models
# ==============================================================================
set -eo pipefail

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

# Conda environment
EVAL_CONDA_ENV=${EVAL_CONDA_ENV:-dvlm}
if [ -f "${HOME}/.bashrc" ]; then
    source "${HOME}/.bashrc"
fi
if command -v conda >/dev/null 2>&1; then
    conda activate "${EVAL_CONDA_ENV}" || true
fi

# Configuration
MODEL_PATH=${MODEL_PATH:?"MODEL_PATH must be set"}
PROCESSOR_PATH=${PROCESSOR_PATH:-Qwen/Qwen2.5-VL-3B-Instruct}
TASK=${TASK:-humaneval}          # humaneval or mbpp
N_SAMPLES=${N_SAMPLES:-1}       # samples per problem (1 for pass@1)
MAX_TOKENS=${MAX_TOKENS:-512}
BLOCK_SIZE=${BLOCK_SIZE:-}
OUTPUT_DIR=${OUTPUT_DIR:-}

# PYTHONPATH
DVLM_DIR="${SCRIPT_DIR}/../.."
export PYTHONPATH="${DVLM_DIR}/models:${PROJECT_ROOT}/src:${PROJECT_ROOT}/third_party/bigcode-eval:${PYTHONPATH:-}"

# Build args
ARGS="--model_path ${MODEL_PATH} --processor_path ${PROCESSOR_PATH}"
ARGS="${ARGS} --task ${TASK} --n_samples ${N_SAMPLES} --max_tokens ${MAX_TOKENS}"
[ -n "${OUTPUT_DIR}" ] && ARGS="${ARGS} --output_dir ${OUTPUT_DIR}"
[ -n "${BLOCK_SIZE}" ] && ARGS="${ARGS} --block_size ${BLOCK_SIZE}"

echo "Running ${TASK} evaluation..."
echo "MODEL_PATH: ${MODEL_PATH}"
echo "TASK: ${TASK}"
echo "N_SAMPLES: ${N_SAMPLES}"
echo ""

python3 "${SCRIPT_DIR}/eval_humaneval_mbpp.py" ${ARGS}
