From f64faaf403743192c0c3a71f94f3b3de8bcefd28 Mon Sep 17 00:00:00 2001 From: Daniel van Strien Date: Thu, 13 Aug 2026 17:20:44 +0100 Subject: [PATCH] Honor source splits throughout benchmark runs --- src/ocr_bench/cli.py | 3 +++ src/ocr_bench/publish.py | 6 +++++ src/ocr_bench/run.py | 4 +++ src/ocr_bench/viewer.py | 20 ++++++++++----- src/ocr_bench/web.py | 8 +++++- tests/test_cli.py | 17 +++++++++++++ tests/test_publish.py | 25 ++++++++++++++++++ tests/test_run.py | 34 +++++++++++++++++++++++-- tests/test_viewer.py | 43 +++++++++++++++++++++++++++++++ tests/test_web.py | 55 ++++++++++++++++++++++++++++++++++++++++ 10 files changed, 206 insertions(+), 9 deletions(-) diff --git a/src/ocr_bench/cli.py b/src/ocr_bench/cli.py index 648da56..41de068 100644 --- a/src/ocr_bench/cli.py +++ b/src/ocr_bench/cli.py @@ -1418,6 +1418,7 @@ def _judge_batch(batch_comps: list[Comparison]) -> list[ComparisonResult]: if results_repo: metadata = EvalMetadata( source_dataset=args.dataset, + source_split=args.split, judge_models=[], seed=args.seed, max_samples=args.max_samples or len(ds), @@ -1545,6 +1546,7 @@ def _judge_batch(batch_comps: list[Comparison]) -> list[ComparisonResult]: if results_repo: metadata = EvalMetadata( source_dataset=args.dataset, + source_split=args.split, judge_models=[j.name for j in judges], seed=args.seed, max_samples=args.max_samples or len(ds), @@ -1659,6 +1661,7 @@ def cmd_run(args: argparse.Namespace) -> list[JobRun]: args.input_dataset, args.output_repo, slug, + split=args.split, max_samples=args.max_samples, shuffle=args.shuffle, seed=args.seed, diff --git a/src/ocr_bench/publish.py b/src/ocr_bench/publish.py index 45537db..52543cd 100644 --- a/src/ocr_bench/publish.py +++ b/src/ocr_bench/publish.py @@ -49,6 +49,10 @@ class EvalMetadata: max_samples: int total_comparisons: int valid_comparisons: int + # Source dataset split used to build the comparison grid. Historical + # metadata rows predate this field and are interpreted as ``train`` by the + # viewer for backward compatibility. + source_split: str = "train" auto_tied: int = 0 # Global comparison budget for the run (--max-comparisons); None = uncapped. # ``budget_exhausted`` records whether the run stopped because it hit the cap @@ -233,6 +237,7 @@ def build_metadata_row(metadata: EvalMetadata) -> dict: """Convert EvalMetadata into a single row for a Hub dataset.""" return { "source_dataset": metadata.source_dataset, + "source_split": metadata.source_split, "judge_models": json.dumps(metadata.judge_models), "seed": metadata.seed, "max_samples": metadata.max_samples, @@ -558,6 +563,7 @@ def _build_readme( "", f"- **Source dataset**: [`{metadata.source_dataset}`]" f"(https://huggingface.co/datasets/{metadata.source_dataset})", + f"- **Source split**: `{metadata.source_split}`", f"- **Judge**: {judge_str}", f"- **Judge criteria**: {metadata.criteria}", f"- **Judge prompt hash**: `{metadata.prompt_hash or 'unrecorded'}`", diff --git a/src/ocr_bench/run.py b/src/ocr_bench/run.py index 7c321d4..88e33fe 100644 --- a/src/ocr_bench/run.py +++ b/src/ocr_bench/run.py @@ -190,6 +190,7 @@ def build_script_args( output_repo: str, config_name: str, *, + split: str = "train", max_samples: int | None = None, shuffle: bool = False, seed: int = 42, @@ -202,6 +203,8 @@ def build_script_args( "--config", config_name, "--create-pr", + "--split", + split, ] if max_samples is not None: args += ["--max-samples", str(max_samples)] @@ -250,6 +253,7 @@ def launch_ocr_jobs( input_dataset, output_repo, slug, + split=split, max_samples=max_samples, shuffle=shuffle, seed=seed, diff --git a/src/ocr_bench/viewer.py b/src/ocr_bench/viewer.py index 684971b..359cf48 100644 --- a/src/ocr_bench/viewer.py +++ b/src/ocr_bench/viewer.py @@ -65,12 +65,12 @@ def load_results(repo_id: str) -> tuple[list[dict[str, Any]], list[dict[str, Any def _load_source_metadata(repo_id: str) -> dict[str, Any]: - """Load metadata config from results repo to find the source dataset.""" + """Load the latest run metadata used to resolve source images.""" revision = _latest_revision(repo_id) try: meta_ds = load_dataset(repo_id, name="metadata", split="train", revision=revision) if len(meta_ds) > 0: - return dict(meta_ds[0]) + return dict(meta_ds[-1]) except Exception as exc: logger.warning("could_not_load_metadata", repo=repo_id, error=str(exc)) return {} @@ -79,9 +79,15 @@ def _load_source_metadata(repo_id: str) -> dict[str, Any]: class ImageLoader: """Lazy image loader — fetches images from source dataset by sample_idx.""" - def __init__(self, source_dataset: str, from_prs: bool = False): + def __init__( + self, + source_dataset: str, + from_prs: bool = False, + source_split: str = "train", + ): self._source = source_dataset self._from_prs = from_prs + self._source_split = source_split self._cache: dict[int, Any] = {} self._image_col: str | None = None self._pr_revision: str | None = None @@ -105,7 +111,10 @@ def _init_source(self) -> None: self._pr_revision = revisions[first_config] # Probe for image column by loading 1 row - kwargs: dict[str, Any] = {"path": self._source, "split": "train[:1]"} + kwargs: dict[str, Any] = { + "path": self._source, + "split": f"{self._source_split}[:1]", + } if self._pr_revision: # Load from the first PR config first_config = next(iter(revisions)) @@ -133,7 +142,7 @@ def get(self, sample_idx: int) -> Image.Image | None: try: kwargs: dict[str, Any] = { "path": self._source, - "split": f"train[{sample_idx}:{sample_idx + 1}]", + "split": f"{self._source_split}[{sample_idx}:{sample_idx + 1}]", } if self._pr_revision: from ocr_bench.dataset import discover_pr_configs @@ -225,4 +234,3 @@ def _build_pair_summary(comparisons: list[dict[str, Any]]) -> str: parts.append(f"**{short_a}** vs **{short_b}**: {wins}W {losses}L {ties}T") return " | ".join(parts) - diff --git a/src/ocr_bench/web.py b/src/ocr_bench/web.py index 9d8f54f..259eadd 100644 --- a/src/ocr_bench/web.py +++ b/src/ocr_bench/web.py @@ -148,10 +148,16 @@ def create_app( metadata = _load_source_metadata(repo_id) source_dataset = metadata.get("source_dataset", "") from_prs = metadata.get("from_prs", False) + # Results written before source_split was recorded always came from train. + source_split = metadata.get("source_split") or "train" img_loader: ImageLoader | None = None if source_dataset: - img_loader = ImageLoader(source_dataset, from_prs=from_prs) + img_loader = ImageLoader( + source_dataset, + from_prs=from_prs, + source_split=source_split, + ) validation_comps = build_validation_comparisons( comparison_rows, diff --git a/tests/test_cli.py b/tests/test_cli.py index d6e93dc..ad4e5b6 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -935,6 +935,23 @@ def test_missing_repo_arguments_fails_cleanly(self, capsys): assert exc.value.code == 2 assert "requires INPUT_DATASET and OUTPUT_REPO" in capsys.readouterr().out + def test_dry_run_includes_split_in_job_args(self, capsys): + args = build_parser().parse_args( + [ + "run", + "in/ds", + "out/repo", + "--models", + "glm-ocr", + "--split", + "validation", + "--dry-run", + ] + ) + + assert cli.cmd_run(args) == [] + assert "--split validation" in capsys.readouterr().out + def _run(self, monkeypatch, statuses): from ocr_bench.run import JobRun diff --git a/tests/test_publish.py b/tests/test_publish.py index 4420783..8d38e38 100644 --- a/tests/test_publish.py +++ b/tests/test_publish.py @@ -119,9 +119,22 @@ def test_auto_timestamp(self): ) row = build_metadata_row(meta) assert row["source_dataset"] == "repo/data" + assert row["source_split"] == "train" assert row["timestamp"] # auto-filled assert '"judge-a"' in row["judge_models"] + def test_source_split_recorded(self): + meta = EvalMetadata( + source_dataset="repo/data", + source_split="validation", + judge_models=["judge-a"], + seed=42, + max_samples=10, + total_comparisons=30, + valid_comparisons=28, + ) + assert build_metadata_row(meta)["source_split"] == "validation" + def test_preserved_timestamp(self): meta = EvalMetadata( source_dataset="repo/data", @@ -756,6 +769,18 @@ def test_explicit_license_included(self): ) assert "license: cc0-1.0" in readme + def test_source_split_included(self): + from ocr_bench.publish import _build_readme + + board = _make_board() + rows = build_leaderboard_rows(board) + meta = self._make_metadata() + meta.source_split = "validation" + + readme = _build_readme("user/results", rows, board, meta) + + assert "- **Source split**: `validation`" in readme + def test_pipes_in_model_names_escaped(self): from ocr_bench.publish import _build_readme diff --git a/tests/test_run.py b/tests/test_run.py index bc0b5d3..e0f6211 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -112,7 +112,19 @@ def test_returns_sorted_slugs(self): class TestBuildScriptArgs: def test_basic_args(self): args = build_script_args("input/ds", "output/repo", "glm-ocr") - assert args == ["input/ds", "output/repo", "--config", "glm-ocr", "--create-pr"] + assert args == [ + "input/ds", + "output/repo", + "--config", + "glm-ocr", + "--create-pr", + "--split", + "train", + ] + + def test_non_default_split(self): + args = build_script_args("in", "out", "x", split="validation") + assert args[args.index("--split") + 1] == "validation" def test_max_samples(self): args = build_script_args("in", "out", "x", max_samples=50) @@ -173,6 +185,25 @@ def test_launches_subset(self, mock_token): assert jobs[0].model_slug == "glm-ocr" assert jobs[1].model_slug == "dots-ocr" + @patch("ocr_bench.run.get_token", return_value="fake-token") + def test_passes_split_to_every_job(self, mock_token): + mock_api = MagicMock() + mock_job = MagicMock(id="job-1", url="https://example.com") + mock_api.run_uv_job.return_value = mock_job + + launch_ocr_jobs( + "input/ds", + "output/repo", + models=["glm-ocr", "dots-ocr"], + split="validation", + api=mock_api, + ) + + assert mock_api.run_uv_job.call_count == 2 + for call in mock_api.run_uv_job.call_args_list: + script_args = call.kwargs["script_args"] + assert script_args[script_args.index("--split") + 1] == "validation" + @patch("ocr_bench.run.get_token", return_value="fake-token") def test_unknown_model_raises(self, mock_token): mock_api = MagicMock() @@ -332,4 +363,3 @@ def test_run_models_flag(self): parser = build_parser() args = parser.parse_args(["run", "in", "out", "--models", "glm-ocr", "dots-ocr"]) assert args.models == ["glm-ocr", "dots-ocr"] - diff --git a/tests/test_viewer.py b/tests/test_viewer.py index 7177175..564ced3 100644 --- a/tests/test_viewer.py +++ b/tests/test_viewer.py @@ -5,6 +5,7 @@ from unittest.mock import MagicMock, patch from ocr_bench.viewer import ( + ImageLoader, _filter_comparisons, _load_source_metadata, _winner_badge, @@ -111,6 +112,48 @@ def test_pins_metadata_load_to_revision(self, mock_load, mock_rev): mock_rev.assert_called_once_with("user/results") assert mock_load.call_args.kwargs.get("revision") == "sha456" + @patch("ocr_bench.viewer._latest_revision", return_value="sha456") + @patch("ocr_bench.viewer.load_dataset") + def test_returns_latest_metadata_row(self, mock_load, mock_rev): + rows = [ + {"source_dataset": "user/source"}, + {"source_dataset": "user/source", "source_split": "validation"}, + ] + mock_ds = MagicMock() + mock_ds.__len__ = MagicMock(return_value=len(rows)) + mock_ds.__getitem__ = MagicMock(side_effect=rows.__getitem__) + mock_load.return_value = mock_ds + + meta = _load_source_metadata("user/results") + + assert meta["source_split"] == "validation" + + +class TestImageLoader: + @patch("ocr_bench.viewer.load_dataset") + def test_uses_source_split_for_probe_and_row(self, mock_load): + probe = MagicMock() + probe.column_names = ["image"] + mock_load.side_effect = [probe, [{"image": "page-image"}]] + + loader = ImageLoader("user/source", source_split="validation") + + assert loader.get(3) == "page-image" + assert mock_load.call_args_list[0].kwargs["split"] == "validation[:1]" + assert mock_load.call_args_list[1].kwargs["split"] == "validation[3:4]" + + @patch("ocr_bench.viewer.load_dataset") + def test_old_metadata_defaults_to_train(self, mock_load): + probe = MagicMock() + probe.column_names = ["image"] + mock_load.side_effect = [probe, [{"image": "page-image"}]] + + loader = ImageLoader("user/source") + + assert loader.get(0) == "page-image" + assert mock_load.call_args_list[0].kwargs["split"] == "train[:1]" + assert mock_load.call_args_list[1].kwargs["split"] == "train[0:1]" + class TestFilterComparisons: def test_no_filters(self): diff --git a/tests/test_web.py b/tests/test_web.py index 5be8f69..db0635c 100644 --- a/tests/test_web.py +++ b/tests/test_web.py @@ -288,6 +288,61 @@ def test_image_with_loader(self, tmp_path): assert resp.status_code == 200 assert resp.headers["content-type"] == "image/png" + def test_passes_recorded_source_split_to_image_loader(self, tmp_path): + with ( + patch("ocr_bench.web.load_results") as mock_load, + patch("ocr_bench.web._load_source_metadata") as mock_meta, + patch("ocr_bench.web.load_annotations") as mock_ann, + patch("ocr_bench.web.ImageLoader") as mock_loader, + ): + mock_load.return_value = (SAMPLE_LEADERBOARD, SAMPLE_COMPARISONS) + mock_meta.return_value = { + "source_dataset": "user/source", + "source_split": "validation", + "from_prs": False, + } + mock_ann.return_value = ({}, []) + + from ocr_bench.web import create_app + + create_app( + "user/test-results", + output_path=str(tmp_path / "ann.json"), + ) + + mock_loader.assert_called_once_with( + "user/source", + from_prs=False, + source_split="validation", + ) + + def test_old_metadata_defaults_source_split_to_train(self, tmp_path): + with ( + patch("ocr_bench.web.load_results") as mock_load, + patch("ocr_bench.web._load_source_metadata") as mock_meta, + patch("ocr_bench.web.load_annotations") as mock_ann, + patch("ocr_bench.web.ImageLoader") as mock_loader, + ): + mock_load.return_value = (SAMPLE_LEADERBOARD, SAMPLE_COMPARISONS) + mock_meta.return_value = { + "source_dataset": "user/source", + "from_prs": False, + } + mock_ann.return_value = ({}, []) + + from ocr_bench.web import create_app + + create_app( + "user/test-results", + output_path=str(tmp_path / "ann.json"), + ) + + mock_loader.assert_called_once_with( + "user/source", + from_prs=False, + source_split="train", + ) + def test_image_not_found(self, tmp_path): """Test image 404 when loader returns None.""" with (