diff --git a/judgearena/benchmarks/registry.py b/judgearena/benchmarks/registry.py index 7bab0b06..abcce555 100644 --- a/judgearena/benchmarks/registry.py +++ b/judgearena/benchmarks/registry.py @@ -42,6 +42,12 @@ class ResolvedBenchmark: task: ResolvedTaskSpec | None +def _run_tune_judge(cfg, task): + from judgearena.tuning.runner import run_tune_judge + + return run_tune_judge(cfg, task) + + def benchmark_adapters() -> tuple[BenchmarkAdapter, ...]: """Return registered benchmark implementations, specific first.""" from judgearena.benchmarks.elo.runner import run_elo @@ -55,6 +61,7 @@ def benchmark_adapters() -> tuple[BenchmarkAdapter, ...]: BenchmarkAdapter("meta_eval", frozenset(), run_meta_eval), BenchmarkAdapter("mt_bench", frozenset(), run_mt_bench_benchmark), BenchmarkAdapter("mt_bench_101", frozenset(), run_mt_bench_101_benchmark), + BenchmarkAdapter("tune_judge", frozenset(), _run_tune_judge), BenchmarkAdapter("pairwise", None, run_pairwise), ) diff --git a/judgearena/config.py b/judgearena/config.py index f0ac302e..524d0683 100644 --- a/judgearena/config.py +++ b/judgearena/config.py @@ -5,7 +5,7 @@ import argparse import os from pathlib import Path -from typing import Literal +from typing import Any, Literal import yaml from pydantic import BaseModel, ConfigDict, Field, model_validator @@ -26,6 +26,7 @@ EloScoringSpec, MetaEvalProtocol, MTBenchProtocol, + TuneJudgeProtocol, ) # Set by build_run_config() for the duration of RunConfig() construction. @@ -411,6 +412,118 @@ class MetaEvalArgs(BaseModel): """Share of arena prompts assigned to the validation partition.""" +class TuneJudgeArgs(BaseModel): + """NePS search settings and judge-evaluation objectives.""" + + model_config = ConfigDict(use_attribute_docstrings=True, extra="forbid") + + meta_eval_task: str | None = None + """Packaged meta-evaluation task used by every trial.""" + + search_space: dict[str, Any] + """Nested judge parameters, written as lists or bounded ranges.""" + + fidelity: dict[str, dict[str, int]] = Field(default_factory=dict) + optimizer: dict[str, Any] = Field(default_factory=dict) + neps: dict[str, Any] = Field(default_factory=dict) + """Native neps.run arguments; search space and optimizer are separate.""" + + objectives: list[Literal["agreement", "cost"]] | None = None + """Agreement is maximized and expected token cost minimized.""" + + price_per_million_tokens: dict[str, float | dict[str, float]] = Field( + default_factory=dict + ) + """Optional USD per million tokens; scalar or separate input/output rates.""" + + def resolve(self, protocol: TuneJudgeProtocol) -> TuneJudgeArgs: + values = self.model_dump() + values["meta_eval_task"] = ( + self.meta_eval_task or protocol.default_meta_eval_task + ) + values["objectives"] = ( + self.objectives + if self.objectives is not None + else protocol.default_objectives + ) + values["fidelity"] = { + name: {**bounds, **self.fidelity.get(name, {})} + for name, bounds in protocol.default_fidelity.items() + } + if set(self.fidelity) - {"battles_per_model"}: + raise ValueError("Only battles_per_model fidelity is supported") + lower, upper = ( + values["fidelity"]["battles_per_model"][k] for k in ("lower", "upper") + ) + if not 0 < lower < upper: + raise ValueError("Fidelity requires 0 < lower < upper") + defaults = protocol.default_optimizer + values["optimizer"] = ( + {**defaults, **self.optimizer} + if self.optimizer.get("name", defaults["name"]) == defaults["name"] + else self.optimizer + ) + values["neps"] = {"ignore_errors": True, **self.neps} + return TuneJudgeArgs.model_validate(values) + + test_battles_per_model: int | None = Field(default=None, gt=0) + """Test budget; defaults to the highest fidelity.""" + + run_dir: Path | None = None + """Explicit directory to resume or share; otherwise create a fresh run.""" + + search_only: bool = False + """Join an initialized run as a validation-only worker.""" + + @model_validator(mode="after") + def _validate_runtime(self) -> TuneJudgeArgs: + reserved = { + "pipeline_space", + "optimizer", + "evaluate_pipeline", + "root_directory", + "overwrite_root_directory", + "continue_until_max_evaluation_completed", + "total_cost_to_spend", + "worker_cost_to_spend", + } + if forbidden := reserved.intersection(self.neps): + raise ValueError( + f"Unsupported or managed NePS options: {sorted(forbidden)}" + ) + if self.search_only and self.run_dir is None: + raise ValueError("search_only requires an explicit tune_judge.run_dir") + if not self.search_only: + if not any( + self.neps.get(k) + for k in ("total_evaluations_to_spend", "total_fidelities_to_spend") + ): + raise ValueError("The primary worker requires a finite global budget") + if any( + self.neps.get(k) is not None + for k in ("worker_evaluations_to_spend", "worker_fidelities_to_spend") + ): + raise ValueError("Worker-local budgets require search_only") + name = self.optimizer.get("name") + if self.objectives is not None: + if ( + name in {"neps_priorband", "neps_hyperband"} + and len(self.objectives) != 1 + ): + raise ValueError(f"{name} requires a single objective") + if name in {"mo_hyperband", "primo"} and len(self.objectives) != 2: + raise ValueError(f"{name} requires agreement and cost objectives") + if len(set(self.objectives)) != len(self.objectives): + raise ValueError("objectives must not contain duplicates") + for price in self.price_per_million_tokens.values(): + if isinstance(price, dict) and set(price) != {"input", "output"}: + raise ValueError("Token price overrides require input and output rates") + rates = price.values() if isinstance(price, dict) else [price] + if any(rate < 0 for rate in rates): + raise ValueError("Judge token prices must be nonnegative") + return self + + class RunArgs(BaseModel): """Run-level settings: seed, output location, caching, and logging.""" @@ -467,6 +580,9 @@ class RunConfig(BaseSettings): meta_eval: MetaEvalArgs | None = None """Runtime settings used only by tasks with a meta-evaluation protocol.""" + tune_judge: TuneJudgeArgs | None = None + """Tune the judge on a meta-evaluation task instead of scoring it once.""" + run: RunArgs = Field(default_factory=RunArgs) """Run-level settings (seed, output, caching, logging).""" @@ -485,11 +601,22 @@ def _validate(self) -> RunConfig: resolved_task = get_packaged_task(self.task) if resolved_task is None: raise ValueError( - f"Unknown task {self.task!r}; use 'judgearena tasks list' to " - "inspect packaged tasks." + f"Unknown task {self.task!r}; use 'judgearena tasks list'." ) - protocol = resolved_task.spec.protocol + is_tune_judge = isinstance(protocol, TuneJudgeProtocol) + if (self.tune_judge is not None) != is_tune_judge: + raise ValueError( + "tune_judge config is required by, and only valid for, tune-judge." + ) + if is_tune_judge: + self.tune_judge = self.tune_judge.resolve(protocol) + target = get_packaged_task(self.tune_judge.meta_eval_task) + if target is None or not isinstance(target.spec.protocol, MetaEvalProtocol): + raise ValueError( + "tune_judge.meta_eval_task must select a meta-evaluation task." + ) + protocol = target.spec.protocol task_generation = getattr(protocol, "generation", None) if ( "truncate_all_input_chars" not in self.generation.model_fields_set diff --git a/judgearena/tasks/definitions/tuning/tune-judge.yaml b/judgearena/tasks/definitions/tuning/tune-judge.yaml new file mode 100644 index 00000000..aaa86c11 --- /dev/null +++ b/judgearena/tasks/definitions/tuning/tune-judge.yaml @@ -0,0 +1,15 @@ +schema_version: 1 +task: tune-judge +task_version: 1 +description: Tune a judge against human preferences with agreement and cost objectives. +tags: [tuning, judge] +protocol: + runner: tune_judge + default_meta_eval_task: meta-eval-lmarena-140k-en + default_objectives: [agreement, cost] + default_fidelity: + battles_per_model: {lower: 10, upper: 90} + default_optimizer: {name: mo_hyperband, eta: 3, mo_selector: epsnet} +metadata: + reference_implementation: https://github.com/geoalgo/judgetuning + paper: https://arxiv.org/abs/2501.17178 diff --git a/judgearena/tasks/registry.py b/judgearena/tasks/registry.py index 11e8a783..95c60aad 100644 --- a/judgearena/tasks/registry.py +++ b/judgearena/tasks/registry.py @@ -28,6 +28,7 @@ TaskProvenance, TaskSelection, TaskSpec, + TuneJudgeProtocol, ) logger = get_logger(__name__) @@ -317,6 +318,10 @@ def _discover_tasks( def _validate_adapter_ids(resolved: ResolvedTaskSpec, adapters: AdapterCatalog) -> None: spec = resolved.spec + if isinstance(spec.protocol, TuneJudgeProtocol): + if spec.protocol.runner not in adapters.runners: + raise TaskDefinitionError(f"Unknown runner {spec.protocol.runner!r}") + return is_battle_backed = isinstance(spec.protocol, (EloProtocol, MetaEvalProtocol)) dataset_names = ( adapters.battle_datasets if is_battle_backed else adapters.instruction_datasets diff --git a/judgearena/tasks/schema/__init__.py b/judgearena/tasks/schema/__init__.py index 2c82226f..b219511c 100644 --- a/judgearena/tasks/schema/__init__.py +++ b/judgearena/tasks/schema/__init__.py @@ -51,6 +51,7 @@ TaskMetadata, TaskSpec, ) +from judgearena.tasks.schema.tuning import TuneJudgeProtocol __all__ = [ "BaselineSpec", @@ -88,4 +89,5 @@ "TaskProvenance", "TaskSelection", "TaskSpec", + "TuneJudgeProtocol", ] diff --git a/judgearena/tasks/schema/task.py b/judgearena/tasks/schema/task.py index 6f5ab80e..d5a7a449 100644 --- a/judgearena/tasks/schema/task.py +++ b/judgearena/tasks/schema/task.py @@ -17,13 +17,15 @@ from judgearena.tasks.schema.mt_bench import MTBenchProtocol from judgearena.tasks.schema.mt_bench_101 import MTBench101Protocol from judgearena.tasks.schema.pairwise import PairwiseProtocol +from judgearena.tasks.schema.tuning import TuneJudgeProtocol ProtocolSpec = Annotated[ PairwiseProtocol | MTBenchProtocol | MTBench101Protocol | EloProtocol - | MetaEvalProtocol, + | MetaEvalProtocol + | TuneJudgeProtocol, Field(discriminator="runner"), ] @@ -73,7 +75,7 @@ class TaskSpec(StrictFrozenModel): task_version: int = Field(ge=1) description: str = Field(min_length=1) tags: tuple[str, ...] = () - dataset: DatasetSpec + dataset: DatasetSpec | None = None protocol: ProtocolSpec variants: SuffixVariants | None = None metadata: TaskMetadata = Field(default_factory=TaskMetadata) @@ -82,6 +84,14 @@ class TaskSpec(StrictFrozenModel): def _validate_task(self) -> TaskSpec: if len(set(self.tags)) != len(self.tags): raise ValueError("tags must not contain duplicates") + if isinstance(self.protocol, TuneJudgeProtocol): + if self.dataset is not None: + raise ValueError( + "Tuning tasks use their target meta-eval benchmark's dataset." + ) + return self + if self.dataset is None: + raise ValueError("Benchmark tasks require a dataset.") source_names = set(self.dataset.sources) baseline = self.protocol.baseline if ( diff --git a/judgearena/tasks/schema/tuning.py b/judgearena/tasks/schema/tuning.py new file mode 100644 index 00000000..7f0604e7 --- /dev/null +++ b/judgearena/tasks/schema/tuning.py @@ -0,0 +1,15 @@ +"""Task-owned defaults for judge configuration search.""" + +from typing import Any, Literal + +from pydantic import Field + +from judgearena.tasks.schema.base import StrictFrozenModel + + +class TuneJudgeProtocol(StrictFrozenModel): + runner: Literal["tune_judge"] + default_meta_eval_task: str + default_objectives: list[Literal["agreement", "cost"]] + default_fidelity: dict[str, dict[str, int]] + default_optimizer: dict[str, Any] = Field(min_length=1) diff --git a/judgearena/tuning/__init__.py b/judgearena/tuning/__init__.py new file mode 100644 index 00000000..a2fa5b7b --- /dev/null +++ b/judgearena/tuning/__init__.py @@ -0,0 +1 @@ +"""Judge configuration tuning against human-labeled arena battles.""" diff --git a/judgearena/tuning/search_space.py b/judgearena/tuning/search_space.py new file mode 100644 index 00000000..71d2620f --- /dev/null +++ b/judgearena/tuning/search_space.py @@ -0,0 +1,140 @@ +"""Construct NePS domains using its public constructor arguments.""" + +from __future__ import annotations + +import copy +from collections.abc import Iterator, Mapping +from typing import TYPE_CHECKING, Any, get_args + +from judgearena.config import RunConfig + +if TYPE_CHECKING: + import neps + +FIDELITY = "meta_eval.battles_per_model" +DOMAIN_TYPES = frozenset({"Categorical", "Float", "Integer", "IntegerFidelity"}) + + +def _range_type(name: str) -> str: + model = RunConfig + for part in name.split("."): + field = model.model_fields[part] + model = next( + t + for t in (get_args(field.annotation) or (field.annotation,)) + if t is not type(None) + ) + if model not in (int, float): + raise ValueError(f"Range requires a numeric config field: {name}") + return "Integer" if model is int else "Float" + + +def parameter_specs(search_space: Mapping, base: Mapping, fidelity: Mapping) -> dict: + """Flatten nested parameter definitions and resolve value-based priors.""" + specs = {} + + def visit(section, prefix=""): + for key, value in section.items(): + name = f"{prefix}.{key}" if prefix else key + if isinstance(value, Mapping) and not ( + {"choices", "lower", "upper"} & value.keys() + ): + visit(value, name) + continue + if not ( + name.startswith("judge.") + or name == "generation.truncate_judge_input_chars" + ): + raise ValueError( + f"Search space must not override tuner-owned keys: {name}" + ) + spec = {"choices": value} if isinstance(value, list) else dict(value) + node = base + for part in name.split("."): + node = node.get(part) if isinstance(node, Mapping) else None + explicit_prior = "prior" in spec + prior = spec.pop("prior", node) + confidence = spec.pop("prior_confidence", "medium") + if "choices" in spec: + kind = "Categorical" + if explicit_prior and prior not in spec["choices"]: + raise ValueError(f"Prior for {name} is not in its choices") + prior = ( + spec["choices"].index(prior) if prior in spec["choices"] else None + ) + else: + kind = _range_type(name) + if prior is not None and not spec["lower"] <= prior <= spec["upper"]: + if explicit_prior: + raise ValueError(f"Prior for {name} is outside its range") + prior = None + if prior is not None: + spec.update(prior=prior, prior_confidence=confidence) + specs[name] = {"type": kind, **spec} + + visit(search_space) + specs[FIDELITY] = {"type": "IntegerFidelity", **fidelity["battles_per_model"]} + return specs + + +def build_neps_space(specs: Mapping[str, dict]) -> neps.PipelineSpace: + """Translate constructor tags to domains without changing their arguments.""" + import neps + + parameters = {} + for name, spec in specs.items(): + kind = spec["type"] + if kind not in DOMAIN_TYPES: + raise ValueError(f"Unsupported NePS domain: {kind}") + if name == FIDELITY: + if kind != "IntegerFidelity": + raise ValueError(f"{FIDELITY} must use IntegerFidelity") + elif not ( + name.startswith("judge.") or name == "generation.truncate_judge_input_chars" + ): + raise ValueError(f"Search space must not override tuner-owned keys: {name}") + elif kind == "IntegerFidelity": + raise ValueError(f"Only {FIDELITY} may be a fidelity") + kwargs = {key: value for key, value in spec.items() if key != "type"} + if kind == "Categorical": + kwargs["choices"] = tuple(kwargs["choices"]) + parameters[name] = getattr(neps, kind)(**kwargs) + if FIDELITY not in parameters: + raise ValueError(f"Search space requires {FIDELITY}") + return type("JudgeSpace", (neps.PipelineSpace,), parameters)() + + +def axis_overrides(specs: Mapping[str, dict]) -> Iterator[dict[str, object]]: + """Yield candidate values for preflight RunConfig validation.""" + for name, spec in specs.items(): + if name == FIDELITY: + continue + values = ( + spec["choices"] + if spec["type"] == "Categorical" + else [spec["lower"], spec["upper"]] + ) + for value in values: + yield {name: value} + + +def decode_config(config: Mapping[str, Any]) -> tuple[dict[str, object], int]: + """Split the sampled judge settings from the evaluation fidelity.""" + overrides = dict(config) + return overrides, int(overrides.pop(FIDELITY)) + + +def apply_overrides( + base: Mapping[str, object], overrides: Mapping[str, object] +) -> dict[str, object]: + """Return a copy of base with dotted-path overrides assigned.""" + values = copy.deepcopy(dict(base)) + for path, value in overrides.items(): + *parents, leaf = path.split(".") + node = values + for key in parents: + if node.get(key) is None: + node[key] = {} + node = node[key] + node[leaf] = value + return values diff --git a/judgearena/tuning/session.py b/judgearena/tuning/session.py new file mode 100644 index 00000000..931d388f --- /dev/null +++ b/judgearena/tuning/session.py @@ -0,0 +1,93 @@ +"""Explicit tuning sessions and collection of shared NePS state.""" + +from __future__ import annotations + +import hashlib +import json +import time +from datetime import UTC, datetime +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path + +import pandas as pd + +from judgearena.artifacts import prepare_run_directory, safe_filename +from judgearena.config import RunConfig + + +def _config_hash(cfg: RunConfig) -> str: + payload = cfg.model_dump(mode="json") + payload["run"] = {"seed": cfg.run.seed} + tuning = payload["tune_judge"] + tuning.pop("run_dir") + tuning.pop("search_only") + tuning["neps"] = { + key: value + for key, value in tuning["neps"].items() + if not key.endswith("_to_spend") and key != "worker_id" + } + if cfg.judge.prompt is not None: + payload["prompt_contents"] = [ + path.read_text() + for path in (cfg.judge.prompt.system_file, cfg.judge.prompt.user_file) + ] + return hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest() + + +def prepare_session(cfg: RunConfig) -> Path: + """Create a new run or verify an explicitly selected existing run.""" + tuning = cfg.tune_judge + tune_dir = tuning.run_dir or ( + Path(cfg.run.result_folder) + / f"{safe_filename(cfg.task)}-{datetime.now(UTC):%Y%m%d_%H%M%S_%f}" + ) + tune_dir = tune_dir.resolve() + metadata_path = tune_dir / "tuning-metadata.json" + fingerprint = _config_hash(cfg) + if metadata_path.exists(): + metadata = json.loads(metadata_path.read_text()) + if metadata["config_hash"] != fingerprint: + raise ValueError(f"Incompatible configuration for tuning run {tune_dir}") + if not tuning.search_only: + prepare_run_directory(cfg, tune_dir) + elif tuning.search_only: + raise ValueError("Start the primary process before joining its run_dir") + else: + prepare_run_directory(cfg, tune_dir) + versions = {"neural-pipeline-search": version("neural-pipeline-search")} + try: + versions["vllm"] = version("vllm") + except PackageNotFoundError: + pass + metadata = {"config_hash": fingerprint, "versions": versions} + temporary = metadata_path.with_suffix(".tmp") + temporary.write_text(json.dumps(metadata, indent=2)) + temporary.replace(metadata_path) + if tuning.search_only: + if not (tune_dir / "neps" / "pipeline_space.pkl").exists(): + raise ValueError("The primary process has not initialized NePS yet") + return tune_dir + + +def collect_trials(root: Path, *, wait: bool) -> pd.DataFrame: + """Wait for active evaluations and read all workers' trial records.""" + from neps.state import NePSState + + state = NePSState.create_or_load(root, load_only=True) + while True: + trials = state.lock_and_read_trials() + active = [t for t in trials.values() if t.metadata.state == "evaluating"] + if not wait or not active: + break + time.sleep(1) + records = [] + for trial in trials.values(): + if trial.metadata.state in {"success", "failed", "crashed"}: + records.append( + { + "neps_trial_id": trial.id, + "status": "failed", + **trial.report.extra, + } + ) + return pd.DataFrame(records) diff --git a/judgearena/utils/io.py b/judgearena/utils/io.py index aeab573e..6dba1198 100644 --- a/judgearena/utils/io.py +++ b/judgearena/utils/io.py @@ -29,6 +29,8 @@ def download_hf(name: str, local_path: Path): resolved_task = get_packaged_task(name) if resolved_task is None: raise ValueError(f"Unknown task {name!r}.") + if resolved_task.spec.dataset is None: + return from judgearena.datasets.registry import resolve_download_adapter resolve_download_adapter(resolved_task.spec.dataset.adapter).download( @@ -67,8 +69,9 @@ def download_all(): logger.info("Downloading all datasets in %s", data_root) local_path_tables = data_root / "tables" - for task_id in load_tasks(): - download_hf(name=task_id, local_path=local_path_tables) + for task_id, task in load_tasks().items(): + if task.spec.dataset is not None: + download_hf(name=task_id, local_path=local_path_tables) if __name__ == "__main__": diff --git a/pyproject.toml b/pyproject.toml index 2e214bdc..81228bf8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -114,3 +114,6 @@ vllm = [ "openai>=2.30.0", # vLLM 0.25 imports responses.NamespaceTool. ] llamacpp = ["llama-cpp-python>=0.3.0"] +tune = [ + "neural-pipeline-search==0.17.0", +] diff --git a/tests/test_tuning_config.py b/tests/test_tuning_config.py new file mode 100644 index 00000000..fbd4f443 --- /dev/null +++ b/tests/test_tuning_config.py @@ -0,0 +1,54 @@ +import pytest + +from judgearena.config import RunConfig +from judgearena.tuning.search_space import FIDELITY, build_neps_space, parameter_specs + + +def test_nested_space_and_priors(): + neps = pytest.importorskip("neps") + domains = { + "judge": { + "temperature": {"lower": 0.0, "upper": 1.0}, + "model": {"choices": ["Dummy/a", "Dummy/b"], "prior": "Dummy/b"}, + } + } + specs = parameter_specs( + domains, + {"judge": {"temperature": 0.2}}, + {"battles_per_model": {"lower": 1, "upper": 3}}, + ) + attrs = build_neps_space(specs).get_attrs() + assert isinstance(attrs["judge.temperature"], neps.Float) + assert attrs["judge.temperature"].prior == 0.2 + assert attrs["judge.model"].prior == 1 + assert isinstance(attrs[FIDELITY], neps.IntegerFidelity) + with pytest.raises(ValueError, match="tuner-owned"): + parameter_specs({"meta_eval": {"split": ["test"]}}, {}, {}) + + +def test_packaged_defaults_and_target_restrictions(): + settings = { + "search_space": {"judge": {"temperature": [0.0, 1.0]}}, + "optimizer": {"eta": 2}, + "neps": {"total_evaluations_to_spend": 6}, + } + cfg = RunConfig( + task="tune-judge", judge={"model": "Dummy/judge"}, tune_judge=settings + ) + direct = RunConfig( + task=cfg.tune_judge.meta_eval_task, judge={"model": "Dummy/judge"} + ) + assert (cfg.judge, cfg.meta_eval) == (direct.judge, direct.meta_eval) + assert cfg.tune_judge.optimizer == { + "name": "mo_hyperband", + "eta": 2, + "mo_selector": "epsnet", + } + assert cfg.tune_judge.objectives == ["agreement", "cost"] + assert cfg.tune_judge.fidelity == {"battles_per_model": {"lower": 10, "upper": 90}} + with pytest.raises(ValueError): + RunConfig( + task="tune-judge", + judge={"model": "Dummy/judge"}, + tune_judge={**settings, "meta_eval_task": "mt-bench"}, + ) diff --git a/tests/test_utils.py b/tests/test_utils.py index 0c0b10b8..14faae80 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -41,7 +41,9 @@ def test_download_all_dispatches_registered_tasks(monkeypatch, tmp_path): ) utils_io.download_all() - assert [name for _, name, _ in calls] == list(load_tasks()) + assert [name for _, name, _ in calls] == [ + name for name, task in load_tasks().items() if task.spec.dataset is not None + ] assert {path for _, _, path in calls} == {tmp_path / "tables"}