qwen_test_runner/vision/run_vlmbench.py

"""
run_vlmbench.py — CLI entry point (`qwen-vlmbench`).

Examples:
  # Offline CPU smoke (no torch, no network) — Phase-0 acceptance:
  qwen-vlmbench --runner stub --dataset smoke \
      --categories image_classification bbox_grounding ocr_text

  # Real run on Colab (single RTX 6000 Pro):
  qwen-vlmbench --runner vlm --dataset full --n 200 \
      --models qwen3.5-9b qwen3vl-8b --reasoning instruct thinking \
      --modes json_mode constrained --categories image_classification bbox_grounding ocr_text
"""

from __future__ import annotations

import argparse
import json

from .bench import BenchConfig, run_bench
from .tasks_vision import category_names, pilot_categories


def _build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(prog="qwen-vlmbench", description="Qwen VLM image→JSON labeler benchmark")
    p.add_argument("--models", nargs="+", default=["qwen3.5-0.8b-json-captioner"],
                   help="model keys from the model registry (or any label for --runner stub)")
    p.add_argument("--categories", nargs="+", default=pilot_categories(),
                   help=f"vision categories. all: {category_names()}")
    p.add_argument("--reasoning", nargs="+", default=["instruct"], choices=["instruct", "thinking"])
    p.add_argument("--modes", nargs="+", default=["json_mode"],
                   choices=["json_mode", "constrained", "tool_use", "free"])
    p.add_argument("--n", type=int, default=50, help="samples per category")
    p.add_argument("--dataset", default="smoke", choices=["smoke", "full"])
    p.add_argument("--runner", default="stub", choices=["stub", "vlm"])
    p.add_argument("--precision", default="bf16", choices=["bf16", "fp8", "int4"])
    p.add_argument("--stub-behavior", default="perfect", choices=["perfect", "fragile", "random"])
    p.add_argument("--output-root", default="runs/vision")
    p.add_argument("--gpu-hourly-rate", type=float, default=2.0)
    p.add_argument("--clear-cache-after-model", action="store_true",
                   help="delete each model's HF cache after use (full-array sweeps on a tight SSD)")
    return p


def main(argv: list[str] | None = None) -> int:
    args = _build_parser().parse_args(argv)
    config = BenchConfig(
        models=args.models,
        categories=args.categories,
        reasonings=args.reasoning,
        modes=args.modes,
        n=args.n,
        dataset=args.dataset,
        runner=args.runner,
        precision=args.precision,
        stub_behavior=args.stub_behavior,
        output_root=args.output_root,
        gpu_hourly_rate=args.gpu_hourly_rate,
        clear_cache_after_model=args.clear_cache_after_model,
    )
    summary = run_bench(config)
    print(json.dumps(summary, indent=2))
    print(f"\nLeaderboard: {summary['run_dir']}/leaderboard.md")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())