diff --git a/README.md b/README.md index a61a2eb..d118d7c 100644 --- a/README.md +++ b/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 --models glm-ocr lighton-ocr-2 +ocr-bench run --models tesseract falcon-ocr ovis-ocr2 ``` ## Example results diff --git a/src/ocr_bench/cli.py b/src/ocr_bench/cli.py index 11c13eb..b7890b0 100644 --- a/src/ocr_bench/cli.py +++ b/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: diff --git a/src/ocr_bench/run.py b/src/ocr_bench/run.py index 3119a55..7c321d4 100644 --- a/src/ocr_bench/run.py +++ b/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", diff --git a/tests/test_cli.py b/tests/test_cli.py index 8f78f91..35c0235 100644 --- a/tests/test_cli.py +++ b/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 diff --git a/tests/test_run.py b/tests/test_run.py index 7faf913..bc0b5d3 100644 --- a/tests/test_run.py +++ b/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