run_minicpm_multi.sh

#!/usr/bin/env bash
# ============================================================
# StreamArena - MiniCPM-o-4.5 unified multi-GPU evaluation
#
# Shards questions by video across N GPUs, merges shard*.jsonl, optionally
# runs the shared LLM-as-Judge scorer.
#
# Required env:
#   MODEL_PATH        local MiniCPM-o checkpoint dir
#
# Optional env:
#   DATASET           HF dataset id or local dataset dir (default hkuzxc/StreamArena)
#   VIDEO_DIR         where <video_id>/<video_id>.mp4 or <video_id>.tar live
#   OUT_ROOT          output root (default ./results)
#   RUN_NAME          subfolder name (default: model basename)
#   NUM_GPUS / GPU_OFFSET      shard count / starting CUDA id (defaults 8 / 0)
#   INCLUDE_E                  set to 1 to also evaluate Pro (E-class)
#   E_WINDOW                   Pro answering window seconds (default 3)
#   AUTO_JUDGE                 1 to score with JUDGE_URL / JUDGE_API_KEY
#   SERPER_API_KEY             only needed if the model calls tools
# ============================================================
set -euo pipefail

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

: "${MODEL_PATH:?Set MODEL_PATH to your local MiniCPM-o-4.5 checkpoint}"
DATASET="${DATASET:-hkuzxc/StreamArena}"
LANGUAGE_ARG="${STREAMARENA_LANGUAGE:-en}"
VIDEO_DIR="${VIDEO_DIR:-${STREAMARENA_VIDEO_DIR:-./videos}}"
NUM_GPUS="${NUM_GPUS:-8}"
GPU_OFFSET="${GPU_OFFSET:-0}"
INCLUDE_E="${INCLUDE_E:-1}"    # default on, matches upstream run_minicpmo_unified_multi.sh
E_WINDOW="${E_WINDOW:-3}"
OUT_ROOT="${OUT_ROOT:-$SCRIPT_DIR/results}"
RUN_NAME="${RUN_NAME:-$(basename "$MODEL_PATH")}"
METHOD=minicpm

OUT_DIR="$OUT_ROOT/$RUN_NAME"
MERGED="$OUT_DIR/${METHOD}_output.jsonl"
mkdir -p "$OUT_DIR/logs"

INCLUDE_E_FLAG=()
[ "$INCLUDE_E" = "1" ] && INCLUDE_E_FLAG+=("--include_e")

echo "=== MiniCPM-o | run=$RUN_NAME | $NUM_GPUS GPUs (offset $GPU_OFFSET) ==="
echo "    dataset=$DATASET  video_dir=$VIDEO_DIR"
echo "    out=$OUT_DIR"
for i in $(seq 0 $((NUM_GPUS-1))); do
    GPU_ID=$((GPU_OFFSET + i))
    OUTPUT="$OUT_DIR/shard${i}.jsonl"
    LOG="$OUT_DIR/logs/shard${i}.log"
    CUDA_VISIBLE_DEVICES=$GPU_ID python "$SCRIPT_DIR/run_minicpm.py" \
        --model_path "$MODEL_PATH" \
        --dataset "$DATASET" --language "$LANGUAGE_ARG" \
        --video_dir "$VIDEO_DIR" \
        --output "$OUTPUT" \
        --device "cuda:0" \
        --e_window_sec "$E_WINDOW" \
        --resume \
        --shard_id $i --num_shards $NUM_GPUS \
        "${INCLUDE_E_FLAG[@]+"${INCLUDE_E_FLAG[@]}"}" \
        "$@" > "$LOG" 2>&1 &
    echo "  shard $i -> GPU $GPU_ID -> $OUTPUT"
done

echo "Waiting for all $NUM_GPUS shards..."
wait
echo "=== All shards done. Merging ==="

python -c "
import glob
with open('$MERGED','w') as out:
    n = 0
    for f in sorted(glob.glob('$OUT_DIR/shard*.jsonl')):
        for line in open(f):
            out.write(line); n += 1
print(f'Merged {n} records -> $MERGED')
"

if [ "${AUTO_JUDGE:-0}" = "1" ]; then
    : "${JUDGE_URL:?Set JUDGE_URL for --auto-judge}"
    : "${JUDGE_API_KEY:?Set JUDGE_API_KEY for --auto-judge}"
    echo "=== Auto-judge ==="
    bash "$SCRIPT_DIR/../../judge/scripts/run_judge.sh" "$MERGED"
fi

echo "=== DONE  merged=$MERGED  judged=${MERGED}.judge.jsonl ==="