Add Falcon OCR and OvisOCR2 to the model registry

README.md

@@ -75,7 +75,7 @@ To avoid wasting judge calls on uninformative pairs, the judge skips comparisons
 
 ## Available models
 
-ocr-bench ships with 5 OCR models ready to run:
+ocr-bench ships with 16 registered OCR models. Five form the default set:
 
 | Model           | Size | Best for                   | Notes                        |
 | --------------- | ---- | -------------------------- | ---------------------------- |
@@ -85,12 +85,18 @@ ocr-bench ships with 5 OCR models ready to run:
 | `deepseek-ocr`  | 4B   | Diverse documents          | Most consistent across types |
 | `dots-ocr`      | 1.7B | General                    | Struggles on historical text |
 
+The opt-in registry also includes compact models at different quality/cost points, including `falcon-ocr` (0.3B) and `ovis-ocr2` (0.9B). List every available model with:
+
+```bash
+ocr-bench run --list-models
+```
+
 All model scripts are available at [uv-scripts/ocr](https://huggingface.co/datasets/uv-scripts/ocr) on the Hub.
 
-By default all 5 run. To pick specific models:
+By default all 5 models in the table run. To pick specific models:
 
 ```bash
-ocr-bench run <dataset> <output> --models glm-ocr lighton-ocr-2
+ocr-bench run <dataset> <output> --models tesseract falcon-ocr ovis-ocr2
 ```
 
 ## Example results

src/ocr_bench/cli.py

@@ -249,10 +249,18 @@ def build_parser() -> argparse.ArgumentParser:
 
     # --- run subcommand ---
     run = sub.add_parser("run", help="Launch OCR models on a dataset via HF Jobs")
-    run.add_argument("input_dataset", help="HF dataset repo id with images")
-    run.add_argument("output_repo", help="Output dataset repo (all models push here)")
     run.add_argument(
-        "--models", nargs="+", default=None, help="Model slugs to run (default: all 4 core)"
+        "input_dataset",
+        nargs="?",
+        help="HF dataset repo id with images (not needed with --list-models)",
+    )
+    run.add_argument(
+        "output_repo",
+        nargs="?",
+        help="Output dataset repo (not needed with --list-models)",
+    )
+    run.add_argument(
+        "--models", nargs="+", default=None, help="Model slugs to run (default: core set)"
     )
     run.add_argument("--max-samples", type=int, default=None, help="Per-model sample limit")
     run.add_argument("--split", default="train", help="Dataset split (default: train)")
@@ -1349,6 +1357,13 @@ def cmd_run(args: argparse.Namespace) -> list[JobRun]:
         console.print(f"\nDefault set: {', '.join(DEFAULT_MODELS)}")
         return []
 
+    if not args.input_dataset or not args.output_repo:
+        console.print(
+            "[red]Error:[/red] run requires INPUT_DATASET and OUTPUT_REPO "
+            "unless --list-models is used."
+        )
+        raise SystemExit(2)
+
     selected = args.models or DEFAULT_MODELS
     for slug in selected:
         if slug not in MODEL_REGISTRY:

src/ocr_bench/run.py

@@ -44,6 +44,12 @@ class ModelConfig:
 
 
 MODEL_REGISTRY: dict[str, ModelConfig] = {
+    "falcon-ocr": ModelConfig(
+        script="https://huggingface.co/datasets/uv-scripts/ocr/raw/main/falcon-ocr.py",
+        model_id="tiiuae/Falcon-OCR",
+        size="0.3B",
+        default_flavor="l4x1",
+    ),
     "glm-ocr": ModelConfig(
         script="https://huggingface.co/datasets/uv-scripts/ocr/raw/main/glm-ocr.py",
         model_id="zai-org/GLM-OCR",
@@ -63,6 +69,12 @@ class ModelConfig:
         size="1B",
         default_flavor="a100-large",
     ),
+    "ovis-ocr2": ModelConfig(
+        script="https://huggingface.co/datasets/uv-scripts/ocr/raw/main/ovis-ocr2.py",
+        model_id="ATH-MaaS/OvisOCR2",
+        size="0.9B",
+        default_flavor="l4x1",
+    ),
     "dots-ocr": ModelConfig(
         script="https://huggingface.co/datasets/uv-scripts/ocr/raw/main/dots-ocr.py",
         model_id="rednote-hilab/dots.ocr",

tests/test_cli.py

@@ -872,6 +872,20 @@ def test_aborts_when_a_job_fails(self, monkeypatch, capsys):
 class TestCmdRun:
     """cmd_run returns the launched jobs and prints a per-job status summary."""
 
+    def test_list_models_needs_no_repo_arguments(self, capsys):
+        args = build_parser().parse_args(["run", "--list-models"])
+        assert cli.cmd_run(args) == []
+        out = capsys.readouterr().out
+        assert "falcon-ocr" in out
+        assert "ovis-ocr2" in out
+
+    def test_missing_repo_arguments_fails_cleanly(self, capsys):
+        args = build_parser().parse_args(["run"])
+        with pytest.raises(SystemExit) as exc:
+            cli.cmd_run(args)
+        assert exc.value.code == 2
+        assert "requires INPUT_DATASET and OUTPUT_REPO" in capsys.readouterr().out
+
     def _run(self, monkeypatch, statuses):
         from ocr_bench.run import JobRun
 

tests/test_run.py

@@ -49,7 +49,7 @@ def test_image_mode_fields_default_none(self):
 
 class TestModelRegistry:
     def test_has_core_models(self):
-        assert len(MODEL_REGISTRY) == 14
+        assert len(MODEL_REGISTRY) == 16
 
     def test_default_models_exist_in_registry(self):
         for slug in DEFAULT_MODELS:
@@ -67,6 +67,26 @@ def test_deepseek_has_prompt_mode_free(self):
         assert "--prompt-mode" in cfg.default_args
         assert "free" in cfg.default_args
 
+    def test_compact_models_use_standard_l4_launch(self):
+        expected = {
+            "falcon-ocr": ("tiiuae/Falcon-OCR", "0.3B", "falcon-ocr.py"),
+            "ovis-ocr2": ("ATH-MaaS/OvisOCR2", "0.9B", "ovis-ocr2.py"),
+        }
+        for slug, (model_id, size, script_name) in expected.items():
+            cfg = MODEL_REGISTRY[slug]
+            assert cfg.model_id == model_id
+            assert cfg.size == size
+            assert cfg.script.endswith(f"/{script_name}")
+            assert cfg.default_flavor == "l4x1"
+            assert cfg.default_args == []
+            assert cfg.image is None
+            assert cfg.python is None
+            assert cfg.env is None
+
+    def test_compact_models_are_opt_in(self):
+        assert "falcon-ocr" not in DEFAULT_MODELS
+        assert "ovis-ocr2" not in DEFAULT_MODELS
+
     def test_image_mode_models_configured(self):
         # NuExtract3 and PaddleOCR-VL-1.6 need the prebuilt-kernel image on a100.
         for slug in ("nuextract3", "paddleocr-vl-1.6"):
@@ -290,12 +310,14 @@ def test_run_subcommand_parses(self):
         assert args.output_repo == "output/repo"
         assert args.max_samples == 50
 
-    def test_run_list_models(self):
+    def test_run_list_models_without_repo_arguments(self):
         from ocr_bench.cli import build_parser
 
         parser = build_parser()
-        args = parser.parse_args(["run", "in", "out", "--list-models"])
+        args = parser.parse_args(["run", "--list-models"])
         assert args.list_models is True
+        assert args.input_dataset is None
+        assert args.output_repo is None
 
     def test_run_dry_run(self):
         from ocr_bench.cli import build_parser