rpx_benchmark.vqa.roster

  1"""Frozen single-image model roster supplied for the RPX VQA benchmark."""
  2
  3from __future__ import annotations
  4
  5from dataclasses import dataclass
  6
  7from ..exceptions import ConfigError
  8
  9
 10@dataclass(frozen=True)
 11class ModelSpec:
 12    key: str
 13    display_name: str
 14    size_b: float
 15    vram_gb: int
 16    bbox_mode: str
 17    capabilities: frozenset[str]
 18    diagnostic_capabilities: frozenset[str] = frozenset()
 19
 20
 21ALL_TASKS = frozenset({"bbox"})
 22FLORENCE_DIAGNOSTICS = frozenset({"semantic_label", "phrase_grounding"})
 23PALIGEMMA_DIAGNOSTICS = frozenset({"semantic_label", "object_detection"})
 24
 25MODELS = (
 26    # Native localization models receive the original question as their
 27    # referring expression and return a bbox in one scored generation.
 28    ModelSpec(
 29        "florence2-base",
 30        "Florence 2 Base",
 31        0.23,
 32        4,
 33        "native_loc",
 34        ALL_TASKS,
 35        FLORENCE_DIAGNOSTICS,
 36    ),
 37    ModelSpec(
 38        "florence2-large",
 39        "Florence 2 Large",
 40        0.77,
 41        6,
 42        "native_loc",
 43        ALL_TASKS,
 44        FLORENCE_DIAGNOSTICS,
 45    ),
 46    ModelSpec(
 47        "paligemma2-3b",
 48        "PaliGemma 2 3B",
 49        3.0,
 50        7,
 51        "native_loc",
 52        ALL_TASKS,
 53        PALIGEMMA_DIAGNOSTICS,
 54    ),
 55    ModelSpec("qwen2.5-vl-3b", "Qwen2.5-VL 3B", 3.0, 7, "native", ALL_TASKS),
 56    ModelSpec("qwen3-vl-2b", "Qwen3-VL 2B", 2.0, 7, "native", ALL_TASKS),
 57    ModelSpec(
 58        "cosmos-reason2-2b",
 59        "Cosmos Reason 2 2B",
 60        2.44,
 61        24,
 62        "native_loc",
 63        ALL_TASKS,
 64    ),
 65    ModelSpec(
 66        "deepseek-vl2-tiny",
 67        "DeepSeek-VL2 Tiny (1B active)",
 68        1.0,
 69        16,
 70        "native_loc",
 71        ALL_TASKS,
 72    ),
 73    ModelSpec("internvl3.5-1b", "InternVL 3.5 1B", 1.1, 8, "native_loc", ALL_TASKS),
 74    ModelSpec("gemma4-e4b", "Gemma 4 E4B", 4.5, 18, "prompted", ALL_TASKS),
 75    ModelSpec("phi-3.5-vision-4b", "Phi-3.5-Vision 4.2B", 4.2, 10, "prompted", ALL_TASKS),
 76    ModelSpec("llava-onevision-7b", "LLaVA-OneVision 7B", 7.0, 16, "prompted", ALL_TASKS),
 77    ModelSpec("qwen2.5-vl-7b", "Qwen2.5-VL 7B", 7.0, 15, "native", ALL_TASKS),
 78    ModelSpec("qwen3-vl-8b", "Qwen3-VL 8B", 8.0, 18, "native", ALL_TASKS),
 79    ModelSpec(
 80        "cosmos-reason2-8b",
 81        "Cosmos Reason 2 8B",
 82        8.0,
 83        32,
 84        "native_loc",
 85        ALL_TASKS,
 86    ),
 87    ModelSpec("idefics3-8b", "Idefics3 8B", 8.0, 18, "prompted", ALL_TASKS),
 88    ModelSpec("internvl2.5-8b", "InternVL 2.5 8B", 8.0, 17, "prompted", ALL_TASKS),
 89    ModelSpec("internvl3.5-14b", "InternVL 3.5 14B", 15.1, 36, "native_loc_tp2", ALL_TASKS),
 90    ModelSpec(
 91        "deepseek-vl2",
 92        "DeepSeek-VL2 (4.5B active)",
 93        4.5,
 94        64,
 95        "native_loc_tp2",
 96        ALL_TASKS,
 97    ),
 98    ModelSpec(
 99        "paligemma2-10b",
100        "PaliGemma 2 10B",
101        10.0,
102        21,
103        "native_loc",
104        ALL_TASKS,
105        PALIGEMMA_DIAGNOSTICS,
106    ),
107    ModelSpec("gemma4-12b", "Gemma 4 12B", 12.0, 28, "prompted", ALL_TASKS),
108)
109
110
111def get_model(key: str) -> ModelSpec:
112    try:
113        return next(model for model in MODELS if model.key == key)
114    except StopIteration as exc:
115        raise ConfigError(
116            f"unknown VQA model {key!r}", hint="choose a key from the frozen model roster"
117        ) from exc