ocr: add surya-ocr.py recipe (Surya OCR 2, 650M, structured OCR + layout/table) (#54)

Datalab's Surya OCR 2 (datalab-to/surya-ocr-2, 650M, qwen3_5) as a self-contained
UV script. Unlike the flat-markdown OCR scripts, it returns structure: a flattened
text column plus a `surya_blocks` JSON column (per-block HTML + bbox + reading order
+ confidence). `--task ocr|layout|table`, image and PDF input.

Engine: Surya normally spawns a vLLM server (Docker), which can't run in a Job. We
inject a custom in-process Backend that runs vLLM's offline LLM().chat() engine and
let Surya keep ownership of prompting, image scaling, and HTML/bbox parsing — offline
batch, no server. Runs on the pinned vllm/vllm-openai:v0.20.1 image (its qwen3_5
arch is recent/version-sensitive; vLLM lives at site-packages there, not dist-packages).

Validated on HF Jobs (l4x1): ocr 5/5, layout 3/3, table 3/3, pdf 1/1 (real arXiv PDF,
pages concatenated in text + per-page in surya_blocks). Weights are modified OpenRAIL-M.

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: davanstrien

ocr/CLAUDE.md

@@ -273,6 +273,80 @@ repo (same multi-config pattern as the other OCR scripts).
 
 Requires Python ≥3.12 (lift-pdf constraint) — fine on the standard images.
 
+### Surya OCR 2 (`surya-ocr.py`)
+✅ **OCR + layout + table validated on Jobs** (added 2026-06-22)
+
+Datalab's **Surya OCR 2** (`datalab-to/surya-ocr-2`, 650M, Qwen3.5-style) for **structured** OCR.
+Unlike the flat-markdown scripts, it returns per-block HTML + bounding boxes + reading order. The
+recipe writes **two columns**: `--output-column` (default `markdown`, flattened reading-order text)
+**and** `surya_blocks` (the full structured result as JSON, one entry per page). `--task` switches
+between `ocr` (RecognitionPredictor, full-page), `layout` (LayoutPredictor), and `table`
+(TableRecPredictor; `--table-mode full` → HTML, `simple` → rows/cols/cells).
+
+**Engine — offline vLLM batch, NO server (the whole trick).** Surya normally runs its VLM through a
+**spawned server**: on GPU it `docker run`s `vllm/vllm-openai`, on CPU a `llama-server` subprocess
+(`surya/inference/backends/{vllm,llamacpp}.py`). Docker-in-Docker isn't available inside a Job, so
+the default path can't work. Instead we subclass Surya's `Backend` ABC
+(`surya/inference/backends/base.py`: `start`/`stop`/`generate(batch)->List[BatchOutputItem]`) with an
+in-process `OfflineVLLMBackend` that runs vLLM's offline `LLM().chat()` and inject it via
+`manager.backend = ...` (bypassing `SuryaInferenceManager.__init__`'s autodetect). **Surya still owns
+everything else** — prompts (`PROMPT_MAPPING`), image scaling (`scale_to_fit`), HTML/bbox parsing, the
+repeat-loop fallback, the 0–1000→pixel bbox rescale, and the layout/table predictors — so we only swap
+the transport. We reuse Surya's own `_build_messages`/`scale_to_fit`/`PROMPT_MAPPING` so the offline
+path matches the server byte-for-byte. `mm_processor_kwargs={min_pixels:3136,max_pixels:6291456}`,
+`dtype=bfloat16`, `max_model_len=18000`, sampling `temperature=0.0/top_p=0.1`, `logprobs=1` →
+`mean_token_prob` → Surya's per-block `confidence`. Guided JSON (layout's `LAYOUT_JSON_SCHEMA`) maps to
+`StructuredOutputsParams`/`GuidedDecodingParams` (same shim as `ocr-vllm-judge.py`). **Not mirrored:**
+Surya's per-item repeat-token retry — its recognition layer already detects loops and falls back to
+layout+block OCR, so the backend stays simple (like lift).
+
+**⚠️ Image gotcha — pin `vllm/vllm-openai:v0.20.1` AND use the `site-packages` path.** Surya-2 is the
+recent, **version-sensitive, hybrid (linear-attention) `qwen3_5`** architecture; v0.20.1 is Surya's
+known-good vLLM. Unlike the other vLLM recipes (which use the unversioned image at
+`/usr/bin/python3` + `dist-packages`), the **`:v0.20.1`** image puts python at `/usr/local/bin/python3`
+and vLLM/torch at **`/usr/local/lib/python3.12/site-packages`**. The first smoke run used the old
+`dist-packages` path → `No module named 'vllm'` → 0/5. Correct flags:
+```bash
+hf jobs uv run --flavor l4x1 -s HF_TOKEN \
+    --image vllm/vllm-openai:v0.20.1 --python /usr/local/bin/python3 \
+    -e PYTHONPATH=/usr/local/lib/python3.12/site-packages \
+    ./ocr/surya-ocr.py davanstrien/ufo-ColPali OUTPUT --max-samples 5
+```
+`PYTHONPATH` is prepended ahead of the uv venv, so the **image's** torch 2.11.0+cu130 / transformers /
+vLLM 0.20.1 win at import even though `surya-ocr` pulls its own torch into the venv (harmless, just a
+wasted download). Confirmed via a probe job: vLLM at `…/site-packages/vllm`, python 3.12.13.
+
+**Naming gotcha:** must be `surya-ocr.py`, never `surya.py` (would shadow the `surya` package on
+import). Checked: no other `surya*` file in the repo.
+
+**Smoke-test results (2026-06-22, `davanstrien/ufo-ColPali`, l4x1, `vllm/vllm-openai:v0.20.1`):**
+- **ocr** (5 samples): 5/5 OK, 3.7 min (vLLM engine init ~113s incl. 34s compile + CUDA-graph capture,
+  then inference). `markdown` clean reading-order text; `surya_blocks` valid JSON with **pixel-space**
+  bboxes (e.g. `[21.6,65.5,30.9,343.4]` within `image_bbox=[0,0,618,1007]`), sequential `reading_order`,
+  canonical labels (PageHeader/SectionHeader/Text/…), `confidence` ~0.94 (logprobs path works), per-block
+  HTML (`<h1>`, `<sup>`, `<br/>`). Output `davanstrien/surya-smoke-ocr`. Resolved `vllm==0.20.1,
+  torch==2.11.0+cu130, transformers==5.7.0, surya-ocr==0.20.0`.
+- **layout** (3 samples): 3/3 OK; `surya_blocks` = `LayoutResult` per page (bboxes with `label`/
+  `position`/`count`/`confidence`, guided-JSON enforced). Output `davanstrien/surya-smoke-layout`.
+- **table** `--table-mode full` (3 samples): 3/3 OK; `TableResult` with `html` populated (rows/cols/cells
+  empty in full mode, by design). ufo-ColPali has no real tables, so use a table dataset for meaningful
+  output — the code path is what's validated. Output `davanstrien/surya-smoke-table`.
+
+- **pdf** (`--pdf-column`/`--page-range`, real 14.8MB arXiv PDF, pages 0–2): 1/1 OK. Text
+  concatenates the 3 pages (title/authors/abstract of arXiv:2606.17162 extracted in reading order);
+  `surya_blocks` has **3 page entries** (`image_bbox=[0,0,1632,2112]` at 192 DPI) with sensible labels
+  (PageHeader/SectionHeader/Text/Picture/Diagram/Caption/ListGroup/…). Source built by wrapping the PDF
+  bytes into a `Value("binary")` column. Output `davanstrien/surya-smoke-pdf`.
+
+**Still untested (low risk):** `--table-mode simple` (rows/cols/cells). Larger GPUs (l4x1 confirmed
+comfortable for 650M). A **bucket** variant (`surya-ocr-bucket.py`, mount a bucket of real PDF *files* →
+output, like `glm-ocr-bucket.py`) is a natural follow-up — distinct I/O from this dataset-column recipe.
+
+**License:** code Apache-2.0, **weights modified OpenRAIL-M** (research/personal/<$5M, no competitive use
+vs Datalab's API). Surfaced in the docstring, README entry, and output dataset card.
+
+**Benchmark/compare:** `--config`/`--create-pr` push the same multi-config pattern as the other scripts.
+
 ---
 
 ## Future: OCR Smoke Test Dataset

ocr/README.md

@@ -47,6 +47,7 @@ _Sorted by model size:_
 |--------|-------|------|---------|-------|
 | `falcon-ocr.py` | [Falcon-OCR](https://huggingface.co/tiiuae/Falcon-OCR) | 0.3B | falcon-perception | Smallest in collection. #1 on multi-column docs and tables (olmOCR), Apache 2.0 |
 | `smoldocling-ocr.py` | [SmolDocling](https://huggingface.co/ds4sd/SmolDocling-256M-preview) | 256M | Transformers | DocTags structured output |
+| `surya-ocr.py` | [Surya OCR 2](https://huggingface.co/datalab-to/surya-ocr-2) | 0.65B | vLLM | **Structured** OCR + `--task layout\|table`: per-block HTML with bboxes & reading order in an extra `surya_blocks` column. 91 langs, top-under-3B on olmOCR-Bench. Modified OpenRAIL-M license. Needs the **pinned** `vllm/vllm-openai:v0.20.1` image |
 | `glm-ocr.py` | [GLM-OCR](https://huggingface.co/zai-org/GLM-OCR) | 0.9B | vLLM | 94.62% OmniDocBench V1.5 |
 | `paddleocr-vl.py` | [PaddleOCR-VL](https://huggingface.co/PaddlePaddle/PaddleOCR-VL) | 0.9B | Transformers | 4 task modes (ocr/table/formula/chart) |
 | `paddleocr-vl-1.5.py` | [PaddleOCR-VL-1.5](https://huggingface.co/PaddlePaddle/PaddleOCR-VL-1.5) | 0.9B | Transformers | 94.5% OmniDocBench, 6 task modes |
@@ -71,6 +72,8 @@ _Sorted by model size:_
 
 **Variants & tools** (same models, different I/O): `glm-ocr-v2.py` adds checkpoint/resume for very large jobs · `glm-ocr-bucket.py` and `falcon-ocr-bucket.py` read images/PDFs from a mounted bucket and write one `.md` per page · `ocr-vllm-judge.py` runs pairwise OCR-quality comparisons.
 
+`surya-ocr.py` is the structured outlier: besides the flattened text column it writes a `surya_blocks` JSON column (per-block HTML + bounding boxes + reading order), and `--task` switches between OCR, `layout`, and `table`. It runs as **offline vLLM batch** (no server) and must use the **pinned** `vllm/vllm-openai:v0.20.1` image — its `qwen3_5` architecture is recent and version-sensitive, and that image puts vLLM at `/usr/local/lib/python3.12/site-packages` (use `--python /usr/local/bin/python3`; the exact command is in the script's docstring). Weights are **modified OpenRAIL-M**.
+
 ## Structured extraction (image or text → JSON)
 
 Most scripts here output markdown. These take a **schema** and return **structured data** instead — give them the fields you want, they fill them in:
@@ -194,6 +197,7 @@ Beyond the shared flags, some models add their own. Run `--help` on any script f
 
 | Script | Extra options |
 |--------|---------------|
+| `surya-ocr.py` | `--task ocr\|layout\|table`, `--table-mode full\|simple`, `--pdf-column`/`--page-range`, `--blocks-column` |
 | `glm-ocr.py` | `--task ocr\|formula\|table` |
 | `paddleocr-vl.py` | `--task-mode ocr\|table\|formula\|chart` |
 | `paddleocr-vl-1.5.py` | `--task-mode ocr\|table\|formula\|chart\|spotting\|seal` |