From b07201c01247d529933ce11eac01ff88c66b8b04 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 27 Jun 2026 09:14:27 +0000 Subject: [PATCH 1/2] Add Baidu Unlimited-OCR engine adapter Add UnlimitedOCREngine, wrapping Baidu's open-weight one-shot long-horizon document-parsing VLM (loaded from HuggingFace via trust_remote_code). Supports gundam/base modes, images, and PDFs (rendered to images via PyMuPDF), emitting Markdown/HTML/JSON. - New engine: src/docfold/engines/unlimited_ocr_engine.py - Register in router priority (PDF + image) and default fallback - Register in CLI _build_router() - Add unlimited-ocr optional-dependency group (+ to [all]) - Tests: TestUnlimitedOCREngine + interface parametrize entry - Docs: proposal, README tables, benchmarks profile, CHANGELOG Follows the TDD workflow (proposal -> failing tests -> implementation). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P23UQav6r1QsPGKzogma26 --- CHANGELOG.md | 1 + README.md | 3 + docs/benchmarks.md | 11 + docs/tasks/UNLIMITED_OCR_ENGINE.md | 98 ++++++++ pyproject.toml | 11 +- src/docfold/cli.py | 6 + src/docfold/engines/router.py | 10 +- src/docfold/engines/unlimited_ocr_engine.py | 255 ++++++++++++++++++++ tests/engines/test_adapters.py | 128 ++++++++++ 9 files changed, 517 insertions(+), 6 deletions(-) create mode 100644 docs/tasks/UNLIMITED_OCR_ENGINE.md create mode 100644 src/docfold/engines/unlimited_ocr_engine.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 04a6868..09eeddd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- **Unlimited-OCR engine adapter** — wraps Baidu's open-weight [`Unlimited-OCR`](https://github.com/baidu/Unlimited-OCR) document-parsing VLM (loaded from HuggingFace via `trust_remote_code`). Builds on DeepSeek-OCR with Reference Sliding Window Attention (R-SWA) for one-shot long-horizon parsing. Supports `gundam` and `base` modes, images, and PDFs (rendered to images via PyMuPDF), emitting Markdown/HTML/JSON. Registered in the router for PDF and image inputs. Install: `pip install docfold[unlimited-ocr]` (requires a CUDA GPU). - **MarkItDown engine adapter** — wraps Microsoft's [`markitdown`](https://github.com/microsoft/markitdown) pure-Python library that converts Office files, PDFs, HTML, images, CSV/JSON/XML, ePub, audio, and ZIP archives into LLM-friendly Markdown. Added to the `benchmark.py` harness alongside the other local engines. Install: `pip install docfold[markitdown]`. - **Non-PDF benchmark fixtures** — `benchmark.py` now also generates synthetic DOCX (built with stdlib `zipfile` + minimal Office Open XML, no extra deps), HTML, and CSV documents, and filters engines per-doc by `supported_extensions` so PyMuPDF / OCR engines no longer log spurious errors on Office or web fixtures. - **OpenDataLoader PDF engine adapter** — wraps the Java-based [`opendataloader-pdf`](https://github.com/opendataloader-project/opendataloader-pdf) tool (via its bundled-JAR Python wheel). Local, deterministic extraction with typed structural elements (heading, paragraph, table, list, header, footer) and per-element bounding boxes. Install: `pip install docfold[opendataloader]` (also requires Java 11+). diff --git a/README.md b/README.md index c623586..df0797e 100644 --- a/README.md +++ b/README.md @@ -35,6 +35,7 @@ Docfold is the open-source extraction engine from [Datatera.ai](https://datatera | [**Azure Doc Intel**](https://azure.microsoft.com/en-us/products/ai-services/ai-document-intelligence) | ✅ | SaaS | Paid | ★★★ | ★★★ | ★★★ | ✅ | ✅ | Fast | $$ | | [**Nougat**](https://github.com/facebookresearch/nougat) | ✅ | Local | MIT | ★★★ | ★★☆ | ★★☆ | — | — | Slow | Free | | [**Surya**](https://github.com/VikParuchuri/surya) | ✅ | Local | GPL | ★★☆ | ★★★ | ★★☆ | ✅ | ✅ | Medium | Free | +| [**Unlimited-OCR**](https://github.com/baidu/Unlimited-OCR) | ✅ | Local/VLM | MIT | ★★★ | ★★★ | ★★★ | — | — | Slow | Free | | [**MarkItDown**](https://github.com/microsoft/markitdown) | ✅ | Local | MIT | ★★☆ | ★☆☆ | ★★☆ | — | — | Fast | Free | **★★★** Excellent **★★☆** Good **★☆☆** Basic **☆☆☆** Not supported — **$$** ~$1-3/1K pages **$$$** ~$5-15/1K pages — **BBox** Bounding boxes — **Conf** Confidence scores @@ -109,6 +110,7 @@ for name, res in results.items(): | [**Azure Doc Intel**](https://azure.microsoft.com/en-us/products/ai-services/ai-document-intelligence) | SaaS | Paid | PDF, Office, HTML, images | N/A | `pip install docfold[azure-docint]` | | [**Nougat**](https://github.com/facebookresearch/nougat) | Local | MIT (code) | PDF | Recommended | `pip install docfold[nougat]` | | [**Surya**](https://github.com/VikParuchuri/surya) | Local | GPL-3.0 | PDF, images | Optional | `pip install docfold[surya]` | +| [**Unlimited-OCR**](https://github.com/baidu/Unlimited-OCR) | Local/VLM | MIT | PDF, images | Required | `pip install docfold[unlimited-ocr]` | | [**MarkItDown**](https://github.com/microsoft/markitdown) | Local | MIT | PDF, Office, HTML, images, CSV/JSON/XML, ePub, audio, ZIP | No | `pip install docfold[markitdown]` | > **Adding your own engine?** Implement the `DocumentEngine` interface — see [Adding a Custom Engine](#adding-a-custom-engine) below. @@ -335,6 +337,7 @@ Docfold builds on and integrates with these excellent projects: | [Zerox](https://github.com/getomni-ai/zerox) | Model-agnostic Vision LLM OCR | | [Nougat](https://github.com/facebookresearch/nougat) | Meta's academic PDF to Markdown model | | [Surya](https://github.com/VikParuchuri/surya) | Multilingual OCR + layout analysis | +| [Unlimited-OCR](https://github.com/baidu/Unlimited-OCR) | Baidu's one-shot long-horizon document-parsing VLM | ### Built by diff --git a/docs/benchmarks.md b/docs/benchmarks.md index 720e46f..ebb5a82 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -23,6 +23,7 @@ This guide helps you choose the right document processing engine for your use ca | **Zerox** | VLM | MIT | ★★★ | ★★★ | ★★☆ | ★★☆ | ★★★ | Slow | VLM API cost | | **Nougat** | Local | MIT | ★★★ | ★★☆ | ★★☆ | ★★★ | ★☆☆ | Slow | Free | | **Surya** | Local | GPL-3.0 | ★★☆ | ★★★ | ★★☆ | ★☆☆ | ★★★ (90+) | Medium | Free | +| **Unlimited-OCR** | Local/VLM | MIT | ★★★ | ★★★ | ★★★ | ★★☆ | ★★☆ | Slow | Free | | AWS Textract | SaaS | Paid | ★★★ | ★★★ | ★★★ | ★☆☆ | ★★☆ | Fast | ~$1.50/1K pages | | Google Document AI | SaaS | Paid | ★★★ | ★★★ | ★★★ | ★★☆ | ★★★ | Fast | ~$1.50/1K pages | | Azure Document Intelligence | SaaS | Paid | ★★★ | ★★★ | ★★★ | ★★☆ | ★★★ | Fast | ~$1.50/1K pages | @@ -238,6 +239,16 @@ This guide helps you choose the right document processing engine for your use ca - **Install:** `pip install docfold[surya]` - **Links:** [GitHub](https://github.com/VikParuchuri/surya) +### Unlimited-OCR (Baidu) + +**Best for:** Long, multi-page documents parsed in one shot — the headline "long-horizon" use case. + +- **Strengths:** Open-weight VLM (MIT) released June 2026. Builds on DeepSeek-OCR and replaces decoder attention with **Reference Sliding Window Attention (R-SWA)** for a *constant* KV cache, so it can transcribe dozens of pages in a single 32K-context forward pass. Strong Markdown output with headings, tables, and reading order. Two modes: `gundam` (dynamic crop, single images) and `base` (full resolution, multi-page). +- **Weaknesses:** Requires a CUDA GPU and heavy deps (torch, transformers). Loaded via `trust_remote_code` from HuggingFace. Newer project — API may evolve. The docfold adapter processes page-by-page; native single-pass `infer_multi` long-horizon batching is a future enhancement. +- **GPU:** Required (CUDA), bfloat16. +- **Install:** `pip install docfold[unlimited-ocr]` +- **Links:** [GitHub](https://github.com/baidu/Unlimited-OCR) | [HuggingFace](https://huggingface.co/baidu/Unlimited-OCR) | [Paper](https://arxiv.org/abs/2606.23050) + --- ## Feature Coverage Matrix diff --git a/docs/tasks/UNLIMITED_OCR_ENGINE.md b/docs/tasks/UNLIMITED_OCR_ENGINE.md new file mode 100644 index 0000000..f0eb9bf --- /dev/null +++ b/docs/tasks/UNLIMITED_OCR_ENGINE.md @@ -0,0 +1,98 @@ +--- +purpose: "Add Baidu Unlimited-OCR as a local VLM document-parsing engine" +status: "OPEN" +priority: "P2" +created: "2026-06-27" +--- + +# Feature: Unlimited-OCR Engine + +## Problem +[Baidu Unlimited-OCR](https://github.com/baidu/Unlimited-OCR) (released June 2026, +10.7k★, MIT-licensed code) is a new open-weight document-parsing VLM. It takes +DeepSeek-OCR as a baseline and replaces the decoder attention with **Reference +Sliding Window Attention (R-SWA)**, giving a *constant* KV cache across decoding. +Combined with DeepSeek-OCR's high-compression encoder, it can transcribe **dozens +of pages in a single forward pass** under a 32K context — the headline "one-shot +long-horizon parsing" capability. It outputs structured Markdown (headings, +tables, reading order) and is competitive on document-parsing benchmarks. + +docfold already ships 20+ engine adapters (Chandra, Surya, MinerU, …) behind a +unified `DocumentEngine` interface. Adding Unlimited-OCR gives users a fresh, +free, locally-runnable VLM option — particularly for long, multi-page documents +where its constant-KV-cache design is an advantage. "Now" because the model just +shipped and there is clear user interest. + +## Proposed Solution +Add a new `UnlimitedOCREngine` adapter that wraps the upstream HuggingFace model +(`baidu/Unlimited-OCR`, loaded with `trust_remote_code=True`), mirroring the +existing `ChandraEngine`/`SuryaEngine` adapters: + +- Lazy-load the model + tokenizer on first `process()` call (never at import, + construction, `is_available()`, or router-registration time). +- Run blocking inference inside `loop.run_in_executor` like the other local + engines. +- Support the upstream **`gundam`** (base_size=1024, image_size=640, + crop_mode=True) and **`base`** (base_size=1024, image_size=1024, + crop_mode=False) modes via a `mode` constructor parameter. +- Accept images directly and render PDF pages to images (via PyMuPDF/`fitz`) + before inference, like `ChandraEngine`. +- Emit Markdown / HTML / JSON / text from the unified `EngineResult`. + +## Affected Files +- `src/docfold/engines/unlimited_ocr_engine.py` — NEW adapter. +- `src/docfold/engines/router.py` — add `unlimited_ocr` to PDF + image priority + lists and the default fallback. +- `src/docfold/cli.py` — register the engine in `_build_router()`. +- `pyproject.toml` — add `unlimited-ocr` optional-dependency group; include in + `all`. +- `tests/engines/test_adapters.py` — `TestUnlimitedOCREngine` + add to the + `TestAllEnginesImplementInterface` parametrize list. +- `README.md`, `docs/benchmarks.md`, `CHANGELOG.md` — document the engine. + +## Test Plan + +### Unit / Functional Tests +- [ ] `test_name` — `name == "unlimited_ocr"`. +- [ ] `test_supported_extensions` — pdf + common image extensions present. +- [ ] `test_is_available_when_missing` — returns `bool` (False) when `torch` + is unavailable. +- [ ] `test_config_stored` — constructor params stored (`_mode`, `_model`, + `_max_length`, `_prompt`, `_device`). +- [ ] `test_default_mode_is_gundam` — `_mode` defaults to `"gundam"`. +- [ ] `test_mode_params` — `gundam` → (1024, 640, True); `base` → (1024, 1024, False). +- [ ] `test_capabilities` — table_structure / heading_detection / reading_order + True; bounding_boxes / confidence False. +- [ ] `test_process_returns_engine_result` — mocked model produces a valid + `EngineResult` (markdown). +- [ ] `test_process_json_output` — JSON output format wraps per-page text. +- [ ] Added to `TestAllEnginesImplementInterface` parametrize list. + +### Integration / E2E Tests +- [ ] (manual, GPU) Run a real image and a multi-page PDF through the engine and + verify Markdown output and page count. + +### Test Commands +```bash +# Run the new engine tests +pytest tests/engines/test_adapters.py -k UnlimitedOCR -v + +# Full suite (no regressions) +pytest tests/ -m "not slow" +``` + +## Edge Cases +- `torch` / `transformers` not installed → `is_available()` returns `False`, + engine is simply not registered. +- PDF input with no PyMuPDF → raise a clear `ImportError` from the render step. +- Empty / unreadable page → model returns empty string; adapter keeps the page + slot so page numbering stays correct. +- Model emits text only when `save_results=False`; adapter relies on the return + value and uses a temp dir for any side-effect output. + +## Out of Scope +- Native multi-page `infer_multi` long-horizon batching (single forward pass over + many pages). The first cut processes page-by-page for determinism and + testability; long-horizon batching is a future enhancement. +- Bounding-box extraction (the simple parse path returns Markdown only). +- vLLM serving backend. diff --git a/pyproject.toml b/pyproject.toml index fb25214..a509103 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -105,6 +105,15 @@ surya = [ "torch>=2.0", "Pillow>=10.0", ] +unlimited-ocr = [ + # Model is loaded from HuggingFace via trust_remote_code — no dedicated package. + "torch>=2.0", + "torchvision>=0.15", + "transformers>=4.57", + "Pillow>=10.0", + "einops>=0.8", + "PyMuPDF>=1.23", # render PDF pages to images +] firecrawl = [ "firecrawl-py>=1.0", ] @@ -119,7 +128,7 @@ evaluation = [ "psutil>=5.9", # Memory measurement ] all = [ - "docfold[docling,mineru,marker,pymupdf,paddleocr,tesseract,easyocr,unstructured,llamaparse,liteparse,opendataloader,mistral-ocr,textract,google-docai,azure-docint,nougat,chandra,surya,firecrawl,markitdown,evaluation]", + "docfold[docling,mineru,marker,pymupdf,paddleocr,tesseract,easyocr,unstructured,llamaparse,liteparse,opendataloader,mistral-ocr,textract,google-docai,azure-docint,nougat,chandra,surya,unlimited-ocr,firecrawl,markitdown,evaluation]", # Note: zerox excluded from [all] — py-zerox requires Python 3.11+ # Install separately: pip install docfold[zerox] ] diff --git a/src/docfold/cli.py b/src/docfold/cli.py index 8671289..935d760 100644 --- a/src/docfold/cli.py +++ b/src/docfold/cli.py @@ -190,6 +190,12 @@ def _build_router(): except Exception: pass + try: + from docfold.engines.unlimited_ocr_engine import UnlimitedOCREngine + router.register(UnlimitedOCREngine()) + except Exception: + pass + try: from docfold.engines.firecrawl_engine import FirecrawlEngine router.register(FirecrawlEngine()) diff --git a/src/docfold/engines/router.py b/src/docfold/engines/router.py index 8cafa77..14040ac 100644 --- a/src/docfold/engines/router.py +++ b/src/docfold/engines/router.py @@ -23,15 +23,15 @@ # the first *available* engine that supports the extension. _IMAGE_PRIORITY = [ - "chandra", "surya", "paddleocr", "tesseract", "easyocr", "docling", "liteparse", - "mistral_ocr", "google_docai", "textract", "azure_docint", "zerox", "marker", - "markitdown", + "chandra", "unlimited_ocr", "surya", "paddleocr", "tesseract", "easyocr", + "docling", "liteparse", "mistral_ocr", "google_docai", "textract", + "azure_docint", "zerox", "marker", "markitdown", ] _EXTENSION_PRIORITY: dict[str, list[str]] = { # --- PDF --- "pdf": [ - "docling", "mineru", "chandra", "unstructured", "marker", + "docling", "mineru", "chandra", "unlimited_ocr", "unstructured", "marker", "llamaparse", "liteparse", "mistral_ocr", "firecrawl", "google_docai", "azure_docint", "textract", "zerox", "nougat", "surya", "pymupdf", "paddleocr", "tesseract", "easyocr", "markitdown", @@ -93,7 +93,7 @@ # Ultimate fallback when extension is unknown or missing from the map. _DEFAULT_FALLBACK = [ - "docling", "mineru", "chandra", "unstructured", "marker", + "docling", "mineru", "chandra", "unlimited_ocr", "unstructured", "marker", "llamaparse", "liteparse", "mistral_ocr", "google_docai", "azure_docint", "textract", "zerox", "nougat", "surya", "pymupdf", "paddleocr", "tesseract", "easyocr", "markitdown", diff --git a/src/docfold/engines/unlimited_ocr_engine.py b/src/docfold/engines/unlimited_ocr_engine.py new file mode 100644 index 0000000..b90b351 --- /dev/null +++ b/src/docfold/engines/unlimited_ocr_engine.py @@ -0,0 +1,255 @@ +"""Unlimited-OCR engine adapter — Baidu's one-shot long-horizon document parser. + +Install: ``pip install docfold[unlimited-ocr]`` + +Unlimited-OCR is an open-weight document-parsing Vision-Language Model from +Baidu (released June 2026). It takes DeepSeek-OCR as a baseline and replaces the +decoder's attention layers with **Reference Sliding Window Attention (R-SWA)**, +keeping a constant KV cache across decoding. Combined with DeepSeek-OCR's +high-compression encoder, it can transcribe dozens of pages of a document in a +single forward pass under a 32K context window — the "one-shot long-horizon +parsing" capability that gives the project its name. + +The model is loaded from HuggingFace (``baidu/Unlimited-OCR``) via +``trust_remote_code=True``; there is no separate pip package to import. + +Two inference modes are supported, matching upstream: + +- **gundam** (default) — ``base_size=1024, image_size=640, crop_mode=True``; + dynamic cropping, best for single high-resolution images. +- **base** — ``base_size=1024, image_size=1024, crop_mode=False``; full + resolution, best for multi-page documents. + +Model & code license: MIT. +""" + +from __future__ import annotations + +import logging +import time +from pathlib import Path +from typing import Any + +from docfold.engines.base import DocumentEngine, EngineCapabilities, EngineResult, OutputFormat + +logger = logging.getLogger(__name__) + +_SUPPORTED_EXTENSIONS = {"pdf", "png", "jpg", "jpeg", "tiff", "tif", "bmp", "webp"} + +# mode -> (base_size, image_size, crop_mode) +_MODE_PARAMS: dict[str, tuple[int, int, bool]] = { + "gundam": (1024, 640, True), + "base": (1024, 1024, False), +} + +_DTYPE_MAP = { + "bfloat16": "bfloat16", + "float16": "float16", + "float32": "float32", +} + + +class UnlimitedOCREngine(DocumentEngine): + """Adapter for Baidu Unlimited-OCR (document → Markdown/HTML/JSON). + + Loads the ``baidu/Unlimited-OCR`` model from HuggingFace with + ``trust_remote_code=True`` and runs ``model.infer`` per page. PDFs are + rendered to images (via PyMuPDF) before inference. + + See https://github.com/baidu/Unlimited-OCR + """ + + def __init__( + self, + mode: str = "gundam", + model: str = "baidu/Unlimited-OCR", + prompt: str = "document parsing.", + max_length: int = 32768, + torch_dtype: str = "bfloat16", + device: str = "cuda", + dpi: int = 200, + ) -> None: + self._mode = mode + self._model = model + self._prompt = prompt + self._max_length = max_length + self._torch_dtype = torch_dtype + self._device = device + self._dpi = dpi + # Lazily populated on first process() call. + self._loaded_model: Any = None + self._tokenizer: Any = None + + @property + def name(self) -> str: + return "unlimited_ocr" + + @property + def supported_extensions(self) -> set[str]: + return _SUPPORTED_EXTENSIONS + + @property + def capabilities(self) -> EngineCapabilities: + return EngineCapabilities( + table_structure=True, + heading_detection=True, + reading_order=True, + ) + + def is_available(self) -> bool: + try: + import torch # noqa: F401 + import transformers # noqa: F401 + from PIL import Image # noqa: F401 + + return True + except Exception: + return False + + def _mode_params(self) -> tuple[int, int, bool]: + """Return ``(base_size, image_size, crop_mode)`` for the configured mode.""" + if self._mode not in _MODE_PARAMS: + raise ValueError( + f"Unknown mode '{self._mode}'. Choose one of: {sorted(_MODE_PARAMS)}" + ) + return _MODE_PARAMS[self._mode] + + def _ensure_model(self) -> tuple[Any, Any]: + """Lazy-load the model + tokenizer, caching them on the instance.""" + if self._loaded_model is not None and self._tokenizer is not None: + return self._loaded_model, self._tokenizer + + import torch + from transformers import AutoModel, AutoTokenizer + + dtype = getattr(torch, _DTYPE_MAP.get(self._torch_dtype, "bfloat16")) + + tokenizer = AutoTokenizer.from_pretrained(self._model, trust_remote_code=True) + model = AutoModel.from_pretrained( + self._model, + trust_remote_code=True, + use_safetensors=True, + torch_dtype=dtype, + ) + model = model.eval().to(self._device) + + self._loaded_model = model + self._tokenizer = tokenizer + return model, tokenizer + + async def process( + self, + file_path: str, + output_format: OutputFormat = OutputFormat.MARKDOWN, + **kwargs: Any, + ) -> EngineResult: + import asyncio + + start = time.perf_counter() + + loop = asyncio.get_running_loop() + content, page_count = await loop.run_in_executor( + None, self._do_process, file_path, output_format, + ) + + elapsed_ms = int((time.perf_counter() - start) * 1000) + + return EngineResult( + content=content, + format=output_format, + engine_name=self.name, + pages=page_count, + processing_time_ms=elapsed_ms, + metadata={ + "model": self._model, + "mode": self._mode, + "max_length": self._max_length, + }, + ) + + def _do_process( + self, file_path: str, output_format: OutputFormat, + ) -> tuple[str, int]: + import tempfile + + model, tokenizer = self._ensure_model() + base_size, image_size, crop_mode = self._mode_params() + + image_paths, cleanup_dir = self._to_image_paths(file_path) + try: + pages_text: list[str] = [] + # ``output_path`` is required by infer for any side-effect output; + # use a throwaway dir and rely on the returned text instead. + with tempfile.TemporaryDirectory() as out_dir: + for img_path in image_paths: + text = model.infer( + tokenizer, + prompt=self._prompt, + image_file=img_path, + output_path=out_dir, + base_size=base_size, + image_size=image_size, + crop_mode=crop_mode, + max_length=self._max_length, + save_results=False, + ) + pages_text.append((text or "").strip()) + finally: + if cleanup_dir is not None: + cleanup_dir.cleanup() + + page_count = len(pages_text) + content = self._format_output(pages_text, output_format) + return content, page_count + + def _to_image_paths(self, file_path: str) -> tuple[list[str], Any]: + """Return image file paths for the input plus an optional temp-dir handle. + + Images are used as-is. PDFs are rendered page-by-page to PNG files in a + :class:`tempfile.TemporaryDirectory`, returned as the second element so + the caller can clean it up. + """ + ext = Path(file_path).suffix.lstrip(".").lower() + if ext != "pdf": + return [file_path], None + + import tempfile + + tmp = tempfile.TemporaryDirectory() + paths = self._render_pdf(file_path, tmp.name) + return paths, tmp + + def _render_pdf(self, file_path: str, out_dir: str) -> list[str]: + """Render each PDF page to a PNG file under *out_dir*.""" + import fitz # PyMuPDF + + paths: list[str] = [] + doc = fitz.open(file_path) + try: + for i, page in enumerate(doc): + pix = page.get_pixmap(dpi=self._dpi) + out_path = str(Path(out_dir) / f"page_{i + 1:04d}.png") + pix.save(out_path) + paths.append(out_path) + finally: + doc.close() + return paths + + def _format_output(self, pages_text: list[str], output_format: OutputFormat) -> str: + if output_format == OutputFormat.JSON: + import json + + return json.dumps( + {"pages": [{"page": i + 1, "text": t} for i, t in enumerate(pages_text)]}, + ensure_ascii=False, + ) + + if output_format == OutputFormat.HTML: + html_parts = [ + f"
{t}
" + for i, t in enumerate(pages_text) + ] + return "" + "\n".join(html_parts) + "" + + # MARKDOWN / TEXT + return "\n\n".join(pages_text) diff --git a/tests/engines/test_adapters.py b/tests/engines/test_adapters.py index 4d9b1a1..c61e34f 100644 --- a/tests/engines/test_adapters.py +++ b/tests/engines/test_adapters.py @@ -1282,6 +1282,133 @@ def test_capabilities(self): assert caps.confidence is False +class TestUnlimitedOCREngine: + def test_name(self): + from docfold.engines.unlimited_ocr_engine import UnlimitedOCREngine + e = UnlimitedOCREngine() + assert e.name == "unlimited_ocr" + + def test_supported_extensions(self): + from docfold.engines.unlimited_ocr_engine import UnlimitedOCREngine + e = UnlimitedOCREngine() + exts = e.supported_extensions + assert "pdf" in exts + assert "png" in exts + assert "jpg" in exts + assert "jpeg" in exts + assert "tiff" in exts + assert "bmp" in exts + assert "webp" in exts + + def test_is_available_when_missing(self): + from docfold.engines.unlimited_ocr_engine import UnlimitedOCREngine + e = UnlimitedOCREngine() + with patch.dict("sys.modules", {"torch": None}): + result = e.is_available() + assert isinstance(result, bool) + + def test_config_stored(self): + from docfold.engines.unlimited_ocr_engine import UnlimitedOCREngine + e = UnlimitedOCREngine( + mode="base", + model="baidu/Unlimited-OCR", + max_length=16384, + prompt="parse.", + device="cpu", + ) + assert e._mode == "base" + assert e._model == "baidu/Unlimited-OCR" + assert e._max_length == 16384 + assert e._prompt == "parse." + assert e._device == "cpu" + + def test_default_mode_is_gundam(self): + from docfold.engines.unlimited_ocr_engine import UnlimitedOCREngine + e = UnlimitedOCREngine() + assert e._mode == "gundam" + + def test_mode_params(self): + from docfold.engines.unlimited_ocr_engine import UnlimitedOCREngine + gundam = UnlimitedOCREngine(mode="gundam")._mode_params() + assert gundam == (1024, 640, True) + base = UnlimitedOCREngine(mode="base")._mode_params() + assert base == (1024, 1024, False) + + def test_capabilities(self): + from docfold.engines.unlimited_ocr_engine import UnlimitedOCREngine + caps = UnlimitedOCREngine().capabilities + assert caps.table_structure is True + assert caps.heading_detection is True + assert caps.reading_order is True + assert caps.bounding_boxes is False + assert caps.confidence is False + + @pytest.mark.asyncio + async def test_process_returns_engine_result(self): + """Unlimited-OCR processes an image and returns a valid EngineResult.""" + import os + import tempfile + from unittest.mock import MagicMock, patch + + from docfold.engines.base import EngineResult, OutputFormat + from docfold.engines.unlimited_ocr_engine import UnlimitedOCREngine + + mock_model = MagicMock() + mock_model.infer.return_value = "# Hello\n\nExtracted content" + mock_tokenizer = MagicMock() + + prefix = "docfold.engines.unlimited_ocr_engine" + with patch( + f"{prefix}.UnlimitedOCREngine._ensure_model", + return_value=(mock_model, mock_tokenizer), + ): + e = UnlimitedOCREngine(device="cpu") + with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f: + f.write(b"\x89PNG\r\n\x1a\n fake") + try: + result = await e.process(f.name, OutputFormat.MARKDOWN) + assert isinstance(result, EngineResult) + assert result.engine_name == "unlimited_ocr" + assert result.content == "# Hello\n\nExtracted content" + assert result.format == OutputFormat.MARKDOWN + assert result.pages == 1 + assert result.processing_time_ms >= 0 + mock_model.infer.assert_called_once() + finally: + os.unlink(f.name) + + @pytest.mark.asyncio + async def test_process_json_output(self): + """Unlimited-OCR returns per-page JSON when output_format is JSON.""" + import json + import os + import tempfile + from unittest.mock import MagicMock, patch + + from docfold.engines.base import OutputFormat + from docfold.engines.unlimited_ocr_engine import UnlimitedOCREngine + + mock_model = MagicMock() + mock_model.infer.return_value = "page text" + mock_tokenizer = MagicMock() + + prefix = "docfold.engines.unlimited_ocr_engine" + with patch( + f"{prefix}.UnlimitedOCREngine._ensure_model", + return_value=(mock_model, mock_tokenizer), + ): + e = UnlimitedOCREngine(device="cpu") + with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f: + f.write(b"\x89PNG\r\n\x1a\n fake") + try: + result = await e.process(f.name, OutputFormat.JSON) + assert result.format == OutputFormat.JSON + parsed = json.loads(result.content) + assert parsed["pages"][0]["text"] == "page text" + finally: + os.unlink(f.name) + + class TestAllEnginesImplementInterface: """Verify every adapter satisfies the DocumentEngine ABC.""" @@ -1306,6 +1433,7 @@ class TestAllEnginesImplementInterface: "docfold.engines.firecrawl_engine.FirecrawlEngine", "docfold.engines.chandra_engine.ChandraEngine", "docfold.engines.marker_local_engine.MarkerLocalEngine", + "docfold.engines.unlimited_ocr_engine.UnlimitedOCREngine", ]) def test_has_required_attributes(self, engine_cls_path): module_path, cls_name = engine_cls_path.rsplit(".", 1) From 496e4f534bbca846a406bfe49002c1e1481b8950 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 27 Jun 2026 09:26:00 +0000 Subject: [PATCH 2/2] Include Unlimited-OCR in benchmark.py harness Add UnlimitedOCREngine to the local-engine candidate list so it is benchmarked on GPU machines. It auto-skips (is_available() == False) where torch/transformers are not installed, so CPU runs are unaffected. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01P23UQav6r1QsPGKzogma26 --- benchmark.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/benchmark.py b/benchmark.py index e70427f..5e0738b 100644 --- a/benchmark.py +++ b/benchmark.py @@ -483,6 +483,7 @@ async def main(): from docfold.engines.pymupdf_engine import PyMuPDFEngine from docfold.engines.surya_engine import SuryaEngine from docfold.engines.tesseract_engine import TesseractEngine + from docfold.engines.unlimited_ocr_engine import UnlimitedOCREngine from docfold.engines.unstructured_engine import UnstructuredEngine # All local/open-source engines to benchmark @@ -496,6 +497,7 @@ async def main(): (MinerUEngine(), "pip install docfold[mineru]"), (MarkerLocalEngine(), "pip install marker-pdf"), (SuryaEngine(), "pip install surya-ocr"), + (UnlimitedOCREngine(), "pip install docfold[unlimited-ocr] (needs CUDA GPU)"), (DoclingEngine(), "pip install docling"), (EasyOCREngine(gpu=False), "pip install easyocr"), (NougatEngine(), "pip install nougat-ocr"),