(4/n) Package Arena-Hard Task Variants

judgearena/benchmarks/pairwise/baselines.py

@@ -4,7 +4,6 @@
 
 from collections.abc import Mapping
 
-from judgearena.datasets.arena_hard import ARENA_HARD_BASELINES
 from judgearena.datasets.m_arenahard import (
     M_ARENA_HARD_BASELINES,
     split_m_arena_hard_dataset,
@@ -14,7 +13,6 @@
 from judgearena.tasks.schema import CategoryDefaultsBaseline, TaskDefaultBaseline
 
 LEGACY_PAIRWISE_BASELINES: dict[str, str | Mapping[str, str]] = {
-    **ARENA_HARD_BASELINES,
     **M_ARENA_HARD_BASELINES,
     **MT_BENCH_BASELINES,
 }

judgearena/benchmarks/pairwise/runner.py

@@ -16,10 +16,6 @@
 from judgearena.benchmarks.mt_bench.mt_bench_utils import run_mt_bench
 from judgearena.benchmarks.pairwise.baselines import native_pairwise_baseline
 from judgearena.datasets import load_instructions
-from judgearena.datasets.arena_hard import (
-    download_arena_hard,
-    is_arena_hard_dataset,
-)
 from judgearena.datasets.fluency import is_fluency_task as task_is_fluency
 from judgearena.datasets.fluency import load_fluency_contexts
 from judgearena.evaluate import judge_and_parse_prefs, resolve_run_judge_prompt
@@ -58,17 +54,12 @@ def try_load_dataset_completions(
     local_path_tables = data_root / "tables"
     resolved_task = get_packaged_task(dataset)
     if resolved_task is not None:
-        from judgearena.datasets.judgearena_tables import load_task_model_outputs
+        from judgearena.datasets.registry import resolve_dataset_adapter
 
-        df_outputs = load_task_model_outputs(resolved_task, local_path_tables)
+        adapter = resolve_dataset_adapter(resolved_task.spec.dataset.adapter)
+        df_outputs = adapter.load_model_outputs(resolved_task, local_path_tables)
         if df_outputs is None:
             return None
-    elif is_arena_hard_dataset(dataset):
-        download_arena_hard(dataset=dataset, local_tables_path=local_path_tables)
-        output_path = local_path_tables / "model_outputs" / f"{dataset}.csv.zip"
-        if not output_path.exists():
-            return None
-        df_outputs = read_df(output_path)
     else:
         download_hf(name=dataset, local_path=local_path_tables)
         output_path = local_path_tables / "model_outputs" / f"{dataset}.csv.zip"

judgearena/datasets/__init__.py

@@ -1,9 +1,5 @@
 import pandas as pd
 
-from judgearena.datasets.arena_hard import (
-    download_arena_hard,
-    is_arena_hard_dataset,
-)
 from judgearena.datasets.m_arenahard import (
     load_m_arenahard,
     split_m_arena_hard_dataset,
@@ -18,9 +14,10 @@ def load_instructions(dataset: str, n_instructions: int | None = None) -> pd.Dat
     resolved_task = get_packaged_task(dataset)
     if resolved_task is not None:
         from judgearena import utils as judgearena_utils
-        from judgearena.datasets.judgearena_tables import load_task_instructions
+        from judgearena.datasets.registry import resolve_dataset_adapter
 
-        df_instructions = load_task_instructions(
+        adapter = resolve_dataset_adapter(resolved_task.spec.dataset.adapter)
+        df_instructions = adapter.load_instructions(
             resolved_task, judgearena_utils.data_root / "tables"
         )
 
@@ -55,20 +52,7 @@ def load_instructions(dataset: str, n_instructions: int | None = None) -> pd.Dat
         )
 
     else:
-        assert dataset in [
-            "arena-hard-v0.1",
-            "arena-hard-v2.0",
-        ]
-        from judgearena import utils as judgearena_utils
-
-        local_path_tables = judgearena_utils.data_root / "tables"
-        if is_arena_hard_dataset(dataset):
-            download_arena_hard(dataset=dataset, local_tables_path=local_path_tables)
-        else:
-            judgearena_utils.download_hf(name=dataset, local_path=local_path_tables)
-        df_instructions = judgearena_utils.read_df(
-            local_path_tables / "instructions" / f"{dataset}.csv"
-        )
+        raise ValueError(f"Unsupported instruction dataset {dataset!r}.")
 
     df_instructions = df_instructions.set_index("instruction_index").sort_index()
     logger.info("Loaded %d instructions for %s.", len(df_instructions), dataset)

judgearena/datasets/arena_hard.py

@@ -5,91 +5,70 @@
 import pandas as pd
 from huggingface_hub import snapshot_download
 
-from judgearena.dataset_revisions import hf_revision
-
-ARENA_HARD_HF_REPO_ID = "lmarena-ai/arena-hard-auto"
-
-# Mirrors upstream's `JUDGE_SETTINGS` baseline assignment in
-# `arena-hard-auto/utils/judge_utils.py` verbatim: v0.1 has a single flat
-# baseline, v2.0 routes per question category. `is_arena_hard_dataset` and
-# the dispatcher in `generate_and_evaluate.py` key off this map.
-#
-# Note: the released v2.0 `question.jsonl` only tags rows as `hard_prompt`
-# (500) or `creative_writing` (250); `coding` and `math` are inert keys
-# upstream ships for forward compatibility (no question carries those
-# labels, so the dispatcher never looks them up). We keep them so any
-# future re-tagging upstream lights up automatically without a code
-# change here.
-ARENA_HARD_BASELINES: dict[str, str | Mapping[str, str]] = {
-    "arena-hard-v0.1": "gpt-4-0314",
-    "arena-hard-v2.0": {
-        "hard_prompt": "o3-mini-2025-01-31",
-        "coding": "o3-mini-2025-01-31",
-        "math": "o3-mini-2025-01-31",
-        "creative_writing": "gemini-2.0-flash-001",
-    },
-}
-
-# Dataset name -> upstream HF `data/<variant>/` directory. Kept private so the
-# public API of this module is just the baseline map and helpers below.
-_ARENA_HARD_HF_VARIANTS: dict[str, str] = {
-    "arena-hard-v0.1": "arena-hard-v0.1",
-    "arena-hard-v2.0": "arena-hard-v2.0",
-}
+from judgearena.tasks.registry import get_packaged_task
+from judgearena.tasks.schema import HuggingFaceDatasetSource, ResolvedTaskSpec
 
 
 def is_arena_hard_dataset(dataset: str) -> bool:
-    return dataset in ARENA_HARD_BASELINES
+    task = get_packaged_task(dataset)
+    return task is not None and task.spec.dataset.adapter == "arena_hard"
 
 
 def arena_hard_native_baseline(
     dataset: str,
 ) -> str | Mapping[str, str] | None:
-    """Dataset-native baseline assignment.
+    """Return the YAML-defined baseline for an Arena-Hard task."""
+    if not is_arena_hard_dataset(dataset):
+        return None
+    from judgearena.benchmarks.pairwise.baselines import native_pairwise_baseline
 
-    Returns a plain string for flat datasets (v0.1), a `{category: model}`
-    mapping for per-category datasets (v2.0), or `None` for datasets that
-    don't ship a native baseline.
-    """
-    return ARENA_HARD_BASELINES.get(dataset)
+    return native_pairwise_baseline(dataset)
 
 
 def normalize_official_arena_hard(
     raw_df: pd.DataFrame, dataset: str
 ) -> tuple[pd.DataFrame, pd.DataFrame | None]:
-    if dataset not in _ARENA_HARD_HF_VARIANTS:
+    if not is_arena_hard_dataset(dataset):
         raise ValueError(f"Unsupported Arena-Hard dataset: {dataset}")
     df_instructions = _build_instructions(raw_df)
     df_model_outputs = _build_model_outputs(raw_df)
     return df_instructions, df_model_outputs
 
 
-def download_arena_hard(dataset: str, local_tables_path: Path) -> None:
-    """Populate `{dataset}.csv` and `{dataset}.csv.zip` on disk if missing.
+def _source(task: ResolvedTaskSpec) -> HuggingFaceDatasetSource:
+    source = task.spec.dataset.sources.get("examples")
+    if not isinstance(source, HuggingFaceDatasetSource) or source.config is None:
+        raise ValueError(
+            f"Task {task.task!r} must define an 'examples' Hugging Face source "
+            "with its Arena-Hard data variant in 'config'."
+        )
+    return source
+
+
+def download_task_sources(task: ResolvedTaskSpec, local_tables_path: Path) -> None:
+    """Populate canonical instruction and output tables for one task.
 
     Pulls the raw jsonl files directly via `snapshot_download` and reads them
     with pandas: upstream's per-row `messages[].content` oscillates between
     string and dict across answer files, so `datasets.load_dataset` can't
     materialize them into a single Arrow schema.
-
     """
-    if dataset not in _ARENA_HARD_HF_VARIANTS:
-        return
+    if task.spec.dataset.adapter != "arena_hard":
+        raise ValueError(f"Task {task.task!r} does not use the Arena-Hard adapter.")
+    dataset = task.task
     instructions_path = local_tables_path / "instructions" / f"{dataset}.csv"
     model_outputs_path = local_tables_path / "model_outputs" / f"{dataset}.csv.zip"
     if instructions_path.exists() and model_outputs_path.exists():
         return
 
-    variant = _ARENA_HARD_HF_VARIANTS[dataset]
+    source = _source(task)
+    variant = source.config
     snapshot_root = snapshot_download(
-        repo_id=ARENA_HARD_HF_REPO_ID,
+        repo_id=source.repo_id,
         repo_type="dataset",
-        allow_patterns=[
-            f"data/{variant}/question.jsonl",
-            f"data/{variant}/model_answer/*.jsonl",
-        ],
+        allow_patterns=list(source.allow_patterns) or None,
         force_download=False,
-        revision=hf_revision(ARENA_HARD_HF_REPO_ID),
+        revision=source.revision,
     )
     raw_df = _read_arena_hard_jsonl_frames(
         variant_dir=Path(snapshot_root) / "data" / variant
@@ -104,6 +83,30 @@ def download_arena_hard(dataset: str, local_tables_path: Path) -> None:
         df_model_outputs.to_csv(model_outputs_path, index=False)
 
 
+def load_task_instructions(
+    task: ResolvedTaskSpec, local_tables_path: Path
+) -> pd.DataFrame:
+    """Load normalized instructions for a registered Arena-Hard task."""
+    download_task_sources(task, local_tables_path)
+    return pd.read_csv(local_tables_path / "instructions" / f"{task.task}.csv")
+
+
+def load_task_model_outputs(
+    task: ResolvedTaskSpec, local_tables_path: Path
+) -> pd.DataFrame | None:
+    """Load normalized reference outputs for a registered Arena-Hard task."""
+    download_task_sources(task, local_tables_path)
+    path = local_tables_path / "model_outputs" / f"{task.task}.csv.zip"
+    return pd.read_csv(path) if path.exists() else None
+
+
+def download_arena_hard(dataset: str, local_tables_path: Path) -> None:
+    """Compatibility wrapper around the registered Arena-Hard adapter."""
+    task = get_packaged_task(dataset)
+    if task is not None and task.spec.dataset.adapter == "arena_hard":
+        download_task_sources(task, local_tables_path)
+
+
 def _read_arena_hard_jsonl_frames(variant_dir: Path) -> pd.DataFrame:
     frames: list[pd.DataFrame] = []
     question_path = variant_dir / "question.jsonl"

judgearena/datasets/registry.py

@@ -0,0 +1,50 @@
+"""Registry connecting task dataset-adapter IDs to implementations."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+
+import pandas as pd
+
+from judgearena.tasks.schema import ResolvedTaskSpec
+
+TaskDataFunction = Callable[[ResolvedTaskSpec, Path], pd.DataFrame | None]
+TaskDownloadFunction = Callable[[ResolvedTaskSpec, Path], None]
+
+
+@dataclass(frozen=True)
+class DatasetAdapter:
+    name: str
+    download: TaskDownloadFunction
+    load_instructions: TaskDataFunction
+    load_model_outputs: TaskDataFunction
+
+
+def dataset_adapters() -> tuple[DatasetAdapter, ...]:
+    """Return registered dataset implementations."""
+    from judgearena.datasets import arena_hard, judgearena_tables
+
+    return (
+        DatasetAdapter(
+            "judgearena_tables",
+            judgearena_tables.download_task_sources,
+            judgearena_tables.load_task_instructions,
+            judgearena_tables.load_task_model_outputs,
+        ),
+        DatasetAdapter(
+            "arena_hard",
+            arena_hard.download_task_sources,
+            arena_hard.load_task_instructions,
+            arena_hard.load_task_model_outputs,
+        ),
+    )
+
+
+def resolve_dataset_adapter(name: str) -> DatasetAdapter:
+    """Return the implementation registered under ``name``."""
+    for adapter in dataset_adapters():
+        if adapter.name == name:
+            return adapter
+    raise ValueError(f"Unknown task dataset adapter {name!r}.")

judgearena/prompts/registry.py

@@ -93,8 +93,6 @@ def metadata(self) -> dict[str, str | bool | None]:
 JUDGE_PROMPT_PRESETS = tuple(PRESETS)
 
 TASK_DEFAULT_PRESET: dict[str, str] = {
-    "arena-hard-v0.1": DEFAULT_JUDGE_PROMPT_PRESET,
-    "arena-hard-v2.0": DEFAULT_JUDGE_PROMPT_PRESET,
     "mt-bench": FASTCHAT_PAIRWISE_PROMPT_PRESET,
 }
 

judgearena/tasks/definitions/arena_hard/_base.yaml

@@ -0,0 +1,31 @@
+schema_version: 1
+task_version: 1
+tags: [pairwise, instruction-following, arena-hard]
+
+dataset:
+  adapter: arena_hard
+  sources:
+    examples:
+      type: huggingface_dataset
+      repo_id: lmarena-ai/arena-hard-auto
+      revision: "15f3746e21432264ce9b453999bde4f3c946d2e6"
+  fields:
+    id: instruction_index
+    instruction: instruction
+
+protocol:
+  runner: pairwise
+  generation:
+    mode: single_turn_chat
+  judge:
+    default_prompt: default
+    parser: pairwise_preference
+    default_swap_mode: fixed
+    allowed_swap_modes: [fixed, both]
+  scoring:
+    adapter: pairwise_win_rate
+    primary_metric: winrate
+    higher_is_better: true
+
+metadata:
+  reference_implementation: https://github.com/lmarena/arena-hard-auto

judgearena/tasks/definitions/arena_hard/arena-hard-v0.1.yaml

@@ -0,0 +1,17 @@
+extends: _base.yaml
+task: arena-hard-v0.1
+description: Pairwise evaluation on Arena-Hard v0.1.
+
+dataset:
+  sources:
+    examples:
+      config: arena-hard-v0.1
+      allow_patterns:
+        - data/arena-hard-v0.1/question.jsonl
+        - data/arena-hard-v0.1/model_answer/*.jsonl
+
+protocol:
+  baseline:
+    strategy: task_default
+    reference_id: gpt-4-0314
+    allow_runtime_override: true

judgearena/tasks/definitions/arena_hard/arena-hard-v2.0.yaml

@@ -0,0 +1,24 @@
+extends: _base.yaml
+task: arena-hard-v2.0
+description: Pairwise evaluation on Arena-Hard v2.0.
+
+dataset:
+  sources:
+    examples:
+      config: arena-hard-v2.0
+      allow_patterns:
+        - data/arena-hard-v2.0/question.jsonl
+        - data/arena-hard-v2.0/model_answer/*.jsonl
+  fields:
+    category: category
+
+protocol:
+  baseline:
+    strategy: category_defaults
+    category_field: category
+    references:
+      hard_prompt: o3-mini-2025-01-31
+      coding: o3-mini-2025-01-31
+      math: o3-mini-2025-01-31
+      creative_writing: gemini-2.0-flash-001
+    allow_runtime_override: true

judgearena/tasks/registry.py

@@ -17,7 +17,7 @@ class AdapterCatalog:
     """Component IDs that task YAML files may reference."""
 
     runners: frozenset[str] = frozenset({"pairwise"})
-    datasets: frozenset[str] = frozenset({"judgearena_tables"})
+    datasets: frozenset[str] = frozenset({"arena_hard", "judgearena_tables"})
     prompts: frozenset[str] = frozenset(JUDGE_PROMPT_PRESETS)
     parsers: frozenset[str] = frozenset({"pairwise_preference"})
     scorers: frozenset[str] = frozenset({"pairwise_win_rate"})