(7/n) Complete declarative task runtime and shared pairwise orchestration

judgearena/benchmarks/mt_bench/runner.py

@@ -19,6 +19,7 @@
 )
 from judgearena.benchmarks.mt_bench.preset_judging import judge_mt_bench_with_preset
 from judgearena.benchmarks.pairwise.baselines import native_pairwise_baseline
+from judgearena.benchmarks.pairwise.scoring import PAIRWISE_SCORERS
 from judgearena.datasets import load_instructions
 from judgearena.datasets.mt_bench import (
     load_mt_bench_model_answers,
@@ -27,11 +28,9 @@
 from judgearena.log import get_logger
 from judgearena.models import is_thinking_model, make_model
 from judgearena.prompts.registry import ResolvedJudgePrompt, resolve_run_judge_prompt
-from judgearena.tasks.registry import get_packaged_task
 from judgearena.tasks.schema import MTBenchProtocol
 from judgearena.utils import (
     cache_function_dataframe,
-    compute_pref_summary,
     generation_cache_token,
 )
 from judgearena.utils.eval import BattleReport, _compute_grouped_stats
@@ -40,13 +39,7 @@
 
 if TYPE_CHECKING:
     from judgearena.config import RunConfig
-
-
-def _task_protocol(task_id: str) -> MTBenchProtocol:
-    task = get_packaged_task(task_id)
-    if task is None or not isinstance(task.spec.protocol, MTBenchProtocol):
-        raise ValueError(f"Task {task_id!r} does not define an MT-Bench protocol.")
-    return task.spec.protocol
+    from judgearena.tasks.schema import ResolvedTaskSpec
 
 
 def _align_mt_bench_completions(
@@ -90,10 +83,10 @@ def _build_mt_bench_generation_kwargs(
 
 def _generate_mt_bench_completions(
     cfg: RunConfig,
+    protocol: MTBenchProtocol,
     questions_df: pd.DataFrame,
 ) -> tuple[pd.DataFrame, pd.DataFrame]:
     cache_prefix = cfg.task
-    protocol = _task_protocol(cfg.task)
 
     def _run_generation(
         model_name: str, *, generation_kwargs: dict[str, object]
@@ -198,6 +191,7 @@ def _save_mt_bench_results(
 def _finalize_mt_bench_run(
     *,
     cfg: RunConfig,
+    protocol: MTBenchProtocol,
     res_folder: Path,
     result_name: str,
     prefs: pd.Series,
@@ -210,7 +204,8 @@ def _finalize_mt_bench_run(
     started_at_utc: datetime,
     extra_result_fields: dict[str, object] | None = None,
 ) -> pd.Series:
-    stats = compute_pref_summary(prefs)
+    scorer = PAIRWISE_SCORERS[protocol.scoring.adapter]
+    stats = scorer.summarize(prefs)
     report = BattleReport(
         task=cfg.task,
         model_a=cfg.model.name,
@@ -253,17 +248,16 @@ def _finalize_mt_bench_run(
 def _run_mt_bench_fastchat(
     *,
     cfg: RunConfig,
+    protocol: MTBenchProtocol,
     res_folder: Path,
     result_name: str,
     questions_df: pd.DataFrame,
     completions_a: pd.DataFrame,
     completions_b: pd.DataFrame,
     judge_chat_model,
     resolved_prompt: ResolvedJudgePrompt,
-    fastchat_prompt_preset: str,
     started_at_utc: datetime,
 ) -> pd.Series:
-    protocol = _task_protocol(cfg.task)
     prefs, annotations, combined_metadata, num_inconsistent = (
         judge_mt_bench_pairwise_fastchat(
             judge_chat_model=judge_chat_model,
@@ -278,12 +272,13 @@ def _run_mt_bench_fastchat(
             truncate_input_chars=cfg.generation.truncate_judge_input_chars,
             use_tqdm=cfg.run.use_tqdm,
             reference_categories=protocol.judge.reference_categories,
-            prompt_preset=fastchat_prompt_preset,
+            prompt_preset=protocol.judge.fastchat_prompt_preset,
             strip_thinking_before_judging=cfg.judge.strip_thinking_before_judging,
         )
     )
     return _finalize_mt_bench_run(
         cfg=cfg,
+        protocol=protocol,
         res_folder=res_folder,
         result_name=result_name,
         prefs=prefs,
@@ -301,6 +296,7 @@ def _run_mt_bench_fastchat(
 def _run_mt_bench_preset(
     *,
     cfg: RunConfig,
+    protocol: MTBenchProtocol,
     res_folder: Path,
     result_name: str,
     questions_df: pd.DataFrame,
@@ -310,7 +306,6 @@ def _run_mt_bench_preset(
     resolved_prompt: ResolvedJudgePrompt,
     started_at_utc: datetime,
 ) -> pd.Series:
-    protocol = _task_protocol(cfg.task)
     prefs, annotations, combined_metadata = judge_mt_bench_with_preset(
         judge_chat_model=judge_chat_model,
         judge_model=cfg.judge.model,
@@ -332,6 +327,7 @@ def _run_mt_bench_preset(
     )
     return _finalize_mt_bench_run(
         cfg=cfg,
+        protocol=protocol,
         res_folder=res_folder,
         result_name=result_name,
         prefs=prefs,
@@ -345,10 +341,12 @@ def _run_mt_bench_preset(
     )
 
 
-def run_mt_bench_benchmark(cfg: RunConfig):
+def run_mt_bench_benchmark(cfg: RunConfig, task: ResolvedTaskSpec | None = None):
     """Run the registered MT-Bench generation, judging, and reporting lifecycle."""
     run_started_at = datetime.now(UTC)
-    protocol = _task_protocol(cfg.task)
+    protocol = task.spec.protocol if task is not None else None
+    if not isinstance(protocol, MTBenchProtocol):
+        raise ValueError(f"Task {cfg.task!r} does not define an MT-Bench protocol.")
     if cfg.model.baseline is None:
         baseline = native_pairwise_baseline(cfg.task)
         cfg.model.baseline = baseline if isinstance(baseline, str) else None
@@ -357,7 +355,6 @@ def run_mt_bench_benchmark(cfg: RunConfig):
             f"--model_B is required for dataset '{cfg.task}'; "
             "no dataset-native baseline registered."
         )
-
     result_name = (
         f"{cfg.task}-{cfg.model.name}-{cfg.model.baseline}-{cfg.judge.model}-"
         f"{cfg.judge.swap_mode}"
@@ -378,6 +375,7 @@ def run_mt_bench_benchmark(cfg: RunConfig):
     )
     completions_a, completions_b = _generate_mt_bench_completions(
         cfg=cfg,
+        protocol=protocol,
         questions_df=questions_df,
     )
     resolved_prompt = resolve_run_judge_prompt(cfg.task, cfg.judge, multi_turn=True)
@@ -398,18 +396,19 @@ def run_mt_bench_benchmark(cfg: RunConfig):
     if resolved_prompt.delegated:
         return _run_mt_bench_fastchat(
             cfg=cfg,
+            protocol=protocol,
             res_folder=res_folder,
             result_name=result_name,
             questions_df=questions_df,
             completions_a=completions_a,
             completions_b=completions_b,
             judge_chat_model=judge_chat_model,
             resolved_prompt=resolved_prompt,
-            fastchat_prompt_preset=protocol.judge.fastchat_prompt_preset,
             started_at_utc=run_started_at,
         )
     return _run_mt_bench_preset(
         cfg=cfg,
+        protocol=protocol,
         res_folder=res_folder,
         result_name=result_name,
         questions_df=questions_df,

judgearena/benchmarks/pairwise/baselines.py

@@ -3,9 +3,103 @@
 from __future__ import annotations
 
 from collections.abc import Mapping
+from dataclasses import dataclass
+
+import pandas as pd
 
 from judgearena.tasks.registry import get_packaged_task
-from judgearena.tasks.schema import CategoryDefaultsBaseline, TaskDefaultBaseline
+from judgearena.tasks.schema import (
+    CategoryDefaultsBaseline,
+    ResolvedTaskSpec,
+    TaskDefaultBaseline,
+)
+
+
+@dataclass(frozen=True)
+class BaselinePlan:
+    """Row-aligned baseline assignment for model B."""
+
+    baseline_by_index: pd.Series
+
+    @classmethod
+    def flat(cls, model: str, *, index: pd.Index) -> BaselinePlan:
+        return cls(
+            baseline_by_index=pd.Series(model, index=index, name="model_B", dtype=str)
+        )
+
+    @classmethod
+    def per_row(cls, series: pd.Series) -> BaselinePlan:
+        return cls(baseline_by_index=series.astype(str).rename("model_B"))
+
+    @property
+    def unique_models(self) -> list[str]:
+        return sorted(self.baseline_by_index.dropna().unique().tolist())
+
+    @property
+    def is_single_model(self) -> bool:
+        return len(self.unique_models) == 1
+
+    @property
+    def single_model(self) -> str:
+        if not self.is_single_model:
+            raise ValueError(
+                "BaselinePlan is per-row; use baseline_by_index for row-level lookups."
+            )
+        return self.unique_models[0]
+
+    @property
+    def display_name(self) -> str:
+        return (
+            self.single_model if self.is_single_model else "+".join(self.unique_models)
+        )
+
+    def aligned_to(self, index: pd.Index) -> pd.Series:
+        return self.baseline_by_index.loc[index]
+
+
+def resolve_baseline_plan(
+    *,
+    task_id: str,
+    task: ResolvedTaskSpec | None,
+    runtime_baseline: str | None,
+    instructions: pd.DataFrame,
+) -> BaselinePlan:
+    """Resolve a runtime override or the baseline declared by a task."""
+    if runtime_baseline is not None:
+        return BaselinePlan.flat(runtime_baseline, index=instructions.index)
+
+    if task is None:
+        raise ValueError(
+            f"model.baseline is required for task {task_id!r}; no task baseline "
+            "is registered."
+        )
+
+    baseline = task.spec.protocol.baseline
+    if isinstance(baseline, TaskDefaultBaseline):
+        return BaselinePlan.flat(baseline.reference_id, index=instructions.index)
+
+    if isinstance(baseline, CategoryDefaultsBaseline):
+        if "category" not in instructions.columns:
+            raise ValueError(
+                f"{task_id} requires a 'category' column for per-category "
+                "baseline routing; re-run dataset download to regenerate the "
+                "instructions table."
+            )
+        per_row = instructions["category"].map(baseline.references)
+        if per_row.isna().any():
+            unknown = sorted(
+                instructions.loc[per_row.isna(), "category"].unique().tolist()
+            )
+            raise ValueError(
+                f"Unknown baseline categories for {task_id}: {unknown}. "
+                f"Known: {sorted(baseline.references)}"
+            )
+        return BaselinePlan.per_row(per_row)
+
+    raise ValueError(
+        f"model.baseline is required for task {task_id!r}; its "
+        f"{baseline.strategy!r} baseline policy does not provide a model."
+    )
 
 
 def native_pairwise_baseline(task: str) -> str | Mapping[str, str] | None: