Honor source splits throughout benchmark runs

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,

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'}`",

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,

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)
 
-

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,

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
 

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
 

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"]
-

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):

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 (