I need pull

AGENTS.md

@@ -1,109 +0,0 @@
-# For coding agents
-
-This repo is a curated collection of ready-to-run OCR scripts — each one self-contained
-via UV inline metadata, runnable over the network via `hf jobs uv run`. No clone, no
-install, no setup.
-
-## Don't rely on this doc — discover the current state
-
-This file will go stale. Prefer these sources of truth:
-
-- `hf jobs uv run --help` — job submission flags (volumes, secrets, flavors, timeouts)
-- `hf jobs hardware` — current GPU flavors and pricing
-- `hf auth whoami` — check HF token is set
-- `hf jobs ps` / `hf jobs logs <id>` — monitor running jobs
-- `ls` the repo to see which scripts actually exist (bucket variants especially)
-- [README.md](./README.md) — the table of scripts with model sizes and notes
-
-## Picking a script
-
-The [README.md](./README.md) table lists every script with model size, backend, and
-a short note. Axes that matter:
-
-- **Model size** vs accuracy vs GPU cost. Smaller = cheaper per doc.
-- **Backend**: vLLM scripts are usually fastest at scale. `transformers` and
-  `falcon-perception` are alternatives for specific models.
-- **Task support**: most scripts do plain text; some expose `--task-mode`
-  (table, formula, layout, etc.) — check the script's own docstring.
-- **Language coverage**: [`models.json`](./models.json) maps every script to its model,
-  params, backend, required image pins, and what the model card claims about languages —
-  with an `evidence` level (`per-language-benchmark` beats a bare `multilingual` tag).
-  Human-readable version: [LANGUAGES.md](./LANGUAGES.md).
-
-For the authoritative benchmark numbers on any model in the table, query the model
-card programmatically — every OCR model publishes eval results on its card:
-
-    from huggingface_hub import HfApi
-    info = HfApi().model_info("tiiuae/Falcon-OCR", expand=["evalResults"])
-    for r in info.eval_results:
-        print(r.dataset_id, r.value)
-
-See the [leaderboard data guide](https://huggingface.co/docs/hub/en/leaderboard-data-guide)
-for the full API. This is more reliable than any markdown table that might drift.
-
-## Getting help from a specific script
-
-Each script has a docstring at the top with a description and usage examples. To read it
-without downloading:
-
-    curl -s https://huggingface.co/datasets/uv-scripts/ocr/raw/main/<script>.py | head -100
-
-Or open the URL in a browser. Running `uv run <url> --help` locally may fail if the
-script has GPU-only dependencies — reading the docstring is more reliable.
-
-## The main pattern: dataset → dataset
-
-Most scripts take an input HF dataset ID and push results to an output HF dataset ID:
-
-    hf jobs uv run --flavor l4x1 -s HF_TOKEN \
-      https://huggingface.co/datasets/uv-scripts/ocr/raw/main/<script>.py \
-      <input-dataset-id> <output-dataset-id> [--max-samples N] [--shuffle]
-
-The script adds a `markdown` column to the input dataset and pushes the merged result
-to the output dataset ID on the Hub.
-
-## Alternative: directory → directory (bucket variants)
-
-A couple of scripts have `-bucket.py` variants (currently `falcon-ocr-bucket.py` and
-`glm-ocr-bucket.py`) that read from a mounted directory and write one `.md` per image
-(or per PDF page). Useful with HF Buckets via `-v`:
-
-    hf jobs uv run --flavor l4x1 -s HF_TOKEN \
-      -v hf://buckets/<user>/<input>:/input:ro \
-      -v hf://buckets/<user>/<output>:/output \
-      https://huggingface.co/datasets/uv-scripts/ocr/raw/main/<script>-bucket.py \
-      /input /output
-
-`ls` the repo to check whether a `-bucket.py` variant exists for the model you want
-before assuming it's available.
-
-The `-v` source can also be a **local directory** (`huggingface_hub` ≥ 1.22) — the CLI
-syncs it to a private bucket and mounts it, so local images/PDFs need no upload step:
-
-    hf jobs uv run --flavor l4x1 -s HF_TOKEN \
-      -v ./pages:/input \
-      -v hf://buckets/<user>/<output>:/output \
-      https://huggingface.co/datasets/uv-scripts/ocr/raw/main/<script>-bucket.py \
-      /input /output
-
-Re-runs only sync new/changed files. Local mounts are read-only by default; a local
-*output* dir works too with `:rw` — pull results back with the `hf buckets sync`
-command the CLI prints at launch.
-
-## Common flags across dataset-mode scripts
-
-Most scripts support: `--max-samples`, `--shuffle`, `--seed`, `--split`, `--image-column`,
-`--output-column`, `--private`, `--config`, `--create-pr`, `--verbose`. Read the script's
-docstring for the authoritative list — individual scripts may add model-specific options
-like `--task-mode`.
-
-## Gotchas
-
-- **Secrets**: pass `-s HF_TOKEN` to forward the user's token into the job.
-- **GPU required**: all scripts exit if CUDA isn't available. `l4x1` is the cheapest
-  GPU flavor and works for models up to ~3B. Check `hf jobs hardware` for current options.
-- **First run is slow**: model download + `torch.compile` / vLLM warmup dominates small
-  runs. Cost per doc drops sharply past a few hundred images — test with `--max-samples 10`
-  first, then scale.
-- **Don't poll jobs**: jobs run async. Submit once, check status later with
-  `hf jobs ps` or `hf jobs logs <id>`.

CLAUDE.md

@@ -1,288 +1,213 @@
-# OCR Scripts — Development Notes
+# OCR Scripts - Development Notes
 
-Dev notes for the `ocr/` recipes: the **conventions** to follow, the **per-script gotchas** (the
-"why" behind each script's quirks), and the **internal tooling**. Runnable examples live in each
-script's docstring and `README.md`; benchmark result tables live in `OCR-BENCHMARK.md`.
+## Active Scripts
 
-- [Conventions & invariants](#conventions--invariants) — read before adding/changing a recipe
-- [Script status](#script-status) · [Per-script gotchas](#per-script-gotchas)
-- [Internal tooling](#internal-tooling) · [Deferred / tracked](#deferred--tracked) · [Change log](#change-log)
+### DeepSeek-OCR v1 (`deepseek-ocr-vllm.py`)
+✅ **Production Ready**
+- Fully supported by vLLM
+- Fast batch processing
+- Tested and working on HF Jobs
 
----
-
-## Conventions & invariants
-
-Read this before adding or changing a recipe. Each rule maps to a failure we've actually hit; the
-planned **self-review skill** (see [Deferred](#deferred--tracked)) just enforces this list.
-
-- **Catalog entry.** Adding or changing a recipe means updating [`models.json`](./models.json)
-  (script → model, params, backend, image pins, language claim) and, if the language claim
-  changed, [LANGUAGES.md](./LANGUAGES.md). Language fields record what the **model card claims**
-  (with an evidence level), never inferred coverage. Hand-maintained for now; if drift becomes a
-  problem, the follow-up is generating the README table from the JSON.
-- **Self-contained single file.** Each recipe is one PEP 723 UV script runnable from a raw URL
-  (`hf jobs uv run <url>`). No shared *importable local* module (the job env only gets the one file).
-  Extra pip deps are fine — **pin them**. A heavy/stable/shared subsystem may become an opt-in *package*
-  dep (e.g. bucket I/O → `bucketbag`, [#67](https://github.com/davanstrien/uv-scripts-for-ai/issues/67)),
-  never a local import. Recipes = inline; internal tooling may share freely.
-- **GPU + output.** Check `torch.cuda.is_available()` and exit clearly if absent. Write to the Hub
-  (`push_to_hub`) or a bucket (`-v hf://…`), never bare local paths (Jobs disk is ephemeral; local *input*
-  dirs mount via `-v ./dir:/mnt`, synced to a bucket automatically).
-- **vLLM image + fail-fast preflight.** If the model's arch isn't in a stable vLLM wheel, **omit
-  `vllm`/`torch` from deps** and run on the pinned `vllm/vllm-openai` image via
-  `--image … --python … -e PYTHONPATH=…`; add a preflight that `sys.exit(1)`s naming the exact flags
-  (surya-class — see gotchas). Pinned-image scripts: `surya-ocr` (`:v0.20.1`, **site-packages**),
-  `nanonets-ocr2` (`:v0.10.2`), `unlimited-ocr` (`:unlimited-ocr`), `deepseek-ocr2`/`glm-ocr` (nightly).
-- **Pins are temporary.** An image/version pin (`:v0.10.2`, `:v0.20.1`, a nightly, `surya-ocr==0.20.0`,
-  …) is a workaround for a *current* ecosystem gap — a decode regression, an arch not yet in a stable
-  wheel, a resolver backtrack. In the recipe, record **why** the pin exists and **what would loosen it**
-  (e.g. "move back to the default image when a newer vLLM ships a Qwen2.5-VL decode fix"; "drop the
-  nightly once the arch lands in a stable release"). Re-test periodically and relax when the gap closes —
-  that's what the `bump-vllm-pins` skill is for; prefer floors over exact pins once you can.
-- **Env guards on the bare image.** Set `VLLM_USE_FLASHINFER_SAMPLER=0` (and `VLLM_USE_DEEP_GEMM=0`
-  for nightly vLLM) **before** importing `vllm` — both JIT paths need `nvcc`, which the bare uv image lacks.
-- **Context-length invariant.** `--max-tokens` ≤ `--max-model-len` ≤ the model's real max context
-  (`config.json` `max_position_embeddings`; VLMs → `text_config`, mind `rope_scaling`). vLLM refuses to
-  start if `max_model_len` is over (we don't set `VLLM_ALLOW_LONG_MAX_MODEL_LEN`); output can't fit if
-  `max_tokens` > `max_model_len`. Check:
-  `curl -s https://huggingface.co/<model>/raw/main/config.json | python -c "import json,sys;c=json.load(sys.stdin);t=c.get('text_config',c);print(t.get('max_position_embeddings'),t.get('rope_scaling'))"`.
-- **Output-column collision guard.** Every recipe that adds an output column calls
-  `ensure_output_columns_free(dataset, [cols], overwrite)` (or the inline sink guard for pp-*) so it
-  fails fast instead of duplicating/clobbering an input column; `--overwrite` opts into replacing it.
-  Default `--output-column` is `markdown` (never a bare `text`). ([#66](https://github.com/davanstrien/uv-scripts-for-ai/pull/66))
-- **Bound large images.** Full-page recipes cap input pixels / resize, or size `max_model_len` to fit —
-  a 7–9 MP page is ~14k image tokens. Prefer bounding the input (deterministic) over a giant context;
-  don't auto-size `max_model_len` from images (it's fixed at engine init, before images are seen).
-- **Dep sanity.** No stale version *caps* that drag a transitive lib back — e.g. `pyarrow<18` forced an
-  old `datasets` lacking the `Json` feature → `load_dataset` crashed (the glm-ocr bug). Use floors, not ceilings.
-- **Error signalling (known gap).** Scripts currently write sentinels (`[OCR ERROR]`, `[SURYA GENERATE
-  ERROR]`) *into* the output column, so partial failures are silent. A companion `ocr_error` status
-  column is the deferred fix (see [Deferred](#deferred--tracked)).
+### LightOnOCR-2-1B (`lighton-ocr2.py`)
+✅ **Production Ready** (Fixed 2026-01-29)
 
----
+**Status:** Working with vLLM nightly
 
-## Script status
-
-Legend: ✅ production-ready · ⚠️ works only with a required pinned image · 🧪 experimental/on-hold.
-"+image" = needs a `--image vllm/vllm-openai:<tag>` override (not the default uv image).
-
-| Script | | Backend | Flavor | Note |
-|--------|--|---------|--------|------|
-| `deepseek-ocr-vllm.py` | ✅ | vLLM (stable) | l4x1 | `NGramPerReqLogitsProcessor` anti-repeat |
-| `deepseek-ocr2-vllm.py` | ✅ | vLLM (nightly) | l4x1 | arch needs nightly; `addict`+`matplotlib` deps |
-| `lighton-ocr2.py` | ✅ | vLLM | a100-large / l4x1 | resize 1540px; `--max-model-len` 16384 |
-| `paddleocr-vl-1.5.py` | ✅ | transformers | l4x1 | not vLLM (server-only upstream); single-image |
-| `paddleocr-vl-1.6.py` | ✅ | vLLM | l4x1 | smart-resize ~1M px; SOTA OmniDocBench |
-| `paddleocr-vl.py` | ✅ | vLLM | l4x1 | |
-| `dots-ocr.py` | ✅ | vLLM (stable) | l4x1 | `--max-model-len` 32768; no internal resize |
-| `dots-ocr-1.5.py` | ✅ | vLLM 0.17.1 | l4x1 | see gotcha (`content_format`, mirror, bbox space) |
-| `glm-ocr.py` | ✅ | vLLM (nightly) | l4x1 | `VLLM_USE_DEEP_GEMM=0`; no pyarrow cap |
-| `ovis-ocr2.py` | ✅ | vLLM (stable ≥0.22.1) | l4x1 / a10g-small | `gdn_prefill_backend="triton"` (card); card-exact prompt via `enable_thinking=False` template + `llm.generate`; `<img>` region tags filtered by default (`--keep-image-tags`) |
-| `glm-ocr-v2.py` | 🧪 | vLLM (nightly) | l4x1 | CommitScheduler incremental — on hold (see Deferred) |
-| `nanonets-ocr.py` | ✅ | vLLM | a10g-small | `--max-model-len` 32768 (`--max-tokens` 15000) |
-| `nanonets-ocr2.py` | ⚠️+image | vLLM `:v0.10.2` | a10g-small | Qwen2.5-VL ≥0.11 regression → pure `!` |
-| `unlimited-ocr-vllm.py` | ✅+image | vLLM `:unlimited-ocr` | l4x1 | single-image; multi-page → serve |
-| `surya-ocr.py` | ✅+image | vLLM `:v0.20.1` | l4x1 | offline backend inject; site-packages PYTHONPATH |
-| `surya-ocr-bucket.py` | ✅+image | vLLM `:v0.20.1` | l4x1 | bucket I/O; pin `surya-ocr==0.20.0` |
-| `lift-extract.py` | ✅ | hf / vLLM | a100-large | schema-constrained extraction; naming gotcha |
-| `nuextract3.py`, `lfm2-extract.py`, `lfm2-vl-extract.py` | ✅ | vLLM | l4x1 | structured extraction |
-| `hunyuan-ocr.py` | ✅ | vLLM | l4x1 | **1.0, revision-pinned** (root repo became 1.5 in-place); `transformers<5.13` cap — see gotcha |
-| `hunyuan-ocr-1.5.py` | ✅ | vLLM | l4x1 | tracks repo root (=1.5); task-locked prompts; `transformers<5.13` cap — see gotcha |
-| `rolm-ocr.py`, `smoldocling-ocr.py`, `numarkdown-ocr.py`, `qianfan-ocr.py`, `firered-ocr.py`, `abot-ocr.py`, `falcon-ocr.py`, `olmocr2-vllm.py`, `dots-mocr.py` | ✅ | vLLM | varies | see `README.md` for flags |
-| `pp-ocrv6.py`, `pp-doclayout.py` | ✅ | PaddleOCR / PaddleX | l4x1 | classical det+rec; dataset **or** bucket I/O |
-
-**License note:** Surya and `lift` ship code as Apache-2.0 but **weights under a modified OpenRAIL-M**
-(research/personal/<$5M, no competitive use vs Datalab's API) — surfaced in each docstring + card.
-HunyuanOCR (1.0 + 1.5) is under the [Tencent Hunyuan Community License](https://huggingface.co/tencent/HunyuanOCR/blob/main/LICENSE)
-(territory excludes EU/UK/South Korea; standard across Tencent's Hunyuan releases) — surfaced in each docstring + card.
+**What was fixed:**
+- Root cause was NOT vLLM - it was the deprecated `HF_HUB_ENABLE_HF_TRANSFER=1` env var
+- The script was setting this env var but `hf_transfer` package no longer exists
+- This caused download failures that manifested as "Can't load image processor" errors
+- Fix: Removed the `HF_HUB_ENABLE_HF_TRANSFER=1` setting from the script
 
----
+**Test results (2026-01-29):**
+- 10/10 samples processed successfully
+- Clean markdown output with proper headers and paragraphs
+- Output dataset: `davanstrien/lighton-ocr2-test-v4`
 
-## Per-script gotchas
-
-Only the scripts with load-bearing quirks; the rest are unremarkable (`README.md` covers flags).
-
-### `surya-ocr.py` / `surya-ocr-bucket.py` — pinned image + site-packages path
-Surya-2 (`datalab-to/surya-ocr-2`, 650M, `qwen3_5`) needs its **known-good** vLLM build, and the
-`:v0.20.1` image puts python at `/usr/local/bin/python3` and libs at
-**`/usr/local/lib/python3.12/site-packages`** (not the usual `dist-packages`) — wrong path →
-`No module named 'vllm'` → 0/5. It can't Docker-in-Docker Surya's normal server, so it **injects an
-in-process `OfflineVLLMBackend`** into `SuryaInferenceManager` (subclassing Surya's `Backend` ABC) and
-reuses Surya's own prompts/`scale_to_fit`/HTML+bbox parsing so the offline path matches the server.
-`mm_processor_kwargs={min_pixels:3136, max_pixels:6291456}`, `max_model_len=18000`, `logprobs=1` →
-per-block `confidence`. Writes two columns (`--output-column` markdown + `surya_blocks` JSON). Never
-name the file `surya.py` (shadows the package). The recipe now **fails fast** if `vllm` isn't importable,
-naming the required flags. **Bucket variant:** pin `surya-ocr==0.20.0` (loosening it, or adding
-`huggingface-hub>=1.6.0`, lets uv backtrack to a surya without `surya.inference`); **copy beats mount**
-for bucket reads (FUSE `rglob` is ~26× slower on a 38k-file bucket; mount also hit a transient CSI flake);
-`.jp2` via an `imagecodecs` fallback (Pillow lacks OpenJPEG); resume-by-skip on the output `.json`.
-
-### `nanonets-ocr2.py` — pinned `:v0.10.2` image
-Nanonets-OCR2-3B is **Qwen2.5-VL**, which has a vLLM **≥0.11 decode regression** (outputs pure `!` on
-every page) — [vllm#27775](https://github.com/vllm-project/vllm/issues/27775). 0.9.2/0.10.1/**0.10.2**
-are known-good. Not context length (still `!` at 32768) and not torch.compile. Pip-pinning `vllm==0.10.2`
-clashes with modern `transformers` (old tokenizer API), so run on the **`:v0.10.2` image** (ships a
-consistent vLLM 0.10.2 + transformers 4.56.1); `vllm`/`torch` omitted from deps. `--max-model-len` 32768
-(the 15000 `--max-tokens` can't fit 8192). Re-test the default image when a newer vLLM ships a Qwen2.5-VL
-decode fix.
-
-### `dots-ocr-1.5.py` — `content_format="string"` + resized-bbox space
-Must pass `chat_template_content_format="string"` to `llm.chat()` — the model's `tokenizer_config.json`
-template expects string content; without it you get ~1 token then EOS (empty output). The v1.5 weights
-aren't on HF from the authors — **mirrored to `davanstrien/dots.ocr-1.5`** from ModelScope (MIT-based).
-Layout bboxes are in the **resized** image space (`Qwen2VLImageProcessor.smart_resize`,
-`max_pixels=11,289,600`, `factor=28`); map back with:
-```python
-import math
-def smart_resize(h, w, factor=28, min_pixels=3136, max_pixels=11289600):
-    h_bar, w_bar = max(factor, round(h/factor)*factor), max(factor, round(w/factor)*factor)
-    if h_bar*w_bar > max_pixels:
-        beta = math.sqrt((h*w)/max_pixels); h_bar, w_bar = math.floor(h/beta/factor)*factor, math.floor(w/beta/factor)*factor
-    elif h_bar*w_bar < min_pixels:
-        beta = math.sqrt(min_pixels/(h*w)); h_bar, w_bar = math.ceil(h*beta/factor)*factor, math.ceil(w*beta/factor)*factor
-    return h_bar, w_bar
-# orig_x = bbox_x * (orig_w / w_bar);  orig_y = bbox_y * (orig_h / h_bar)
+**Example usage:**
+```bash
+hf jobs uv run --flavor a100-large \
+    -s HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/lighton-ocr2.py \
+    davanstrien/ufo-ColPali output-dataset \
+    --max-samples 10 --shuffle --seed 42
 ```
-(Same `smart_resize`/`max_pixels=11.29M` applies to `dots-ocr.py` v1's processor cap — but the
-`content_format="string"` fix does **not**: v1 works with the auto-detected `openai` chat format.)
-
-### `unlimited-ocr-vllm.py` — dedicated image, single-image batch
-Baidu `baidu/Unlimited-OCR` (3.3B, DeepSeek-OCR descendant); arch is in **no stable vLLM wheel**, so it
-runs on **`vllm/vllm-openai:unlimited-ocr`** (`:unlimited-ocr-cu129` on Hopper), standard `/usr/bin/python3`
-+ `dist-packages`. `NGramPerReqLogitsProcessor` (re-exported via `unlimited_ocr`), prompt
-`<image>document parsing.`, `limit_mm_per_prompt={"image":1}`, `--strip-grounding` drops `<|det|>`/`<|ref|>`.
-**Batch recipe stays single-image**; multi-page is finicky offline (one `<image>` per page; degrades on
-hard scans) and belongs to **serving** — both engines read clean multi-page docs, but **SGLang is more
-robust on hard/degraded scans**. Serving setup (SGLang pin `lmsysorg/sglang:v0.5.10.post1`, a100+flashinfer
-— HF `h200` nodes fail with `CUDA error 802`) is in `serving-unlimited-ocr.md`.
-
-### `lift-extract.py` — naming + backends
-Datalab `lift` (9B, Qwen3.5), schema-constrained image/PDF → JSON; the only recipe ingesting PDFs directly.
-**Must not be named `lift.py`** (shadows the installed `lift` package → ImportError). Two in-process backends
-via `--method`: `hf` (default image, plain `model.generate`) and `vllm` (needs the `vllm/vllm-openai` image;
-reproduces lift's own recipe: `mm_processor_kwargs={min_pixels:3136,max_pixels:861696}`, guided JSON schema,
-`temperature=0.0,top_p=0.1,max_tokens=12384`). Pin `--model datalab-to/lift` via the `MODEL_CHECKPOINT` env
-(settings read env at import).
-
-### `deepseek-ocr-vllm.py` / `deepseek-ocr2-vllm.py`
-v1 uses the official offline pattern (`llm.generate()` + `NGramPerReqLogitsProcessor` for repetition).
-**Known bug** (hit on vLLM *nightly*, 2026-02-12; unverified on the stable wheels v1 now resolves):
-some aspect ratios trip `images_crop dim[2] expected 1024, got 640` (gundam-mode
-default vs a validator expecting 1024²) — hit 2/10 on `ufo-ColPali`, aspect-ratio dependent, no upstream
-issue filed ([vllm#28160](https://github.com/vllm-project/vllm/issues/28160) is the related request). v2
-needs **nightly** vLLM (`DeepseekOCR2ForCausalLM` not in stable) + `addict`/`matplotlib` (its HF custom
-code), plus `limit_mm_per_prompt={"image":1}`.
-
-### `hunyuan-ocr.py` / `hunyuan-ocr-1.5.py` — upstream replaced the repo root in-place
-On 2026-07-06 Tencent pushed HunyuanOCR-1.5 **into the same repo** `tencent/HunyuanOCR` (1.5 at root,
-1.0 archived under `v1.0/`, DFlash draft under `dflash/`, **no 1.0 tag or branch**). vLLM can't load a
-repo subfolder, so `hunyuan-ocr.py` pins the last 1.0 commit by `revision` (`f6af82ee…`) — that pin is
-the 1.0 *identity*, never loosen it to `main`; 1.5 is its own recipe (`hunyuan-ocr-1.5.py`, tracks root,
-has `--revision` as insurance against the next in-place swap). Separate breakage, both scripts: stable
-vLLM ≤0.24.0's `hunyuan_vl_image.py` does a string-key `AutoImageProcessor.register(...)` which
-transformers 5.13 rejects (`'str' object has no attribute '__module__'` → "architectures failed to be
-inspected" at engine init) — hence the `transformers<5.13` cap in both; drop it when the vllm#47872 fix
-ships in a stable wheel. 1.5 prompts are task-locked (12 types, Chinese wording, from the official
-client's `hunyuan_tasks.py`) and sampling is card-locked (temp 0.0, rep-penalty 1.08) — don't
-"improve" either; upstream observed hand-tweaked prompts silently degrade quality.
-
-### `glm-ocr.py`
-Chatty on blank pages / can emit degenerate repeats — that's **model quality, not a crash**; don't
-re-debug it as a recipe bug. (The actual historical crash was the `pyarrow<18` cap — see Conventions.)
-
-### `lighton-ocr2.py`
-The original breakage was **not** vLLM — it was a dead `HF_HUB_ENABLE_HF_TRANSFER=1` (the `hf_transfer`
-package is gone), which surfaced as "Can't load image processor". Removed. Pixtral ViT + Qwen3, RLVR-trained,
-resize 1540px @200 DPI, `--max-model-len` 16384. `paddleocr-vl-1.5.py` uses the **transformers** backend
-(single-image) because PaddleOCR-VL only supports vLLM in server mode.
 
----
+**Model Info:**
+- Model: `lightonai/LightOnOCR-2-1B`
+- Architecture: Pixtral ViT encoder + Qwen3 LLM
+- Training: RLVR (Reinforcement Learning with Verifiable Rewards)
+- Performance: 83.2% on OlmOCR-Bench, 42.8 pages/sec on H100
+
+### PaddleOCR-VL-1.5 (`paddleocr-vl-1.5.py`)
+✅ **Production Ready** (Added 2026-01-30)
+
+**Status:** Working with transformers
 
-## Internal tooling
+**Note:** Uses transformers backend (not vLLM) because PaddleOCR-VL only supports vLLM in server mode, which doesn't fit the single-command UV script pattern. Images are processed one at a time for stability.
 
-Not user recipes — benchmark/eval infra. **Result tables + validation history → `OCR-BENCHMARK.md`.**
+**Test results (2026-01-30):**
+- 10/10 samples processed successfully
+- Processing time: ~50s per image on L4 GPU
+- Output dataset: `davanstrien/paddleocr-vl15-final-test`
 
-### `ocr-bench-run.py` — coordinator
-Launches N OCR models on the same dataset, each pushing to a shared repo as a separate config via
-`--config/--create-pr`. Eval separately with `ocr-vllm-judge.py` / `ocr-elo-bench.py`. Registry (4 models):
-`glm-ocr`, `deepseek-ocr` (auto `--prompt-mode free`), `lighton-ocr-2`, `dots-ocr`; each has a `default_args`.
+**Example usage:**
 ```bash
-uv run ocr-bench-run.py source-dataset --output my-bench --max-samples 50
-uv run ocr-bench-run.py --list-models              # registry table
-uv run ocr-bench-run.py ... --models glm-ocr dots-ocr --dry-run
-uv run ocr-elo-bench.py my-bench --from-prs --mode both   # eval from PRs, no merge
+hf jobs uv run --flavor l4x1 \
+    -s HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/paddleocr-vl-1.5.py \
+    davanstrien/ufo-ColPali output-dataset \
+    --max-samples 10 --shuffle --seed 42
 ```
 
-### `ocr-vllm-judge.py` — offline vLLM jury judge
-Pairwise OCR-quality comparisons via vLLM offline `LLM()`; jury mode (multiple models vote, majority
-aggregation) with **0 parse failures** (structured output via the `StructuredOutputsParams` → `GuidedDecodingParams`
-→ prompt shim). Prefer over the API judge (`ocr-jury-bench.py`) for batch eval — no rate limits, reproducible.
-`--from-prs` loads configs from open PRs without merging; `--save-results REPO` persists comparisons/leaderboard/metadata.
-A100 recommended for jury mode; L4 works for a single 7B judge.
-```bash
-hf jobs uv run --flavor a100-large -s HF_TOKEN ocr-vllm-judge.py davanstrien/ocr-bench-nls-50 --from-prs \
-    --judge-model Qwen/Qwen2.5-VL-7B-Instruct --judge-model Qwen/Qwen3-VL-8B-Instruct --max-samples 50
+**Task modes:**
+- `ocr` (default): General text extraction to markdown
+- `table`: Table extraction to HTML format
+- `formula`: Mathematical formula recognition to LaTeX
+- `chart`: Chart and diagram analysis
+- `spotting`: Text spotting with localization (uses higher resolution)
+- `seal`: Seal and stamp recognition
+
+**Model Info:**
+- Model: `PaddlePaddle/PaddleOCR-VL-1.5`
+- Size: 0.9B parameters (ultra-compact)
+- Performance: 94.5% SOTA on OmniDocBench v1.5
+- Backend: Transformers (single image processing)
+- Requires: `transformers>=5.0.0`
+
+## Pending Development
+
+### DeepSeek-OCR-2 (Visual Causal Flow Architecture)
+
+**Status:** ⏳ Waiting for vLLM upstream support
+
+**Context:**
+DeepSeek-OCR-2 is the next generation OCR model (3B parameters) with Visual Causal Flow architecture offering improved quality. We attempted to create a UV script (`deepseek-ocr2-vllm.py`) but encountered a blocker.
+
+**Blocker:**
+vLLM does not yet support `DeepseekOCR2ForCausalLM` architecture in the official release.
+
+**PR to Watch:**
+🔗 https://github.com/vllm-project/vllm/pull/33165
+
+This PR adds DeepSeek-OCR-2 support but is currently:
+- ⚠️ **Open** (not merged)
+- Has unresolved review comments
+- Pre-commit checks failing
+- Issues: hardcoded parameters, device mismatch bugs, missing error handling
+
+**What's Needed:**
+1. PR #33165 needs to be reviewed, fixed, and merged
+2. vLLM needs to release a version including the merge
+3. Then we can add these dependencies to our script:
+   ```python
+   # dependencies = [
+   #     "datasets>=4.0.0",
+   #     "huggingface-hub",
+   #     "pillow",
+   #     "vllm",
+   #     "tqdm",
+   #     "toolz",
+   #     "torch",
+   #     "addict",
+   #     "matplotlib",
+   # ]
+   ```
+
+**Implementation Progress:**
+- ✅ Created `deepseek-ocr2-vllm.py` script
+- ✅ Fixed dependency issues (pyarrow, datasets>=4.0.0)
+- ✅ Tested script structure on HF Jobs
+- ❌ Blocked: vLLM doesn't recognize architecture
+
+**Partial Implementation:**
+The file `deepseek-ocr2-vllm.py` exists in this repo but is **not functional** until vLLM support lands. Consider it a draft.
+
+**Testing Evidence:**
+When we ran on HF Jobs, we got:
+```
+ValidationError: Model architectures ['DeepseekOCR2ForCausalLM'] are not supported for now.
+Supported architectures: [...'DeepseekOCRForCausalLM'...]
 ```
 
-### `ocr-human-eval.py` — blind human A/B
-Gradio app for blind A/B with ELO + optional agreement-vs-judge analysis (split-jury comparisons shown
-first; round-robin image variety). Resume-safe (atomic JSON per vote).
-```bash
-uv run ocr-human-eval.py davanstrien/ocr-bench-rubenstein --from-prs \
-    --judge-results davanstrien/ocr-bench-rubenstein-judge --max-samples 5
+**Next Steps (when PR merges):**
+1. Update `deepseek-ocr2-vllm.py` dependencies to include `addict` and `matplotlib`
+2. Test on HF Jobs with small dataset (10 samples)
+3. Verify output quality
+4. Update README.md with DeepSeek-OCR-2 section
+5. Document v1 vs v2 differences
+
+**Alternative Approaches (if urgent):**
+- Create transformers-based script (slower, no vLLM batching)
+- Use DeepSeek's official repo setup (complex, not UV-script compatible)
+
+**Model Information:**
+- Model ID: `deepseek-ai/DeepSeek-OCR-2`
+- Model Card: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2
+- GitHub: https://github.com/deepseek-ai/DeepSeek-OCR-2
+- Parameters: 3B
+- Resolution: (0-6)×768×768 + 1×1024×1024 patches
+- Key improvement: Visual Causal Flow architecture
+
+**Resolution Modes (for v2):**
+```python
+RESOLUTION_MODES = {
+    "tiny": {"base_size": 512, "image_size": 512, "crop_mode": False},
+    "small": {"base_size": 640, "image_size": 640, "crop_mode": False},
+    "base": {"base_size": 1024, "image_size": 768, "crop_mode": False},  # v2 optimized
+    "large": {"base_size": 1280, "image_size": 1024, "crop_mode": False},
+    "gundam": {"base_size": 1024, "image_size": 768, "crop_mode": True},  # v2 optimized
+}
 ```
 
-**Validation headline** (full tables in `OCR-BENCHMARK.md`): DeepSeek-OCR is consistently #1 across
-datasets and eval methods; middle-pack rankings are dataset-dependent; a jury of small models gives 0
-parse failures; **Kimi K2.5 (170B)** is the only judge matching the human's #1 (small judges overrate
-LightOnOCR-2's commentary style).
+## Other OCR Scripts
+
+### Nanonets OCR (`nanonets-ocr.py`, `nanonets-ocr2.py`)
+✅ Both versions working
+
+### PaddleOCR-VL (`paddleocr-vl.py`)
+✅ Working
 
 ---
 
-## Deferred / tracked
-
-- **bucketbag adoption** — evaluate adopting `bucketbag` for the bucket recipes (slim recipes / harden
-  bucketbag / find other beneficiaries) → [#67](https://github.com/davanstrien/uv-scripts-for-ai/issues/67).
-- **Self-review skill** (spark) — a dev-only skill (sibling to `bump-vllm-pins`) that reviews a recipe/diff
-  against the [Conventions](#conventions--invariants) block: context-length, collision guard, vLLM-image +
-  preflight, env guards, dep sanity, image bounding, optional Jobs smoke. Enforces that list; build after
-  the current work.
-- **Error-signalling** — companion `ocr_error` status column (null cell + truncated exception) instead of
-  sentinels in the output column, so "read nothing" ≠ "run errored". Touches ~all recipes; deferred.
-- **OCR smoke-test dataset** — a tiny curated set (~20–30 images across doc-type/quality/language/layout,
-  ground truth where possible) for fast CI-style regression checks after dep bumps. Pairs with the skill.
-- **Multi-page batch (Unlimited-OCR)** — an SGLang-server-in-job recipe for robust multi-page at scale
-  (single-image vLLM stays the batch default). Gated on a real corpus-scale need + the `h200`/`fa3` infra
-  fix; see `serving-unlimited-ocr.md`.
-- **ALTO XML export** — from `surya_blocks` (block-level bbox→`HPOS/VPOS/…`, label→`TextBlock`/`Illustration`);
-  the surya-ocr-bucket test bucket ships CA's own ALTO `.xml` as a diff target.
-- **Incremental uploads** — superseded by HF Buckets / `bucketbag` ([#67](https://github.com/davanstrien/uv-scripts-for-ai/issues/67));
-  `glm-ocr-v2.py` keeps the older CommitScheduler resume path for very large jobs today (do not port it — on hold).
-- **Leaderboard Space** — public ELO/pointwise view fed by the benchmark datasets. Idea only.
-- **MonkeyOCRv2 — evaluated, not added** (2026-07-14). `zenosai/MonkeyOCRv2-S-Parsing` (0.6B) needs the
-  [GitHub repo's](https://github.com/Yuliang-Liu/MonkeyOCRv2) custom multi-stage pipeline (`parsing/parse.py`,
-  structure detection) — not pip-packaged, pins `vllm==0.11.2`, and the card text says "academic research and
-  non-commercial use only" while the metadata tag says `apache-2.0` (conflicting license signals). Revisit if
-  it gets packaged / vLLM-native support or the license is clarified.
-
-**Watch:** `deepseek-ocr2` / `glm-ocr` stay on nightly vLLM until their arch lands in a stable release.
-The nightly index (`https://wheels.vllm.ai/nightly`) occasionally has transient build issues (e.g. only
-ARM wheels) — if a nightly-recipe install fails on resolution, wait and retry before debugging the recipe.
+## Future: OCR Smoke Test Dataset
+
+**Status:** Idea (noted 2026-02-12)
+
+Build a small curated dataset (`uv-scripts/ocr-smoke-test`?) with ~2-5 samples from diverse sources. Purpose: fast CI-style verification that scripts still work after dep updates, without downloading full datasets.
+
+**Design goals:**
+- Tiny (~20-30 images total) so download is seconds not minutes
+- Covers the axes that break things: document type, image quality, language, layout complexity
+- Has ground truth text where possible for quality regression checks
+- All permissively licensed (CC0/CC-BY preferred)
+
+**Candidate sources:**
+
+| Source | What it covers | Why |
+|--------|---------------|-----|
+| `NationalLibraryOfScotland/medical-history-of-british-india` | Historical English, degraded scans | Has hand-corrected `text` column for comparison. CC0. Already tested with GLM-OCR. |
+| `davanstrien/ufo-ColPali` | Mixed modern documents | Already used as our go-to test set. Varied layouts. |
+| Something with **tables** | Structured data extraction | Tests `--task table` modes. Maybe a financial report or census page. |
+| Something with **formulas/LaTeX** | Math notation | Tests `--task formula`. arXiv pages or textbook scans. |
+| Something **multilingual** (CJK, Arabic, etc.) | Non-Latin scripts | GLM-OCR claims zh/ja/ko support. Good to verify. |
+| Something **handwritten** | Handwriting recognition | Edge case that reveals model limits. |
+
+**How it would work:**
+```bash
+# Quick smoke test for any script
+uv run glm-ocr.py uv-scripts/ocr-smoke-test smoke-out --max-samples 5
+# Or a dedicated test runner that checks all scripts against it
+```
+
+**Open questions:**
+- Build as a proper HF dataset, or just a folder of images in the repo?
+- Should we include expected output for regression testing (fragile if models change)?
+- Could we add a `--smoke-test` flag to each script that auto-uses this dataset?
+- Worth adding to HF Jobs scheduled runs for ongoing monitoring?
 
 ---
 
-## Change log
-
-- **2026-07-14** — added `ovis-ocr2.py` (`ATH-MaaS/OvisOCR2`, 0.9B Qwen3.5, 96.58 OmniDocBench v1.6,
-  Apache-2.0; stable vLLM ≥0.22.1, `gdn_prefill_backend="triton"`, card-exact prompt/postprocessing).
-  Smoke-tested green on the default uv image, a10g-small, resolved vLLM 0.25.1 (5/5 pages, tags filtered,
-  stamp OK). Evaluated MonkeyOCRv2 alongside it — deferred (see Deferred / tracked).
-- **2026-07-08** — HunyuanOCR upstream repo swap: pinned `hunyuan-ocr.py` to the last 1.0 revision +
-  added `hunyuan-ocr-1.5.py` (12 task types, locked sampling); `transformers<5.13` cap in both for the
-  stable-vLLM HunyuanVL register breakage (vllm#47872). See the hunyuan gotcha above.
-- **2026-07-01** — large full-page scan fixes ([#65](https://github.com/davanstrien/uv-scripts-for-ai/pull/65)):
-  surya vLLM-missing preflight; `dots` 8192→32768 + `--max-pixels`; `lighton-ocr2` 8192→16384; `glm` dropped
-  `pyarrow<18` (→ `datasets` `Json` load crash) + `VLLM_USE_DEEP_GEMM=0` + `--max-pixels`; `pp-ocrv6`
-  `--output-column` + collision guard. Then the output-column collision-guard + `--overwrite` sweep across
-  the recipes ([#66](https://github.com/davanstrien/uv-scripts-for-ai/pull/66)); `nanonets-ocr` 8192→32768.
-- **Earlier** — per-script fixes are recorded in git history + the gotchas above; benchmark runs in `OCR-BENCHMARK.md`.
+**Last Updated:** 2026-02-12
+**Watch PRs:**
+- DeepSeek-OCR-2: https://github.com/vllm-project/vllm/pull/33165

LANGUAGES.md

@@ -1,38 +0,0 @@
-# Language support
-
-What each model's own card claims about language coverage — compiled from the model cards as of **2026-07-15**. Treat these as **claims, not guarantees**: only one model (Surya) publishes per-language scores, and a "multilingual" tag tells you very little. Machine-readable version (plus params, backend, image pins): [`models.json`](models.json).
-
-The fastest way to find out whether a model handles *your* language is to run it on a few of your own pages — `--max-samples 10` costs a few cents on a T4/L4:
-
-```bash
-hf jobs uv run --flavor l4x1 -s HF_TOKEN \
-    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/surya-ocr.py \
-    your-dataset your-output --max-samples 10
-```
-
-_Sorted by strength of evidence, then coverage:_
-
-| Model | Card claim | Evidence |
-|-------|-----------|----------|
-| [Surya OCR 2](https://huggingface.co/datalab-to/surya-ocr-2) | **91 languages** | The only **per-language benchmark** published ([full table](https://github.com/datalab-to/surya/blob/master/static/docs/multilingual.md)): 38/91 score ≥90%, 76/91 ≥80%. Weakest of the top-15: Arabic 72.7%, Vietnamese 73.2% |
-| [dots.ocr](https://huggingface.co/rednote-hilab/dots.ocr) | 100 languages | Explicit **low-resource** claim, backed by an in-house benchmark (1,493 PDFs across 100 languages) — aggregate scores only, not per-language |
-| [Tesseract 5](https://github.com/tesseract-ocr/tesseract) | 125 traineddata packs (~100+ languages + script models) | Fully **enumerable** ([tessdata_best](https://github.com/tesseract-ocr/tessdata_best)) — the broadest *named* coverage here, incl. many low-resource languages; script models (Latin, Cyrillic, Devanagari, …) cover languages without a dedicated pack |
-| [Qianfan-OCR](https://huggingface.co/baidu/Qianfan-OCR) | 192 languages | Headline claim; no list or per-language numbers |
-| [PaddleOCR-VL](https://huggingface.co/PaddlePaddle/PaddleOCR-VL) / [1.5](https://huggingface.co/PaddlePaddle/PaddleOCR-VL-1.5) / [1.6](https://huggingface.co/PaddlePaddle/PaddleOCR-VL-1.6) | 109 languages; 1.5 adds Tibetan + Bengali | Names scripts (Cyrillic, Arabic, Devanagari, Thai) and claims handwriting + historical documents; no per-language numbers |
-| [PP-OCRv6](https://huggingface.co/collections/PaddlePaddle/pp-ocrv6) | 48 languages | Official card claim |
-| [Nanonets-OCR2-3B](https://huggingface.co/nanonets/Nanonets-OCR2-3B) | 11 named + "many more" (incl. Arabic + CJK) | Illustrative list, no benchmark — but the only card claiming multilingual **handwriting** |
-| [LightOnOCR-2-1B](https://huggingface.co/lightonai/LightOnOCR-2-1B) | 11 (9 European + zh/ja) | Declared, not benchmarked. [v1](https://huggingface.co/lightonai/LightOnOCR-1B-1025) is explicitly European/Latin-script only (9) |
-| [GLM-OCR](https://huggingface.co/zai-org/GLM-OCR) | 8 (zh en fr es ru de ja ko) | Declared in card metadata only; no per-language evidence in the card body |
-| [HunyuanOCR-1.5](https://huggingface.co/tencent/HunyuanOCR) | "Multilingual" | Nothing enumerated, but explicitly targets **low-resource + ancient-script OCR** as a design goal (new in 1.5; the pinned 1.0 revision doesn't claim this) |
-| [dots.mocr](https://huggingface.co/rednote-hilab/dots.mocr) · [DeepSeek-OCR](https://huggingface.co/deepseek-ai/DeepSeek-OCR) / [-2](https://huggingface.co/deepseek-ai/DeepSeek-OCR-2) · [Unlimited-OCR](https://huggingface.co/baidu/Unlimited-OCR) · [NuExtract3](https://huggingface.co/numind/NuExtract3) | "Multilingual", unspecified | A tag or one-liner only — no count, list, or benchmark |
-| [olmOCR-2](https://huggingface.co/allenai/olmOCR-2-7B-1025-FP8) · [Nanonets-OCR-s](https://huggingface.co/nanonets/Nanonets-OCR-s) · [SmolDocling](https://huggingface.co/ds4sd/SmolDocling-256M-preview) · [LFM2.5-VL-Extract](https://huggingface.co/LiquidAI/LFM2.5-VL-1.6B-Extract) | English only | Stated on the card |
-| [Falcon-OCR](https://huggingface.co/tiiuae/Falcon-OCR) · [OvisOCR2](https://huggingface.co/ATH-MaaS/OvisOCR2) · [FireRed-OCR](https://huggingface.co/FireRedTeam/FireRed-OCR) · [ABot-OCR](https://huggingface.co/acvlab/ABot-OCR) · [RolmOCR](https://huggingface.co/reducto/RolmOCR) · [NuMarkdown-8B](https://huggingface.co/numind/NuMarkdown-8B-Thinking) · [lift](https://huggingface.co/datalab-to/lift) | Not stated | No language information on the card at all |
-
-Text-only extraction: [LFM2-1.2B-Extract](https://huggingface.co/LiquidAI/LFM2-1.2B-Extract) (chains after OCR) names 9 languages: English, Arabic, Chinese, French, German, Japanese, Korean, Portuguese, Spanish.
-
-## Reading the table for low-resource work
-
-- **Only Surya lets you check your language before running anything** — its [benchmark table](https://github.com/datalab-to/surya/blob/master/static/docs/multilingual.md) has a row per language.
-- **Tesseract's coverage is enumerable** — if your language has a [traineddata pack](https://github.com/tesseract-ocr/tessdata_best), it's supported (quality varies; it's the legacy baseline, not the quality leader).
-- **dots.ocr and HunyuanOCR-1.5 are the VLMs that talk about low-resource languages at all**; PaddleOCR-VL-1.5 names Tibetan and Bengali specifically.
-- A big claimed number (192, 109) without a list or per-language scores means you're testing it yourself either way — which is cheap (see the command above).

OCR-BENCHMARK.md

@@ -1,93 +0,0 @@
-# OCR Benchmark — results & history
-
-Result tables and validation history for the OCR benchmark tooling. **How to *run* the tools**
-(`ocr-bench-run.py`, `ocr-vllm-judge.py`, `ocr-human-eval.py`) lives in `CLAUDE.md` → "Internal
-tooling"; this file is the accumulated *evidence*.
-
-## Model registry (as benchmarked)
-
-| Slug | Model | Size | GPU | Notes |
-|------|-------|------|-----|-------|
-| `glm-ocr` | `zai-org/GLM-OCR` | 0.9B | l4x1 | |
-| `deepseek-ocr` | `deepseek-ai/DeepSeek-OCR` | 4B | l4x1 | auto `--prompt-mode free` (no grounding tags) |
-| `lighton-ocr-2` | `lightonai/LightOnOCR-2-1B` | 1B | a100-large | |
-| `dots-ocr` | `rednote-hilab/dots.ocr` | 1.7B | l4x1 | stable vLLM (>=0.9.1) |
-
-## Run 1 — NLS Medical History (2026-02-14, pilot)
-
-`NationalLibraryOfScotland/medical-history-of-british-india`, 10 samples, seed 42. Judge:
-`Qwen2.5-VL-72B` via Inference Providers. Historical English, degraded scans.
-
-- **ELO (pairwise, 5 samples):** DoTS 1540 (67%) · DeepSeek 1539 (57%) · LightOnOCR-2 1486 (50%) · GLM 1436 (29%)
-- **Pointwise (5):** DeepSeek 5.0 · GLM 4.6 · LightOnOCR-2 4.4 · DoTS 4.2
-- **Key finding:** DeepSeek's `--prompt-mode document` emits grounding tags (`<|ref|>`/`<|det|>`) the
-  judge penalises heavily; switching to `--prompt-mode free` moved it last→top-2 (now the registry default).
-- **Caveat:** 5 samples is far too few for stable rankings.
-
-## Run 2 — Rubenstein Manuscript Catalog (2026-02-15, first full run)
-
-`biglam/rubenstein-manuscript-catalog`, 50 samples, seed 42. Judge: jury of `Qwen2.5-VL-7B` +
-`Qwen3-VL-8B` on A100 (`ocr-vllm-judge.py`). ~48K typewritten + handwritten cards (Duke, CC0).
-
-**ELO (50 samples, 300 comparisons, 0 parse failures):**
-
-| Rank | Model | ELO | W | L | T | Win% |
-|------|-------|-----|---|---|---|------|
-| 1 | LightOnOCR-2-1B | 1595 | 100 | 50 | 0 | 67% |
-| 2 | DeepSeek-OCR | 1497 | 73 | 77 | 0 | 49% |
-| 3 | GLM-OCR | 1471 | 57 | 93 | 0 | 38% |
-| 4 | dots.ocr | 1437 | 70 | 80 | 0 | 47% |
-
-Job times (50 samples): dots 5.3 min (L4) · deepseek 5.6 (L4) · glm 5.7 (L4) · lighton 6.4 (A100).
-
-**Findings:** LightOnOCR-2 dominates on manuscript cards (very different from the NLS pilot) — rankings
-are **dataset-dependent**; a jury of small models works well (0 parse failures via vLLM structured output);
-50 samples gives meaningful separation.
-
-## Run 3 — UFO-ColPali (2026-02-15, cross-dataset validation)
-
-`davanstrien/ufo-ColPali`, 50 samples, seed 42. Judge: `Qwen3-VL-30B-A3B` on A100 (updated prompt).
-Mixed modern documents.
-
-**ELO (50 samples, 294 comparisons):**
-
-| Rank | Model | ELO | W | L | T | Win% |
-|------|-------|-----|---|---|---|------|
-| 1 | DeepSeek-OCR | 1827 | 130 | 17 | 0 | 88% |
-| 2 | dots.ocr | 1510 | 64 | 83 | 0 | 44% |
-| 3 | LightOnOCR-2-1B | 1368 | 77 | 70 | 0 | 52% |
-| 4 | GLM-OCR | 1294 | 23 | 124 | 0 | 16% |
-
-**Human validation (30 comparisons):** DeepSeek #1 (matches judge), LightOnOCR-2 #3 (matches). Middle
-pack (GLM, dots) shuffled between human and judge.
-
-## Cross-dataset comparison (human-validated)
-
-| Model | Rubenstein Human | Rubenstein Kimi | UFO Human | UFO 30B |
-|-------|:---:|:---:|:---:|:---:|
-| DeepSeek-OCR | **#1** | **#1** | **#1** | **#1** |
-| GLM-OCR | #2 | #3 | #2 | #4 |
-| LightOnOCR-2 | #4 | #2 | #3 | #3 |
-| dots.ocr | #3 | #4 | #4 | #2 |
-
-**Conclusion:** DeepSeek-OCR is consistently #1 across datasets and eval methods; middle-pack rankings
-are dataset-dependent. (NLS pilot omitted — 5 samples / 72B API judge, not comparable with the newer
-methodology.)
-
-## Judge validation — `ocr-vllm-judge.py` (2026-02-15)
-
-- **Test 1** (single judge, 1 sample, L4): `Qwen2.5-VL-7B`, 6/6 comparisons, 0 parse failures, ~3 min.
-- **Test 2** (jury of 2, 3 samples, A100): `Qwen2.5-VL-7B` + `Qwen3-VL-8B`, 15/15, 0 failures; GPU cleanup
-  between models OK; majority-vote aggregation working (`[2/2]` unanimous, `[1/2]` split).
-- **Test 3** (full, 50 samples, A100, Rubenstein): 300/300 comparisons, 0 parse failures; clear ELO
-  separation. First saved dataset: `davanstrien/ocr-bench-rubenstein-judge`.
-
-Structured output via a compatibility shim: `StructuredOutputsParams` (vLLM ≥0.12) → `GuidedDecodingParams`
-(older) → prompt-based fallback. Position bias mitigated by A/B randomisation. A100 recommended for jury mode.
-
-## Human eval — `ocr-human-eval.py` first validation (Rubenstein, 30 annotations)
-
-Tested 3 judge configs against 30 human annotations. **Kimi K2.5 (170B) via Novita + the updated prompt**
-is the only judge to match the human's #1 (DeepSeek-OCR); it's now the default in `ocr-jury-bench.py`.
-Small models (7B/8B/30B) overrate LightOnOCR-2 (bias toward its commentary style); the updated prompt
-(faithfulness > completeness > accuracy) helps, but model size is the bigger factor.

README.md

@@ -1,17 +1,15 @@
 ---
 viewer: false
-tags: [uv-script, ocr, extraction, vision-language-model, document-processing, hf-jobs]
+tags: [uv-script, ocr, vision-language-model, document-processing, hf-jobs]
 ---
 
 # OCR UV Scripts
 
-<a href="https://huggingface.co/uv-scripts"><picture><source media="(prefers-color-scheme: dark)" srcset="https://huggingface.co/datasets/huggingface/badges/resolve/main/follow-us-on-hf-md-dark.svg"><img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/follow-us-on-hf-md.svg" alt="Follow uv-scripts on Hugging Face"></picture></a>
+> Part of [uv-scripts](https://huggingface.co/uv-scripts) - ready-to-run ML tools powered by UV and HuggingFace Jobs.
 
-> Part of [uv-scripts](https://huggingface.co/uv-scripts) — self-contained UV scripts you run on Hugging Face Jobs in one command.
+14 OCR models from 0.9B to 8B parameters. Pick a model, point at your dataset, get markdown — no setup required.
 
-A model zoo of OCR scripts — one per model — that add a `markdown` column to an image dataset. Pick a model from the table below, point it at your dataset, and run it on a GPU with one command. A few recipes do **structured extraction** instead — image *or* text → JSON given a schema (see [Structured extraction](#structured-extraction-image-or-text--json) below). Two more companions sit alongside: `pp-doclayout.py` detects layout regions (bboxes for text/title/table/figure/…) instead of text, and `ocr-vllm-judge.py` compares model outputs head-to-head.
-
-## Quick Start
+## 🚀 Quick Start
 
 Run OCR on any dataset without needing your own GPU:
 
@@ -24,243 +22,589 @@ hf jobs uv run --flavor l4x1 \
     --max-samples 10
 ```
 
-This will:
+That's it! The script will:
 
-- Process the first 10 images from your dataset
+- Process first 10 images from your dataset
 - Add OCR results as a new `markdown` column
 - Push the results to a new dataset
 - View results at: `https://huggingface.co/datasets/[your-output-dataset]`
 
-## Serve a model as a live endpoint
+<details><summary>All scripts at a glance (sorted by model size)</summary>
 
-The recipes here run as batch jobs. Some models also have a **`-server.py` sibling recipe** that runs the same dataset→dataset batch job through an in-job `vllm serve` + concurrent driver — measurably faster (continuous batching stays fed) and more robust (one bad image fails one request, not a whole batch); see [SERVING.md](SERVING.md) for the architecture, A/B numbers, and which models officially document server mode. To call a model interactively, from an agent, or with concurrent ad-hoc requests, you can instead run it as a temporary endpoint: [HF Jobs serving](https://huggingface.co/docs/hub/jobs-serving) exposes a port on a GPU Job, giving an OpenAI-compatible endpoint that runs until the job is cancelled or its `--timeout` is reached. See [serving-unlimited-ocr.md](serving-unlimited-ocr.md) for a worked example serving Baidu's [Unlimited-OCR](https://huggingface.co/baidu/Unlimited-OCR) — with vLLM (official image) or SGLang. To OCR a whole corpus of single-page images instead, the batch recipe `unlimited-ocr-vllm.py` is the better fit (it's single-image only). **Multi-page** documents need a server: both vLLM and SGLang read clean multi-page docs, but **SGLang is the more robust** — on hard/degraded scans vLLM multi-page hallucinated in our tests while SGLang held up.
+| Script | Model | Size | Backend | Notes |
+|--------|-------|------|---------|-------|
+| `smoldocling-ocr.py` | [SmolDocling](https://huggingface.co/ds4sd/SmolDocling-256M-preview) | 256M | Transformers | DocTags structured output |
+| `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 |
+| `lighton-ocr.py` | [LightOnOCR-1B](https://huggingface.co/lightonai/LightOnOCR-1B-1025) | 1B | vLLM | Fast, 3 vocab sizes |
+| `lighton-ocr2.py` | [LightOnOCR-2-1B](https://huggingface.co/lightonai/LightOnOCR-2-1B) | 1B | vLLM | 7× faster than v1, RLVR trained |
+| `hunyuan-ocr.py` | [HunyuanOCR](https://huggingface.co/tencent/HunyuanOCR) | 1B | vLLM | Lightweight VLM |
+| `dots-ocr.py` | [DoTS.ocr](https://huggingface.co/Tencent/DoTS.ocr) | 1.7B | vLLM | 100+ languages |
+| `firered-ocr.py` | [FireRed-OCR](https://huggingface.co/FireRedTeam/FireRed-OCR) | 2.1B | vLLM | Qwen3-VL fine-tune, Apache 2.0 |
+| `nanonets-ocr.py` | [Nanonets-OCR-s](https://huggingface.co/nanonets/Nanonets-OCR-s) | 2B | vLLM | LaTeX, tables, forms |
+| `dots-ocr-1.5.py` | [DoTS.ocr-1.5](https://huggingface.co/Tencent/DoTS.ocr-1.5) | 3B | vLLM | Updated multilingual model |
+| `nanonets-ocr2.py` | [Nanonets-OCR2-3B](https://huggingface.co/nanonets/Nanonets-OCR2-s) | 3B | vLLM | Next-gen, Qwen2.5-VL base |
+| `deepseek-ocr-vllm.py` | [DeepSeek-OCR](https://huggingface.co/deepseek-ai/DeepSeek-OCR) | 4B | vLLM | 5 resolution + 5 prompt modes |
+| `deepseek-ocr.py` | [DeepSeek-OCR](https://huggingface.co/deepseek-ai/DeepSeek-OCR) | 4B | Transformers | Same model, Transformers backend |
+| `deepseek-ocr2-vllm.py` | [DeepSeek-OCR-2](https://huggingface.co/deepseek-ai/DeepSeek-OCR-2) | 3B | vLLM | Newer, requires nightly vLLM |
+| `olmocr2-vllm.py` | [olmOCR-2-7B](https://huggingface.co/allenai/olmOCR-2-7B-1025-FP8) | 7B | vLLM | 82.4% olmOCR-Bench |
+| `rolm-ocr.py` | [RolmOCR](https://huggingface.co/reducto/RolmOCR) | 7B | vLLM | Qwen2.5-VL based, general-purpose |
+| `numarkdown-ocr.py` | [NuMarkdown-8B](https://huggingface.co/numind/NuMarkdown-8B-Thinking) | 8B | vLLM | Reasoning-based OCR |
+
+</details>
 
-## Models at a glance
+## Common Options
 
-**Start here:** for a quick first run, try **`lighton-ocr2.py`** (1B, very fast), **`paddleocr-vl-1.6.py`** (0.9B, 96.33 OmniDocBench) or **`ovis-ocr2.py`** (0.9B, 96.58 OmniDocBench — current SOTA); for the smallest footprint, **`falcon-ocr.py`** (0.3B, strong on tables). Reach for a 7–8B model only when quality demands it. Several of these models sit on the public [olmOCR-Bench](https://huggingface.co/datasets/allenai/olmOCR-bench) — pull the live ranking from your terminal in one command:
+All scripts accept the same core flags. Model-specific defaults (batch size, context length, temperature) are tuned per model based on model card recommendations and can be overridden.
+
+| Option | Description |
+|--------|-------------|
+| `--image-column` | Column containing images (default: `image`) |
+| `--output-column` | Output column name (default: `markdown`) |
+| `--split` | Dataset split (default: `train`) |
+| `--max-samples` | Limit number of samples (useful for testing) |
+| `--private` | Make output dataset private |
+| `--shuffle` | Shuffle dataset before processing |
+| `--seed` | Random seed for shuffling (default: `42`) |
+| `--batch-size` | Images per batch (default varies per model) |
+| `--max-model-len` | Max context length (default varies per model) |
+| `--max-tokens` | Max output tokens (default varies per model) |
+| `--gpu-memory-utilization` | GPU memory fraction (default: `0.8`) |
+| `--config` | Config name for Hub push (for benchmarking) |
+| `--create-pr` | Push as PR instead of direct commit |
+| `--verbose` | Log resolved package versions after run |
+
+Every script supports `--help` to see all available options:
 
 ```bash
-hf datasets leaderboard allenai/olmOCR-bench
+uv run glm-ocr.py --help
 ```
 
-But which model wins on *your* documents is still document-dependent — so [ocr-bench](https://github.com/davanstrien/ocr-bench) builds a **per-collection leaderboard** for your own data (pairwise VLM-as-judge, optionally human-validated), using these scripts under the hood.
+## Example: GLM-OCR
 
-**Language coverage:** [LANGUAGES.md](LANGUAGES.md) lists what each model's card claims (and how much evidence backs it). Machine-readable catalog for agents — script → model, params, backend, image pins, languages: [`models.json`](models.json).
+[GLM-OCR](https://huggingface.co/zai-org/GLM-OCR) (0.9B) scores 94.62% on OmniDocBench V1.5 and supports OCR, formula, and table extraction:
 
-_Sorted by model size:_
+```bash
+# Basic OCR
+hf jobs uv run --flavor l4x1 -s HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/glm-ocr.py \
+    my-documents my-ocr-output
 
-| Script | Model | Size | Backend | Notes |
-|--------|-------|------|---------|-------|
-| [`tesseract-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/tesseract-ocr.py) | [Tesseract 5](https://github.com/tesseract-ocr/tesseract) | n/a (classical) | pytesseract (CPU) | **The legacy baseline** — no GPU at all, runs on `cpu-upgrade`. Plain-text output, `--lang`/`--psm`/`--oem` exposed, 100+ language packs via apt. Apache 2.0 |
-| [`pp-ocrv6.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/pp-ocrv6.py) | [PP-OCRv6](https://huggingface.co/collections/PaddlePaddle/pp-ocrv6) | 1.5–34.5M | PaddleOCR (paddle) | **Smallest neural** — classical det+rec pipeline, not a VLM. Three tiers (`--model-tier tiny\|small\|medium`), plain-text output (not markdown). 48 langs. Runs on `t4-small`. Apache 2.0 |
-| [`falcon-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/falcon-ocr.py) | [Falcon-OCR](https://huggingface.co/tiiuae/Falcon-OCR) | 0.3B | falcon-perception | Smallest VLM in collection. #1 on multi-column docs and tables (olmOCR), Apache 2.0 |
-| [`smoldocling-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/smoldocling-ocr.py) | [SmolDocling](https://huggingface.co/ds4sd/SmolDocling-256M-preview) | 256M | Transformers | DocTags structured output |
-| [`surya-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/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`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/glm-ocr.py) | [GLM-OCR](https://huggingface.co/zai-org/GLM-OCR) | 0.9B | vLLM | 94.62% OmniDocBench V1.5 |
-| [`paddleocr-vl.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/paddleocr-vl.py) | [PaddleOCR-VL](https://huggingface.co/PaddlePaddle/PaddleOCR-VL) | 0.9B | vLLM | 4 task modes (ocr/table/formula/chart) |
-| [`paddleocr-vl-1.5.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/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 |
-| [`paddleocr-vl-1.6.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/paddleocr-vl-1.6.py) | [PaddleOCR-VL-1.6](https://huggingface.co/PaddlePaddle/PaddleOCR-VL-1.6) | 0.9B | vLLM | **96.33% OmniDocBench v1.6**, drop-in upgrade of 1.5 |
-| [`ovis-ocr2.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/ovis-ocr2.py) | [OvisOCR2](https://huggingface.co/ATH-MaaS/OvisOCR2) | 0.9B | vLLM | **96.58 OmniDocBench v1.6** (SOTA; first end-to-end model to top it). Qwen3.5 base; markdown + LaTeX + HTML tables. Apache 2.0 |
-| [`ovis-ocr2-server.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/ovis-ocr2-server.py) | [OvisOCR2](https://huggingface.co/ATH-MaaS/OvisOCR2) | 0.9B | vLLM server | **Server-mode sibling** of `ovis-ocr2.py`: in-job `vllm serve` + concurrent driver — ~1.7× its throughput, per-image failure isolation. See [SERVING.md](SERVING.md) |
-| [`lighton-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/lighton-ocr.py) | [LightOnOCR-1B](https://huggingface.co/lightonai/LightOnOCR-1B-1025) | 1B | vLLM | Fast, 3 vocab sizes |
-| [`lighton-ocr2.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/lighton-ocr2.py) | [LightOnOCR-2-1B](https://huggingface.co/lightonai/LightOnOCR-2-1B) | 1B | vLLM | 7× faster than v1, RLVR trained |
-| [`lighton-ocr2-server.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/lighton-ocr2-server.py) | [LightOnOCR-2-1B](https://huggingface.co/lightonai/LightOnOCR-2-1B) | 1B | vLLM server | **Server-mode sibling** of `lighton-ocr2.py` (the card's own documented path) — ~1.8× its throughput. See [SERVING.md](SERVING.md) |
-| [`hunyuan-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/hunyuan-ocr.py) | [HunyuanOCR 1.0](https://huggingface.co/tencent/HunyuanOCR/tree/f6af82ee007fe6091b29fb3bb287b491ead41c82) | 1B | vLLM | Lightweight VLM. Pinned to the last 1.0 revision (repo root became 1.5 in-place on 2026-07-06). [Hunyuan Community License](https://huggingface.co/tencent/HunyuanOCR/blob/main/LICENSE) (excludes EU/UK/KR) |
-| [`hunyuan-ocr-1.5.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/hunyuan-ocr-1.5.py) | [HunyuanOCR-1.5](https://huggingface.co/tencent/HunyuanOCR) | 1B | vLLM | 128K context, 4K images, 12 task types, ancient scripts. ~4-5× faster/page than dots.ocr & DeepSeek-OCR-2 (tech report). [Hunyuan Community License](https://huggingface.co/tencent/HunyuanOCR/blob/main/LICENSE) (excludes EU/UK/KR) |
-| [`dots-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/dots-ocr.py) | [dots.ocr](https://huggingface.co/rednote-hilab/dots.ocr) | 1.7B | vLLM | 100 languages (in-house bench), explicit low-resource claim |
-| [`firered-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/firered-ocr.py) | [FireRed-OCR](https://huggingface.co/FireRedTeam/FireRed-OCR) | 2.1B | vLLM | Qwen3-VL fine-tune, Apache 2.0 |
-| [`abot-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/abot-ocr.py) | [ABot-OCR](https://huggingface.co/acvlab/ABot-OCR) | 2B | vLLM | Qwen3-VL based, doc→Markdown (text/LaTeX/HTML tables). Needs `vllm/vllm-openai` image. [paper](https://arxiv.org/abs/2605.27978) |
-| [`nanonets-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/nanonets-ocr.py) | [Nanonets-OCR-s](https://huggingface.co/nanonets/Nanonets-OCR-s) | 2B | vLLM | LaTeX, tables, forms |
-| [`dots-mocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/dots-mocr.py) | [dots.mocr](https://huggingface.co/rednote-hilab/dots.mocr) | 3B | vLLM | 8 prompt modes incl. SVG generation, layout + bbox, 100+ languages |
-| [`nanonets-ocr2.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/nanonets-ocr2.py) | [Nanonets-OCR2-3B](https://huggingface.co/nanonets/Nanonets-OCR2-3B) | 3B | vLLM | Next-gen, Qwen2.5-VL base. **Pin `--image vllm/vllm-openai:v0.10.2`** (vLLM ≥0.11 breaks Qwen2.5-VL → all `!`) |
-| [`deepseek-ocr-vllm.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/deepseek-ocr-vllm.py) | [DeepSeek-OCR](https://huggingface.co/deepseek-ai/DeepSeek-OCR) | 4B | vLLM | 5 resolution + 5 prompt modes |
-| [`deepseek-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/deepseek-ocr.py) | [DeepSeek-OCR](https://huggingface.co/deepseek-ai/DeepSeek-OCR) | 4B | Transformers | Same model, Transformers backend |
-| [`deepseek-ocr2-vllm.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/deepseek-ocr2-vllm.py) | [DeepSeek-OCR-2](https://huggingface.co/deepseek-ai/DeepSeek-OCR-2) | 3B | vLLM | Newer; needs nightly vLLM **+ the `vllm/vllm-openai` image** ([why](#if-a-vllm-script-crashes-at-startup-the-nvcc--nvrtc-error)) |
-| [`unlimited-ocr-vllm.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/unlimited-ocr-vllm.py) | [Unlimited-OCR](https://huggingface.co/baidu/Unlimited-OCR) | 3.3B | vLLM | DeepSeek-OCR-based; layout-grounded markdown (`--strip-grounding` for clean text). Single-image batch — needs Baidu's **dedicated `vllm/vllm-openai:unlimited-ocr` image** (`-cu129` on Hopper). Multi-page "long-horizon" parsing → serve it ([doc](serving-unlimited-ocr.md)); both engines do clean docs, **SGLang more robust** on hard scans. MIT |
-| [`nuextract3.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/nuextract3.py) | [NuExtract3](https://huggingface.co/numind/NuExtract3) | 4B | vLLM | Markdown OCR **+ schema-guided JSON extraction** (template/Pydantic). Needs `vllm/vllm-openai` image |
-| [`qianfan-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/qianfan-ocr.py) | [Qianfan-OCR](https://huggingface.co/baidu/Qianfan-OCR) | 4.7B | vLLM | #1 OmniDocBench v1.5 (93.12), Layout-as-Thought, 192 languages |
-| [`olmocr2-vllm.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/olmocr2-vllm.py) | [olmOCR-2-7B](https://huggingface.co/allenai/olmOCR-2-7B-1025-FP8) | 7B | vLLM | 82.4% olmOCR-Bench |
-| [`rolm-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/rolm-ocr.py) | [RolmOCR](https://huggingface.co/reducto/RolmOCR) | 7B | vLLM | Qwen2.5-VL based, general-purpose |
-| [`numarkdown-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/numarkdown-ocr.py) | [NuMarkdown-8B](https://huggingface.co/numind/NuMarkdown-8B-Thinking) | 8B | vLLM | Reasoning-based OCR |
-
-**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 · `surya-ocr-bucket.py` is the structured bucket recipe — OCR a bucket of files (no dataset round-trip) via either a FUSE mount **or** `huggingface_hub` batch-copy (`--io-mode mount|copy`), writing per-page `.md` + `.json` (`surya_blocks`) back to a bucket (resumable) and/or a pushed dataset · `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:
-
-| Script | Model | Size | Input | Output |
-|--------|-------|------|-------|--------|
-| [`lfm2-vl-extract.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/lfm2-vl-extract.py) | [LFM2.5-VL-1.6B-Extract](https://huggingface.co/LiquidAI/LFM2.5-VL-1.6B-Extract) | 1.6B | image | JSON |
-| [`nuextract3.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/nuextract3.py) | [NuExtract3](https://huggingface.co/numind/NuExtract3) | 4B | image | markdown **or** JSON |
-| [`lfm2-extract.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/lfm2-extract.py) | [LFM2-1.2B-Extract](https://huggingface.co/LiquidAI/LFM2-1.2B-Extract) | 1.2B | **text** | JSON / XML / YAML |
-| [`lift-extract.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/lift-extract.py) | [lift](https://huggingface.co/datalab-to/lift) | 9B | image **or** PDF | JSON |
-
-Pass `--schema` (inline JSON, a URL, or a file path). The LFM models are small and fast; run them on the `vllm/vllm-openai` image so the CUDA toolkit is present (each script's docstring has the exact command). Because `lfm2-extract.py` works on a **text** column, you can **chain it after OCR**: a recipe above turns a page into `markdown`, then `lfm2-extract.py` turns that markdown into fields.
-
-`lift-extract.py` is the one outlier: a 9B model that also reads **multi-page PDFs** (`--pdf-column`, `--page-range`) and runs on either Transformers (`--method hf`) or vLLM (`--method vllm`). Its weights are **modified OpenRAIL-M** (free for research, personal use, and startups under $5M; no competitive use against Datalab's API) — the only non-permissive license here, so check the terms.
+# Table extraction
+hf jobs uv run --flavor l4x1 -s HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/glm-ocr.py \
+    my-documents my-tables --task table
 
-```bash
-# image → JSON directly
-hf jobs uv run --flavor l4x1 --secrets HF_TOKEN \
-    --image vllm/vllm-openai --python /usr/bin/python3 \
-    -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages \
-    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/lfm2-vl-extract.py \
-    my-images my-fields --schema '{"title": "the document title", "date": "any date shown"}'
+# Test on 10 samples first
+hf jobs uv run --flavor l4x1 -s HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/glm-ocr.py \
+    my-documents my-test --max-samples 10
 ```
 
-## Layout detection (not OCR)
+<details><summary>Detailed per-model documentation</summary>
 
-`pp-doclayout.py` runs PaddleOCR's [PP-DocLayout-L](https://huggingface.co/PaddlePaddle/PP-DocLayout-L) (or M / S / plus-L) and emits per-image **bounding boxes + region classes** (text, title, table, figure, formula, list, header, footer, ...) — it does NOT extract text. Useful for filtering pages, cropping regions for downstream OCR, dataset analysis, and training-data prep.
+### PaddleOCR-VL-1.5 (`paddleocr-vl-1.5.py`) — 6 task modes
 
-| Script | Model | Size | Backend | Notes |
-|--------|-------|------|---------|-------|
-| [`pp-doclayout.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/pp-doclayout.py) | [PP-DocLayout-L](https://huggingface.co/PaddlePaddle/PP-DocLayout-L) | 123M | paddleocr | Layout bboxes (no text). Bucket support: incremental parquet shards, resumable. |
+OCR using [PaddlePaddle/PaddleOCR-VL-1.5](https://huggingface.co/PaddlePaddle/PaddleOCR-VL-1.5) with 94.5% accuracy:
+
+- **94.5% on OmniDocBench v1.5** (0.9B parameters)
+- 🧩 **Ultra-compact** - Only 0.9B parameters
+- 📝 **OCR mode** - General text extraction to markdown
+- 📊 **Table mode** - HTML table recognition
+- 📐 **Formula mode** - LaTeX mathematical notation
+- 📈 **Chart mode** - Chart and diagram analysis
+- 🔍 **Spotting mode** - Text spotting with localization (higher resolution)
+- 🔖 **Seal mode** - Seal and stamp recognition
+- 🌍 **Multilingual** - Support for multiple languages
+
+**Task Modes:**
+
+- `ocr`: General text extraction (default)
+- `table`: Table extraction to HTML
+- `formula`: Mathematical formula to LaTeX
+- `chart`: Chart and diagram analysis
+- `spotting`: Text spotting with localization
+- `seal`: Seal and stamp recognition
+
+**Quick start:**
 
 ```bash
-hf jobs uv run --flavor l4x1 -s HF_TOKEN \
-    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/pp-doclayout.py \
-    your-dataset your-layout-output --max-samples 10
+# Basic OCR mode
+hf jobs uv run --flavor l4x1 \
+    -s HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/paddleocr-vl-1.5.py \
+    your-input-dataset your-output-dataset \
+    --max-samples 100
+
+# Table extraction
+hf jobs uv run --flavor l4x1 \
+    -s HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/paddleocr-vl-1.5.py \
+    documents tables-extracted \
+    --task-mode table
+
+# Seal recognition
+hf jobs uv run --flavor l4x1 \
+    -s HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/paddleocr-vl-1.5.py \
+    documents seals-extracted \
+    --task-mode seal
 ```
 
-Source/sink can be either an HF dataset repo OR an `hf://buckets/...` URL (auto-detected). Bucket output writes incremental zstd parquet shards via the buckets API — resumable across runs (snapshot-backed source listing) and no git/commit overhead. See the script's `--help` for all flags.
+### PaddleOCR-VL (`paddleocr-vl.py`) 🎯 Smallest model with task-specific modes!
 
-## If a vLLM script crashes at startup (the `nvcc` / `nvrtc` error)
+Ultra-compact OCR using [PaddlePaddle/PaddleOCR-VL](https://huggingface.co/PaddlePaddle/PaddleOCR-VL) with only 0.9B parameters:
 
-The vLLM recipes run on the **default** Jobs image and carry a guard (`VLLM_USE_FLASHINFER_SAMPLER=0`) so they work there with the plain command. But some — especially nightly-vLLM ones — JIT-compile a CUDA kernel at engine init and crash on the default image with one of:
+- 🎯 **Smallest model** - Only 0.9B parameters (even smaller than LightOnOCR!)
+- 📝 **OCR mode** - General text extraction to markdown
+- 📊 **Table mode** - HTML table recognition and extraction
+- 📐 **Formula mode** - LaTeX mathematical notation
+- 📈 **Chart mode** - Structured chart and diagram analysis
+- 🌍 **Multilingual** - Support for multiple languages
+- ⚡ **Fast initialization** - Tiny model size for quick startup
+- 🔧 **ERNIE-4.5 based** - Different architecture from Qwen models
 
+**Task Modes:**
+
+- `ocr`: General text extraction (default)
+- `table`: Table extraction to HTML
+- `formula`: Mathematical formula to LaTeX
+- `chart`: Chart and diagram analysis
+
+**Quick start:**
+
+```bash
+# Basic OCR mode
+hf jobs uv run --flavor l4x1 \
+    -s HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/paddleocr-vl.py \
+    your-input-dataset your-output-dataset \
+    --max-samples 100
+
+# Table extraction
+hf jobs uv run --flavor l4x1 \
+    -s HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/paddleocr-vl.py \
+    documents tables-extracted \
+    --task-mode table \
+    --batch-size 32
 ```
-RuntimeError: Could not find nvcc and default cuda_home='/usr/local/cuda' doesn't exist
-nvrtc: error: failed to open libnvrtc-builtins.so...
-```
 
-Run those on the **`vllm/vllm-openai` image**, which ships the full CUDA toolkit. Add these flags to any recipe — they point `import vllm` at the image's CUDA-matched build:
+### GLM-OCR (`glm-ocr.py`) 🏆 SOTA on OmniDocBench V1.5!
+
+Compact high-performance OCR using [zai-org/GLM-OCR](https://huggingface.co/zai-org/GLM-OCR) with 0.9B parameters:
+
+- 🏆 **94.62% on OmniDocBench V1.5** - #1 overall ranking
+- 🧠 **Multi-Token Prediction** - MTP loss + stable full-task RL for quality
+- 📝 **Text recognition** - Clean markdown output
+- 📐 **Formula recognition** - LaTeX mathematical notation
+- 📊 **Table recognition** - Structured table extraction
+- 🌍 **Multilingual** - zh, en, fr, es, ru, de, ja, ko
+- ⚡ **Compact** - Only 0.9B parameters, MIT licensed
+- 🔧 **CogViT + GLM** - Visual encoder with efficient token downsampling
+
+**Task Modes:**
+
+- `ocr`: Text recognition (default)
+- `formula`: LaTeX formula recognition
+- `table`: Table extraction
+
+**Quick start:**
 
 ```bash
-hf jobs uv run --flavor l4x1 --secrets HF_TOKEN \
-    --image vllm/vllm-openai --python /usr/bin/python3 \
-    -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages \
-    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/<script>.py \
-    INPUT OUTPUT --max-samples 10
+# Basic OCR
+hf jobs uv run --flavor l4x1 \
+    -s HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/glm-ocr.py \
+    your-input-dataset your-output-dataset \
+    --max-samples 100
+
+# Formula recognition
+hf jobs uv run --flavor l4x1 \
+    -s HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/glm-ocr.py \
+    scientific-papers formulas-extracted \
+    --task formula
+
+# Table extraction
+hf jobs uv run --flavor l4x1 \
+    -s HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/glm-ocr.py \
+    documents tables-extracted \
+    --task table
 ```
 
-This is **required** for a few scripts (e.g. `deepseek-ocr2-vllm.py`, `abot-ocr.py`, `nuextract3.py`) and a safe fallback for any vLLM recipe that crashes at startup. (It's also the more robust way to run any vLLM recipe — full CUDA toolkit, ABI-matched build. It isn't a speed-up: uv still reinstalls the script's deps either way.)
+### LightOnOCR (`lighton-ocr.py`) ⚡ Good one to test first since it's small and fast!
 
-## Common Options
+Fast and compact OCR using [lightonai/LightOnOCR-1B-1025](https://huggingface.co/lightonai/LightOnOCR-1B-1025):
 
-The scripts aim to expose a **consistent interface**: every OCR model script takes `input-dataset output-dataset` as positional arguments, accepts the shared core flags below, and writes a `markdown` column — so switching models is usually just swapping the script URL. Models differ where they need to, though: some add their own flags (task modes, resolution presets, `--think`, vocab sizes), a few need a specific Docker image, and per-model defaults (batch size, context length, temperature) are tuned to each model card. Always check a script's `--help` for its specifics.
+- ⚡ **Fastest**: 5.71 pages/sec on H100, ~6.25 images/sec on A100 with batch_size=4096
+- 🎯 **Compact**: Only 1B parameters - quick to download and initialize
+- 🌍 **Multilingual**: 3 vocabulary sizes for different use cases
+- 📐 **LaTeX formulas**: Mathematical notation in LaTeX format
+- 📊 **Table extraction**: Markdown table format
+- 📝 **Document structure**: Preserves hierarchy and layout
+- 🚀 **Production-ready**: 76.1% benchmark score, used in production
 
-| Option | Description |
-|--------|-------------|
-| `--image-column` | Column containing images (default: `image`) |
-| `--output-column` | Output column name (default: `markdown`) |
-| `--split` | Dataset split (default: `train`) |
-| `--max-samples` | Limit number of samples (useful for testing) |
-| `--private` | Make output dataset private |
-| `--shuffle` | Shuffle dataset before processing |
-| `--seed` | Random seed for shuffling (default: `42`) |
-| `--batch-size` | Images per batch (default varies per model) |
-| `--max-model-len` | Max context length (default varies per model) |
-| `--max-tokens` | Max output tokens (default varies per model) |
-| `--gpu-memory-utilization` | GPU memory fraction (default: `0.8`) |
-| `--config` | Config name for Hub push (for benchmarking) |
-| `--create-pr` | Push as PR instead of direct commit |
-| `--verbose` | Log resolved package versions after run |
+**Vocabulary sizes:**
 
-Every script supports `--help` to see all available options:
+- `151k`: Full vocabulary, all languages (default)
+- `32k`: European languages, ~12% faster decoding
+- `16k`: European languages, ~12% faster decoding
+
+**Quick start:**
 
 ```bash
-uv run glm-ocr.py --help
+# Test on 100 samples with English text (32k vocab is fastest for European languages)
+hf jobs uv run --flavor l4x1 \
+    -s HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/lighton-ocr.py \
+    your-input-dataset your-output-dataset \
+    --vocab-size 32k \
+    --batch-size 32 \
+    --max-samples 100
+
+# Full production run on A100 (can handle huge batches!)
+hf jobs uv run --flavor a100-large \
+    -s HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/lighton-ocr.py \
+    your-input-dataset your-output-dataset \
+    --vocab-size 32k \
+    --batch-size 4096 \
+    --temperature 0.0
 ```
 
-## NuExtract3: markdown OCR + structured extraction
+### LightOnOCR-2 (`lighton-ocr2.py`) ⚡ Fastest OCR model!
+
+Next-generation fast OCR using [lightonai/LightOnOCR-2-1B](https://huggingface.co/lightonai/LightOnOCR-2-1B) with RLVR training:
 
-[NuExtract3](https://huggingface.co/numind/NuExtract3) (4B, Apache-2.0) is the one script here that does both document-to-markdown OCR *and* schema-guided JSON extraction. Give it a template (or a JSON Schema / Pydantic model) and it returns JSON shaped to match.
+- ⚡ **7× faster than v1**: 42.8 pages/sec on H100 (vs 5.71 for v1)
+- 🎯 **Higher accuracy**: 83.2% on OlmOCR-Bench (+7.1% vs v1)
+- 🧠 **RLVR trained**: Eliminates repetition loops and formatting errors
+- 📚 **Better dataset**: 2.5× larger training data with cleaner annotations
+- 🌍 **Multilingual**: Optimized for European languages
+- 📐 **LaTeX formulas**: Mathematical notation support
+- 📊 **Table extraction**: Markdown table format
+- 💪 **Production-ready**: Outperforms models 9× larger
 
-> **Run it with the `vllm/vllm-openai` image.** NuExtract3's Qwen3.5 architecture needs the image's prebuilt CUDA kernels — the default uv-script image lacks `nvcc`, so flashinfer's JIT compile fails at engine warmup. Use `--image vllm/vllm-openai:latest --python /usr/bin/python3 -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages` on `a100-large`.
+**Quick start:**
 
 ```bash
-# Markdown OCR (default mode)
+# Test on 100 samples
 hf jobs uv run --flavor a100-large \
-    --image vllm/vllm-openai:latest \
-    --python /usr/bin/python3 \
-    -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages \
     -s HF_TOKEN \
-    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/nuextract3.py \
-    my-documents my-markdown --max-samples 10
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/lighton-ocr2.py \
+    your-input-dataset your-output-dataset \
+    --batch-size 32 \
+    --max-samples 100
 
-# Structured extraction with an inline template
+# Full production run
 hf jobs uv run --flavor a100-large \
-    --image vllm/vllm-openai:latest \
-    --python /usr/bin/python3 \
-    -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages \
     -s HF_TOKEN \
-    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/nuextract3.py \
-    receipts extracted \
-    --template '{"store": "verbatim-string", "date": "date", "total": "number"}'
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/lighton-ocr2.py \
+    your-input-dataset your-output-dataset \
+    --batch-size 32
+```
+
+### DeepSeek-OCR (`deepseek-ocr-vllm.py`)
+
+Advanced document OCR using [deepseek-ai/DeepSeek-OCR](https://huggingface.co/deepseek-ai/DeepSeek-OCR) with visual-text compression:
+
+- 📐 **LaTeX equations** - Mathematical formulas in LaTeX format
+- 📊 **Tables** - Extracted as HTML/markdown
+- 📝 **Document structure** - Headers, lists, formatting preserved
+- 🖼️ **Image grounding** - Spatial layout with bounding boxes
+- 🔍 **Complex layouts** - Multi-column and hierarchical structures
+- 🌍 **Multilingual** - Multiple language support
+- 🎚️ **Resolution modes** - 5 presets for speed/quality trade-offs
+- 💬 **Prompt modes** - 5 presets for different OCR tasks
+- ⚡ **Fast batch processing** - vLLM acceleration
+
+**Resolution Modes:**
+
+- `tiny` (512×512): Fast, 64 vision tokens
+- `small` (640×640): Balanced, 100 vision tokens
+- `base` (1024×1024): High quality, 256 vision tokens
+- `large` (1280×1280): Maximum quality, 400 vision tokens
+- `gundam` (dynamic): Adaptive multi-tile (default)
+
+**Prompt Modes:**
+
+- `document`: Convert to markdown with grounding (default)
+- `image`: OCR any image with grounding
+- `free`: Fast OCR without layout
+- `figure`: Parse figures from documents
+- `describe`: Detailed image descriptions
+
+### RolmOCR (`rolm-ocr.py`)
+
+Fast general-purpose OCR using [reducto/RolmOCR](https://huggingface.co/reducto/RolmOCR) based on Qwen2.5-VL-7B:
+
+- 🚀 **Fast extraction** - Optimized for speed and efficiency
+- 📄 **Plain text output** - Clean, natural text representation
+- 💪 **General-purpose** - Works well on various document types
+- 🔥 **Large context** - Handles up to 16K tokens
+- ⚡ **Batch optimized** - Efficient processing with vLLM
+
+### Nanonets OCR (`nanonets-ocr.py`)
+
+State-of-the-art document OCR using [nanonets/Nanonets-OCR-s](https://huggingface.co/nanonets/Nanonets-OCR-s) that handles:
+
+- 📐 **LaTeX equations** - Mathematical formulas preserved
+- 📊 **Tables** - Extracted as HTML format
+- 📝 **Document structure** - Headers, lists, formatting maintained
+- 🖼️ **Images** - Captions and descriptions included
+- ☑️ **Forms** - Checkboxes rendered as ☐/☑
+
+### Nanonets OCR2 (`nanonets-ocr2.py`)
+
+Next-generation Nanonets OCR using [nanonets/Nanonets-OCR2-3B](https://huggingface.co/nanonets/Nanonets-OCR2-3B) with improved accuracy:
+
+- 🎯 **Enhanced quality** - 3.75B parameters for superior OCR accuracy
+- 📐 **LaTeX equations** - Mathematical formulas preserved in LaTeX format
+- 📊 **Advanced tables** - Improved HTML table extraction
+- 📝 **Document structure** - Headers, lists, formatting maintained
+- 🖼️ **Smart image captions** - Intelligent descriptions and captions
+- ☑️ **Forms** - Checkboxes rendered as ☐/☑
+- 🌍 **Multilingual** - Enhanced language support
+- 🔧 **Based on Qwen2.5-VL** - Built on state-of-the-art vision-language model
+
+### SmolDocling (`smoldocling-ocr.py`)
+
+Ultra-compact document understanding using [ds4sd/SmolDocling-256M-preview](https://huggingface.co/ds4sd/SmolDocling-256M-preview) with only 256M parameters:
+
+- 🏷️ **DocTags format** - Efficient XML-like representation
+- 💻 **Code blocks** - Preserves indentation and syntax
+- 🔢 **Formulas** - Mathematical expressions with layout
+- 📊 **Tables & charts** - Structured data extraction
+- 📐 **Layout preservation** - Bounding boxes and spatial info
+- ⚡ **Ultra-fast** - Tiny model size for quick inference
+
+### NuMarkdown (`numarkdown-ocr.py`)
+
+Advanced reasoning-based OCR using [numind/NuMarkdown-8B-Thinking](https://huggingface.co/numind/NuMarkdown-8B-Thinking) that analyzes documents before converting to markdown:
+
+- 🧠 **Reasoning Process** - Thinks through document layout before generation
+- 📊 **Complex Tables** - Superior table extraction and formatting
+- 📐 **Mathematical Formulas** - Accurate LaTeX/math notation preservation
+- 🔍 **Multi-column Layouts** - Handles complex document structures
+- ✨ **Thinking Traces** - Optional inclusion of reasoning process with `--include-thinking`
+
+### DoTS.ocr (`dots-ocr.py`)
+
+Compact multilingual OCR using [rednote-hilab/dots.ocr](https://huggingface.co/rednote-hilab/dots.ocr) with only 1.7B parameters:
+
+- 🌍 **100+ Languages** - Extensive multilingual support
+- 📝 **Simple OCR** - Clean text extraction (default mode)
+- 📊 **Layout Analysis** - Optional structured output with bboxes and categories
+- 📐 **Formula recognition** - LaTeX format support
+- 🎯 **Compact** - Only 1.7B parameters, efficient on smaller GPUs
+- 🔀 **Flexible prompts** - Switch between OCR, layout-all, and layout-only modes
+
+### FireRed-OCR (`firered-ocr.py`)
+
+Document OCR using [FireRedTeam/FireRed-OCR](https://huggingface.co/FireRedTeam/FireRed-OCR), a 2.1B model fine-tuned from Qwen3-VL-2B-Instruct:
+
+- 📝 **Structured Markdown** - Preserves headings, paragraphs, lists
+- 📐 **LaTeX formulas** - Inline and block math support
+- 📊 **HTML tables** - Table extraction with `<table>` tags
+- 🪶 **Lightweight** - 2.1B parameters, runs on L4 GPU
+- 📜 **Apache 2.0** - Permissive license
+
+**Quick start:**
+
+```bash
+hf jobs uv run --flavor l4x1 \
+    -s HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/firered-ocr.py \
+    your-input-dataset your-output-dataset \
+    --max-samples 100
+```
+
+### olmOCR2 (`olmocr2-vllm.py`)
+
+High-quality document OCR using [allenai/olmOCR-2-7B-1025-FP8](https://huggingface.co/allenai/olmOCR-2-7B-1025-FP8) optimized with GRPO reinforcement learning:
+
+- 🎯 **High accuracy** - 82.4 ± 1.1 on olmOCR-Bench (84.9% on math)
+- 📐 **LaTeX equations** - Mathematical formulas in LaTeX format
+- 📊 **Table extraction** - Structured table recognition
+- 📑 **Multi-column layouts** - Complex document structures
+- 🗜️ **FP8 quantized** - Efficient 8B model for faster inference
+- 📜 **Degraded scans** - Works well on old/historical documents
+- 📝 **Long text extraction** - Headers, footers, and full document content
+- 🧩 **YAML metadata** - Structured front matter (language, rotation, content type)
+- 🚀 **Based on Qwen2.5-VL-7B** - Fine-tuned with reinforcement learning
+
+## 🆕 New Features
+
+### Multi-Model Comparison Support
+
+All scripts now include `inference_info` tracking for comparing multiple OCR models:
+
+```bash
+# First model
+uv run rolm-ocr.py my-dataset my-dataset --max-samples 100
+
+# Second model (appends to same dataset)
+uv run nanonets-ocr.py my-dataset my-dataset --max-samples 100
+
+# View all models used
+python -c "import json; from datasets import load_dataset; ds = load_dataset('my-dataset'); print(json.loads(ds[0]['inference_info']))"
+```
+
+### Random Sampling
+
+Get representative samples with the new `--shuffle` flag:
+
+```bash
+# Random 50 samples instead of first 50
+uv run rolm-ocr.py ordered-dataset output --max-samples 50 --shuffle
+
+# Reproducible random sampling
+uv run nanonets-ocr.py dataset output --max-samples 100 --shuffle --seed 42
 ```
 
-**Templates** (`--template`) and **JSON Schemas** (`--schema`) each accept **inline JSON, a URL, or a file path**, so a schema can be hosted once and reused. Add `--enable-thinking` for harder layouts (slower; reasoning trace stored in a `<output-column>_reasoning` column). Template field names act as the model's extraction instructions, so name them descriptively — overly leading names can prompt over-generation, so verify against a few examples.
-
-## Model-specific modes & flags
-
-Beyond the shared flags, some models add their own. Run `--help` on any script for the full list; the common ones:
-
-| Script | Extra options |
-|--------|---------------|
-| [`surya-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/surya-ocr.py) | `--task ocr\|layout\|table`, `--table-mode full\|simple`, `--pdf-column`/`--page-range`, `--blocks-column` |
-| [`pp-ocrv6.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/pp-ocrv6.py) | `--model-tier tiny\|small\|medium` (1.5M–34.5M params) |
-| [`glm-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/glm-ocr.py) | `--task ocr\|formula\|table` |
-| [`ovis-ocr2.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/ovis-ocr2.py) | `--keep-image-tags` (retain visual-region `<img>` bbox tags, filtered by default), `--min-pixels`/`--max-pixels` (processor bounds, card defaults 448²/2880²) |
-| [`paddleocr-vl.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/paddleocr-vl.py) | `--task-mode ocr\|table\|formula\|chart` |
-| [`paddleocr-vl-1.5.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/paddleocr-vl-1.5.py) | `--task-mode ocr\|table\|formula\|chart\|spotting\|seal` |
-| [`paddleocr-vl-1.6.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/paddleocr-vl-1.6.py) | `--task-mode ocr\|table\|formula` |
-| [`lighton-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/lighton-ocr.py) | `--vocab-size 151k\|32k\|16k` (smaller = faster on European languages) |
-| [`deepseek-ocr-vllm.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/deepseek-ocr-vllm.py) | `--resolution-mode tiny\|small\|base\|large\|gundam`, `--prompt-mode document\|image\|free\|figure\|describe`; pass `-e UV_TORCH_BACKEND=auto` |
-| [`dots-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/dots-ocr.py) | `--prompt-mode ocr\|layout-all\|layout-only` |
-| [`dots-mocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/dots-mocr.py) | `--prompt-mode` (8: ocr, layout-all, layout-only, web-parsing, scene-spotting, grounding-ocr, svg, general); SVG: `--model rednote-hilab/dots.mocr-svg --prompt-mode svg` |
-| [`qianfan-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/qianfan-ocr.py) | `--prompt-mode ocr\|table\|formula\|chart\|scene\|kie`, `--think` (Layout-as-Thought); `kie` needs `--custom-prompt` |
-| [`unlimited-ocr-vllm.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/unlimited-ocr-vllm.py) | `--strip-grounding` (drop `<\|det\|>`/`<\|ref\|>` grounding tags); needs the **`vllm/vllm-openai:unlimited-ocr`** image |
-| [`numarkdown-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/numarkdown-ocr.py) | `--include-thinking` (store the reasoning trace) |
-| [`nuextract3.py`](https://huggingface.co/datasets/uv-scripts/ocr/blob/main/nuextract3.py) | `--template` / `--schema` / `--enable-thinking` — see the NuExtract3 section above |
-
-**Image-mode models** — `abot-ocr.py` and `nuextract3.py` (Qwen3.5 architecture) need the `vllm/vllm-openai` image because the default uv-script image lacks `nvcc`. Add `--image vllm/vllm-openai:latest --python /usr/bin/python3 -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages` (see the NuExtract3 example above for the full command). `unlimited-ocr-vllm.py` is a special case — its architecture isn't in any stable vLLM wheel, so it needs Baidu's **dedicated** `vllm/vllm-openai:unlimited-ocr` image (tag `:unlimited-ocr-cu129` on Hopper), e.g. `--image vllm/vllm-openai:unlimited-ocr --python /usr/bin/python3 -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages` (its docstring has the full command).
-
-## Output & features
-
-- **Markdown column** — each run adds an `--output-column` (default `markdown`) with the OCR result.
-- **Multi-model comparison** — every script records `inference_info`, so you can run several models into the *same* dataset and compare. Point a second model at the same output repo:
-  ```bash
-  uv run rolm-ocr.py     my-dataset my-dataset --max-samples 100
-  uv run nanonets-ocr.py my-dataset my-dataset --max-samples 100   # appends
-  ```
-- **Reproducible sampling** — `--shuffle` (with `--seed`, default 42) draws a representative sample instead of the first N rows.
-- **Automatic dataset cards** — every run writes a card with the model config, processing stats, column descriptions, and a reproduction command.
-
-## More examples
+### Automatic Dataset Cards
+
+Every OCR run now generates comprehensive dataset documentation including:
+
+- Model configuration and parameters
+- Processing statistics
+- Column descriptions
+- Reproduction instructions
+
+## 💻 Usage Examples
+
+### Run on HuggingFace Jobs (Recommended)
+
+No GPU? No problem! Run on HF infrastructure:
 
 ```bash
-# DeepSeek-OCR on historical scans, large resolution mode
-hf jobs uv run --flavor a100-large -s HF_TOKEN -e UV_TORCH_BACKEND=auto \
+# PaddleOCR-VL - Smallest model (0.9B) with task modes
+hf jobs uv run --flavor l4x1 \
+    --secrets HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/paddleocr-vl.py \
+    your-input-dataset your-output-dataset \
+    --task-mode ocr \
+    --max-samples 100
+
+# PaddleOCR-VL - Extract tables from documents
+hf jobs uv run --flavor l4x1 \
+    --secrets HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/paddleocr-vl.py \
+    documents tables-dataset \
+    --task-mode table
+
+# PaddleOCR-VL - Formula recognition
+hf jobs uv run --flavor l4x1 \
+    --secrets HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/paddleocr-vl.py \
+    scientific-papers formulas-extracted \
+    --task-mode formula \
+    --batch-size 32
+
+# GLM-OCR - SOTA 0.9B model (94.62% OmniDocBench)
+hf jobs uv run --flavor l4x1 \
+    -s HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/glm-ocr.py \
+    your-input-dataset your-output-dataset \
+    --batch-size 16 \
+    --max-samples 100
+
+# DeepSeek-OCR - Real-world example (National Library of Scotland handbooks)
+hf jobs uv run --flavor a100-large \
+    -s HF_TOKEN \
+    -e UV_TORCH_BACKEND=auto \
     https://huggingface.co/datasets/uv-scripts/ocr/raw/main/deepseek-ocr-vllm.py \
-    NationalLibraryOfScotland/Britain-and-UK-Handbooks-Dataset out \
-    --max-samples 100 --shuffle --resolution-mode large
+    NationalLibraryOfScotland/Britain-and-UK-Handbooks-Dataset \
+    davanstrien/handbooks-deep-ocr \
+    --max-samples 100 \
+    --shuffle \
+    --resolution-mode large
 
-# dots.mocr — SVG generation from charts/figures
-hf jobs uv run --flavor l4x1 -s HF_TOKEN \
-    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/dots-mocr.py \
-    your-charts svg-output --prompt-mode svg --model rednote-hilab/dots.mocr-svg
+# DeepSeek-OCR - Fast testing with tiny mode
+hf jobs uv run --flavor l4x1 \
+    -s HF_TOKEN \
+    -e UV_TORCH_BACKEND=auto \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/deepseek-ocr-vllm.py \
+    your-input-dataset your-output-dataset \
+    --max-samples 10 \
+    --resolution-mode tiny
 
-# Qianfan — key-information extraction
-hf jobs uv run --flavor l4x1 -s HF_TOKEN \
-    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/qianfan-ocr.py \
-    invoices extracted-fields \
-    --prompt-mode kie --custom-prompt "Extract: name, date, total. Output as JSON."
+# DeepSeek-OCR - Parse figures from scientific papers
+hf jobs uv run --flavor a100-large \
+    -s HF_TOKEN \
+    -e UV_TORCH_BACKEND=auto \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/deepseek-ocr-vllm.py \
+    scientific-papers figures-extracted \
+    --prompt-mode figure
+
+# Basic OCR job with Nanonets
+hf jobs uv run --flavor l4x1 \
+    --secrets HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/nanonets-ocr.py \
+    your-input-dataset your-output-dataset
+
+# DoTS.ocr - Multilingual OCR with compact 1.7B model
+hf jobs uv run --flavor a100-large \
+    --secrets HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/dots-ocr.py \
+    davanstrien/ufo-ColPali \
+    your-username/ufo-ocr \
+    --batch-size 256 \
+    --max-samples 1000 \
+    --shuffle
+
+# Real example with UFO dataset 🛸
+hf jobs uv run \
+    --flavor a10g-large \
+    --secrets HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/nanonets-ocr.py \
+    davanstrien/ufo-ColPali \
+    your-username/ufo-ocr \
+    --image-column image \
+    --max-model-len 16384 \
+    --batch-size 128
+
+# Nanonets OCR2 - Next-gen quality with 3B model
+hf jobs uv run \
+    --flavor l4x1 \
+    --secrets HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/nanonets-ocr2.py \
+    your-input-dataset \
+    your-output-dataset \
+    --batch-size 16
+
+# NuMarkdown with reasoning traces for complex documents
+hf jobs uv run \
+    --flavor l4x4 \
+    --secrets HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/numarkdown-ocr.py \
+    your-input-dataset your-output-dataset \
+    --max-samples 50 \
+    --include-thinking \
+    --shuffle
+
+# olmOCR2 - High-quality OCR with YAML metadata
+hf jobs uv run \
+    --flavor a100-large \
+    --secrets HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/olmocr2-vllm.py \
+    your-input-dataset your-output-dataset \
+    --batch-size 16 \
+    --max-samples 100
+
+# Private dataset with custom settings
+hf jobs uv run --flavor l40sx1 \
+    --secrets HF_TOKEN \
+    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/nanonets-ocr.py \
+    private-input private-output \
+    --private \
+    --batch-size 32
 ```
 
-**Python API:**
+### Python API
 
 ```python
 from huggingface_hub import run_uv_job
@@ -268,17 +612,36 @@ from huggingface_hub import run_uv_job
 job = run_uv_job(
     "https://huggingface.co/datasets/uv-scripts/ocr/raw/main/nanonets-ocr.py",
     args=["input-dataset", "output-dataset", "--batch-size", "16"],
-    flavor="l4x1",
+    flavor="l4x1"
 )
 ```
 
-**Run locally** (needs your own GPU) — same scripts, run directly from the URL:
+### Run Locally (Requires GPU)
 
 ```bash
-uv run https://huggingface.co/datasets/uv-scripts/ocr/raw/main/glm-ocr.py \
+# Clone and run
+git clone https://huggingface.co/datasets/uv-scripts/ocr
+cd ocr
+uv run nanonets-ocr.py input-dataset output-dataset
+
+# Or run directly from URL
+uv run https://huggingface.co/datasets/uv-scripts/ocr/raw/main/nanonets-ocr.py \
     input-dataset output-dataset
+
+# PaddleOCR-VL for task-specific OCR (smallest model!)
+uv run paddleocr-vl.py documents extracted --task-mode ocr
+uv run paddleocr-vl.py papers tables --task-mode table  # Extract tables
+uv run paddleocr-vl.py textbooks formulas --task-mode formula  # LaTeX formulas
+
+# RolmOCR for fast text extraction
+uv run rolm-ocr.py documents extracted-text
+uv run rolm-ocr.py images texts --shuffle --max-samples 100  # Random sample
+
+# Nanonets OCR2 for highest quality
+uv run nanonets-ocr2.py documents ocr-results
+
 ```
 
----
+</details>
 
-Works with any Hugging Face dataset containing images — documents, forms, receipts, books, handwriting.
+Works with any HuggingFace dataset containing images — documents, forms, receipts, books, handwriting.

SERVING.md

@@ -1,93 +0,0 @@
-# Server-mode OCR: official support per model + the `-server.py` recipes
-
-The recipes in this folder come in two shapes:
-
-- **Offline batch** (`<model>.py`): the script owns the vLLM engine (`LLM.generate`)
-  and processes the dataset in fixed-size batches. One `hf jobs uv run` command.
-- **Server + driver** (`<model>-server.py`): one job starts `vllm serve` in the
-  background, then a lightweight driver (no torch/vllm deps — installs in seconds)
-  posts images concurrently to the OpenAI-compatible endpoint on localhost.
-
-Why server mode? Two measured reasons (100-page historical-scan A/B on `l4x1`,
-same data, same sampling, driver concurrency 32):
-
-| Model | Offline (inference-only) | Server | Speedup | Output parity |
-|---|---|---|---|---|
-| OvisOCR2 0.9B | 0.83 img/s | 1.44 img/s | ~1.7× | 94/100 byte-identical, rest ≥0.978 similar |
-| LightOnOCR-2 1B | 0.33 img/s | 0.59 img/s | ~1.8× | 64/100 byte-identical at temp 0.2, rest median 0.998 |
-| Nanonets-OCR2 3B | 0.31 img/s (steady-state) | 0.36 img/s | ~1.2× | 52/100 byte-identical, median 0.999; 3 hard plate pages fork under greedy (worst case was the *offline* arm degenerating into a 22k-char repetition loop) |
-
-1. **Throughput**: offline `llm.generate` drains at every batch boundary (GPU idles
-   while the CPU decodes the next batch of images); a concurrent driver keeps vLLM's
-   continuous batching fed. The speedup is a floor, not a ceiling — at concurrency 32
-   the server's KV cache was <10% used.
-2. **Failure isolation**: offline, one bad image (e.g. a None cell) fails its whole
-   batch of 16; server mode fails that one request. On a real dataset where ~half the
-   image cells were empty, the offline recipe produced 0 usable outputs and the server
-   recipe produced all of the valid ones.
-
-The job command stays the standard `hf jobs uv run` shape: the driver spawns
-`vllm serve` itself as a subprocess when no server is reachable (flags live in the
-script's `SERVE_ARGS`, taken from the model's own card where they exist), so the only
-thing to get right is `--image vllm/vllm-openai:<tag>` — which provides the `vllm`
-binary. Without it the script fails fast, printing the exact correct command. Pass
-`--server URL` to use an already-running or remote endpoint instead (nothing is
-spawned when the server is reachable, so the moss-style explicit `bash -c` serve
-command keeps working too).
-
-For **interactive/agent** use (a live endpoint instead of a batch run), see
-[serving-unlimited-ocr.md](serving-unlimited-ocr.md) — `hf jobs run --expose` gives an
-OpenAI-compatible URL that outlives a single script.
-
-## The recurring official serve pattern for OCR
-
-Three flags recur across independent vendors' official serve commands (DeepSeek's vLLM
-recipe, LightOn's card, Paddle's vLLM recipe, Unlimited-OCR's recipe):
-
-```
---no-enable-prefix-caching --mm-processor-cache-gb 0 --limit-mm-per-prompt '{"image": 1}'
-```
-
-OCR workloads never reuse images, so prefix/multimodal caches only cost memory.
-
-## Version pins carry over — via the image tag
-
-Where an offline recipe pins an engine version, the server variant needs the same pin
-as a `vllm/vllm-openai:<tag>` image (e.g. Nanonets-OCR2-3B uses `v0.10.2`, matching its
-offline recipe). [`models.json`](models.json) records the required image per script.
-When trying a new model/image combination, sanity-check the first few outputs before
-scaling: a mismatched combination can produce degenerate output while looking like
-normal load from the outside (full GPU utilisation, no errors).
-
-## Which models document server mode themselves? (surveyed 2026-07-16)
-
-Server mode is the officially documented path for most models in this collection —
-for several of them it's the *only* documented vLLM path. Summary of each model's own
-card/docs (verbatim commands live in the linked sources):
-
-| Model | Official server example | Notes |
-|---|---|---|
-| lightonai/LightOnOCR-1B / 2-1B | ✅ vLLM | Card ships the serve command + client; ≥0.11.1 for v1; images longest-dim 1540px |
-| nanonets/OCR-s / OCR2-3B | ✅ vLLM | Bare `vllm serve` + client in card; **pin v0.10.2 for OCR2** (see above) |
-| rednote-hilab/dots.mocr | ✅ vLLM ≥0.11 | Direct from Hub id on `vllm/vllm-openai:v0.11.0`+ |
-| rednote-hilab/dots.ocr | ✅ (GitHub) | HF card shows a legacy pre-0.11 hack; GitHub README: integrated upstream since 0.11 |
-| tencent/HunyuanOCR | ✅ vLLM 0.18.1 | serve.sh in repo; nightly adds DFlash speculative decoding |
-| zai-org/GLM-OCR | ✅ vLLM + SGLang + Ollama | Only card here with an SGLang serve example; needs vLLM nightly |
-| numind/NuExtract3 | ✅ vLLM | Production-grade recipe: MTP speculative decoding, per-request `chat_template_kwargs` |
-| numind/NuMarkdown-8B-Thinking | ✅ vLLM | Thinking always on — parse `<think>`/`<answer>` |
-| baidu/Unlimited-OCR | ✅ vLLM + SGLang | Custom image / dev wheel; see [serving-unlimited-ocr.md](serving-unlimited-ocr.md) |
-| baidu/Qianfan-OCR | ✅ vLLM (minimal) | Needs `--hf-overrides '{"architectures": ["InternVLChatModel"]}'` |
-| PaddlePaddle/PaddleOCR-VL 1.x | ✅ paddle `genai_server` + vLLM recipe | Serves the 0.9B VLM only; raw serve skips the layout stage (official quality warning) |
-| deepseek-ai/DeepSeek-OCR / -2 | ⚠️ vLLM recipe pages only | docs.vllm.ai recipes; needs the DeepSeek n-gram logits processor flags |
-| allenai/olmOCR-2-7B-FP8 | ✅ (GitHub) | olmocr toolkit spawns `vllm serve` itself; YAML-front-matter prompt required |
-| reducto/RolmOCR | ✅ vLLM | Card serve + client (client's model string has a typo — pass `--served-model-name`) |
-| tiiuae/Falcon-OCR | ✅ vLLM in official Docker | `ghcr.io/tiiuae/falcon-ocr`; task-token prompts |
-| datalab-to/surya-ocr-2, lift | ⚠️ wrapper-mediated | Server-native but driven via their own managers (`SURYA_INFERENCE_URL`, `lift_vllm`) |
-| LiquidAI LFM2.5-VL-Extract | ⚠️ family docs | vLLM ≥0.23 + SGLang cookbooks, not Extract-specific |
-| ATH-MaaS/OvisOCR2 | ❌ offline vLLM only | `ovis-ocr2-server.py` is our translation of the card's offline args (parity-validated, table above) |
-| FireRedTeam/FireRed-OCR | ❌ transformers only | |
-| acvlab/ABot-OCR | ❌ offline script only | pinned vllm 0.18.0 |
-
-SGLang reality check: only GLM-OCR and Unlimited-OCR ship SGLang serve paths (olmOCR
-dropped SGLang for vLLM in v0.1.75). For this collection, "server mode" effectively
-means a vLLM OpenAI-compatible endpoint.

abot-ocr.py

@@ -1,600 +0,0 @@
-# /// script
-# requires-python = ">=3.11"
-# dependencies = [
-#     "datasets>=4.0.0",
-#     "huggingface-hub",
-#     "pillow",
-#     "vllm>=0.15.1",
-#     "tqdm",
-#     "toolz",
-#     "torch",
-# ]
-#
-# ///
-
-"""
-Convert document images to Markdown using ABot-OCR with vLLM.
-
-ABot-OCR (acvlab/ABot-OCR) is a compact Qwen3-VL-based document OCR model that
-converts PDF/document page images into structured Markdown, including:
-- Text with original document structure (headings, paragraphs, lists)
-- Mathematical formulas in LaTeX (inline \\( \\) and block \\[ \\])
-- Tables in HTML (<table>…</table>)
-
-Model:  https://huggingface.co/acvlab/ABot-OCR
-Paper:  https://arxiv.org/abs/2605.27978
-Code:   https://github.com/amap-cvlab/ABot-OCR
-
-HF Jobs note: ABot-OCR (Qwen3-VL) uses vLLM's flashinfer sampler, which needs CUDA
-kernels the default uv-script image can't build ("Could not find nvcc"). Run on Jobs
-with the pre-built vLLM image so the kernels are reused (same image-mode pattern as
-nuextract3.py / paddleocr-vl-1.6.py):
-    --image vllm/vllm-openai:latest
-    --python /usr/bin/python3 -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages
-"""
-
-import argparse
-import base64
-import io
-import json
-import logging
-import os
-import sys
-from typing import Any, Dict, List, Union
-from datetime import datetime
-
-import torch
-from datasets import load_dataset
-from huggingface_hub import DatasetCard, login
-from PIL import Image
-from toolz import partition_all
-from tqdm.auto import tqdm
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
-from vllm import LLM, SamplingParams
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# ABot-OCR's recommended document→Markdown prompt (from the model's reference
-# inference script; LaTeX delimiters de-mangled to proper \( \) and \[ \]).
-DEFAULT_PROMPT = r"""You are an AI assistant specialized in converting PDF images to Markdown format. Please follow these instructions for the conversion:
-
-1. Text Processing:
-- Accurately recognize all text content in the PDF image without guessing or inferring.
-- Convert the recognized text into Markdown format.
-- Maintain the original document structure, including headings, paragraphs, lists, etc.
-
-2. Mathematical Formula Processing:
-- Convert all mathematical formulas to LaTeX format.
-- Enclose inline formulas with \( \). For example: This is an inline formula \( E = mc^2 \)
-- Enclose block formulas with \[ \]. For example: \[ \frac{-b \pm \sqrt{b^2 - 4ac}}{2a} \]
-
-3. Table Processing:
-- Convert tables to HTML format.
-- Wrap the entire table with <table> and </table>.
-
-4. Figure Handling:
-- Ignore figures content in the PDF image. Do not attempt to describe or convert images.
-
-5. Output Format:
-- Ensure the output Markdown document has a clear structure with appropriate line breaks between elements.
-- For complex layouts, try to maintain the original document's structure and format as closely as possible.
-
-Please strictly follow these guidelines to ensure accuracy and consistency in the conversion. Your task is to accurately convert the content of the PDF image into Markdown format without adding any extra explanations or comments."""
-
-
-def check_cuda_availability():
-    """Check if CUDA is available and exit if not."""
-    if not torch.cuda.is_available():
-        logger.error("CUDA is not available. This script requires a GPU.")
-        logger.error("Please run on a machine with a CUDA-capable GPU.")
-        sys.exit(1)
-    else:
-        logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
-
-
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
-def post_process_text(text: str, threshold: int = 8000) -> str:
-    """Trim runaway repetition loops that VLM OCR can fall into on dense pages.
-
-    Mirrors ABot-OCR's reference post-processing: if the tail of a long output
-    repeats the same short substring many times, collapse it to a single copy.
-    """
-    n = len(text)
-    if n < threshold:
-        return text
-    for length in range(2, n // 10 + 1):
-        candidate = text[-length:]
-        count = 0
-        i = n - length
-        while i >= 0 and text[i:i + length] == candidate:
-            count += 1
-            i -= length
-        if count >= 10:
-            return text[: n - length * (count - 1)]
-    return text
-
-
-def make_ocr_message(
-    image: Union[Image.Image, Dict[str, Any], str],
-    prompt: str = DEFAULT_PROMPT,
-) -> List[Dict]:
-    """Create a vLLM chat message for OCR processing."""
-    # Convert to PIL Image if needed
-    if isinstance(image, Image.Image):
-        pil_img = image
-    elif isinstance(image, dict) and "bytes" in image:
-        pil_img = Image.open(io.BytesIO(image["bytes"]))
-    elif isinstance(image, str):
-        pil_img = Image.open(image)
-    else:
-        raise ValueError(f"Unsupported image type: {type(image)}")
-
-    if pil_img.mode != "RGB":
-        pil_img = pil_img.convert("RGB")
-
-    # Convert to base64 data URI
-    buf = io.BytesIO()
-    pil_img.save(buf, format="PNG")
-    data_uri = f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}"
-
-    return [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image_url", "image_url": {"url": data_uri}},
-                {"type": "text", "text": prompt},
-            ],
-        }
-    ]
-
-
-def create_dataset_card(
-    source_dataset: str,
-    model: str,
-    num_samples: int,
-    processing_time: str,
-    batch_size: int,
-    max_model_len: int,
-    max_tokens: int,
-    gpu_memory_utilization: float,
-    image_column: str = "image",
-    split: str = "train",
-) -> str:
-    """Create a dataset card documenting the OCR process."""
-    model_name = model.split("/")[-1]
-
-    return f"""---
-tags:
-- ocr
-- document-processing
-- abot
-- abot-ocr
-- markdown
-- uv-script
-- generated
----
-
-# Document OCR using {model_name}
-
-This dataset contains Markdown-formatted OCR results from images in [{source_dataset}](https://huggingface.co/datasets/{source_dataset}) using [ABot-OCR](https://huggingface.co/{model}).
-
-## Processing Details
-
-- **Source Dataset**: [{source_dataset}](https://huggingface.co/datasets/{source_dataset})
-- **Model**: [{model}](https://huggingface.co/{model})
-- **Paper**: [arxiv.org/abs/2605.27978](https://arxiv.org/abs/2605.27978)
-- **Number of Samples**: {num_samples:,}
-- **Processing Time**: {processing_time}
-- **Processing Date**: {datetime.now().strftime("%Y-%m-%d %H:%M UTC")}
-
-### Configuration
-
-- **Image Column**: `{image_column}`
-- **Output Column**: `markdown`
-- **Dataset Split**: `{split}`
-- **Batch Size**: {batch_size}
-- **Max Model Length**: {max_model_len:,} tokens
-- **Max Output Tokens**: {max_tokens:,}
-- **GPU Memory Utilization**: {gpu_memory_utilization:.1%}
-
-## Model Information
-
-ABot-OCR is a compact Qwen3-VL-based document OCR model that converts page images to Markdown:
-- 📐 **LaTeX equations** — inline `\\( \\)` and block `\\[ \\]`
-- 📊 **Tables** — extracted as HTML (`<table>…</table>`)
-- 📝 **Document structure** — headings, paragraphs, and lists preserved
-
-## Dataset Structure
-
-The dataset contains all original columns plus:
-- `markdown`: The extracted text in Markdown format with preserved structure
-- `inference_info`: JSON list tracking all OCR models applied to this dataset
-
-## Usage
-
-```python
-from datasets import load_dataset
-
-dataset = load_dataset("{{{{output_dataset_id}}}}", split="{split}")
-for example in dataset:
-    print(example["markdown"])
-    break
-```
-
-## Reproduction
-
-This dataset was generated using the [uv-scripts/ocr](https://huggingface.co/datasets/uv-scripts/ocr) ABot-OCR script:
-
-```bash
-uv run https://huggingface.co/datasets/uv-scripts/ocr/raw/main/abot-ocr.py \\
-    {source_dataset} \\
-    <output-dataset> \\
-    --image-column {image_column} \\
-    --batch-size {batch_size} \\
-    --max-model-len {max_model_len} \\
-    --max-tokens {max_tokens} \\
-    --gpu-memory-utilization {gpu_memory_utilization}
-```
-
-Generated with 🤖 [UV Scripts](https://huggingface.co/uv-scripts)
-"""
-
-
-def main(
-    input_dataset: str,
-    output_dataset: str,
-    image_column: str = "image",
-    batch_size: int = 16,
-    model: str = "acvlab/ABot-OCR",
-    max_model_len: int = 16384,
-    max_tokens: int = 8192,
-    gpu_memory_utilization: float = 0.8,
-    hf_token: str = None,
-    split: str = "train",
-    max_samples: int = None,
-    private: bool = False,
-    shuffle: bool = False,
-    seed: int = 42,
-    output_column: str = "markdown",
-    overwrite: bool = False,
-    verbose: bool = False,
-):
-    """Process images from a HF dataset through the ABot-OCR model."""
-
-    # Check CUDA availability first
-    check_cuda_availability()
-
-    # Track processing start time
-    start_time = datetime.now()
-
-    # Enable high-performance Xet downloads
-    os.environ["HF_XET_HIGH_PERFORMANCE"] = "1"
-
-    # Login to HF if token provided
-    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")
-    if HF_TOKEN:
-        login(token=HF_TOKEN)
-
-    # Load dataset
-    logger.info(f"Loading dataset: {input_dataset}")
-    dataset = load_dataset(input_dataset, split=split)
-
-    # Validate image column
-    if image_column not in dataset.column_names:
-        raise ValueError(
-            f"Column '{image_column}' not found. Available: {dataset.column_names}"
-        )
-
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
-    # Shuffle if requested
-    if shuffle:
-        logger.info(f"Shuffling dataset with seed {seed}")
-        dataset = dataset.shuffle(seed=seed)
-
-    # Limit samples if requested
-    if max_samples:
-        dataset = dataset.select(range(min(max_samples, len(dataset))))
-        logger.info(f"Limited to {len(dataset)} samples")
-
-    # Initialize vLLM
-    logger.info(f"Initializing vLLM with model: {model}")
-    llm = LLM(
-        model=model,
-        trust_remote_code=True,
-        max_model_len=max_model_len,
-        gpu_memory_utilization=gpu_memory_utilization,
-        limit_mm_per_prompt={"image": 1},
-        enforce_eager=True,  # avoid warmup kernel compilation (no nvcc in the image)
-    )
-
-    sampling_params = SamplingParams(
-        temperature=0.0,  # Deterministic for OCR
-        max_tokens=max_tokens,
-    )
-
-    # Process images in batches
-    all_markdown = []
-
-    logger.info(f"Processing {len(dataset)} images in batches of {batch_size}")
-
-    for batch_indices in tqdm(
-        partition_all(batch_size, range(len(dataset))),
-        total=(len(dataset) + batch_size - 1) // batch_size,
-        desc="OCR processing",
-    ):
-        batch_indices = list(batch_indices)
-        batch_images = [dataset[i][image_column] for i in batch_indices]
-
-        try:
-            batch_messages = [make_ocr_message(img) for img in batch_images]
-            outputs = llm.chat(batch_messages, sampling_params)
-            for output in outputs:
-                markdown_text = post_process_text(output.outputs[0].text.strip())
-                all_markdown.append(markdown_text)
-        except Exception as e:
-            logger.error(f"Error processing batch: {e}")
-            all_markdown.extend(["[OCR FAILED]"] * len(batch_images))
-
-    # Add output column to dataset
-    logger.info(f"Adding '{output_column}' column to dataset")
-    dataset = dataset.add_column(output_column, all_markdown)
-
-    # Handle inference_info tracking
-    logger.info("Updating inference_info...")
-
-    inference_entry = {
-        "model_id": model,
-        "model_name": "ABot-OCR",
-        "column_name": output_column,
-        "timestamp": datetime.now().isoformat(),
-        "batch_size": batch_size,
-        "max_tokens": max_tokens,
-        "gpu_memory_utilization": gpu_memory_utilization,
-        "max_model_len": max_model_len,
-        "script": "abot-ocr.py",
-        "script_url": "https://huggingface.co/datasets/uv-scripts/ocr/raw/main/abot-ocr.py",
-    }
-
-    if "inference_info" in dataset.column_names:
-        logger.info("Updating existing inference_info column")
-
-        def update_inference_info(example):
-            try:
-                existing_info = (
-                    json.loads(example["inference_info"])
-                    if example["inference_info"]
-                    else []
-                )
-            except (json.JSONDecodeError, TypeError):
-                existing_info = []
-            existing_info.append(inference_entry)
-            return {"inference_info": json.dumps(existing_info)}
-
-        dataset = dataset.map(update_inference_info)
-    else:
-        logger.info("Creating new inference_info column")
-        inference_list = [json.dumps([inference_entry])] * len(dataset)
-        dataset = dataset.add_column("inference_info", inference_list)
-
-    # Push to hub
-    logger.info(f"Pushing to {output_dataset}")
-    dataset.push_to_hub(output_dataset, private=private, token=HF_TOKEN)
-
-    # Calculate processing time
-    end_time = datetime.now()
-    processing_duration = end_time - start_time
-    processing_time = f"{processing_duration.total_seconds() / 60:.1f} minutes"
-
-    # Create and push dataset card
-    logger.info("Creating dataset card...")
-    card_content = create_dataset_card(
-        source_dataset=input_dataset,
-        model=model,
-        num_samples=len(dataset),
-        processing_time=processing_time,
-        batch_size=batch_size,
-        max_model_len=max_model_len,
-        max_tokens=max_tokens,
-        gpu_memory_utilization=gpu_memory_utilization,
-        image_column=image_column,
-        split=split,
-    )
-
-    card = DatasetCard(card_content)
-    card.push_to_hub(output_dataset, token=HF_TOKEN)
-    logger.info("✅ Dataset card created and pushed!")
-
-    logger.info("✅ OCR conversion complete!")
-    logger.info(
-        f"Dataset available at: https://huggingface.co/datasets/{output_dataset}"
-    )
-
-    if verbose:
-        import importlib.metadata
-
-        logger.info("--- Resolved package versions ---")
-        for pkg in ["vllm", "transformers", "torch", "datasets", "pyarrow", "pillow"]:
-            try:
-                logger.info(f"  {pkg}=={importlib.metadata.version(pkg)}")
-            except importlib.metadata.PackageNotFoundError:
-                logger.info(f"  {pkg}: not installed")
-        logger.info("--- End versions ---")
-
-
-if __name__ == "__main__":
-    # Show example usage if no arguments
-    if len(sys.argv) == 1:
-        print("=" * 80)
-        print("ABot-OCR to Markdown Converter")
-        print("=" * 80)
-        print("\nConverts document images to structured Markdown using the")
-        print("ABot-OCR model (Qwen3-VL based) with vLLM acceleration.")
-        print("\nFeatures:")
-        print("- Document structure preserved (headings, paragraphs, lists)")
-        print("- LaTeX equation recognition (inline \\( \\) and block \\[ \\])")
-        print("- Table extraction as HTML")
-        print("\nExample usage:")
-        print("\n1. Basic OCR conversion:")
-        print("   uv run abot-ocr.py document-images markdown-docs")
-        print("\n2. Process a subset for testing:")
-        print("   uv run abot-ocr.py large-dataset test-output --max-samples 10")
-        print("\n3. Running on HF Jobs (use the pre-built vLLM image for flashinfer kernels):")
-        print("   hf jobs uv run --flavor l4x1 \\")
-        print("     --image vllm/vllm-openai:latest \\")
-        print(
-            "     --python /usr/bin/python3 -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages \\"
-        )
-        print("     -s HF_TOKEN \\")
-        print(
-            "     https://huggingface.co/datasets/uv-scripts/ocr/raw/main/abot-ocr.py \\"
-        )
-        print("       your-document-dataset \\")
-        print("       your-markdown-output \\")
-        print("       --max-samples 10")
-        print("\n" + "=" * 80)
-        print("\nFor full help, run: uv run abot-ocr.py --help")
-        sys.exit(0)
-
-    parser = argparse.ArgumentParser(
-        description="OCR images to Markdown using ABot-OCR (Qwen3-VL)",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog="""
-Examples:
-  # Basic usage
-  uv run abot-ocr.py my-images-dataset ocr-results
-
-  # With specific image column
-  uv run abot-ocr.py documents extracted-text --image-column scan
-
-  # Process subset for testing
-  uv run abot-ocr.py large-dataset test-output --max-samples 100
-
-  # Random sample from ordered dataset
-  uv run abot-ocr.py ordered-dataset random-sample --max-samples 50 --shuffle
-        """,
-    )
-
-    parser.add_argument("input_dataset", help="Input dataset ID from Hugging Face Hub")
-    parser.add_argument("output_dataset", help="Output dataset ID for Hugging Face Hub")
-    parser.add_argument(
-        "--image-column",
-        default="image",
-        help="Column containing images (default: image)",
-    )
-    parser.add_argument(
-        "--batch-size",
-        type=int,
-        default=16,
-        help="Batch size for processing (default: 16)",
-    )
-    parser.add_argument(
-        "--model",
-        default="acvlab/ABot-OCR",
-        help="Model to use (default: acvlab/ABot-OCR)",
-    )
-    parser.add_argument(
-        "--max-model-len",
-        type=int,
-        default=16384,
-        help="Maximum model context length (default: 16384)",
-    )
-    parser.add_argument(
-        "--max-tokens",
-        type=int,
-        default=8192,
-        help="Maximum tokens to generate (default: 8192)",
-    )
-    parser.add_argument(
-        "--gpu-memory-utilization",
-        type=float,
-        default=0.8,
-        help="GPU memory utilization (default: 0.8)",
-    )
-    parser.add_argument("--hf-token", help="Hugging Face API token")
-    parser.add_argument(
-        "--split", default="train", help="Dataset split to use (default: train)"
-    )
-    parser.add_argument(
-        "--max-samples",
-        type=int,
-        help="Maximum number of samples to process (for testing)",
-    )
-    parser.add_argument(
-        "--private", action="store_true", help="Make output dataset private"
-    )
-    parser.add_argument(
-        "--shuffle",
-        action="store_true",
-        help="Shuffle the dataset before processing (useful for random sampling)",
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=42,
-        help="Random seed for shuffling (default: 42)",
-    )
-    parser.add_argument(
-        "--output-column",
-        default="markdown",
-        help="Column name for the OCR output text (default: markdown)",
-    )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
-    parser.add_argument(
-        "--verbose",
-        action="store_true",
-        help="Log resolved package versions after processing (useful for pinning deps)",
-    )
-
-    args = parser.parse_args()
-
-    main(
-        input_dataset=args.input_dataset,
-        output_dataset=args.output_dataset,
-        image_column=args.image_column,
-        batch_size=args.batch_size,
-        model=args.model,
-        max_model_len=args.max_model_len,
-        max_tokens=args.max_tokens,
-        gpu_memory_utilization=args.gpu_memory_utilization,
-        hf_token=args.hf_token,
-        split=args.split,
-        max_samples=args.max_samples,
-        private=args.private,
-        shuffle=args.shuffle,
-        seed=args.seed,
-        output_column=args.output_column,
-        overwrite=args.overwrite,
-        verbose=args.verbose,
-    )

deepseek-ocr-vllm.py

@@ -48,10 +48,6 @@ from huggingface_hub import DatasetCard, login
 from PIL import Image
 from toolz import partition_all
 from tqdm.auto import tqdm
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
 from vllm import LLM, SamplingParams
 from vllm.model_executor.models.deepseek_ocr import NGramPerReqLogitsProcessor
 
@@ -78,28 +74,6 @@ def check_cuda_availability():
         logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
 
 
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
 def to_pil(image: Union[Image.Image, Dict[str, Any], str]) -> Image.Image:
     """Convert various image formats to PIL Image."""
     if isinstance(image, Image.Image):
@@ -233,8 +207,6 @@ def main(
     private: bool = False,
     shuffle: bool = False,
     seed: int = 42,
-    output_column: str = "markdown",
-    overwrite: bool = False,
     config: str = None,
     create_pr: bool = False,
     verbose: bool = False,
@@ -277,9 +249,6 @@ def main(
             f"Column '{image_column}' not found. Available: {dataset.column_names}"
         )
 
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
     # Shuffle if requested
     if shuffle:
         logger.info(f"Shuffling dataset with seed {seed}")
@@ -353,9 +322,9 @@ def main(
     processing_duration = datetime.now() - start_time
     processing_time_str = f"{processing_duration.total_seconds() / 60:.1f} min"
 
-    # Add output column to dataset
-    logger.info(f"Adding '{output_column}' column to dataset")
-    dataset = dataset.add_column(output_column, all_markdown)
+    # Add markdown column to dataset
+    logger.info("Adding markdown column to dataset")
+    dataset = dataset.add_column("markdown", all_markdown)
 
     # Handle inference_info tracking
     logger.info("Updating inference_info...")
@@ -363,7 +332,7 @@ def main(
     inference_entry = {
         "model_id": model,
         "model_name": "DeepSeek-OCR",
-        "column_name": output_column,
+        "column_name": "markdown",
         "timestamp": datetime.now().isoformat(),
         "prompt_mode": prompt_mode if prompt is None else "custom",
         "batch_size": batch_size,
@@ -596,17 +565,6 @@ Examples:
         default=42,
         help="Random seed for shuffling (default: 42)",
     )
-    parser.add_argument(
-        "--output-column",
-        default="markdown",
-        help="Column name for the OCR output text (default: markdown)",
-    )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
     parser.add_argument(
         "--verbose",
         action="store_true",
@@ -632,8 +590,6 @@ Examples:
         private=args.private,
         shuffle=args.shuffle,
         seed=args.seed,
-        output_column=args.output_column,
-        overwrite=args.overwrite,
         config=args.config,
         create_pr=args.create_pr,
         verbose=args.verbose,

deepseek-ocr.py

@@ -1,8 +1,8 @@
 # /// script
 # requires-python = ">=3.11"
 # dependencies = [
-#     "datasets>=4.0.0",
-#     "huggingface-hub",
+#     "datasets",
+#     "huggingface-hub[hf_transfer]",
 #     "pillow",
 #     "torch",
 #     "torchvision",
@@ -76,28 +76,6 @@ def check_cuda_availability():
         logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
 
 
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
 def create_dataset_card(
     source_dataset: str,
     model: str,
@@ -265,8 +243,6 @@ def main(
     private: bool = False,
     shuffle: bool = False,
     seed: int = 42,
-    output_column: str = "markdown",
-    overwrite: bool = False,
 ):
     """Process images from HF dataset through DeepSeek-OCR model."""
 
@@ -276,7 +252,8 @@ def main(
     # Track processing start time
     start_time = datetime.now()
 
-
+    # Enable HF_TRANSFER for faster downloads
+    os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 
     # Login to HF if token provided
     HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")
@@ -318,9 +295,6 @@ def main(
             f"Column '{image_column}' not found. Available: {dataset.column_names}"
         )
 
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
     # Shuffle if requested
     if shuffle:
         logger.info(f"Shuffling dataset with seed {seed}")
@@ -406,12 +380,12 @@ def main(
         try:
             shutil.rmtree(temp_dir)
             shutil.rmtree(temp_output_dir)
-        except Exception:
+        except:
             pass
 
-    # Add output column to dataset
-    logger.info(f"Adding '{output_column}' column to dataset")
-    dataset = dataset.add_column(output_column, all_markdown)
+    # Add markdown column to dataset
+    logger.info("Adding markdown column to dataset")
+    dataset = dataset.add_column("markdown", all_markdown)
 
     # Handle inference_info tracking
     logger.info("Updating inference_info...")
@@ -430,7 +404,7 @@ def main(
 
     # Add new inference info
     new_info = {
-        "column_name": output_column,
+        "column_name": "markdown",
         "model_id": model,
         "processing_date": datetime.now().isoformat(),
         "resolution_mode": resolution_mode,
@@ -608,17 +582,6 @@ Examples:
         default=42,
         help="Random seed for shuffling (default: 42)",
     )
-    parser.add_argument(
-        "--output-column",
-        default="markdown",
-        help="Column name for the OCR output text (default: markdown)",
-    )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
 
     args = parser.parse_args()
 
@@ -638,6 +601,4 @@ Examples:
         private=args.private,
         shuffle=args.shuffle,
         seed=args.seed,
-        output_column=args.output_column,
-        overwrite=args.overwrite,
     )

deepseek-ocr2-vllm.py

@@ -59,10 +59,6 @@ from huggingface_hub import DatasetCard, login
 from PIL import Image
 from toolz import partition_all
 from tqdm.auto import tqdm
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
 from vllm import LLM, SamplingParams
 from vllm.model_executor.models.deepseek_ocr import NGramPerReqLogitsProcessor
 
@@ -86,28 +82,6 @@ def check_cuda_availability():
         logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
 
 
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
 def to_pil(image: Union[Image.Image, Dict[str, Any], str]) -> Image.Image:
     """Convert various image formats to PIL Image."""
     if isinstance(image, Image.Image):
@@ -247,7 +221,6 @@ def main(
     shuffle: bool = False,
     seed: int = 42,
     output_column: str = "markdown",
-    overwrite: bool = False,
     config: str = None,
     create_pr: bool = False,
     verbose: bool = False,
@@ -290,9 +263,6 @@ def main(
             f"Column '{image_column}' not found. Available: {dataset.column_names}"
         )
 
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
     # Shuffle if requested
     if shuffle:
         logger.info(f"Shuffling dataset with seed {seed}")
@@ -600,12 +570,6 @@ Examples:
         default="markdown",
         help="Column name for OCR output (default: markdown). Use a different name to add alongside existing OCR.",
     )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
     parser.add_argument(
         "--shuffle",
         action="store_true",
@@ -652,7 +616,6 @@ Examples:
         shuffle=args.shuffle,
         seed=args.seed,
         output_column=args.output_column,
-        overwrite=args.overwrite,
         config=args.config,
         create_pr=args.create_pr,
         verbose=args.verbose,

dots-mocr.py → dots-ocr-1.5.py

@@ -1,42 +1,35 @@
 # /// script
 # requires-python = ">=3.11"
 # dependencies = [
-#     "datasets>=4.0.0",
+#     "datasets",
 #     "huggingface-hub",
 #     "pillow",
 #     "vllm>=0.15.1",
 #     "tqdm",
 #     "toolz",
 #     "torch",
-#     # dots.mocr's remote code (AutoProcessor.register with a string key)
-#     # breaks on transformers v5; vllm pulls transformers unpinned.
-#     "transformers<5",
 # ]
 #
 # ///
 
 """
-Convert document images to markdown using dots.mocr with vLLM.
+Convert document images to markdown using DoTS.ocr-1.5 with vLLM.
 
-dots.mocr is a 3B multilingual document parsing model with SOTA performance
-on 100+ languages. It excels at converting structured graphics (charts, UI
-layouts, scientific figures) directly into SVG code. Core capabilities include
-grounding, recognition, semantic understanding, and interactive dialogue.
+DoTS.ocr-1.5 is a 3B multilingual document parsing model with SOTA performance
+on 100+ languages. Compared to v1 (1.7B), it adds web screen parsing, scene text
+spotting, SVG code generation, and stronger multilingual document parsing.
 
 Features:
 - Multilingual support (100+ languages)
 - Table extraction and formatting
 - Formula recognition
 - Layout-aware text extraction
-- Web screen parsing
-- Scene text spotting
-- SVG code generation (use --prompt-mode svg, or --model rednote-hilab/dots.mocr-svg for best results)
-
-Model: rednote-hilab/dots.mocr
-SVG variant: rednote-hilab/dots.mocr-svg
-vLLM: Officially integrated since v0.11.0
-GitHub: https://github.com/rednote-hilab/dots.mocr
-Paper: https://arxiv.org/abs/2603.13032
+- Web screen parsing (NEW in v1.5)
+- Scene text spotting (NEW in v1.5)
+- SVG code generation (requires dots.ocr-1.5-svg variant)
+
+Model: rednote-hilab/dots.ocr-1.5
+vLLM: Officially supported (same DotsOCRForCausalLM architecture as v1)
 """
 
 import argparse
@@ -53,13 +46,9 @@ from typing import Any, Dict, List, Union
 import torch
 from datasets import load_dataset
 from huggingface_hub import DatasetCard, login
-from PIL import Image, UnidentifiedImageError
+from PIL import Image
 from toolz import partition_all
 from tqdm.auto import tqdm
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
 from vllm import LLM, SamplingParams
 
 logging.basicConfig(level=logging.INFO)
@@ -67,8 +56,8 @@ logger = logging.getLogger(__name__)
 
 
 # ────────────────────────────────────────────────────────────────
-# dots.mocr Prompt Templates
-# Source: https://github.com/rednote-hilab/dots.mocr/blob/master/dots_mocr/utils/prompts.py
+# DoTS OCR 1.5 Prompt Templates (from official dots.ocr repo)
+# Source: https://github.com/rednote-hilab/dots.ocr/blob/master/dots_ocr/utils/prompts.py
 # ────────────────────────────────────────────────────────────────
 
 PROMPT_TEMPLATES = {
@@ -91,19 +80,11 @@ PROMPT_TEMPLATES = {
 
 5. Final Output: The entire output must be a single JSON object.
 """,
-    # NOTE: Bboxes from layout-all/layout-only are in the resized image coordinate
-    # space (Qwen2VLImageProcessor smart_resize: max_pixels=11289600, factor=28),
-    # NOT original image coordinates. To map back, compute:
-    #   resized_h, resized_w = smart_resize(orig_h, orig_w)
-    #   scale_x, scale_y = orig_w / resized_w, orig_h / resized_h
     "layout-only": """Please output the layout information from this PDF image, including each layout's bbox and its category. The bbox should be in the format [x1, y1, x2, y2]. The layout categories for the PDF document include ['Caption', 'Footnote', 'Formula', 'List-item', 'Page-footer', 'Page-header', 'Picture', 'Section-header', 'Table', 'Text', 'Title']. Do not output the corresponding text. The layout result should be in JSON format.""",
+    # NEW in v1.5:
     "web-parsing": """Parsing the layout info of this webpage image with format json:\n""",
     "scene-spotting": """Detect and recognize the text in the image.""",
     "grounding-ocr": """Extract text from the given bounding box on the image (format: [x1, y1, x2, y2]).\nBounding Box:\n""",
-    # SVG code generation — {width} and {height} are replaced with actual image dimensions.
-    # For best results, use --model rednote-hilab/dots.mocr-svg
-    # Uses higher temperature (0.9) and top_p (1.0) per official recommendation.
-    "svg": """Please generate the SVG code based on the image. viewBox="0 0 {width} {height}" """,
     "general": """ """,
 }
 
@@ -118,28 +99,6 @@ def check_cuda_availability():
         logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
 
 
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
 def make_ocr_message(
     image: Union[Image.Image, Dict[str, Any], str],
     prompt: str = PROMPT_TEMPLATES["ocr"],
@@ -158,12 +117,6 @@ def make_ocr_message(
     # Convert to RGB
     pil_img = pil_img.convert("RGB")
 
-    # For SVG mode, inject actual image dimensions into the prompt
-    if "{width}" in prompt and "{height}" in prompt:
-        prompt = prompt.replace("{width}", str(pil_img.width)).replace(
-            "{height}", str(pil_img.height)
-        )
-
     # Convert to base64 data URI
     buf = io.BytesIO()
     pil_img.save(buf, format="PNG")
@@ -201,7 +154,7 @@ def create_dataset_card(
 tags:
 - ocr
 - document-processing
-- dots-mocr
+- dots-ocr-1.5
 - multilingual
 - markdown
 - uv-script
@@ -210,7 +163,7 @@ tags:
 
 # Document OCR using {model_name}
 
-This dataset contains OCR results from images in [{source_dataset}](https://huggingface.co/datasets/{source_dataset}) using dots.mocr, a 3B multilingual model with SOTA document parsing and SVG generation.
+This dataset contains OCR results from images in [{source_dataset}](https://huggingface.co/datasets/{source_dataset}) using DoTS.ocr-1.5, a 3B multilingual model with SOTA document parsing.
 
 ## Processing Details
 
@@ -233,14 +186,13 @@ This dataset contains OCR results from images in [{source_dataset}](https://hugg
 
 ## Model Information
 
-dots.mocr is a 3B multilingual document parsing model that excels at:
+DoTS.ocr-1.5 is a 3B multilingual document parsing model that excels at:
 - 100+ Languages — Multilingual document support
 - Table extraction — Structured data recognition
 - Formulas — Mathematical notation preservation
 - Layout-aware — Reading order and structure preservation
 - Web screen parsing — Webpage layout analysis
 - Scene text spotting — Text detection in natural scenes
-- SVG code generation — Charts, UI layouts, scientific figures to SVG
 
 ## Dataset Structure
 
@@ -270,10 +222,10 @@ for info in inference_info:
 
 ## Reproduction
 
-This dataset was generated using the [uv-scripts/ocr](https://huggingface.co/datasets/uv-scripts/ocr) dots.mocr script:
+This dataset was generated using the [uv-scripts/ocr](https://huggingface.co/datasets/uv-scripts/ocr) DoTS OCR 1.5 script:
 
 ```bash
-uv run https://huggingface.co/datasets/uv-scripts/ocr/raw/main/dots-mocr.py \\
+uv run https://huggingface.co/datasets/uv-scripts/ocr/raw/main/dots-ocr-1.5.py \\
     {source_dataset} \\
     <output-dataset> \\
     --image-column {image_column} \\
@@ -293,7 +245,7 @@ def main(
     output_dataset: str,
     image_column: str = "image",
     batch_size: int = 16,
-    model: str = "rednote-hilab/dots.mocr",
+    model: str = "rednote-hilab/dots.ocr-1.5",
     max_model_len: int = 24000,
     max_tokens: int = 24000,
     gpu_memory_utilization: float = 0.9,
@@ -306,14 +258,13 @@ def main(
     prompt_mode: str = "ocr",
     custom_prompt: str = None,
     output_column: str = "markdown",
-    overwrite: bool = False,
     config: str = None,
     create_pr: bool = False,
     temperature: float = 0.1,
     top_p: float = 0.9,
     verbose: bool = False,
 ):
-    """Process images from HF dataset through dots.mocr model."""
+    """Process images from HF dataset through DoTS.ocr-1.5 model."""
 
     # Check CUDA availability first
     check_cuda_availability()
@@ -344,9 +295,6 @@ def main(
             f"Column '{image_column}' not found. Available: {dataset.column_names}"
         )
 
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
     # Shuffle if requested
     if shuffle:
         logger.info(f"Shuffling dataset with seed {seed}")
@@ -367,12 +315,6 @@ def main(
         gpu_memory_utilization=gpu_memory_utilization,
     )
 
-    # SVG mode uses higher temperature/top_p per official recommendation
-    if prompt_mode == "svg" and temperature == 0.1 and top_p == 0.9:
-        logger.info("SVG mode: using recommended temperature=0.9, top_p=1.0")
-        temperature = 0.9
-        top_p = 1.0
-
     sampling_params = SamplingParams(
         temperature=temperature,
         top_p=top_p,
@@ -388,35 +330,17 @@ def main(
     for batch_indices in tqdm(
         partition_all(batch_size, range(len(dataset))),
         total=(len(dataset) + batch_size - 1) // batch_size,
-        desc="dots.mocr processing",
+        desc="DoTS.ocr-1.5 processing",
     ):
         batch_indices = list(batch_indices)
-
-        # Fetch images first, with per-batch fallback for unreadable files.
-        # One corrupt image used to take down the entire run via the list
-        # comprehension; now we mark the whole batch as skipped and continue.
-        try:
-            batch_images = [dataset[i][image_column] for i in batch_indices]
-        except (UnidentifiedImageError, OSError) as e:
-            logger.warning(
-                f"Skipping batch of {len(batch_indices)} — unreadable image "
-                f"in batch: {type(e).__name__}: {e}"
-            )
-            all_outputs.extend(
-                ["[OCR SKIPPED — UNREADABLE IMAGE]"] * len(batch_indices)
-            )
-            continue
+        batch_images = [dataset[i][image_column] for i in batch_indices]
 
         try:
             # Create messages for batch
             batch_messages = [make_ocr_message(img, prompt) for img in batch_images]
 
-            # Process with vLLM (dots.mocr needs "string" content format)
-            outputs = llm.chat(
-                batch_messages,
-                sampling_params,
-                chat_template_content_format="string",
-            )
+            # Process with vLLM
+            outputs = llm.chat(batch_messages, sampling_params)
 
             # Extract outputs
             for output in outputs:
@@ -439,7 +363,7 @@ def main(
     # Handle inference_info tracking (for multi-model comparisons)
     inference_entry = {
         "model_id": model,
-        "model_name": "dots.mocr",
+        "model_name": "DoTS.ocr-1.5",
         "column_name": output_column,
         "timestamp": datetime.now().isoformat(),
         "prompt_mode": prompt_mode if not custom_prompt else "custom",
@@ -520,7 +444,7 @@ def main(
     card = DatasetCard(card_content)
     card.push_to_hub(output_dataset, token=HF_TOKEN)
 
-    logger.info("dots.mocr processing complete!")
+    logger.info("DoTS.ocr-1.5 processing complete!")
     logger.info(
         f"Dataset available at: https://huggingface.co/datasets/{output_dataset}"
     )
@@ -542,83 +466,77 @@ if __name__ == "__main__":
     # Show example usage if no arguments
     if len(sys.argv) == 1:
         print("=" * 80)
-        print("dots.mocr Document Processing")
+        print("DoTS.ocr-1.5 Document Processing")
         print("=" * 80)
-        print("\n3B multilingual OCR model with SVG generation")
+        print("\n3B multilingual OCR model supporting 100+ languages")
         print("\nFeatures:")
         print("- Multilingual support (100+ languages)")
         print("- Fast processing with vLLM")
         print("- Table extraction and formatting")
         print("- Formula recognition")
         print("- Layout-aware text extraction")
-        print("- Web screen parsing")
-        print("- Scene text spotting")
-        print("- SVG code generation (charts, UI, figures)")
+        print("- Web screen parsing (NEW in v1.5)")
+        print("- Scene text spotting (NEW in v1.5)")
         print("\nPrompt modes:")
-        print("  ocr            - Text extraction (default)")
-        print("  layout-all     - Layout + bboxes + text (JSON)")
-        print("  layout-only    - Layout + bboxes only (JSON)")
-        print("  web-parsing    - Webpage layout analysis (JSON)")
+        print("  ocr           - Text extraction (default)")
+        print("  layout-all    - Layout + bboxes + text (JSON)")
+        print("  layout-only   - Layout + bboxes only (JSON)")
+        print("  web-parsing   - Webpage layout analysis (JSON)")
         print("  scene-spotting - Scene text detection")
-        print("  grounding-ocr  - Text from bounding box region")
-        print("  svg            - SVG code generation")
-        print("  general        - Free-form (use with --custom-prompt)")
+        print("  grounding-ocr - Text from bounding box region")
+        print("  general       - Free-form (use with --custom-prompt)")
         print("\nExample usage:")
         print("\n1. Basic OCR:")
-        print("   uv run dots-mocr.py input-dataset output-dataset")
-        print("\n2. SVG generation:")
-        print(
-            "   uv run dots-mocr.py charts svg-output --prompt-mode svg --model rednote-hilab/dots.mocr-svg"
-        )
-        print("\n3. Web screen parsing:")
-        print("   uv run dots-mocr.py screenshots parsed --prompt-mode web-parsing")
+        print("   uv run dots-ocr-1.5.py input-dataset output-dataset")
+        print("\n2. Web screen parsing:")
+        print("   uv run dots-ocr-1.5.py screenshots parsed --prompt-mode web-parsing")
+        print("\n3. Scene text spotting:")
+        print("   uv run dots-ocr-1.5.py photos detected --prompt-mode scene-spotting")
         print("\n4. Layout analysis with structure:")
-        print("   uv run dots-mocr.py papers analyzed --prompt-mode layout-all")
+        print("   uv run dots-ocr-1.5.py papers analyzed --prompt-mode layout-all")
         print("\n5. Running on HF Jobs:")
         print("   hf jobs uv run --flavor l4x1 \\")
         print("     -s HF_TOKEN \\")
         print(
-            "     https://huggingface.co/datasets/uv-scripts/ocr/raw/main/dots-mocr.py \\"
+            "     https://huggingface.co/datasets/uv-scripts/ocr/raw/main/dots-ocr-1.5.py \\"
         )
         print("       input-dataset output-dataset")
         print("\n" + "=" * 80)
-        print("\nFor full help, run: uv run dots-mocr.py --help")
+        print("\nFor full help, run: uv run dots-ocr-1.5.py --help")
         sys.exit(0)
 
     parser = argparse.ArgumentParser(
-        description="Document OCR using dots.mocr (3B multilingual model with SVG generation)",
+        description="Document OCR using DoTS.ocr-1.5 (3B multilingual model)",
         formatter_class=argparse.RawDescriptionHelpFormatter,
         epilog="""
-Prompt Modes (official dots.mocr prompts):
+Prompt Modes (official DoTS.ocr-1.5 prompts):
   ocr            - Simple text extraction (default)
   layout-all     - Layout analysis with bboxes, categories, and text (JSON output)
   layout-only    - Layout detection with bboxes and categories only (JSON output)
-  web-parsing    - Webpage layout analysis (JSON output)
-  scene-spotting - Scene text detection and recognition
-  grounding-ocr  - Extract text from bounding box region
-  svg            - SVG code generation (auto-injects image dimensions into viewBox)
-  general        - Free-form QA (use with --custom-prompt)
+  web-parsing    - Webpage layout analysis (JSON output) [NEW in v1.5]
+  scene-spotting - Scene text detection and recognition [NEW in v1.5]
+  grounding-ocr  - Extract text from bounding box region [NEW in v1.5]
+  general        - Free-form QA (use with --custom-prompt) [NEW in v1.5]
 
 SVG Code Generation:
-  Use --prompt-mode svg for SVG output. For best results, combine with
-  --model rednote-hilab/dots.mocr-svg (the SVG-optimized variant).
-  SVG mode automatically uses temperature=0.9, top_p=1.0 unless overridden.
+  For SVG output, use --model rednote-hilab/dots.ocr-1.5-svg with:
+  --custom-prompt 'Please generate the SVG code based on the image.'
 
 Examples:
   # Basic text OCR (default)
-  uv run dots-mocr.py my-docs analyzed-docs
-
-  # SVG generation with optimized variant
-  uv run dots-mocr.py charts svg-out --prompt-mode svg --model rednote-hilab/dots.mocr-svg
+  uv run dots-ocr-1.5.py my-docs analyzed-docs
 
   # Web screen parsing
-  uv run dots-mocr.py screenshots parsed --prompt-mode web-parsing
+  uv run dots-ocr-1.5.py screenshots parsed --prompt-mode web-parsing
+
+  # Scene text spotting
+  uv run dots-ocr-1.5.py photos spotted --prompt-mode scene-spotting
 
   # Full layout analysis with structure
-  uv run dots-mocr.py papers structured --prompt-mode layout-all
+  uv run dots-ocr-1.5.py papers structured --prompt-mode layout-all
 
   # Random sampling for testing
-  uv run dots-mocr.py large-dataset test --max-samples 50 --shuffle
+  uv run dots-ocr-1.5.py large-dataset test --max-samples 50 --shuffle
         """,
     )
 
@@ -637,8 +555,8 @@ Examples:
     )
     parser.add_argument(
         "--model",
-        default="rednote-hilab/dots.mocr",
-        help="Model to use (default: rednote-hilab/dots.mocr, or rednote-hilab/dots.mocr-svg for SVG)",
+        default="rednote-hilab/dots.ocr-1.5",
+        help="Model to use (default: rednote-hilab/dots.ocr-1.5)",
     )
     parser.add_argument(
         "--max-model-len",
@@ -694,12 +612,6 @@ Examples:
         default="markdown",
         help="Column name for output text (default: markdown)",
     )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
     parser.add_argument(
         "--config",
         help="Config/subset name when pushing to Hub (for benchmarking multiple models in one repo)",
@@ -747,7 +659,6 @@ Examples:
         prompt_mode=args.prompt_mode,
         custom_prompt=args.custom_prompt,
         output_column=args.output_column,
-        overwrite=args.overwrite,
         config=args.config,
         create_pr=args.create_pr,
         temperature=args.temperature,

dots-ocr.py

@@ -1,17 +1,13 @@
 # /// script
 # requires-python = ">=3.11"
 # dependencies = [
-#     "datasets>=4.0.0",
+#     "datasets",
 #     "huggingface-hub",
 #     "pillow",
 #     "vllm>=0.15.1",
 #     "tqdm",
 #     "toolz",
 #     "torch",
-#     # dots.ocr's remote code (AutoProcessor.register with a string key)
-#     # breaks on transformers v5; vllm pulls transformers unpinned.
-#     # Same fix as dots-mocr (#76).
-#     "transformers<5",
 # ]
 #
 # ///
@@ -49,10 +45,6 @@ from huggingface_hub import DatasetCard, login
 from PIL import Image
 from toolz import partition_all
 from tqdm.auto import tqdm
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
 from vllm import LLM, SamplingParams
 
 logging.basicConfig(level=logging.INFO)
@@ -97,28 +89,6 @@ def check_cuda_availability():
         logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
 
 
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
 def make_ocr_message(
     image: Union[Image.Image, Dict[str, Any], str],
     prompt: str = PROMPT_TEMPLATES["ocr"],
@@ -265,8 +235,7 @@ def main(
     image_column: str = "image",
     batch_size: int = 16,
     model: str = "rednote-hilab/dots.ocr",
-    max_model_len: int = 32768,
-    max_pixels: int = None,
+    max_model_len: int = 8192,
     max_tokens: int = 8192,
     gpu_memory_utilization: float = 0.8,
     hf_token: str = None,
@@ -278,7 +247,6 @@ def main(
     prompt_mode: str = "ocr",
     custom_prompt: str = None,
     output_column: str = "markdown",
-    overwrite: bool = False,
     config: str = None,
     create_pr: bool = False,
 ):
@@ -313,9 +281,6 @@ def main(
             f"Column '{image_column}' not found. Available: {dataset.column_names}"
         )
 
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
     # Shuffle if requested
     if shuffle:
         logger.info(f"Shuffling dataset with seed {seed}")
@@ -329,17 +294,13 @@ def main(
     # Initialize vLLM model
     logger.info(f"Initializing vLLM with model: {model}")
     logger.info("This may take a few minutes on first run...")
-    llm_kwargs = dict(
+    llm = LLM(
         model=model,
         trust_remote_code=True,
         max_model_len=max_model_len,
         gpu_memory_utilization=gpu_memory_utilization,
         limit_mm_per_prompt={"image": 1},
     )
-    if max_pixels is not None:
-        logger.info(f"Capping input images to max_pixels={max_pixels}")
-        llm_kwargs["mm_processor_kwargs"] = {"max_pixels": max_pixels}
-    llm = LLM(**llm_kwargs)
 
     sampling_params = SamplingParams(
         temperature=0.0,  # Deterministic for OCR
@@ -544,23 +505,8 @@ Examples:
     parser.add_argument(
         "--max-model-len",
         type=int,
-        default=32768,
-        help=(
-            "Maximum model context length (default: 32768). dots.ocr does NOT resize "
-            "input images, so a full page can reach ~14k image tokens (the model's "
-            "11.29M-px processor cap); the old 8192 default rejected such requests and "
-            "the row was written as '[OCR ERROR]'. Pair with --max-pixels to cap memory."
-        ),
-    )
-    parser.add_argument(
-        "--max-pixels",
-        type=int,
-        default=None,
-        help=(
-            "Optional cap on input image pixels (width*height) passed to vLLM's "
-            "mm_processor. Lower this (e.g. 4000000) to bound image tokens and GPU "
-            "memory on very large scans. Default: model's own cap (~11.29M px)."
-        ),
+        default=8192,
+        help="Maximum model context length (default: 8192)",
     )
     parser.add_argument(
         "--max-tokens",
@@ -610,12 +556,6 @@ Examples:
         default="markdown",
         help="Column name for output text (default: markdown)",
     )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
     parser.add_argument(
         "--config",
         help="Config/subset name when pushing to Hub (for benchmarking multiple models in one repo)",
@@ -635,7 +575,6 @@ Examples:
         batch_size=args.batch_size,
         model=args.model,
         max_model_len=args.max_model_len,
-        max_pixels=args.max_pixels,
         max_tokens=args.max_tokens,
         gpu_memory_utilization=args.gpu_memory_utilization,
         hf_token=args.hf_token,
@@ -647,7 +586,6 @@ Examples:
         prompt_mode=args.prompt_mode,
         custom_prompt=args.custom_prompt,
         output_column=args.output_column,
-        overwrite=args.overwrite,
         config=args.config,
         create_pr=args.create_pr,
     )

examples/nls-index-card-v2.json

@@ -1,13 +0,0 @@
-{
-  "image_type": ["index_card", "verso", "cover", "blank", "other"],
-  "heading": "verbatim-string",
-  "heading_type": ["person", "family", "corporate", "geographic", "subject"],
-  "epithet": "string",
-  "entries": [
-    {
-      "ms_no": "verbatim-string",
-      "folios": ["verbatim-string"],
-      "description": "string"
-    }
-  ]
-}

examples/nls-index-card-verbose.json

@@ -1,13 +0,0 @@
-{
-  "image_type": ["index_card", "verso", "cover", "blank", "other"],
-  "main_heading_name": "verbatim-string",
-  "heading_category": ["person", "family", "corporate", "geographic", "subject"],
-  "epithet_title_or_occupation": "string",
-  "manuscript_references": [
-    {
-      "manuscript_number": "verbatim-string",
-      "folio_references": ["verbatim-string"],
-      "entry_description_with_date": "string"
-    }
-  ]
-}

falcon-ocr-bucket.py

@@ -1,303 +0,0 @@
-# /// script
-# requires-python = ">=3.11"
-# dependencies = [
-#     "pillow",
-#     "pymupdf",
-#     "torch>=2.5",
-#     "torchvision",
-#     "falcon-perception[ocr]",
-# ]
-# ///
-
-"""
-OCR images and PDFs from a directory using Falcon OCR, writing markdown files.
-
-Designed to work with HF Buckets mounted as volumes via `hf jobs uv run -v ...`.
-Reads images/PDFs from INPUT_DIR, runs Falcon OCR via the optimized falcon-perception
-engine (CUDA graphs + paged inference), and writes one .md file per image (or per
-PDF page) to OUTPUT_DIR, preserving directory structure.
-
-Input:                          Output:
-  /input/page1.png        ->      /output/page1.md
-  /input/report.pdf       ->      /output/report/page_001.md
-  (3 pages)                       /output/report/page_002.md
-                                  /output/report/page_003.md
-  /input/sub/photo.jpg    ->      /output/sub/photo.md
-
-Examples:
-
-  # Local test
-  uv run falcon-ocr-bucket.py ./test-images ./test-output
-
-  # HF Jobs with bucket volumes
-  hf jobs uv run --flavor l4x1 \\
-      -s HF_TOKEN \\
-      -v hf://buckets/user/ocr-input:/input:ro \\
-      -v hf://buckets/user/ocr-output:/output \\
-      https://huggingface.co/datasets/uv-scripts/ocr/raw/main/falcon-ocr-bucket.py \\
-      /input /output
-
-Model: tiiuae/Falcon-OCR (0.3B, 80.3% olmOCR, Apache 2.0)
-Backend: falcon-perception (OCRInferenceEngine with CUDA graphs)
-"""
-
-import argparse
-import logging
-import sys
-import time
-from pathlib import Path
-
-import torch
-from PIL import Image
-
-logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
-logger = logging.getLogger(__name__)
-
-MODEL_ID = "tiiuae/Falcon-OCR"
-IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".tiff", ".tif", ".bmp", ".webp"}
-
-
-def check_cuda_availability():
-    if not torch.cuda.is_available():
-        logger.error("CUDA is not available. This script requires a GPU.")
-        sys.exit(1)
-    logger.info(f"CUDA available. GPU: {torch.cuda.get_device_name(0)}")
-
-
-def discover_files(input_dir: Path, limit: int | None = None) -> list[Path]:
-    """Discover image and PDF files under input_dir.
-
-    Without `limit`, returns the full sorted list (deterministic order).
-    With `limit`, stops scanning once `limit` matching files are found
-    and returns them in filesystem order (much faster on huge mounted
-    buckets, but ordering is not deterministic).
-    """
-    files = []
-    iterator = (
-        input_dir.rglob("*") if limit is not None else sorted(input_dir.rglob("*"))
-    )
-    for path in iterator:
-        if not path.is_file():
-            continue
-        ext = path.suffix.lower()
-        if ext in IMAGE_EXTENSIONS or ext == ".pdf":
-            files.append(path)
-            if limit is not None and len(files) >= limit:
-                break
-    return files
-
-
-def prepare_images(
-    files: list[Path], input_dir: Path, output_dir: Path, pdf_dpi: int
-) -> list[tuple[Image.Image, Path]]:
-    import fitz  # pymupdf
-
-    items: list[tuple[Image.Image, Path]] = []
-
-    for file_path in files:
-        rel = file_path.relative_to(input_dir)
-        ext = file_path.suffix.lower()
-
-        if ext == ".pdf":
-            pdf_output_dir = output_dir / rel.with_suffix("")
-            try:
-                doc = fitz.open(file_path)
-                num_pages = len(doc)
-                logger.info(f"PDF: {rel} ({num_pages} pages)")
-                for page_num in range(num_pages):
-                    page = doc[page_num]
-                    zoom = pdf_dpi / 72.0
-                    mat = fitz.Matrix(zoom, zoom)
-                    pix = page.get_pixmap(matrix=mat)
-                    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-                    md_path = pdf_output_dir / f"page_{page_num + 1:03d}.md"
-                    items.append((img, md_path))
-                doc.close()
-            except Exception as e:
-                logger.error(f"Failed to open PDF {rel}: {e}")
-        else:
-            try:
-                img = Image.open(file_path).convert("RGB")
-                md_path = output_dir / rel.with_suffix(".md")
-                items.append((img, md_path))
-            except Exception as e:
-                logger.error(f"Failed to open image {rel}: {e}")
-
-    return items
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="OCR images/PDFs from a directory using Falcon OCR, output markdown files.",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog=__doc__,
-    )
-    parser.add_argument("input_dir", help="Directory containing images and/or PDFs")
-    parser.add_argument("output_dir", help="Directory to write markdown output files")
-    parser.add_argument(
-        "--batch-size", type=int, default=8, help="Images per batch (default: 8)",
-    )
-    parser.add_argument(
-        "--pdf-dpi", type=int, default=300,
-        help="DPI for PDF page rendering (default: 300)",
-    )
-    parser.add_argument(
-        "--no-compile", action="store_true", help="Disable torch.compile",
-    )
-    parser.add_argument(
-        "--no-cudagraph", action="store_true", help="Disable CUDA graph capture",
-    )
-    parser.add_argument(
-        "--max-samples", type=int, default=None,
-        help="Limit number of input files to discover. Stops scanning early "
-             "once the limit is reached (much faster on large mounted buckets). "
-             "Applied before PDF page expansion. With --max-samples set, file "
-             "ordering is filesystem-dependent rather than sorted.",
-    )
-    parser.add_argument(
-        "--verbose", action="store_true", help="Print resolved package versions",
-    )
-
-    args = parser.parse_args()
-
-    check_cuda_availability()
-
-    input_dir = Path(args.input_dir)
-    output_dir = Path(args.output_dir)
-
-    if not input_dir.is_dir():
-        logger.error(f"Input directory does not exist: {input_dir}")
-        sys.exit(1)
-
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    start_time = time.time()
-
-    # Discover files
-    if args.max_samples is not None:
-        logger.info(
-            f"Scanning {input_dir} for up to {args.max_samples} images/PDFs "
-            f"(early termination, --max-samples)..."
-        )
-    else:
-        logger.info(f"Scanning {input_dir} for images and PDFs...")
-    files = discover_files(input_dir, limit=args.max_samples)
-    if not files:
-        logger.error(f"No image or PDF files found in {input_dir}")
-        sys.exit(1)
-
-    pdf_count = sum(1 for f in files if f.suffix.lower() == ".pdf")
-    img_count = len(files) - pdf_count
-    logger.info(f"Found {img_count} image(s) and {pdf_count} PDF(s)")
-
-    # Prepare images
-    logger.info("Preparing images (rendering PDFs)...")
-    items = prepare_images(files, input_dir, output_dir, args.pdf_dpi)
-    if not items:
-        logger.error("No processable images after preparation")
-        sys.exit(1)
-
-    logger.info(f"Total images to OCR: {len(items)}")
-
-    # Load model
-    logger.info(f"Loading {MODEL_ID} via falcon-perception engine...")
-    from falcon_perception import load_and_prepare_model
-    from falcon_perception.data import ImageProcessor
-    from falcon_perception.paged_ocr_inference import OCRInferenceEngine
-
-    do_compile = not args.no_compile
-    do_cudagraph = not args.no_cudagraph
-
-    model, tokenizer, model_args = load_and_prepare_model(
-        hf_model_id=MODEL_ID,
-        device="cuda",
-        dtype="bfloat16",
-        compile=do_compile,
-    )
-
-    image_processor = ImageProcessor(patch_size=16, merge_size=1)
-    engine = OCRInferenceEngine(
-        model, tokenizer, image_processor, capture_cudagraph=do_cudagraph
-    )
-    logger.info(f"Engine loaded. compile={do_compile}, cudagraph={do_cudagraph}")
-
-    # Process in batches
-    errors = 0
-    processed = 0
-    total = len(items)
-    batch_size = args.batch_size
-
-    for batch_start in range(0, total, batch_size):
-        batch_end = min(batch_start + batch_size, total)
-        batch = items[batch_start:batch_end]
-        batch_num = batch_start // batch_size + 1
-        total_batches = (total + batch_size - 1) // batch_size
-
-        logger.info(f"Batch {batch_num}/{total_batches} ({processed}/{total} done)")
-
-        try:
-            batch_images = [img for img, _ in batch]
-            texts = engine.generate_plain(images=batch_images, use_tqdm=False)
-
-            for (_, md_path), text in zip(batch, texts):
-                md_path.parent.mkdir(parents=True, exist_ok=True)
-                md_path.write_text(text.strip(), encoding="utf-8")
-                processed += 1
-
-        except Exception as e:
-            logger.error(f"Batch {batch_num} failed: {e}")
-            for _, md_path in batch:
-                md_path.parent.mkdir(parents=True, exist_ok=True)
-                md_path.write_text(f"[OCR ERROR: {e}]", encoding="utf-8")
-            errors += len(batch)
-            processed += len(batch)
-
-    elapsed = time.time() - start_time
-    elapsed_str = f"{elapsed / 60:.1f} min" if elapsed > 60 else f"{elapsed:.1f}s"
-
-    logger.info("=" * 50)
-    logger.info(f"Done! Processed {total} images in {elapsed_str}")
-    logger.info(f"  Output: {output_dir}")
-    logger.info(f"  Errors: {errors}")
-    if total > 0:
-        logger.info(f"  Speed: {total / elapsed:.2f} images/sec")
-
-    if args.verbose:
-        import importlib.metadata
-
-        logger.info("--- Package versions ---")
-        for pkg in ["falcon-perception", "torch", "pillow", "pymupdf"]:
-            try:
-                logger.info(f"  {pkg}=={importlib.metadata.version(pkg)}")
-            except importlib.metadata.PackageNotFoundError:
-                logger.info(f"  {pkg}: not installed")
-
-
-if __name__ == "__main__":
-    if len(sys.argv) == 1:
-        print("=" * 60)
-        print("Falcon OCR Bucket Script")
-        print("=" * 60)
-        print(f"\nModel: {MODEL_ID} (0.3B, Apache 2.0)")
-        print("OCR images/PDFs from a directory -> markdown files.")
-        print("Designed for HF Buckets mounted as volumes.")
-        print()
-        print("Usage:")
-        print("  uv run falcon-ocr-bucket.py INPUT_DIR OUTPUT_DIR")
-        print()
-        print("Examples:")
-        print("  uv run falcon-ocr-bucket.py ./images ./output")
-        print()
-        print("HF Jobs with bucket volumes:")
-        print("  hf jobs uv run --flavor l4x1 -s HF_TOKEN \\")
-        print("      -v hf://buckets/user/ocr-input:/input:ro \\")
-        print("      -v hf://buckets/user/ocr-output:/output \\")
-        print(
-            "      https://huggingface.co/datasets/uv-scripts/ocr/raw/main/falcon-ocr-bucket.py \\"
-        )
-        print("      /input /output")
-        print()
-        print("For full help: uv run falcon-ocr-bucket.py --help")
-        sys.exit(0)
-
-    main()

falcon-ocr.py

@@ -1,466 +0,0 @@
-# /// script
-# requires-python = ">=3.11"
-# dependencies = [
-#     "datasets>=4.0.0",
-#     "huggingface-hub",
-#     "pillow",
-#     "torch>=2.5",
-#     "torchvision",
-#     "falcon-perception",
-# ]
-# ///
-
-"""
-Convert document images to text using Falcon OCR with the falcon-perception engine.
-
-Uses the optimized OCRInferenceEngine with CUDA graphs and paged inference
-for much faster throughput than the raw transformers API.
-
-Features:
-- Compact: Only 0.3B parameters
-- Fast: Optimized inference with CUDA graphs
-- Multi-format: Plain text, LaTeX formulas, HTML tables
-- Layout-aware: Optional 2-stage pipeline (layout detection + per-region OCR)
-
-Model: tiiuae/Falcon-OCR
-Backend: falcon-perception (OCRInferenceEngine)
-License: Apache 2.0
-
-Examples:
-    # Basic text OCR
-    uv run falcon-ocr.py input-dataset output-dataset
-
-    # Test with small sample
-    uv run falcon-ocr.py dataset test --max-samples 5 --shuffle
-
-    # Run on HF Jobs with GPU
-    hf jobs uv run --flavor l4x1 \\
-        -s HF_TOKEN \\
-        falcon-ocr.py \\
-        input-dataset output-dataset --max-samples 10
-"""
-
-import argparse
-import io
-import json
-import logging
-import os
-import sys
-import time
-from datetime import datetime
-from typing import Any, Dict, Union
-
-import torch
-from datasets import load_dataset
-from huggingface_hub import DatasetCard, login
-from PIL import Image
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-MODEL_ID = "tiiuae/Falcon-OCR"
-
-TASK_MODES = {
-    "plain": "Full-page text extraction",
-}
-
-
-def check_cuda_availability():
-    if not torch.cuda.is_available():
-        logger.error("CUDA is not available. This script requires a GPU.")
-        logger.error("For cloud execution, use HF Jobs with --flavor l4x1 or similar.")
-        sys.exit(1)
-    else:
-        logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
-
-
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
-def prepare_image(image: Union[Image.Image, Dict[str, Any], str]) -> Image.Image:
-    if isinstance(image, Image.Image):
-        pil_img = image
-    elif isinstance(image, dict) and "bytes" in image:
-        pil_img = Image.open(io.BytesIO(image["bytes"]))
-    elif isinstance(image, str):
-        pil_img = Image.open(image)
-    else:
-        raise ValueError(f"Unsupported image type: {type(image)}")
-    return pil_img.convert("RGB")
-
-
-def create_dataset_card(
-    source_dataset: str,
-    task_mode: str,
-    num_samples: int,
-    processing_time: str,
-    image_column: str = "image",
-    split: str = "train",
-) -> str:
-    task_description = TASK_MODES[task_mode]
-    return f"""---
-tags:
-- ocr
-- document-processing
-- falcon-ocr
-- {task_mode}
-- uv-script
-- generated
----
-
-# Document Processing using Falcon OCR ({task_mode} mode)
-
-This dataset contains OCR results from images in [{source_dataset}](https://huggingface.co/datasets/{source_dataset}) using [Falcon OCR](https://huggingface.co/tiiuae/Falcon-OCR), a 0.3B early-fusion vision-language model.
-
-## Processing Details
-
-- **Source Dataset**: [{source_dataset}](https://huggingface.co/datasets/{source_dataset})
-- **Model**: [{MODEL_ID}](https://huggingface.co/{MODEL_ID})
-- **Task Mode**: `{task_mode}` - {task_description}
-- **Number of Samples**: {num_samples:,}
-- **Processing Time**: {processing_time}
-- **Processing Date**: {datetime.now().strftime("%Y-%m-%d %H:%M UTC")}
-- **Backend**: falcon-perception (OCRInferenceEngine)
-
-## Reproduction
-
-```bash
-uv run https://huggingface.co/datasets/uv-scripts/ocr/raw/main/falcon-ocr.py \\
-    {source_dataset} \\
-    <output-dataset> \\
-    --task-mode {task_mode} \\
-    --image-column {image_column}
-```
-
-Generated with [UV Scripts](https://huggingface.co/uv-scripts)
-"""
-
-
-def main(
-    input_dataset: str,
-    output_dataset: str,
-    image_column: str = "image",
-    task_mode: str = "plain",
-    hf_token: str = None,
-    split: str = "train",
-    max_samples: int = None,
-    private: bool = False,
-    shuffle: bool = False,
-    seed: int = 42,
-    output_column: str = "markdown",
-    overwrite: bool = False,
-    config: str = None,
-    create_pr: bool = False,
-    compile: bool = True,
-    cudagraph: bool = True,
-    progress: bool = False,
-    verbose: bool = False,
-):
-    check_cuda_availability()
-    start_time = datetime.now()
-
-    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")
-    if HF_TOKEN:
-        login(token=HF_TOKEN)
-
-    if task_mode not in TASK_MODES:
-        raise ValueError(
-            f"Invalid task_mode '{task_mode}'. Choose from: {list(TASK_MODES.keys())}"
-        )
-
-    logger.info(f"Task mode: {task_mode} - {TASK_MODES[task_mode]}")
-    logger.info(f"Output column: {output_column}")
-
-    # Load dataset
-    logger.info(f"Loading dataset: {input_dataset}")
-    dataset = load_dataset(input_dataset, split=split)
-
-    if image_column not in dataset.column_names:
-        raise ValueError(
-            f"Column '{image_column}' not found. Available: {dataset.column_names}"
-        )
-
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
-    if shuffle:
-        logger.info(f"Shuffling dataset with seed {seed}")
-        dataset = dataset.shuffle(seed=seed)
-
-    if max_samples:
-        dataset = dataset.select(range(min(max_samples, len(dataset))))
-        logger.info(f"Limited to {len(dataset)} samples")
-
-    # Load model using falcon-perception
-    logger.info(f"Loading model: {MODEL_ID} via falcon-perception engine")
-    from falcon_perception import load_and_prepare_model
-    from falcon_perception.data import ImageProcessor
-    from falcon_perception.paged_ocr_inference import OCRInferenceEngine
-
-    model, tokenizer, model_args = load_and_prepare_model(
-        hf_model_id=MODEL_ID,
-        device="cuda",
-        dtype="bfloat16",
-        compile=compile,
-    )
-
-    image_processor = ImageProcessor(patch_size=16, merge_size=1)
-    engine = OCRInferenceEngine(
-        model, tokenizer, image_processor, capture_cudagraph=cudagraph
-    )
-    logger.info(f"Engine loaded. compile={compile}, cudagraph={cudagraph}")
-
-    # Prepare all images
-    logger.info(f"Processing {len(dataset)} images...")
-    all_outputs = []
-
-    # Batch plain OCR for better throughput
-    batch_size = 8
-    total_batches = (len(dataset) + batch_size - 1) // batch_size
-    for batch_idx, batch_start in enumerate(range(0, len(dataset), batch_size), 1):
-        batch_end = min(batch_start + batch_size, len(dataset))
-        logger.info(f"Batch {batch_idx}/{total_batches} ({batch_start}/{len(dataset)} done)")
-        batch_images = []
-        for i in range(batch_start, batch_end):
-            try:
-                batch_images.append(prepare_image(dataset[i][image_column]))
-            except Exception as e:
-                logger.error(f"Error preparing image {i}: {e}")
-                batch_images.append(Image.new("RGB", (100, 100)))
-
-        try:
-            texts = engine.generate_plain(
-                images=batch_images, use_tqdm=progress
-            )
-            all_outputs.extend(texts)
-        except Exception as e:
-            logger.error(f"Error processing batch {batch_start}-{batch_end}: {e}")
-            all_outputs.extend(
-                [f"[OCR ERROR: {str(e)[:200]}]"] * len(batch_images)
-            )
-
-    # Calculate processing time
-    processing_duration = datetime.now() - start_time
-    processing_time_str = f"{processing_duration.total_seconds() / 60:.1f} min"
-
-    # Add output column
-    logger.info(f"Adding '{output_column}' column to dataset")
-    dataset = dataset.add_column(output_column, all_outputs)
-
-    # Track inference info
-    inference_entry = {
-        "model_id": MODEL_ID,
-        "model_name": "Falcon-OCR",
-        "model_size": "0.3B",
-        "task_mode": task_mode,
-        "column_name": output_column,
-        "timestamp": datetime.now().isoformat(),
-        "backend": "falcon-perception",
-    }
-
-    if "inference_info" in dataset.column_names:
-        def update_inference_info(example):
-            try:
-                existing_info = (
-                    json.loads(example["inference_info"])
-                    if example["inference_info"]
-                    else []
-                )
-            except (json.JSONDecodeError, TypeError):
-                existing_info = []
-            existing_info.append(inference_entry)
-            return {"inference_info": json.dumps(existing_info)}
-
-        dataset = dataset.map(update_inference_info)
-    else:
-        inference_list = [json.dumps([inference_entry])] * len(dataset)
-        dataset = dataset.add_column("inference_info", inference_list)
-
-    # Push to hub
-    logger.info(f"Pushing to {output_dataset}")
-    max_retries = 3
-    for attempt in range(1, max_retries + 1):
-        try:
-            if attempt > 1:
-                logger.warning("Disabling XET (fallback to HTTP upload)")
-                os.environ["HF_HUB_DISABLE_XET"] = "1"
-            dataset.push_to_hub(
-                output_dataset,
-                private=private,
-                token=HF_TOKEN,
-                max_shard_size="500MB",
-                **({"config_name": config} if config else {}),
-                create_pr=create_pr,
-                commit_message=f"Add {MODEL_ID} OCR results ({len(dataset)} samples)"
-                + (f" [{config}]" if config else ""),
-            )
-            break
-        except Exception as e:
-            logger.error(f"Upload attempt {attempt}/{max_retries} failed: {e}")
-            if attempt < max_retries:
-                delay = 30 * (2 ** (attempt - 1))
-                logger.info(f"Retrying in {delay}s...")
-                time.sleep(delay)
-            else:
-                logger.error("All upload attempts failed. OCR results are lost.")
-                sys.exit(1)
-
-    # Create and push dataset card
-    logger.info("Creating dataset card")
-    card_content = create_dataset_card(
-        source_dataset=input_dataset,
-        task_mode=task_mode,
-        num_samples=len(dataset),
-        processing_time=processing_time_str,
-        image_column=image_column,
-        split=split,
-    )
-    card = DatasetCard(card_content)
-    card.push_to_hub(output_dataset, token=HF_TOKEN)
-
-    logger.info("Falcon OCR processing complete!")
-    logger.info(
-        f"Dataset available at: https://huggingface.co/datasets/{output_dataset}"
-    )
-    logger.info(f"Processing time: {processing_time_str}")
-    logger.info(
-        f"Speed: {len(dataset) / processing_duration.total_seconds():.2f} images/sec"
-    )
-
-    if verbose:
-        import importlib.metadata
-
-        logger.info("--- Resolved package versions ---")
-        for pkg in [
-            "falcon-perception", "transformers", "torch", "datasets", "pillow"
-        ]:
-            try:
-                logger.info(f"  {pkg}=={importlib.metadata.version(pkg)}")
-            except importlib.metadata.PackageNotFoundError:
-                logger.info(f"  {pkg}: not installed")
-
-
-if __name__ == "__main__":
-    if len(sys.argv) == 1:
-        print("=" * 70)
-        print("Falcon OCR - 0.3B Document OCR (falcon-perception engine)")
-        print("=" * 70)
-        print(f"\nModel: {MODEL_ID}")
-        print("License: Apache 2.0")
-        print("\nTask Modes:")
-        for mode, description in TASK_MODES.items():
-            print(f"  {mode:10} - {description}")
-        print("\nExamples:")
-        print("  uv run falcon-ocr.py input-dataset output-dataset")
-        print("  uv run falcon-ocr.py dense-docs output --task-mode layout")
-        print("\nFor full help: uv run falcon-ocr.py --help")
-        sys.exit(0)
-
-    parser = argparse.ArgumentParser(
-        description="Document OCR using Falcon OCR (0.3B, falcon-perception engine)",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog=__doc__,
-    )
-
-    parser.add_argument("input_dataset", help="Input dataset ID from Hugging Face Hub")
-    parser.add_argument("output_dataset", help="Output dataset ID for Hugging Face Hub")
-    parser.add_argument(
-        "--image-column", default="image",
-        help="Column containing images (default: image)",
-    )
-    parser.add_argument(
-        "--task-mode", choices=list(TASK_MODES.keys()), default="plain",
-        help="Task type: plain (default), layout",
-    )
-    parser.add_argument("--hf-token", help="Hugging Face API token")
-    parser.add_argument(
-        "--split", default="train", help="Dataset split (default: train)",
-    )
-    parser.add_argument(
-        "--max-samples", type=int,
-        help="Maximum number of samples to process (for testing)",
-    )
-    parser.add_argument(
-        "--private", action="store_true", help="Make output dataset private",
-    )
-    parser.add_argument(
-        "--shuffle", action="store_true", help="Shuffle dataset before processing",
-    )
-    parser.add_argument(
-        "--seed", type=int, default=42, help="Random seed for shuffling (default: 42)",
-    )
-    parser.add_argument(
-        "--output-column", default="markdown",
-        help="Column name for output text (default: markdown)",
-    )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
-    parser.add_argument(
-        "--config",
-        help="Config/subset name for Hub (for benchmarking multiple models)",
-    )
-    parser.add_argument(
-        "--create-pr", action="store_true",
-        help="Create a pull request instead of pushing directly",
-    )
-    parser.add_argument(
-        "--no-compile", action="store_true",
-        help="Disable torch.compile",
-    )
-    parser.add_argument(
-        "--no-cudagraph", action="store_true",
-        help="Disable CUDA graph capture",
-    )
-    parser.add_argument(
-        "--progress", action="store_true",
-        help="Show per-image progress bar from the inference engine",
-    )
-    parser.add_argument(
-        "--verbose", action="store_true", help="Log resolved package versions",
-    )
-
-    args = parser.parse_args()
-
-    main(
-        input_dataset=args.input_dataset,
-        output_dataset=args.output_dataset,
-        image_column=args.image_column,
-        task_mode=args.task_mode,
-        hf_token=args.hf_token,
-        split=args.split,
-        max_samples=args.max_samples,
-        private=args.private,
-        shuffle=args.shuffle,
-        seed=args.seed,
-        output_column=args.output_column,
-        overwrite=args.overwrite,
-        config=args.config,
-        create_pr=args.create_pr,
-        compile=not args.no_compile,
-        cudagraph=not args.no_cudagraph,
-        progress=args.progress,
-        verbose=args.verbose,
-    )

firered-ocr.py

@@ -1,7 +1,7 @@
 # /// script
 # requires-python = ">=3.11"
 # dependencies = [
-#     "datasets>=4.0.0",
+#     "datasets",
 #     "huggingface-hub",
 #     "pillow",
 #     "vllm>=0.15.1",
@@ -39,10 +39,6 @@ from huggingface_hub import DatasetCard, login
 from PIL import Image
 from toolz import partition_all
 from tqdm.auto import tqdm
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
 from vllm import LLM, SamplingParams
 
 logging.basicConfig(level=logging.INFO)
@@ -90,28 +86,6 @@ def check_cuda_availability():
         logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
 
 
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
 def make_ocr_message(
     image: Union[Image.Image, Dict[str, Any], str],
     prompt: str = FIRERED_OCR_PROMPT,
@@ -130,10 +104,10 @@ def make_ocr_message(
     # Convert to RGB
     pil_img = pil_img.convert("RGB")
 
-    # Convert to base64 data URI (JPEG is faster than PNG for encoding)
+    # Convert to base64 data URI
     buf = io.BytesIO()
-    pil_img.save(buf, format="JPEG", quality=95)
-    data_uri = f"data:image/jpeg;base64,{base64.b64encode(buf.getvalue()).decode()}"
+    pil_img.save(buf, format="PNG")
+    data_uri = f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}"
 
     # Return message in vLLM format
     return [
@@ -254,7 +228,7 @@ def main(
     image_column: str = "image",
     batch_size: int = 16,
     model: str = "FireRedTeam/FireRed-OCR",
-    max_model_len: int = 32768,
+    max_model_len: int = 8192,
     max_tokens: int = 8192,
     gpu_memory_utilization: float = 0.8,
     hf_token: str = None,
@@ -264,7 +238,6 @@ def main(
     shuffle: bool = False,
     seed: int = 42,
     output_column: str = "markdown",
-    overwrite: bool = False,
     config: str = None,
     create_pr: bool = False,
 ):
@@ -291,9 +264,6 @@ def main(
             f"Column '{image_column}' not found. Available: {dataset.column_names}"
         )
 
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
     # Shuffle if requested
     if shuffle:
         logger.info(f"Shuffling dataset with seed {seed}")
@@ -365,10 +335,7 @@ def main(
     processing_duration = datetime.now() - start_time
     processing_time_str = f"{processing_duration.total_seconds() / 60:.1f} min"
 
-    # Add output column to dataset (remove existing column if present)
-    if output_column in dataset.column_names:
-        logger.info(f"Removing existing '{output_column}' column before adding new results")
-        dataset = dataset.remove_columns([output_column])
+    # Add output column to dataset
     logger.info(f"Adding '{output_column}' column to dataset")
     dataset = dataset.add_column(output_column, all_outputs)
 
@@ -513,8 +480,8 @@ Examples:
     parser.add_argument(
         "--max-model-len",
         type=int,
-        default=32768,
-        help="Maximum model context length (default: 32768)",
+        default=8192,
+        help="Maximum model context length (default: 8192)",
     )
     parser.add_argument(
         "--max-tokens",
@@ -554,12 +521,6 @@ Examples:
         default="markdown",
         help="Column name for output text (default: markdown)",
     )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
     parser.add_argument(
         "--config",
         help="Config/subset name when pushing to Hub (for benchmarking multiple models in one repo)",
@@ -588,7 +549,6 @@ Examples:
         shuffle=args.shuffle,
         seed=args.seed,
         output_column=args.output_column,
-        overwrite=args.overwrite,
         config=args.config,
         create_pr=args.create_pr,
     )

glm-ocr-bucket.py

@@ -1,369 +0,0 @@
-# /// script
-# requires-python = ">=3.11"
-# dependencies = [
-#     "pillow",
-#     "pymupdf",
-#     "vllm",
-#     "torch",
-# ]
-#
-# [[tool.uv.index]]
-# url = "https://wheels.vllm.ai/nightly/cu129"
-#
-# [tool.uv]
-# prerelease = "allow"
-# override-dependencies = ["transformers>=5.1.0"]
-# ///
-
-"""
-OCR images and PDFs from a directory using GLM-OCR, writing markdown files.
-
-Designed to work with HF Buckets mounted as volumes via `hf jobs uv run -v ...`
-(requires huggingface_hub with PR #3936 volume mounting support).
-
-The script reads images/PDFs from INPUT_DIR, runs GLM-OCR via vLLM, and writes
-one .md file per image (or per PDF page) to OUTPUT_DIR, preserving directory structure.
-
-Input:                          Output:
-  /input/page1.png        →      /output/page1.md
-  /input/report.pdf       →      /output/report/page_001.md
-  (3 pages)                       /output/report/page_002.md
-                                  /output/report/page_003.md
-  /input/sub/photo.jpg    →      /output/sub/photo.md
-
-Examples:
-
-  # Local test
-  uv run glm-ocr-bucket.py ./test-images ./test-output
-
-  # HF Jobs with bucket volumes (PR #3936)
-  hf jobs uv run --flavor l4x1 \\
-      -s HF_TOKEN \\
-      -v bucket/user/ocr-input:/input:ro \\
-      -v bucket/user/ocr-output:/output \\
-      glm-ocr-bucket.py /input /output
-
-Model: zai-org/GLM-OCR (0.9B, 94.62% OmniDocBench V1.5, MIT licensed)
-"""
-
-import argparse
-import base64
-import io
-import logging
-import sys
-import time
-from pathlib import Path
-
-import torch
-from PIL import Image
-import os
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
-from vllm import LLM, SamplingParams
-
-logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
-logger = logging.getLogger(__name__)
-
-MODEL = "zai-org/GLM-OCR"
-
-TASK_PROMPTS = {
-    "ocr": "Text Recognition:",
-    "formula": "Formula Recognition:",
-    "table": "Table Recognition:",
-}
-
-IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".tiff", ".tif", ".bmp", ".webp"}
-
-
-def check_cuda_availability():
-    if not torch.cuda.is_available():
-        logger.error("CUDA is not available. This script requires a GPU.")
-        sys.exit(1)
-    logger.info(f"CUDA available. GPU: {torch.cuda.get_device_name(0)}")
-
-
-def make_ocr_message(image: Image.Image, task: str = "ocr") -> list[dict]:
-    """Create chat message for GLM-OCR from a PIL Image."""
-    image = image.convert("RGB")
-    buf = io.BytesIO()
-    image.save(buf, format="PNG")
-    data_uri = f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}"
-
-    return [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image_url", "image_url": {"url": data_uri}},
-                {"type": "text", "text": TASK_PROMPTS.get(task, TASK_PROMPTS["ocr"])},
-            ],
-        }
-    ]
-
-
-def discover_files(input_dir: Path) -> list[Path]:
-    """Walk input_dir recursively, returning sorted list of image and PDF files."""
-    files = []
-    for path in sorted(input_dir.rglob("*")):
-        if not path.is_file():
-            continue
-        ext = path.suffix.lower()
-        if ext in IMAGE_EXTENSIONS or ext == ".pdf":
-            files.append(path)
-    return files
-
-
-def prepare_images(
-    files: list[Path], input_dir: Path, output_dir: Path, pdf_dpi: int
-) -> list[tuple[Image.Image, Path]]:
-    """
-    Convert discovered files into (PIL.Image, output_md_path) pairs.
-
-    Images map 1:1. PDFs expand to one image per page in a subdirectory.
-    """
-    import fitz  # pymupdf
-
-    items: list[tuple[Image.Image, Path]] = []
-
-    for file_path in files:
-        rel = file_path.relative_to(input_dir)
-        ext = file_path.suffix.lower()
-
-        if ext == ".pdf":
-            # PDF → one .md per page in a subdirectory named after the PDF
-            pdf_output_dir = output_dir / rel.with_suffix("")
-            try:
-                doc = fitz.open(file_path)
-                num_pages = len(doc)
-                logger.info(f"PDF: {rel} ({num_pages} pages)")
-                for page_num in range(num_pages):
-                    page = doc[page_num]
-                    # Render at specified DPI
-                    zoom = pdf_dpi / 72.0
-                    mat = fitz.Matrix(zoom, zoom)
-                    pix = page.get_pixmap(matrix=mat)
-                    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-                    md_path = pdf_output_dir / f"page_{page_num + 1:03d}.md"
-                    items.append((img, md_path))
-                doc.close()
-            except Exception as e:
-                logger.error(f"Failed to open PDF {rel}: {e}")
-        else:
-            # Image �� single .md
-            try:
-                img = Image.open(file_path).convert("RGB")
-                md_path = output_dir / rel.with_suffix(".md")
-                items.append((img, md_path))
-            except Exception as e:
-                logger.error(f"Failed to open image {rel}: {e}")
-
-    return items
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="OCR images/PDFs from a directory using GLM-OCR, output markdown files.",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog="""
-Task modes:
-  ocr      Text recognition to markdown (default)
-  formula  LaTeX formula recognition
-  table    Table extraction (HTML)
-
-Examples:
-  uv run glm-ocr-bucket.py ./images ./output
-  uv run glm-ocr-bucket.py /input /output --task table --pdf-dpi 200
-
-HF Jobs with bucket volumes (requires huggingface_hub PR #3936):
-  hf jobs uv run --flavor l4x1 -s HF_TOKEN \\
-      -v bucket/user/input-bucket:/input:ro \\
-      -v bucket/user/output-bucket:/output \\
-      glm-ocr-bucket.py /input /output
-        """,
-    )
-    parser.add_argument("input_dir", help="Directory containing images and/or PDFs")
-    parser.add_argument("output_dir", help="Directory to write markdown output files")
-    parser.add_argument(
-        "--task",
-        choices=["ocr", "formula", "table"],
-        default="ocr",
-        help="OCR task mode (default: ocr)",
-    )
-    parser.add_argument(
-        "--batch-size", type=int, default=16, help="Batch size for vLLM (default: 16)"
-    )
-    parser.add_argument(
-        "--max-model-len",
-        type=int,
-        default=8192,
-        help="Max model context length (default: 8192)",
-    )
-    parser.add_argument(
-        "--max-tokens",
-        type=int,
-        default=8192,
-        help="Max output tokens (default: 8192)",
-    )
-    parser.add_argument(
-        "--gpu-memory-utilization",
-        type=float,
-        default=0.8,
-        help="GPU memory utilization (default: 0.8)",
-    )
-    parser.add_argument(
-        "--pdf-dpi",
-        type=int,
-        default=300,
-        help="DPI for PDF page rendering (default: 300)",
-    )
-    parser.add_argument(
-        "--temperature",
-        type=float,
-        default=0.01,
-        help="Sampling temperature (default: 0.01)",
-    )
-    parser.add_argument(
-        "--top-p", type=float, default=0.00001, help="Top-p sampling (default: 0.00001)"
-    )
-    parser.add_argument(
-        "--repetition-penalty",
-        type=float,
-        default=1.1,
-        help="Repetition penalty (default: 1.1)",
-    )
-    parser.add_argument(
-        "--verbose",
-        action="store_true",
-        help="Print resolved package versions",
-    )
-
-    args = parser.parse_args()
-
-    check_cuda_availability()
-
-    input_dir = Path(args.input_dir)
-    output_dir = Path(args.output_dir)
-
-    if not input_dir.is_dir():
-        logger.error(f"Input directory does not exist: {input_dir}")
-        sys.exit(1)
-
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    # Discover and prepare
-    start_time = time.time()
-
-    logger.info(f"Scanning {input_dir} for images and PDFs...")
-    files = discover_files(input_dir)
-    if not files:
-        logger.error(f"No image or PDF files found in {input_dir}")
-        sys.exit(1)
-
-    pdf_count = sum(1 for f in files if f.suffix.lower() == ".pdf")
-    img_count = len(files) - pdf_count
-    logger.info(f"Found {img_count} image(s) and {pdf_count} PDF(s)")
-
-    logger.info("Preparing images (rendering PDFs)...")
-    items = prepare_images(files, input_dir, output_dir, args.pdf_dpi)
-    if not items:
-        logger.error("No processable images after preparation")
-        sys.exit(1)
-
-    logger.info(f"Total images to OCR: {len(items)}")
-
-    # Init vLLM
-    logger.info(f"Initializing vLLM with {MODEL}...")
-    llm = LLM(
-        model=MODEL,
-        trust_remote_code=True,
-        max_model_len=args.max_model_len,
-        gpu_memory_utilization=args.gpu_memory_utilization,
-        limit_mm_per_prompt={"image": 1},
-    )
-
-    sampling_params = SamplingParams(
-        temperature=args.temperature,
-        top_p=args.top_p,
-        max_tokens=args.max_tokens,
-        repetition_penalty=args.repetition_penalty,
-    )
-
-    # Process in batches
-    errors = 0
-    processed = 0
-    total = len(items)
-
-    for batch_start in range(0, total, args.batch_size):
-        batch_end = min(batch_start + args.batch_size, total)
-        batch = items[batch_start:batch_end]
-        batch_num = batch_start // args.batch_size + 1
-        total_batches = (total + args.batch_size - 1) // args.batch_size
-
-        logger.info(f"Batch {batch_num}/{total_batches} ({processed}/{total} done)")
-
-        try:
-            messages = [make_ocr_message(img, task=args.task) for img, _ in batch]
-            outputs = llm.chat(messages, sampling_params)
-
-            for (_, md_path), output in zip(batch, outputs):
-                text = output.outputs[0].text.strip()
-                md_path.parent.mkdir(parents=True, exist_ok=True)
-                md_path.write_text(text, encoding="utf-8")
-                processed += 1
-
-        except Exception as e:
-            logger.error(f"Batch {batch_num} failed: {e}")
-            # Write error markers for failed batch
-            for _, md_path in batch:
-                md_path.parent.mkdir(parents=True, exist_ok=True)
-                md_path.write_text(f"[OCR ERROR: {e}]", encoding="utf-8")
-            errors += len(batch)
-            processed += len(batch)
-
-    elapsed = time.time() - start_time
-    elapsed_str = f"{elapsed / 60:.1f} min" if elapsed > 60 else f"{elapsed:.1f}s"
-
-    logger.info("=" * 50)
-    logger.info(f"Done! Processed {total} images in {elapsed_str}")
-    logger.info(f"  Output: {output_dir}")
-    logger.info(f"  Errors: {errors}")
-    if total > 0:
-        logger.info(f"  Speed: {total / elapsed:.2f} images/sec")
-
-    if args.verbose:
-        import importlib.metadata
-
-        logger.info("--- Package versions ---")
-        for pkg in ["vllm", "transformers", "torch", "pillow", "pymupdf"]:
-            try:
-                logger.info(f"  {pkg}=={importlib.metadata.version(pkg)}")
-            except importlib.metadata.PackageNotFoundError:
-                logger.info(f"  {pkg}: not installed")
-
-
-if __name__ == "__main__":
-    if len(sys.argv) == 1:
-        print("=" * 60)
-        print("GLM-OCR Bucket Script")
-        print("=" * 60)
-        print("\nOCR images/PDFs from a directory → markdown files.")
-        print("Designed for HF Buckets mounted as volumes (PR #3936).")
-        print()
-        print("Usage:")
-        print("  uv run glm-ocr-bucket.py INPUT_DIR OUTPUT_DIR")
-        print()
-        print("Examples:")
-        print("  uv run glm-ocr-bucket.py ./images ./output")
-        print("  uv run glm-ocr-bucket.py /input /output --task table")
-        print()
-        print("HF Jobs with bucket volumes:")
-        print("  hf jobs uv run --flavor l4x1 -s HF_TOKEN \\")
-        print("      -v bucket/user/ocr-input:/input:ro \\")
-        print("      -v bucket/user/ocr-output:/output \\")
-        print("      glm-ocr-bucket.py /input /output")
-        print()
-        print("For full help: uv run glm-ocr-bucket.py --help")
-        sys.exit(0)
-
-    main()

glm-ocr-v2.py

@@ -69,10 +69,6 @@ from datasets import load_dataset
 from huggingface_hub import CommitScheduler, DatasetCard, HfApi, login
 from PIL import Image
 from toolz import partition_all
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
 from vllm import LLM, SamplingParams
 
 logging.basicConfig(level=logging.INFO)
@@ -136,28 +132,6 @@ def check_cuda_availability():
         logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
 
 
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
 def make_ocr_message(
     image: Union[Image.Image, Dict[str, Any], str],
     task: str = "ocr",
@@ -408,7 +382,6 @@ def main(
     shuffle: bool = False,
     seed: int = 42,
     output_column: str = "markdown",
-    overwrite: bool = False,
     verbose: bool = False,
     config: str = None,
     create_pr: bool = False,
@@ -461,11 +434,6 @@ def main(
             f"Column '{image_column}' not found. Available: {dataset.column_names}"
         )
 
-    # Fail fast if the output column would collide with an existing input column.
-    # The incremental path derives batch datasets via dataset.select(...), so dropping
-    # it here (on --overwrite) also keeps those per-batch add_columns clean.
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
     if shuffle:
         logger.info(f"Shuffling dataset with seed {seed}")
         dataset = dataset.shuffle(seed=seed)
@@ -983,12 +951,6 @@ Examples:
         default="markdown",
         help="Column name for output text (default: markdown)",
     )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
     parser.add_argument(
         "--verbose",
         action="store_true",
@@ -1033,7 +995,6 @@ Examples:
         shuffle=args.shuffle,
         seed=args.seed,
         output_column=args.output_column,
-        overwrite=args.overwrite,
         verbose=args.verbose,
         config=args.config,
         create_pr=args.create_pr,

glm-ocr.py

@@ -2,6 +2,7 @@
 # requires-python = ">=3.11"
 # dependencies = [
 #     "datasets>=3.1.0",
+#     "pyarrow>=17.0.0,<18.0.0",
 #     "huggingface-hub",
 #     "pillow",
 #     "vllm",
@@ -52,21 +53,13 @@ import os
 import sys
 import time
 from datetime import datetime
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Union
 
 import torch
 from datasets import load_dataset
 from huggingface_hub import DatasetCard, login
 from PIL import Image
 from toolz import partition_all
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
-# Same story for DeepGEMM (nightly vLLM): its init calls _find_cuda_home, which asserts on the
-# nvcc-less base image (a non-fatal warning that clutters the log and hides the real traceback).
-# Greedy OCR doesn't need the DeepGEMM JIT path, so disable it explicitly.
-os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
 from vllm import LLM, SamplingParams
 
 logging.basicConfig(level=logging.INFO)
@@ -92,49 +85,9 @@ def check_cuda_availability():
         logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
 
 
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
-def downscale_to_max_pixels(img: Image.Image, max_pixels: Optional[int]) -> Image.Image:
-    """Shrink an image so width*height <= max_pixels, preserving aspect ratio.
-
-    GLM-OCR does no internal resizing and its card gives no resolution guidance. Capping
-    input pixels bounds both image tokens and vision-encoder memory, a safety valve for very
-    large (multi-MP) scans that can pressure GPU memory at high batch sizes. No-op when
-    max_pixels is None or the image is already small enough (never upscales)."""
-    if not max_pixels:
-        return img
-    w, h = img.size
-    if w * h <= max_pixels:
-        return img
-    scale = (max_pixels / (w * h)) ** 0.5
-    new_size = (max(1, int(w * scale)), max(1, int(h * scale)))
-    return img.resize(new_size, Image.Resampling.LANCZOS)
-
-
 def make_ocr_message(
     image: Union[Image.Image, Dict[str, Any], str],
     task: str = "ocr",
-    max_pixels: Optional[int] = None,
 ) -> List[Dict]:
     """
     Create chat message for OCR processing.
@@ -155,9 +108,6 @@ def make_ocr_message(
     # Convert to RGB
     pil_img = pil_img.convert("RGB")
 
-    # Optionally cap resolution to protect the vision encoder from OOM on huge scans
-    pil_img = downscale_to_max_pixels(pil_img, max_pixels)
-
     # Convert to base64 data URI
     buf = io.BytesIO()
     pil_img.save(buf, format="PNG")
@@ -199,15 +149,6 @@ def create_dataset_card(
         "table": "table recognition",
     }
 
-    # Canonical provenance stamp (see AGENTS.md): Jobs claim gated on JOB_ID, set by HF Jobs in-container.
-    on_jobs = os.environ.get("JOB_ID") is not None
-    hw = os.environ.get("ACCELERATOR") or ""
-    origin = (
-        "Produced on [Hugging Face Jobs](https://huggingface.co/docs/huggingface_hub/guides/jobs)"
-        + (f" (`{hw}`)" if hw else "")
-    ) if on_jobs else "Generated"
-    jobs_tag = "\n- hf-jobs" if on_jobs else ""
-
     return f"""---
 tags:
 - ocr
@@ -215,7 +156,7 @@ tags:
 - glm-ocr
 - markdown
 - uv-script
-- generated{jobs_tag}
+- generated
 ---
 
 # Document OCR using {model_name}
@@ -261,16 +202,16 @@ The dataset contains all original columns plus:
 
 ## Reproduction
 
-{origin} with the [`glm-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/raw/main/glm-ocr.py) recipe from [uv-scripts](https://huggingface.co/uv-scripts). Run it yourself:
-
 ```bash
-hf jobs uv run https://huggingface.co/datasets/uv-scripts/ocr/raw/main/glm-ocr.py \\
+uv run https://huggingface.co/datasets/uv-scripts/ocr/raw/main/glm-ocr.py \\
     {source_dataset} \\
     <output-dataset> \\
     --image-column {image_column} \\
     --batch-size {batch_size} \\
     --task {task}
 ```
+
+Generated with [UV Scripts](https://huggingface.co/uv-scripts)
 """
 
 
@@ -280,7 +221,6 @@ def main(
     image_column: str = "image",
     batch_size: int = 16,
     max_model_len: int = 8192,
-    max_pixels: Optional[int] = None,
     max_tokens: int = 8192,
     temperature: float = 0.01,
     top_p: float = 0.00001,
@@ -294,7 +234,6 @@ def main(
     shuffle: bool = False,
     seed: int = 42,
     output_column: str = "markdown",
-    overwrite: bool = False,
     verbose: bool = False,
     config: str = None,
     create_pr: bool = False,
@@ -326,9 +265,6 @@ def main(
             f"Column '{image_column}' not found. Available: {dataset.column_names}"
         )
 
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
     if shuffle:
         logger.info(f"Shuffling dataset with seed {seed}")
         dataset = dataset.shuffle(seed=seed)
@@ -379,10 +315,7 @@ def main(
         )
 
         try:
-            batch_messages = [
-                make_ocr_message(img, task=task, max_pixels=max_pixels)
-                for img in batch_images
-            ]
+            batch_messages = [make_ocr_message(img, task=task) for img in batch_images]
 
             outputs = llm.chat(batch_messages, sampling_params)
 
@@ -572,17 +505,6 @@ Examples:
         default=8192,
         help="Maximum model context length (default: 8192)",
     )
-    parser.add_argument(
-        "--max-pixels",
-        type=int,
-        default=None,
-        help=(
-            "Optional cap on input image pixels (width*height); larger scans are "
-            "downscaled (aspect preserved) before OCR. GLM-OCR does no internal resizing, "
-            "so this bounds vision-encoder memory on very large scans — set e.g. 4000000 "
-            "if you hit a GPU OOM at high batch sizes on a big-page corpus. Default: no cap."
-        ),
-    )
     parser.add_argument(
         "--max-tokens",
         type=int,
@@ -654,12 +576,6 @@ Examples:
         default="markdown",
         help="Column name for output text (default: markdown)",
     )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
     parser.add_argument(
         "--verbose",
         action="store_true",
@@ -674,7 +590,6 @@ Examples:
         image_column=args.image_column,
         batch_size=args.batch_size,
         max_model_len=args.max_model_len,
-        max_pixels=args.max_pixels,
         max_tokens=args.max_tokens,
         temperature=args.temperature,
         top_p=args.top_p,
@@ -688,7 +603,6 @@ Examples:
         shuffle=args.shuffle,
         seed=args.seed,
         output_column=args.output_column,
-        overwrite=args.overwrite,
         verbose=args.verbose,
         config=args.config,
         create_pr=args.create_pr,

hunyuan-ocr-1.5.py

@@ -1,884 +0,0 @@
-# /// script
-# requires-python = ">=3.11"
-# dependencies = [
-#     "datasets>=4.0.0",
-#     "huggingface-hub",
-#     "pillow",
-#     "vllm>=0.18.1",
-#     "transformers<5.13",  # vLLM ≤0.24.0's HunyuanVL processor breaks on transformers 5.13
-#                           # (string-key AutoImageProcessor.register; fixed in vllm#47872).
-#                           # Drop this cap once that fix ships in a stable vLLM release.
-#     "tqdm",
-#     "toolz",
-#     "torch",
-# ]
-# ///
-
-"""
-Convert document images to markdown using HunyuanOCR-1.5 with vLLM.
-
-HunyuanOCR-1.5 is a lightweight ~1B-parameter, end-to-end OCR-specialized VLM
-from Tencent. It keeps the validated 1.0 backbone but extends the max image
-resolution to 4K and the context window to 128K, and adds targeted long-tail
-capabilities (low-resource / ancient-script OCR, multi-image text QA). Per the
-technical report (arXiv:2607.04884) it is faster than dots.ocr / DeepSeek-OCR-2
-and top-tier on OmniDocBench v1.6. This script runs it offline via vLLM.
-
-Features:
-- 📝 End-to-end document parsing to markdown (tables → HTML, formulas → LaTeX)
-- 🧩 Structured / layout-aware parsing
-- 📍 Text spotting with coordinates (JSON or Hunyuan format)
-- 📐 Formula (LaTeX) and 📊 table (HTML) recognition
-- 📈 Chart parsing (Mermaid / Markdown)
-- 🌐 Document + general-scene translation (→ zh / → en)
-- 🎯 Compact model (~1B parameters)
-
-Model: tencent/HunyuanOCR
-  On 2026-07-06 Tencent replaced the repo root in-place with HunyuanOCR-1.5
-  (1.0 archived under `v1.0/`, no git tag). So the repo *root* — the default
-  here — is now 1.5. The sibling recipe `hunyuan-ocr.py` pins the last 1.0
-  commit by revision to keep the 1.0 behavior; this script deliberately tracks
-  root (1.5).
-
-vLLM: 0.18.1 (release) is the first stable wheel with native
-  `HunYuanVLForConditionalGeneration` support for autoregressive decoding — no
-  nightly or patch needed for batch OCR. The floor stays at 0.18.1; a bare
-  `vllm` resolves to the latest stable (0.24.0 as of 2026-07), which also works
-  once transformers is capped <5.13 (see the deps block for why). The DFlash
-  speculative-decoding draft (a per-request *latency* win that needs a vLLM
-  nightly) is intentionally NOT implemented: it does not change offline batch
-  throughput or output distribution.
-
-trust_remote_code=True per the model card (the processor ships custom code).
-
-License: Tencent Hunyuan Community License (territory excludes EU/UK/South Korea)
-https://huggingface.co/tencent/HunyuanOCR/blob/main/LICENSE
-
-Note: batch_size defaults to 16 (untested on this arch as of writing — 1.5 on
-vLLM ≥0.18.1 should batch fine, unlike the 1.0 V1 batching issue; will be
-smoke-tested). Lower it if you hit engine errors.
-
-Post-processing note: only the shared tail-repetition cleanup
-(`clean_repeated_substrings`, byte-for-byte from the official toolkit) is
-ported. The upstream doc_parse-only markdown normalization (10 OmniDocBench
-GT-alignment regex passes in `hunyuan_utils.process_one`) is intentionally NOT
-ported — it is benchmark-GT alignment, not general OCR, and would bloat this
-self-contained recipe. For bench-exact output, use Tencent's toolkit directly.
-"""
-
-import argparse
-import base64
-import io
-import json
-import logging
-import os
-import sys
-import time
-from datetime import datetime
-from typing import Any, Dict, List, Union
-
-import torch
-from datasets import load_dataset
-from huggingface_hub import DatasetCard, login
-from PIL import Image
-from toolz import partition_all
-from tqdm.auto import tqdm
-
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
-from vllm import LLM, SamplingParams
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-# ────────────────────────────────────────────────────────────────
-# HunyuanOCR-1.5 official task prompts.
-# Reproduced VERBATIM from the shipped client's `hunyuan_tasks.py`:
-#   https://github.com/Tencent-Hunyuan/HunyuanOCR (inference/*/hunyuan_tasks.py)
-# The model card only prints the `doc_parse` prompt; the other 11 live only in
-# the client. Prompts are FIXED per task type — upstream deliberately does NOT
-# expose free-form prompt editing because hand-tweaked instructions were
-# observed to silently degrade quality (users pick a *task*, not a prompt).
-# All prompts are Chinese-language, including for English documents — this is
-# the officially recommended wording; the card provides no English variants.
-# ────────────────────────────────────────────────────────────────
-
-TASK_PROMPTS = {
-    # 端到端文档解析
-    "doc_parse": "提取文档图片中正文的所有信息用markdown格式表示，其中页眉、页脚部分忽略，"
-    "表格用html格式表达，文档中公式用latex格式表示，按照阅读顺序组织进行解析。",
-    # 结构化解析（古文、街景等非文档结构化场景）
-    "structured_parse": "提取图中的文字。",
-    # Spotting — JSON 格式
-    "spotting_json": "检测并识别图中所有的文字行，请按从上到下、从左到右的阅读顺序进行识别。 "
-    "输出格式为 JSON 数组，每个元素必须包含："
-    '"box": [xmin, ymin, xmax, ymax]（坐标需归一化到 [0, 1000] 范围内）；'
-    '"text": "识别出的文字内容"。 '
-    "注意：请直接输出 JSON 数组，不要包含任何多余的描述性文字。",
-    # Spotting — Hunyuan 模式
-    "spotting_hunyuan": "检测并识别图片中的文字，将文本坐标格式化输出。",
-    # 版式分析
-    "layout": "按照阅读顺序解析图中的版式信息。",
-    # 版式分析 + 解析
-    "layout_parse": "提取文档图片中所有内容用markdown格式表示，表格用html格式表达，"
-    "文档中公式用latex格式表示，请按照阅读顺序组织进行全文解析，并输出版式分析信息。",
-    # 图表解析
-    "chart_parse": "解析图中的图表，对于流程图使用Mermaid格式表示，其他图表使用Markdown格式表示。",
-    # 公式解析
-    "formula": "识别图片中的公式，用LaTeX格式表示。",
-    # 表格解析
-    "table": "把图中的表格解析为HTML。",
-    # 文档英译中
-    "doc_trans_en2zh": "先解析文档，再将文档内容翻译为中文，其中页眉、页脚忽略，"
-    "公式用latex格式表示，表格用html格式表示。",
-    # 通用场景翻译 other2en
-    "trans_other2en": "按照阅读顺序，提取图中文字，公式用latex格式表示，表格用markdown格式表示，"
-    "再将文字内容翻译为英文。",
-    # 通用场景翻译 other2zh
-    "trans_other2zh": "按照阅读顺序，提取图中文字，公式用latex格式表示，表格用markdown格式表示，"
-    "再将文字内容翻译为中文。",
-}
-
-# English glosses for --help / the no-args banner (upstream ships Chinese ones).
-TASK_DESCRIPTIONS = {
-    "doc_parse": "End-to-end doc parse (body→markdown, tables→HTML, formulas→LaTeX, headers/footers ignored). Default.",
-    "structured_parse": "Structured parse for non-document scenes (ancient scripts, street signs) — extract all text.",
-    "spotting_json": "Text detect+recognize as a JSON array (box normalized to 0-1000 + text).",
-    "spotting_hunyuan": "Text detect+recognize in Hunyuan coordinate format.",
-    "layout": "Layout analysis in reading order.",
-    "layout_parse": "Layout analysis + full-document parse (markdown/HTML/LaTeX).",
-    "chart_parse": "Chart parsing (flowcharts→Mermaid, other charts→Markdown).",
-    "formula": "Formula recognition → LaTeX.",
-    "table": "Table parsing → HTML.",
-    "doc_trans_en2zh": "Document translation to Chinese (parse then translate; formulas LaTeX, tables HTML).",
-    "trans_other2en": "General-scene extraction + translation to English.",
-    "trans_other2zh": "General-scene extraction + translation to Chinese.",
-}
-
-DEFAULT_TASK = "doc_parse"
-
-# Sampling params LOCKED by the model card across all official setups so outputs
-# are comparable: temperature=0.0, top_p=1.0, top_k=-1, repetition_penalty=1.08.
-# Only repetition_penalty is exposed as a flag (the others are fixed for
-# deterministic OCR); repetition_penalty is the model's built-in anti-repeat.
-DEFAULT_REPETITION_PENALTY = 1.08
-
-
-def clean_repeated_substrings(text: str, min_repeats: int = 10) -> str:
-    """Trim a long repeated suffix as a final safety net against greedy-decoding
-    degeneration. Byte-for-byte from the official `hunyuan_utils.py`.
-    """
-    n = len(text)
-    if n < 2000:
-        return text
-    for length in range(2, n // min_repeats + 1):
-        candidate = text[-length:]
-        count = 0
-        i = n - length
-        while i >= 0 and text[i : i + length] == candidate:
-            count += 1
-            i -= length
-        if count >= min_repeats:
-            return text[: n - length * (count - 1)]
-    return text
-
-
-def check_cuda_availability():
-    """Check if CUDA is available and exit if not."""
-    if not torch.cuda.is_available():
-        logger.error("CUDA is not available. This script requires a GPU.")
-        logger.error("Please run on a machine with a CUDA-capable GPU.")
-        sys.exit(1)
-    else:
-        logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
-
-
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error(
-        "Choose a different --output-column, or pass --overwrite to replace them."
-    )
-    sys.exit(1)
-
-
-def get_prompt(task_type: str) -> str:
-    """Return the official prompt for a task type."""
-    if task_type not in TASK_PROMPTS:
-        raise ValueError(
-            f"Unknown task type: {task_type}. Available: {list(TASK_PROMPTS.keys())}"
-        )
-    return TASK_PROMPTS[task_type]
-
-
-def make_ocr_message(
-    image: Union[Image.Image, Dict[str, Any], str],
-    prompt: str,
-) -> List[Dict]:
-    """Create the chat messages for one image + prompt.
-
-    Mirrors the official client: an empty system message followed by a user turn
-    with the image *before* the text. The empty system content pins "no system
-    prompt" (matching how the model is served) rather than letting the chat
-    template inject a default.
-    """
-    # Convert to PIL Image if needed
-    if isinstance(image, Image.Image):
-        pil_img = image
-    elif isinstance(image, dict) and "bytes" in image:
-        pil_img = Image.open(io.BytesIO(image["bytes"]))
-    elif isinstance(image, str):
-        pil_img = Image.open(image)
-    else:
-        raise ValueError(f"Unsupported image type: {type(image)}")
-
-    # Convert to RGB
-    pil_img = pil_img.convert("RGB")
-
-    # Convert to base64 data URI
-    buf = io.BytesIO()
-    pil_img.save(buf, format="PNG")
-    data_uri = f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}"
-
-    return [
-        {"role": "system", "content": ""},
-        {
-            "role": "user",
-            "content": [
-                {"type": "image_url", "image_url": {"url": data_uri}},
-                {"type": "text", "text": prompt},
-            ],
-        },
-    ]
-
-
-def create_dataset_card(
-    source_dataset: str,
-    model: str,
-    num_samples: int,
-    processing_time: str,
-    batch_size: int,
-    max_model_len: int,
-    max_tokens: int,
-    repetition_penalty: float,
-    gpu_memory_utilization: float,
-    image_column: str = "image",
-    output_column: str = "markdown",
-    split: str = "train",
-    task_type: str = "doc_parse",
-) -> str:
-    """Create a dataset card documenting the OCR process."""
-    model_name = model.split("/")[-1]
-
-    return f"""---
-tags:
-- ocr
-- document-processing
-- hunyuan-ocr-1.5
-- multilingual
-- markdown
-- uv-script
-- generated
----
-
-# Document OCR using {model_name} (HunyuanOCR-1.5)
-
-This dataset contains OCR results from images in [{source_dataset}](https://huggingface.co/datasets/{source_dataset}) using HunyuanOCR-1.5, a lightweight ~1B end-to-end OCR VLM from Tencent (128K context, 4K max image resolution).
-
-Model license: [Tencent Hunyuan Community License](https://huggingface.co/tencent/HunyuanOCR/blob/main/LICENSE) (territory excludes EU/UK/South Korea).
-
-## Processing Details
-
-- **Source Dataset**: [{source_dataset}](https://huggingface.co/datasets/{source_dataset})
-- **Model**: [{model}](https://huggingface.co/{model})
-- **Number of Samples**: {num_samples:,}
-- **Processing Time**: {processing_time}
-- **Processing Date**: {datetime.now().strftime("%Y-%m-%d %H:%M UTC")}
-
-### Configuration
-
-- **Image Column**: `{image_column}`
-- **Output Column**: `{output_column}`
-- **Dataset Split**: `{split}`
-- **Task Type**: `{task_type}`
-- **Batch Size**: {batch_size}
-- **Max Model Length**: {max_model_len:,} tokens
-- **Max Output Tokens**: {max_tokens:,}
-- **Repetition Penalty**: {repetition_penalty}
-- **GPU Memory Utilization**: {gpu_memory_utilization:.1%}
-
-## Model Information
-
-HunyuanOCR-1.5 is a lightweight end-to-end OCR VLM that excels at:
-- 📝 **Document Parsing** - Full markdown extraction in reading order
-- 🧩 **Structured / Layout Parsing** - Layout-aware full-document parse
-- 📊 **Table Extraction** - HTML format tables
-- 📐 **Formula Recognition** - LaTeX format formulas
-- 📈 **Chart Parsing** - Mermaid / Markdown format
-- 📍 **Text Spotting** - Detection with coordinates (JSON / Hunyuan)
-- 🌐 **Translation** - Document and general-scene translation (→ zh / → en)
-
-Per the technical report ([arXiv:2607.04884](https://arxiv.org/pdf/2607.04884)),
-1.5 is faster than dots.ocr / DeepSeek-OCR-2 and top-tier on OmniDocBench v1.6.
-
-## Task Types Available
-
-- `doc_parse` - End-to-end document parsing (default)
-- `structured_parse` - Non-document structured scenes (ancient scripts, street signs)
-- `spotting_json` - Text detection + recognition as JSON array (box 0-1000 + text)
-- `spotting_hunyuan` - Text detection + recognition, Hunyuan coordinate format
-- `layout` - Layout analysis in reading order
-- `layout_parse` - Layout analysis + full-document parse
-- `chart_parse` - Chart parsing (flowcharts → Mermaid, others → Markdown)
-- `formula` - Formula recognition → LaTeX
-- `table` - Table parsing → HTML
-- `doc_trans_en2zh` - Document translation to Chinese
-- `trans_other2en` - General-scene extraction + translation to English
-- `trans_other2zh` - General-scene extraction + translation to Chinese
-
-## Dataset Structure
-
-The dataset contains all original columns plus:
-- `{output_column}`: The extracted text (markdown for `doc_parse`, else the task's format)
-- `inference_info`: JSON list tracking all OCR models applied to this dataset
-
-## Usage
-
-```python
-from datasets import load_dataset
-import json
-
-# Load the dataset
-dataset = load_dataset("{{output_dataset_id}}", split="{split}")
-
-# Access the extracted text
-for example in dataset:
-    print(example["{output_column}"])
-    break
-
-# View all OCR models applied to this dataset
-inference_info = json.loads(dataset[0]["inference_info"])
-for info in inference_info:
-    print(f"Column: {{info['column_name']}} - Model: {{info['model_id']}}")
-```
-
-## Reproduction
-
-This dataset was generated using the [uv-scripts/ocr](https://huggingface.co/datasets/uv-scripts/ocr) HunyuanOCR-1.5 script:
-
-```bash
-uv run https://huggingface.co/datasets/uv-scripts/ocr/raw/main/hunyuan-ocr-1.5.py \\
-    {source_dataset} \\
-    <output-dataset> \\
-    --image-column {image_column} \\
-    --batch-size {batch_size} \\
-    --task-type {task_type} \\
-    --max-model-len {max_model_len} \\
-    --max-tokens {max_tokens} \\
-    --gpu-memory-utilization {gpu_memory_utilization}
-```
-
-Generated with [UV Scripts](https://huggingface.co/uv-scripts)
-"""
-
-
-def main(
-    input_dataset: str,
-    output_dataset: str,
-    image_column: str = "image",
-    batch_size: int = 16,
-    model: str = "tencent/HunyuanOCR",
-    revision: str = None,
-    max_model_len: int = 32768,
-    max_tokens: int = 8192,
-    repetition_penalty: float = DEFAULT_REPETITION_PENALTY,
-    gpu_memory_utilization: float = 0.8,
-    hf_token: str = None,
-    split: str = "train",
-    max_samples: int = None,
-    private: bool = False,
-    shuffle: bool = False,
-    seed: int = 42,
-    task_type: str = DEFAULT_TASK,
-    custom_prompt: str = None,
-    output_column: str = "markdown",
-    overwrite: bool = False,
-    clean_output: bool = True,
-    config: str = None,
-    create_pr: bool = False,
-    verbose: bool = False,
-):
-    """Process images from an HF dataset through HunyuanOCR-1.5."""
-
-    # Check CUDA availability first
-    check_cuda_availability()
-
-    # Context-length invariant (config.json): text max_position_embeddings=131072;
-    # the vision processor caps a single image at img_max_token_num=16384. So the
-    # default budget holds without resizing: 16384 (image) + prompt + 8192 (output)
-    # ≈ 25k ≤ 32768 (default max_model_len). Enforce max_tokens ≤ max_model_len ≤ 131072.
-    if max_model_len > 131072:
-        logger.error(
-            f"--max-model-len {max_model_len} exceeds the model's max context (131072)."
-        )
-        sys.exit(1)
-    if max_tokens > max_model_len:
-        logger.error(
-            f"--max-tokens ({max_tokens}) cannot exceed --max-model-len ({max_model_len})."
-        )
-        sys.exit(1)
-
-    # Track processing start time
-    start_time = datetime.now()
-
-    # Login to HF if token provided
-    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")
-    if HF_TOKEN:
-        login(token=HF_TOKEN)
-
-    # Determine prompt to use
-    if custom_prompt:
-        prompt = custom_prompt
-        logger.warning(
-            "Using --custom-prompt. Note: upstream deliberately locks prompts per "
-            "task type — hand-tweaked instructions can silently degrade quality."
-        )
-        logger.info(f"Custom prompt: {prompt[:60]}...")
-    else:
-        prompt = get_prompt(task_type)
-        logger.info(f"Using task type: {task_type}")
-
-    # Load dataset
-    logger.info(f"Loading dataset: {input_dataset}")
-    dataset = load_dataset(input_dataset, split=split)
-
-    # Validate image column
-    if image_column not in dataset.column_names:
-        raise ValueError(
-            f"Column '{image_column}' not found. Available: {dataset.column_names}"
-        )
-
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
-    # Shuffle if requested
-    if shuffle:
-        logger.info(f"Shuffling dataset with seed {seed}")
-        dataset = dataset.shuffle(seed=seed)
-
-    # Limit samples if requested
-    if max_samples:
-        dataset = dataset.select(range(min(max_samples, len(dataset))))
-        logger.info(f"Limited to {len(dataset)} samples")
-
-    # Initialize vLLM model
-    logger.info(f"Initializing vLLM with model: {model}")
-    logger.info("This may take a few minutes on first run...")
-
-    llm = LLM(
-        model=model,
-        revision=revision,
-        trust_remote_code=True,
-        max_model_len=max_model_len,
-        gpu_memory_utilization=gpu_memory_utilization,
-        limit_mm_per_prompt={"image": 1},
-    )
-
-    # Locked sampling per the model card (deterministic OCR); only repetition_penalty
-    # is user-tunable.
-    sampling_params = SamplingParams(
-        temperature=0.0,
-        top_p=1.0,
-        top_k=-1,
-        repetition_penalty=repetition_penalty,
-        max_tokens=max_tokens,
-        skip_special_tokens=True,
-    )
-
-    logger.info(f"Processing {len(dataset)} images in batches of {batch_size}")
-    logger.info(f"Output will be written to column: {output_column}")
-
-    # Process images in batches
-    all_outputs = []
-
-    for batch_indices in tqdm(
-        partition_all(batch_size, range(len(dataset))),
-        total=(len(dataset) + batch_size - 1) // batch_size,
-        desc="HunyuanOCR-1.5 processing",
-    ):
-        batch_indices = list(batch_indices)
-        batch_images = [dataset[i][image_column] for i in batch_indices]
-
-        try:
-            # Create messages for batch
-            batch_messages = [make_ocr_message(img, prompt) for img in batch_images]
-
-            # Process with vLLM
-            outputs = llm.chat(batch_messages, sampling_params)
-
-            # Extract outputs
-            for output in outputs:
-                text = output.outputs[0].text.strip()
-                # Clean repeated substrings if enabled
-                if clean_output:
-                    text = clean_repeated_substrings(text)
-                all_outputs.append(text)
-
-        except Exception as e:
-            logger.error(f"Error processing batch: {e}")
-            # Add error placeholders for failed batch
-            all_outputs.extend(["[OCR ERROR]"] * len(batch_images))
-
-    # Calculate processing time
-    processing_duration = datetime.now() - start_time
-    processing_time_str = f"{processing_duration.total_seconds() / 60:.1f} min"
-
-    # Add output column to dataset
-    logger.info(f"Adding '{output_column}' column to dataset")
-    dataset = dataset.add_column(output_column, all_outputs)
-
-    # Handle inference_info tracking (for multi-model comparisons)
-    inference_entry = {
-        "model_id": model,
-        "model_name": "HunyuanOCR-1.5",
-        "model_revision": revision or "main",
-        "column_name": output_column,
-        "timestamp": datetime.now().isoformat(),
-        "task_type": task_type if not custom_prompt else "custom",
-        "repetition_penalty": repetition_penalty,
-    }
-
-    if "inference_info" in dataset.column_names:
-        # Append to existing inference info
-        logger.info("Updating existing inference_info column")
-
-        def update_inference_info(example):
-            try:
-                existing_info = (
-                    json.loads(example["inference_info"])
-                    if example["inference_info"]
-                    else []
-                )
-            except (json.JSONDecodeError, TypeError):
-                existing_info = []
-
-            existing_info.append(inference_entry)
-            return {"inference_info": json.dumps(existing_info)}
-
-        dataset = dataset.map(update_inference_info)
-    else:
-        # Create new inference_info column
-        logger.info("Creating new inference_info column")
-        inference_list = [json.dumps([inference_entry])] * len(dataset)
-        dataset = dataset.add_column("inference_info", inference_list)
-
-    # Push to hub with retry and XET fallback
-    logger.info(f"Pushing to {output_dataset}")
-    commit_msg = f"Add HunyuanOCR-1.5 OCR results ({len(dataset)} samples)" + (
-        f" [{config}]" if config else ""
-    )
-    max_retries = 3
-    for attempt in range(1, max_retries + 1):
-        try:
-            if attempt > 1:
-                logger.warning("Disabling XET (fallback to HTTP upload)")
-                os.environ["HF_HUB_DISABLE_XET"] = "1"
-            dataset.push_to_hub(
-                output_dataset,
-                private=private,
-                token=HF_TOKEN,
-                max_shard_size="500MB",
-                **({"config_name": config} if config else {}),
-                create_pr=create_pr,
-                commit_message=commit_msg,
-            )
-            break
-        except Exception as e:
-            logger.error(f"Upload attempt {attempt}/{max_retries} failed: {e}")
-            if attempt < max_retries:
-                delay = 30 * (2 ** (attempt - 1))
-                logger.info(f"Retrying in {delay}s...")
-                time.sleep(delay)
-            else:
-                logger.error("All upload attempts failed. OCR results are lost.")
-                sys.exit(1)
-
-    # Create and push dataset card (skip when creating PR to avoid conflicts)
-    if not create_pr:
-        logger.info("Creating dataset card")
-        card_content = create_dataset_card(
-            source_dataset=input_dataset,
-            model=model,
-            num_samples=len(dataset),
-            processing_time=processing_time_str,
-            batch_size=batch_size,
-            max_model_len=max_model_len,
-            max_tokens=max_tokens,
-            repetition_penalty=repetition_penalty,
-            gpu_memory_utilization=gpu_memory_utilization,
-            image_column=image_column,
-            output_column=output_column,
-            split=split,
-            task_type=task_type if not custom_prompt else "custom",
-        )
-
-        card = DatasetCard(card_content)
-        card.push_to_hub(output_dataset, token=HF_TOKEN)
-
-    logger.info("HunyuanOCR-1.5 processing complete!")
-    logger.info(
-        f"Dataset available at: https://huggingface.co/datasets/{output_dataset}"
-    )
-    logger.info(f"Processing time: {processing_time_str}")
-
-    if verbose:
-        import importlib.metadata
-
-        logger.info("--- Resolved package versions ---")
-        for pkg in [
-            "vllm",
-            "transformers",
-            "torch",
-            "datasets",
-            "pyarrow",
-            "pillow",
-        ]:
-            try:
-                logger.info(f"  {pkg}=={importlib.metadata.version(pkg)}")
-            except importlib.metadata.PackageNotFoundError:
-                logger.info(f"  {pkg}: not installed")
-        logger.info("--- End versions ---")
-
-
-if __name__ == "__main__":
-    # Show example usage if no arguments
-    if len(sys.argv) == 1:
-        print("=" * 80)
-        print("HunyuanOCR-1.5 Document Processing")
-        print("=" * 80)
-        print(
-            "\nLightweight ~1B end-to-end OCR VLM from Tencent (128K context, 4K images)"
-        )
-        print("\nFeatures:")
-        print("- 📝 End-to-end document parsing to markdown")
-        print("- 📊 Table extraction (HTML format)")
-        print("- 📐 Formula recognition (LaTeX format)")
-        print("- 📍 Text spotting with coordinates (JSON / Hunyuan)")
-        print("- 📈 Chart parsing (Mermaid / Markdown)")
-        print("- 🌐 Document + general-scene translation (→ zh / → en)")
-        print("\nExample usage:")
-        print("\n1. Basic document parsing:")
-        print("   uv run hunyuan-ocr-1.5.py input-dataset output-dataset")
-        print("\n2. Formula extraction:")
-        print("   uv run hunyuan-ocr-1.5.py math-docs formulas --task-type formula")
-        print("\n3. Table extraction:")
-        print("   uv run hunyuan-ocr-1.5.py docs tables --task-type table")
-        print("\n4. Text spotting as JSON (box + text):")
-        print("   uv run hunyuan-ocr-1.5.py images spotted --task-type spotting_json")
-        print("\n5. Translate a document to Chinese:")
-        print(
-            "   uv run hunyuan-ocr-1.5.py en-docs zh-docs --task-type doc_trans_en2zh"
-        )
-        print("\n6. Running on HF Jobs:")
-        print("   hf jobs uv run --flavor l4x1 \\")
-        print(
-            '     -e HF_TOKEN=$(python3 -c "from huggingface_hub import get_token; print(get_token())") \\'
-        )
-        print(
-            "     https://huggingface.co/datasets/uv-scripts/ocr/raw/main/hunyuan-ocr-1.5.py \\"
-        )
-        print("       input-dataset output-dataset")
-        print("\n" + "=" * 80)
-        print("\nFor full help, run: uv run hunyuan-ocr-1.5.py --help")
-        sys.exit(0)
-
-    task_help = "\n".join(f"  {k:18s}- {TASK_DESCRIPTIONS[k]}" for k in TASK_PROMPTS)
-    parser = argparse.ArgumentParser(
-        description="Document OCR using HunyuanOCR-1.5 (lightweight ~1B end-to-end OCR VLM)",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog=f"""
-Task Types (official HunyuanOCR-1.5 prompts, all Chinese-language):
-{task_help}
-
-Examples:
-  # Basic document OCR (default)
-  uv run hunyuan-ocr-1.5.py my-docs analyzed-docs
-
-  # Extract formulas as LaTeX
-  uv run hunyuan-ocr-1.5.py math-papers formulas --task-type formula
-
-  # Extract tables as HTML
-  uv run hunyuan-ocr-1.5.py reports tables --task-type table
-
-  # Text spotting as JSON (box normalized 0-1000 + text)
-  uv run hunyuan-ocr-1.5.py images spotted --task-type spotting_json
-
-  # Translate documents to Chinese
-  uv run hunyuan-ocr-1.5.py en-docs translated --task-type doc_trans_en2zh
-
-  # Random sampling for testing
-  uv run hunyuan-ocr-1.5.py large-dataset test --max-samples 50 --shuffle
-        """,
-    )
-
-    parser.add_argument("input_dataset", help="Input dataset ID from Hugging Face Hub")
-    parser.add_argument("output_dataset", help="Output dataset ID for Hugging Face Hub")
-    parser.add_argument(
-        "--image-column",
-        default="image",
-        help="Column containing images (default: image)",
-    )
-    parser.add_argument(
-        "--batch-size",
-        type=int,
-        default=16,
-        help="Batch size for processing (default: 16; lower it if you hit engine errors)",
-    )
-    parser.add_argument(
-        "--model",
-        default="tencent/HunyuanOCR",
-        help="Model to use (default: tencent/HunyuanOCR — repo root is 1.5)",
-    )
-    parser.add_argument(
-        "--revision",
-        default=None,
-        help="Model repo revision (default: main). Tencent has replaced this repo's "
-        "root in-place before (1.0 → 1.5); pin a commit hash for reproducible runs.",
-    )
-    parser.add_argument(
-        "--max-model-len",
-        type=int,
-        default=32768,
-        help="Maximum model context length (default: 32768; max 131072). A single "
-        "image is capped at ~16384 tokens by the vision processor, so 32768 fits "
-        "image + 8192 output; raise for very long outputs.",
-    )
-    parser.add_argument(
-        "--max-tokens",
-        type=int,
-        default=8192,
-        help="Maximum tokens to generate (default: 8192; must be ≤ --max-model-len). "
-        "Dense pages may need more — raise toward 32768.",
-    )
-    parser.add_argument(
-        "--repetition-penalty",
-        type=float,
-        default=DEFAULT_REPETITION_PENALTY,
-        help=f"Repetition penalty (default: {DEFAULT_REPETITION_PENALTY}, the model card's locked value)",
-    )
-    parser.add_argument(
-        "--gpu-memory-utilization",
-        type=float,
-        default=0.8,
-        help="GPU memory utilization (default: 0.8)",
-    )
-    parser.add_argument("--hf-token", help="Hugging Face API token")
-    parser.add_argument(
-        "--split", default="train", help="Dataset split to use (default: train)"
-    )
-    parser.add_argument(
-        "--max-samples",
-        type=int,
-        help="Maximum number of samples to process (for testing)",
-    )
-    parser.add_argument(
-        "--private", action="store_true", help="Make output dataset private"
-    )
-    parser.add_argument(
-        "--shuffle", action="store_true", help="Shuffle dataset before processing"
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=42,
-        help="Random seed for shuffling (default: 42)",
-    )
-    parser.add_argument(
-        "--task-type",
-        choices=list(TASK_PROMPTS.keys()),
-        default=DEFAULT_TASK,
-        metavar="TASK",
-        help=f"Official task type (default: {DEFAULT_TASK}). See the epilog for all types.",
-    )
-    parser.add_argument(
-        "--custom-prompt",
-        help="Custom prompt text (overrides --task-type; may degrade quality — upstream "
-        "locks prompts per task)",
-    )
-    parser.add_argument(
-        "--output-column",
-        default="markdown",
-        help="Column name for output text (default: markdown)",
-    )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
-    parser.add_argument(
-        "--no-clean-output",
-        action="store_true",
-        help="Disable cleaning of repeated substrings in output",
-    )
-    parser.add_argument(
-        "--config",
-        help="Dataset config name for multi-model benchmarks",
-    )
-    parser.add_argument(
-        "--create-pr",
-        action="store_true",
-        help="Push results as a pull request instead of direct commit",
-    )
-    parser.add_argument(
-        "--verbose",
-        action="store_true",
-        help="Log resolved package versions at the end of the run",
-    )
-
-    args = parser.parse_args()
-
-    main(
-        input_dataset=args.input_dataset,
-        output_dataset=args.output_dataset,
-        image_column=args.image_column,
-        batch_size=args.batch_size,
-        model=args.model,
-        revision=args.revision,
-        max_model_len=args.max_model_len,
-        max_tokens=args.max_tokens,
-        repetition_penalty=args.repetition_penalty,
-        gpu_memory_utilization=args.gpu_memory_utilization,
-        hf_token=args.hf_token,
-        split=args.split,
-        max_samples=args.max_samples,
-        private=args.private,
-        shuffle=args.shuffle,
-        seed=args.seed,
-        task_type=args.task_type,
-        custom_prompt=args.custom_prompt,
-        output_column=args.output_column,
-        overwrite=args.overwrite,
-        clean_output=not args.no_clean_output,
-        config=args.config,
-        create_pr=args.create_pr,
-        verbose=args.verbose,
-    )

hunyuan-ocr.py

@@ -5,9 +5,6 @@
 #     "huggingface-hub",
 #     "pillow",
 #     "vllm>=0.15.1",
-#     "transformers<5.13",  # vLLM ≤0.24.0's HunyuanVL processor breaks on transformers 5.13
-#                           # (string-key AutoImageProcessor.register; fixed in vllm#47872).
-#                           # Drop this cap once that fix ships in a stable vLLM release.
 #     "tqdm",
 #     "toolz",
 #     "torch",
@@ -15,18 +12,11 @@
 # ///
 
 """
-Convert document images to markdown using HunyuanOCR (v1.0) with vLLM.
+Convert document images to markdown using HunyuanOCR with vLLM.
 
 HunyuanOCR is a lightweight 1B parameter VLM from Tencent designed for complex
 multilingual document parsing. This script uses vLLM for processing.
 
-⚠️ This recipe targets HunyuanOCR **1.0**. On 2026-07-06 Tencent pushed
-HunyuanOCR-1.5 into the *same* repo (new weights at root, 1.0 archived under
-`v1.0/`, no git tag), so loading `tencent/HunyuanOCR@main` now yields 1.5 —
-different context length, prompts, and transformers/vLLM requirements. We pin
-the last 1.0 commit by revision to keep this script's validated behavior.
-For 1.5, use the sibling `hunyuan-ocr-1.5.py`.
-
 Features:
 - 📝 Full document parsing to markdown
 - 📊 Table extraction (HTML format)
@@ -36,10 +26,8 @@ Features:
 - 🌐 Photo translation
 - 🎯 Compact model (1B parameters)
 
-Model: tencent/HunyuanOCR (pinned to the last 1.0 revision)
-
-License: Tencent Hunyuan Community License (territory excludes EU/UK/South Korea)
-https://huggingface.co/tencent/HunyuanOCR/blob/main/LICENSE
+Model: tencent/HunyuanOCR
+vLLM: Requires nightly build (trust_remote_code=True)
 
 Note: Due to vLLM V1 engine batching issues with HunyuanOCR, batch_size defaults to 1.
 """
@@ -61,21 +49,11 @@ from huggingface_hub import DatasetCard, login
 from PIL import Image
 from toolz import partition_all
 from tqdm.auto import tqdm
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
 from vllm import LLM, SamplingParams
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Last commit where the repo root held the 1.0 weights (2026-01-13). Not a temporary
-# workaround pin: upstream replaced root with 1.5 in-place (2026-07-06) and left no
-# 1.0 tag/branch, so this revision IS the 1.0 identity. Never loosen to "main" here —
-# 1.5 lives in hunyuan-ocr-1.5.py. Override with --revision only to reproduce old runs.
-DEFAULT_REVISION = "f6af82ee007fe6091b29fb3bb287b491ead41c82"
-
 
 # ────────────────────────────────────────────────────────────────
 # HunyuanOCR Prompt Templates (from official README)
@@ -168,28 +146,6 @@ def check_cuda_availability():
         logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
 
 
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
 def get_prompt(
     prompt_mode: str,
     use_chinese: bool = False,
@@ -266,7 +222,6 @@ def make_ocr_message(
 def create_dataset_card(
     source_dataset: str,
     model: str,
-    revision: str,
     num_samples: int,
     processing_time: str,
     batch_size: int,
@@ -296,13 +251,10 @@ tags:
 
 This dataset contains OCR results from images in [{source_dataset}](https://huggingface.co/datasets/{source_dataset}) using HunyuanOCR, a lightweight 1B VLM from Tencent.
 
-Model license: [Tencent Hunyuan Community License](https://huggingface.co/tencent/HunyuanOCR/blob/main/LICENSE) (territory excludes EU/UK/South Korea).
-
 ## Processing Details
 
 - **Source Dataset**: [{source_dataset}](https://huggingface.co/datasets/{source_dataset})
 - **Model**: [{model}](https://huggingface.co/{model})
-- **Model Revision**: `{revision}` (HunyuanOCR 1.0 — the repo root holds 1.5 since 2026-07-06)
 - **Number of Samples**: {num_samples:,}
 - **Processing Time**: {processing_time}
 - **Processing Date**: {datetime.now().strftime("%Y-%m-%d %H:%M UTC")}
@@ -394,7 +346,6 @@ def main(
     image_column: str = "image",
     batch_size: int = 1,  # Default to 1 due to vLLM V1 batching issues with HunyuanOCR
     model: str = "tencent/HunyuanOCR",
-    revision: str = None,
     max_model_len: int = 16384,
     max_tokens: int = 16384,
     gpu_memory_utilization: float = 0.8,
@@ -411,7 +362,6 @@ def main(
     use_chinese: bool = False,
     custom_prompt: str = None,
     output_column: str = "markdown",
-    overwrite: bool = False,
     clean_output: bool = True,
     config: str = None,
     create_pr: bool = False,
@@ -455,9 +405,6 @@ def main(
             f"Column '{image_column}' not found. Available: {dataset.column_names}"
         )
 
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
     # Shuffle if requested
     if shuffle:
         logger.info(f"Shuffling dataset with seed {seed}")
@@ -468,21 +415,14 @@ def main(
         dataset = dataset.select(range(min(max_samples, len(dataset))))
         logger.info(f"Limited to {len(dataset)} samples")
 
-    # Pin the 1.0 revision only for the default repo — a user-supplied --model
-    # must not inherit a foreign commit hash.
-    if revision is None and model == "tencent/HunyuanOCR":
-        revision = DEFAULT_REVISION
-        logger.info(f"Pinning tencent/HunyuanOCR to 1.0 revision {revision[:8]}")
-
     # Initialize vLLM model
-    logger.info(f"Initializing vLLM: {model} (revision: {revision or 'main'})")
+    logger.info(f"Initializing vLLM with model: {model}")
     logger.info("This may take a few minutes on first run...")
 
     # Note: HunyuanOCR has batching issues with vLLM V1 engine when batch_size > 1
     # Using limit_mm_per_prompt for stability
     llm = LLM(
         model=model,
-        revision=revision,
         trust_remote_code=True,
         max_model_len=max_model_len,
         gpu_memory_utilization=gpu_memory_utilization,
@@ -541,7 +481,6 @@ def main(
     inference_entry = {
         "model_id": model,
         "model_name": "HunyuanOCR",
-        "model_revision": revision or "main",
         "column_name": output_column,
         "timestamp": datetime.now().isoformat(),
         "prompt_mode": prompt_mode if not custom_prompt else "custom",
@@ -609,7 +548,6 @@ def main(
         card_content = create_dataset_card(
             source_dataset=input_dataset,
             model=model,
-            revision=revision or "main",
             num_samples=len(dataset),
             processing_time=processing_time_str,
             batch_size=batch_size,
@@ -763,13 +701,6 @@ Examples:
         default="tencent/HunyuanOCR",
         help="Model to use (default: tencent/HunyuanOCR)",
     )
-    parser.add_argument(
-        "--revision",
-        default=None,
-        help="Model repo revision. Defaults to the last 1.0 commit when --model is "
-        "tencent/HunyuanOCR (the repo root now holds 1.5 — see hunyuan-ocr-1.5.py); "
-        "defaults to main for any other model.",
-    )
     parser.add_argument(
         "--max-model-len",
         type=int,
@@ -851,12 +782,6 @@ Examples:
         default="markdown",
         help="Column name for output text (default: markdown)",
     )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
     parser.add_argument(
         "--no-clean-output",
         action="store_true",
@@ -890,7 +815,6 @@ Examples:
         image_column=args.image_column,
         batch_size=args.batch_size,
         model=args.model,
-        revision=args.revision,
         max_model_len=args.max_model_len,
         max_tokens=args.max_tokens,
         gpu_memory_utilization=args.gpu_memory_utilization,
@@ -907,7 +831,6 @@ Examples:
         use_chinese=args.use_chinese_prompts,
         custom_prompt=args.custom_prompt,
         output_column=args.output_column,
-        overwrite=args.overwrite,
         clean_output=not args.no_clean_output,
         config=args.config,
         create_pr=args.create_pr,

lfm2-extract.py

@@ -1,327 +0,0 @@
-# /// script
-# requires-python = ">=3.11"
-# dependencies = [
-#     "datasets>=4.0.0",
-#     "huggingface-hub",
-#     "vllm",
-#     "transformers",
-#     "tqdm",
-#     "toolz",
-#     "torch",
-# ]
-# ///
-"""
-Extract structured data (JSON / XML / YAML) from text using LiquidAI's LFM2-1.2B-Extract.
-
-LFM2-1.2B-Extract is a compact 1.2B text-only model purpose-built for turning unstructured
-documents into structured data: give it a schema, it returns JSON, XML, or YAML. It reports
-beating Gemma 3 27B (22x larger) on syntax validity / format accuracy / faithfulness, and
-is multilingual (en, ar, zh, fr, de, ja, ko, pt, es).
-
-This is the *text* counterpart to `lfm2-vl-extract.py` (which extracts from images). Pair them:
-OCR a page to markdown with one of the OCR recipes, then extract fields from that text here.
-
-Pass `--schema` as inline text/JSON, a URL, or a file path describing the structure to extract:
-
-    --schema '{"invoice_number": "string", "total": "number", "line_items": "array"}'
-
-Model:  https://huggingface.co/LiquidAI/LFM2-1.2B-Extract
-Docs:   https://docs.liquid.ai/deployment/gpu-inference/vllm
-
-HF Jobs note: run on the vLLM image so the CUDA toolkit + prebuilt FlashInfer kernels are
-present and startup is fast (it reuses the image's CUDA-matched vLLM build):
-
-    hf jobs uv run --flavor l4x1 --secrets HF_TOKEN \
-        --image vllm/vllm-openai --python /usr/bin/python3 \
-        -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages \
-        https://huggingface.co/datasets/uv-scripts/ocr/raw/main/lfm2-extract.py \
-        INPUT OUTPUT --text-column text --schema '{"field": "description"}'
-
-It also runs on the default uv image, just with a slower first-time vLLM build. Deps are left
-unpinned so uv resolves a recent vLLM; FlashInfer sampling is disabled (see below) so the engine
-never JIT-compiles a kernel that needs nvcc — absent from the default image.
-"""
-
-import argparse
-import json
-import logging
-import os
-import sys
-from datetime import datetime, timezone
-from typing import List, Optional
-from urllib.request import urlopen
-
-# Disable vLLM's FlashInfer sampler before the engine starts: it JIT-compiles at warmup and
-# needs nvcc (absent from the default uv image). Harmless for greedy decoding.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
-
-import torch
-from datasets import load_dataset
-from huggingface_hub import DatasetCard, login
-from toolz import partition_all
-from tqdm import tqdm
-from vllm import LLM, SamplingParams
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-DEFAULT_MODEL = "LiquidAI/LFM2-1.2B-Extract"
-FORMATS = {"json": "JSON", "xml": "XML", "yaml": "YAML"}
-
-
-def check_cuda_availability() -> None:
-    if not torch.cuda.is_available():
-        logger.error("CUDA is not available. This script requires a GPU.")
-        logger.error("Run on Hugging Face Jobs with: hf jobs uv run --flavor l4x1 ...")
-        sys.exit(1)
-    logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name()}")
-
-
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
-def load_text_arg(value: str) -> str:
-    """Resolve --schema (inline text/JSON, URL, or file path) into a string."""
-    text = value.strip()
-    if text.startswith("http://") or text.startswith("https://"):
-        logger.info(f"Loading schema from URL: {text}")
-        return urlopen(text).read().decode("utf-8").strip()
-    if os.path.exists(text):
-        logger.info(f"Loading schema from file: {text}")
-        with open(text) as f:
-            return f.read().strip()
-    return text
-
-
-def build_system_prompt(schema_text: str, fmt: str) -> str:
-    return f"Return data as a {FORMATS[fmt]} object with the following schema:\n\n{schema_text}"
-
-
-def parse_output(text: str, fmt: str) -> tuple[str, bool]:
-    """Strip code fences; for JSON, validate. Returns (cleaned_text, is_valid)."""
-    stripped = text.strip()
-    if stripped.startswith("```"):
-        stripped = stripped.split("\n", 1)[-1]
-        if stripped.endswith("```"):
-            stripped = stripped.rsplit("```", 1)[0]
-    stripped = stripped.strip()
-    if fmt == "json":
-        try:
-            return json.dumps(json.loads(stripped), ensure_ascii=False), True
-        except (json.JSONDecodeError, ValueError):
-            return stripped, False
-    return stripped, True  # xml/yaml: store as-is (no strict validator)
-
-
-def main(
-    input_dataset: str,
-    output_dataset: str,
-    schema: str,
-    text_column: str = "text",
-    output_column: str = "extraction",
-    overwrite: bool = False,
-    output_format: str = "json",
-    split: str = "train",
-    max_samples: Optional[int] = None,
-    shuffle: bool = False,
-    seed: int = 42,
-    batch_size: int = 32,
-    model: str = DEFAULT_MODEL,
-    max_model_len: int = 8192,
-    max_tokens: int = 4096,
-    private: bool = False,
-    hf_token: Optional[str] = None,
-) -> None:
-    check_cuda_availability()
-    if output_format not in FORMATS:
-        logger.error(f"--format must be one of {list(FORMATS)}; got {output_format}")
-        sys.exit(1)
-
-    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")
-    if HF_TOKEN:
-        login(token=HF_TOKEN)
-
-    schema_text = load_text_arg(schema)
-    system_prompt = build_system_prompt(schema_text, output_format)
-
-    logger.info(f"Loading dataset: {input_dataset} (split={split})")
-    dataset = load_dataset(input_dataset, split=split)
-
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
-    if shuffle:
-        dataset = dataset.shuffle(seed=seed)
-    if max_samples:
-        dataset = dataset.select(range(min(max_samples, len(dataset))))
-    logger.info(f"Processing {len(dataset)} examples; format={output_format}")
-
-    if text_column not in dataset.column_names:
-        logger.error(f"Text column '{text_column}' not found. Columns: {dataset.column_names}")
-        sys.exit(1)
-
-    logger.info(f"Loading model: {model}")
-    llm = LLM(model=model, max_model_len=max_model_len, enforce_eager=True)
-    sampling_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
-
-    all_outputs: List[str] = []
-    n_valid = 0
-    texts = dataset[text_column]
-    for batch in tqdm(list(partition_all(batch_size, texts)), desc="Extracting"):
-        batch_messages = [
-            [
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": str(doc)},
-            ]
-            for doc in batch
-        ]
-        outputs = llm.chat(batch_messages, sampling_params)
-        for out in outputs:
-            cleaned, ok = parse_output(out.outputs[0].text, output_format)
-            n_valid += int(ok)
-            all_outputs.append(cleaned)
-
-    logger.info(f"Valid {output_format.upper()}: {n_valid}/{len(all_outputs)}")
-    dataset = dataset.add_column(output_column, all_outputs)
-
-    inference_entry = {
-        "model": model,
-        "column_name": output_column,
-        "task": "structured extraction",
-        "format": output_format,
-        "timestamp": datetime.now(timezone.utc).isoformat(),
-        "script": "lfm2-extract.py",
-    }
-    if "inference_info" in dataset.column_names:
-        def update_info(example):
-            try:
-                existing = json.loads(example["inference_info"]) if example["inference_info"] else []
-            except (json.JSONDecodeError, TypeError):
-                existing = []
-            existing.append(inference_entry)
-            return {"inference_info": json.dumps(existing)}
-        dataset = dataset.map(update_info)
-    else:
-        dataset = dataset.add_column(
-            "inference_info", [json.dumps([inference_entry])] * len(dataset)
-        )
-
-    logger.info(f"Pushing to {output_dataset}")
-    dataset.push_to_hub(output_dataset, private=private, token=HF_TOKEN)
-
-    card_text = f"""---
-tags:
-- uv-script
-- extraction
-- lfm2
-- {output_format}
----
-
-# Structured extraction with LFM2-1.2B-Extract
-
-`{output_format.upper()}` extracted from the `{text_column}` column of
-[{input_dataset}](https://huggingface.co/datasets/{input_dataset})
-using [{model}](https://huggingface.co/{model}).
-
-- **Source**: `{input_dataset}` (split `{split}`, column `{text_column}`)
-- **Model**: `{model}`
-- **Format**: `{output_format}`
-- **Output column**: `{output_column}`
-- **Valid {output_format.upper()}**: {n_valid}/{len(all_outputs)}
-- **Date**: {datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC")}
-
-Generated with the [uv-scripts/ocr](https://huggingface.co/datasets/uv-scripts/ocr) `lfm2-extract.py` script.
-"""
-    try:
-        DatasetCard(card_text).push_to_hub(output_dataset, token=HF_TOKEN)
-    except Exception as e:
-        logger.warning(f"Could not push dataset card: {e}")
-
-    logger.info("Done! Extraction complete.")
-    logger.info(f"Dataset: https://huggingface.co/datasets/{output_dataset}")
-
-
-if __name__ == "__main__":
-    if len(sys.argv) == 1:
-        print("LFM2-1.2B-Extract — structured extraction (JSON/XML/YAML) from text")
-        print("\nUsage:")
-        print("  uv run lfm2-extract.py INPUT OUTPUT --schema SCHEMA [--text-column text] [--format json]")
-        print("\nExample:")
-        print('  uv run lfm2-extract.py my-docs my-fields \\')
-        print('    --text-column markdown \\')
-        print('    --schema \'{"title": "the title", "date": "any date", "summary": "one sentence"}\'')
-        print("\n  --schema accepts inline text/JSON, a URL, or a file path.")
-        print("\nFor full help: uv run lfm2-extract.py --help")
-        sys.exit(0)
-
-    parser = argparse.ArgumentParser(
-        description="Structured extraction (JSON/XML/YAML) from text using LFM2-1.2B-Extract",
-    )
-    parser.add_argument("input_dataset", help="Input dataset ID (with a text column)")
-    parser.add_argument("output_dataset", help="Output dataset ID")
-    parser.add_argument(
-        "--schema", required=True,
-        help="Structure to extract: inline text/JSON, a URL, or a file path",
-    )
-    parser.add_argument("--text-column", default="text", help="Text column (default: text)")
-    parser.add_argument("--output-column", default="extraction", help="Output column (default: extraction)")
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
-    parser.add_argument(
-        "--format", dest="output_format", default="json", choices=list(FORMATS),
-        help="Output format (default: json)",
-    )
-    parser.add_argument("--split", default="train", help="Dataset split (default: train)")
-    parser.add_argument("--max-samples", type=int, help="Limit number of samples")
-    parser.add_argument("--shuffle", action="store_true", help="Shuffle before sampling")
-    parser.add_argument("--seed", type=int, default=42, help="Shuffle seed (default: 42)")
-    parser.add_argument("--batch-size", type=int, default=32, help="Batch size (default: 32)")
-    parser.add_argument("--model", default=DEFAULT_MODEL, help=f"Model (default: {DEFAULT_MODEL})")
-    parser.add_argument("--max-model-len", type=int, default=8192, help="Max context length (default: 8192)")
-    parser.add_argument("--max-tokens", type=int, default=4096, help="Max output tokens (default: 4096)")
-    parser.add_argument("--private", action="store_true", help="Make output dataset private")
-    parser.add_argument("--hf-token", help="HF token (or set HF_TOKEN)")
-    args = parser.parse_args()
-
-    main(
-        input_dataset=args.input_dataset,
-        output_dataset=args.output_dataset,
-        schema=args.schema,
-        text_column=args.text_column,
-        output_column=args.output_column,
-        overwrite=args.overwrite,
-        output_format=args.output_format,
-        split=args.split,
-        max_samples=args.max_samples,
-        shuffle=args.shuffle,
-        seed=args.seed,
-        batch_size=args.batch_size,
-        model=args.model,
-        max_model_len=args.max_model_len,
-        max_tokens=args.max_tokens,
-        private=args.private,
-        hf_token=args.hf_token,
-    )

lfm2-vl-extract.py

@@ -1,358 +0,0 @@
-# /// script
-# requires-python = ">=3.11"
-# dependencies = [
-#     "datasets>=4.0.0",
-#     "huggingface-hub",
-#     "pillow",
-#     "vllm",
-#     "transformers",
-#     "tqdm",
-#     "toolz",
-#     "torch",
-# ]
-# ///
-"""
-Extract structured JSON from images using LiquidAI's LFM2.5-VL-1.6B-Extract with vLLM.
-
-LFM2.5-VL-1.6B-Extract (1.6B = LFM2 1.2B LM + SigLIP2 0.4B vision) is a compact
-vision-language model purpose-built for *schema-guided* extraction: you give it a
-list of fields, it returns a flat JSON object with those fields filled from the image.
-It reports 99.6 JSON-validity / F1 on its benchmark, beating similarly-sized VLMs.
-
-Unlike the markdown-OCR scripts here, this one needs a SCHEMA (a field list). Pass
-`--schema` as inline JSON, a URL, or a file path, mapping field names to short
-descriptions:
-
-    --schema '{"invoice_number": "the invoice number", "total": "the total amount"}'
-
-Model:  https://huggingface.co/LiquidAI/LFM2.5-VL-1.6B-Extract
-Docs:   https://docs.liquid.ai/deployment/gpu-inference/vllm
-
-HF Jobs note: run on the vLLM image so the CUDA toolkit + prebuilt FlashInfer kernels
-are present and startup is fast (it reuses the image's CUDA-matched vLLM build):
-
-    hf jobs uv run --flavor l4x1 --secrets HF_TOKEN \
-        --image vllm/vllm-openai --python /usr/bin/python3 \
-        -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages \
-        https://huggingface.co/datasets/uv-scripts/ocr/raw/main/lfm2-vl-extract.py \
-        INPUT OUTPUT --schema '{"title": "the document title", "date": "any date shown"}'
-
-It also runs on the default uv image, just with a slower first-time vLLM build. Deps are
-left unpinned so uv resolves a vLLM that supports the LFM2-VL (transformers 5) architecture,
-and FlashInfer sampling is disabled (VLLM_USE_FLASHINFER_SAMPLER=0, see below) so the engine
-never JIT-compiles a kernel that needs nvcc — absent from the default image.
-"""
-
-import argparse
-import base64
-import io
-import json
-import logging
-import os
-import sys
-from datetime import datetime, timezone
-from typing import Any, Dict, List, Optional, Union
-from urllib.request import urlopen
-
-# Disable vLLM's FlashInfer top-k/top-p sampler before the engine starts: it JIT-compiles
-# at warmup and needs nvcc (absent from the default uv image). Harmless for greedy decoding.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
-
-import torch
-from datasets import load_dataset
-from huggingface_hub import DatasetCard, login
-from PIL import Image
-from toolz import partition_all
-from tqdm import tqdm
-from vllm import LLM, SamplingParams
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-DEFAULT_MODEL = "LiquidAI/LFM2.5-VL-1.6B-Extract"
-
-
-def check_cuda_availability() -> None:
-    if not torch.cuda.is_available():
-        logger.error("CUDA is not available. This script requires a GPU.")
-        logger.error("Run on Hugging Face Jobs with: hf jobs uv run --flavor l4x1 ...")
-        sys.exit(1)
-    logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name()}")
-
-
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
-def load_schema_arg(value: str) -> Dict[str, str]:
-    """Resolve --schema (inline JSON, URL, or file path) into a {field: description} dict."""
-    text = value.strip()
-    if text.startswith("http://") or text.startswith("https://"):
-        logger.info(f"Loading schema from URL: {text}")
-        text = urlopen(text).read().decode("utf-8")
-    elif not text.startswith("{") and not text.startswith("["):
-        if os.path.exists(text):
-            logger.info(f"Loading schema from file: {text}")
-            with open(text) as f:
-                text = f.read()
-    parsed = json.loads(text)
-    # Accept {"field": "description"} or ["field1", "field2"]
-    if isinstance(parsed, list):
-        return {str(field): "" for field in parsed}
-    if isinstance(parsed, dict):
-        return {str(k): str(v) for k, v in parsed.items()}
-    raise ValueError("--schema must be a JSON object {field: description} or a JSON list of field names.")
-
-
-def build_system_prompt(schema: Dict[str, str]) -> str:
-    """LFM2.5-VL-Extract prompt: a field list in the system message → flat JSON out."""
-    lines = []
-    for field, desc in schema.items():
-        lines.append(f"{field}: {desc}" if desc else field)
-    fields_block = "\n".join(lines)
-    return (
-        f"Extract the following from the image:\n\n{fields_block}\n\n"
-        "Respond with only a JSON object."
-    )
-
-
-def image_to_data_uri(image: Union[Image.Image, Dict[str, Any], str]) -> str:
-    if isinstance(image, dict) and "bytes" in image:
-        image = Image.open(io.BytesIO(image["bytes"]))
-    elif isinstance(image, str):
-        image = Image.open(image)
-    if image.mode != "RGB":
-        image = image.convert("RGB")
-    buf = io.BytesIO()
-    image.save(buf, format="PNG")
-    return f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}"
-
-
-def make_message(image: Any, system_prompt: str) -> List[Dict]:
-    data_uri = image_to_data_uri(image)
-    return [
-        {"role": "system", "content": system_prompt},
-        {"role": "user", "content": [{"type": "image_url", "image_url": {"url": data_uri}}]},
-    ]
-
-
-def parse_json_output(text: str) -> tuple[Optional[Any], bool]:
-    """Return (parsed, ok). Strips ```json fences if present."""
-    stripped = text.strip()
-    if stripped.startswith("```"):
-        stripped = stripped.split("\n", 1)[-1]
-        if stripped.endswith("```"):
-            stripped = stripped.rsplit("```", 1)[0]
-    stripped = stripped.strip()
-    try:
-        return json.loads(stripped), True
-    except (json.JSONDecodeError, ValueError):
-        return None, False
-
-
-def main(
-    input_dataset: str,
-    output_dataset: str,
-    schema: str,
-    image_column: str = "image",
-    output_column: str = "extraction",
-    overwrite: bool = False,
-    split: str = "train",
-    max_samples: Optional[int] = None,
-    shuffle: bool = False,
-    seed: int = 42,
-    batch_size: int = 16,
-    model: str = DEFAULT_MODEL,
-    max_model_len: int = 4096,
-    max_tokens: int = 1024,
-    private: bool = False,
-    hf_token: Optional[str] = None,
-) -> None:
-    check_cuda_availability()
-
-    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")
-    if HF_TOKEN:
-        login(token=HF_TOKEN)
-
-    schema_dict = load_schema_arg(schema)
-    system_prompt = build_system_prompt(schema_dict)
-    logger.info(f"Extraction fields: {list(schema_dict.keys())}")
-
-    logger.info(f"Loading dataset: {input_dataset} (split={split})")
-    dataset = load_dataset(input_dataset, split=split)
-
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
-    if shuffle:
-        dataset = dataset.shuffle(seed=seed)
-    if max_samples:
-        dataset = dataset.select(range(min(max_samples, len(dataset))))
-    logger.info(f"Processing {len(dataset)} examples")
-
-    if image_column not in dataset.column_names:
-        logger.error(f"Image column '{image_column}' not found. Columns: {dataset.column_names}")
-        sys.exit(1)
-
-    logger.info(f"Loading model: {model}")
-    llm = LLM(
-        model=model,
-        max_model_len=max_model_len,
-        limit_mm_per_prompt={"image": 1},
-        enforce_eager=True,
-    )
-    sampling_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
-
-    all_outputs: List[str] = []
-    n_valid = 0
-    images = dataset[image_column]
-    for batch in tqdm(list(partition_all(batch_size, images)), desc="Extracting"):
-        batch_messages = [make_message(img, system_prompt) for img in batch]
-        outputs = llm.chat(batch_messages, sampling_params)
-        for out in outputs:
-            text = out.outputs[0].text.strip()
-            parsed, ok = parse_json_output(text)
-            if ok:
-                n_valid += 1
-                all_outputs.append(json.dumps(parsed, ensure_ascii=False))
-            else:
-                all_outputs.append(text)  # keep raw on parse failure
-
-    logger.info(f"Valid JSON: {n_valid}/{len(all_outputs)}")
-
-    dataset = dataset.add_column(output_column, all_outputs)
-
-    inference_entry = {
-        "model": model,
-        "column_name": output_column,
-        "task": "schema-guided extraction",
-        "fields": list(schema_dict.keys()),
-        "timestamp": datetime.now(timezone.utc).isoformat(),
-        "script": "lfm2-vl-extract.py",
-    }
-    if "inference_info" in dataset.column_names:
-        def update_info(example):
-            try:
-                existing = json.loads(example["inference_info"]) if example["inference_info"] else []
-            except (json.JSONDecodeError, TypeError):
-                existing = []
-            existing.append(inference_entry)
-            return {"inference_info": json.dumps(existing)}
-        dataset = dataset.map(update_info)
-    else:
-        dataset = dataset.add_column(
-            "inference_info", [json.dumps([inference_entry])] * len(dataset)
-        )
-
-    logger.info(f"Pushing to {output_dataset}")
-    dataset.push_to_hub(output_dataset, private=private, token=HF_TOKEN)
-
-    card_text = f"""---
-tags:
-- uv-script
-- extraction
-- lfm2-vl
-- json
----
-
-# Structured extraction with LFM2.5-VL-1.6B-Extract
-
-JSON fields extracted from images in [{input_dataset}](https://huggingface.co/datasets/{input_dataset})
-using [{model}](https://huggingface.co/{model}).
-
-- **Source**: `{input_dataset}` (split `{split}`)
-- **Model**: `{model}`
-- **Fields**: {", ".join(f"`{k}`" for k in schema_dict.keys())}
-- **Output column**: `{output_column}` (JSON string per row)
-- **Valid JSON**: {n_valid}/{len(all_outputs)}
-- **Date**: {datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC")}
-
-Generated with the [uv-scripts/ocr](https://huggingface.co/datasets/uv-scripts/ocr) `lfm2-vl-extract.py` script.
-"""
-    try:
-        card = DatasetCard(card_text)
-        card.push_to_hub(output_dataset, token=HF_TOKEN)
-    except Exception as e:
-        logger.warning(f"Could not push dataset card: {e}")
-
-    logger.info("Done! Extraction complete.")
-    logger.info(f"Dataset: https://huggingface.co/datasets/{output_dataset}")
-
-
-if __name__ == "__main__":
-    if len(sys.argv) == 1:
-        print("LFM2.5-VL-1.6B-Extract — schema-guided JSON extraction from images")
-        print("\nUsage:")
-        print("  uv run lfm2-vl-extract.py INPUT OUTPUT --schema SCHEMA [options]")
-        print("\nExample:")
-        print('  uv run lfm2-vl-extract.py my-images my-extractions \\')
-        print('    --schema \'{"title": "the document title", "date": "any date shown"}\'')
-        print("\n  --schema accepts inline JSON, a URL, or a file path.")
-        print("\nFor full help: uv run lfm2-vl-extract.py --help")
-        sys.exit(0)
-
-    parser = argparse.ArgumentParser(
-        description="Schema-guided JSON extraction from images using LFM2.5-VL-1.6B-Extract",
-    )
-    parser.add_argument("input_dataset", help="Input dataset ID (with images)")
-    parser.add_argument("output_dataset", help="Output dataset ID")
-    parser.add_argument(
-        "--schema", required=True,
-        help="Fields to extract: inline JSON {field: description}, a URL, or a file path",
-    )
-    parser.add_argument("--image-column", default="image", help="Image column (default: image)")
-    parser.add_argument("--output-column", default="extraction", help="Output column (default: extraction)")
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
-    parser.add_argument("--split", default="train", help="Dataset split (default: train)")
-    parser.add_argument("--max-samples", type=int, help="Limit number of samples")
-    parser.add_argument("--shuffle", action="store_true", help="Shuffle before sampling")
-    parser.add_argument("--seed", type=int, default=42, help="Shuffle seed (default: 42)")
-    parser.add_argument("--batch-size", type=int, default=16, help="Batch size (default: 16)")
-    parser.add_argument("--model", default=DEFAULT_MODEL, help=f"Model (default: {DEFAULT_MODEL})")
-    parser.add_argument("--max-model-len", type=int, default=4096, help="Max context length (default: 4096)")
-    parser.add_argument("--max-tokens", type=int, default=1024, help="Max output tokens (default: 1024)")
-    parser.add_argument("--private", action="store_true", help="Make output dataset private")
-    parser.add_argument("--hf-token", help="HF token (or set HF_TOKEN)")
-    args = parser.parse_args()
-
-    main(
-        input_dataset=args.input_dataset,
-        output_dataset=args.output_dataset,
-        schema=args.schema,
-        image_column=args.image_column,
-        output_column=args.output_column,
-        overwrite=args.overwrite,
-        split=args.split,
-        max_samples=args.max_samples,
-        shuffle=args.shuffle,
-        seed=args.seed,
-        batch_size=args.batch_size,
-        model=args.model,
-        max_model_len=args.max_model_len,
-        max_tokens=args.max_tokens,
-        private=args.private,
-        hf_token=args.hf_token,
-    )

lift-extract.py

@@ -1,846 +0,0 @@
-# /// script
-# requires-python = ">=3.12"
-# dependencies = [
-#     "lift-pdf[hf]",
-#     "datasets>=3.1.0",
-#     "huggingface-hub",
-#     "pillow",
-#     "toolz",
-#     "tqdm",
-# ]
-# ///
-"""
-Extract structured JSON from document images OR multi-page PDFs using Datalab's
-`lift` model (`datalab-to/lift`, 9B, Qwen3.5-based).
-
-Unlike the markdown-OCR scripts here, lift does *schema-constrained* extraction:
-you give it a JSON Schema, it returns a JSON object matching that schema. It
-natively handles multi-page documents — a whole PDF is collapsed into a single
-extraction.
-
-Two in-process backends, selected with `--method` (no server, single command):
-
-  --method hf    (default)  Transformers via the `lift-pdf` package. Runs on the
-                            default uv image. Simplest path; best for small jobs.
-  --method vllm             vLLM's offline `LLM()` engine (`llm.chat`) with
-                            structured-output decoding — the fast batched path the
-                            other vLLM OCR scripts here use. Needs the
-                            `vllm/vllm-openai` image (which ships vLLM). Reproduces
-                            lift's own prompt + guided-JSON recipe against the
-                            offline engine. Wins on large jobs via continuous batching.
-
-Benchmark the two by pushing each to one repo with `--config hf` / `--config vllm`.
-
-Input is one document per row:
-  --image-column COL   (default `image`)  one image per row  -> one extraction
-  --pdf-column COL                        PDF bytes per row -> one extraction
-                                          (multi-page; respects --page-range)
-
-Pass `--schema` as inline JSON, a URL, or a file path (standard JSON Schema):
-
-    --schema '{"type":"object","properties":{"invoice_number":{"type":"string"},
-               "total":{"type":"number"}},"required":["invoice_number"]}'
-
-LICENSE NOTE: lift's *code* is Apache-2.0 but the *weights* are a modified
-OpenRAIL-M license — free for research, personal use, and startups under $5M
-funding/revenue, but restricted from competitive use against Datalab's API.
-Confirm you are within those terms before using it. https://huggingface.co/datalab-to/lift
-
-HF Jobs — HF backend (default image is fine; 9B needs a roomy GPU):
-
-    hf jobs uv run --flavor a100-large -s HF_TOKEN \\
-        https://huggingface.co/datasets/uv-scripts/ocr/raw/main/lift-extract.py \\
-        INPUT_DATASET OUTPUT_DATASET \\
-        --schema '{"type":"object","properties":{"title":{"type":"string"}}}' \\
-        --max-samples 5 --shuffle --seed 42
-
-HF Jobs — vLLM offline backend (use the vllm image so vLLM is present):
-
-    hf jobs uv run --flavor a100-large -s HF_TOKEN \\
-        --image vllm/vllm-openai --python /usr/bin/python3 \\
-        -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages \\
-        https://huggingface.co/datasets/uv-scripts/ocr/raw/main/lift-extract.py \\
-        INPUT_DATASET OUTPUT_DATASET --method vllm \\
-        --schema '{"type":"object","properties":{"title":{"type":"string"}}}' \\
-        --max-samples 5
-
-Model: datalab-to/lift  (package: lift-pdf, https://github.com/datalab-to/lift)
-"""
-
-import argparse
-import base64
-import io
-import json
-import logging
-import os
-import sys
-import tempfile
-import time
-from datetime import datetime, timezone
-from typing import Any, Dict, List, Optional, Tuple
-from urllib.request import urlopen
-
-from datasets import load_dataset
-from huggingface_hub import DatasetCard, login
-from PIL import Image
-from toolz import partition_all
-from tqdm import tqdm
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# The package default checkpoint drifts between releases (e.g. "datalab-to/lift-extract");
-# pin to the canonical card repo so the script is stable across lift-pdf versions.
-DEFAULT_MODEL = "datalab-to/lift"
-DEFAULT_MAX_TOKENS = 12384  # lift-pdf's own MAX_OUTPUT_TOKENS default
-
-# A processed document: (parsed JSON or None, error flag, raw model text).
-DocResult = Tuple[Optional[Any], bool, str]
-
-
-def check_cuda_availability() -> None:
-    """Exit early with a clear message if there's no GPU."""
-    import torch
-
-    if not torch.cuda.is_available():
-        logger.error("CUDA is not available. This script requires a GPU.")
-        logger.error(
-            "Run on Hugging Face Jobs with: hf jobs uv run --flavor a100-large ..."
-        )
-        sys.exit(1)
-    logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
-
-
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
-def load_schema_arg(value: str) -> Dict[str, Any]:
-    """Resolve --schema (inline JSON, a URL, or a file path) into a JSON Schema dict."""
-    text = value.strip()
-    if text.startswith(("http://", "https://")):
-        logger.info(f"Loading schema from URL: {text}")
-        text = urlopen(text).read().decode("utf-8")  # noqa: S310
-    elif not text.startswith("{"):
-        # Looks like a path (inline JSON would start with "{"); read it if it exists.
-        if os.path.isfile(text):
-            logger.info(f"Loading schema from file: {text}")
-            with open(text) as f:
-                text = f.read()
-    try:
-        parsed = json.loads(text)
-    except json.JSONDecodeError as e:
-        raise ValueError(
-            f"Could not parse --schema as JSON (tried URL/path/inline): {e}"
-        ) from e
-    if not isinstance(parsed, dict):
-        raise ValueError("--schema must be a JSON object (a JSON Schema).")
-    return parsed
-
-
-def cell_to_bytes(cell: Any) -> bytes:
-    """Normalize an HF dataset cell (image or document) to raw file bytes.
-
-    Handles decoded PIL images (Image feature), {"bytes"/"path"} dicts, raw bytes
-    (e.g. a binary PDF column), and string paths/URLs.
-    """
-    if isinstance(cell, Image.Image):
-        buf = io.BytesIO()
-        cell.convert("RGB").save(buf, format="PNG")
-        return buf.getvalue()
-    if isinstance(cell, dict):
-        if cell.get("bytes"):
-            return cell["bytes"]
-        if cell.get("path"):
-            with open(cell["path"], "rb") as f:
-                return f.read()
-        raise ValueError(
-            f"Unsupported image/document dict (no bytes/path): {list(cell)}"
-        )
-    if isinstance(cell, (bytes, bytearray)):
-        return bytes(cell)
-    if isinstance(cell, str):
-        if cell.startswith(("http://", "https://")):
-            return urlopen(cell).read()  # noqa: S310
-        with open(cell, "rb") as f:
-            return f.read()
-    raise ValueError(f"Unsupported cell type: {type(cell)}")
-
-
-def load_document_images(
-    load_file, cell: Any, page_range: Optional[str]
-) -> List[Image.Image]:
-    """Render one dataset cell into the page images lift expects.
-
-    Reuses lift's own `load_file`, which auto-detects PDF vs image by content
-    (pypdfium2 for PDFs, with the model's DPI/min-dim and page-range handling).
-    """
-    data = cell_to_bytes(cell)
-    # load_file detects type from content, so the temp file needs no extension.
-    with tempfile.NamedTemporaryFile(delete=False) as tmp:
-        tmp.write(data)
-        path = tmp.name
-    try:
-        config = {"page_range": page_range} if page_range else {}
-        return load_file(path, config)
-    finally:
-        os.unlink(path)
-
-
-def pil_to_data_uri(img: Image.Image) -> str:
-    """PNG data URI for an OpenAI-format image content block."""
-    if img.mode != "RGB":
-        img = img.convert("RGB")
-    buf = io.BytesIO()
-    img.save(buf, format="PNG")
-    return f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}"
-
-
-def parse_json_output(text: str) -> Tuple[Optional[Any], bool]:
-    """Return (parsed, ok). Strips ```json fences if present."""
-    stripped = text.strip()
-    if stripped.startswith("```"):
-        stripped = stripped.split("\n", 1)[-1] if "\n" in stripped else stripped[3:]
-        if stripped.endswith("```"):
-            stripped = stripped[:-3].rstrip()
-    try:
-        return json.loads(stripped), True
-    except (json.JSONDecodeError, ValueError):
-        return None, False
-
-
-# --- HF backend (lift-pdf package, in-process Transformers) ---
-def make_hf_processor(schema: Dict[str, Any], max_tokens: Optional[int]):
-    """Load lift via the package's HF backend; return a batch-processing closure."""
-    from lift.model import InferenceManager
-    from lift.model.schema import BatchInputItem
-
-    logger.info("Loading lift via Transformers (method=hf)...")
-    manager = InferenceManager(method="hf")
-
-    def process(image_lists: List[List[Image.Image]]) -> List[DocResult]:
-        items = [
-            BatchInputItem(images=imgs, schema=schema, prompt_type="direct")
-            for imgs in image_lists
-        ]
-        results = manager.generate(items, max_output_tokens=max_tokens)
-        return [(r.extraction, bool(r.error), r.raw) for r in results]
-
-    return process
-
-
-# --- vLLM backend (offline LLM() engine + structured outputs) ---
-def build_guided_schema(schema: Dict[str, Any]) -> Dict[str, Any]:
-    """Reproduce lift's vLLM guided-decoding schema: JSON Schema -> pydantic ->
-    json_schema with every leaf made nullable (so absent fields can be null,
-    matching lift's own server-side behavior)."""
-    from json_schema_to_pydantic import create_model
-    from lift.model.vllm import make_properties_nullable
-
-    schema_model = create_model(schema)
-    json_schema = schema_model.model_json_schema()
-    make_properties_nullable(json_schema)
-    return json_schema
-
-
-def make_sampling_params(json_schema: Dict[str, Any], max_tokens: int):
-    """SamplingParams with structured JSON output, across vLLM API versions.
-
-    lift uses greedy-ish decoding (temperature 0.0, top_p 0.1).
-    """
-    from vllm import SamplingParams
-
-    # vLLM >= 0.12
-    try:
-        from vllm.sampling_params import StructuredOutputsParams
-
-        return SamplingParams(
-            temperature=0.0,
-            top_p=0.1,
-            max_tokens=max_tokens,
-            structured_outputs=StructuredOutputsParams(json=json_schema),
-        )
-    except (ImportError, TypeError):
-        pass
-    # Older vLLM
-    try:
-        from vllm.sampling_params import GuidedDecodingParams
-
-        return SamplingParams(
-            temperature=0.0,
-            top_p=0.1,
-            max_tokens=max_tokens,
-            guided_decoding=GuidedDecodingParams(json=json_schema),
-        )
-    except (ImportError, TypeError):
-        pass
-    logger.warning(
-        "Structured output unavailable in this vLLM version; relying on lift's "
-        "training to emit valid JSON."
-    )
-    return SamplingParams(temperature=0.0, top_p=0.1, max_tokens=max_tokens)
-
-
-def make_vllm_processor(
-    schema: Dict[str, Any],
-    model: str,
-    max_tokens: Optional[int],
-    max_model_len: int,
-    gpu_memory_utilization: float,
-    max_images_per_doc: int,
-):
-    """Load lift into vLLM's offline engine; return a batch-processing closure."""
-    try:
-        from vllm import LLM
-    except ImportError as e:
-        raise RuntimeError(
-            "--method vllm needs vLLM. Run on the vllm/vllm-openai image: "
-            "--image vllm/vllm-openai --python /usr/bin/python3 "
-            "-e PYTHONPATH=/usr/local/lib/python3.12/dist-packages"
-        ) from e
-    from lift.model.util import scale_to_fit
-    from lift.prompts import PROMPT_MAPPING
-
-    json_schema = build_guided_schema(schema)
-    prompt = PROMPT_MAPPING["direct"].replace("{schema}", json.dumps(schema, indent=2))
-
-    logger.info("Loading lift via vLLM offline engine (method=vllm)...")
-    llm = LLM(
-        model=model,
-        trust_remote_code=True,
-        max_model_len=max_model_len,
-        gpu_memory_utilization=gpu_memory_utilization,
-        limit_mm_per_prompt={"image": max_images_per_doc},
-        # lift's own server-side image bounds, applied by the offline processor too.
-        mm_processor_kwargs={"min_pixels": 3136, "max_pixels": 861696},
-    )
-    sampling_params = make_sampling_params(
-        json_schema, max_tokens or DEFAULT_MAX_TOKENS
-    )
-
-    def process(image_lists: List[List[Image.Image]]) -> List[DocResult]:
-        messages = []
-        for imgs in image_lists:
-            content = [
-                {
-                    "type": "image_url",
-                    "image_url": {"url": pil_to_data_uri(scale_to_fit(img))},
-                }
-                for img in imgs
-            ]
-            content.append({"type": "text", "text": prompt})
-            messages.append([{"role": "user", "content": content}])
-        outputs = llm.chat(
-            messages, sampling_params, chat_template_content_format="openai"
-        )
-        results: List[DocResult] = []
-        for o in outputs:
-            raw = o.outputs[0].text
-            parsed, ok = parse_json_output(raw)
-            results.append((parsed if ok else None, not ok, raw))
-        return results
-
-    return process
-
-
-def create_dataset_card(
-    source_dataset: str,
-    model: str,
-    method: str,
-    schema: Dict[str, Any],
-    num_samples: int,
-    n_valid: int,
-    source_column: str,
-    is_pdf: bool,
-    page_range: Optional[str],
-    output_column: str,
-    split: str,
-    processing_time: str,
-) -> str:
-    """Build the output dataset card documenting the lift run."""
-    schema_block = json.dumps(schema, indent=2)
-    input_kind = "PDF documents" if is_pdf else "images"
-    col_desc = "PDF" if is_pdf else "image"
-    if page_range:
-        col_desc += f", pages {page_range}"
-    backend_desc = (
-        "vLLM offline engine" if method == "vllm" else "Transformers (lift-pdf)"
-    )
-    return f"""---
-tags:
-- ocr
-- structured-extraction
-- document-processing
-- lift
-- json
-- uv-script
-- generated
----
-
-# lift structured extraction on {source_dataset}
-
-Schema-constrained JSON extracted from {input_kind} in
-[{source_dataset}](https://huggingface.co/datasets/{source_dataset}) using
-[lift](https://huggingface.co/{model}) (9B, Qwen3.5-based) by Datalab, via the
-[`lift-pdf`](https://github.com/datalab-to/lift) package.
-
-## Processing Details
-
-- **Source Dataset**: [{source_dataset}](https://huggingface.co/datasets/{source_dataset})
-- **Model**: [{model}](https://huggingface.co/{model})
-- **Backend**: `{method}` ({backend_desc})
-- **Input column**: `{source_column}` ({col_desc})
-- **Output column**: `{output_column}` (JSON string per row)
-- **Split**: `{split}`
-- **Samples**: {num_samples:,}
-- **Valid JSON**: {n_valid:,} / {num_samples:,}
-- **Processing time**: {processing_time}
-- **Date**: {datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC")}
-
-### Extraction Schema
-
-```json
-{schema_block}
-```
-
-## License note
-
-lift's code is Apache-2.0, but the model **weights** use a modified OpenRAIL-M
-license: free for research, personal use, and startups under $5M funding/revenue,
-restricted from competitive use against Datalab's API. See the
-[model card](https://huggingface.co/{model}).
-
-## Dataset Structure
-
-Original columns plus:
-- `{output_column}`: lift output (JSON string; raw text kept on parse failure)
-- `inference_info`: JSON list tracking models applied to this dataset
-
-Generated with [UV Scripts](https://huggingface.co/uv-scripts).
-"""
-
-
-def main(
-    input_dataset: str,
-    output_dataset: str,
-    schema_arg: str,
-    image_column: str = "image",
-    pdf_column: Optional[str] = None,
-    output_column: str = "extraction",
-    overwrite: bool = False,
-    method: str = "hf",
-    page_range: Optional[str] = None,
-    split: str = "train",
-    max_samples: Optional[int] = None,
-    shuffle: bool = False,
-    seed: int = 42,
-    batch_size: int = 8,
-    max_tokens: Optional[int] = None,
-    max_model_len: int = 32768,
-    gpu_memory_utilization: float = 0.9,
-    max_images_per_doc: Optional[int] = None,
-    model: str = DEFAULT_MODEL,
-    private: bool = False,
-    config: Optional[str] = None,
-    create_pr: bool = False,
-    hf_token: Optional[str] = None,
-    verbose: bool = False,
-) -> None:
-    # Unlock full Xet bandwidth for the 9B (~19GB) model download (repo convention).
-    os.environ["HF_XET_HIGH_PERFORMANCE"] = "1"
-    check_cuda_availability()
-    start_time = datetime.now(timezone.utc)
-
-    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")
-    if HF_TOKEN:
-        login(token=HF_TOKEN)
-
-    schema = load_schema_arg(schema_arg)
-
-    # lift reads the checkpoint from env (pydantic-settings) at import time; set it first.
-    os.environ["MODEL_CHECKPOINT"] = model
-
-    # Import lift only after env is set so settings pick up the right checkpoint.
-    from lift import resolve_schema
-    from lift.input import load_file
-
-    schema = resolve_schema(schema)  # validates and normalizes
-    fields = list(schema.get("properties", {}).keys())
-
-    source_column = pdf_column or image_column
-    is_pdf = pdf_column is not None
-    # vLLM caps images per prompt at init; PDFs need headroom for multiple pages.
-    if max_images_per_doc is None:
-        max_images_per_doc = 30 if is_pdf else 1
-
-    logger.info(f"Model: {model}  Backend: {method}")
-    logger.info(f"Schema top-level fields: {fields}")
-
-    logger.info(f"Loading dataset: {input_dataset} (split={split})")
-    dataset = load_dataset(input_dataset, split=split)
-    if source_column not in dataset.column_names:
-        logger.error(
-            f"Column '{source_column}' not found. Available: {dataset.column_names}"
-        )
-        sys.exit(1)
-
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
-    if shuffle:
-        dataset = dataset.shuffle(seed=seed)
-    if max_samples:
-        dataset = dataset.select(range(min(max_samples, len(dataset))))
-    logger.info(f"Processing {len(dataset)} documents from column '{source_column}'")
-
-    if method == "vllm":
-        process_batch = make_vllm_processor(
-            schema,
-            model,
-            max_tokens,
-            max_model_len,
-            gpu_memory_utilization,
-            max_images_per_doc,
-        )
-    else:
-        process_batch = make_hf_processor(schema, max_tokens)
-
-    extractions: List[Optional[str]] = [None] * len(dataset)
-    error_flags: List[bool] = [True] * len(dataset)
-
-    chunks = list(partition_all(batch_size, range(len(dataset))))
-    for chunk in tqdm(chunks, desc="Extracting"):
-        chunk = list(chunk)
-        rendered: Dict[int, List[Image.Image]] = {}
-        for i in chunk:
-            try:
-                rendered[i] = load_document_images(
-                    load_file, dataset[i][source_column], page_range
-                )
-            except Exception as e:
-                logger.warning(f"Row {i}: failed to load document: {e}")
-                extractions[i] = f"[LIFT LOAD ERROR] {e}"
-                error_flags[i] = True
-        if not rendered:
-            continue
-
-        idxs = list(rendered.keys())
-        try:
-            results = process_batch([rendered[i] for i in idxs])
-        except Exception as e:
-            logger.error(f"Batch generate failed: {e}")
-            for i in idxs:
-                extractions[i] = "[LIFT GENERATE ERROR]"
-                error_flags[i] = True
-            continue
-
-        for i, (parsed, err, raw) in zip(idxs, results):
-            if parsed is not None and not err:
-                extractions[i] = json.dumps(parsed, ensure_ascii=False)
-                error_flags[i] = False
-            else:
-                extractions[i] = raw if raw else "[LIFT EMPTY OUTPUT]"
-                error_flags[i] = True
-
-    n_valid = sum(not f for f in error_flags)
-    logger.info(f"Valid JSON: {n_valid}/{len(dataset)}")
-
-    dataset = dataset.add_column(output_column, extractions)
-
-    inference_entry = {
-        "model": model,
-        "model_name": "lift",
-        "column_name": output_column,
-        "task": "schema-constrained extraction",
-        "backend": method,
-        "fields": fields,
-        "page_range": page_range,
-        "parse_error_rate": (len(dataset) - n_valid) / len(dataset)
-        if len(dataset)
-        else 0.0,
-        "timestamp": datetime.now(timezone.utc).isoformat(),
-        "script": "lift-extract.py",
-    }
-    if "inference_info" in dataset.column_names:
-
-        def update_info(example):
-            try:
-                existing = (
-                    json.loads(example["inference_info"])
-                    if example["inference_info"]
-                    else []
-                )
-            except (json.JSONDecodeError, TypeError):
-                existing = []
-            existing.append(inference_entry)
-            return {"inference_info": json.dumps(existing)}
-
-        dataset = dataset.map(update_info)
-    else:
-        dataset = dataset.add_column(
-            "inference_info", [json.dumps([inference_entry])] * len(dataset)
-        )
-
-    processing_time = (
-        f"{(datetime.now(timezone.utc) - start_time).total_seconds() / 60:.1f} min"
-    )
-
-    logger.info(f"Pushing to {output_dataset}")
-    max_retries = 3
-    for attempt in range(1, max_retries + 1):
-        try:
-            if attempt > 1:
-                logger.warning("Disabling XET (fallback to HTTP upload)")
-                os.environ["HF_HUB_DISABLE_XET"] = "1"
-            dataset.push_to_hub(
-                output_dataset,
-                private=private,
-                token=HF_TOKEN,
-                max_shard_size="500MB",
-                create_pr=create_pr,
-                **({"config_name": config} if config else {}),
-                commit_message=f"Add lift extraction results ({len(dataset)} samples)"
-                + (f" [{config}]" if config else ""),
-            )
-            break
-        except Exception as e:
-            logger.error(f"Upload attempt {attempt}/{max_retries} failed: {e}")
-            if attempt < max_retries:
-                delay = 30 * (2 ** (attempt - 1))
-                logger.info(f"Retrying in {delay}s...")
-                time.sleep(delay)
-            else:
-                logger.error("All upload attempts failed. Results are lost.")
-                sys.exit(1)
-
-    try:
-        card = DatasetCard(
-            create_dataset_card(
-                source_dataset=input_dataset,
-                model=model,
-                method=method,
-                schema=schema,
-                num_samples=len(dataset),
-                n_valid=n_valid,
-                source_column=source_column,
-                is_pdf=is_pdf,
-                page_range=page_range,
-                output_column=output_column,
-                split=split,
-                processing_time=processing_time,
-            )
-        )
-        card.push_to_hub(output_dataset, token=HF_TOKEN)
-    except Exception as e:
-        logger.warning(f"Could not push dataset card: {e}")
-
-    logger.info("Done! lift extraction complete.")
-    logger.info(f"Dataset: https://huggingface.co/datasets/{output_dataset}")
-    logger.info(f"Processing time: {processing_time}")
-
-    if verbose:
-        import importlib.metadata
-
-        logger.info("--- Resolved package versions ---")
-        pkgs = ["lift-pdf", "transformers", "torch", "datasets", "pillow", "openai"]
-        if method == "vllm":
-            pkgs.append("vllm")
-        for pkg in pkgs:
-            try:
-                logger.info(f"  {pkg}=={importlib.metadata.version(pkg)}")
-            except importlib.metadata.PackageNotFoundError:
-                logger.info(f"  {pkg}: not installed")
-
-
-if __name__ == "__main__":
-    if len(sys.argv) == 1:
-        print("lift — schema-constrained JSON extraction from images & PDFs (9B)")
-        print("\nUsage:")
-        print("  uv run lift-extract.py INPUT OUTPUT --schema SCHEMA [options]")
-        print("\nExamples:")
-        print("  # image column -> JSON")
-        print("  uv run lift-extract.py my-images my-fields \\")
-        print(
-            '    --schema \'{"type":"object","properties":{"title":{"type":"string"}}}\''
-        )
-        print("\n  # multi-page PDFs -> JSON (one extraction per document)")
-        print(
-            "  uv run lift-extract.py my-pdfs my-fields --pdf-column pdf --page-range 0-5 \\"
-        )
-        print("    --schema schema.json")
-        print("\n  --schema accepts inline JSON, a URL, or a file path.")
-        print(
-            "  --method hf (default) | vllm (offline LLM engine; needs the vllm image)"
-        )
-        print("\nFor full help: uv run lift-extract.py --help")
-        sys.exit(0)
-
-    parser = argparse.ArgumentParser(
-        description="Schema-constrained JSON extraction from images & PDFs using datalab-to/lift",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog="""
-Backends (both in-process, single command):
-  --method hf     Transformers via lift-pdf (default). Simplest; default image.
-  --method vllm   vLLM offline LLM() engine with structured outputs. Faster on
-                  large jobs. Needs the vllm/vllm-openai image.
-
-Input (one document per row):
-  --image-column COL   one image per row   (default: image)
-  --pdf-column COL     PDF bytes per row    (multi-page; honors --page-range)
-""",
-    )
-    parser.add_argument(
-        "input_dataset", help="Input dataset ID from the Hugging Face Hub"
-    )
-    parser.add_argument(
-        "output_dataset", help="Output dataset ID for the Hugging Face Hub"
-    )
-    parser.add_argument(
-        "--schema",
-        required=True,
-        help="JSON Schema: inline JSON, a URL, or a file path",
-    )
-    parser.add_argument(
-        "--image-column", default="image", help="Image column (default: image)"
-    )
-    parser.add_argument(
-        "--pdf-column",
-        default=None,
-        help="PDF column (bytes/path). Mutually exclusive with --image-column.",
-    )
-    parser.add_argument(
-        "--output-column",
-        default="extraction",
-        help="Output column (default: extraction)",
-    )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
-    parser.add_argument(
-        "--method",
-        choices=["hf", "vllm"],
-        default="hf",
-        help="Inference backend (default: hf)",
-    )
-    parser.add_argument(
-        "--page-range",
-        default=None,
-        help="Pages to extract from PDFs, e.g. '0-5,7' (PDF column only)",
-    )
-    parser.add_argument(
-        "--split", default="train", help="Dataset split (default: train)"
-    )
-    parser.add_argument(
-        "--max-samples", type=int, help="Limit number of documents (for testing)"
-    )
-    parser.add_argument(
-        "--shuffle", action="store_true", help="Shuffle before sampling"
-    )
-    parser.add_argument(
-        "--seed", type=int, default=42, help="Shuffle seed (default: 42)"
-    )
-    parser.add_argument(
-        "--batch-size",
-        type=int,
-        default=8,
-        help="Documents per generate() call (default: 8; lower for big multi-page PDFs)",
-    )
-    parser.add_argument(
-        "--max-tokens",
-        type=int,
-        default=None,
-        help=f"Max output tokens (default: lift's {DEFAULT_MAX_TOKENS})",
-    )
-    parser.add_argument(
-        "--max-model-len",
-        type=int,
-        default=32768,
-        help="vLLM context length (default: 32768; raise for long multi-page PDFs)",
-    )
-    parser.add_argument(
-        "--gpu-memory-utilization",
-        type=float,
-        default=0.9,
-        help="vLLM GPU memory fraction (default: 0.9)",
-    )
-    parser.add_argument(
-        "--max-images-per-doc",
-        type=int,
-        default=None,
-        help="vLLM images-per-prompt cap (default: 1 for images, 30 for PDFs)",
-    )
-    parser.add_argument(
-        "--model", default=DEFAULT_MODEL, help=f"Model ID (default: {DEFAULT_MODEL})"
-    )
-    parser.add_argument(
-        "--private", action="store_true", help="Make output dataset private"
-    )
-    parser.add_argument(
-        "--config",
-        default=None,
-        help="Config/subset name when pushing (for benchmarking backends in one repo)",
-    )
-    parser.add_argument(
-        "--create-pr",
-        action="store_true",
-        help="Push as a pull request instead of directly (for parallel benchmarking)",
-    )
-    parser.add_argument("--hf-token", help="Hugging Face API token (or set HF_TOKEN)")
-    parser.add_argument(
-        "--verbose",
-        action="store_true",
-        help="Log resolved package versions after processing",
-    )
-
-    args = parser.parse_args()
-
-    if args.pdf_column and args.image_column != "image":
-        parser.error("--image-column and --pdf-column are mutually exclusive.")
-
-    main(
-        input_dataset=args.input_dataset,
-        output_dataset=args.output_dataset,
-        schema_arg=args.schema,
-        image_column=args.image_column,
-        pdf_column=args.pdf_column,
-        output_column=args.output_column,
-        overwrite=args.overwrite,
-        method=args.method,
-        page_range=args.page_range,
-        split=args.split,
-        max_samples=args.max_samples,
-        shuffle=args.shuffle,
-        seed=args.seed,
-        batch_size=args.batch_size,
-        max_tokens=args.max_tokens,
-        max_model_len=args.max_model_len,
-        gpu_memory_utilization=args.gpu_memory_utilization,
-        max_images_per_doc=args.max_images_per_doc,
-        model=args.model,
-        private=args.private,
-        config=args.config,
-        create_pr=args.create_pr,
-        hf_token=args.hf_token,
-        verbose=args.verbose,
-    )

lighton-ocr.py

@@ -1,7 +1,7 @@
 # /// script
 # requires-python = ">=3.11"
 # dependencies = [
-#     "datasets>=4.0.0",
+#     "datasets",
 #     "huggingface-hub",
 #     "pillow",
 #     "vllm",
@@ -57,10 +57,6 @@ from huggingface_hub import DatasetCard, login
 from PIL import Image
 from toolz import partition_all
 from tqdm.auto import tqdm
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
 from vllm import LLM, SamplingParams
 
 logging.basicConfig(level=logging.INFO)
@@ -85,28 +81,6 @@ def check_cuda_availability():
         logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
 
 
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
 def resize_image_to_target(image: Image.Image, target_size: int = 1540) -> Image.Image:
     """
     Resize image so longest dimension is target_size while maintaining aspect ratio.
@@ -317,7 +291,6 @@ def main(
     shuffle: bool = False,
     seed: int = 42,
     output_column: str = "markdown",
-    overwrite: bool = False,
 ):
     """Process images from HF dataset through LightOnOCR model."""
 
@@ -350,9 +323,6 @@ def main(
             f"Column '{image_column}' not found. Available: {dataset.column_names}"
         )
 
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
     # Shuffle if requested
     if shuffle:
         logger.info(f"Shuffling dataset with seed {seed}")
@@ -651,12 +621,6 @@ Examples:
         default="markdown",
         help="Column name for output text (default: markdown)",
     )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
 
     args = parser.parse_args()
 
@@ -680,5 +644,4 @@ Examples:
         shuffle=args.shuffle,
         seed=args.seed,
         output_column=args.output_column,
-        overwrite=args.overwrite,
     )

lighton-ocr2-server.py

@@ -1,651 +0,0 @@
-# /// script
-# requires-python = ">=3.11"
-# dependencies = [
-#     "datasets>=4.0.0",
-#     "huggingface-hub",
-#     "pillow",
-#     "requests",
-# ]
-# ///
-
-"""
-Convert document images to markdown using LightOnOCR-2 via an in-job vLLM server.
-
-Same model, message shape, and sampling as lighton-ocr2.py, but serves the
-model behind `vllm serve` inside the job (the model card's own documented
-path) and posts images concurrently — continuous batching stays fed instead
-of draining at each offline batch barrier, and a bad image fails one request
-instead of a whole batch. Measured ~1.8x the offline recipe's inference
-throughput on a 100-page historical-scan smoke test (l4x1, concurrency 32).
-
-This script is the *driver* half: it expects the server on localhost
-(started by the job command below), loads the input dataset, posts images
-concurrently, and pushes the result dataset. The driver has no torch/vllm
-deps, so `uv run` starts in seconds while the server warms up in parallel.
-
-Run on HF Jobs (standard uv-run shape — the script starts `vllm serve` itself
-as a subprocess when no server is already reachable; the only thing to get
-right is the --image flag, which provides the `vllm` binary):
-
-  hf jobs uv run --detach --flavor l4x1 -s HF_TOKEN --timeout 4h \\
-      --image vllm/vllm-openai:v0.22.1 \\
-      https://huggingface.co/datasets/uv-scripts/ocr/raw/main/lighton-ocr2-server.py \\
-      <input-dataset> <output-dataset>
-
-To use an already-running or remote endpoint instead, pass --server URL — the
-script only spawns a server when the (default localhost) URL is unreachable.
-The serve flags live in SERVE_ARGS below (the model card's own recommended
-command: `--limit-mm-per-prompt`, `--mm-processor-cache-gb 0`,
-`--no-enable-prefix-caching` — OCR never reuses images, so the caches only
-cost memory).
-
-Model: lightonai/LightOnOCR-2-1B (1B, Apache-2.0)
-- Message is the image ONLY (no text prompt) — LightOnOCR-2's trained format.
-- Images resized client-side so the longest dimension is 1540px (training
-  resolution at 200 DPI), same as the offline recipe.
-- Sampling per the card: temperature 0.2, top_p 0.9, max_tokens 4096.
-"""
-
-import argparse
-import atexit
-import base64
-import concurrent.futures
-import io
-import json
-import logging
-import os
-import shutil
-import subprocess
-import sys
-import threading
-import time
-from datetime import datetime
-from typing import Any, Dict, Union
-from urllib.parse import urlparse
-
-import requests
-from datasets import load_dataset
-from huggingface_hub import DatasetCard, login
-from PIL import Image
-
-logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
-logger = logging.getLogger(__name__)
-
-MODEL = "lightonai/LightOnOCR-2-1B"
-
-DEFAULT_TARGET_SIZE = 1540  # longest dimension; LightOnOCR-2 training resolution
-
-# The serve command this script spawns when no server is reachable — the model
-# card's own recommended flags; single source of truth for the serving config.
-SERVE_ARGS = [
-    "vllm", "serve", MODEL,
-    "--limit-mm-per-prompt", '{"image": 1}',
-    "--mm-processor-cache-gb", "0",
-    "--no-enable-prefix-caching",
-    "--max-model-len", "8192",
-    "--gpu-memory-utilization", "0.8",
-    "--port", "8000",
-]
-
-RUN_COMMAND = (
-    "hf jobs uv run --detach --flavor l4x1 -s HF_TOKEN --timeout 4h \\\n"
-    "    --image vllm/vllm-openai:v0.22.1 \\\n"
-    "    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/lighton-ocr2-server.py \\\n"
-    "    <input-dataset> <output-dataset>"
-)
-
-
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
-def to_pil_image(image: Union[Image.Image, Dict[str, Any], str]) -> Image.Image:
-    """Convert a dataset image cell (PIL image, bytes dict, or path) to RGB PIL."""
-    if isinstance(image, Image.Image):
-        pil_img = image
-    elif isinstance(image, dict) and "bytes" in image:
-        pil_img = Image.open(io.BytesIO(image["bytes"]))
-    elif isinstance(image, str):
-        pil_img = Image.open(image)
-    else:
-        raise ValueError(f"Unsupported image type: {type(image)}")
-    return pil_img.convert("RGB")
-
-
-def encode_image(image, target_size: int) -> str:
-    """RGB-convert, resize longest dimension to target_size, return base64 PNG."""
-    img = to_pil_image(image)
-    if target_size:
-        w, h = img.size
-        if max(w, h) != target_size:
-            scale = target_size / max(w, h)
-            img = img.resize(
-                (max(1, int(w * scale)), max(1, int(h * scale))),
-                Image.Resampling.LANCZOS,
-            )
-    buf = io.BytesIO()
-    img.save(buf, format="PNG")
-    return base64.b64encode(buf.getvalue()).decode()
-
-
-def server_alive(server: str) -> bool:
-    try:
-        return requests.get(f"{server}/health", timeout=5).status_code == 200
-    except requests.RequestException:
-        return False
-
-
-def wait_for_server(server: str, timeout_s: int, proc: "subprocess.Popen | None" = None):
-    logger.info(f"Waiting for server at {server}...")
-    deadline = time.time() + timeout_s
-    while time.time() < deadline:
-        if server_alive(server):
-            logger.info("Server is ready")
-            return
-        if proc is not None and proc.poll() is not None:
-            logger.error(f"Spawned vllm serve exited with code {proc.returncode} before becoming ready")
-            sys.exit(1)
-        time.sleep(10)
-    logger.error(f"Server did not become ready within {timeout_s}s")
-    sys.exit(1)
-
-
-def ensure_server(server: str, timeout_s: int = 1800):
-    """Use a reachable server; otherwise spawn `vllm serve` ourselves; else fail fast.
-
-    Spawning is only attempted for a localhost URL — a remote --server that is
-    down is the user's to fix, not ours to shadow with a local model.
-    """
-    if server_alive(server):
-        logger.info(f"Using already-running server at {server}")
-        return
-
-    host = urlparse(server).hostname or ""
-    if host not in ("127.0.0.1", "localhost", "0.0.0.0"):
-        logger.info(f"Remote server {server} not up yet — waiting for it")
-        wait_for_server(server, timeout_s)
-        return
-
-    if shutil.which("vllm") is None:
-        logger.error("No server is running and the `vllm` binary is not on PATH.")
-        logger.error("Run this script on a vLLM image so it can start the server itself:\n")
-        logger.error(RUN_COMMAND)
-        logger.error("\n(or start `vllm serve` yourself / pass --server URL of a running endpoint)")
-        sys.exit(1)
-
-    logger.info(f"Starting server: {' '.join(SERVE_ARGS)}")
-    proc = subprocess.Popen(SERVE_ARGS)  # logs interleave with ours on stdout/stderr
-    atexit.register(proc.terminate)  # don't leave a GPU server behind on local runs
-    wait_for_server(server, timeout_s, proc=proc)
-
-
-def ocr_one(
-    server: str,
-    image,
-    target_size: int,
-    max_tokens: int,
-    temperature: float,
-    top_p: float,
-    timeout_s: int,
-    retries: int = 2,
-) -> str:
-    """OCR a single image via the chat completions API. Returns raw model text."""
-    b64 = encode_image(image, target_size)
-    payload = {
-        "model": MODEL,
-        "messages": [
-            {
-                "role": "user",
-                "content": [
-                    # Image ONLY — LightOnOCR-2 uses no text prompt.
-                    {
-                        "type": "image_url",
-                        "image_url": {"url": f"data:image/png;base64,{b64}"},
-                    },
-                ],
-            }
-        ],
-        "temperature": temperature,
-        "top_p": top_p,
-        "max_tokens": max_tokens,
-    }
-    last_err = None
-    for attempt in range(retries + 1):
-        try:
-            resp = requests.post(
-                f"{server}/v1/chat/completions", json=payload, timeout=timeout_s
-            )
-            resp.raise_for_status()
-            return resp.json()["choices"][0]["message"]["content"]
-        except Exception as e:
-            last_err = e
-            if attempt < retries:
-                time.sleep(10 * (attempt + 1))
-    raise RuntimeError(f"request failed after {retries + 1} attempts: {last_err}")
-
-
-def create_dataset_card(
-    source_dataset: str,
-    model: str,
-    num_samples: int,
-    num_errors: int,
-    processing_time: str,
-    images_per_sec: float,
-    concurrency: int,
-    max_tokens: int,
-    temperature: float,
-    target_size: int,
-    image_column: str = "image",
-    split: str = "train",
-) -> str:
-    """Create a dataset card documenting the OCR process."""
-    model_name = model.split("/")[-1]
-
-    # Canonical provenance stamp (see AGENTS.md): Jobs claim gated on JOB_ID, set by HF Jobs in-container.
-    on_jobs = os.environ.get("JOB_ID") is not None
-    hw = os.environ.get("ACCELERATOR") or ""
-    origin = (
-        "Produced on [Hugging Face Jobs](https://huggingface.co/docs/huggingface_hub/guides/jobs)"
-        + (f" (`{hw}`)" if hw else "")
-    ) if on_jobs else "Generated"
-    jobs_tag = "\n- hf-jobs" if on_jobs else ""
-
-    return f"""---
-tags:
-- ocr
-- document-processing
-- lighton-ocr
-- markdown
-- uv-script
-- generated{jobs_tag}
----
-
-# Document OCR using {model_name} (server mode)
-
-This dataset contains OCR results from images in [{source_dataset}](https://huggingface.co/datasets/{source_dataset}) using LightOnOCR-2 (1B), served behind an in-job vLLM server with concurrent requests (continuous batching).
-
-## Processing Details
-
-- **Source Dataset**: [{source_dataset}](https://huggingface.co/datasets/{source_dataset})
-- **Model**: [{model}](https://huggingface.co/{model})
-- **Number of Samples**: {num_samples:,}
-- **Failed Requests**: {num_errors:,} (marked `[OCR ERROR]`)
-- **Processing Time**: {processing_time}
-- **Throughput**: {images_per_sec:.2f} images/sec
-- **Processing Date**: {datetime.now().strftime("%Y-%m-%d %H:%M UTC")}
-
-### Configuration
-
-- **Mode**: vLLM server (`vllm serve`) + concurrent driver, {concurrency} concurrent requests
-- **Image Column**: `{image_column}`
-- **Dataset Split**: `{split}`
-- **Target Image Size**: {target_size}px (longest dimension)
-- **Max Output Tokens**: {max_tokens:,}
-- **Temperature**: {temperature}
-
-## Dataset Structure
-
-The dataset contains all original columns plus:
-- `markdown`: The extracted text in markdown format
-- `inference_info`: JSON list tracking all OCR models applied to this dataset
-
-## Reproduction
-
-{origin} with the [`lighton-ocr2-server.py`](https://huggingface.co/datasets/uv-scripts/ocr/raw/main/lighton-ocr2-server.py) recipe from [uv-scripts](https://huggingface.co/uv-scripts) — see the script docstring for the single `hf jobs run` command that starts the server and driver together. The offline-vLLM sibling recipe is [`lighton-ocr2.py`](https://huggingface.co/datasets/uv-scripts/ocr/raw/main/lighton-ocr2.py).
-"""
-
-
-def main(
-    input_dataset: str,
-    output_dataset: str,
-    image_column: str = "image",
-    server: str = "http://127.0.0.1:8000",
-    concurrency: int = 32,
-    max_tokens: int = 4096,
-    temperature: float = 0.2,
-    top_p: float = 0.9,
-    target_size: int = DEFAULT_TARGET_SIZE,
-    request_timeout: int = 1800,
-    hf_token: str = None,
-    split: str = "train",
-    max_samples: int = None,
-    private: bool = False,
-    shuffle: bool = False,
-    seed: int = 42,
-    output_column: str = "markdown",
-    overwrite: bool = False,
-    config: str = None,
-    create_pr: bool = False,
-):
-    """Process images from HF dataset through a LightOnOCR-2 vLLM server."""
-
-    start_time = datetime.now()
-
-    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")
-    if HF_TOKEN:
-        login(token=HF_TOKEN)
-
-    logger.info(f"Using model: {MODEL} via server {server}")
-
-    logger.info(f"Loading dataset: {input_dataset}")
-    dataset = load_dataset(input_dataset, split=split)
-
-    if image_column not in dataset.column_names:
-        raise ValueError(
-            f"Column '{image_column}' not found. Available: {dataset.column_names}"
-        )
-
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
-    if shuffle:
-        logger.info(f"Shuffling dataset with seed {seed}")
-        dataset = dataset.shuffle(seed=seed)
-
-    if max_samples:
-        dataset = dataset.select(range(min(max_samples, len(dataset))))
-        logger.info(f"Limited to {len(dataset)} samples")
-
-    # Reuse a reachable server, else spawn `vllm serve` (needs the vllm binary,
-    # i.e. a vllm/vllm-openai image), else fail fast with the correct command.
-    ensure_server(server)
-
-    n = len(dataset)
-    logger.info(f"Processing {n} images, concurrency {concurrency}")
-    all_outputs = [None] * n
-    errors = 0
-    done = 0
-    inference_start = time.time()
-    lock = threading.Lock()
-
-    def worker(i: int) -> None:
-        nonlocal errors, done
-        try:
-            text = ocr_one(
-                server,
-                dataset[i][image_column],
-                target_size,
-                max_tokens,
-                temperature,
-                top_p,
-                request_timeout,
-            )
-            all_outputs[i] = text.strip()
-        except Exception as e:
-            logger.error(f"Image {i} failed: {e}")
-            all_outputs[i] = "[OCR ERROR]"
-            with lock:
-                errors += 1
-        with lock:
-            done += 1
-            if done % 25 == 0 or done == n:
-                rate = done / max(time.time() - inference_start, 1e-9)
-                logger.info(f"{done}/{n} done ({rate:.2f} img/s, {errors} errors)")
-
-    with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as pool:
-        list(pool.map(worker, range(n)))
-
-    inference_secs = time.time() - inference_start
-    processing_duration = datetime.now() - start_time
-    processing_time_str = f"{processing_duration.total_seconds() / 60:.1f} min"
-    images_per_sec = n / inference_secs if inference_secs else 0.0
-
-    logger.info(f"Adding '{output_column}' column to dataset")
-    dataset = dataset.add_column(output_column, all_outputs)
-
-    # Inference info tracking
-    inference_entry = {
-        "model_id": MODEL,
-        "model_name": "LightOnOCR-2-1B",
-        "column_name": output_column,
-        "timestamp": datetime.now().isoformat(),
-        "temperature": temperature,
-        "top_p": top_p,
-        "max_tokens": max_tokens,
-        "target_size": target_size,
-        "mode": "vllm-server",
-        "concurrency": concurrency,
-    }
-
-    if "inference_info" in dataset.column_names:
-        logger.info("Updating existing inference_info column")
-
-        def update_inference_info(example):
-            try:
-                existing_info = (
-                    json.loads(example["inference_info"])
-                    if example["inference_info"]
-                    else []
-                )
-            except (json.JSONDecodeError, TypeError):
-                existing_info = []
-            existing_info.append(inference_entry)
-            return {"inference_info": json.dumps(existing_info)}
-
-        dataset = dataset.map(update_inference_info)
-    else:
-        logger.info("Creating new inference_info column")
-        inference_list = [json.dumps([inference_entry])] * len(dataset)
-        dataset = dataset.add_column("inference_info", inference_list)
-
-    # Push to hub with retry and XET fallback
-    logger.info(f"Pushing to {output_dataset}")
-    max_retries = 3
-    for attempt in range(1, max_retries + 1):
-        try:
-            if attempt > 1:
-                logger.warning("Disabling XET (fallback to HTTP upload)")
-                os.environ["HF_HUB_DISABLE_XET"] = "1"
-            dataset.push_to_hub(
-                output_dataset,
-                private=private,
-                token=HF_TOKEN,
-                max_shard_size="500MB",
-                **({"config_name": config} if config else {}),
-                create_pr=create_pr,
-                commit_message=f"Add {MODEL} OCR results ({len(dataset)} samples, server mode)"
-                + (f" [{config}]" if config else ""),
-            )
-            break
-        except Exception as e:
-            logger.error(f"Upload attempt {attempt}/{max_retries} failed: {e}")
-            if attempt < max_retries:
-                delay = 30 * (2 ** (attempt - 1))
-                logger.info(f"Retrying in {delay}s...")
-                time.sleep(delay)
-            else:
-                logger.error("All upload attempts failed. OCR results are lost.")
-                sys.exit(1)
-
-    logger.info("Creating dataset card")
-    card_content = create_dataset_card(
-        source_dataset=input_dataset,
-        model=MODEL,
-        num_samples=len(dataset),
-        num_errors=errors,
-        processing_time=processing_time_str,
-        images_per_sec=images_per_sec,
-        concurrency=concurrency,
-        max_tokens=max_tokens,
-        temperature=temperature,
-        target_size=target_size,
-        image_column=image_column,
-        split=split,
-    )
-
-    card = DatasetCard(card_content)
-    card.push_to_hub(output_dataset, token=HF_TOKEN)
-
-    logger.info("Done! LightOnOCR-2 server-mode processing complete.")
-    logger.info(
-        f"Dataset available at: https://huggingface.co/datasets/{output_dataset}"
-    )
-    logger.info(f"Processing time: {processing_time_str}")
-    logger.info(
-        f"Throughput: {images_per_sec:.2f} images/sec "
-        f"(inference only, excl. dataset load/push; {errors} errors)"
-    )
-
-
-if __name__ == "__main__":
-    if len(sys.argv) == 1:
-        print("=" * 70)
-        print("LightOnOCR-2 Document Processing (vLLM server mode)")
-        print("=" * 70)
-        print("\nSame model + outputs as lighton-ocr2.py, but drives an in-job")
-        print("`vllm serve` with concurrent requests — no batch barriers,")
-        print("per-image (not per-batch) failure isolation.")
-        print("\nThe server must already be running (the job command starts")
-        print("both — see the module docstring for the full `hf jobs run`).")
-        print("\nExamples:")
-        print("\n1. Basic OCR (server on localhost:8000):")
-        print("   uv run lighton-ocr2-server.py input-dataset output-dataset")
-        print("\n2. Test with a small sample:")
-        print("   uv run lighton-ocr2-server.py large-dataset test --max-samples 10 --shuffle")
-        print("\nFor full help: uv run lighton-ocr2-server.py --help")
-        sys.exit(0)
-
-    parser = argparse.ArgumentParser(
-        description="Document OCR using LightOnOCR-2 via an in-job vLLM server",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog="""
-Examples:
-  uv run lighton-ocr2-server.py my-docs analyzed-docs
-  uv run lighton-ocr2-server.py large-dataset test --max-samples 50 --shuffle
-See the module docstring for the full `hf jobs run` command (server + driver in one job).
-        """,
-    )
-
-    parser.add_argument("input_dataset", help="Input dataset ID from Hugging Face Hub")
-    parser.add_argument("output_dataset", help="Output dataset ID for Hugging Face Hub")
-    parser.add_argument(
-        "--image-column",
-        default="image",
-        help="Column containing images (default: image)",
-    )
-    parser.add_argument(
-        "--server",
-        default="http://127.0.0.1:8000",
-        help="vLLM server base URL (default: in-job localhost:8000)",
-    )
-    parser.add_argument(
-        "--concurrency",
-        type=int,
-        default=32,
-        help="Concurrent OCR requests (default: 32; vLLM queues excess internally, "
-        "so this mainly needs to be high enough to keep continuous batching fed)",
-    )
-    parser.add_argument(
-        "--max-tokens",
-        type=int,
-        default=4096,
-        help="Maximum tokens to generate (default: 4096, the model card value)",
-    )
-    parser.add_argument(
-        "--temperature",
-        type=float,
-        default=0.2,
-        help="Sampling temperature (default: 0.2, the model card value)",
-    )
-    parser.add_argument(
-        "--top-p",
-        type=float,
-        default=0.9,
-        help="Top-p sampling (default: 0.9, the model card value)",
-    )
-    parser.add_argument(
-        "--target-size",
-        type=int,
-        default=DEFAULT_TARGET_SIZE,
-        help=f"Resize images so the longest dimension is this many pixels before upload "
-        f"(default: {DEFAULT_TARGET_SIZE}, the model's training resolution); 0 disables",
-    )
-    parser.add_argument(
-        "--request-timeout",
-        type=int,
-        default=1800,
-        help="Per-request timeout in seconds (default: 1800)",
-    )
-    parser.add_argument("--hf-token", help="Hugging Face API token")
-    parser.add_argument(
-        "--split", default="train", help="Dataset split to use (default: train)"
-    )
-    parser.add_argument(
-        "--max-samples",
-        type=int,
-        help="Maximum number of samples to process (for testing)",
-    )
-    parser.add_argument(
-        "--private", action="store_true", help="Make output dataset private"
-    )
-    parser.add_argument(
-        "--config",
-        help="Config/subset name when pushing to Hub (for benchmarking multiple models in one repo)",
-    )
-    parser.add_argument(
-        "--create-pr",
-        action="store_true",
-        help="Create a pull request instead of pushing directly (for parallel benchmarking)",
-    )
-    parser.add_argument(
-        "--shuffle", action="store_true", help="Shuffle dataset before processing"
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=42,
-        help="Random seed for shuffling (default: 42)",
-    )
-    parser.add_argument(
-        "--output-column",
-        default="markdown",
-        help="Column name for output text (default: markdown)",
-    )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
-
-    args = parser.parse_args()
-
-    main(
-        input_dataset=args.input_dataset,
-        output_dataset=args.output_dataset,
-        image_column=args.image_column,
-        server=args.server,
-        concurrency=args.concurrency,
-        max_tokens=args.max_tokens,
-        temperature=args.temperature,
-        top_p=args.top_p,
-        target_size=args.target_size,
-        request_timeout=args.request_timeout,
-        hf_token=args.hf_token,
-        split=args.split,
-        max_samples=args.max_samples,
-        private=args.private,
-        shuffle=args.shuffle,
-        seed=args.seed,
-        output_column=args.output_column,
-        overwrite=args.overwrite,
-        config=args.config,
-        create_pr=args.create_pr,
-    )

lighton-ocr2.py

@@ -53,10 +53,6 @@ from huggingface_hub import DatasetCard, login
 from PIL import Image
 from toolz import partition_all
 from tqdm.auto import tqdm
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
 from vllm import LLM, SamplingParams
 
 logging.basicConfig(level=logging.INFO)
@@ -77,28 +73,6 @@ def check_cuda_availability():
         logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
 
 
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
 def resize_image_to_target(image: Image.Image, target_size: int = 1540) -> Image.Image:
     """
     Resize image so longest dimension is target_size while maintaining aspect ratio.
@@ -308,7 +282,6 @@ def main(
     shuffle: bool = False,
     seed: int = 42,
     output_column: str = "markdown",
-    overwrite: bool = False,
     config: str = None,
     create_pr: bool = False,
     verbose: bool = False,
@@ -338,9 +311,6 @@ def main(
             f"Column '{image_column}' not found. Available: {dataset.column_names}"
         )
 
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
     # Shuffle if requested
     if shuffle:
         logger.info(f"Shuffling dataset with seed {seed}")
@@ -584,12 +554,8 @@ Examples:
     parser.add_argument(
         "--max-model-len",
         type=int,
-        default=16384,
-        help=(
-            "Maximum model context length (default: 16384). A full page resized to "
-            "1540px is ~6k image tokens; with --max-tokens 4096 output that overflows "
-            "the old 8192 default at admission and vLLM rejects the request."
-        ),
+        default=8192,
+        help="Maximum model context length (default: 8192)",
     )
     parser.add_argument(
         "--max-tokens",
@@ -661,12 +627,6 @@ Examples:
         default="markdown",
         help="Column name for output text (default: markdown)",
     )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
     parser.add_argument(
         "--verbose",
         action="store_true",
@@ -694,7 +654,6 @@ Examples:
         shuffle=args.shuffle,
         seed=args.seed,
         output_column=args.output_column,
-        overwrite=args.overwrite,
         config=args.config,
         create_pr=args.create_pr,
         verbose=args.verbose,

models.json

@@ -1,386 +0,0 @@
-{
-  "_meta": {
-    "description": "Machine-readable catalog of the OCR recipes in this directory: script -> model, params, backend, required image pins, and what the model's own card claims about language coverage. Human-readable language notes: LANGUAGES.md. Language evidence levels, strongest first: per-language-benchmark | count-claim | named-list | multilingual-unspecified | english-only | not-stated | not-applicable. Claims come from model cards (compiled 2026-07-15) - treat as claims, not guarantees.",
-    "updated": "2026-07-15",
-    "run_pattern": "hf jobs uv run --flavor <flavor> -s HF_TOKEN https://huggingface.co/datasets/uv-scripts/ocr/raw/main/<script> <input-dataset> <output-dataset>"
-  },
-  "tesseract-ocr.py": {
-    "model_id": null,
-    "model_name": "Tesseract 5",
-    "params": null,
-    "backend": "pytesseract (CPU)",
-    "task": "ocr",
-    "output": "plain text",
-    "license": "apache-2.0",
-    "languages": {
-      "evidence": "named-list",
-      "count": "125 traineddata packs (~100+ languages + script models)",
-      "list_url": "https://github.com/tesseract-ocr/tessdata_best",
-      "notes": "Broadest named coverage; script models (Latin, Cyrillic, Devanagari, ...) cover languages without a dedicated pack. Legacy baseline quality."
-    }
-  },
-  "pp-ocrv6.py": {
-    "model_id": "PaddlePaddle/PP-OCRv6",
-    "params": "1.5M-34.5M",
-    "backend": "paddleocr",
-    "task": "ocr",
-    "output": "plain text",
-    "license": "apache-2.0",
-    "languages": { "evidence": "count-claim", "count": 48 }
-  },
-  "falcon-ocr.py": {
-    "model_id": "tiiuae/Falcon-OCR",
-    "params": "0.3B",
-    "backend": "falcon-perception",
-    "task": "ocr",
-    "output": "markdown",
-    "license": "apache-2.0",
-    "languages": { "evidence": "not-stated" }
-  },
-  "smoldocling-ocr.py": {
-    "model_id": "ds4sd/SmolDocling-256M-preview",
-    "params": "256M",
-    "backend": "transformers",
-    "task": "ocr",
-    "output": "DocTags",
-    "languages": { "evidence": "english-only" }
-  },
-  "surya-ocr.py": {
-    "model_id": "datalab-to/surya-ocr-2",
-    "params": "0.65B",
-    "backend": "vllm",
-    "image": "vllm/vllm-openai:v0.20.1",
-    "task": "ocr | layout | table",
-    "output": "markdown + surya_blocks (per-block HTML, bboxes, reading order)",
-    "license": "modified OpenRAIL-M",
-    "languages": {
-      "evidence": "per-language-benchmark",
-      "count": 91,
-      "benchmark_url": "https://github.com/datalab-to/surya/blob/master/static/docs/multilingual.md",
-      "notes": "38/91 languages score >=90%, 76/91 >=80%; weakest of top-15: Arabic 72.7%, Vietnamese 73.2%."
-    }
-  },
-  "glm-ocr.py": {
-    "model_id": "zai-org/GLM-OCR",
-    "params": "0.9B",
-    "backend": "vllm",
-    "task": "ocr",
-    "output": "markdown",
-    "languages": {
-      "evidence": "named-list",
-      "count": 8,
-      "named": ["zh", "en", "fr", "es", "ru", "de", "ja", "ko"],
-      "notes": "Declared in card metadata only; no per-language evidence in the card body."
-    }
-  },
-  "paddleocr-vl.py": {
-    "model_id": "PaddlePaddle/PaddleOCR-VL",
-    "params": "0.9B",
-    "backend": "vllm",
-    "task": "ocr | table | formula | chart",
-    "output": "markdown",
-    "languages": {
-      "evidence": "count-claim",
-      "count": 109,
-      "notes": "Names scripts (Cyrillic, Arabic, Devanagari, Thai); claims handwriting + historical documents; no per-language numbers."
-    }
-  },
-  "paddleocr-vl-1.5.py": {
-    "model_id": "PaddlePaddle/PaddleOCR-VL-1.5",
-    "params": "0.9B",
-    "backend": "transformers",
-    "task": "ocr (6 modes)",
-    "output": "markdown",
-    "languages": {
-      "evidence": "count-claim",
-      "count": "109+ (adds Tibetan, Bengali)",
-      "notes": "Claims improved rare-character and ancient-text recognition."
-    }
-  },
-  "paddleocr-vl-1.6.py": {
-    "model_id": "PaddlePaddle/PaddleOCR-VL-1.6",
-    "params": "0.9B",
-    "backend": "vllm",
-    "task": "ocr",
-    "output": "markdown",
-    "languages": {
-      "evidence": "count-claim",
-      "count": "109+ (inherited from 1.5, not restated)",
-      "notes": "Upgrade-focused card; language-specific claims are Chinese ancient documents and rare characters."
-    }
-  },
-  "ovis-ocr2.py": {
-    "model_id": "ATH-MaaS/OvisOCR2",
-    "params": "0.9B",
-    "backend": "vllm",
-    "task": "ocr",
-    "output": "markdown + LaTeX + HTML tables",
-    "license": "apache-2.0",
-    "languages": { "evidence": "not-stated" }
-  },
-  "ovis-ocr2-server.py": {
-    "model_id": "ATH-MaaS/OvisOCR2",
-    "params": "0.9B",
-    "backend": "vllm-server (in-job vllm serve + concurrent driver)",
-    "image": "vllm/vllm-openai:v0.22.1",
-    "task": "ocr",
-    "output": "markdown + LaTeX + HTML tables",
-    "license": "apache-2.0",
-    "notes": "Server-mode sibling of ovis-ocr2.py: ~1.7x its inference throughput, per-request failure isolation. See SERVING.md.",
-    "languages": { "evidence": "not-stated" }
-  },
-  "lighton-ocr.py": {
-    "model_id": "lightonai/LightOnOCR-1B-1025",
-    "params": "1B",
-    "backend": "vllm",
-    "task": "ocr",
-    "output": "markdown",
-    "languages": {
-      "evidence": "named-list",
-      "count": 9,
-      "named": ["en", "fr", "de", "es", "it", "nl", "pt", "sv", "da"],
-      "notes": "Explicitly European / Latin-alphabet focused."
-    }
-  },
-  "lighton-ocr2.py": {
-    "model_id": "lightonai/LightOnOCR-2-1B",
-    "params": "1B",
-    "backend": "vllm",
-    "task": "ocr",
-    "output": "markdown",
-    "languages": {
-      "evidence": "named-list",
-      "count": 11,
-      "named": ["en", "fr", "de", "es", "it", "nl", "pt", "sv", "da", "zh", "ja"],
-      "notes": "Adds zh/ja over v1; declared, not benchmarked per language."
-    }
-  },
-  "lighton-ocr2-server.py": {
-    "model_id": "lightonai/LightOnOCR-2-1B",
-    "params": "1B",
-    "backend": "vllm-server (in-job vllm serve + concurrent driver)",
-    "image": "vllm/vllm-openai:v0.22.1",
-    "task": "ocr",
-    "output": "markdown",
-    "notes": "Server-mode sibling of lighton-ocr2.py (the model card's own documented path): ~1.8x its inference throughput. See SERVING.md.",
-    "languages": {
-      "evidence": "named-list",
-      "count": 11,
-      "named": ["en", "fr", "de", "es", "it", "nl", "pt", "sv", "da", "zh", "ja"],
-      "notes": "Adds zh/ja over v1; declared, not benchmarked per language."
-    }
-  },
-  "hunyuan-ocr.py": {
-    "model_id": "tencent/HunyuanOCR",
-    "revision": "f6af82ee007fe6091b29fb3bb287b491ead41c82",
-    "params": "1B",
-    "backend": "vllm",
-    "task": "ocr",
-    "output": "markdown",
-    "license": "Hunyuan Community License (excludes EU/UK/KR)",
-    "languages": { "evidence": "multilingual-unspecified" }
-  },
-  "hunyuan-ocr-1.5.py": {
-    "model_id": "tencent/HunyuanOCR",
-    "params": "1B",
-    "backend": "vllm",
-    "task": "ocr | 12 task types",
-    "output": "markdown",
-    "license": "Hunyuan Community License (excludes EU/UK/KR)",
-    "languages": {
-      "evidence": "multilingual-unspecified",
-      "notes": "Explicitly targets low-resource + ancient-script OCR as a design goal; nothing enumerated."
-    }
-  },
-  "dots-ocr.py": {
-    "model_id": "rednote-hilab/dots.ocr",
-    "params": "1.7B",
-    "backend": "vllm",
-    "task": "ocr + layout",
-    "output": "markdown",
-    "languages": {
-      "evidence": "count-claim",
-      "count": 100,
-      "notes": "Explicit low-resource claim backed by in-house dots.ocr-bench (1,493 PDFs across 100 languages); aggregate scores only."
-    }
-  },
-  "firered-ocr.py": {
-    "model_id": "FireRedTeam/FireRed-OCR",
-    "params": "2.1B",
-    "backend": "vllm",
-    "task": "ocr",
-    "output": "markdown",
-    "license": "apache-2.0",
-    "languages": { "evidence": "not-stated" }
-  },
-  "abot-ocr.py": {
-    "model_id": "acvlab/ABot-OCR",
-    "params": "2B",
-    "backend": "vllm",
-    "image": "vllm/vllm-openai",
-    "task": "ocr",
-    "output": "markdown (text, LaTeX, HTML tables)",
-    "languages": { "evidence": "not-stated" }
-  },
-  "nanonets-ocr.py": {
-    "model_id": "nanonets/Nanonets-OCR-s",
-    "params": "2B",
-    "backend": "vllm",
-    "task": "ocr",
-    "output": "markdown (LaTeX, tables, forms)",
-    "languages": { "evidence": "english-only" }
-  },
-  "dots-mocr.py": {
-    "model_id": "rednote-hilab/dots.mocr",
-    "params": "3B",
-    "backend": "vllm",
-    "task": "ocr | 8 prompt modes incl. SVG, layout + bbox",
-    "output": "markdown",
-    "languages": { "evidence": "multilingual-unspecified" }
-  },
-  "nanonets-ocr2.py": {
-    "model_id": "nanonets/Nanonets-OCR2-3B",
-    "params": "3B",
-    "backend": "vllm",
-    "image": "vllm/vllm-openai:v0.10.2",
-    "task": "ocr",
-    "output": "markdown",
-    "languages": {
-      "evidence": "named-list",
-      "count": "11 named + 'many more'",
-      "named": ["en", "zh", "fr", "es", "pt", "de", "it", "ru", "ja", "ko", "ar"],
-      "notes": "Only card claiming multilingual handwriting; list illustrative, no per-language benchmark."
-    }
-  },
-  "deepseek-ocr-vllm.py": {
-    "model_id": "deepseek-ai/DeepSeek-OCR",
-    "params": "4B",
-    "backend": "vllm",
-    "task": "ocr | 5 resolution + 5 prompt modes",
-    "output": "markdown",
-    "license": "mit",
-    "languages": { "evidence": "multilingual-unspecified" }
-  },
-  "deepseek-ocr.py": {
-    "model_id": "deepseek-ai/DeepSeek-OCR",
-    "params": "4B",
-    "backend": "transformers",
-    "task": "ocr",
-    "output": "markdown",
-    "license": "mit",
-    "languages": { "evidence": "multilingual-unspecified" }
-  },
-  "deepseek-ocr2-vllm.py": {
-    "model_id": "deepseek-ai/DeepSeek-OCR-2",
-    "params": "3B",
-    "backend": "vllm (nightly)",
-    "image": "vllm/vllm-openai",
-    "task": "ocr",
-    "output": "markdown",
-    "license": "apache-2.0",
-    "languages": { "evidence": "multilingual-unspecified" }
-  },
-  "unlimited-ocr-vllm.py": {
-    "model_id": "baidu/Unlimited-OCR",
-    "params": "3.3B",
-    "backend": "vllm",
-    "image": "vllm/vllm-openai:unlimited-ocr",
-    "task": "ocr (layout-grounded markdown; single-image batch)",
-    "output": "markdown",
-    "license": "mit",
-    "languages": { "evidence": "multilingual-unspecified" }
-  },
-  "nuextract3.py": {
-    "model_id": "numind/NuExtract3",
-    "params": "4B",
-    "backend": "vllm",
-    "image": "vllm/vllm-openai",
-    "task": "ocr | schema-guided extraction",
-    "output": "markdown or JSON",
-    "languages": { "evidence": "multilingual-unspecified" }
-  },
-  "qianfan-ocr.py": {
-    "model_id": "baidu/Qianfan-OCR",
-    "params": "4.7B",
-    "backend": "vllm",
-    "task": "ocr",
-    "output": "markdown",
-    "languages": {
-      "evidence": "count-claim",
-      "count": 192,
-      "notes": "Headline claim; no list or per-language numbers."
-    }
-  },
-  "olmocr2-vllm.py": {
-    "model_id": "allenai/olmOCR-2-7B-1025-FP8",
-    "params": "7B",
-    "backend": "vllm",
-    "task": "ocr",
-    "output": "markdown",
-    "languages": { "evidence": "english-only" }
-  },
-  "rolm-ocr.py": {
-    "model_id": "reducto/RolmOCR",
-    "params": "7B",
-    "backend": "vllm",
-    "task": "ocr",
-    "output": "plain text",
-    "languages": { "evidence": "not-stated" }
-  },
-  "numarkdown-ocr.py": {
-    "model_id": "numind/NuMarkdown-8B-Thinking",
-    "params": "8B",
-    "backend": "vllm",
-    "task": "ocr (reasoning)",
-    "output": "markdown",
-    "languages": { "evidence": "not-stated" }
-  },
-  "lfm2-vl-extract.py": {
-    "model_id": "LiquidAI/LFM2.5-VL-1.6B-Extract",
-    "params": "1.6B",
-    "backend": "vllm",
-    "image": "vllm/vllm-openai",
-    "task": "schema-guided extraction (image -> JSON)",
-    "output": "JSON",
-    "languages": {
-      "evidence": "english-only",
-      "notes": "Vision card declares en only; the text sibling's 9-language list does NOT carry over."
-    }
-  },
-  "lfm2-extract.py": {
-    "model_id": "LiquidAI/LFM2-1.2B-Extract",
-    "params": "1.2B",
-    "backend": "vllm",
-    "image": "vllm/vllm-openai",
-    "task": "schema-guided extraction (text -> structured); chain after any OCR recipe",
-    "output": "JSON / XML / YAML",
-    "languages": {
-      "evidence": "named-list",
-      "count": 9,
-      "named": ["en", "ar", "zh", "fr", "de", "ja", "ko", "pt", "es"],
-      "notes": "Text-only model (not OCR). Card prose names 9 incl. Portuguese; card YAML lists 8 (pt missing)."
-    }
-  },
-  "lift-extract.py": {
-    "model_id": "datalab-to/lift",
-    "params": "9B",
-    "backend": "transformers | vllm",
-    "task": "schema-guided extraction (image or multi-page PDF -> JSON)",
-    "output": "JSON",
-    "license": "modified OpenRAIL-M",
-    "languages": { "evidence": "not-stated" }
-  },
-  "pp-doclayout.py": {
-    "model_id": "PaddlePaddle/PP-DocLayout-L",
-    "params": "123M",
-    "backend": "paddleocr",
-    "task": "layout detection (no text extraction)",
-    "output": "bounding boxes + region classes",
-    "languages": { "evidence": "not-applicable" }
-  },
-  "glm-ocr-v2.py": { "variant_of": "glm-ocr.py", "notes": "Adds checkpoint/resume for very large jobs." },
-  "glm-ocr-bucket.py": { "variant_of": "glm-ocr.py", "notes": "Reads images/PDFs from a mounted bucket, writes one .md per page." },
-  "falcon-ocr-bucket.py": { "variant_of": "falcon-ocr.py", "notes": "Reads images/PDFs from a mounted bucket, writes one .md per page." },
-  "surya-ocr-bucket.py": { "variant_of": "surya-ocr.py", "notes": "Bucket-to-bucket structured OCR (--io-mode mount|copy), resumable." },
-  "ocr-vllm-judge.py": { "model_id": "configurable", "task": "pairwise OCR-quality judging (VLM-as-judge)", "languages": { "evidence": "not-applicable" } }
-}

nanonets-ocr.py

@@ -40,10 +40,6 @@ from huggingface_hub import DatasetCard, login
 from PIL import Image
 from toolz import partition_all
 from tqdm.auto import tqdm
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
 from vllm import LLM, SamplingParams
 from datetime import datetime
 
@@ -61,28 +57,6 @@ def check_cuda_availability():
         logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
 
 
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
 def make_ocr_message(
     image: Union[Image.Image, Dict[str, Any], str],
     prompt: str = "Extract the text from the above document as if you were reading it naturally. Return the tables in html format. Return the equations in LaTeX representation. If there is an image in the document and image caption is not present, add a small description of the image inside the <img></img> tag; otherwise, add the image caption inside <img></img>. Watermarks should be wrapped in brackets. Ex: <watermark>OFFICIAL COPY</watermark>. Page numbers should be wrapped in brackets. Ex: <page_number>14</page_number> or <page_number>9/22</page_number>. Prefer using ☐ and ☑ for check boxes.",
@@ -130,15 +104,6 @@ def create_dataset_card(
     """Create a dataset card documenting the OCR process."""
     model_name = model.split("/")[-1]
 
-    # Canonical provenance stamp (see AGENTS.md): Jobs claim gated on JOB_ID, set by HF Jobs in-container.
-    on_jobs = os.environ.get("JOB_ID") is not None
-    hw = os.environ.get("ACCELERATOR") or ""
-    origin = (
-        "Produced on [Hugging Face Jobs](https://huggingface.co/docs/huggingface_hub/guides/jobs)"
-        + (f" (`{hw}`)" if hw else "")
-    ) if on_jobs else "Generated"
-    jobs_tag = "\n- hf-jobs" if on_jobs else ""
-
     return f"""---
 viewer: false
 tags:
@@ -147,7 +112,7 @@ tags:
 - nanonets
 - markdown
 - uv-script
-- generated{jobs_tag}
+- generated
 ---
 
 # Document OCR using {model_name}
@@ -211,10 +176,10 @@ for info in inference_info:
 
 ## Reproduction
 
-{origin} with the [`nanonets-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/raw/main/nanonets-ocr.py) recipe from [uv-scripts](https://huggingface.co/uv-scripts). Run it yourself:
+This dataset was generated using the [uv-scripts/ocr](https://huggingface.co/datasets/uv-scripts/ocr) Nanonets OCR script:
 
 ```bash
-hf jobs uv run https://huggingface.co/datasets/uv-scripts/ocr/raw/main/nanonets-ocr.py \\
+uv run https://huggingface.co/datasets/uv-scripts/ocr/raw/main/nanonets-ocr.py \\
     {source_dataset} \\
     <output-dataset> \\
     --image-column {image_column} \\
@@ -228,6 +193,8 @@ hf jobs uv run https://huggingface.co/datasets/uv-scripts/ocr/raw/main/nanonets-
 
 - **Processing Speed**: ~{num_samples / (float(processing_time.split()[0]) * 60):.1f} images/second
 - **GPU Configuration**: vLLM with {gpu_memory_utilization:.0%} GPU memory utilization
+
+Generated with 🤖 [UV Scripts](https://huggingface.co/uv-scripts)
 """
 
 
@@ -237,7 +204,7 @@ def main(
     image_column: str = "image",
     batch_size: int = 32,
     model: str = "nanonets/Nanonets-OCR-s",
-    max_model_len: int = 32768,
+    max_model_len: int = 8192,
     max_tokens: int = 15000,
     gpu_memory_utilization: float = 0.8,
     hf_token: str = None,
@@ -246,8 +213,6 @@ def main(
     private: bool = False,
     shuffle: bool = False,
     seed: int = 42,
-    output_column: str = "markdown",
-    overwrite: bool = False,
     verbose: bool = False,
 ):
     """Process images from HF dataset through OCR model."""
@@ -276,9 +241,6 @@ def main(
             f"Column '{image_column}' not found. Available: {dataset.column_names}"
         )
 
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
     # Shuffle if requested
     if shuffle:
         logger.info(f"Shuffling dataset with seed {seed}")
@@ -335,9 +297,9 @@ def main(
             # Add error placeholders for failed batch
             all_markdown.extend(["[OCR FAILED]"] * len(batch_images))
 
-    # Add output column to dataset
-    logger.info(f"Adding '{output_column}' column to dataset")
-    dataset = dataset.add_column(output_column, all_markdown)
+    # Add markdown column to dataset
+    logger.info("Adding markdown column to dataset")
+    dataset = dataset.add_column("markdown", all_markdown)
 
     # Handle inference_info tracking
     logger.info("Updating inference_info...")
@@ -345,7 +307,7 @@ def main(
     inference_entry = {
         "model_id": model,
         "model_name": "Nanonets-OCR-s",
-        "column_name": output_column,
+        "column_name": "markdown",
         "timestamp": datetime.now().isoformat(),
         "batch_size": batch_size,
         "max_tokens": max_tokens,
@@ -500,10 +462,8 @@ Examples:
     parser.add_argument(
         "--max-model-len",
         type=int,
-        default=32768,
-        help="Maximum model context length (default: 32768; must exceed --max-tokens "
-        "15000 plus the image/prompt tokens — the old 8192 default couldn't hold the "
-        "output budget). Model max is 128k, so there's ample room.",
+        default=8192,
+        help="Maximum model context length (default: 8192)",
     )
     parser.add_argument(
         "--max-tokens",
@@ -540,17 +500,6 @@ Examples:
         default=42,
         help="Random seed for shuffling (default: 42)",
     )
-    parser.add_argument(
-        "--output-column",
-        default="markdown",
-        help="Column name for the OCR output text (default: markdown)",
-    )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
     parser.add_argument(
         "--verbose",
         action="store_true",
@@ -574,7 +523,5 @@ Examples:
         private=args.private,
         shuffle=args.shuffle,
         seed=args.seed,
-        output_column=args.output_column,
-        overwrite=args.overwrite,
         verbose=args.verbose,
     )

nanonets-ocr2.py

@@ -4,8 +4,10 @@
 #     "datasets>=4.0.0",
 #     "huggingface-hub",
 #     "pillow",
+#     "vllm>=0.15.1",
 #     "tqdm",
 #     "toolz",
+#     "torch",
 # ]
 #
 # ///
@@ -24,20 +26,6 @@ Features:
 - Signature and watermark detection
 - Checkbox recognition
 - Multilingual support
-
-⚠️ Pinned vLLM image required. Nanonets-OCR2-3B is a Qwen2.5-VL model, and vLLM
->=0.11 regressed Qwen2.5-VL decoding to a stream of "!" (vllm-project/vllm#27775,
-#14126); 0.10.2 is the last known-good. So `vllm` and `torch` are intentionally
-omitted from the deps above and come from the `vllm/vllm-openai:v0.10.2` image at
-runtime via PYTHONPATH. Run it like:
-
-    hf jobs uv run --flavor a10g-small -s HF_TOKEN \
-        --image vllm/vllm-openai:v0.10.2 --python /usr/bin/python3 \
-        -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages \
-        nanonets-ocr2.py INPUT_DATASET OUTPUT_DATASET --max-samples 10 --shuffle
-
-Re-test the pin when a newer vLLM ships a Qwen2.5-VL decode fix; it can then move
-back to a plain `vllm>=...` dep on the default image.
 """
 
 import argparse
@@ -50,33 +38,12 @@ import sys
 from typing import Any, Dict, List, Union
 from datetime import datetime
 
-# Preflight: torch/vllm are deliberately NOT PEP 723 deps (see docstring) — they come from
-# the pinned vllm/vllm-openai:v0.10.2 image. On the bare uv image the imports below would
-# die with a raw ModuleNotFoundError; fail fast naming the exact flags instead.
-import importlib.util
-
-if importlib.util.find_spec("vllm") is None or importlib.util.find_spec("torch") is None:
-    sys.stderr.write(
-        "ERROR: vllm/torch not importable — this recipe cannot run on the bare uv image.\n"
-        "Nanonets-OCR2-3B needs vLLM 0.10.2 (>=0.11 regresses Qwen2.5-VL decoding to '!').\n"
-        "Re-run with the pinned image + interpreter flags:\n"
-        "    hf jobs uv run --flavor a10g-small -s HF_TOKEN \\\n"
-        "        --image vllm/vllm-openai:v0.10.2 --python /usr/bin/python3 \\\n"
-        "        -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages \\\n"
-        "        nanonets-ocr2.py INPUT_DATASET OUTPUT_DATASET ...\n"
-    )
-    sys.exit(1)
-
 import torch
 from datasets import load_dataset
 from huggingface_hub import DatasetCard, login
 from PIL import Image
 from toolz import partition_all
 from tqdm.auto import tqdm
-# Disable vLLM's FlashInfer sampler (JIT-compiles a CUDA kernel; greedy OCR doesn't use it).
-# Harmless no-op on the pinned vllm/vllm-openai image; kept in case the pin is ever relaxed
-# back to an environment without nvcc.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
 from vllm import LLM, SamplingParams
 
 logging.basicConfig(level=logging.INFO)
@@ -93,28 +60,6 @@ def check_cuda_availability():
         logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
 
 
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
 def make_ocr_message(
     image: Union[Image.Image, Dict[str, Any], str],
     prompt: str = "Extract the text from the above document as if you were reading it naturally. Return the tables in html format. Return the equations in LaTeX representation. If there is an image in the document and image caption is not present, add a small description of the image inside the <img></img> tag; otherwise, add the image caption inside <img></img>. Watermarks should be wrapped in brackets. Ex: <watermark>OFFICIAL COPY</watermark>. Page numbers should be wrapped in brackets. Ex: <page_number>14</page_number> or <page_number>9/22</page_number>. Prefer using ☐ and ☑ for check boxes.",
@@ -265,7 +210,7 @@ def main(
     image_column: str = "image",
     batch_size: int = 16,
     model: str = "nanonets/Nanonets-OCR2-3B",
-    max_model_len: int = 32768,
+    max_model_len: int = 8192,
     max_tokens: int = 15000,
     gpu_memory_utilization: float = 0.8,
     hf_token: str = None,
@@ -274,11 +219,7 @@ def main(
     private: bool = False,
     shuffle: bool = False,
     seed: int = 42,
-    output_column: str = "markdown",
-    overwrite: bool = False,
     verbose: bool = False,
-    config: str = None,
-    create_pr: bool = False,
 ):
     """Process images from HF dataset through Nanonets-OCR2-3B model."""
 
@@ -306,9 +247,6 @@ def main(
             f"Column '{image_column}' not found. Available: {dataset.column_names}"
         )
 
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
     # Shuffle if requested
     if shuffle:
         logger.info(f"Shuffling dataset with seed {seed}")
@@ -365,9 +303,9 @@ def main(
             # Add error placeholders for failed batch
             all_markdown.extend(["[OCR FAILED]"] * len(batch_images))
 
-    # Add output column to dataset
-    logger.info(f"Adding '{output_column}' column to dataset")
-    dataset = dataset.add_column(output_column, all_markdown)
+    # Add markdown column to dataset
+    logger.info("Adding markdown column to dataset")
+    dataset = dataset.add_column("markdown", all_markdown)
 
     # Handle inference_info tracking
     logger.info("Updating inference_info...")
@@ -375,7 +313,7 @@ def main(
     inference_entry = {
         "model_id": model,
         "model_name": "Nanonets-OCR2-3B",
-        "column_name": output_column,
+        "column_name": "markdown",
         "timestamp": datetime.now().isoformat(),
         "batch_size": batch_size,
         "max_tokens": max_tokens,
@@ -408,15 +346,7 @@ def main(
 
     # Push to hub
     logger.info(f"Pushing to {output_dataset}")
-    commit_info = dataset.push_to_hub(
-        output_dataset,
-        private=private,
-        token=HF_TOKEN,
-        **({"config_name": config} if config else {}),
-        create_pr=create_pr,
-        commit_message=f"Add {model} OCR results ({len(dataset)} samples)"
-        + (f" [{config}]" if config else ""),
-    )
+    dataset.push_to_hub(output_dataset, private=private, token=HF_TOKEN)
 
     # Calculate processing time
     end_time = datetime.now()
@@ -446,8 +376,6 @@ def main(
     logger.info(
         f"Dataset available at: https://huggingface.co/datasets/{output_dataset}"
     )
-    if create_pr and getattr(commit_info, "pr_url", None):
-        logger.info(f"Pull request created: {commit_info.pr_url}")
 
     if verbose:
         import importlib.metadata
@@ -544,8 +472,8 @@ Examples:
     parser.add_argument(
         "--max-model-len",
         type=int,
-        default=32768,
-        help="Maximum model context length (default: 32768)",
+        default=8192,
+        help="Maximum model context length (default: 8192)",
     )
     parser.add_argument(
         "--max-tokens",
@@ -571,15 +499,6 @@ Examples:
     parser.add_argument(
         "--private", action="store_true", help="Make output dataset private"
     )
-    parser.add_argument(
-        "--config",
-        help="Config/subset name when pushing to Hub (for benchmarking multiple models in one repo)",
-    )
-    parser.add_argument(
-        "--create-pr",
-        action="store_true",
-        help="Create a pull request instead of pushing directly (for parallel benchmarking)",
-    )
     parser.add_argument(
         "--shuffle",
         action="store_true",
@@ -591,17 +510,6 @@ Examples:
         default=42,
         help="Random seed for shuffling (default: 42)",
     )
-    parser.add_argument(
-        "--output-column",
-        default="markdown",
-        help="Column name for the OCR output text (default: markdown)",
-    )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
     parser.add_argument(
         "--verbose",
         action="store_true",
@@ -625,9 +533,5 @@ Examples:
         private=args.private,
         shuffle=args.shuffle,
         seed=args.seed,
-        output_column=args.output_column,
-        overwrite=args.overwrite,
         verbose=args.verbose,
-        config=args.config,
-        create_pr=args.create_pr,
     )

nuextract3.py

@@ -1,786 +0,0 @@
-# /// script
-# requires-python = ">=3.11"
-# dependencies = [
-#     "datasets>=3.1.0",
-#     "huggingface-hub",
-#     "pillow",
-#     "vllm",
-#     "toolz",
-#     "torch",
-#     "numind",
-# ]
-# ///
-
-"""
-Convert document images to markdown OR extract structured JSON using NuExtract3 with vLLM.
-
-NuExtract3 is a 4B Qwen3.5-based VLM for document understanding. It does two things:
-
-1. Document-to-Markdown OCR (default): images -> clean markdown with HTML tables,
-   LaTeX math, and <figure> tags.
-2. Schema-guided structured extraction: images + a JSON template -> JSON output
-   shaped exactly like the template. Useful for invoices, receipts, forms, contracts.
-
-Modes are selected via flags:
-- (no flags)         -> markdown OCR
-- --mode content     -> plain-content extraction
-- --template SOURCE  -> structured extraction with a NuExtract template
-- --schema SOURCE    -> structured extraction with a JSON Schema
-                        (auto-converted via numind.nuextract_utils)
-- --instructions STR -> free-text guidance passed through to the model
-                        (output-format rules, branch routing, etc.).
-                        Combines with any of the modes above.
-                        See https://huggingface.co/numind/NuExtract3#instructions
-
---template / --schema each accept inline JSON, a URL, or a local file path, so a
-schema can be hosted (e.g. on an HF dataset's raw URL) and reused across jobs:
-    --template https://huggingface.co/datasets/ORG/REPO/raw/main/card.json
-
-HF Jobs invocation (recommended): use the vllm/vllm-openai:latest image so the
-pre-built CUDA kernels (flashinfer etc.) are reused — the default uv-script
-image lacks nvcc and flashinfer's JIT compile fails at engine warmup.
-
-    hf jobs uv run \\
-        --image vllm/vllm-openai:latest \\
-        --flavor a100-large \\
-        --python /usr/bin/python3 \\
-        -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages \\
-        -s HF_TOKEN \\
-        https://huggingface.co/datasets/uv-scripts/ocr/raw/main/nuextract3.py \\
-        INPUT_DATASET OUTPUT_DATASET --max-samples 5 --shuffle --seed 42
-
-Model: numind/NuExtract3
-License: Apache-2.0
-"""
-
-import argparse
-import base64
-import io
-import json
-import logging
-import os
-import sys
-import time
-from datetime import datetime
-from pathlib import Path
-from typing import Any, Dict, List, Optional, Union
-
-import torch
-from datasets import load_dataset
-from huggingface_hub import DatasetCard, login
-from PIL import Image
-from toolz import partition_all
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
-from vllm import LLM, SamplingParams
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-MODEL_DEFAULT = "numind/NuExtract3"
-MODEL_NAME = "NuExtract3"
-
-
-def check_cuda_availability():
-    """Check if CUDA is available and exit if not."""
-    if not torch.cuda.is_available():
-        logger.error("CUDA is not available. This script requires a GPU.")
-        logger.error("Please run on a machine with a CUDA-capable GPU.")
-        sys.exit(1)
-    else:
-        logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
-
-
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
-def load_template_arg(value: Optional[str]) -> Optional[Dict[str, Any]]:
-    """Load a NuExtract template/JSON Schema from inline JSON, a URL, or a file path."""
-    if value is None:
-        return None
-    text = value
-    if value.startswith(("http://", "https://")):
-        import urllib.request
-
-        with urllib.request.urlopen(value) as resp:  # noqa: S310
-            text = resp.read().decode("utf-8")
-    elif "{" not in value:
-        # Inline JSON often exceeds the OS filename limit, so only probe the
-        # filesystem when the value doesn't look like JSON; treat OSError as
-        # "not a path".
-        try:
-            candidate_path = Path(value)
-            if candidate_path.is_file():
-                text = candidate_path.read_text()
-        except OSError:
-            pass
-    try:
-        return json.loads(text)
-    except json.JSONDecodeError as e:
-        raise ValueError(
-            f"Could not parse template/schema as JSON (tried URL/path/inline): {e}"
-        ) from e
-
-
-def resolve_template(
-    template_arg: Optional[str],
-    schema_arg: Optional[str],
-) -> Optional[Dict[str, Any]]:
-    """Resolve --template / --schema into a NuExtract template dict, or None."""
-    if template_arg and schema_arg:
-        raise ValueError("--template and --schema are mutually exclusive.")
-
-    if template_arg is not None:
-        return load_template_arg(template_arg)
-
-    if schema_arg is not None:
-        schema = load_template_arg(schema_arg)
-        try:
-            from numind.nuextract_utils import convert_json_schema_to_nuextract_template
-        except ImportError as e:
-            raise RuntimeError(
-                "--schema requires the `numind` package. "
-                "It should be listed in this script's PEP 723 dependencies."
-            ) from e
-        template, dropped = convert_json_schema_to_nuextract_template(schema)
-        if dropped:
-            logger.warning(
-                f"numind dropped {len(dropped)} unsupported branches from the JSON Schema: "
-                f"{dropped}"
-            )
-        return template
-
-    return None
-
-
-def image_to_data_uri(image: Union[Image.Image, Dict[str, Any], str]) -> str:
-    """Normalize an HF dataset image cell to a PNG data URI."""
-    if isinstance(image, Image.Image):
-        pil_img = image
-    elif isinstance(image, dict) and "bytes" in image:
-        pil_img = Image.open(io.BytesIO(image["bytes"]))
-    elif isinstance(image, str):
-        pil_img = Image.open(image)
-    else:
-        raise ValueError(f"Unsupported image type: {type(image)}")
-
-    pil_img = pil_img.convert("RGB")
-    buf = io.BytesIO()
-    pil_img.save(buf, format="PNG")
-    return f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}"
-
-
-def make_message(image: Union[Image.Image, Dict[str, Any], str]) -> List[Dict]:
-    """Build an OpenAI-format chat message containing one image."""
-    data_uri = image_to_data_uri(image)
-    return [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image_url", "image_url": {"url": data_uri}},
-            ],
-        }
-    ]
-
-
-def split_thinking(text: str) -> tuple[Optional[str], str]:
-    """Return (reasoning, answer) if <think>...</think> is present, else (None, text)."""
-    if "<think>" in text and "</think>" in text:
-        reasoning = text.split("<think>", 1)[1].split("</think>", 1)[0].strip()
-        answer = text.split("</think>", 1)[1].strip()
-        return reasoning, answer
-    return None, text.strip()
-
-
-def parse_json_output(text: str) -> tuple[Optional[Any], bool]:
-    """Parse an extraction output; strip ``` fences as the model card describes.
-
-    Returns (parsed_value, parse_error). On failure, parsed_value is None.
-    """
-    stripped = text.strip()
-    if stripped.startswith("```"):
-        stripped = stripped.split("\n", 1)[-1] if "\n" in stripped else stripped[3:]
-        if stripped.endswith("```"):
-            stripped = stripped[:-3].rstrip()
-    try:
-        return json.loads(stripped), False
-    except json.JSONDecodeError:
-        return None, True
-
-
-def create_dataset_card(
-    source_dataset: str,
-    model: str,
-    num_samples: int,
-    processing_time: str,
-    mode_label: str,
-    template: Optional[Dict[str, Any]],
-    enable_thinking: bool,
-    temperature: float,
-    output_column: str,
-    image_column: str,
-    split: str,
-) -> str:
-    """Create a dataset card documenting the NuExtract3 run."""
-    model_name = model.split("/")[-1]
-    template_block = ""
-    if template is not None:
-        template_block = (
-            "\n### Extraction Template\n\n```json\n"
-            + json.dumps(template, indent=2)
-            + "\n```\n"
-        )
-
-    return f"""---
-tags:
-- ocr
-- structured-extraction
-- document-processing
-- nuextract3
-- markdown
-- uv-script
-- generated
----
-
-# {model_name} on {source_dataset}
-
-This dataset contains outputs from [{source_dataset}](https://huggingface.co/datasets/{source_dataset}) processed with [NuExtract3](https://huggingface.co/{model}), a 4B vision-language model for document understanding.
-
-## Processing Details
-
-- **Source Dataset**: [{source_dataset}](https://huggingface.co/datasets/{source_dataset})
-- **Model**: [{model}](https://huggingface.co/{model})
-- **Mode**: {mode_label}
-- **Number of Samples**: {num_samples:,}
-- **Processing Time**: {processing_time}
-- **Processing Date**: {datetime.now().strftime("%Y-%m-%d %H:%M UTC")}
-
-### Configuration
-
-- **Image Column**: `{image_column}`
-- **Output Column**: `{output_column}`
-- **Dataset Split**: `{split}`
-- **Temperature**: {temperature}
-- **Thinking Mode**: {"enabled" if enable_thinking else "disabled"}
-{template_block}
-## Dataset Structure
-
-Original columns plus:
-- `{output_column}`: NuExtract3 output ({"JSON string" if template else "markdown"})
-- `inference_info`: JSON list tracking models applied to this dataset
-{"- `" + output_column + "_reasoning`: model's thinking trace (when enabled)" if enable_thinking else ""}
-
-Generated with [UV Scripts](https://huggingface.co/uv-scripts)
-"""
-
-
-def main(
-    input_dataset: str,
-    output_dataset: str,
-    image_column: str = "image",
-    batch_size: int = 16,
-    max_model_len: int = 16384,
-    max_tokens: int = 8192,
-    gpu_memory_utilization: float = 0.8,
-    mode: str = "markdown",
-    template_arg: Optional[str] = None,
-    schema_arg: Optional[str] = None,
-    enable_thinking: bool = False,
-    instructions: Optional[str] = None,
-    temperature: Optional[float] = None,
-    model: str = MODEL_DEFAULT,
-    hf_token: str = None,
-    split: str = "train",
-    max_samples: int = None,
-    private: bool = False,
-    shuffle: bool = False,
-    seed: int = 42,
-    output_column: Optional[str] = None,
-    overwrite: bool = False,
-    verbose: bool = False,
-    config: str = None,
-    create_pr: bool = False,
-):
-    """Process images from an HF dataset through NuExtract3."""
-
-    check_cuda_availability()
-    start_time = datetime.now()
-
-    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")
-    if HF_TOKEN:
-        login(token=HF_TOKEN)
-
-    template = resolve_template(template_arg, schema_arg)
-    extraction_mode = template is not None
-    mode_label = "structured-extraction" if extraction_mode else mode
-
-    if output_column is None:
-        output_column = "extraction" if extraction_mode else "markdown"
-
-    if temperature is None:
-        temperature = 0.6 if enable_thinking else 0.2
-
-    logger.info(f"Using model: {model}")
-    logger.info(f"Mode: {mode_label}")
-    logger.info(f"Thinking: {enable_thinking}  Temperature: {temperature}")
-    if extraction_mode:
-        logger.info(f"Template: {json.dumps(template, indent=2)}")
-
-    logger.info(f"Loading dataset: {input_dataset}")
-    dataset = load_dataset(input_dataset, split=split)
-
-    if image_column not in dataset.column_names:
-        raise ValueError(
-            f"Column '{image_column}' not found. Available: {dataset.column_names}"
-        )
-
-    # Fail fast if an output column would collide with an existing input column.
-    # With --enable-thinking the script also writes "{output_column}_reasoning".
-    out_cols = [output_column]
-    if enable_thinking:
-        out_cols.append(f"{output_column}_reasoning")
-    dataset = ensure_output_columns_free(dataset, out_cols, overwrite=overwrite)
-
-    if shuffle:
-        logger.info(f"Shuffling dataset with seed {seed}")
-        dataset = dataset.shuffle(seed=seed)
-
-    if max_samples:
-        dataset = dataset.select(range(min(max_samples, len(dataset))))
-        logger.info(f"Limited to {len(dataset)} samples")
-
-    logger.info("Initializing vLLM with NuExtract3")
-    logger.info("This may take a few minutes on first run...")
-    llm = LLM(
-        model=model,
-        trust_remote_code=True,
-        max_model_len=max_model_len,
-        gpu_memory_utilization=gpu_memory_utilization,
-        limit_mm_per_prompt={"image": 1},
-    )
-
-    sampling_params = SamplingParams(
-        temperature=temperature,
-        max_tokens=max_tokens,
-    )
-
-    chat_template_kwargs: Dict[str, Any] = {"enable_thinking": enable_thinking}
-    if extraction_mode:
-        chat_template_kwargs["template"] = json.dumps(template, indent=4)
-    else:
-        chat_template_kwargs["mode"] = mode
-    if instructions:
-        chat_template_kwargs["instructions"] = instructions
-
-    logger.info(f"Processing {len(dataset)} images in batches of {batch_size}")
-    logger.info(f"Output will be written to column: {output_column}")
-
-    all_outputs: List[str] = []
-    all_reasoning: List[Optional[str]] = []
-    all_parse_errors: List[bool] = []
-    total_batches = (len(dataset) + batch_size - 1) // batch_size
-    processed = 0
-
-    for batch_num, batch_indices in enumerate(
-        partition_all(batch_size, range(len(dataset))), 1
-    ):
-        batch_indices = list(batch_indices)
-        batch_images = [dataset[i][image_column] for i in batch_indices]
-
-        logger.info(
-            f"Batch {batch_num}/{total_batches} "
-            f"({processed}/{len(dataset)} images done)"
-        )
-
-        try:
-            batch_messages = [make_message(img) for img in batch_images]
-            outputs = llm.chat(
-                batch_messages,
-                sampling_params,
-                chat_template_kwargs=chat_template_kwargs,
-                chat_template_content_format="openai",
-            )
-
-            for output in outputs:
-                raw_text = output.outputs[0].text
-                reasoning, answer = split_thinking(raw_text)
-
-                if extraction_mode:
-                    parsed, parse_error = parse_json_output(answer)
-                    stored = (
-                        json.dumps(parsed, ensure_ascii=False)
-                        if parsed is not None
-                        else answer
-                    )
-                    all_outputs.append(stored)
-                    all_parse_errors.append(parse_error)
-                else:
-                    all_outputs.append(answer)
-                    all_parse_errors.append(False)
-
-                all_reasoning.append(reasoning)
-
-            processed += len(batch_images)
-
-        except Exception as e:
-            logger.error(f"Error processing batch: {e}")
-            all_outputs.extend(["[NUEXTRACT3 ERROR]"] * len(batch_images))
-            all_reasoning.extend([None] * len(batch_images))
-            all_parse_errors.extend([True] * len(batch_images))
-            processed += len(batch_images)
-
-    processing_duration = datetime.now() - start_time
-    processing_time_str = f"{processing_duration.total_seconds() / 60:.1f} min"
-
-    logger.info(f"Adding '{output_column}' column to dataset")
-    dataset = dataset.add_column(output_column, all_outputs)
-
-    if enable_thinking and any(r is not None for r in all_reasoning):
-        reasoning_col = f"{output_column}_reasoning"
-        logger.info(f"Adding '{reasoning_col}' column to dataset")
-        dataset = dataset.add_column(reasoning_col, all_reasoning)
-
-    if extraction_mode:
-        parse_error_count = sum(all_parse_errors)
-        if parse_error_count:
-            logger.warning(
-                f"{parse_error_count}/{len(all_parse_errors)} extractions failed to parse as JSON"
-            )
-
-    inference_entry = {
-        "model_id": model,
-        "model_name": MODEL_NAME,
-        "column_name": output_column,
-        "timestamp": datetime.now().isoformat(),
-        "mode": mode_label,
-        "has_template": extraction_mode,
-        "enable_thinking": enable_thinking,
-        "temperature": temperature,
-        "max_tokens": max_tokens,
-    }
-    if extraction_mode:
-        inference_entry["parse_error_rate"] = (
-            sum(all_parse_errors) / len(all_parse_errors) if all_parse_errors else 0.0
-        )
-
-    if "inference_info" in dataset.column_names:
-        logger.info("Updating existing inference_info column")
-
-        def update_inference_info(example):
-            try:
-                existing_info = (
-                    json.loads(example["inference_info"])
-                    if example["inference_info"]
-                    else []
-                )
-            except (json.JSONDecodeError, TypeError):
-                existing_info = []
-            existing_info.append(inference_entry)
-            return {"inference_info": json.dumps(existing_info)}
-
-        dataset = dataset.map(update_inference_info)
-    else:
-        logger.info("Creating new inference_info column")
-        inference_list = [json.dumps([inference_entry])] * len(dataset)
-        dataset = dataset.add_column("inference_info", inference_list)
-
-    logger.info(f"Pushing to {output_dataset}")
-    max_retries = 3
-    for attempt in range(1, max_retries + 1):
-        try:
-            if attempt > 1:
-                logger.warning("Disabling XET (fallback to HTTP upload)")
-                os.environ["HF_HUB_DISABLE_XET"] = "1"
-            dataset.push_to_hub(
-                output_dataset,
-                private=private,
-                token=HF_TOKEN,
-                max_shard_size="500MB",
-                **({"config_name": config} if config else {}),
-                create_pr=create_pr,
-                commit_message=f"Add {model} {mode_label} results ({len(dataset)} samples)"
-                + (f" [{config}]" if config else ""),
-            )
-            break
-        except Exception as e:
-            logger.error(f"Upload attempt {attempt}/{max_retries} failed: {e}")
-            if attempt < max_retries:
-                delay = 30 * (2 ** (attempt - 1))
-                logger.info(f"Retrying in {delay}s...")
-                time.sleep(delay)
-            else:
-                logger.error("All upload attempts failed. Results are lost.")
-                sys.exit(1)
-
-    logger.info("Creating dataset card")
-    card_content = create_dataset_card(
-        source_dataset=input_dataset,
-        model=model,
-        num_samples=len(dataset),
-        processing_time=processing_time_str,
-        mode_label=mode_label,
-        template=template,
-        enable_thinking=enable_thinking,
-        temperature=temperature,
-        output_column=output_column,
-        image_column=image_column,
-        split=split,
-    )
-    card = DatasetCard(card_content)
-    card.push_to_hub(output_dataset, token=HF_TOKEN)
-
-    logger.info("Done! NuExtract3 processing complete.")
-    logger.info(
-        f"Dataset available at: https://huggingface.co/datasets/{output_dataset}"
-    )
-    logger.info(f"Processing time: {processing_time_str}")
-    logger.info(
-        f"Processing speed: {len(dataset) / processing_duration.total_seconds():.2f} images/sec"
-    )
-
-    if verbose:
-        import importlib.metadata
-
-        logger.info("--- Resolved package versions ---")
-        for pkg in [
-            "vllm",
-            "transformers",
-            "torch",
-            "datasets",
-            "pyarrow",
-            "pillow",
-            "numind",
-        ]:
-            try:
-                logger.info(f"  {pkg}=={importlib.metadata.version(pkg)}")
-            except importlib.metadata.PackageNotFoundError:
-                logger.info(f"  {pkg}: not installed")
-        logger.info("--- End versions ---")
-
-
-if __name__ == "__main__":
-    if len(sys.argv) == 1:
-        print("=" * 70)
-        print("NuExtract3 - Document-to-Markdown + Structured Extraction (4B)")
-        print("=" * 70)
-        print("\nModes:")
-        print("  markdown          - Image -> markdown (default)")
-        print("  content           - Image -> plain content")
-        print("  --template / --schema  - Image -> JSON shaped like the template")
-        print("\nExamples:")
-        print("\n1. Markdown OCR:")
-        print("   uv run nuextract3.py input-dataset output-dataset")
-        print("\n2. Structured extraction with an inline template:")
-        print("   uv run nuextract3.py input output \\")
-        print('     --template \'{"title": "verbatim-string", "date": "date"}\'')
-        print("\n3. Structured extraction from a JSON Schema (e.g. Pydantic):")
-        print("   uv run nuextract3.py input output --schema schema.json")
-        print("\n   (--template / --schema also accept a URL or a local file path)")
-        print("\n4. Reasoning mode for harder documents:")
-        print("   uv run nuextract3.py input output --enable-thinking")
-        print("\n5. Test with 10 samples:")
-        print("   uv run nuextract3.py large-ds test --max-samples 10 --shuffle")
-        print("\n6. Running on HF Jobs (use vllm/vllm-openai image for built kernels):")
-        print("   hf jobs uv run --flavor a100-large \\")
-        print("     --image vllm/vllm-openai:latest \\")
-        print("     --python /usr/bin/python3 \\")
-        print("     -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages \\")
-        print("     -s HF_TOKEN \\")
-        print(
-            "     https://huggingface.co/datasets/uv-scripts/ocr/raw/main/nuextract3.py \\"
-        )
-        print("       input-dataset output-dataset --batch-size 16")
-        print("\nFor full help: uv run nuextract3.py --help")
-        sys.exit(0)
-
-    parser = argparse.ArgumentParser(
-        description="NuExtract3: document-to-markdown + schema-guided JSON extraction (4B VLM)",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog="""
-Modes:
-  (default)     Markdown OCR (image -> clean markdown)
-  --mode content
-                Plain-content extraction (less structured than markdown)
-  --template PATH_OR_JSON
-                Structured extraction with a NuExtract template
-  --schema PATH_OR_JSON
-                Structured extraction from a JSON Schema
-                (e.g. Pydantic Model.model_json_schema())
-
-Examples:
-  uv run nuextract3.py my-docs analyzed-docs
-  uv run nuextract3.py receipts extracted \\
-      --template '{"store": "verbatim-string", "total": "number"}'
-  uv run nuextract3.py contracts extracted --schema contract_schema.json
-  uv run nuextract3.py hard-docs out --enable-thinking
-        """,
-    )
-
-    parser.add_argument("input_dataset", help="Input dataset ID from Hugging Face Hub")
-    parser.add_argument("output_dataset", help="Output dataset ID for Hugging Face Hub")
-    parser.add_argument(
-        "--image-column",
-        default="image",
-        help="Column containing images (default: image)",
-    )
-    parser.add_argument(
-        "--batch-size",
-        type=int,
-        default=16,
-        help="Batch size for processing (default: 16)",
-    )
-    parser.add_argument(
-        "--max-model-len",
-        type=int,
-        default=16384,
-        help="Maximum model context length (default: 16384)",
-    )
-    parser.add_argument(
-        "--max-tokens",
-        type=int,
-        default=8192,
-        help="Maximum tokens to generate (default: 8192)",
-    )
-    parser.add_argument(
-        "--gpu-memory-utilization",
-        type=float,
-        default=0.8,
-        help="GPU memory utilization (default: 0.8)",
-    )
-    parser.add_argument(
-        "--mode",
-        choices=["markdown", "content"],
-        default="markdown",
-        help="OCR mode when no template/schema is given (default: markdown)",
-    )
-    parser.add_argument(
-        "--template",
-        help="NuExtract template: inline JSON, a URL, or a file path",
-    )
-    parser.add_argument(
-        "--schema",
-        help="JSON Schema to auto-convert: inline JSON, a URL, or a file path",
-    )
-    parser.add_argument(
-        "--enable-thinking",
-        action="store_true",
-        help="Enable reasoning mode (slower, better on hard documents)",
-    )
-    parser.add_argument(
-        "--instructions",
-        default=None,
-        help=(
-            "Free-text instructions passed to NuExtract via "
-            "chat_template_kwargs.instructions (e.g. routing guidance across "
-            "optional schema branches, output-format rules). "
-            "See https://huggingface.co/numind/NuExtract3#instructions"
-        ),
-    )
-    parser.add_argument(
-        "--temperature",
-        type=float,
-        default=None,
-        help="Sampling temperature (default: 0.2 non-thinking, 0.6 thinking)",
-    )
-    parser.add_argument(
-        "--model",
-        default=MODEL_DEFAULT,
-        help=f"Model ID (default: {MODEL_DEFAULT})",
-    )
-    parser.add_argument("--hf-token", help="Hugging Face API token")
-    parser.add_argument(
-        "--split", default="train", help="Dataset split to use (default: train)"
-    )
-    parser.add_argument(
-        "--max-samples",
-        type=int,
-        help="Maximum number of samples to process (for testing)",
-    )
-    parser.add_argument(
-        "--private", action="store_true", help="Make output dataset private"
-    )
-    parser.add_argument(
-        "--config",
-        help="Config/subset name when pushing to Hub (for benchmarking multiple models in one repo)",
-    )
-    parser.add_argument(
-        "--create-pr",
-        action="store_true",
-        help="Create a pull request instead of pushing directly (for parallel benchmarking)",
-    )
-    parser.add_argument(
-        "--shuffle", action="store_true", help="Shuffle dataset before processing"
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=42,
-        help="Random seed for shuffling (default: 42)",
-    )
-    parser.add_argument(
-        "--output-column",
-        default=None,
-        help="Column name for output (default: 'markdown' in OCR mode, 'extraction' in template mode)",
-    )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
-    parser.add_argument(
-        "--verbose",
-        action="store_true",
-        help="Log resolved package versions after processing",
-    )
-
-    args = parser.parse_args()
-
-    main(
-        input_dataset=args.input_dataset,
-        output_dataset=args.output_dataset,
-        image_column=args.image_column,
-        batch_size=args.batch_size,
-        max_model_len=args.max_model_len,
-        max_tokens=args.max_tokens,
-        gpu_memory_utilization=args.gpu_memory_utilization,
-        mode=args.mode,
-        template_arg=args.template,
-        schema_arg=args.schema,
-        enable_thinking=args.enable_thinking,
-        instructions=args.instructions,
-        temperature=args.temperature,
-        model=args.model,
-        hf_token=args.hf_token,
-        split=args.split,
-        max_samples=args.max_samples,
-        private=args.private,
-        shuffle=args.shuffle,
-        seed=args.seed,
-        output_column=args.output_column,
-        overwrite=args.overwrite,
-        verbose=args.verbose,
-        config=args.config,
-        create_pr=args.create_pr,
-    )

numarkdown-ocr.py

@@ -47,10 +47,6 @@ from huggingface_hub import DatasetCard, login
 from PIL import Image
 from toolz import partition_all
 from tqdm.auto import tqdm
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
 from vllm import LLM, SamplingParams
 
 logging.basicConfig(level=logging.INFO)
@@ -75,28 +71,6 @@ def check_gpu_availability() -> int:
     return num_gpus
 
 
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
 def validate_and_resize_image(
     image: Image.Image,
     min_pixels: int = 100 * 28 * 28,
@@ -339,7 +313,6 @@ def main(
     temperature: float = 0.0,
     custom_prompt: Optional[str] = None,
     output_column: str = "markdown",
-    overwrite: bool = False,
     config: str = None,
     create_pr: bool = False,
     verbose: bool = False,
@@ -384,9 +357,6 @@ def main(
             f"Column '{image_column}' not found. Available: {dataset.column_names}"
         )
 
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
     # Shuffle if requested
     if shuffle:
         logger.info(f"Shuffling dataset with seed {seed}")
@@ -726,12 +696,6 @@ Examples:
         default="markdown",
         help="Column name for output text (default: markdown)",
     )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
     parser.add_argument(
         "--config",
         help="Config/subset name when pushing to Hub (for benchmarking multiple models in one repo)",
@@ -769,7 +733,6 @@ Examples:
         temperature=args.temperature,
         custom_prompt=args.custom_prompt,
         output_column=args.output_column,
-        overwrite=args.overwrite,
         config=args.config,
         create_pr=args.create_pr,
         verbose=args.verbose,

olmocr2-vllm.py

@@ -1,7 +1,7 @@
 # /// script
 # requires-python = ">=3.11"
 # dependencies = [
-#     "datasets>=4.0.0",
+#     "datasets",
 #     "huggingface-hub",
 #     "pillow",
 #     "vllm",
@@ -50,11 +50,8 @@ from huggingface_hub import DatasetCard, login
 from PIL import Image
 from toolz import partition_all
 from tqdm.auto import tqdm
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
 from vllm import LLM, SamplingParams
+from vllm.sampling_params import GuidedDecodingParams
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -81,28 +78,6 @@ def check_cuda_availability():
         logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
 
 
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
 def parse_yaml_frontmatter(text: str) -> tuple[dict, str]:
     """
     Parse YAML front matter from olmOCR output.
@@ -302,7 +277,6 @@ def main(
     output_dataset: str,
     image_column: str = "image",
     output_column: str = "markdown",
-    overwrite: bool = False,
     batch_size: int = 16,
     model: str = "allenai/olmOCR-2-7B-1025-FP8",
     max_model_len: int = 16384,
@@ -316,8 +290,6 @@ def main(
     private: bool = False,
     shuffle: bool = False,
     seed: int = 42,
-    config: str = None,
-    create_pr: bool = False,
 ):
     """
     Process a dataset of document images through olmOCR-2 to extract markdown.
@@ -358,9 +330,6 @@ def main(
     logger.info(f"Loading dataset: {input_dataset}")
     ds = load_dataset(input_dataset, split=split)
 
-    # Fail fast if the output column would collide with an existing input column
-    ds = ensure_output_columns_free(ds, [output_column], overwrite=overwrite)
-
     # Shuffle if requested
     if shuffle:
         logger.info(f"Shuffling dataset with seed {seed}")
@@ -396,25 +365,13 @@ def main(
         "stop": ["<|im_end|>", "<|endoftext|>"],
     }
 
-    # Add guided decoding if requested (enforces YAML front matter structure).
-    # vLLM 0.22.x renamed this API: GuidedDecodingParams (guided_decoding=) →
-    # StructuredOutputsParams (structured_outputs=). Import lazily + shim so the
-    # default path (guided_decoding=False) never touches the moved symbol.
+    # Add guided decoding if requested (enforces YAML front matter structure)
     if guided_decoding:
         logger.info("Enabling guided decoding with YAML front matter regex")
-        front_matter_regex = r"---\nprimary_language: (?:[a-z]{2}|null)\nis_rotation_valid: (?:True|False|true|false)\nrotation_correction: (?:0|90|180|270)\nis_table: (?:True|False|true|false)\nis_diagram: (?:True|False|true|false)\n(?:---|---\n[\s\S]+)"
-        try:
-            from vllm.sampling_params import StructuredOutputsParams
-
-            sampling_params_kwargs["structured_outputs"] = StructuredOutputsParams(
-                regex=front_matter_regex
-            )
-        except ImportError:
-            from vllm.sampling_params import GuidedDecodingParams
-
-            sampling_params_kwargs["guided_decoding"] = GuidedDecodingParams(
-                regex=front_matter_regex
-            )
+        guided_params = GuidedDecodingParams(
+            regex=r"---\nprimary_language: (?:[a-z]{2}|null)\nis_rotation_valid: (?:True|False|true|false)\nrotation_correction: (?:0|90|180|270)\nis_table: (?:True|False|true|false)\nis_diagram: (?:True|False|true|false)\n(?:---|---\n[\s\S]+)"
+        )
+        sampling_params_kwargs["guided_decoding"] = guided_params
 
     sampling_params = SamplingParams(**sampling_params_kwargs)
 
@@ -463,41 +420,36 @@ def main(
         ds = ds.remove_columns([metadata_column_name])
     ds = ds.add_column(metadata_column_name, all_metadata)
 
-    # Add inference information — standard schema: a JSON-string LIST of entries, each carrying
-    # model_id (used as the leaderboard label) and column_name (the output text column). The
-    # script's own extras (version, timestamp, sampling params) ride along as extra keys.
-    inference_entry = {
-        "model_id": model,
-        "column_name": output_column,
-        "script": "olmocr2-vllm.py",
-        "version": "1.0.0",
-        "timestamp": datetime.now().isoformat(),
-        "batch_size": batch_size,
-        "max_tokens": max_tokens,
-        "temperature": temperature,
-    }
+    # Add inference information
+    inference_info = json.dumps(
+        {
+            "model": model,
+            "script": "olmocr2-vllm.py",
+            "version": "1.0.0",
+            "timestamp": datetime.now().isoformat(),
+            "batch_size": batch_size,
+            "max_tokens": max_tokens,
+            "temperature": temperature,
+        }
+    )
 
-    # Handle existing inference_info column (append to the list; create a fresh list otherwise)
+    # Handle existing inference_info column
     if inference_info_column in ds.column_names:
-        logger.info("Updating existing inference_info column")
-
+        # Parse existing, append new model info
         def update_inference_info(example):
             try:
-                existing_info = (
-                    json.loads(example[inference_info_column])
-                    if example[inference_info_column]
-                    else []
-                )
-            except (json.JSONDecodeError, TypeError):
-                existing_info = []
-            existing_info.append(inference_entry)
-            return {inference_info_column: json.dumps(existing_info)}
+                existing = json.loads(example[inference_info_column])
+                if not isinstance(existing, list):
+                    existing = [existing]
+            except (json.JSONDecodeError, KeyError):
+                existing = []
+
+            existing.append(json.loads(inference_info))
+            return {inference_info_column: json.dumps(existing)}
 
         ds = ds.map(update_inference_info)
     else:
-        logger.info("Creating new inference_info column")
-        inference_list = [json.dumps([inference_entry])] * len(ds)
-        ds = ds.add_column(inference_info_column, inference_list)
+        ds = ds.add_column(inference_info_column, [inference_info] * len(ds))
 
     # Calculate processing time
     elapsed_time = time.time() - start_time
@@ -522,13 +474,9 @@ def main(
 
     # Push to hub
     logger.info(f"Pushing to HuggingFace Hub: {output_dataset}")
-    commit_info = ds.push_to_hub(
+    ds.push_to_hub(
         output_dataset,
         private=private,
-        **({"config_name": config} if config else {}),
-        create_pr=create_pr,
-        commit_message=f"Add {model} OCR results ({len(ds)} samples)"
-        + (f" [{config}]" if config else ""),
     )
 
     # Update dataset card
@@ -539,8 +487,6 @@ def main(
     logger.info(f"✓ Dataset: https://huggingface.co/datasets/{output_dataset}")
     logger.info(f"✓ Processing time: {processing_time}")
     logger.info(f"✓ Samples processed: {len(ds):,}")
-    if create_pr and getattr(commit_info, "pr_url", None):
-        logger.info(f"✓ Pull request created: {commit_info.pr_url}")
 
 
 if __name__ == "__main__":
@@ -597,12 +543,6 @@ Examples:
         default="markdown",
         help="Column name for markdown output (default: markdown)",
     )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
     parser.add_argument(
         "--batch-size",
         type=int,
@@ -662,15 +602,6 @@ Examples:
         action="store_true",
         help="Make output dataset private",
     )
-    parser.add_argument(
-        "--config",
-        help="Config/subset name when pushing to Hub (for benchmarking multiple models in one repo)",
-    )
-    parser.add_argument(
-        "--create-pr",
-        action="store_true",
-        help="Create a pull request instead of pushing directly (for parallel benchmarking)",
-    )
     parser.add_argument(
         "--shuffle",
         action="store_true",
@@ -689,7 +620,6 @@ Examples:
         output_dataset=args.output_dataset,
         image_column=args.image_column,
         output_column=args.output_column,
-        overwrite=args.overwrite,
         batch_size=args.batch_size,
         model=args.model,
         max_model_len=args.max_model_len,
@@ -703,6 +633,4 @@ Examples:
         private=args.private,
         shuffle=args.shuffle,
         seed=args.seed,
-        config=args.config,
-        create_pr=args.create_pr,
     )

ovis-ocr2-server.py

@@ -1,725 +0,0 @@
-# /// script
-# requires-python = ">=3.11"
-# dependencies = [
-#     "datasets>=4.0.0",
-#     "huggingface-hub",
-#     "pillow",
-#     "requests",
-# ]
-# ///
-
-"""
-Convert document images to markdown using OvisOCR2 via an in-job vLLM server.
-
-Same model, prompt, and outputs as ovis-ocr2.py, but serves the model behind
-`vllm serve` inside the job and posts images concurrently — continuous batching
-stays fed instead of draining at each offline `llm.generate()` batch barrier,
-and a bad image fails one request instead of a whole batch of 16.
-
-This script is the *driver* half: it expects the server on localhost (started
-by the job command below), loads the input dataset, posts images concurrently,
-postprocesses (repeat-trim + img-tag filter, as in ovis-ocr2.py), and pushes
-the result dataset. The driver has no torch/vllm deps, so `uv run` starts in
-seconds while the server warms up in parallel.
-
-NOTE: OvisOCR2's card documents offline vLLM only (checked 2026-07-16) — this
-server translation is ours, not the authors'. The serve flags below mirror the
-card's offline args; treat A/B output parity with ovis-ocr2.py as part of any
-benchmark run (same --max-samples slice, diff the markdown columns).
-
-Run on HF Jobs (standard uv-run shape — the script starts `vllm serve` itself
-as a subprocess when no server is already reachable; the only thing to get
-right is the --image flag, which provides the `vllm` binary):
-
-  hf jobs uv run --detach --flavor l4x1 -s HF_TOKEN --timeout 4h \\
-      --image vllm/vllm-openai:v0.22.1 \\
-      https://huggingface.co/datasets/uv-scripts/ocr/raw/main/ovis-ocr2-server.py \\
-      <input-dataset> <output-dataset>
-
-To use an already-running or remote endpoint instead, pass --server URL — the
-script only spawns a server when the (default localhost) URL is unreachable.
-The serve flags live in SERVE_ARGS below, the single source of truth.
-
-Serve-flag provenance (the card only shows offline `LLM(...)` args):
-- `--no-enable-prefix-caching --mm-processor-cache-gb 0`: the recurring official
-  OCR-serving pattern (DeepSeek-OCR vLLM recipe, LightOnOCR, PaddleOCR-VL recipe)
-  — OCR never reuses images, the caches only cost memory.
-- `--mm-processor-kwargs`: server-side equivalent of the card's per-request
-  `images_kwargs` pixel bounds. The driver ALSO downscales oversized images
-  client-side to the same max_pixels, so outputs match even if the server flag
-  is dropped.
-- The card's offline `gdn_prefill_backend="triton"` (JIT/nvcc workaround for the
-  bare uv image) is NOT needed here: the vllm-openai image ships the full CUDA
-  toolchain. Add `--gdn-prefill-backend triton` to the serve command if the
-  default backend misbehaves.
-- `enable_thinking=False` (card-mandated, Qwen3.5 templates can inject a
-  thinking preamble) is passed per-request via `chat_template_kwargs`.
-
-Model: ATH-MaaS/OvisOCR2 (0.9B, Apache-2.0, 96.58 OmniDocBench v1.6)
-vLLM: stock Qwen3_5ForConditionalGeneration arch, needs vllm >= 0.22.1.
-"""
-
-import argparse
-import atexit
-import base64
-import concurrent.futures
-import io
-import json
-import logging
-import math
-import os
-import shutil
-import subprocess
-import sys
-import threading
-import time
-from datetime import datetime
-from typing import Any, Dict, Union
-from urllib.parse import urlparse
-
-import requests
-from datasets import load_dataset
-from huggingface_hub import DatasetCard, login
-from PIL import Image
-
-logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
-logger = logging.getLogger(__name__)
-
-MODEL = "ATH-MaaS/OvisOCR2"
-
-# Fixed instruction prompt, verbatim from the model card (including the leading newline).
-# The card warns outputs are tuned to this exact wording — don't "improve" it.
-OCR_PROMPT = (
-    "\nExtract all readable content from the image in natural human reading order "
-    "and output the result as a single Markdown document. For charts or images, "
-    'represent them using an HTML image tag: <img src="images/bbox_{left}_{top}_'
-    '{right}_{bottom}.jpg" />, where left, top, right, bottom are bounding box '
-    "coordinates scaled to [0, 1000). Format formulas as LaTeX. Format tables as "
-    "HTML: <table>...</table>. Transcribe all other text as standard Markdown. "
-    "Preserve the original text without translation or paraphrasing."
-)
-
-# Image bounds from the model card's parser.
-DEFAULT_MIN_PIXELS = 448 * 448  # 200,704
-DEFAULT_MAX_PIXELS = 2880 * 2880  # 8,294,400
-
-# The serve command this script spawns when no server is reachable — single
-# source of truth for the serving configuration (see docstring for provenance).
-SERVE_ARGS = [
-    "vllm", "serve", MODEL,
-    "--max-model-len", "32768",
-    "--gpu-memory-utilization", "0.85",
-    "--limit-mm-per-prompt", '{"image": 1}',
-    "--no-enable-prefix-caching",
-    "--mm-processor-cache-gb", "0",
-    "--mm-processor-kwargs",
-    f'{{"images_kwargs": {{"min_pixels": {DEFAULT_MIN_PIXELS}, "max_pixels": {DEFAULT_MAX_PIXELS}}}}}',
-    "--port", "8000",
-]
-
-RUN_COMMAND = (
-    "hf jobs uv run --detach --flavor l4x1 -s HF_TOKEN --timeout 4h \\\n"
-    "    --image vllm/vllm-openai:v0.22.1 \\\n"
-    "    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/ovis-ocr2-server.py \\\n"
-    "    <input-dataset> <output-dataset>"
-)
-
-
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
-def to_pil_image(image: Union[Image.Image, Dict[str, Any], str]) -> Image.Image:
-    """Convert a dataset image cell (PIL image, bytes dict, or path) to RGB PIL."""
-    if isinstance(image, Image.Image):
-        pil_img = image
-    elif isinstance(image, dict) and "bytes" in image:
-        pil_img = Image.open(io.BytesIO(image["bytes"]))
-    elif isinstance(image, str):
-        pil_img = Image.open(image)
-    else:
-        raise ValueError(f"Unsupported image type: {type(image)}")
-    return pil_img.convert("RGB")
-
-
-def encode_image(image, max_pixels: int) -> str:
-    """RGB-convert, downscale to max_pixels if oversized, return base64 JPEG.
-
-    The server's processor clamps to the same bound, so this only changes where
-    the downscale happens — doing it client-side shrinks the request payload
-    (a 30MP scan is ~4x the bytes of its 8.3MP clamp) and keeps outputs
-    identical even if the serve command omits --mm-processor-kwargs.
-    """
-    img = to_pil_image(image)
-    w, h = img.size
-    if w * h > max_pixels:
-        scale = math.sqrt(max_pixels / (w * h))
-        img = img.resize((max(1, int(w * scale)), max(1, int(h * scale))), Image.LANCZOS)
-    buf = io.BytesIO()
-    img.save(buf, format="JPEG", quality=95)
-    return base64.b64encode(buf.getvalue()).decode()
-
-
-def clean_truncated_repeats(
-    text: str,
-    min_text_len: int = 8000,
-    max_period: int = 200,
-    min_period: int = 1,
-    min_repeat_chars: int = 100,
-    min_repeat_times: int = 5,
-) -> str:
-    """Trim degenerate trailing repetition (verbatim port of the model card's cleanup).
-
-    Long outputs that hit max_tokens can end in a repeated unit (a char, phrase, or
-    table row); this detects the shortest repeating tail unit and keeps one copy.
-    """
-    n = len(text)
-    if n < min_text_len:
-        return text
-
-    max_period = min(max_period, n - 1)
-    for unit_len in range(min_period, max_period + 1):
-        if text[n - 1] != text[n - 1 - unit_len]:
-            continue
-
-        match_len = 1
-        idx = n - 2
-        while idx >= unit_len and text[idx] == text[idx - unit_len]:
-            match_len += 1
-            idx -= 1
-
-        total_len = match_len + unit_len
-        repeat_times = total_len // unit_len
-        tail_len = total_len % unit_len
-
-        if repeat_times >= min_repeat_times and total_len >= min_repeat_chars:
-            return text[: n - total_len + unit_len] + text[n - tail_len :]
-
-    return text
-
-
-def filter_image_tags(text: str) -> str:
-    """Drop visual-region <img> blocks (upstream parser's default behaviour)."""
-    return "\n\n".join(
-        block
-        for block in text.split("\n\n")
-        if not block.strip().startswith('<img src="images/bbox_')
-    )
-
-
-def postprocess_output(text: str, keep_image_tags: bool) -> str:
-    text = text.strip()
-    if not keep_image_tags:
-        text = filter_image_tags(text)
-    return clean_truncated_repeats(text)
-
-
-def server_alive(server: str) -> bool:
-    try:
-        return requests.get(f"{server}/health", timeout=5).status_code == 200
-    except requests.RequestException:
-        return False
-
-
-def wait_for_server(server: str, timeout_s: int, proc: "subprocess.Popen | None" = None):
-    logger.info(f"Waiting for server at {server}...")
-    deadline = time.time() + timeout_s
-    while time.time() < deadline:
-        if server_alive(server):
-            logger.info("Server is ready")
-            return
-        if proc is not None and proc.poll() is not None:
-            logger.error(f"Spawned vllm serve exited with code {proc.returncode} before becoming ready")
-            sys.exit(1)
-        time.sleep(10)
-    logger.error(f"Server did not become ready within {timeout_s}s")
-    sys.exit(1)
-
-
-def ensure_server(server: str, timeout_s: int = 1800):
-    """Use a reachable server; otherwise spawn `vllm serve` ourselves; else fail fast.
-
-    Spawning is only attempted for a localhost URL — a remote --server that is
-    down is the user's to fix, not ours to shadow with a local model.
-    """
-    if server_alive(server):
-        logger.info(f"Using already-running server at {server}")
-        return
-
-    host = urlparse(server).hostname or ""
-    if host not in ("127.0.0.1", "localhost", "0.0.0.0"):
-        logger.info(f"Remote server {server} not up yet — waiting for it")
-        wait_for_server(server, timeout_s)
-        return
-
-    if shutil.which("vllm") is None:
-        logger.error("No server is running and the `vllm` binary is not on PATH.")
-        logger.error("Run this script on a vLLM image so it can start the server itself:\n")
-        logger.error(RUN_COMMAND)
-        logger.error("\n(or start `vllm serve` yourself / pass --server URL of a running endpoint)")
-        sys.exit(1)
-
-    logger.info(f"Starting server: {' '.join(SERVE_ARGS)}")
-    proc = subprocess.Popen(SERVE_ARGS)  # logs interleave with ours on stdout/stderr
-    atexit.register(proc.terminate)  # don't leave a GPU server behind on local runs
-    wait_for_server(server, timeout_s, proc=proc)
-
-
-def ocr_one(
-    server: str,
-    image,
-    max_pixels: int,
-    max_tokens: int,
-    timeout_s: int,
-    retries: int = 2,
-) -> str:
-    """OCR a single image via the chat completions API. Returns raw model text."""
-    b64 = encode_image(image, max_pixels)
-    payload = {
-        "model": MODEL,
-        "messages": [
-            {
-                "role": "user",
-                "content": [
-                    # Image first, then text — same order the card's chat template uses.
-                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}},
-                    {"type": "text", "text": OCR_PROMPT},
-                ],
-            }
-        ],
-        "temperature": 0.0,
-        "max_tokens": max_tokens,
-        # Card-mandated: Qwen3.5 templates can inject a thinking preamble otherwise.
-        "chat_template_kwargs": {"enable_thinking": False},
-    }
-    last_err = None
-    for attempt in range(retries + 1):
-        try:
-            resp = requests.post(
-                f"{server}/v1/chat/completions", json=payload, timeout=timeout_s
-            )
-            resp.raise_for_status()
-            return resp.json()["choices"][0]["message"]["content"]
-        except Exception as e:
-            last_err = e
-            if attempt < retries:
-                time.sleep(10 * (attempt + 1))
-    raise RuntimeError(f"request failed after {retries + 1} attempts: {last_err}")
-
-
-def create_dataset_card(
-    source_dataset: str,
-    model: str,
-    num_samples: int,
-    num_errors: int,
-    processing_time: str,
-    images_per_sec: float,
-    concurrency: int,
-    max_tokens: int,
-    keep_image_tags: bool,
-    image_column: str = "image",
-    split: str = "train",
-) -> str:
-    """Create a dataset card documenting the OCR process."""
-    model_name = model.split("/")[-1]
-
-    # Canonical provenance stamp (see AGENTS.md): Jobs claim gated on JOB_ID, set by HF Jobs in-container.
-    on_jobs = os.environ.get("JOB_ID") is not None
-    hw = os.environ.get("ACCELERATOR") or ""
-    origin = (
-        "Produced on [Hugging Face Jobs](https://huggingface.co/docs/huggingface_hub/guides/jobs)"
-        + (f" (`{hw}`)" if hw else "")
-    ) if on_jobs else "Generated"
-    jobs_tag = "\n- hf-jobs" if on_jobs else ""
-
-    return f"""---
-tags:
-- ocr
-- document-processing
-- ovis-ocr2
-- markdown
-- uv-script
-- generated{jobs_tag}
----
-
-# Document OCR using {model_name} (server mode)
-
-This dataset contains OCR results from images in [{source_dataset}](https://huggingface.co/datasets/{source_dataset}) using OvisOCR2, a compact 0.9B document parsing model (96.58 on OmniDocBench v1.6), served behind an in-job vLLM server with concurrent requests (continuous batching).
-
-## Processing Details
-
-- **Source Dataset**: [{source_dataset}](https://huggingface.co/datasets/{source_dataset})
-- **Model**: [{model}](https://huggingface.co/{model})
-- **Number of Samples**: {num_samples:,}
-- **Failed Requests**: {num_errors:,} (marked `[OCR ERROR]`)
-- **Processing Time**: {processing_time}
-- **Throughput**: {images_per_sec:.2f} images/sec
-- **Processing Date**: {datetime.now().strftime("%Y-%m-%d %H:%M UTC")}
-
-### Configuration
-
-- **Mode**: vLLM server (`vllm serve`) + concurrent driver, {concurrency} concurrent requests
-- **Image Column**: `{image_column}`
-- **Dataset Split**: `{split}`
-- **Max Output Tokens**: {max_tokens:,}
-- **Temperature**: 0.0 (greedy, per model card)
-- **Visual-region image tags**: {"kept" if keep_image_tags else "filtered (default)"}
-
-## Dataset Structure
-
-The dataset contains all original columns plus:
-- `markdown`: The extracted text in markdown format
-- `inference_info`: JSON list tracking all OCR models applied to this dataset
-
-## Reproduction
-
-{origin} with the [`ovis-ocr2-server.py`](https://huggingface.co/datasets/uv-scripts/ocr/raw/main/ovis-ocr2-server.py) recipe from [uv-scripts](https://huggingface.co/uv-scripts) — see the script docstring for the single `hf jobs run` command that starts the server and driver together. The offline-vLLM sibling recipe is [`ovis-ocr2.py`](https://huggingface.co/datasets/uv-scripts/ocr/raw/main/ovis-ocr2.py).
-"""
-
-
-def main(
-    input_dataset: str,
-    output_dataset: str,
-    image_column: str = "image",
-    server: str = "http://127.0.0.1:8000",
-    concurrency: int = 32,
-    max_tokens: int = 16384,
-    max_pixels: int = DEFAULT_MAX_PIXELS,
-    request_timeout: int = 1800,
-    keep_image_tags: bool = False,
-    hf_token: str = None,
-    split: str = "train",
-    max_samples: int = None,
-    private: bool = False,
-    shuffle: bool = False,
-    seed: int = 42,
-    output_column: str = "markdown",
-    overwrite: bool = False,
-    config: str = None,
-    create_pr: bool = False,
-):
-    """Process images from HF dataset through an OvisOCR2 vLLM server."""
-
-    start_time = datetime.now()
-
-    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")
-    if HF_TOKEN:
-        login(token=HF_TOKEN)
-
-    logger.info(f"Using model: {MODEL} via server {server}")
-
-    logger.info(f"Loading dataset: {input_dataset}")
-    dataset = load_dataset(input_dataset, split=split)
-
-    if image_column not in dataset.column_names:
-        raise ValueError(
-            f"Column '{image_column}' not found. Available: {dataset.column_names}"
-        )
-
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
-    if shuffle:
-        logger.info(f"Shuffling dataset with seed {seed}")
-        dataset = dataset.shuffle(seed=seed)
-
-    if max_samples:
-        dataset = dataset.select(range(min(max_samples, len(dataset))))
-        logger.info(f"Limited to {len(dataset)} samples")
-
-    # Reuse a reachable server, else spawn `vllm serve` (needs the vllm binary,
-    # i.e. a vllm/vllm-openai image), else fail fast with the correct command.
-    ensure_server(server)
-
-    n = len(dataset)
-    logger.info(f"Processing {n} images, concurrency {concurrency}")
-    all_outputs = [None] * n
-    errors = 0
-    done = 0
-    inference_start = time.time()
-    lock = threading.Lock()
-
-    def worker(i: int) -> None:
-        nonlocal errors, done
-        try:
-            text = ocr_one(
-                server,
-                dataset[i][image_column],
-                max_pixels,
-                max_tokens,
-                request_timeout,
-            )
-            all_outputs[i] = postprocess_output(text, keep_image_tags)
-        except Exception as e:
-            logger.error(f"Image {i} failed: {e}")
-            all_outputs[i] = "[OCR ERROR]"
-            with lock:
-                errors += 1
-        with lock:
-            done += 1
-            if done % 25 == 0 or done == n:
-                rate = done / max(time.time() - inference_start, 1e-9)
-                logger.info(f"{done}/{n} done ({rate:.2f} img/s, {errors} errors)")
-
-    with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as pool:
-        list(pool.map(worker, range(n)))
-
-    inference_secs = time.time() - inference_start
-    processing_duration = datetime.now() - start_time
-    processing_time_str = f"{processing_duration.total_seconds() / 60:.1f} min"
-    images_per_sec = n / inference_secs if inference_secs else 0.0
-
-    logger.info(f"Adding '{output_column}' column to dataset")
-    dataset = dataset.add_column(output_column, all_outputs)
-
-    # Inference info tracking
-    inference_entry = {
-        "model_id": MODEL,
-        "model_name": "OvisOCR2",
-        "column_name": output_column,
-        "timestamp": datetime.now().isoformat(),
-        "temperature": 0.0,
-        "max_tokens": max_tokens,
-        "max_pixels": max_pixels,
-        "keep_image_tags": keep_image_tags,
-        "mode": "vllm-server",
-        "concurrency": concurrency,
-    }
-
-    if "inference_info" in dataset.column_names:
-        logger.info("Updating existing inference_info column")
-
-        def update_inference_info(example):
-            try:
-                existing_info = (
-                    json.loads(example["inference_info"])
-                    if example["inference_info"]
-                    else []
-                )
-            except (json.JSONDecodeError, TypeError):
-                existing_info = []
-            existing_info.append(inference_entry)
-            return {"inference_info": json.dumps(existing_info)}
-
-        dataset = dataset.map(update_inference_info)
-    else:
-        logger.info("Creating new inference_info column")
-        inference_list = [json.dumps([inference_entry])] * len(dataset)
-        dataset = dataset.add_column("inference_info", inference_list)
-
-    # Push to hub with retry and XET fallback
-    logger.info(f"Pushing to {output_dataset}")
-    max_retries = 3
-    for attempt in range(1, max_retries + 1):
-        try:
-            if attempt > 1:
-                logger.warning("Disabling XET (fallback to HTTP upload)")
-                os.environ["HF_HUB_DISABLE_XET"] = "1"
-            dataset.push_to_hub(
-                output_dataset,
-                private=private,
-                token=HF_TOKEN,
-                max_shard_size="500MB",
-                **({"config_name": config} if config else {}),
-                create_pr=create_pr,
-                commit_message=f"Add {MODEL} OCR results ({len(dataset)} samples, server mode)"
-                + (f" [{config}]" if config else ""),
-            )
-            break
-        except Exception as e:
-            logger.error(f"Upload attempt {attempt}/{max_retries} failed: {e}")
-            if attempt < max_retries:
-                delay = 30 * (2 ** (attempt - 1))
-                logger.info(f"Retrying in {delay}s...")
-                time.sleep(delay)
-            else:
-                logger.error("All upload attempts failed. OCR results are lost.")
-                sys.exit(1)
-
-    logger.info("Creating dataset card")
-    card_content = create_dataset_card(
-        source_dataset=input_dataset,
-        model=MODEL,
-        num_samples=len(dataset),
-        num_errors=errors,
-        processing_time=processing_time_str,
-        images_per_sec=images_per_sec,
-        concurrency=concurrency,
-        max_tokens=max_tokens,
-        keep_image_tags=keep_image_tags,
-        image_column=image_column,
-        split=split,
-    )
-
-    card = DatasetCard(card_content)
-    card.push_to_hub(output_dataset, token=HF_TOKEN)
-
-    logger.info("Done! OvisOCR2 server-mode processing complete.")
-    logger.info(
-        f"Dataset available at: https://huggingface.co/datasets/{output_dataset}"
-    )
-    logger.info(f"Processing time: {processing_time_str}")
-    logger.info(
-        f"Throughput: {images_per_sec:.2f} images/sec "
-        f"(inference only, excl. dataset load/push; {errors} errors)"
-    )
-
-
-if __name__ == "__main__":
-    if len(sys.argv) == 1:
-        print("=" * 70)
-        print("OvisOCR2 Document Processing (vLLM server mode)")
-        print("=" * 70)
-        print("\nSame model + outputs as ovis-ocr2.py, but drives an in-job")
-        print("`vllm serve` with concurrent requests — no batch barriers,")
-        print("per-image (not per-batch) failure isolation.")
-        print("\nThe server must already be running (the job command starts")
-        print("both — see the module docstring for the full `hf jobs run`).")
-        print("\nExamples:")
-        print("\n1. Basic OCR (server on localhost:8000):")
-        print("   uv run ovis-ocr2-server.py input-dataset output-dataset")
-        print("\n2. Test with a small sample:")
-        print("   uv run ovis-ocr2-server.py large-dataset test --max-samples 10 --shuffle")
-        print("\n3. Throughput A/B vs the offline recipe:")
-        print("   run both scripts on the same --max-samples slice and compare")
-        print("   the images/sec lines + diff the markdown columns")
-        print("\nFor full help: uv run ovis-ocr2-server.py --help")
-        sys.exit(0)
-
-    parser = argparse.ArgumentParser(
-        description="Document OCR using OvisOCR2 via an in-job vLLM server",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog="""
-Examples:
-  uv run ovis-ocr2-server.py my-docs analyzed-docs
-  uv run ovis-ocr2-server.py large-dataset test --max-samples 50 --shuffle
-See the module docstring for the full `hf jobs run` command (server + driver in one job).
-        """,
-    )
-
-    parser.add_argument("input_dataset", help="Input dataset ID from Hugging Face Hub")
-    parser.add_argument("output_dataset", help="Output dataset ID for Hugging Face Hub")
-    parser.add_argument(
-        "--image-column",
-        default="image",
-        help="Column containing images (default: image)",
-    )
-    parser.add_argument(
-        "--server",
-        default="http://127.0.0.1:8000",
-        help="vLLM server base URL (default: in-job localhost:8000)",
-    )
-    parser.add_argument(
-        "--concurrency",
-        type=int,
-        default=32,
-        help="Concurrent OCR requests (default: 32; vLLM queues excess internally, "
-        "so this mainly needs to be high enough to keep continuous batching fed)",
-    )
-    parser.add_argument(
-        "--max-tokens",
-        type=int,
-        default=16384,
-        help="Maximum tokens to generate (default: 16384, the model card value)",
-    )
-    parser.add_argument(
-        "--max-pixels",
-        type=int,
-        default=DEFAULT_MAX_PIXELS,
-        help=f"Maximum image pixels; larger images are downscaled client-side before "
-        f"upload (default: {DEFAULT_MAX_PIXELS}, = 2880*2880, the model card value)",
-    )
-    parser.add_argument(
-        "--request-timeout",
-        type=int,
-        default=1800,
-        help="Per-request timeout in seconds (default: 1800)",
-    )
-    parser.add_argument(
-        "--keep-image-tags",
-        action="store_true",
-        help="Keep visual-region <img src=\"images/bbox_...\"> tags in the output "
-        "(default: filtered, matching the upstream parser)",
-    )
-    parser.add_argument("--hf-token", help="Hugging Face API token")
-    parser.add_argument(
-        "--split", default="train", help="Dataset split to use (default: train)"
-    )
-    parser.add_argument(
-        "--max-samples",
-        type=int,
-        help="Maximum number of samples to process (for testing)",
-    )
-    parser.add_argument(
-        "--private", action="store_true", help="Make output dataset private"
-    )
-    parser.add_argument(
-        "--config",
-        help="Config/subset name when pushing to Hub (for benchmarking multiple models in one repo)",
-    )
-    parser.add_argument(
-        "--create-pr",
-        action="store_true",
-        help="Create a pull request instead of pushing directly (for parallel benchmarking)",
-    )
-    parser.add_argument(
-        "--shuffle", action="store_true", help="Shuffle dataset before processing"
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=42,
-        help="Random seed for shuffling (default: 42)",
-    )
-    parser.add_argument(
-        "--output-column",
-        default="markdown",
-        help="Column name for output text (default: markdown)",
-    )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
-
-    args = parser.parse_args()
-
-    main(
-        input_dataset=args.input_dataset,
-        output_dataset=args.output_dataset,
-        image_column=args.image_column,
-        server=args.server,
-        concurrency=args.concurrency,
-        max_tokens=args.max_tokens,
-        max_pixels=args.max_pixels,
-        request_timeout=args.request_timeout,
-        keep_image_tags=args.keep_image_tags,
-        hf_token=args.hf_token,
-        split=args.split,
-        max_samples=args.max_samples,
-        private=args.private,
-        shuffle=args.shuffle,
-        seed=args.seed,
-        output_column=args.output_column,
-        overwrite=args.overwrite,
-        config=args.config,
-        create_pr=args.create_pr,
-    )

ovis-ocr2.py

@@ -1,682 +0,0 @@
-# /// script
-# requires-python = ">=3.11"
-# dependencies = [
-#     "datasets>=4.0.0",
-#     "huggingface-hub",
-#     "pillow",
-#     "vllm>=0.22.1",
-#     "toolz",
-#     "torch",
-# ]
-# ///
-
-"""
-Convert document images to markdown using OvisOCR2 with vLLM.
-
-OvisOCR2 is a compact 0.9B end-to-end document parsing model (post-trained from
-Qwen3.5-0.8B with SFT + RL + OPD). It scores 96.58 on OmniDocBench v1.6 — the
-first end-to-end model to top that leaderboard — and 75.06 Avg3 on PureDocBench.
-Outputs a single Markdown document in natural reading order: LaTeX formulas,
-HTML tables, and (optionally) HTML <img> tags marking chart/image regions with
-bounding boxes scaled to [0, 1000).
-
-Model: ATH-MaaS/OvisOCR2 (Apache-2.0)
-vLLM: stock Qwen3_5ForConditionalGeneration arch, in stable vLLM >= 0.22.1
-      (the version the model card installs); no trust_remote_code needed.
-
-Features:
-- 0.9B parameters (ultra-compact, runs on l4x1)
-- Markdown output with LaTeX formulas + HTML tables
-- Visual-region <img> tags filtered by default (upstream parser default);
-  keep them with --keep-image-tags for downstream crop extraction
-- Upstream trailing-repeat cleanup applied to each output
-"""
-
-import argparse
-import io
-import json
-import logging
-import os
-import sys
-import time
-from datetime import datetime
-from typing import Any, Dict, Union
-
-import torch
-from datasets import load_dataset
-from huggingface_hub import DatasetCard, login
-from PIL import Image
-from toolz import partition_all
-
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
-# DeepGEMM's init calls _find_cuda_home, which asserts on the nvcc-less base image — a
-# non-fatal warning traceback that clutters the log (seen in the 2026-07-14 smoke run on
-# vLLM 0.25.1). Greedy OCR doesn't need the DeepGEMM JIT path, so disable it explicitly.
-os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
-from vllm import LLM, SamplingParams
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-MODEL = "ATH-MaaS/OvisOCR2"
-
-# Fixed instruction prompt, verbatim from the model card (including the leading newline).
-# The card warns outputs are tuned to this exact wording — don't "improve" it.
-OCR_PROMPT = (
-    "\nExtract all readable content from the image in natural human reading order "
-    "and output the result as a single Markdown document. For charts or images, "
-    'represent them using an HTML image tag: <img src="images/bbox_{left}_{top}_'
-    '{right}_{bottom}.jpg" />, where left, top, right, bottom are bounding box '
-    "coordinates scaled to [0, 1000). Format formulas as LaTeX. Format tables as "
-    "HTML: <table>...</table>. Transcribe all other text as standard Markdown. "
-    "Preserve the original text without translation or paraphrasing."
-)
-
-# Image bounds from the model card's parser (passed per-request as images_kwargs).
-DEFAULT_MIN_PIXELS = 448 * 448  # 200,704
-DEFAULT_MAX_PIXELS = 2880 * 2880  # 8,294,400
-
-
-def check_cuda_availability():
-    """Check if CUDA is available and exit if not."""
-    if not torch.cuda.is_available():
-        logger.error("CUDA is not available. This script requires a GPU.")
-        logger.error("Please run on a machine with a CUDA-capable GPU.")
-        sys.exit(1)
-    else:
-        logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
-
-
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
-def to_pil_image(image: Union[Image.Image, Dict[str, Any], str]) -> Image.Image:
-    """Convert a dataset image cell (PIL image, bytes dict, or path) to RGB PIL."""
-    if isinstance(image, Image.Image):
-        pil_img = image
-    elif isinstance(image, dict) and "bytes" in image:
-        pil_img = Image.open(io.BytesIO(image["bytes"]))
-    elif isinstance(image, str):
-        pil_img = Image.open(image)
-    else:
-        raise ValueError(f"Unsupported image type: {type(image)}")
-    return pil_img.convert("RGB")
-
-
-def clean_truncated_repeats(
-    text: str,
-    min_text_len: int = 8000,
-    max_period: int = 200,
-    min_period: int = 1,
-    min_repeat_chars: int = 100,
-    min_repeat_times: int = 5,
-) -> str:
-    """Trim degenerate trailing repetition (verbatim port of the model card's cleanup).
-
-    Long outputs that hit max_tokens can end in a repeated unit (a char, phrase, or
-    table row); this detects the shortest repeating tail unit and keeps one copy.
-    """
-    n = len(text)
-    if n < min_text_len:
-        return text
-
-    max_period = min(max_period, n - 1)
-    for unit_len in range(min_period, max_period + 1):
-        if text[n - 1] != text[n - 1 - unit_len]:
-            continue
-
-        match_len = 1
-        idx = n - 2
-        while idx >= unit_len and text[idx] == text[idx - unit_len]:
-            match_len += 1
-            idx -= 1
-
-        total_len = match_len + unit_len
-        repeat_times = total_len // unit_len
-        tail_len = total_len % unit_len
-
-        if repeat_times >= min_repeat_times and total_len >= min_repeat_chars:
-            return text[: n - total_len + unit_len] + text[n - tail_len :]
-
-    return text
-
-
-def filter_image_tags(text: str) -> str:
-    """Drop visual-region <img> blocks (upstream parser's default behaviour)."""
-    return "\n\n".join(
-        block
-        for block in text.split("\n\n")
-        if not block.strip().startswith('<img src="images/bbox_')
-    )
-
-
-def postprocess_output(text: str, keep_image_tags: bool) -> str:
-    text = text.strip()
-    if not keep_image_tags:
-        text = filter_image_tags(text)
-    return clean_truncated_repeats(text)
-
-
-def create_dataset_card(
-    source_dataset: str,
-    model: str,
-    num_samples: int,
-    processing_time: str,
-    batch_size: int,
-    max_model_len: int,
-    max_tokens: int,
-    gpu_memory_utilization: float,
-    keep_image_tags: bool,
-    image_column: str = "image",
-    split: str = "train",
-) -> str:
-    """Create a dataset card documenting the OCR process."""
-    model_name = model.split("/")[-1]
-
-    # Canonical provenance stamp (see AGENTS.md): Jobs claim gated on JOB_ID, set by HF Jobs in-container.
-    on_jobs = os.environ.get("JOB_ID") is not None
-    hw = os.environ.get("ACCELERATOR") or ""
-    origin = (
-        "Produced on [Hugging Face Jobs](https://huggingface.co/docs/huggingface_hub/guides/jobs)"
-        + (f" (`{hw}`)" if hw else "")
-    ) if on_jobs else "Generated"
-    jobs_tag = "\n- hf-jobs" if on_jobs else ""
-
-    return f"""---
-tags:
-- ocr
-- document-processing
-- ovis-ocr2
-- markdown
-- uv-script
-- generated{jobs_tag}
----
-
-# Document OCR using {model_name}
-
-This dataset contains OCR results from images in [{source_dataset}](https://huggingface.co/datasets/{source_dataset}) using OvisOCR2, a compact 0.9B document parsing model (96.58 on OmniDocBench v1.6).
-
-## Processing Details
-
-- **Source Dataset**: [{source_dataset}](https://huggingface.co/datasets/{source_dataset})
-- **Model**: [{model}](https://huggingface.co/{model})
-- **Number of Samples**: {num_samples:,}
-- **Processing Time**: {processing_time}
-- **Processing Date**: {datetime.now().strftime("%Y-%m-%d %H:%M UTC")}
-
-### Configuration
-
-- **Image Column**: `{image_column}`
-- **Dataset Split**: `{split}`
-- **Batch Size**: {batch_size}
-- **Max Model Length**: {max_model_len:,} tokens
-- **Max Output Tokens**: {max_tokens:,}
-- **Temperature**: 0.0 (greedy, per model card)
-- **GPU Memory Utilization**: {gpu_memory_utilization:.1%}
-- **Visual-region image tags**: {"kept" if keep_image_tags else "filtered (default)"}
-
-## Model Information
-
-OvisOCR2 is a compact, high-performance document parsing model:
-- 0.9B parameters (post-trained from Qwen3.5-0.8B with SFT + RL + OPD)
-- 96.58 on OmniDocBench v1.6 (first end-to-end model to top the leaderboard)
-- Markdown output in natural reading order
-- LaTeX formula recognition, HTML table extraction
-- Apache-2.0 licensed
-
-## Dataset Structure
-
-The dataset contains all original columns plus:
-- `markdown`: The extracted text in markdown format
-- `inference_info`: JSON list tracking all OCR models applied to this dataset
-
-## Reproduction
-
-{origin} with the [`ovis-ocr2.py`](https://huggingface.co/datasets/uv-scripts/ocr/raw/main/ovis-ocr2.py) recipe from [uv-scripts](https://huggingface.co/uv-scripts). Run it yourself:
-
-```bash
-hf jobs uv run https://huggingface.co/datasets/uv-scripts/ocr/raw/main/ovis-ocr2.py \\
-    {source_dataset} \\
-    <output-dataset> \\
-    --image-column {image_column} \\
-    --batch-size {batch_size}
-```
-"""
-
-
-def main(
-    input_dataset: str,
-    output_dataset: str,
-    image_column: str = "image",
-    batch_size: int = 16,
-    max_model_len: int = 32768,
-    max_tokens: int = 16384,
-    min_pixels: int = DEFAULT_MIN_PIXELS,
-    max_pixels: int = DEFAULT_MAX_PIXELS,
-    gpu_memory_utilization: float = 0.8,
-    keep_image_tags: bool = False,
-    hf_token: str = None,
-    split: str = "train",
-    max_samples: int = None,
-    private: bool = False,
-    shuffle: bool = False,
-    seed: int = 42,
-    output_column: str = "markdown",
-    overwrite: bool = False,
-    verbose: bool = False,
-    config: str = None,
-    create_pr: bool = False,
-):
-    """Process images from HF dataset through OvisOCR2."""
-
-    check_cuda_availability()
-
-    start_time = datetime.now()
-
-    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")
-    if HF_TOKEN:
-        login(token=HF_TOKEN)
-
-    logger.info(f"Using model: {MODEL}")
-
-    # Load dataset
-    logger.info(f"Loading dataset: {input_dataset}")
-    dataset = load_dataset(input_dataset, split=split)
-
-    if image_column not in dataset.column_names:
-        raise ValueError(
-            f"Column '{image_column}' not found. Available: {dataset.column_names}"
-        )
-
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
-    if shuffle:
-        logger.info(f"Shuffling dataset with seed {seed}")
-        dataset = dataset.shuffle(seed=seed)
-
-    if max_samples:
-        dataset = dataset.select(range(min(max_samples, len(dataset))))
-        logger.info(f"Limited to {len(dataset)} samples")
-
-    # Initialize vLLM
-    logger.info("Initializing vLLM with OvisOCR2")
-    logger.info("This may take a few minutes on first run...")
-    # gdn_prefill_backend="triton" is from the model card: Qwen3.5 uses gated-delta-net
-    # linear attention, and the non-triton GDN prefill path needs a JIT/CUDA toolchain
-    # the bare uv image lacks (same class of problem as the FLASHINFER guard above).
-    llm = LLM(
-        model=MODEL,
-        max_model_len=max_model_len,
-        gpu_memory_utilization=gpu_memory_utilization,
-        limit_mm_per_prompt={"image": 1},
-        gdn_prefill_backend="triton",
-    )
-
-    # The model card builds the prompt once via the chat template with
-    # enable_thinking=False — Qwen3.5 templates can otherwise inject a thinking
-    # preamble, which is wrong for OCR. Mirror it exactly.
-    prompt = llm.get_tokenizer().apply_chat_template(
-        [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image"},
-                    {"type": "text", "text": OCR_PROMPT},
-                ],
-            }
-        ],
-        tokenize=False,
-        add_generation_prompt=True,
-        enable_thinking=False,
-    )
-
-    # Card-locked sampling: greedy, 16384 max tokens.
-    sampling_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
-
-    logger.info(f"Processing {len(dataset)} images in batches of {batch_size}")
-    logger.info(f"Output will be written to column: {output_column}")
-
-    all_outputs = []
-    total_batches = (len(dataset) + batch_size - 1) // batch_size
-    processed = 0
-
-    for batch_num, batch_indices in enumerate(
-        partition_all(batch_size, range(len(dataset))), 1
-    ):
-        batch_indices = list(batch_indices)
-        batch_images = [dataset[i][image_column] for i in batch_indices]
-
-        logger.info(
-            f"Batch {batch_num}/{total_batches} "
-            f"({processed}/{len(dataset)} images done)"
-        )
-
-        try:
-            # Per-request inputs exactly as the model card's parser builds them
-            # (PIL image + images_kwargs pixel bounds), via llm.generate not llm.chat
-            # so the enable_thinking=False prompt above is used verbatim.
-            vllm_inputs = [
-                {
-                    "prompt": prompt,
-                    "multi_modal_data": {"image": to_pil_image(img)},
-                    "mm_processor_kwargs": {
-                        "images_kwargs": {
-                            "min_pixels": min_pixels,
-                            "max_pixels": max_pixels,
-                        }
-                    },
-                }
-                for img in batch_images
-            ]
-
-            outputs = llm.generate(vllm_inputs, sampling_params)
-
-            for output in outputs:
-                text = output.outputs[0].text
-                all_outputs.append(postprocess_output(text, keep_image_tags))
-
-            processed += len(batch_images)
-
-        except Exception as e:
-            logger.error(f"Error processing batch: {e}")
-            all_outputs.extend(["[OCR ERROR]"] * len(batch_images))
-            processed += len(batch_images)
-
-    processing_duration = datetime.now() - start_time
-    processing_time_str = f"{processing_duration.total_seconds() / 60:.1f} min"
-
-    logger.info(f"Adding '{output_column}' column to dataset")
-    dataset = dataset.add_column(output_column, all_outputs)
-
-    # Inference info tracking
-    inference_entry = {
-        "model_id": MODEL,
-        "model_name": "OvisOCR2",
-        "column_name": output_column,
-        "timestamp": datetime.now().isoformat(),
-        "temperature": 0.0,
-        "max_tokens": max_tokens,
-        "min_pixels": min_pixels,
-        "max_pixels": max_pixels,
-        "keep_image_tags": keep_image_tags,
-    }
-
-    if "inference_info" in dataset.column_names:
-        logger.info("Updating existing inference_info column")
-
-        def update_inference_info(example):
-            try:
-                existing_info = (
-                    json.loads(example["inference_info"])
-                    if example["inference_info"]
-                    else []
-                )
-            except (json.JSONDecodeError, TypeError):
-                existing_info = []
-            existing_info.append(inference_entry)
-            return {"inference_info": json.dumps(existing_info)}
-
-        dataset = dataset.map(update_inference_info)
-    else:
-        logger.info("Creating new inference_info column")
-        inference_list = [json.dumps([inference_entry])] * len(dataset)
-        dataset = dataset.add_column("inference_info", inference_list)
-
-    # Push to hub with retry and XET fallback
-    logger.info(f"Pushing to {output_dataset}")
-    max_retries = 3
-    for attempt in range(1, max_retries + 1):
-        try:
-            if attempt > 1:
-                logger.warning("Disabling XET (fallback to HTTP upload)")
-                os.environ["HF_HUB_DISABLE_XET"] = "1"
-            dataset.push_to_hub(
-                output_dataset,
-                private=private,
-                token=HF_TOKEN,
-                max_shard_size="500MB",
-                **({"config_name": config} if config else {}),
-                create_pr=create_pr,
-                commit_message=f"Add {MODEL} OCR results ({len(dataset)} samples)"
-                + (f" [{config}]" if config else ""),
-            )
-            break
-        except Exception as e:
-            logger.error(f"Upload attempt {attempt}/{max_retries} failed: {e}")
-            if attempt < max_retries:
-                delay = 30 * (2 ** (attempt - 1))
-                logger.info(f"Retrying in {delay}s...")
-                time.sleep(delay)
-            else:
-                logger.error("All upload attempts failed. OCR results are lost.")
-                sys.exit(1)
-
-    # Create and push dataset card
-    logger.info("Creating dataset card")
-    card_content = create_dataset_card(
-        source_dataset=input_dataset,
-        model=MODEL,
-        num_samples=len(dataset),
-        processing_time=processing_time_str,
-        batch_size=batch_size,
-        max_model_len=max_model_len,
-        max_tokens=max_tokens,
-        gpu_memory_utilization=gpu_memory_utilization,
-        keep_image_tags=keep_image_tags,
-        image_column=image_column,
-        split=split,
-    )
-
-    card = DatasetCard(card_content)
-    card.push_to_hub(output_dataset, token=HF_TOKEN)
-
-    logger.info("Done! OvisOCR2 processing complete.")
-    logger.info(
-        f"Dataset available at: https://huggingface.co/datasets/{output_dataset}"
-    )
-    logger.info(f"Processing time: {processing_time_str}")
-    logger.info(
-        f"Processing speed: {len(dataset) / processing_duration.total_seconds():.2f} images/sec"
-    )
-
-    if verbose:
-        import importlib.metadata
-
-        logger.info("--- Resolved package versions ---")
-        for pkg in ["vllm", "transformers", "torch", "datasets", "pyarrow", "pillow"]:
-            try:
-                logger.info(f"  {pkg}=={importlib.metadata.version(pkg)}")
-            except importlib.metadata.PackageNotFoundError:
-                logger.info(f"  {pkg}: not installed")
-        logger.info("--- End versions ---")
-
-
-if __name__ == "__main__":
-    if len(sys.argv) == 1:
-        print("=" * 70)
-        print("OvisOCR2 Document Processing")
-        print("=" * 70)
-        print("\n0.9B document parsing model - 96.58 on OmniDocBench v1.6")
-        print("\nOutputs markdown in natural reading order:")
-        print("  - LaTeX formulas, HTML tables")
-        print("  - Visual-region <img> tags filtered by default")
-        print("    (--keep-image-tags to retain them)")
-        print("\nExamples:")
-        print("\n1. Basic OCR:")
-        print("   uv run ovis-ocr2.py input-dataset output-dataset")
-        print("\n2. Keep visual-region image tags:")
-        print("   uv run ovis-ocr2.py docs results --keep-image-tags")
-        print("\n3. Test with small sample:")
-        print("   uv run ovis-ocr2.py large-dataset test --max-samples 10 --shuffle")
-        print("\n4. Running on HF Jobs:")
-        print("   hf jobs uv run --flavor l4x1 \\")
-        print("     -s HF_TOKEN \\")
-        print(
-            "     https://huggingface.co/datasets/uv-scripts/ocr/raw/main/ovis-ocr2.py \\"
-        )
-        print("       input-dataset output-dataset --batch-size 16")
-        print("\nFor full help: uv run ovis-ocr2.py --help")
-        sys.exit(0)
-
-    parser = argparse.ArgumentParser(
-        description="Document OCR using OvisOCR2 (0.9B, 96.58 OmniDocBench v1.6)",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog="""
-Examples:
-  uv run ovis-ocr2.py my-docs analyzed-docs
-  uv run ovis-ocr2.py docs results --keep-image-tags
-  uv run ovis-ocr2.py large-dataset test --max-samples 50 --shuffle
-        """,
-    )
-
-    parser.add_argument("input_dataset", help="Input dataset ID from Hugging Face Hub")
-    parser.add_argument("output_dataset", help="Output dataset ID for Hugging Face Hub")
-    parser.add_argument(
-        "--image-column",
-        default="image",
-        help="Column containing images (default: image)",
-    )
-    parser.add_argument(
-        "--batch-size",
-        type=int,
-        default=16,
-        help="Batch size for processing (default: 16)",
-    )
-    parser.add_argument(
-        "--max-model-len",
-        type=int,
-        default=32768,
-        help="Maximum model context length (default: 32768; model supports 262144)",
-    )
-    parser.add_argument(
-        "--max-tokens",
-        type=int,
-        default=16384,
-        help="Maximum tokens to generate (default: 16384, the model card value)",
-    )
-    parser.add_argument(
-        "--min-pixels",
-        type=int,
-        default=DEFAULT_MIN_PIXELS,
-        help=f"Minimum image pixels for the processor (default: {DEFAULT_MIN_PIXELS}, "
-        "= 448*448, the model card value)",
-    )
-    parser.add_argument(
-        "--max-pixels",
-        type=int,
-        default=DEFAULT_MAX_PIXELS,
-        help=f"Maximum image pixels for the processor; larger images are downscaled "
-        f"internally (default: {DEFAULT_MAX_PIXELS}, = 2880*2880, the model card value)",
-    )
-    parser.add_argument(
-        "--gpu-memory-utilization",
-        type=float,
-        default=0.8,
-        help="GPU memory utilization (default: 0.8)",
-    )
-    parser.add_argument(
-        "--keep-image-tags",
-        action="store_true",
-        help="Keep visual-region <img src=\"images/bbox_...\"> tags in the output "
-        "(default: filtered, matching the upstream parser)",
-    )
-    parser.add_argument("--hf-token", help="Hugging Face API token")
-    parser.add_argument(
-        "--split", default="train", help="Dataset split to use (default: train)"
-    )
-    parser.add_argument(
-        "--max-samples",
-        type=int,
-        help="Maximum number of samples to process (for testing)",
-    )
-    parser.add_argument(
-        "--private", action="store_true", help="Make output dataset private"
-    )
-    parser.add_argument(
-        "--config",
-        help="Config/subset name when pushing to Hub (for benchmarking multiple models in one repo)",
-    )
-    parser.add_argument(
-        "--create-pr",
-        action="store_true",
-        help="Create a pull request instead of pushing directly (for parallel benchmarking)",
-    )
-    parser.add_argument(
-        "--shuffle", action="store_true", help="Shuffle dataset before processing"
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=42,
-        help="Random seed for shuffling (default: 42)",
-    )
-    parser.add_argument(
-        "--output-column",
-        default="markdown",
-        help="Column name for output text (default: markdown)",
-    )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
-    parser.add_argument(
-        "--verbose",
-        action="store_true",
-        help="Log resolved package versions after processing (useful for pinning deps)",
-    )
-
-    args = parser.parse_args()
-
-    if args.max_tokens > args.max_model_len:
-        parser.error(
-            f"--max-tokens ({args.max_tokens}) must be <= --max-model-len ({args.max_model_len})"
-        )
-
-    main(
-        input_dataset=args.input_dataset,
-        output_dataset=args.output_dataset,
-        image_column=args.image_column,
-        batch_size=args.batch_size,
-        max_model_len=args.max_model_len,
-        max_tokens=args.max_tokens,
-        min_pixels=args.min_pixels,
-        max_pixels=args.max_pixels,
-        gpu_memory_utilization=args.gpu_memory_utilization,
-        keep_image_tags=args.keep_image_tags,
-        hf_token=args.hf_token,
-        split=args.split,
-        max_samples=args.max_samples,
-        private=args.private,
-        shuffle=args.shuffle,
-        seed=args.seed,
-        output_column=args.output_column,
-        overwrite=args.overwrite,
-        verbose=args.verbose,
-        config=args.config,
-        create_pr=args.create_pr,
-    )

paddleocr-vl-1.5.py

@@ -1,7 +1,7 @@
 # /// script
 # requires-python = ">=3.11"
 # dependencies = [
-#     "datasets>=4.0.0",
+#     "datasets",
 #     "huggingface-hub",
 #     "pillow",
 #     "torch",
@@ -101,28 +101,6 @@ def check_cuda_availability():
         logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
 
 
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
 def prepare_image(
     image: Union[Image.Image, Dict[str, Any], str],
 ) -> Image.Image:
@@ -304,7 +282,6 @@ def main(
     shuffle: bool = False,
     seed: int = 42,
     output_column: str = "markdown",
-    overwrite: bool = False,
     config: str = None,
     create_pr: bool = False,
     verbose: bool = False,
@@ -341,9 +318,6 @@ def main(
             f"Column '{image_column}' not found. Available: {dataset.column_names}"
         )
 
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
     # Shuffle if requested
     if shuffle:
         logger.info(f"Shuffling dataset with seed {seed}")
@@ -667,12 +641,6 @@ Backend: Transformers batch inference (not vLLM)
         default="markdown",
         help="Column name for output text (default: markdown)",
     )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
     parser.add_argument(
         "--config",
         help="Config/subset name when pushing to Hub (for benchmarking multiple models in one repo)",
@@ -703,7 +671,6 @@ Backend: Transformers batch inference (not vLLM)
         shuffle=args.shuffle,
         seed=args.seed,
         output_column=args.output_column,
-        overwrite=args.overwrite,
         config=args.config,
         create_pr=args.create_pr,
         verbose=args.verbose,

paddleocr-vl-1.6.py

@@ -1,831 +0,0 @@
-# /// script
-# requires-python = ">=3.11"
-# dependencies = [
-#     "datasets>=4.0.0",
-#     "huggingface-hub",
-#     "pillow",
-#     "vllm>=0.15.1",
-#     "tqdm",
-#     "toolz",
-#     "torch",
-#     "pyarrow",
-#     "transformers",
-# ]
-# ///
-
-"""
-Convert document images to text/tables/formulas using PaddleOCR-VL-1.6 with vLLM.
-
-PaddleOCR-VL-1.6 is a compact 0.9B OCR model that reaches a new SOTA of 96.33% on
-OmniDocBench v1.6. It combines a NaViT-style dynamic resolution visual encoder with
-the ERNIE-4.5-0.3B language model and is a plug-and-play upgrade of PaddleOCR-VL-1.5.
-
-Features:
-- 🎯 SOTA: 96.33% on OmniDocBench v1.6 (0.9B params, smallest top-tier OCR model)
-- 📝 OCR mode: General text extraction to markdown
-- 📊 Table mode: HTML table recognition and extraction
-- 📐 Formula mode: LaTeX mathematical notation
-- 📈 Chart mode: Structured chart analysis
-- 🔍 Spotting mode: Text spotting with localization
-- 🔖 Seal mode: Seal/stamp recognition
-- 🌍 Multilingual support (en/zh + more)
-- 🔧 Based on ERNIE-4.5 (different from Qwen-based models)
-
-Model: PaddlePaddle/PaddleOCR-VL-1.6
-Backend: vLLM offline (batch inference)
-
-HF Jobs note: PaddleOCR-VL-1.6 is supported by stable vLLM, but on HF Jobs you must run
-with the pre-built vLLM image so flashinfer's CUDA kernels are reused. The default
-uv-script image has the CUDA runtime but no `nvcc`, so vLLM's flashinfer sampler crashes
-at warmup with "Could not find nvcc". Use image-mode (see the example at the bottom):
-    --image vllm/vllm-openai:latest --flavor a100-large
-    --python /usr/bin/python3 -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages
-This is the same image-mode pattern as nuextract3.py. Verified end-to-end on a100-large
-(2026-06-01): 5/5 clean markdown on davanstrien/ufo-ColPali, ~194 tok/s, 0 errors.
-"""
-
-import argparse
-import base64
-import io
-import json
-import logging
-import math
-import os
-import sys
-import time
-from typing import Any, Dict, List, Union
-from datetime import datetime
-
-import torch
-from datasets import load_dataset
-from huggingface_hub import DatasetCard, login
-from PIL import Image
-from toolz import partition_all
-from tqdm.auto import tqdm
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
-from vllm import LLM, SamplingParams
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-MODEL_ID = "PaddlePaddle/PaddleOCR-VL-1.6"
-
-# Task mode configurations from official PaddleOCR-VL documentation
-TASK_MODES = {
-    "ocr": "OCR:",
-    "table": "Table Recognition:",
-    "formula": "Formula Recognition:",
-    "chart": "Chart Recognition:",
-    "spotting": "Spotting:",
-    "seal": "Seal Recognition:",
-}
-
-# Task descriptions for dataset card
-TASK_DESCRIPTIONS = {
-    "ocr": "General text extraction to markdown format",
-    "table": "Table extraction to HTML format",
-    "formula": "Mathematical formula recognition to LaTeX",
-    "chart": "Chart and diagram analysis",
-    "spotting": "Text spotting with localization",
-    "seal": "Seal and stamp recognition",
-}
-
-
-def check_cuda_availability():
-    """Check if CUDA is available and exit if not."""
-    if not torch.cuda.is_available():
-        logger.error("CUDA is not available. This script requires a GPU.")
-        logger.error("Please run on a machine with a CUDA-capable GPU.")
-        sys.exit(1)
-    else:
-        logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
-
-
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
-def smart_resize(
-    height: int,
-    width: int,
-    factor: int = 28,
-    min_pixels: int = 28 * 28 * 130,
-    max_pixels: int = 28 * 28 * 1280,
-) -> tuple[int, int]:
-    """
-    PaddleOCR-VL's intelligent resize logic.
-
-    Rescales the image so that:
-    1. Both dimensions are divisible by 'factor' (28)
-    2. Total pixels are within [min_pixels, max_pixels]
-    3. Aspect ratio is maintained as closely as possible
-
-    Args:
-        height: Original image height
-        width: Original image width
-        factor: Dimension divisibility factor (default: 28)
-        min_pixels: Minimum total pixels (default: 100,880)
-        max_pixels: Maximum total pixels (default: 1,003,520)
-
-    Returns:
-        Tuple of (new_height, new_width)
-    """
-    if height < factor:
-        width = round((width * factor) / height)
-        height = factor
-
-    if width < factor:
-        height = round((height * factor) / width)
-        width = factor
-
-    if max(height, width) / min(height, width) > 200:
-        logger.warning(
-            f"Extreme aspect ratio detected: {max(height, width) / min(height, width):.1f}"
-        )
-        # Continue anyway, but warn about potential issues
-
-    h_bar = round(height / factor) * factor
-    w_bar = round(width / factor) * factor
-
-    if h_bar * w_bar > max_pixels:
-        beta = math.sqrt((height * width) / max_pixels)
-        h_bar = math.floor(height / beta / factor) * factor
-        w_bar = math.floor(width / beta / factor) * factor
-    elif h_bar * w_bar < min_pixels:
-        beta = math.sqrt(min_pixels / (height * width))
-        h_bar = math.ceil(height * beta / factor) * factor
-        w_bar = math.ceil(width * beta / factor) * factor
-
-    return h_bar, w_bar
-
-
-def make_ocr_message(
-    image: Union[Image.Image, Dict[str, Any], str],
-    task_mode: str = "ocr",
-    apply_smart_resize: bool = True,
-) -> List[Dict]:
-    """
-    Create chat message for PaddleOCR-VL processing.
-
-    PaddleOCR-VL expects a specific format with the task prefix after the image.
-    """
-    # Convert to PIL Image if needed
-    if isinstance(image, Image.Image):
-        pil_img = image
-    elif isinstance(image, dict) and "bytes" in image:
-        pil_img = Image.open(io.BytesIO(image["bytes"]))
-    elif isinstance(image, str):
-        pil_img = Image.open(image)
-    else:
-        raise ValueError(f"Unsupported image type: {type(image)}")
-
-    # Convert to RGB
-    pil_img = pil_img.convert("RGB")
-
-    # Apply smart resize if requested. Spotting benefits from higher resolution
-    # (per the model card), so allow more pixels in that mode.
-    if apply_smart_resize:
-        original_size = pil_img.size
-        max_pixels = 28 * 28 * (2048 if task_mode == "spotting" else 1280)
-        new_height, new_width = smart_resize(
-            pil_img.height, pil_img.width, max_pixels=max_pixels
-        )
-        if (new_width, new_height) != (pil_img.width, pil_img.height):
-            pil_img = pil_img.resize((new_width, new_height), Image.Resampling.LANCZOS)
-            logger.debug(f"Resized image from {original_size} to {pil_img.size}")
-
-    # Convert to base64 data URI
-    buf = io.BytesIO()
-    pil_img.save(buf, format="PNG")
-    data_uri = f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}"
-
-    # PaddleOCR-VL message format: image first, then task prefix
-    return [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image_url", "image_url": {"url": data_uri}},
-                {"type": "text", "text": TASK_MODES[task_mode]},
-            ],
-        }
-    ]
-
-
-def create_dataset_card(
-    source_dataset: str,
-    model: str,
-    task_mode: str,
-    num_samples: int,
-    processing_time: str,
-    batch_size: int,
-    max_model_len: int,
-    max_tokens: int,
-    gpu_memory_utilization: float,
-    temperature: float,
-    apply_smart_resize: bool,
-    image_column: str = "image",
-    split: str = "train",
-) -> str:
-    """Create a dataset card documenting the OCR process."""
-    task_description = TASK_DESCRIPTIONS[task_mode]
-
-    return f"""---
-tags:
-- ocr
-- document-processing
-- paddleocr-vl
-- paddleocr-vl-1.6
-- {task_mode}
-- uv-script
-- generated
----
-
-# Document Processing using PaddleOCR-VL-1.6 ({task_mode.upper()} mode)
-
-This dataset contains {task_mode.upper()} results from images in [{source_dataset}](https://huggingface.co/datasets/{source_dataset}) using PaddleOCR-VL-1.6, an ultra-compact 0.9B OCR model (96.33% SOTA on OmniDocBench v1.6).
-
-## Processing Details
-
-- **Source Dataset**: [{source_dataset}](https://huggingface.co/datasets/{source_dataset})
-- **Model**: [{model}](https://huggingface.co/{model})
-- **Task Mode**: `{task_mode}` - {task_description}
-- **Number of Samples**: {num_samples:,}
-- **Processing Time**: {processing_time}
-- **Processing Date**: {datetime.now().strftime("%Y-%m-%d %H:%M UTC")}
-
-### Configuration
-
-- **Image Column**: `{image_column}`
-- **Output Column**: `markdown`
-- **Dataset Split**: `{split}`
-- **Batch Size**: {batch_size}
-- **Smart Resize**: {"Enabled" if apply_smart_resize else "Disabled"}
-- **Max Model Length**: {max_model_len:,} tokens
-- **Max Output Tokens**: {max_tokens:,}
-- **Temperature**: {temperature}
-- **GPU Memory Utilization**: {gpu_memory_utilization:.1%}
-
-## Model Information
-
-PaddleOCR-VL-1.6 is a state-of-the-art, resource-efficient model tailored for document parsing:
-- 🎯 **SOTA** - 96.33% on OmniDocBench v1.6
-- 🧩 **Ultra-compact** - Only 0.9B parameters
-- 📝 **OCR mode** - General text extraction
-- 📊 **Table mode** - HTML table recognition
-- 📐 **Formula mode** - LaTeX mathematical notation
-- 📈 **Chart mode** - Structured chart analysis
-- 🔍 **Spotting mode** - Text spotting with localization
-- 🔖 **Seal mode** - Seal/stamp recognition
-- 🌍 **Multilingual** - Support for multiple languages
-- 🔧 **ERNIE-4.5 based** - Different architecture from Qwen models
-
-### Task Modes
-
-- **OCR**: Extract text content to markdown format
-- **Table Recognition**: Extract tables to HTML format
-- **Formula Recognition**: Extract mathematical formulas to LaTeX
-- **Chart Recognition**: Analyze and describe charts/diagrams
-- **Spotting**: Text spotting with localization
-- **Seal Recognition**: Seal and stamp recognition
-
-## Dataset Structure
-
-The dataset contains all original columns plus:
-- `markdown`: The extracted content based on task mode
-- `inference_info`: JSON list tracking all OCR models applied to this dataset
-
-## Usage
-
-```python
-from datasets import load_dataset
-import json
-
-# Load the dataset
-dataset = load_dataset("{{output_dataset_id}}", split="{split}")
-
-# Access the extracted content
-for example in dataset:
-    print(example["markdown"])
-    break
-
-# View all OCR models applied to this dataset
-inference_info = json.loads(dataset[0]["inference_info"])
-for info in inference_info:
-    print(f"Task: {{info['task_mode']}} - Model: {{info['model_id']}}")
-```
-
-## Reproduction
-
-This dataset was generated using the [uv-scripts/ocr](https://huggingface.co/datasets/uv-scripts/ocr) PaddleOCR-VL-1.6 script.
-On HF Jobs, run with the pre-built vLLM image (image-mode) so flashinfer kernels are reused:
-
-```bash
-hf jobs uv run \\
-    --image vllm/vllm-openai:latest --flavor a100-large \\
-    --python /usr/bin/python3 -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages \\
-    -s HF_TOKEN \\
-    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/paddleocr-vl-1.6.py \\
-    {source_dataset} \\
-    <output-dataset> \\
-    --task-mode {task_mode} \\
-    --image-column {image_column} \\
-    --batch-size {batch_size} \\
-    --max-model-len {max_model_len} \\
-    --max-tokens {max_tokens} \\
-    --gpu-memory-utilization {gpu_memory_utilization}
-```
-
-## Performance
-
-- **Model Size**: 0.9B parameters (smallest among top-tier OCR models)
-- **Processing Speed**: ~{num_samples / (float(processing_time.split()[0]) * 60):.2f} images/second
-- **Architecture**: NaViT visual encoder + ERNIE-4.5-0.3B language model
-
-Generated with 🤖 [UV Scripts](https://huggingface.co/uv-scripts)
-"""
-
-
-def main(
-    input_dataset: str,
-    output_dataset: str,
-    image_column: str = "image",
-    batch_size: int = 16,
-    task_mode: str = "ocr",
-    max_model_len: int = 8192,
-    max_tokens: int = 4096,
-    temperature: float = 0.0,
-    gpu_memory_utilization: float = 0.8,
-    apply_smart_resize: bool = True,
-    hf_token: str = None,
-    split: str = "train",
-    max_samples: int = None,
-    private: bool = False,
-    shuffle: bool = False,
-    seed: int = 42,
-    output_column: str = None,
-    overwrite: bool = False,
-    config: str = None,
-    create_pr: bool = False,
-    verbose: bool = False,
-):
-    """Process images from HF dataset through PaddleOCR-VL-1.6 model."""
-
-    # Check CUDA availability first
-    check_cuda_availability()
-
-    # Track processing start time
-    start_time = datetime.now()
-
-    # Enable high-performance Xet downloads
-    os.environ["HF_XET_HIGH_PERFORMANCE"] = "1"
-
-    # Login to HF if token provided
-    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")
-    if HF_TOKEN:
-        login(token=HF_TOKEN)
-
-    # Validate task mode
-    if task_mode not in TASK_MODES:
-        raise ValueError(
-            f"Invalid task_mode '{task_mode}'. Choose from: {list(TASK_MODES.keys())}"
-        )
-
-    # Default output column is 'markdown' for consistency across scripts
-    if output_column is None:
-        output_column = "markdown"
-
-    logger.info(f"Using task mode: {task_mode} - {TASK_DESCRIPTIONS[task_mode]}")
-    logger.info(f"Output will be written to column: {output_column}")
-
-    # Load dataset
-    logger.info(f"Loading dataset: {input_dataset}")
-    dataset = load_dataset(input_dataset, split=split)
-
-    # Validate image column
-    if image_column not in dataset.column_names:
-        raise ValueError(
-            f"Column '{image_column}' not found. Available: {dataset.column_names}"
-        )
-
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
-    # Shuffle if requested
-    if shuffle:
-        logger.info(f"Shuffling dataset with seed {seed}")
-        dataset = dataset.shuffle(seed=seed)
-
-    # Limit samples if requested
-    if max_samples:
-        dataset = dataset.select(range(min(max_samples, len(dataset))))
-        logger.info(f"Limited to {len(dataset)} samples")
-
-    # Initialize vLLM model
-    logger.info(f"Initializing vLLM with {MODEL_ID}")
-    logger.info("This may take a minute on first run (model is only 0.9B)...")
-
-    try:
-        llm = LLM(
-            model=MODEL_ID,
-            trust_remote_code=True,
-            max_model_len=max_model_len,
-            gpu_memory_utilization=gpu_memory_utilization,
-            limit_mm_per_prompt={"image": 1},
-            max_num_batched_tokens=16384,
-            enable_prefix_caching=False,
-            enforce_eager=True,
-        )
-    except Exception as e:
-        logger.error(f"Failed to initialize PaddleOCR-VL-1.6 with vLLM: {e}")
-        logger.error(
-            "On HF Jobs, run with the pre-built vLLM image so flashinfer kernels are "
-            "reused (the default uv-script image has no nvcc):"
-        )
-        logger.error(
-            "  --image vllm/vllm-openai:latest --flavor a100-large "
-            "--python /usr/bin/python3 -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages"
-        )
-        sys.exit(1)
-
-    # Sampling parameters - deterministic for OCR
-    sampling_params = SamplingParams(
-        temperature=temperature,
-        max_tokens=max_tokens,
-    )
-
-    logger.info(f"Processing {len(dataset)} images in batches of {batch_size}")
-    if apply_smart_resize:
-        logger.info("Smart resize enabled (PaddleOCR-VL's adaptive resolution)")
-
-    # Process images in batches
-    all_outputs = []
-
-    for batch_indices in tqdm(
-        partition_all(batch_size, range(len(dataset))),
-        total=(len(dataset) + batch_size - 1) // batch_size,
-        desc=f"PaddleOCR-VL-1.6 {task_mode.upper()} processing",
-    ):
-        batch_indices = list(batch_indices)
-        batch_images = [dataset[i][image_column] for i in batch_indices]
-
-        try:
-            # Create messages for batch with task-specific prefix
-            batch_messages = [
-                make_ocr_message(
-                    img, task_mode=task_mode, apply_smart_resize=apply_smart_resize
-                )
-                for img in batch_images
-            ]
-
-            # Process with vLLM
-            outputs = llm.chat(batch_messages, sampling_params)
-
-            # Extract outputs
-            for output in outputs:
-                text = output.outputs[0].text.strip()
-                all_outputs.append(text)
-
-        except Exception as e:
-            logger.error(f"Error processing batch: {e}")
-            # Add error placeholders for failed batch
-            all_outputs.extend([f"[{task_mode.upper()} ERROR]"] * len(batch_images))
-
-    # Calculate processing time
-    processing_duration = datetime.now() - start_time
-    processing_time_str = f"{processing_duration.total_seconds() / 60:.1f} min"
-
-    # Add output column to dataset
-    logger.info(f"Adding '{output_column}' column to dataset")
-    dataset = dataset.add_column(output_column, all_outputs)
-
-    # Handle inference_info tracking (for multi-model comparisons)
-    inference_entry = {
-        "model_id": MODEL_ID,
-        "model_name": "PaddleOCR-VL-1.6",
-        "model_size": "0.9B",
-        "task_mode": task_mode,
-        "column_name": output_column,
-        "timestamp": datetime.now().isoformat(),
-        "temperature": temperature,
-        "max_tokens": max_tokens,
-        "smart_resize": apply_smart_resize,
-        "backend": "vllm",
-    }
-
-    if "inference_info" in dataset.column_names:
-        # Append to existing inference info
-        logger.info("Updating existing inference_info column")
-
-        def update_inference_info(example):
-            try:
-                existing_info = (
-                    json.loads(example["inference_info"])
-                    if example["inference_info"]
-                    else []
-                )
-            except (json.JSONDecodeError, TypeError):
-                existing_info = []
-
-            existing_info.append(inference_entry)
-            return {"inference_info": json.dumps(existing_info)}
-
-        dataset = dataset.map(update_inference_info)
-    else:
-        # Create new inference_info column
-        logger.info("Creating new inference_info column")
-        inference_list = [json.dumps([inference_entry])] * len(dataset)
-        dataset = dataset.add_column("inference_info", inference_list)
-
-    # Push to hub with retry and XET fallback
-    logger.info(f"Pushing to {output_dataset}")
-    max_retries = 3
-    for attempt in range(1, max_retries + 1):
-        try:
-            if attempt > 1:
-                logger.warning("Disabling XET (fallback to HTTP upload)")
-                os.environ["HF_HUB_DISABLE_XET"] = "1"
-            dataset.push_to_hub(
-                output_dataset,
-                private=private,
-                token=HF_TOKEN,
-                max_shard_size="500MB",
-                **({"config_name": config} if config else {}),
-                create_pr=create_pr,
-                commit_message=f"Add {MODEL_ID} OCR results ({len(dataset)} samples)"
-                + (f" [{config}]" if config else ""),
-            )
-            break
-        except Exception as e:
-            logger.error(f"Upload attempt {attempt}/{max_retries} failed: {e}")
-            if attempt < max_retries:
-                delay = 30 * (2 ** (attempt - 1))
-                logger.info(f"Retrying in {delay}s...")
-                time.sleep(delay)
-            else:
-                logger.error("All upload attempts failed. OCR results are lost.")
-                sys.exit(1)
-
-    # Create and push dataset card (skip when creating a PR to avoid touching main)
-    if not create_pr:
-        logger.info("Creating dataset card")
-        card_content = create_dataset_card(
-            source_dataset=input_dataset,
-            model=MODEL_ID,
-            task_mode=task_mode,
-            num_samples=len(dataset),
-            processing_time=processing_time_str,
-            batch_size=batch_size,
-            max_model_len=max_model_len,
-            max_tokens=max_tokens,
-            gpu_memory_utilization=gpu_memory_utilization,
-            temperature=temperature,
-            apply_smart_resize=apply_smart_resize,
-            image_column=image_column,
-            split=split,
-        )
-
-        card = DatasetCard(card_content)
-        card.push_to_hub(output_dataset, token=HF_TOKEN)
-
-    logger.info("✅ PaddleOCR-VL-1.6 processing complete!")
-    logger.info(
-        f"Dataset available at: https://huggingface.co/datasets/{output_dataset}"
-    )
-    logger.info(f"Processing time: {processing_time_str}")
-    logger.info(f"Task mode: {task_mode} - {TASK_DESCRIPTIONS[task_mode]}")
-
-    if verbose:
-        import importlib.metadata
-
-        logger.info("--- Resolved package versions ---")
-        for pkg in ["vllm", "transformers", "torch", "datasets", "pyarrow", "pillow"]:
-            try:
-                logger.info(f"  {pkg}=={importlib.metadata.version(pkg)}")
-            except importlib.metadata.PackageNotFoundError:
-                logger.info(f"  {pkg}: not installed")
-        logger.info("--- End versions ---")
-
-
-if __name__ == "__main__":
-    # Show example usage if no arguments
-    if len(sys.argv) == 1:
-        print("=" * 80)
-        print("PaddleOCR-VL-1.6 Document Processing")
-        print("=" * 80)
-        print("\nUltra-compact 0.9B OCR model (96.33% SOTA on OmniDocBench v1.6)")
-        print("\nFeatures:")
-        print("- 🎯 SOTA - 96.33% on OmniDocBench v1.6 (0.9B params)")
-        print("- 📝 OCR mode - General text extraction")
-        print("- 📊 Table mode - HTML table recognition")
-        print("- 📐 Formula mode - LaTeX mathematical notation")
-        print("- 📈 Chart mode - Structured chart analysis")
-        print("- 🔍 Spotting mode - Text spotting with localization")
-        print("- 🔖 Seal mode - Seal/stamp recognition")
-        print("- 🌍 Multilingual support")
-        print("- 🔧 Based on ERNIE-4.5 (unique architecture)")
-        print("\nTask Modes:")
-        for mode, description in TASK_DESCRIPTIONS.items():
-            print(f"  {mode:8} - {description}")
-        print("\nExample usage:")
-        print("\n1. Basic OCR (default mode):")
-        print("   uv run paddleocr-vl-1.6.py input-dataset output-dataset")
-        print("\n2. Table extraction:")
-        print("   uv run paddleocr-vl-1.6.py docs tables-extracted --task-mode table")
-        print("\n3. Formula recognition:")
-        print(
-            "   uv run paddleocr-vl-1.6.py papers formulas --task-mode formula --batch-size 32"
-        )
-        print("\n4. Chart analysis:")
-        print("   uv run paddleocr-vl-1.6.py diagrams charts-analyzed --task-mode chart")
-        print("\n5. Test with small sample:")
-        print("   uv run paddleocr-vl-1.6.py dataset test --max-samples 10 --shuffle")
-        print("\n6. Running on HF Jobs (image-mode required — see note below):")
-        print("   hf jobs uv run \\")
-        print("     --image vllm/vllm-openai:latest --flavor a100-large \\")
-        print(
-            "     --python /usr/bin/python3 -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages \\"
-        )
-        print("     -s HF_TOKEN \\")
-        print(
-            "     https://huggingface.co/datasets/uv-scripts/ocr/raw/main/paddleocr-vl-1.6.py \\"
-        )
-        print("       input-dataset output-dataset --task-mode ocr")
-        print("\n   NOTE: the default uv-script image has no nvcc, so vLLM's flashinfer")
-        print("   sampler crashes at warmup. The vllm/vllm-openai image ships the kernels.")
-        print("\n" + "=" * 80)
-        print("\nFor full help, run: uv run paddleocr-vl-1.6.py --help")
-        sys.exit(0)
-
-    parser = argparse.ArgumentParser(
-        description="Document processing using PaddleOCR-VL-1.6 (0.9B SOTA OCR model)",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog="""
-Task Modes:
-  ocr       General text extraction to markdown (default)
-  table     Table extraction to HTML format
-  formula   Mathematical formula recognition to LaTeX
-  chart     Chart and diagram analysis
-  spotting  Text spotting with localization
-  seal      Seal and stamp recognition
-
-Examples:
-  # Basic text OCR
-  uv run paddleocr-vl-1.6.py my-docs analyzed-docs
-
-  # Extract tables from documents
-  uv run paddleocr-vl-1.6.py papers tables --task-mode table
-
-  # Recognize mathematical formulas
-  uv run paddleocr-vl-1.6.py textbooks formulas --task-mode formula
-
-  # Analyze charts and diagrams
-  uv run paddleocr-vl-1.6.py reports charts --task-mode chart
-
-  # Test with random sampling
-  uv run paddleocr-vl-1.6.py large-dataset test --max-samples 50 --shuffle --task-mode ocr
-
-  # Disable smart resize for original resolution
-  uv run paddleocr-vl-1.6.py images output --no-smart-resize
-        """,
-    )
-
-    parser.add_argument("input_dataset", help="Input dataset ID from Hugging Face Hub")
-    parser.add_argument("output_dataset", help="Output dataset ID for Hugging Face Hub")
-    parser.add_argument(
-        "--image-column",
-        default="image",
-        help="Column containing images (default: image)",
-    )
-    parser.add_argument(
-        "--batch-size",
-        type=int,
-        default=16,
-        help="Batch size for processing (default: 16)",
-    )
-    parser.add_argument(
-        "--task-mode",
-        choices=list(TASK_MODES.keys()),
-        default="ocr",
-        help="Task type: ocr (default), table, formula, chart, spotting, or seal",
-    )
-    parser.add_argument(
-        "--max-model-len",
-        type=int,
-        default=8192,
-        help="Maximum model context length (default: 8192)",
-    )
-    parser.add_argument(
-        "--max-tokens",
-        type=int,
-        default=4096,
-        help="Maximum tokens to generate (default: 4096)",
-    )
-    parser.add_argument(
-        "--temperature",
-        type=float,
-        default=0.0,
-        help="Sampling temperature (default: 0.0 for deterministic)",
-    )
-    parser.add_argument(
-        "--gpu-memory-utilization",
-        type=float,
-        default=0.8,
-        help="GPU memory utilization (default: 0.8)",
-    )
-    parser.add_argument(
-        "--no-smart-resize",
-        action="store_true",
-        help="Disable PaddleOCR-VL's smart resize, use original image size",
-    )
-    parser.add_argument("--hf-token", help="Hugging Face API token")
-    parser.add_argument(
-        "--split", default="train", help="Dataset split to use (default: train)"
-    )
-    parser.add_argument(
-        "--max-samples",
-        type=int,
-        help="Maximum number of samples to process (for testing)",
-    )
-    parser.add_argument(
-        "--private", action="store_true", help="Make output dataset private"
-    )
-    parser.add_argument(
-        "--shuffle", action="store_true", help="Shuffle dataset before processing"
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=42,
-        help="Random seed for shuffling (default: 42)",
-    )
-    parser.add_argument(
-        "--output-column",
-        help="Column name for output (default: markdown)",
-    )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
-    parser.add_argument(
-        "--config",
-        help="Config/subset name when pushing to Hub (for benchmarking multiple models in one repo)",
-    )
-    parser.add_argument(
-        "--create-pr",
-        action="store_true",
-        help="Create a pull request instead of pushing directly (for parallel benchmarking)",
-    )
-    parser.add_argument(
-        "--verbose",
-        action="store_true",
-        help="Log resolved package versions after processing (useful for pinning deps)",
-    )
-
-    args = parser.parse_args()
-
-    main(
-        input_dataset=args.input_dataset,
-        output_dataset=args.output_dataset,
-        image_column=args.image_column,
-        batch_size=args.batch_size,
-        task_mode=args.task_mode,
-        max_model_len=args.max_model_len,
-        max_tokens=args.max_tokens,
-        temperature=args.temperature,
-        gpu_memory_utilization=args.gpu_memory_utilization,
-        apply_smart_resize=not args.no_smart_resize,
-        hf_token=args.hf_token,
-        split=args.split,
-        max_samples=args.max_samples,
-        private=args.private,
-        shuffle=args.shuffle,
-        seed=args.seed,
-        output_column=args.output_column,
-        overwrite=args.overwrite,
-        config=args.config,
-        create_pr=args.create_pr,
-        verbose=args.verbose,
-    )

paddleocr-vl.py

@@ -56,10 +56,6 @@ from huggingface_hub import DatasetCard, login
 from PIL import Image
 from toolz import partition_all
 from tqdm.auto import tqdm
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
 from vllm import LLM, SamplingParams
 
 logging.basicConfig(level=logging.INFO)
@@ -93,28 +89,6 @@ def check_cuda_availability():
         logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
 
 
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
 def smart_resize(
     height: int,
     width: int,
@@ -359,7 +333,6 @@ def main(
     shuffle: bool = False,
     seed: int = 42,
     output_column: str = None,
-    overwrite: bool = False,
     verbose: bool = False,
 ):
     """Process images from HF dataset through PaddleOCR-VL model."""
@@ -408,9 +381,6 @@ def main(
             f"Column '{image_column}' not found. Available: {dataset.column_names}"
         )
 
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
     # Shuffle if requested
     if shuffle:
         logger.info(f"Shuffling dataset with seed {seed}")
@@ -732,12 +702,6 @@ Examples:
         "--output-column",
         help="Column name for output (default: markdown)",
     )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
     parser.add_argument(
         "--verbose",
         action="store_true",
@@ -764,6 +728,5 @@ Examples:
         shuffle=args.shuffle,
         seed=args.seed,
         output_column=args.output_column,
-        overwrite=args.overwrite,
         verbose=args.verbose,
     )

pp-doclayout.py

@@ -1,1226 +0,0 @@
-# /// script
-# requires-python = ">=3.10"
-# dependencies = [
-#     "paddlepaddle-gpu>=3.0.0",
-#     "paddleocr>=3.0.0",
-#     "opencv-contrib-python-headless",
-#     "datasets>=4.0.0",
-#     "huggingface-hub>=1.6.0",
-#     "pyarrow>=15.0",
-#     "pillow",
-#     "numpy",
-#     "tqdm",
-# ]
-#
-# [tool.uv]
-# # PaddleOCR/PaddleX pull in opencv-contrib-python (full) which needs system
-# # libGL.so.1 — not present in the slim uv-on-bookworm image used by HF Jobs.
-# # Swap to the headless cv2 variant (same `import cv2`, no GUI deps).
-# override-dependencies = [
-#     "opencv-contrib-python ; python_version < '0'",
-#     "opencv-python ; python_version < '0'",
-# ]
-#
-# [[tool.uv.index]]
-# name = "paddle"
-# url = "https://www.paddlepaddle.org.cn/packages/stable/cu126/"
-# explicit = true
-#
-# [tool.uv.sources]
-# paddlepaddle-gpu = { index = "paddle" }
-# ///
-
-"""
-Detect document layout regions (text/title/table/figure/formula/...) with PP-DocLayout-L.
-
-Runs PaddleOCR's PP-DocLayout-L (or M / S / plus-L variant) over an image source
-and emits per-image bounding-box predictions. Unlike the OCR scripts in this repo
-this does NOT extract text — it only locates and classifies regions.
-
-Source can be:
-  - HF dataset repo (default): "namespace/dataset"
-  - HF bucket of image files: "hf://buckets/namespace/bucket/optional/prefix"
-
-Sink can be:
-  - HF dataset repo (default): "namespace/dataset" (one push at end + dataset card)
-  - HF bucket: "hf://buckets/namespace/bucket/run-name" (incremental parquet
-    shards, resumable, no git overhead)
-
-Output schema (column `layout` is a JSON string):
-  [{"bbox": [x1, y1, x2, y2], "label": "text", "score": 0.97, "cls_id": 2}, ...]
-
-Coordinates are in the original input-image pixel space.
-
-Example commands:
-
-  # Dataset -> dataset (smoke on L4)
-  hf jobs uv run --flavor l4x1 -s HF_TOKEN https://huggingface.co/datasets/uv-scripts/ocr/raw/main/pp-doclayout.py \\
-      davanstrien/ufo-ColPali pp-doclayout-smoke \\
-      --max-samples 3 --shuffle --seed 42 --private
-
-  # Dataset -> bucket (incremental shards, resumable)
-  hf buckets create davanstrien/pp-doclayout-scratch --exist-ok
-  hf jobs uv run --flavor l4x1 -s HF_TOKEN https://huggingface.co/datasets/uv-scripts/ocr/raw/main/pp-doclayout.py \\
-      davanstrien/ufo-ColPali \\
-      hf://buckets/davanstrien/pp-doclayout-scratch/run1 \\
-      --max-samples 20 --shard-size 5
-
-  # Bucket of images -> dataset
-  hf jobs uv run --flavor l4x1 -s HF_TOKEN https://huggingface.co/datasets/uv-scripts/ocr/raw/main/pp-doclayout.py \\
-      hf://buckets/davanstrien/pp-doclayout-images \\
-      pp-doclayout-from-bucket --private
-"""
-
-import argparse
-import io
-import json
-import logging
-import os
-import sys
-import time
-from dataclasses import dataclass
-from datetime import datetime, timezone
-from pathlib import Path
-from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
-
-import numpy as np
-from PIL import Image, UnidentifiedImageError
-from tqdm.auto import tqdm
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-# ---------------------------------------------------------------------------
-# Constants
-# ---------------------------------------------------------------------------
-
-VALID_MODELS = [
-    "PP-DocLayout-L",
-    "PP-DocLayout-M",
-    "PP-DocLayout-S",
-    "PP-DocLayout_plus-L",
-]
-
-MODEL_SIZES = {
-    "PP-DocLayout-L": "~123M params (RT-DETR-L backbone)",
-    "PP-DocLayout-M": "~22M params (PicoDet-M)",
-    "PP-DocLayout-S": "~4M params (PicoDet-S)",
-    "PP-DocLayout_plus-L": "~123M params, 20-class plus variant",
-}
-
-IMAGE_EXTENSIONS = {
-    ".jpg", ".jpeg", ".png", ".tif", ".tiff", ".webp", ".bmp", ".jp2", ".j2k",
-}
-
-BUCKET_PREFIX = "hf://buckets/"
-
-
-# ---------------------------------------------------------------------------
-# URL helpers
-# ---------------------------------------------------------------------------
-
-
-def is_bucket_url(s: str) -> bool:
-    return s.startswith(BUCKET_PREFIX)
-
-
-def parse_bucket_url(url: str) -> Tuple[str, str]:
-    """Split `hf://buckets/ns/bucket/path/in/bucket` into (`ns/bucket`, `path/in/bucket`)."""
-    if not is_bucket_url(url):
-        raise ValueError(f"Not a bucket URL: {url}")
-    rest = url[len(BUCKET_PREFIX) :].strip("/")
-    parts = rest.split("/", 2)
-    if len(parts) < 2:
-        raise ValueError(
-            f"Bucket URL must include namespace and bucket name: {url}"
-        )
-    bucket_id = f"{parts[0]}/{parts[1]}"
-    prefix = parts[2] if len(parts) > 2 else ""
-    return bucket_id, prefix
-
-
-# ---------------------------------------------------------------------------
-# Image helpers
-# ---------------------------------------------------------------------------
-
-
-def to_pil(image: Union[Image.Image, Dict[str, Any], str, bytes]) -> Image.Image:
-    if isinstance(image, Image.Image):
-        return image.convert("RGB")
-    if isinstance(image, dict) and "bytes" in image:
-        return Image.open(io.BytesIO(image["bytes"])).convert("RGB")
-    if isinstance(image, (bytes, bytearray)):
-        return Image.open(io.BytesIO(image)).convert("RGB")
-    if isinstance(image, str):
-        return Image.open(image).convert("RGB")
-    raise ValueError(f"Unsupported image type: {type(image)}")
-
-
-def pil_to_array(pil_img: Image.Image) -> np.ndarray:
-    """RGB PIL -> uint8 ndarray. PaddleOCR's predict() accepts numpy arrays directly."""
-    return np.asarray(pil_img, dtype=np.uint8)
-
-
-# ---------------------------------------------------------------------------
-# Result extraction
-# ---------------------------------------------------------------------------
-
-
-def extract_detections(result: Any) -> List[Dict[str, Any]]:
-    """Pull a clean list of detections out of a paddleocr LayoutDetection result."""
-    payload = result.json
-    res = payload.get("res", payload) if isinstance(payload, dict) else {}
-    boxes = res.get("boxes", []) if isinstance(res, dict) else []
-    detections = []
-    for box in boxes:
-        coord = box.get("coordinate") or box.get("bbox") or []
-        coord = [float(x) for x in coord]
-        detections.append(
-            {
-                "bbox": coord,
-                "label": box.get("label"),
-                "score": float(box.get("score", 0.0)),
-                "cls_id": int(box.get("cls_id", -1)),
-            }
-        )
-    return detections
-
-
-# ---------------------------------------------------------------------------
-# Sources
-# ---------------------------------------------------------------------------
-
-
-@dataclass
-class SourceItem:
-    key: str  # stable identifier per image (used for dedup/resume)
-    image: Optional[Image.Image]  # None for an unreadable row (placeholder)
-    extras: Dict[str, Any]  # original row fields (only populated for dataset source)
-
-
-def iter_dataset_images(
-    dataset_id: str,
-    image_column: str,
-    split: str,
-    shuffle: bool,
-    seed: int,
-    max_samples: Optional[int],
-):
-    """Iterate (key, PIL) pairs from an HF dataset repo.
-
-    Returns: (iterator, total, dataset_reference). The dataset reference is the
-    post-shuffle/post-select Dataset, kept around so the dataset-repo sink can
-    `add_column("layout", ...)` and preserve the original schema (especially
-    Image-type columns).
-    """
-    from datasets import load_dataset
-
-    logger.info(f"Loading dataset: {dataset_id} (split={split})")
-    ds = load_dataset(dataset_id, split=split)
-
-    if image_column not in ds.column_names:
-        raise ValueError(
-            f"Column '{image_column}' not found. Available: {ds.column_names}"
-        )
-
-    if shuffle:
-        logger.info(f"Shuffling with seed {seed}")
-        ds = ds.shuffle(seed=seed)
-    if max_samples:
-        ds = ds.select(range(min(max_samples, len(ds))))
-        logger.info(f"Limited to {len(ds)} samples")
-
-    total = len(ds)
-
-    def gen() -> Iterator[SourceItem]:
-        failed = 0
-        for i in range(total):
-            try:
-                row = ds[i]
-                image = to_pil(row[image_column])
-            except (UnidentifiedImageError, OSError) as e:
-                # Still yield a placeholder so the output row stays aligned with
-                # the source row (the dataset sink writes layouts positionally).
-                failed += 1
-                logger.warning(
-                    f"Unreadable image at row {i}: "
-                    f"{type(e).__name__}: {e} — writing empty layout"
-                )
-                yield SourceItem(
-                    key=f"row-{i:08d}",
-                    image=None,
-                    extras={"failed": True},
-                )
-                continue
-            yield SourceItem(
-                key=f"row-{i:08d}",
-                image=image,
-                extras={},  # original schema is preserved by the sink via the dataset ref
-            )
-        if failed:
-            logger.info(f"{failed} unreadable image(s) written as empty layouts")
-
-    return gen(), total, ds
-
-
-SOURCE_PATHS_SNAPSHOT = "_source_paths.json"
-
-
-def _bucket_snapshot_path(output_url: str) -> Tuple[str, str]:
-    """Return (bucket_id, key) for the source-paths snapshot inside an output bucket."""
-    out_bucket_id, out_prefix = parse_bucket_url(output_url)
-    snapshot_key = (
-        f"{out_prefix}/{SOURCE_PATHS_SNAPSHOT}".lstrip("/")
-        if out_prefix
-        else SOURCE_PATHS_SNAPSHOT
-    )
-    return out_bucket_id, snapshot_key
-
-
-def iter_bucket_images(
-    bucket_url: str,
-    shuffle: bool,
-    seed: int,
-    max_samples: Optional[int],
-    hf_token: Optional[str],
-    output_url: Optional[str] = None,
-) -> Tuple[Iterator[SourceItem], int]:
-    """Glob image files under a bucket prefix and stream them via HfFileSystem.
-
-    If `output_url` is a bucket, the resolved source-path list is snapshotted to
-    `<output>/_source_paths.json` on first run. Subsequent runs against the same
-    output prefix reuse that snapshot, so resume stays consistent even if the
-    source bucket grows or `--shuffle`/`--max-samples` would otherwise pick a
-    different subset on the second run.
-    """
-    from huggingface_hub import HfApi, HfFileSystem
-
-    bucket_id, prefix = parse_bucket_url(bucket_url)
-    fs = HfFileSystem(token=hf_token)
-    base = f"{BUCKET_PREFIX}{bucket_id}/{prefix}".rstrip("/")
-
-    snapshot_bucket_id: Optional[str] = None
-    snapshot_key: Optional[str] = None
-    cached_paths: Optional[List[str]] = None
-
-    if output_url and is_bucket_url(output_url):
-        snapshot_bucket_id, snapshot_key = _bucket_snapshot_path(output_url)
-        snapshot_url = f"{BUCKET_PREFIX}{snapshot_bucket_id}/{snapshot_key}"
-        try:
-            with fs.open(snapshot_url, "rb") as f:
-                snapshot = json.load(f)
-            if snapshot.get("source_url") != bucket_url:
-                logger.warning(
-                    f"Output prefix already has a snapshot referencing a "
-                    f"different source ({snapshot.get('source_url')!r} vs "
-                    f"{bucket_url!r}). Ignoring and re-listing."
-                )
-            else:
-                cached_paths = snapshot["paths"]
-                logger.info(
-                    f"Reusing existing snapshot of {len(cached_paths)} source paths "
-                    f"(written {snapshot.get('created_at', 'unknown')})"
-                )
-        except FileNotFoundError:
-            pass
-        except Exception as e:
-            logger.warning(f"Could not read existing snapshot ({e}); re-listing.")
-
-    if cached_paths is not None:
-        all_paths = cached_paths
-    else:
-        logger.info(f"Listing images under {base}")
-        all_paths = []
-        try:
-            for entry in fs.find(base, detail=False):
-                ext = Path(entry).suffix.lower()
-                if ext in IMAGE_EXTENSIONS:
-                    all_paths.append(entry)
-        except FileNotFoundError as e:
-            raise ValueError(f"Bucket prefix not found: {base}") from e
-
-        if not all_paths:
-            raise ValueError(
-                f"No image files (any of {sorted(IMAGE_EXTENSIONS)}) under {base}"
-            )
-
-        all_paths.sort()
-        if shuffle:
-            rng = np.random.default_rng(seed)
-            rng.shuffle(all_paths)
-        if max_samples:
-            all_paths = all_paths[:max_samples]
-
-        # Persist the chosen list so resume runs see exactly this set.
-        if snapshot_bucket_id is not None and snapshot_key is not None:
-            api = HfApi(token=hf_token)
-            payload = {
-                "source_url": bucket_url,
-                "shuffle": shuffle,
-                "seed": seed,
-                "max_samples": max_samples,
-                "created_at": datetime.now(timezone.utc).isoformat(),
-                "paths": all_paths,
-            }
-            api.batch_bucket_files(
-                snapshot_bucket_id,
-                add=[(json.dumps(payload).encode(), snapshot_key)],
-                token=hf_token,
-            )
-            logger.info(
-                f"Wrote source-path snapshot ({len(all_paths)} paths) to "
-                f"hf://buckets/{snapshot_bucket_id}/{snapshot_key}"
-            )
-
-    total = len(all_paths)
-    logger.info(f"Found {total} images in bucket")
-
-    def key_for(path: str) -> str:
-        # Use the full bucket path (`buckets/<id>/<rel>`) as returned by
-        # fs.find. This is stable across reruns (so resume works), and the
-        # stored value in `source_path` is fully addressable — open via
-        # HfFileSystem directly with `hf://` re-prepended.
-        return path
-
-    def gen() -> Iterator[SourceItem]:
-        skipped = 0
-        for path in all_paths:
-            try:
-                with fs.open(path, "rb") as f:
-                    data = f.read()
-                image = to_pil(data)
-            except (UnidentifiedImageError, OSError) as e:
-                skipped += 1
-                logger.warning(
-                    f"Skipping unreadable image {path}: "
-                    f"{type(e).__name__}: {e}"
-                )
-                continue
-            yield SourceItem(
-                key=key_for(path),
-                image=image,
-                extras={"__source_path": key_for(path)},
-            )
-        if skipped:
-            logger.info(f"Skipped {skipped} unreadable image(s) total")
-
-    return gen(), total
-
-
-# ---------------------------------------------------------------------------
-# Sinks
-# ---------------------------------------------------------------------------
-
-
-class DatasetRepoSink:
-    """Buffer all results in memory, push once at end with dataset card + inference_info.
-
-    Two modes:
-      - `original_dataset` provided (dataset-repo source): preserve the source
-        schema (including Image-type columns) and just `add_column("layout", ...)`.
-      - `original_dataset` is None (bucket-image source): build a Dataset from
-        collected rows containing __source_path + layout.
-    """
-
-    def __init__(
-        self,
-        repo_id: str,
-        *,
-        hf_token: Optional[str],
-        private: bool,
-        config: Optional[str],
-        create_pr: bool,
-        source_id: str,
-        original_dataset=None,
-        overwrite: bool = False,
-    ):
-        self.repo_id = repo_id
-        self.hf_token = hf_token
-        self.private = private
-        self.config = config
-        self.create_pr = create_pr
-        self.source_id = source_id
-        self.original_dataset = original_dataset
-        self.overwrite = overwrite
-        # Used when original_dataset is None: row-by-row buffer.
-        self._rows: List[Dict[str, Any]] = []
-        # Used when original_dataset is set: ordered layouts aligned with dataset rows.
-        self._layouts: List[str] = []
-
-    @property
-    def kind(self) -> str:
-        return "dataset"
-
-    def already_done(self) -> set:
-        return set()  # dataset sink does a single push, no resume
-
-    def write(self, key: str, layout: List[Dict[str, Any]], extras: Dict[str, Any]) -> None:
-        layout_json = json.dumps(layout, ensure_ascii=False)
-        if self.original_dataset is not None:
-            self._layouts.append(layout_json)
-            return
-        row = {"__source_key": key, "layout": layout_json}
-        for k, v in extras.items():
-            if isinstance(v, (str, int, float, bool)) or v is None:
-                row[k] = v
-        self._rows.append(row)
-
-    def finalize(self, model_id: str, args_dict: Dict[str, Any]) -> None:
-        from datasets import Dataset
-
-        if self.original_dataset is not None:
-            if len(self._layouts) != len(self.original_dataset):
-                logger.warning(
-                    f"Layout count ({len(self._layouts)}) != dataset rows "
-                    f"({len(self.original_dataset)}); padding with empty layouts."
-                )
-                # Pad to keep add_column happy.
-                while len(self._layouts) < len(self.original_dataset):
-                    self._layouts.append("[]")
-            base = self.original_dataset
-            if "layout" in base.column_names:
-                if not self.overwrite:
-                    raise ValueError(
-                        "Output column 'layout' already exists in the input dataset; "
-                        "pass --overwrite to replace it."
-                    )
-                logger.warning("--overwrite: replacing existing column 'layout'")
-                base = base.remove_columns("layout")
-            ds = base.add_column("layout", self._layouts)
-        else:
-            if not self._rows:
-                logger.warning("No rows produced; nothing to push.")
-                return
-            ds = Dataset.from_list(self._rows)
-            if "__source_key" in ds.column_names:
-                ds = ds.rename_column("__source_key", "source_path")
-
-        inference_entry = build_inference_entry(model_id, args_dict)
-
-        if "inference_info" in ds.column_names:
-            logger.info("Updating existing inference_info column")
-
-            def _update(example):
-                try:
-                    existing = (
-                        json.loads(example["inference_info"])
-                        if example["inference_info"]
-                        else []
-                    )
-                except (json.JSONDecodeError, TypeError):
-                    existing = []
-                existing.append(inference_entry)
-                return {"inference_info": json.dumps(existing)}
-
-            ds = ds.map(_update)
-        else:
-            ds = ds.add_column(
-                "inference_info", [json.dumps([inference_entry])] * len(ds)
-            )
-
-        logger.info(f"Pushing {len(ds)} rows to {self.repo_id}")
-        push_kwargs = {
-            "private": self.private,
-            "token": self.hf_token,
-            "max_shard_size": "500MB",
-            "create_pr": self.create_pr,
-            "commit_message": f"Add PP-DocLayout layout predictions ({len(ds)} samples)"
-            + (f" [{self.config}]" if self.config else ""),
-        }
-        if self.config:
-            push_kwargs["config_name"] = self.config
-
-        max_retries = 3
-        for attempt in range(1, max_retries + 1):
-            try:
-                if attempt > 1:
-                    logger.warning("Disabling XET (fallback to HTTP upload)")
-                    os.environ["HF_HUB_DISABLE_XET"] = "1"
-                ds.push_to_hub(self.repo_id, **push_kwargs)
-                break
-            except Exception as e:
-                logger.error(f"Upload attempt {attempt}/{max_retries} failed: {e}")
-                if attempt == max_retries:
-                    logger.error("All upload attempts failed.")
-                    raise
-                time.sleep(30 * (2 ** (attempt - 1)))
-
-        # Dataset card
-        from huggingface_hub import DatasetCard
-
-        card = DatasetCard(
-            create_dataset_card(
-                source=self.source_id,
-                model_name=args_dict["model_name"],
-                num_samples=len(ds),
-                processing_time=args_dict["processing_time"],
-                output_column="layout",
-                threshold=args_dict["threshold"],
-                layout_nms=args_dict["layout_nms"],
-            )
-        )
-        card.push_to_hub(self.repo_id, token=self.hf_token)
-        logger.info(
-            f"Done: https://huggingface.co/datasets/{self.repo_id}"
-        )
-
-
-class BucketShardSink:
-    """Write incremental parquet shards to a bucket prefix. Resumable."""
-
-    METADATA_FILE = "_metadata.json"
-    SHARD_PATTERN = "shard-{:05d}.parquet"
-
-    def __init__(
-        self,
-        bucket_url: str,
-        *,
-        hf_token: Optional[str],
-        shard_size: int,
-        include_images: bool,
-        resume: bool,
-        source_id: str,
-    ):
-        from huggingface_hub import HfApi, HfFileSystem, create_bucket
-
-        self.bucket_url = bucket_url
-        self.bucket_id, self.prefix = parse_bucket_url(bucket_url)
-        self.hf_token = hf_token
-        self.shard_size = shard_size
-        self.include_images = include_images
-        self.resume = resume
-        self.source_id = source_id
-
-        self._api = HfApi(token=hf_token)
-        self._fs = HfFileSystem(token=hf_token)
-
-        # Make sure the bucket exists. Path inside the bucket is created lazily on first write.
-        try:
-            create_bucket(self.bucket_id, exist_ok=True, token=hf_token)
-        except Exception as e:
-            # If we don't have create rights but the bucket already exists, that's fine.
-            logger.warning(f"create_bucket('{self.bucket_id}') warning: {e}")
-
-        self._buffer: List[Dict[str, Any]] = []
-        self._next_shard_idx = self._discover_next_shard_idx()
-        self._completed_keys = self._discover_completed_keys() if resume else set()
-        if self._completed_keys:
-            logger.info(
-                f"Resume: found {len(self._completed_keys)} already-processed keys, will skip them"
-            )
-
-    @property
-    def kind(self) -> str:
-        return "bucket"
-
-    def already_done(self) -> set:
-        return self._completed_keys
-
-    # --- internal helpers ---
-
-    def _shard_path(self, idx: int) -> str:
-        return self._join(self.SHARD_PATTERN.format(idx))
-
-    def _join(self, name: str) -> str:
-        return f"{self.prefix}/{name}".lstrip("/") if self.prefix else name
-
-    def _list_existing_shards(self) -> List[str]:
-        try:
-            tree = self._api.list_bucket_tree(
-                self.bucket_id, prefix=self.prefix or None, recursive=True
-            )
-        except Exception:
-            return []
-        shards: List[str] = []
-        for item in tree:
-            path = getattr(item, "path", None)
-            ftype = getattr(item, "type", None)
-            if not path or ftype not in (None, "file"):
-                continue
-            base = Path(path).name
-            if base.startswith("shard-") and base.endswith(".parquet"):
-                shards.append(path)
-        return sorted(shards)
-
-    def _discover_next_shard_idx(self) -> int:
-        shards = self._list_existing_shards()
-        max_idx = -1
-        for s in shards:
-            stem = Path(s).stem  # shard-00007
-            try:
-                max_idx = max(max_idx, int(stem.split("-")[-1]))
-            except ValueError:
-                continue
-        return max_idx + 1
-
-    def _discover_completed_keys(self) -> set:
-        import pyarrow.parquet as pq
-
-        keys: set = set()
-        for shard_path in self._list_existing_shards():
-            full = f"{BUCKET_PREFIX}{self.bucket_id}/{shard_path}"
-            try:
-                with self._fs.open(full, "rb") as f:
-                    table = pq.read_table(f, columns=["__source_key"])
-                keys.update(table.column("__source_key").to_pylist())
-            except Exception as e:
-                logger.warning(f"Could not read keys from {shard_path}: {e}")
-        return keys
-
-    def _flush(self) -> None:
-        if not self._buffer:
-            return
-        import pyarrow as pa
-        import pyarrow.parquet as pq
-
-        # Build a stable schema. Skip the image column if not requested.
-        columns = ["__source_key", "layout"]
-        if self.include_images:
-            columns.append("__image_bytes")
-        # Carry through any extra string-coercible fields (e.g. __source_path).
-        extra_keys = sorted(
-            {k for row in self._buffer for k in row.keys() if k not in columns}
-        )
-        columns.extend(extra_keys)
-
-        table_dict = {c: [row.get(c) for row in self._buffer] for c in columns}
-        # pyarrow infers types from python objects; strings/bytes/lists handled fine.
-        table = pa.Table.from_pydict(table_dict)
-
-        buf = io.BytesIO()
-        pq.write_table(table, buf, compression="zstd")
-        data = buf.getvalue()
-
-        shard_remote = self._shard_path(self._next_shard_idx)
-        logger.info(
-            f"Writing shard {self._next_shard_idx} ({len(self._buffer)} rows, "
-            f"{len(data) / 1024 / 1024:.1f} MiB) to {shard_remote}"
-        )
-        self._api.batch_bucket_files(
-            self.bucket_id, add=[(data, shard_remote)], token=self.hf_token
-        )
-        self._next_shard_idx += 1
-        self._buffer.clear()
-
-    def write(self, key: str, layout: List[Dict[str, Any]], extras: Dict[str, Any]) -> None:
-        row: Dict[str, Any] = {
-            "__source_key": key,
-            "layout": json.dumps(layout, ensure_ascii=False),
-        }
-        if self.include_images and "__image_bytes" in extras:
-            row["__image_bytes"] = extras["__image_bytes"]
-        # Pass through string/numeric extras (skip raw PIL Image objects which
-        # the dataset source never injects directly into extras anyway).
-        for k, v in extras.items():
-            if k in row or k == "__image_bytes":
-                continue
-            if isinstance(v, (str, int, float, bool)) or v is None:
-                row[k] = v
-        self._buffer.append(row)
-        if len(self._buffer) >= self.shard_size:
-            self._flush()
-
-    def finalize(self, model_id: str, args_dict: Dict[str, Any]) -> None:
-        # Flush trailing rows.
-        self._flush()
-        # Write/update the metadata file alongside the shards.
-        meta = {
-            "model_id": model_id,
-            "model_name": args_dict["model_name"],
-            "task_mode": "layout-detection",
-            "source": self.source_id,
-            "threshold": args_dict["threshold"],
-            "layout_nms": args_dict["layout_nms"],
-            "shard_size": args_dict["shard_size"],
-            "include_images": self.include_images,
-            "last_run_at": datetime.now(timezone.utc).isoformat(),
-            "processing_time": args_dict.get("processing_time"),
-        }
-        meta_bytes = json.dumps(meta, indent=2).encode("utf-8")
-        meta_path = self._join(self.METADATA_FILE)
-        self._api.batch_bucket_files(
-            self.bucket_id, add=[(meta_bytes, meta_path)], token=self.hf_token
-        )
-        logger.info(
-            f"Done: https://huggingface.co/buckets/{self.bucket_id}"
-            + (f"/{self.prefix}" if self.prefix else "")
-        )
-
-
-# ---------------------------------------------------------------------------
-# inference_info + dataset card
-# ---------------------------------------------------------------------------
-
-
-def build_inference_entry(model_id: str, args_dict: Dict[str, Any]) -> Dict[str, Any]:
-    return {
-        "model_id": "PaddlePaddle/" + args_dict["model_name"],
-        "model_name": args_dict["model_name"],
-        "model_size": MODEL_SIZES.get(args_dict["model_name"], "unknown"),
-        "task_mode": "layout-detection",
-        "column_name": "layout",
-        "timestamp": datetime.now(timezone.utc).isoformat(),
-        "threshold": args_dict["threshold"],
-        "layout_nms": args_dict["layout_nms"],
-        "backend": "paddleocr",
-    }
-
-
-def create_dataset_card(
-    source: str,
-    model_name: str,
-    num_samples: int,
-    processing_time: str,
-    output_column: str,
-    threshold: float,
-    layout_nms: bool,
-) -> str:
-    """Render the dataset card markdown for the dataset-repo sink."""
-    if is_bucket_url(source):
-        source_link = f"[{source}]({source})"
-    else:
-        source_link = f"[{source}](https://huggingface.co/datasets/{source})"
-
-    return f"""---
-tags:
-- layout-detection
-- document-processing
-- paddleocr
-- pp-doclayout
-- uv-script
-- generated
-viewer: false
----
-
-# Layout detection with {model_name}
-
-Bounding-box layout predictions for images from {source_link}, produced by
-PaddleOCR's [{model_name}](https://huggingface.co/PaddlePaddle/{model_name}).
-
-## Processing details
-
-- **Source**: {source_link}
-- **Model**: PaddlePaddle/{model_name} ({MODEL_SIZES.get(model_name, "unknown")})
-- **Samples**: {num_samples:,}
-- **Processing time**: {processing_time}
-- **Processing date**: {datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC")}
-- **Confidence threshold**: {threshold}
-- **Layout NMS**: {"on" if layout_nms else "off"}
-- **Output column**: `{output_column}` (JSON-encoded list of detections)
-
-## Schema
-
-Each row contains the original columns plus:
-
-- `{output_column}`: JSON string. List of detections:
-  ```json
-  [
-    {{"bbox": [x1, y1, x2, y2], "label": "text", "score": 0.97, "cls_id": 2}},
-    {{"bbox": [x1, y1, x2, y2], "label": "table", "score": 0.92, "cls_id": 5}}
-  ]
-  ```
-  Coordinates are in **original input-image pixel space** (top-left origin,
-  `[xmin, ymin, xmax, ymax]`).
-- `inference_info`: JSON list tracking every model that has been applied to
-  this dataset (appended on each run).
-
-## Usage
-
-```python
-import json
-from datasets import load_dataset
-
-ds = load_dataset("{{output_dataset_id}}", split="train")
-detections = json.loads(ds[0]["{output_column}"])
-for det in detections:
-    print(det["label"], det["score"], det["bbox"])
-```
-
-## Reproduction
-
-```bash
-hf jobs uv run --flavor l4x1 -s HF_TOKEN \\
-    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/pp-doclayout.py \\
-    {source} <output> --model-name {model_name}
-```
-
-Generated with [UV Scripts](https://huggingface.co/uv-scripts).
-"""
-
-
-# ---------------------------------------------------------------------------
-# Main
-# ---------------------------------------------------------------------------
-
-
-def resolve_device(device: str) -> str:
-    if device == "gpu":
-        try:
-            import paddle  # noqa: F401
-
-            if paddle.device.is_compiled_with_cuda() and paddle.device.cuda.device_count() > 0:
-                logger.info(
-                    f"GPU available: {paddle.device.cuda.device_count()} device(s)"
-                )
-                return "gpu"
-            logger.warning("No CUDA GPU detected; falling back to CPU.")
-            return "cpu"
-        except Exception as e:
-            logger.warning(f"GPU check failed ({e}); falling back to CPU.")
-            return "cpu"
-    return device
-
-
-def main(args: argparse.Namespace) -> None:
-    from huggingface_hub import login
-
-    start_time = datetime.now()
-    hf_token = args.hf_token or os.environ.get("HF_TOKEN")
-    if hf_token:
-        login(token=hf_token)
-
-    device = resolve_device(args.device)
-
-    # ---------- source ----------
-    original_dataset = None
-    if is_bucket_url(args.input_source):
-        src_iter, total = iter_bucket_images(
-            args.input_source,
-            shuffle=args.shuffle,
-            seed=args.seed,
-            max_samples=args.max_samples,
-            hf_token=hf_token,
-            output_url=args.output_target,
-        )
-    else:
-        src_iter, total, original_dataset = iter_dataset_images(
-            args.input_source,
-            image_column=args.image_column,
-            split=args.split,
-            shuffle=args.shuffle,
-            seed=args.seed,
-            max_samples=args.max_samples,
-        )
-        # Fail fast, before minutes of inference, if the 'layout' output column would
-        # collide with an existing input column. --overwrite opts in to replacing it.
-        if original_dataset is not None and "layout" in original_dataset.column_names:
-            if not args.overwrite:
-                logger.error(
-                    "Output column 'layout' already exists in the input dataset "
-                    f"(columns: {original_dataset.column_names})."
-                )
-                logger.error("Pass --overwrite to replace it.")
-                sys.exit(1)
-            logger.warning("--overwrite: will replace existing column 'layout'")
-
-    # ---------- sink ----------
-    if is_bucket_url(args.output_target):
-        sink: Union[BucketShardSink, DatasetRepoSink] = BucketShardSink(
-            args.output_target,
-            hf_token=hf_token,
-            shard_size=args.shard_size,
-            include_images=args.include_images,
-            resume=not args.no_resume,
-            source_id=args.input_source,
-        )
-    else:
-        sink = DatasetRepoSink(
-            args.output_target,
-            hf_token=hf_token,
-            private=args.private,
-            config=args.config,
-            create_pr=args.create_pr,
-            source_id=args.input_source,
-            original_dataset=original_dataset,
-            overwrite=args.overwrite,
-        )
-
-    completed = sink.already_done()
-
-    # ---------- model ----------
-    if args.model_name not in VALID_MODELS:
-        raise ValueError(
-            f"Invalid model {args.model_name!r}. Choose from: {VALID_MODELS}"
-        )
-    logger.info(f"Loading PaddleOCR LayoutDetection model: {args.model_name} on {device}")
-    # PaddleX gates `import cv2` at module load time on
-    # `is_dep_available("opencv-contrib-python")`, which checks
-    # `importlib.metadata.version(...)`. We ship `opencv-contrib-python-headless`
-    # (same `cv2`, no system libGL.so.1 needed) — but that's a different
-    # distribution name, so the gate fails and `cv2` is never bound, causing
-    # NameErrors deep inside paddlex modules. Patch the metadata lookup to
-    # alias the GUI cv2 distros to the headless variant before importing
-    # paddleocr; this lets paddlex's own `import cv2` succeed naturally.
-    import importlib.metadata as _metadata
-
-    _orig_metadata_version = _metadata.version
-
-    def _patched_metadata_version(dep_name):
-        if dep_name in ("opencv-contrib-python", "opencv-python"):
-            for headless_alias in (
-                "opencv-contrib-python-headless",
-                "opencv-python-headless",
-            ):
-                try:
-                    return _orig_metadata_version(headless_alias)
-                except _metadata.PackageNotFoundError:
-                    continue
-        return _orig_metadata_version(dep_name)
-
-    _metadata.version = _patched_metadata_version
-
-    from paddleocr import LayoutDetection
-
-    model = LayoutDetection(model_name=args.model_name, device=device)
-
-    # ---------- loop ----------
-    processed = 0
-    skipped = 0
-    errors = 0
-    pbar = tqdm(src_iter, total=total, desc=f"Layout {args.model_name}")
-    for item in pbar:
-        if item.key in completed:
-            skipped += 1
-            continue
-        if item.extras.get("failed") or item.image is None:
-            # Unreadable source image — write an empty layout in position so the
-            # output stays row-aligned with the source dataset.
-            sink.write(item.key, [], {})
-            errors += 1
-            processed += 1
-            continue
-        try:
-            arr = pil_to_array(item.image)
-            results = model.predict(
-                arr,
-                batch_size=args.batch_size,
-                layout_nms=args.layout_nms,
-            )
-            if not results:
-                detections: List[Dict[str, Any]] = []
-            else:
-                detections = extract_detections(results[0])
-                if args.threshold and args.threshold > 0:
-                    detections = [d for d in detections if d["score"] >= args.threshold]
-        except Exception as e:
-            logger.error(f"Error on {item.key}: {e}")
-            detections = []
-            errors += 1
-
-        extras = dict(item.extras)
-        if isinstance(sink, BucketShardSink) and args.include_images:
-            buf = io.BytesIO()
-            item.image.save(buf, format="PNG")
-            extras["__image_bytes"] = buf.getvalue()
-
-        sink.write(item.key, detections, extras)
-        processed += 1
-
-    duration = datetime.now() - start_time
-    processing_time_str = f"{duration.total_seconds() / 60:.2f} min"
-    logger.info(
-        f"Processed {processed} (skipped {skipped}, errors {errors}) in {processing_time_str}"
-    )
-
-    args_dict = {
-        "model_name": args.model_name,
-        "threshold": args.threshold,
-        "layout_nms": args.layout_nms,
-        "shard_size": args.shard_size,
-        "processing_time": processing_time_str,
-    }
-    sink.finalize(model_id=f"PaddlePaddle/{args.model_name}", args_dict=args_dict)
-
-    if args.verbose:
-        import importlib.metadata
-
-        logger.info("--- Resolved package versions ---")
-        for pkg in [
-            "paddlepaddle",
-            "paddlepaddle-gpu",
-            "paddleocr",
-            "huggingface-hub",
-            "datasets",
-            "pyarrow",
-            "pillow",
-            "numpy",
-        ]:
-            try:
-                logger.info(f"  {pkg}=={importlib.metadata.version(pkg)}")
-            except importlib.metadata.PackageNotFoundError:
-                logger.info(f"  {pkg}: not installed")
-        logger.info("--- End versions ---")
-
-
-# ---------------------------------------------------------------------------
-# CLI
-# ---------------------------------------------------------------------------
-
-
-def _print_usage_banner() -> None:
-    print("=" * 80)
-    print("PP-DocLayout layout detection")
-    print("=" * 80)
-    print(
-        "\nDetect document layout regions (text/title/table/figure/formula/...)"
-    )
-    print("with PaddleOCR's PP-DocLayout-L (or M / S / plus-L variant).")
-    print("\nModels:")
-    for m in VALID_MODELS:
-        print(f"  {m:24s} {MODEL_SIZES.get(m, '')}")
-    print("\nSources:")
-    print("  - HF dataset repo:        namespace/dataset")
-    print("  - HF bucket of images:    hf://buckets/namespace/bucket[/prefix]")
-    print("\nSinks:")
-    print("  - HF dataset repo (one push + dataset card):")
-    print("      namespace/dataset")
-    print("  - HF bucket (incremental shards, resumable):")
-    print("      hf://buckets/namespace/bucket/run-name")
-    print("\nExamples:")
-    print("\n  # Smoke test on L4 (dataset -> dataset)")
-    print("  hf jobs uv run --flavor l4x1 -s HF_TOKEN https://huggingface.co/datasets/uv-scripts/ocr/raw/main/pp-doclayout.py \\")
-    print("      davanstrien/ufo-ColPali pp-doclayout-smoke \\")
-    print("      --max-samples 3 --shuffle --seed 42 --private")
-    print("\n  # Dataset -> bucket (incremental shards)")
-    print(
-        "  hf jobs uv run --flavor l4x1 -s HF_TOKEN https://huggingface.co/datasets/uv-scripts/ocr/raw/main/pp-doclayout.py \\"
-    )
-    print("      davanstrien/ufo-ColPali \\")
-    print(
-        "      hf://buckets/davanstrien/pp-doclayout-scratch/run1 \\"
-    )
-    print("      --max-samples 20 --shard-size 5")
-    print("\n  # Bucket of images -> dataset")
-    print(
-        "  hf jobs uv run --flavor l4x1 -s HF_TOKEN https://huggingface.co/datasets/uv-scripts/ocr/raw/main/pp-doclayout.py \\"
-    )
-    print(
-        "      hf://buckets/davanstrien/pp-doclayout-images \\"
-    )
-    print("      pp-doclayout-from-bucket --private")
-    print("\nFor full help, run: uv run pp-doclayout.py --help")
-    print("=" * 80)
-
-
-def build_parser() -> argparse.ArgumentParser:
-    p = argparse.ArgumentParser(
-        description="PP-DocLayout layout detection over an HF dataset or bucket.",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-    )
-    p.add_argument(
-        "input_source",
-        help="HF dataset id (namespace/dataset) OR hf://buckets/ns/bucket[/prefix]",
-    )
-    p.add_argument(
-        "output_target",
-        help="HF dataset id (namespace/dataset) OR hf://buckets/ns/bucket/run-name",
-    )
-    p.add_argument(
-        "--model-name",
-        default="PP-DocLayout-L",
-        choices=VALID_MODELS,
-        help="PaddleOCR layout model variant (default: PP-DocLayout-L)",
-    )
-    p.add_argument(
-        "--device",
-        default="gpu",
-        choices=["gpu", "cpu"],
-        help="Device for inference (default: gpu, falls back to cpu if CUDA missing)",
-    )
-    p.add_argument(
-        "--batch-size",
-        type=int,
-        default=1,
-        help="Per-image batch size passed to model.predict (default: 1)",
-    )
-    p.add_argument(
-        "--threshold",
-        type=float,
-        default=0.5,
-        help="Drop detections below this confidence (default: 0.5; 0 disables)",
-    )
-    p.add_argument(
-        "--layout-nms",
-        dest="layout_nms",
-        action="store_true",
-        default=True,
-        help="Enable layout NMS (default: on)",
-    )
-    p.add_argument(
-        "--no-layout-nms",
-        dest="layout_nms",
-        action="store_false",
-        help="Disable layout NMS",
-    )
-    # Dataset-source-specific
-    p.add_argument(
-        "--image-column",
-        default="image",
-        help="Column containing images (dataset-repo source only, default: image)",
-    )
-    p.add_argument(
-        "--split",
-        default="train",
-        help="Dataset split (dataset-repo source only, default: train)",
-    )
-    p.add_argument(
-        "--max-samples", type=int, help="Limit number of samples (for testing)"
-    )
-    p.add_argument(
-        "--shuffle", action="store_true", help="Shuffle source before processing"
-    )
-    p.add_argument(
-        "--seed", type=int, default=42, help="Random seed for shuffle (default: 42)"
-    )
-    # Dataset-sink-specific
-    p.add_argument(
-        "--private", action="store_true", help="Private dataset output (dataset sink only)"
-    )
-    p.add_argument(
-        "--config",
-        help="Config/subset name when pushing to Hub (dataset sink only)",
-    )
-    p.add_argument(
-        "--create-pr",
-        action="store_true",
-        help="Create PR instead of direct push (dataset sink only)",
-    )
-    p.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
-    # Bucket-sink-specific
-    p.add_argument(
-        "--shard-size",
-        type=int,
-        default=256,
-        help="Rows per parquet shard for bucket sink (default: 256)",
-    )
-    p.add_argument(
-        "--include-images",
-        action="store_true",
-        help="Embed source image bytes in bucket output shards (off by default)",
-    )
-    p.add_argument(
-        "--no-resume",
-        action="store_true",
-        help="Disable resume scan when writing to a bucket sink",
-    )
-    # Auth + diagnostics
-    p.add_argument("--hf-token", help="Hugging Face API token (else uses HF_TOKEN env)")
-    p.add_argument(
-        "--verbose",
-        action="store_true",
-        help="Log resolved package versions at the end",
-    )
-    return p
-
-
-if __name__ == "__main__":
-    if len(sys.argv) == 1:
-        _print_usage_banner()
-        sys.exit(0)
-    main(build_parser().parse_args())

pp-ocrv6.py

@@ -1,1102 +0,0 @@
-# /// script
-# requires-python = ">=3.10"
-# dependencies = [
-#     "paddlepaddle-gpu>=3.0.0",
-#     "paddleocr>=3.7.0",
-#     "paddlex[ocr]>=3.7.0",
-#     "opencv-contrib-python-headless",
-#     "datasets>=3.1.0",
-#     "huggingface-hub",
-#     "pillow",
-#     "numpy",
-#     "tqdm",
-# ]
-#
-# [tool.uv]
-# # PaddleOCR/PaddleX pull in opencv-contrib-python (full) which needs system
-# # libGL.so.1 — not present in the slim uv-on-bookworm image used by HF Jobs.
-# # Swap to the headless cv2 variant (same `import cv2`, no GUI deps). A matching
-# # importlib.metadata patch in main() makes paddlex recognise the headless name.
-# override-dependencies = [
-#     "opencv-contrib-python ; python_version < '0'",
-#     "opencv-python ; python_version < '0'",
-# ]
-#
-# [[tool.uv.index]]
-# name = "paddle"
-# url = "https://www.paddlepaddle.org.cn/packages/stable/cu126/"
-# explicit = true
-#
-# [tool.uv.sources]
-# paddlepaddle-gpu = { index = "paddle" }
-# ///
-"""
-OCR images with PP-OCRv6 — a lightweight detection+recognition pipeline from
-PaddlePaddle. Three tiers from **1.5M to 34.5M parameters**.
-
-Unlike the VLM-based OCR recipes here, PP-OCRv6 is a **classical det+rec pipeline**
-that outputs **plain text** (not markdown). At 1.5M-34.5M params it's far smaller
-than the VLM OCRs and runs on a cheap t4-small GPU.
-
-Model tiers (pick with `--model-tier`):
-  tiny    1.5M params  (0.4M det + 1.1M rec)  49 languages, ~73% recognition
-  small   7.7M params  (2.5M det + 5.3M rec)  50 languages, ~81% recognition
-  medium  34.5M params (22M det + 19M rec)     50 languages, ~83% recognition
-
-All tiers are Apache 2.0 licensed. Runs via PaddleOCR's default Paddle engine
-(`paddle_static`) — same proven header pattern as `pp-doclayout.py`.
-
-HF Jobs examples:
-
-    # Tiny on a cheap GPU
-    hf jobs uv run --flavor t4-small -s HF_TOKEN \\
-        https://huggingface.co/datasets/uv-scripts/ocr/raw/main/pp-ocrv6.py \\
-        INPUT_DATASET OUTPUT_DATASET \\
-        --model-tier tiny --max-samples 5
-
-    # Medium on a small GPU (recommended for quality)
-    hf jobs uv run --flavor t4-small -s HF_TOKEN \\
-        https://huggingface.co/datasets/uv-scripts/ocr/raw/main/pp-ocrv6.py \\
-        INPUT_DATASET OUTPUT_DATASET \\
-        --model-tier medium --max-samples 10
-
-Models: PaddlePaddle/PP-OCRv6_<tier>_det + PP-OCRv6_<tier>_rec
-Blog: https://huggingface.co/blog/PaddlePaddle/pp-ocrv6
-"""
-
-import argparse
-import io
-import json
-import logging
-import os
-import sys
-import time
-from dataclasses import dataclass
-from datetime import datetime, timezone
-from pathlib import Path
-from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
-
-import numpy as np
-from PIL import Image, UnidentifiedImageError
-from tqdm.auto import tqdm
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-# ---------------------------------------------------------------------------
-# Constants
-# ---------------------------------------------------------------------------
-
-TIER_MODELS = {
-    "tiny":   ("PP-OCRv6_tiny_det",   "PP-OCRv6_tiny_rec"),
-    "small":  ("PP-OCRv6_small_det",  "PP-OCRv6_small_rec"),
-    "medium": ("PP-OCRv6_medium_det", "PP-OCRv6_medium_rec"),
-}
-
-TIER_PARAMS = {
-    "tiny":   "1.5M (0.4M det + 1.1M rec)",
-    "small":  "7.7M (2.5M det + 5.3M rec)",
-    "medium": "34.5M (22M det + 19M rec)",
-}
-
-TIER_LANGUAGES = {
-    "tiny":   "49 languages (zh, zh-Hant, en + 46 Latin-script — no Japanese)",
-    "small":  "50 languages (zh, zh-Hant, en, ja + 46 Latin-script)",
-    "medium": "50 languages (zh, zh-Hant, en, ja + 46 Latin-script)",
-}
-
-TIER_REC = {
-    "tiny":   73.5,
-    "small":  81.3,
-    "medium": 83.2,
-}
-
-BUCKET_PREFIX = "hf://buckets/"
-
-IMAGE_EXTENSIONS = {
-    ".jpg", ".jpeg", ".png", ".tif", ".tiff", ".webp", ".bmp", ".jp2", ".j2k",
-}
-
-
-# ---------------------------------------------------------------------------
-# URL helpers
-# ---------------------------------------------------------------------------
-
-def is_bucket_url(s: str) -> bool:
-    return s.startswith(BUCKET_PREFIX)
-
-
-def parse_bucket_url(url: str) -> Tuple[str, str]:
-    if not is_bucket_url(url):
-        raise ValueError(f"Not a bucket URL: {url}")
-    rest = url[len(BUCKET_PREFIX):].strip("/")
-    parts = rest.split("/", 2)
-    if len(parts) < 2:
-        raise ValueError(f"Bucket URL must include namespace and bucket name: {url}")
-    bucket_id = f"{parts[0]}/{parts[1]}"
-    prefix = parts[2] if len(parts) > 2 else ""
-    return bucket_id, prefix
-
-
-# ---------------------------------------------------------------------------
-# Image helpers
-# ---------------------------------------------------------------------------
-
-def to_pil(image: Union[Image.Image, Dict[str, Any], str, bytes]) -> Image.Image:
-    if isinstance(image, Image.Image):
-        return image.convert("RGB")
-    if isinstance(image, dict) and "bytes" in image:
-        return Image.open(io.BytesIO(image["bytes"])).convert("RGB")
-    if isinstance(image, (bytes, bytearray)):
-        return Image.open(io.BytesIO(image)).convert("RGB")
-    if isinstance(image, str):
-        return Image.open(image).convert("RGB")
-    raise ValueError(f"Unsupported image type: {type(image)}")
-
-
-def pil_to_array(pil_img: Image.Image) -> np.ndarray:
-    return np.asarray(pil_img, dtype=np.uint8)
-
-
-# ---------------------------------------------------------------------------
-# Result extraction
-# ---------------------------------------------------------------------------
-
-def extract_text(result: Any) -> Tuple[str, List[Dict[str, Any]]]:
-    """Pull text and per-line details from a PaddleOCR predict result.
-
-    Returns (concatenated_text, per_line_details) where per_line_details is
-    a list of dicts with keys: text, score, bbox (4-point detection polygon as
-    [[x1,y1],[x2,y2],[x3,y3],[x4,y4]] in input-image pixel coordinates).
-    """
-    payload = result.json if hasattr(result, "json") else result
-    res = payload.get("res", payload) if isinstance(payload, dict) else {}
-    rec_texts = res.get("rec_texts", []) or []
-    rec_scores = res.get("rec_scores", []) or []
-    dt_polys = res.get("dt_polys", []) or []
-
-    # Concatenate reading-order text lines (PaddleOCR returns them in order)
-    text = "\n".join(rec_texts)
-
-    per_line = []
-    for i, t in enumerate(rec_texts):
-        entry = {"text": t}
-        if i < len(rec_scores):
-            entry["score"] = float(rec_scores[i])
-        if i < len(dt_polys):
-            entry["bbox"] = [[float(c) for c in point] for point in dt_polys[i]]
-        per_line.append(entry)
-
-    return text, per_line
-
-
-# ---------------------------------------------------------------------------
-# Sources
-# ---------------------------------------------------------------------------
-
-@dataclass
-class SourceItem:
-    key: str
-    image: Optional[Image.Image]
-    extras: Dict[str, Any]
-
-
-def iter_dataset_images(
-    dataset_id: str,
-    image_column: str,
-    split: str,
-    shuffle: bool,
-    seed: int,
-    max_samples: Optional[int],
-):
-    from datasets import load_dataset
-
-    logger.info(f"Loading dataset: {dataset_id} (split={split})")
-    ds = load_dataset(dataset_id, split=split)
-
-    if image_column not in ds.column_names:
-        raise ValueError(
-            f"Column '{image_column}' not found. Available: {ds.column_names}"
-        )
-
-    if shuffle:
-        logger.info(f"Shuffling with seed {seed}")
-        ds = ds.shuffle(seed=seed)
-    if max_samples:
-        ds = ds.select(range(min(max_samples, len(ds))))
-        logger.info(f"Limited to {len(ds)} samples")
-
-    total = len(ds)
-
-    def gen() -> Iterator[SourceItem]:
-        failed = 0
-        for i in range(total):
-            try:
-                row = ds[i]
-                image = to_pil(row[image_column])
-            except (UnidentifiedImageError, OSError) as e:
-                # Still yield a placeholder so the output row stays aligned with
-                # the source row (the dataset sink writes results positionally).
-                failed += 1
-                logger.warning(
-                    f"Unreadable image at row {i}: {type(e).__name__}: {e} "
-                    f"— writing empty result"
-                )
-                yield SourceItem(key=f"row-{i:08d}", image=None, extras={"failed": True})
-                continue
-            yield SourceItem(key=f"row-{i:08d}", image=image, extras={})
-        if failed:
-            logger.info(f"{failed} unreadable image(s) written as empty results")
-
-    return gen(), total, ds
-
-
-SOURCE_PATHS_SNAPSHOT = "_source_paths.json"
-
-
-def _bucket_snapshot_path(output_url: str) -> Tuple[str, str]:
-    out_bucket_id, out_prefix = parse_bucket_url(output_url)
-    snapshot_key = (
-        f"{out_prefix}/{SOURCE_PATHS_SNAPSHOT}".lstrip("/")
-        if out_prefix
-        else SOURCE_PATHS_SNAPSHOT
-    )
-    return out_bucket_id, snapshot_key
-
-
-def iter_bucket_images(
-    bucket_url: str,
-    shuffle: bool,
-    seed: int,
-    max_samples: Optional[int],
-    hf_token: Optional[str],
-    output_url: Optional[str] = None,
-) -> Tuple[Iterator[SourceItem], int]:
-    from huggingface_hub import HfApi, HfFileSystem
-
-    bucket_id, prefix = parse_bucket_url(bucket_url)
-    fs = HfFileSystem(token=hf_token)
-    base = f"{BUCKET_PREFIX}{bucket_id}/{prefix}".rstrip("/")
-
-    snapshot_bucket_id: Optional[str] = None
-    snapshot_key: Optional[str] = None
-    cached_paths: Optional[List[str]] = None
-
-    if output_url and is_bucket_url(output_url):
-        snapshot_bucket_id, snapshot_key = _bucket_snapshot_path(output_url)
-        snapshot_url = f"{BUCKET_PREFIX}{snapshot_bucket_id}/{snapshot_key}"
-        try:
-            with fs.open(snapshot_url, "rb") as f:
-                snapshot = json.load(f)
-            mismatches = []
-            if snapshot.get("source_url") != bucket_url:
-                mismatches.append(
-                    f"source_url ({snapshot.get('source_url')!r} vs {bucket_url!r})"
-                )
-            if snapshot.get("shuffle") != shuffle:
-                mismatches.append(f"shuffle ({snapshot.get('shuffle')} vs {shuffle})")
-            if shuffle and snapshot.get("seed") != seed:
-                mismatches.append(f"seed ({snapshot.get('seed')} vs {seed})")
-            if snapshot.get("max_samples") != max_samples:
-                mismatches.append(
-                    f"max_samples ({snapshot.get('max_samples')} vs {max_samples})"
-                )
-            if mismatches:
-                logger.warning(
-                    "Existing snapshot params differ from this run ("
-                    + "; ".join(mismatches)
-                    + "); ignoring snapshot and re-listing."
-                )
-            else:
-                cached_paths = snapshot["paths"]
-                logger.info(
-                    f"Reusing existing snapshot of {len(cached_paths)} source paths "
-                    f"(written {snapshot.get('created_at', 'unknown')})"
-                )
-        except FileNotFoundError:
-            pass
-        except Exception as e:
-            logger.warning(f"Could not read existing snapshot ({e}); re-listing.")
-
-    if cached_paths is not None:
-        all_paths = cached_paths
-    else:
-        logger.info(f"Listing images under {base}")
-        all_paths = []
-        try:
-            for entry in fs.find(base, detail=False):
-                ext = Path(entry).suffix.lower()
-                if ext in IMAGE_EXTENSIONS:
-                    all_paths.append(entry)
-        except FileNotFoundError as e:
-            raise ValueError(f"Bucket prefix not found: {base}") from e
-
-        if not all_paths:
-            raise ValueError(
-                f"No image files (any of {sorted(IMAGE_EXTENSIONS)}) under {base}"
-            )
-
-        all_paths.sort()
-        if shuffle:
-            rng = np.random.default_rng(seed)
-            rng.shuffle(all_paths)
-        if max_samples:
-            all_paths = all_paths[:max_samples]
-
-        if snapshot_bucket_id is not None and snapshot_key is not None:
-            api = HfApi(token=hf_token)
-            payload = {
-                "source_url": bucket_url,
-                "shuffle": shuffle,
-                "seed": seed,
-                "max_samples": max_samples,
-                "created_at": datetime.now(timezone.utc).isoformat(),
-                "paths": all_paths,
-            }
-            api.batch_bucket_files(
-                snapshot_bucket_id,
-                add=[(json.dumps(payload).encode(), snapshot_key)],
-                token=hf_token,
-            )
-            logger.info(
-                f"Wrote source-path snapshot ({len(all_paths)} paths) to "
-                f"hf://buckets/{snapshot_bucket_id}/{snapshot_key}"
-            )
-
-    total = len(all_paths)
-    logger.info(f"Found {total} images in bucket")
-
-    def key_for(path: str) -> str:
-        return path
-
-    def gen() -> Iterator[SourceItem]:
-        skipped = 0
-        for path in all_paths:
-            try:
-                with fs.open(path, "rb") as f:
-                    data = f.read()
-                image = to_pil(data)
-            except (UnidentifiedImageError, OSError) as e:
-                skipped += 1
-                logger.warning(
-                    f"Skipping unreadable image {path}: {type(e).__name__}: {e}"
-                )
-                continue
-            yield SourceItem(key=key_for(path), image=image, extras={})
-        if skipped:
-            logger.info(f"Skipped {skipped} unreadable image(s) total")
-
-    return gen(), total
-
-
-# ---------------------------------------------------------------------------
-# Sinks
-# ---------------------------------------------------------------------------
-
-class DatasetRepoSink:
-    def __init__(
-        self,
-        repo_id: str,
-        *,
-        hf_token: Optional[str],
-        private: bool,
-        config: Optional[str],
-        create_pr: bool,
-        source_id: str,
-        original_dataset=None,
-        output_column: str = "markdown",
-        overwrite: bool = False,
-    ):
-        self.repo_id = repo_id
-        self.hf_token = hf_token
-        self.private = private
-        self.config = config
-        self.create_pr = create_pr
-        self.source_id = source_id
-        self.original_dataset = original_dataset
-        self.output_column = output_column
-        self.overwrite = overwrite
-        self._texts: List[str] = []
-        self._blocks: List[str] = []
-
-    @property
-    def kind(self) -> str:
-        return "dataset"
-
-    def already_done(self) -> set:
-        return set()
-
-    def write(self, key: str, text: str, blocks: List[Dict[str, Any]]) -> None:
-        self._texts.append(text)
-        self._blocks.append(json.dumps(blocks, ensure_ascii=False))
-
-    def finalize(self, tier: str, det_model: str, rec_model: str, args_dict: Dict[str, Any]) -> None:
-        from datasets import Dataset
-
-        if self.original_dataset is not None:
-            if len(self._texts) != len(self.original_dataset):
-                logger.warning(
-                    f"Text count ({len(self._texts)}) != dataset rows "
-                    f"({len(self.original_dataset)}); padding with empty strings."
-                )
-                while len(self._texts) < len(self.original_dataset):
-                    self._texts.append("")
-                    self._blocks.append("[]")
-            # Guard again at save time in case the input column set changed under us.
-            base = self.original_dataset
-            clash = [c for c in (self.output_column, "pp_ocr_blocks") if c in base.column_names]
-            if clash:
-                if not self.overwrite:
-                    raise ValueError(
-                        f"Output column(s) {clash} already exist in the input dataset; "
-                        f"pass a different --output-column, or --overwrite to replace them."
-                    )
-                logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-                base = base.remove_columns(clash)
-            ds = base.add_column(self.output_column, self._texts)
-            ds = ds.add_column("pp_ocr_blocks", self._blocks)
-        else:
-            if not self._texts:
-                logger.warning("No rows produced; nothing to push.")
-                return
-            ds = Dataset.from_list([
-                {"source_path": None, self.output_column: t, "pp_ocr_blocks": b}
-                for t, b in zip(self._texts, self._blocks)
-            ])
-
-        inference_entry = build_inference_entry(tier, det_model, rec_model, args_dict)
-
-        if "inference_info" in ds.column_names:
-            logger.info("Updating existing inference_info column")
-
-            def _update(example):
-                try:
-                    existing = (
-                        json.loads(example["inference_info"])
-                        if example["inference_info"]
-                        else []
-                    )
-                except (json.JSONDecodeError, TypeError):
-                    existing = []
-                existing.append(inference_entry)
-                return {"inference_info": json.dumps(existing)}
-
-            ds = ds.map(_update)
-        else:
-            ds = ds.add_column(
-                "inference_info", [json.dumps([inference_entry])] * len(ds)
-            )
-
-        logger.info(f"Pushing {len(ds)} rows to {self.repo_id}")
-        push_kwargs = {
-            "private": self.private,
-            "token": self.hf_token,
-            "max_shard_size": "500MB",
-            "create_pr": self.create_pr,
-            "commit_message": f"Add PP-OCRv6-{tier} OCR results ({len(ds)} samples)"
-            + (f" [{self.config}]" if self.config else ""),
-        }
-        if self.config:
-            push_kwargs["config_name"] = self.config
-
-        max_retries = 3
-        for attempt in range(1, max_retries + 1):
-            try:
-                if attempt > 1:
-                    logger.warning("Disabling XET (fallback to HTTP upload)")
-                    os.environ["HF_HUB_DISABLE_XET"] = "1"
-                ds.push_to_hub(self.repo_id, **push_kwargs)
-                break
-            except Exception as e:
-                logger.error(f"Upload attempt {attempt}/{max_retries} failed: {e}")
-                if attempt == max_retries:
-                    logger.error("All upload attempts failed.")
-                    raise
-                time.sleep(30 * (2 ** (attempt - 1)))
-
-        from huggingface_hub import DatasetCard
-
-        card = DatasetCard(
-            create_dataset_card(
-                source=self.source_id,
-                tier=tier,
-                det_model=det_model,
-                rec_model=rec_model,
-                num_samples=len(ds),
-                processing_time=args_dict["processing_time"],
-                engine=args_dict.get("engine", "paddle_static"),
-                output_id=self.repo_id,
-                output_column=self.output_column,
-            )
-        )
-        card.push_to_hub(self.repo_id, token=self.hf_token)
-        logger.info(f"Done: https://huggingface.co/datasets/{self.repo_id}")
-
-
-class BucketShardSink:
-    METADATA_FILE = "_metadata.json"
-    SHARD_PATTERN = "shard-{:05d}.parquet"
-
-    def __init__(
-        self,
-        bucket_url: str,
-        *,
-        hf_token: Optional[str],
-        shard_size: int,
-        resume: bool,
-        source_id: str,
-    ):
-        from huggingface_hub import HfApi, HfFileSystem, create_bucket
-
-        self.bucket_url = bucket_url
-        self.bucket_id, self.prefix = parse_bucket_url(bucket_url)
-        self.hf_token = hf_token
-        self.shard_size = shard_size
-        self.resume = resume
-        self.source_id = source_id
-
-        self._api = HfApi(token=hf_token)
-        self._fs = HfFileSystem(token=hf_token)
-
-        try:
-            create_bucket(self.bucket_id, exist_ok=True, token=hf_token)
-        except Exception as e:
-            logger.warning(f"create_bucket('{self.bucket_id}') warning: {e}")
-
-        self._buffer: List[Dict[str, Any]] = []
-        self._next_shard_idx = self._discover_next_shard_idx()
-        self._completed_keys = self._discover_completed_keys() if resume else set()
-        if self._completed_keys:
-            logger.info(
-                f"Resume: found {len(self._completed_keys)} already-processed keys, will skip them"
-            )
-
-    @property
-    def kind(self) -> str:
-        return "bucket"
-
-    def already_done(self) -> set:
-        return self._completed_keys
-
-    def _shard_path(self, idx: int) -> str:
-        return self._join(self.SHARD_PATTERN.format(idx))
-
-    def _join(self, name: str) -> str:
-        return f"{self.prefix}/{name}".lstrip("/") if self.prefix else name
-
-    def _list_existing_shards(self) -> List[str]:
-        try:
-            tree = self._api.list_bucket_tree(
-                self.bucket_id, prefix=self.prefix or None, recursive=True
-            )
-        except Exception:
-            return []
-        shards: List[str] = []
-        for item in tree:
-            path = getattr(item, "path", None)
-            ftype = getattr(item, "type", None)
-            if not path or ftype not in (None, "file"):
-                continue
-            base = Path(path).name
-            if base.startswith("shard-") and base.endswith(".parquet"):
-                shards.append(path)
-        return sorted(shards)
-
-    def _discover_next_shard_idx(self) -> int:
-        shards = self._list_existing_shards()
-        max_idx = -1
-        for s in shards:
-            stem = Path(s).stem
-            try:
-                max_idx = max(max_idx, int(stem.split("-")[-1]))
-            except ValueError:
-                continue
-        return max_idx + 1
-
-    def _discover_completed_keys(self) -> set:
-        import pyarrow.parquet as pq
-
-        keys: set = set()
-        for shard_path in self._list_existing_shards():
-            full = f"{BUCKET_PREFIX}{self.bucket_id}/{shard_path}"
-            try:
-                with self._fs.open(full, "rb") as f:
-                    table = pq.read_table(f, columns=["__source_key"])
-                keys.update(table.column("__source_key").to_pylist())
-            except Exception as e:
-                logger.warning(f"Could not read keys from {shard_path}: {e}")
-        return keys
-
-    def _flush(self) -> None:
-        if not self._buffer:
-            return
-        import pyarrow as pa
-        import pyarrow.parquet as pq
-
-        columns = ["__source_key", "text", "pp_ocr_blocks"]
-        table_dict = {c: [row.get(c) for row in self._buffer] for c in columns}
-        table = pa.Table.from_pydict(table_dict)
-
-        buf = io.BytesIO()
-        pq.write_table(table, buf, compression="zstd")
-        data = buf.getvalue()
-
-        shard_remote = self._shard_path(self._next_shard_idx)
-        logger.info(
-            f"Writing shard {self._next_shard_idx} ({len(self._buffer)} rows, "
-            f"{len(data) / 1024 / 1024:.1f} MiB) to {shard_remote}"
-        )
-        self._api.batch_bucket_files(
-            self.bucket_id, add=[(data, shard_remote)], token=self.hf_token
-        )
-        self._next_shard_idx += 1
-        self._buffer.clear()
-
-    def write(self, key: str, text: str, blocks: List[Dict[str, Any]]) -> None:
-        row: Dict[str, Any] = {
-            "__source_key": key,
-            "text": text,
-            "pp_ocr_blocks": json.dumps(blocks, ensure_ascii=False),
-        }
-        self._buffer.append(row)
-        if len(self._buffer) >= self.shard_size:
-            self._flush()
-
-    def finalize(self, tier: str, det_model: str, rec_model: str, args_dict: Dict[str, Any]) -> None:
-        self._flush()
-        meta = {
-            "model": f"PP-OCRv6_{tier}",
-            "det_model": det_model,
-            "rec_model": rec_model,
-            "tier": tier,
-            "engine": "paddle_static",
-            "source": self.source_id,
-            "shard_size": args_dict["shard_size"],
-            "last_run_at": datetime.now(timezone.utc).isoformat(),
-            "processing_time": args_dict.get("processing_time"),
-        }
-        meta_bytes = json.dumps(meta, indent=2).encode("utf-8")
-        meta_path = self._join(self.METADATA_FILE)
-        self._api.batch_bucket_files(
-            self.bucket_id, add=[(meta_bytes, meta_path)], token=self.hf_token
-        )
-        logger.info(
-            f"Done: https://huggingface.co/buckets/{self.bucket_id}"
-            + (f"/{self.prefix}" if self.prefix else "")
-        )
-
-
-# ---------------------------------------------------------------------------
-# inference_info + dataset card
-# ---------------------------------------------------------------------------
-
-def build_inference_entry(tier: str, det_model: str, rec_model: str, args_dict: Dict[str, Any]) -> Dict[str, Any]:
-    return {
-        "model_id": f"PaddlePaddle/PP-OCRv6_{tier}",
-        "det_model": det_model,
-        "rec_model": rec_model,
-        "tier": tier,
-        "params": TIER_PARAMS.get(tier, "unknown"),
-        "rec_accuracy_pct": TIER_REC.get(tier),
-        "languages": TIER_LANGUAGES.get(tier, ""),
-        "engine": "paddle_static",
-        # column_name is the key ocr-bench's column discovery reads; keep
-        # output_column too for backward compat with existing outputs.
-        "column_name": args_dict.get("output_column", "markdown"),
-        "output_column": args_dict.get("output_column", "markdown"),
-        "blocks_column": "pp_ocr_blocks",
-        "timestamp": datetime.now(timezone.utc).isoformat(),
-    }
-
-
-def create_dataset_card(
-    source: str,
-    tier: str,
-    det_model: str,
-    rec_model: str,
-    num_samples: int,
-    processing_time: str,
-    engine: str,
-    output_id: str,
-    output_column: str = "markdown",
-) -> str:
-    tier_display = tier.upper() if tier == "tiny" else tier.capitalize()
-    if is_bucket_url(source):
-        source_link = f"[{source}]({source})"
-    else:
-        source_link = f"[{source}](https://huggingface.co/datasets/{source})"
-
-    return f"""---
-tags:
-- ocr
-- text-recognition
-- paddleocr
-- pp-ocrv6
-- uv-script
-- generated
----
-
-# OCR with PP-OCRv6 {tier_display}
-
-Plain-text OCR results for images from {source_link}, produced by
-PaddlePaddle's [PP-OCRv6](https://huggingface.co/collections/PaddlePaddle/pp-ocrv6)
-{tier} pipeline ({TIER_PARAMS.get(tier, "unknown")}).
-
-## Processing details
-
-- **Source**: {source_link}
-- **Model**: PP-OCRv6_{tier} ({det_model} + {rec_model})
-- **Tier**: {tier} ({TIER_PARAMS.get(tier, "unknown")})
-- **Recognition accuracy**: {TIER_REC.get(tier, "?"):.1f}%
-- **Languages**: {TIER_LANGUAGES.get(tier, "")}
-- **Engine**: {engine}
-- **Samples**: {num_samples:,}
-- **Processing time**: {processing_time}
-- **Processing date**: {datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC")}
-- **License**: Apache 2.0 (models)
-
-## Schema
-
-Each row contains the original columns plus:
-
-- `{output_column}`: Plain text extracted from the image (reading-order concatenation of
-  detected text lines, newline-separated).
-- `pp_ocr_blocks`: JSON list, one dict per detected text line:
-  ```json
-  [
-    {{
-      "text": "recognized text",
-      "score": 0.987,
-      "bbox": [[x1, y1], [x2, y2], [x3, y3], [x4, y4]]
-    }}
-  ]
-  ```
-  `score` is the recognition confidence and `bbox` is the detection polygon
-  (4-point quadrilateral in input-image pixel coordinates).
-- `inference_info`: JSON list tracking every model applied to this dataset.
-
-> **Note:** PP-OCRv6 is a classical detection+recognition pipeline, not a VLM.
-> It outputs **plain text** rather than markdown. Per-line bounding boxes and
-> confidence scores are available in `pp_ocr_blocks`.
-
-## Usage
-
-```python
-import json
-from datasets import load_dataset
-
-ds = load_dataset("{output_id}", split="train")
-print(ds[0]["{output_column}"])
-for block in json.loads(ds[0]["pp_ocr_blocks"]):
-    print(block["text"], block["score"])
-```
-
-## Reproduction
-
-```bash
-hf jobs uv run --flavor t4-small -s HF_TOKEN \\
-    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/pp-ocrv6.py \\
-    {source} <output> --model-tier {tier}
-```
-
-Generated with [UV Scripts](https://huggingface.co/uv-scripts).
-"""
-
-
-# ---------------------------------------------------------------------------
-# Main
-# ---------------------------------------------------------------------------
-
-def main(args: argparse.Namespace) -> None:
-    from huggingface_hub import login
-
-    start_time = datetime.now()
-    hf_token = args.hf_token or os.environ.get("HF_TOKEN")
-    if hf_token:
-        login(token=hf_token)
-
-    # ---------- tier → model names ----------
-    if args.model_tier not in TIER_MODELS:
-        raise ValueError(
-            f"Invalid tier {args.model_tier!r}. Choose from: {list(TIER_MODELS)}"
-        )
-    det_model, rec_model = TIER_MODELS[args.model_tier]
-    tier = args.model_tier
-    logger.info(f"PP-OCRv6 {tier}: {det_model} + {rec_model}")
-
-    # ---------- source ----------
-    original_dataset = None
-    if is_bucket_url(args.input_source):
-        src_iter, total = iter_bucket_images(
-            args.input_source,
-            shuffle=args.shuffle,
-            seed=args.seed,
-            max_samples=args.max_samples,
-            hf_token=hf_token,
-            output_url=args.output_target,
-        )
-    else:
-        src_iter, total, original_dataset = iter_dataset_images(
-            args.input_source,
-            image_column=args.image_column,
-            split=args.split,
-            shuffle=args.shuffle,
-            seed=args.seed,
-            max_samples=args.max_samples,
-        )
-        # Fail fast, before minutes of inference, if the output column would collide
-        # with an existing input column (e.g. a 'text' ground-truth column). Writing
-        # into it would either crash on push or silently overwrite the input data.
-        # --overwrite opts in to replacing the existing column(s) instead of erroring.
-        if original_dataset is not None:
-            clash = [
-                col
-                for col in (args.output_column, "pp_ocr_blocks")
-                if col in original_dataset.column_names
-            ]
-            if clash and not args.overwrite:
-                logger.error(
-                    f"Output column(s) {clash} already exist in the input dataset "
-                    f"(columns: {original_dataset.column_names})."
-                )
-                logger.error(
-                    "Choose a different --output-column, or pass --overwrite to replace them."
-                )
-                sys.exit(1)
-            if clash:
-                logger.warning(f"--overwrite: will replace existing column(s) {clash}")
-
-    # ---------- sink ----------
-    if is_bucket_url(args.output_target):
-        sink: Union[BucketShardSink, DatasetRepoSink] = BucketShardSink(
-            args.output_target,
-            hf_token=hf_token,
-            shard_size=args.shard_size,
-            resume=not args.no_resume,
-            source_id=args.input_source,
-        )
-    else:
-        sink = DatasetRepoSink(
-            args.output_target,
-            hf_token=hf_token,
-            private=args.private,
-            config=args.config,
-            create_pr=args.create_pr,
-            source_id=args.input_source,
-            original_dataset=original_dataset,
-            output_column=args.output_column,
-            overwrite=args.overwrite,
-        )
-
-    completed = sink.already_done()
-
-    # ---------- model ----------
-    # PaddleX gates `import cv2` at module load time on
-    # `is_dep_available("opencv-contrib-python")`, which checks
-    # `importlib.metadata.version(...)`. We ship `opencv-contrib-python-headless`
-    # (same `cv2`, no system libGL.so.1 needed) — but that's a different
-    # distribution name, so the gate fails and the OCR pipeline's `ocr` extra
-    # check returns False. Patch the metadata lookup to alias the GUI cv2 distros
-    # to the headless variant before importing paddleocr; this lets paddlex's own
-    # `import cv2` succeed and `is_extra_available('ocr')` return True.
-    import importlib.metadata as _metadata
-
-    _orig_metadata_version = _metadata.version
-
-    def _patched_metadata_version(dep_name):
-        if dep_name in ("opencv-contrib-python", "opencv-python"):
-            for headless_alias in (
-                "opencv-contrib-python-headless",
-                "opencv-python-headless",
-            ):
-                try:
-                    return _orig_metadata_version(headless_alias)
-                except _metadata.PackageNotFoundError:
-                    continue
-        return _orig_metadata_version(dep_name)
-
-    _metadata.version = _patched_metadata_version
-
-    # Silence the connectivity check for speed (not needed in a Job)
-    os.environ.setdefault("PADDLE_PDX_DISABLE_MODEL_SOURCE_CHECK", "True")
-
-    from paddleocr import PaddleOCR
-
-    ocr = PaddleOCR(
-        text_detection_model_name=det_model,
-        text_recognition_model_name=rec_model,
-        use_doc_orientation_classify=False,
-        use_doc_unwarping=False,
-        use_textline_orientation=False,
-    )
-
-    # ---------- loop ----------
-    processed = 0
-    skipped = 0
-    errors = 0
-    pbar = tqdm(src_iter, total=total, desc=f"PP-OCRv6 {tier}")
-    for item in pbar:
-        if item.key in completed:
-            skipped += 1
-            continue
-        if item.extras.get("failed") or item.image is None:
-            # Unreadable source image — write an empty result in position so the
-            # output stays row-aligned with the source dataset.
-            sink.write(item.key, "", [])
-            errors += 1
-            processed += 1
-            continue
-        try:
-            arr = pil_to_array(item.image)
-            result = ocr.predict(arr)
-            if result:
-                text, blocks = extract_text(result[0])
-            else:
-                text, blocks = "", []
-        except Exception as e:
-            logger.error(f"Error on {item.key}: {e}")
-            text, blocks = "", []
-            errors += 1
-
-        sink.write(item.key, text, blocks)
-        processed += 1
-
-    duration = datetime.now() - start_time
-    processing_time_str = f"{duration.total_seconds() / 60:.2f} min"
-    logger.info(
-        f"Processed {processed} (skipped {skipped}, errors {errors}) in {processing_time_str}"
-    )
-
-    args_dict = {
-        "tier": tier,
-        "det_model": det_model,
-        "rec_model": rec_model,
-        "engine": "paddle_static",
-        "shard_size": args.shard_size,
-        "processing_time": processing_time_str,
-        "output_column": args.output_column,
-    }
-    sink.finalize(
-        tier=tier,
-        det_model=det_model,
-        rec_model=rec_model,
-        args_dict=args_dict,
-    )
-
-    if args.verbose:
-        import importlib.metadata
-
-        logger.info("--- Resolved package versions ---")
-        for pkg in [
-            "paddleocr",
-            "paddlex",
-            "paddlepaddle-gpu",
-            "huggingface-hub",
-            "datasets",
-            "pillow",
-            "numpy",
-        ]:
-            try:
-                logger.info(f"  {pkg}=={importlib.metadata.version(pkg)}")
-            except importlib.metadata.PackageNotFoundError:
-                logger.info(f"  {pkg}: not installed")
-        logger.info("--- End versions ---")
-
-
-# ---------------------------------------------------------------------------
-# CLI
-# ---------------------------------------------------------------------------
-
-def build_parser() -> argparse.ArgumentParser:
-    p = argparse.ArgumentParser(
-        description="PP-OCRv6 OCR over an HF dataset or bucket of images.",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-    )
-    p.add_argument(
-        "input_source",
-        help="HF dataset id (namespace/dataset) OR hf://buckets/ns/bucket[/prefix]",
-    )
-    p.add_argument(
-        "output_target",
-        help="HF dataset id (namespace/dataset) OR hf://buckets/ns/bucket/run-name",
-    )
-    p.add_argument(
-        "--model-tier",
-        default="medium",
-        choices=list(TIER_MODELS),
-        help="PP-OCRv6 model tier: tiny (1.5M), small (7.7M), medium (34.5M). Default: medium.",
-    )
-    # Dataset-source-specific
-    p.add_argument(
-        "--image-column",
-        default="image",
-        help="Column containing images (dataset-repo source only, default: image)",
-    )
-    p.add_argument(
-        "--split",
-        default="train",
-        help="Dataset split (dataset-repo source only, default: train)",
-    )
-    p.add_argument(
-        "--max-samples", type=int, help="Limit number of samples (for testing)"
-    )
-    p.add_argument(
-        "--shuffle", action="store_true", help="Shuffle source before processing"
-    )
-    p.add_argument(
-        "--seed", type=int, default=42, help="Random seed for shuffle (default: 42)"
-    )
-    # Dataset-sink-specific
-    p.add_argument(
-        "--private", action="store_true", help="Private dataset output (dataset sink only)"
-    )
-    p.add_argument(
-        "--config",
-        help="Config/subset name when pushing to Hub (dataset sink only)",
-    )
-    p.add_argument(
-        "--create-pr",
-        action="store_true",
-        help="Create PR instead of direct push (dataset sink only)",
-    )
-    p.add_argument(
-        "--output-column",
-        default="markdown",
-        help=(
-            "Column name for the recognized text (dataset sink only, default: markdown). "
-            "Must not collide with an existing input column — many corpora already ship a "
-            "'text' ground-truth column, so 'text' would fail on push. Blocks always go to "
-            "'pp_ocr_blocks'."
-        ),
-    )
-    p.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column(s) if they already exist in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
-    # Bucket-sink-specific
-    p.add_argument(
-        "--shard-size",
-        type=int,
-        default=256,
-        help="Rows per parquet shard for bucket sink (default: 256)",
-    )
-    p.add_argument(
-        "--no-resume",
-        action="store_true",
-        help="Disable resume scan when writing to a bucket sink",
-    )
-    # Auth + diagnostics
-    p.add_argument("--hf-token", help="Hugging Face API token (else uses HF_TOKEN env)")
-    p.add_argument(
-        "--verbose",
-        action="store_true",
-        help="Log resolved package versions at the end",
-    )
-    return p
-
-
-if __name__ == "__main__":
-    main(build_parser().parse_args())

qianfan-ocr.py

@@ -1,665 +0,0 @@
-# /// script
-# requires-python = ">=3.11"
-# dependencies = [
-#     "datasets>=4.0.0",
-#     "huggingface-hub",
-#     "pillow",
-#     "vllm>=0.15.1",
-#     "tqdm",
-#     "toolz",
-#     "torch",
-# ]
-# ///
-
-"""
-Convert document images to markdown using Qianfan-OCR with vLLM.
-
-Qianfan-OCR is a 4.7B end-to-end document intelligence model from Baidu,
-built on InternVL architecture with Qianfan-ViT encoder + Qwen3-4B LLM.
-
-Features:
-- #1 end-to-end model on OmniDocBench v1.5 (93.12) and OlmOCR Bench (79.8)
-- Layout-as-Thought: optional reasoning phase for complex layouts via --think
-- 192 language support (Latin, CJK, Arabic, Cyrillic, and more)
-- Multiple task modes: OCR, table (HTML), formula (LaTeX), chart, scene text
-- Key information extraction with custom prompts
-- 1.024 PPS on A100 with W8A8 quantization
-
-Model: baidu/Qianfan-OCR
-License: Apache 2.0
-Paper: https://arxiv.org/abs/2603.13398
-"""
-
-import argparse
-import base64
-import io
-import json
-import logging
-import os
-import sys
-import time
-from datetime import datetime
-from typing import Any, Dict, List, Union
-
-import torch
-from datasets import load_dataset
-from huggingface_hub import DatasetCard, login
-from PIL import Image
-from toolz import partition_all
-from tqdm.auto import tqdm
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
-from vllm import LLM, SamplingParams
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-MODEL = "baidu/Qianfan-OCR"
-
-PROMPT_TEMPLATES = {
-    "ocr": "Parse this document to Markdown.",
-    "table": "Extract tables to HTML format.",
-    "formula": "Extract formulas to LaTeX.",
-    "chart": "What trends are shown in this chart?",
-    "scene": "Extract all visible text from the image.",
-    "kie": None,  # requires --custom-prompt
-}
-
-
-def check_cuda_availability():
-    """Check if CUDA is available and exit if not."""
-    if not torch.cuda.is_available():
-        logger.error("CUDA is not available. This script requires a GPU.")
-        sys.exit(1)
-    else:
-        logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
-
-
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
-def extract_content_from_thinking(text: str, include_thinking: bool = False) -> str:
-    """
-    Extract final content from Qianfan-OCR's Layout-as-Thought output.
-
-    When --think is enabled, the model generates layout analysis inside
-    <think>...</think> tags before the final markdown output.
-    """
-    if include_thinking:
-        return text.strip()
-
-    # If no thinking tags, return as-is
-    if "<think>" not in text:
-        return text.strip()
-
-    # Extract everything after </think>
-    think_end = text.find("</think>")
-    if think_end != -1:
-        return text[think_end + 8 :].strip()
-
-    # Thinking started but never closed — return full text
-    logger.warning("Found <think> but no </think>, returning full text")
-    return text.strip()
-
-
-def make_ocr_message(
-    image: Union[Image.Image, Dict[str, Any], str],
-    prompt: str,
-) -> List[Dict]:
-    """Create vLLM chat message with image and prompt."""
-    if isinstance(image, Image.Image):
-        pil_img = image
-    elif isinstance(image, dict) and "bytes" in image:
-        pil_img = Image.open(io.BytesIO(image["bytes"]))
-    elif isinstance(image, str):
-        pil_img = Image.open(image)
-    else:
-        raise ValueError(f"Unsupported image type: {type(image)}")
-
-    pil_img = pil_img.convert("RGB")
-
-    buf = io.BytesIO()
-    pil_img.save(buf, format="PNG")
-    data_uri = f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}"
-
-    return [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image_url", "image_url": {"url": data_uri}},
-                {"type": "text", "text": prompt},
-            ],
-        }
-    ]
-
-
-def create_dataset_card(
-    source_dataset: str,
-    model: str,
-    num_samples: int,
-    processing_time: str,
-    batch_size: int,
-    max_model_len: int,
-    max_tokens: int,
-    gpu_memory_utilization: float,
-    prompt_mode: str,
-    think: bool,
-    include_thinking: bool,
-    image_column: str = "image",
-    split: str = "train",
-) -> str:
-    """Create a dataset card documenting the OCR process."""
-    model_name = model.split("/")[-1]
-
-    return f"""---
-tags:
-- ocr
-- document-processing
-- qianfan-ocr
-- markdown
-- uv-script
-- generated
----
-
-# Document OCR using {model_name}
-
-This dataset contains OCR results from [{source_dataset}](https://huggingface.co/datasets/{source_dataset}) using Qianfan-OCR, Baidu's 4.7B end-to-end document intelligence model.
-
-## Processing Details
-
-- **Source Dataset**: [{source_dataset}](https://huggingface.co/datasets/{source_dataset})
-- **Model**: [{model}](https://huggingface.co/{model})
-- **Number of Samples**: {num_samples:,}
-- **Processing Time**: {processing_time}
-- **Processing Date**: {datetime.now().strftime("%Y-%m-%d %H:%M UTC")}
-
-### Configuration
-
-- **Image Column**: `{image_column}`
-- **Output Column**: `markdown`
-- **Dataset Split**: `{split}`
-- **Batch Size**: {batch_size}
-- **Prompt Mode**: {prompt_mode}
-- **Layout-as-Thought**: {"Enabled" if think else "Disabled"}
-- **Thinking Traces**: {"Included" if include_thinking else "Excluded"}
-- **Max Model Length**: {max_model_len:,} tokens
-- **Max Output Tokens**: {max_tokens:,}
-- **GPU Memory Utilization**: {gpu_memory_utilization:.1%}
-
-## Model Information
-
-Qianfan-OCR key capabilities:
-- #1 end-to-end model on OmniDocBench v1.5 (93.12)
-- #1 on OlmOCR Bench (79.8)
-- 192 language support
-- Layout-as-Thought reasoning for complex documents
-- Document parsing, table extraction, formula recognition, chart understanding
-- Key information extraction
-
-## Dataset Structure
-
-The dataset contains all original columns plus:
-- `markdown`: The extracted text in markdown format
-- `inference_info`: JSON list tracking all OCR models applied
-
-## Reproduction
-
-```bash
-uv run https://huggingface.co/datasets/uv-scripts/ocr/raw/main/qianfan-ocr.py \\
-    {source_dataset} \\
-    <output-dataset> \\
-    --image-column {image_column} \\
-    --prompt-mode {prompt_mode} \\
-    --batch-size {batch_size}{"  --think" if think else ""}
-```
-
-Generated with [UV Scripts](https://huggingface.co/uv-scripts)
-"""
-
-
-def main(
-    input_dataset: str,
-    output_dataset: str,
-    image_column: str = "image",
-    batch_size: int = 8,
-    max_model_len: int = 16384,
-    max_tokens: int = 8192,
-    temperature: float = 0.0,
-    top_p: float = 1.0,
-    gpu_memory_utilization: float = 0.85,
-    hf_token: str = None,
-    split: str = "train",
-    max_samples: int = None,
-    private: bool = False,
-    shuffle: bool = False,
-    seed: int = 42,
-    prompt_mode: str = "ocr",
-    think: bool = False,
-    include_thinking: bool = False,
-    custom_prompt: str = None,
-    output_column: str = "markdown",
-    overwrite: bool = False,
-    config: str = None,
-    create_pr: bool = False,
-    verbose: bool = False,
-):
-    """Process images from HF dataset through Qianfan-OCR model."""
-
-    check_cuda_availability()
-    start_time = datetime.now()
-
-    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")
-    if HF_TOKEN:
-        login(token=HF_TOKEN)
-
-    # Build prompt
-    if custom_prompt:
-        prompt = custom_prompt
-        logger.info(f"Using custom prompt: {prompt[:80]}...")
-    else:
-        if prompt_mode == "kie":
-            logger.error("--prompt-mode kie requires --custom-prompt")
-            sys.exit(1)
-        prompt = PROMPT_TEMPLATES[prompt_mode]
-        logger.info(f"Using prompt mode: {prompt_mode}")
-
-    if think:
-        prompt = prompt + "<think>"
-        logger.info("Layout-as-Thought enabled (appending <think> to prompt)")
-
-    logger.info(f"Using model: {MODEL}")
-
-    # Load dataset
-    logger.info(f"Loading dataset: {input_dataset}")
-    dataset = load_dataset(input_dataset, split=split)
-
-    if image_column not in dataset.column_names:
-        raise ValueError(
-            f"Column '{image_column}' not found. Available: {dataset.column_names}"
-        )
-
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
-    if shuffle:
-        logger.info(f"Shuffling dataset with seed {seed}")
-        dataset = dataset.shuffle(seed=seed)
-
-    if max_samples:
-        dataset = dataset.select(range(min(max_samples, len(dataset))))
-        logger.info(f"Limited to {len(dataset)} samples")
-
-    # Initialize vLLM
-    logger.info("Initializing vLLM with Qianfan-OCR")
-    logger.info("This may take a few minutes on first run...")
-    llm = LLM(
-        model=MODEL,
-        trust_remote_code=True,
-        max_model_len=max_model_len,
-        gpu_memory_utilization=gpu_memory_utilization,
-        limit_mm_per_prompt={"image": 1},
-        enforce_eager=False,
-    )
-
-    sampling_params = SamplingParams(
-        temperature=temperature,
-        top_p=top_p,
-        max_tokens=max_tokens,
-    )
-
-    logger.info(f"Processing {len(dataset)} images in batches of {batch_size}")
-    logger.info(f"Output will be written to column: {output_column}")
-
-    # Process images in batches
-    all_outputs = []
-
-    for batch_indices in tqdm(
-        partition_all(batch_size, range(len(dataset))),
-        total=(len(dataset) + batch_size - 1) // batch_size,
-        desc="Qianfan-OCR processing",
-    ):
-        batch_indices = list(batch_indices)
-        batch_images = [dataset[i][image_column] for i in batch_indices]
-
-        try:
-            batch_messages = [make_ocr_message(img, prompt) for img in batch_images]
-            outputs = llm.chat(batch_messages, sampling_params)
-
-            for output in outputs:
-                text = output.outputs[0].text.strip()
-                if think:
-                    text = extract_content_from_thinking(text, include_thinking)
-                all_outputs.append(text)
-
-        except Exception as e:
-            logger.error(f"Error processing batch: {e}")
-            all_outputs.extend(["[OCR ERROR]"] * len(batch_images))
-
-    # Calculate processing time
-    processing_duration = datetime.now() - start_time
-    processing_time_str = f"{processing_duration.total_seconds() / 60:.1f} min"
-
-    # Add output column
-    logger.info(f"Adding '{output_column}' column to dataset")
-    dataset = dataset.add_column(output_column, all_outputs)
-
-    # Handle inference_info tracking
-    inference_entry = {
-        "model_id": MODEL,
-        "model_name": "Qianfan-OCR",
-        "column_name": output_column,
-        "timestamp": datetime.now().isoformat(),
-        "prompt_mode": prompt_mode if not custom_prompt else "custom",
-        "think": think,
-        "temperature": temperature,
-        "max_tokens": max_tokens,
-    }
-
-    if "inference_info" in dataset.column_names:
-        logger.info("Updating existing inference_info column")
-
-        def update_inference_info(example):
-            try:
-                existing_info = (
-                    json.loads(example["inference_info"])
-                    if example["inference_info"]
-                    else []
-                )
-            except (json.JSONDecodeError, TypeError):
-                existing_info = []
-            existing_info.append(inference_entry)
-            return {"inference_info": json.dumps(existing_info)}
-
-        dataset = dataset.map(update_inference_info)
-    else:
-        logger.info("Creating new inference_info column")
-        inference_list = [json.dumps([inference_entry])] * len(dataset)
-        dataset = dataset.add_column("inference_info", inference_list)
-
-    # Push to hub with retry and XET fallback
-    logger.info(f"Pushing to {output_dataset}")
-    commit_msg = f"Add Qianfan-OCR results ({len(dataset)} samples)" + (
-        f" [{config}]" if config else ""
-    )
-    max_retries = 3
-    for attempt in range(1, max_retries + 1):
-        try:
-            if attempt > 1:
-                logger.warning("Disabling XET (fallback to HTTP upload)")
-                os.environ["HF_HUB_DISABLE_XET"] = "1"
-            dataset.push_to_hub(
-                output_dataset,
-                private=private,
-                token=HF_TOKEN,
-                max_shard_size="500MB",
-                **({"config_name": config} if config else {}),
-                create_pr=create_pr,
-                commit_message=commit_msg,
-            )
-            break
-        except Exception as e:
-            logger.error(f"Upload attempt {attempt}/{max_retries} failed: {e}")
-            if attempt < max_retries:
-                delay = 30 * (2 ** (attempt - 1))
-                logger.info(f"Retrying in {delay}s...")
-                time.sleep(delay)
-            else:
-                logger.error("All upload attempts failed. OCR results are lost.")
-                sys.exit(1)
-
-    # Create and push dataset card (skip when creating PR to avoid conflicts)
-    if not create_pr:
-        logger.info("Creating dataset card")
-        card_content = create_dataset_card(
-            source_dataset=input_dataset,
-            model=MODEL,
-            num_samples=len(dataset),
-            processing_time=processing_time_str,
-            batch_size=batch_size,
-            max_model_len=max_model_len,
-            max_tokens=max_tokens,
-            gpu_memory_utilization=gpu_memory_utilization,
-            prompt_mode=prompt_mode if not custom_prompt else "custom",
-            think=think,
-            include_thinking=include_thinking,
-            image_column=image_column,
-            split=split,
-        )
-        card = DatasetCard(card_content)
-        card.push_to_hub(output_dataset, token=HF_TOKEN)
-
-    logger.info("Qianfan-OCR processing complete!")
-    logger.info(
-        f"Dataset available at: https://huggingface.co/datasets/{output_dataset}"
-    )
-    logger.info(f"Processing time: {processing_time_str}")
-    logger.info(
-        f"Processing speed: {len(dataset) / processing_duration.total_seconds():.2f} images/sec"
-    )
-
-    if verbose:
-        import importlib.metadata
-
-        logger.info("--- Resolved package versions ---")
-        for pkg in ["vllm", "transformers", "torch", "datasets", "pyarrow", "pillow"]:
-            try:
-                logger.info(f"  {pkg}=={importlib.metadata.version(pkg)}")
-            except importlib.metadata.PackageNotFoundError:
-                logger.info(f"  {pkg}: not installed")
-        logger.info("--- End versions ---")
-
-
-if __name__ == "__main__":
-    if len(sys.argv) == 1:
-        print("=" * 80)
-        print("Qianfan-OCR - End-to-End Document Intelligence")
-        print("=" * 80)
-        print("\n4.7B model from Baidu, #1 on OmniDocBench v1.5 (93.12)")
-        print("\nFeatures:")
-        print("- #1 end-to-end model on OmniDocBench v1.5 and OlmOCR Bench")
-        print("- Layout-as-Thought reasoning for complex documents (--think)")
-        print("- 192 language support")
-        print("- Multiple modes: OCR, table (HTML), formula (LaTeX), chart, scene text")
-        print("- Key information extraction with custom prompts")
-        print("\nExample usage:")
-        print("\n1. Basic OCR:")
-        print("   uv run qianfan-ocr.py input-dataset output-dataset")
-        print("\n2. With Layout-as-Thought (complex documents):")
-        print("   uv run qianfan-ocr.py docs output --think")
-        print("\n3. Table extraction:")
-        print("   uv run qianfan-ocr.py docs output --prompt-mode table")
-        print("\n4. Formula extraction:")
-        print("   uv run qianfan-ocr.py docs output --prompt-mode formula")
-        print("\n5. Key information extraction:")
-        print(
-            '   uv run qianfan-ocr.py invoices output --prompt-mode kie --custom-prompt "Extract: name, date, total. Output JSON."'
-        )
-        print("\n6. Running on HF Jobs:")
-        print("   hf jobs uv run --flavor l4x1 \\")
-        print("     -s HF_TOKEN \\")
-        print(
-            "     https://huggingface.co/datasets/uv-scripts/ocr/raw/main/qianfan-ocr.py \\"
-        )
-        print("       input-dataset output-dataset --max-samples 10")
-        print("\nFor full help, run: uv run qianfan-ocr.py --help")
-        sys.exit(0)
-
-    parser = argparse.ArgumentParser(
-        description="Document OCR using Qianfan-OCR (4.7B, #1 on OmniDocBench v1.5)",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog="""
-Prompt modes:
-  ocr       Document parsing to Markdown (default)
-  table     Table extraction to HTML format
-  formula   Formula recognition to LaTeX
-  chart     Chart understanding and analysis
-  scene     Scene text extraction
-  kie       Key information extraction (requires --custom-prompt)
-
-Examples:
-  uv run qianfan-ocr.py my-docs analyzed-docs
-  uv run qianfan-ocr.py docs output --think --max-samples 50
-  uv run qianfan-ocr.py docs output --prompt-mode table
-  uv run qianfan-ocr.py invoices data --prompt-mode kie --custom-prompt "Extract: name, date, total."
-        """,
-    )
-
-    parser.add_argument("input_dataset", help="Input dataset ID from Hugging Face Hub")
-    parser.add_argument("output_dataset", help="Output dataset ID for Hugging Face Hub")
-    parser.add_argument(
-        "--image-column",
-        default="image",
-        help="Column containing images (default: image)",
-    )
-    parser.add_argument(
-        "--batch-size",
-        type=int,
-        default=8,
-        help="Batch size for processing (default: 8)",
-    )
-    parser.add_argument(
-        "--max-model-len",
-        type=int,
-        default=16384,
-        help="Maximum model context length (default: 16384, reduce to 8192 if OOM on L4)",
-    )
-    parser.add_argument(
-        "--max-tokens",
-        type=int,
-        default=8192,
-        help="Maximum tokens to generate (default: 8192)",
-    )
-    parser.add_argument(
-        "--temperature",
-        type=float,
-        default=0.0,
-        help="Sampling temperature (default: 0.0, deterministic)",
-    )
-    parser.add_argument(
-        "--top-p",
-        type=float,
-        default=1.0,
-        help="Top-p sampling parameter (default: 1.0)",
-    )
-    parser.add_argument(
-        "--gpu-memory-utilization",
-        type=float,
-        default=0.85,
-        help="GPU memory utilization (default: 0.85)",
-    )
-    parser.add_argument("--hf-token", help="Hugging Face API token")
-    parser.add_argument(
-        "--split", default="train", help="Dataset split to use (default: train)"
-    )
-    parser.add_argument(
-        "--max-samples",
-        type=int,
-        help="Maximum number of samples to process (for testing)",
-    )
-    parser.add_argument(
-        "--private", action="store_true", help="Make output dataset private"
-    )
-    parser.add_argument(
-        "--shuffle", action="store_true", help="Shuffle dataset before processing"
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=42,
-        help="Random seed for shuffling (default: 42)",
-    )
-    parser.add_argument(
-        "--prompt-mode",
-        choices=list(PROMPT_TEMPLATES.keys()),
-        default="ocr",
-        help="Prompt mode (default: ocr)",
-    )
-    parser.add_argument(
-        "--think",
-        action="store_true",
-        help="Enable Layout-as-Thought reasoning (appends <think> to prompt)",
-    )
-    parser.add_argument(
-        "--include-thinking",
-        action="store_true",
-        help="Include thinking traces in output (default: only final content)",
-    )
-    parser.add_argument(
-        "--custom-prompt",
-        help="Custom prompt text (overrides --prompt-mode)",
-    )
-    parser.add_argument(
-        "--output-column",
-        default="markdown",
-        help="Column name for output text (default: markdown)",
-    )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
-    parser.add_argument(
-        "--config",
-        help="Config/subset name when pushing to Hub (for benchmarking multiple models)",
-    )
-    parser.add_argument(
-        "--create-pr",
-        action="store_true",
-        help="Create a pull request instead of pushing directly",
-    )
-    parser.add_argument(
-        "--verbose",
-        action="store_true",
-        help="Log resolved package versions after processing",
-    )
-
-    args = parser.parse_args()
-
-    main(
-        input_dataset=args.input_dataset,
-        output_dataset=args.output_dataset,
-        image_column=args.image_column,
-        batch_size=args.batch_size,
-        max_model_len=args.max_model_len,
-        max_tokens=args.max_tokens,
-        temperature=args.temperature,
-        top_p=args.top_p,
-        gpu_memory_utilization=args.gpu_memory_utilization,
-        hf_token=args.hf_token,
-        split=args.split,
-        max_samples=args.max_samples,
-        private=args.private,
-        shuffle=args.shuffle,
-        seed=args.seed,
-        prompt_mode=args.prompt_mode,
-        think=args.think,
-        include_thinking=args.include_thinking,
-        custom_prompt=args.custom_prompt,
-        output_column=args.output_column,
-        overwrite=args.overwrite,
-        config=args.config,
-        create_pr=args.create_pr,
-        verbose=args.verbose,
-    )

rolm-ocr.py

@@ -40,10 +40,6 @@ from huggingface_hub import DatasetCard, login
 from PIL import Image
 from toolz import partition_all
 from tqdm.auto import tqdm
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
 from vllm import LLM, SamplingParams
 from datetime import datetime
 
@@ -61,28 +57,6 @@ def check_cuda_availability():
         logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
 
 
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
 def make_ocr_message(
     image: Union[Image.Image, Dict[str, Any], str],
     prompt: str = "Return the plain text representation of this document as if you were reading it naturally.\n",
@@ -131,15 +105,6 @@ def create_dataset_card(
     """Create a dataset card documenting the OCR process."""
     model_name = model.split("/")[-1]
 
-    # Canonical provenance stamp (see AGENTS.md): Jobs claim gated on JOB_ID, set by HF Jobs in-container.
-    on_jobs = os.environ.get("JOB_ID") is not None
-    hw = os.environ.get("ACCELERATOR") or ""
-    origin = (
-        "Produced on [Hugging Face Jobs](https://huggingface.co/docs/huggingface_hub/guides/jobs)"
-        + (f" (`{hw}`)" if hw else "")
-    ) if on_jobs else "Generated"
-    jobs_tag = "\n- hf-jobs" if on_jobs else ""
-
     return f"""---
 viewer: false
 tags:
@@ -147,7 +112,7 @@ tags:
 - text-extraction
 - rolmocr
 - uv-script
-- generated{jobs_tag}
+- generated
 ---
 
 # OCR Text Extraction using {model_name}
@@ -204,10 +169,10 @@ for info in inference_info:
 
 ## Reproduction
 
-{origin} with the [`rolm-ocr.py`](https://huggingface.co/datasets/uv-scripts/ocr/raw/main/rolm-ocr.py) recipe from [uv-scripts](https://huggingface.co/uv-scripts). Run it yourself:
+This dataset was generated using the [uv-scripts/ocr](https://huggingface.co/datasets/uv-scripts/ocr) RolmOCR script:
 
 ```bash
-hf jobs uv run https://huggingface.co/datasets/uv-scripts/ocr/raw/main/rolm-ocr.py \\
+uv run https://huggingface.co/datasets/uv-scripts/ocr/raw/main/rolm-ocr.py \\
     {source_dataset} \\
     <output-dataset> \\
     --image-column {image_column} \\
@@ -221,6 +186,8 @@ hf jobs uv run https://huggingface.co/datasets/uv-scripts/ocr/raw/main/rolm-ocr.
 
 - **Processing Speed**: ~{num_samples / (float(processing_time.split()[0]) * 60):.1f} images/second
 - **GPU Configuration**: vLLM with {gpu_memory_utilization:.0%} GPU memory utilization
+
+Generated with 🤖 [UV Scripts](https://huggingface.co/uv-scripts)
 """
 
 
@@ -238,7 +205,6 @@ def main(
     max_samples: int = None,
     private: bool = False,
     output_column: str = None,
-    overwrite: bool = False,
     shuffle: bool = False,
     seed: int = 42,
     verbose: bool = False,
@@ -273,9 +239,6 @@ def main(
             f"Column '{image_column}' not found. Available: {dataset.column_names}"
         )
 
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
     # Shuffle if requested
     if shuffle:
         logger.info(f"Shuffling dataset with seed {seed}")
@@ -535,12 +498,6 @@ Examples:
         default=None,
         help="Name of the output column for extracted text (default: markdown)",
     )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
     parser.add_argument(
         "--shuffle",
         action="store_true",
@@ -574,7 +531,6 @@ Examples:
         max_samples=args.max_samples,
         private=args.private,
         output_column=args.output_column,
-        overwrite=args.overwrite,
         shuffle=args.shuffle,
         seed=args.seed,
         verbose=args.verbose,

serving-unlimited-ocr.md

@@ -1,169 +0,0 @@
-# Serve Unlimited-OCR as a live endpoint on HF Jobs
-
-The OCR recipes in this folder run as batch jobs (dataset in → dataset out). To call a model
-interactively, from an agent, or with ad-hoc concurrent requests, you can instead run it as a
-temporary HTTP endpoint. [HF Jobs serving](https://huggingface.co/docs/hub/jobs-serving) exposes a
-port on a GPU Job, giving an OpenAI-compatible endpoint that runs until the job is cancelled or its
-`--timeout` is reached.
-
-This is a worked example for [baidu/Unlimited-OCR](https://huggingface.co/baidu/Unlimited-OCR)
-(3B, MIT, based on DeepSeek-OCR; supports multi-page parsing in a single request). Two server
-options below: **vLLM** on Baidu's official image (the newer official path, OpenAI-compatible), or
-**SGLang** on the stock image with the model's own wheel. Either gives an OpenAI-compatible endpoint.
-
-> **Single-image vs multi-page — pick the engine by task:**
-> - **Single-page** OCR (one image → markdown): both engines work. For a whole corpus, the batch
->   recipe [`unlimited-ocr-vllm.py`](unlimited-ocr-vllm.py) (offline vLLM, resumable, no network) is
->   the better fit than a client loop; for interactive/agent use, serve with **vLLM (Option A)**.
-> - **Multi-page / long-horizon** parsing (the model's headline feature): **both engines do it**
->   (validated 2026-06-28 — a clean 2-page doc read back both pages, `<PAGE>`-separated, on *both* vLLM
->   and SGLang). The difference is **robustness on hard inputs**: on degraded historical scans / newspaper
->   clippings, vLLM multi-page degraded to hallucination in our tests while **SGLang (Option B)** read
->   real content — so SGLang is the **more robust** multi-page path (it's also the authors' documented
->   one, via `images_config`). Use vLLM multi-page for clean docs; reach for SGLang for hard scans.
->   (vLLM's upstream PR [#46564](https://github.com/vllm-project/vllm/pull/46564) benchmarks single-page only.)
-
-## 1. Start the server
-
-### Option A — vLLM (official image)
-
-vLLM support landed upstream; Baidu ships a dedicated image (the architecture isn't in a stable pip
-wheel yet). Use the default `:unlimited-ocr` tag on L4/A100, or `:unlimited-ocr-cu129` on Hopper.
-Runs on `l4x1`, no fa3/Hopper requirement. **Single-image is validated**; **multi-page also works on
-clean docs** (both pages, `<PAGE>`-separated) but degraded to hallucination on hard scans in our tests
-— for hard/degraded inputs prefer Option B (SGLang). For multi-page on vLLM, the request takes one
-`<image>` per page in the text and `window_size=1024` in `vllm_xargs` (it has no `images_config`).
-
-```bash
-hf jobs run --detach --expose 8000 --flavor l4x1 -s HF_TOKEN --timeout 30m \
-  vllm/vllm-openai:unlimited-ocr -- \
-  vllm serve baidu/Unlimited-OCR --served-model-name Unlimited-OCR \
-    --trust-remote-code --max-model-len 32768 --host 0.0.0.0 --port 8000 \
-    --logits_processors vllm.model_executor.models.unlimited_ocr:NGramPerReqLogitsProcessor \
-    --no-enable-prefix-caching --mm-processor-cache-gb 0
-```
-
-Per-request, vLLM takes the no-repeat n-gram knobs via `vllm_xargs` and needs `skip_special_tokens`
-off (it has no `images_config` — that's an SGLang param):
-
-```python
-r = client.chat.completions.create(
-    model="Unlimited-OCR",
-    messages=[{"role": "user", "content": [
-        {"type": "text", "text": "<image>document parsing."},  # literal <image> prefix is required
-        {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{img}"}},
-    ]}],
-    temperature=0,
-    extra_body={"skip_special_tokens": False, "vllm_xargs": {"ngram_size": 35, "window_size": 128}},
-)
-```
-
-### Option B — SGLang (model's own build) · supports multi-page
-
-The model also ships its own SGLang build, installed at startup from a 12 MB wheel. **This is the
-more robust path for multi-page / long-horizon parsing** (§3) — the model authors' documented route
-(`images_config`), and the one that held up on hard scans where vLLM multi-page hallucinated. Two
-pins matter (both learned the hard way, 2026-06-28):
-- **Pin the image to `lmsysorg/sglang:v0.5.10.post1`** — *not* `:latest`. `:latest` drifted to torch
-  2.11 / cu130, incompatible with the wheel (torch 2.9.1 / cuda-python 12.9); v0.5.10.post1 is the last
-  release that matches the wheel exactly.
-- **Run on `a100-large` with `--attention-backend flashinfer`, not `h200`/`fa3`.** `fa3` needs a Hopper
-  GPU, but HF's `h200` nodes currently fail GPU init with `CUDA error 802: system not yet initialized`
-  (3/3 attempts) — an infra issue, not the model. `a100` + `flashinfer` sidesteps it and works.
-
-```bash
-hf jobs run --detach --expose 10000 --flavor a100-large -s HF_TOKEN --timeout 30m \
-  lmsysorg/sglang:v0.5.10.post1 -- \
-  bash -lc 'pip install --no-deps https://github.com/baidu/Unlimited-OCR/raw/main/wheel/sglang-0.0.0.dev11416+g92e8bb79e-py3-none-any.whl \
-    && pip install -q kernels==0.11.7 \
-    && python -m sglang.launch_server --model baidu/Unlimited-OCR --served-model-name Unlimited-OCR \
-       --attention-backend flashinfer --page-size 1 --mem-fraction-static 0.85 --context-length 32768 \
-       --enable-custom-logit-processor --disable-overlap-schedule --skip-server-warmup \
-       --host 0.0.0.0 --port 10000'
-```
-
-Notes:
-- `--` before `bash` is required, or the CLI parses `-lc` as its own flags.
-- `--timeout` stops the endpoint (and billing) at the deadline; `hf jobs cancel <id>` stops it earlier.
-- **Validated 2026-06-28** on `a100-large`: server came up, single-image and multi-page both read
-  correctly (a clean 2-page doc returned both pages verbatim, `<PAGE>`-separated). The model card's
-  "official" backend is `fa3` on Hopper for exact R-SWA — switch back to `--attention-backend fa3
-  --flavor h200` once the h200 `802` infra issue clears; `flashinfer` on `a100` is the working fallback.
-- Follow startup with `hf jobs logs -f <id>`; ready at `The server is fired up` / `Application startup
-  complete` (a few minutes cold; the wheel + model download dominate).
-
-The client examples below use the **SGLang** request format (`images_config` in `extra_body`,
-port 10000). The single-image call (§2) also works on the vLLM server — just use the Option A
-`extra_body` and your exposed port. **Multi-page (§3) is SGLang-only.**
-
-## 2. Call it (OpenAI client; HF token as the API key)
-
-The exposed port is at `https://<job_id>--10000.hf.jobs`; the OpenAI base URL is that plus `/v1`.
-
-```python
-import base64, os
-from openai import OpenAI
-
-client = OpenAI(base_url="https://<job_id>--10000.hf.jobs/v1", api_key=os.environ["HF_TOKEN"])
-img = base64.b64encode(open("page.jpg", "rb").read()).decode()
-
-r = client.chat.completions.create(
-    model="Unlimited-OCR",
-    messages=[{"role": "user", "content": [
-        {"type": "text", "text": "document parsing."},
-        {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{img}"}},
-    ]}],
-    temperature=0,
-    extra_body={"images_config": {"image_mode": "gundam"}},  # "gundam" (crop-tiling) or "base"
-)
-print(r.choices[0].message.content)
-```
-
-Output is layout-grounded markdown: each block is tagged `<|det|>type [x1,y1,x2,y2]<|/det|> text`,
-with coordinates normalized to 0–1000. Remove the tags for plain text
-(`re.sub(r'<\|det\|>.*?<\|/det\|>', '', text)`) or keep them for structure.
-
-## 3. Multi-page / PDF (SGLang shown; vLLM also works on clean docs)
-
-> ✅ This **SGLang** flow (Option B) is **validated working 2026-06-28** (a clean 2-page doc read back
-> both pages verbatim, `<PAGE>`-separated) and follows the model card's multi-page example. The
-> `images_config`/`image_mode` param is SGLang-specific — **vLLM ignores it**; on vLLM, do multi-page
-> with one `<image>` per page in the text + `window_size=1024` in `vllm_xargs` (no `images_config`).
-> Both engines read clean multi-page docs; **SGLang was the more robust on hard/degraded scans**, where
-> vLLM multi-page hallucinated in our tests. (vLLM's upstream
-> [PR #46564](https://github.com/vllm-project/vllm/pull/46564) benchmarks single-page only.)
-
-Send multiple page images in one request with the `Multi page parsing.` prompt and `image_mode="base"`:
-
-```python
-parts = [{"type": "text", "text": "Multi page parsing."}]
-for page_png in page_images:            # e.g. PDF pages rendered with pymupdf at ~150 dpi
-    b64 = base64.b64encode(open(page_png, "rb").read()).decode()
-    parts.append({"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}})
-
-r = client.chat.completions.create(
-    model="Unlimited-OCR",
-    messages=[{"role": "user", "content": parts}],
-    temperature=0, max_tokens=16384,
-    extra_body={"images_config": {"image_mode": "base"}},
-)
-```
-
-Pages are separated by `<PAGE>`; tables are returned as HTML and equations as LaTeX, with reading
-order preserved across pages. The context length is 32k tokens, so split longer documents.
-
-## 4. Concurrency
-
-SGLang batches concurrent requests, so a client can send many requests in parallel to one endpoint;
-the upstream [`infer.py`](https://github.com/baidu/Unlimited-OCR/blob/main/infer.py) uses a
-`ThreadPoolExecutor` at `concurrency=8`. For a large corpus, a batch job that runs next to the data
-(resumable, no network transfer) is usually a better fit than a client-to-endpoint loop.
-
-## 5. Stop it
-
-```bash
-hf jobs cancel <job_id>
-```
-
-Billing is per-minute for the GPU flavor plus a small flat fee for the exposed port; scheduling time
-is not billed. Run `hf jobs hardware` for current flavors and prices.

smoldocling-ocr.py

@@ -1,7 +1,7 @@
 # /// script
 # requires-python = ">=3.11"
 # dependencies = [
-#     "datasets>=4.0.0",
+#     "datasets",
 #     "huggingface-hub",
 #     "pillow",
 #     "vllm",
@@ -42,13 +42,9 @@ from typing import Any, Dict, Union
 import torch
 from datasets import load_dataset
 from huggingface_hub import DatasetCard, login
-from PIL import Image, UnidentifiedImageError
+from PIL import Image
 from toolz import partition_all
 from tqdm.auto import tqdm
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the
-# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this
-# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
 from vllm import LLM, SamplingParams
 
 logging.basicConfig(level=logging.INFO)
@@ -65,28 +61,6 @@ def check_cuda_availability():
         logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
 
 
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
 def prepare_llm_input(
     image: Union[Image.Image, Dict[str, Any], str],
     prompt_text: str = "Convert page to Docling.",
@@ -250,7 +224,6 @@ def main(
     max_samples: int = None,
     private: bool = False,
     output_column: str = "markdown",
-    overwrite: bool = False,
     output_format: str = "markdown",
     shuffle: bool = False,
     seed: int = 42,
@@ -282,9 +255,6 @@ def main(
             f"Column '{image_column}' not found. Available: {dataset.column_names}"
         )
 
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
     # Validate output format
     if output_format not in ["markdown", "doctags"]:
         raise ValueError(
@@ -329,21 +299,7 @@ def main(
         desc="OCR processing",
     ):
         batch_indices = list(batch_indices)
-
-        # Fetch images first, with per-batch fallback for unreadable files.
-        # One corrupt image used to take down the entire run via the list
-        # comprehension; now we mark the whole batch as skipped and continue.
-        try:
-            batch_images = [dataset[i][image_column] for i in batch_indices]
-        except (UnidentifiedImageError, OSError) as e:
-            logger.warning(
-                f"Skipping batch of {len(batch_indices)} — unreadable image "
-                f"in batch: {type(e).__name__}: {e}"
-            )
-            all_output.extend(
-                ["[OCR SKIPPED — UNREADABLE IMAGE]"] * len(batch_indices)
-            )
-            continue
+        batch_images = [dataset[i][image_column] for i in batch_indices]
 
         try:
             # Prepare inputs for batch
@@ -597,12 +553,6 @@ Examples:
         default="markdown",
         help="Column name for output text (default: markdown)",
     )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
     parser.add_argument(
         "--output-format",
         default="markdown",
@@ -656,7 +606,6 @@ Examples:
         max_samples=args.max_samples,
         private=args.private,
         output_column=args.output_column,
-        overwrite=args.overwrite,
         output_format=args.output_format,
         shuffle=args.shuffle,
         seed=args.seed,

surya-ocr-bucket.py

@@ -1,1389 +0,0 @@
-# /// script
-# requires-python = ">=3.10"
-# dependencies = [
-#     "surya-ocr==0.20.0",
-#     "datasets>=3.1.0",
-#     "huggingface-hub",
-#     "pillow",
-#     "imagecodecs",
-#     "toolz",
-#     "tqdm",
-# ]
-#
-# # Pin surya-ocr to the known-good build (has the `surya.inference` engine layout
-# # this recipe injects into); an unpinned/loosened resolve backtracks to an ancient
-# # surya without it. huggingface-hub is left unpinned: at runtime PYTHONPATH puts the
-# # pinned image's hub (with the buckets API) ahead of the venv, so no version tension.
-# ///
-"""
-Structured OCR over a **bucket of document files** (images + PDFs) with Datalab's
-**Surya OCR 2** (`datalab-to/surya-ocr-2`, 650M, Qwen3.5-style) — no dataset
-round-trip. This is the bucket-native sibling of `surya-ocr.py` (which reads a Hub
-dataset column). Point it straight at an HF bucket of `.jp2`/`.png`/`.pdf`/... files.
-
-Like the parent it produces *structured* OCR: per-block HTML + bounding boxes +
-reading order + confidence. `--task` switches between `ocr` (full-page text),
-`layout` (labelled regions), and `table` (HTML / rows-cols-cells).
-
-INPUT — two interchangeable I/O strategies (`--io-mode`, default `auto`):
-  mount  bucket mounted read-only at /in via `-v hf://buckets/<id>:/in:ro`; files
-         are read straight off the FUSE mount. Zero ephemeral disk.
-  copy   take a bucket id directly; the huggingface_hub library LISTs then batch-
-         DOWNLOADS each `--batch-size` chunk to local temp, OCRs it, writes output,
-         then deletes the temp batch. Avoids the known FUSE bulk-read stall; peak
-         disk = one batch. `auto` picks copy for an `hf://buckets/...` input, mount
-         for a local dir.
-
-OUTPUT — one or both (>=1 required):
-  --output-bucket  per page a `.md` (flattened reading-order text) AND a `.json`
-                   (that page's structured `surya_blocks`), mirroring the input dir
-                   structure, into a mounted dir OR an `hf://buckets/...` URL.
-                   Streaming / O(1) memory, with resume-by-skip (a file whose
-                   `.json` already exists is skipped) — the scalable path.
-  --output-dataset a parquet dataset pushed to the Hub (one row per file:
-                   file_name / markdown / surya_blocks / inference_info), like the
-                   parent recipe. Convenient; buffered in memory (no image bytes by
-                   default — use `--include-images` to embed page images).
-
-ENGINE: Surya normally spawns a vLLM **server** (Docker), which can't run inside an
-HF Job. This injects a custom in-process backend into Surya's `SuryaInferenceManager`
-that runs vLLM's offline `LLM().chat()` engine (no server). Surya still owns all the
-prompting, image preprocessing, and HTML/bbox parsing — we only swap the transport.
-
-LICENSE NOTE: Surya's *code* is Apache-2.0 but the *weights* are a modified
-OpenRAIL-M license — free for research, personal use, and startups under $5M
-funding/revenue, restricted from competitive use against Datalab's API. Confirm you
-are within those terms. https://huggingface.co/datalab-to/surya-ocr-2
-
-HF Jobs — MUST use the pinned vLLM image + the site-packages python path (the model
-is the recent, version-sensitive `qwen3_5` architecture; v0.20.1 is Surya's
-known-good build, and it puts python/vLLM under /usr/local, NOT /usr/bin):
-
-    # copy input -> dataset output
-    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 \\
-        https://huggingface.co/datasets/uv-scripts/ocr/raw/main/surya-ocr-bucket.py \\
-        hf://buckets/<ns>/<bucket> --io-mode copy --glob "*.jp2" \\
-        --output-dataset <ns>/<out> --private
-
-    # mount input -> per-file bucket output (mirrors dir structure)
-    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 \\
-        -v hf://buckets/<ns>/<bucket>:/in:ro \\
-        -v hf://buckets/<ns>/<out-bucket>:/out \\
-        https://huggingface.co/datasets/uv-scripts/ocr/raw/main/surya-ocr-bucket.py \\
-        /in --io-mode mount --glob "*.jp2" --output-bucket /out
-
-Model: datalab-to/surya-ocr-2  (package: surya-ocr, https://github.com/datalab-to/surya)
-"""
-
-import argparse
-import json
-import logging
-import math
-import os
-import shutil
-import sys
-import tempfile
-import time
-from contextlib import contextmanager
-from dataclasses import dataclass
-from datetime import datetime, timezone
-from fnmatch import fnmatch
-from pathlib import Path, PurePosixPath
-from typing import Any, Dict, Iterator, List, Optional, Tuple
-
-from PIL import Image, UnidentifiedImageError
-from toolz import partition_all
-from tqdm import tqdm
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-DEFAULT_MODEL = "datalab-to/surya-ocr-2"
-# Surya's own vision-tiling bounds (from its vLLM backend), applied to the
-# offline engine too so preprocessing matches the server path exactly.
-MM_PROCESSOR_KWARGS = {"min_pixels": 3136, "max_pixels": 6291456}
-TASKS = ("ocr", "layout", "table")
-# Extensions read by default. `.jp2`/`.j2k` are first-class: the canonical test
-# corpus (Library of Congress / Chronicling America) is all JPEG-2000.
-DEFAULT_EXTENSIONS = ".jp2,.j2k,.png,.jpg,.jpeg,.tiff,.tif,.bmp,.webp,.pdf"
-JP2_EXTENSIONS = {".jp2", ".j2k"}
-PDF_EXTENSION = ".pdf"
-BUCKET_PREFIX = "hf://buckets/"
-
-
-# ---------------------------------------------------------------------------
-# GPU / page-range helpers (verbatim from surya-ocr.py)
-# ---------------------------------------------------------------------------
-
-
-def check_cuda_availability() -> None:
-    """Exit early with a clear message if there's no GPU."""
-    import torch
-
-    if not torch.cuda.is_available():
-        logger.error("CUDA is not available. This script requires a GPU.")
-        logger.error(
-            "Run on Hugging Face Jobs with: hf jobs uv run --flavor l4x1 "
-            "--image vllm/vllm-openai:v0.20.1 ..."
-        )
-        sys.exit(1)
-    logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
-
-
-def parse_page_range(spec: Optional[str]) -> Optional[List[int]]:
-    """Turn '0-3,5' into [0,1,2,3,5]. None/empty -> None (all pages)."""
-    if not spec:
-        return None
-    pages: List[int] = []
-    for part in spec.split(","):
-        part = part.strip()
-        if not part:
-            continue
-        if "-" in part:
-            lo, hi = part.split("-", 1)
-            pages.extend(range(int(lo), int(hi) + 1))
-        else:
-            pages.append(int(part))
-    return pages or None
-
-
-# --- structured-output shim (vLLM API moved between versions) ---
-def build_structured_outputs(schema: Dict[str, Any]) -> Dict[str, Any]:
-    """SamplingParams kwargs for guided JSON, across vLLM versions (layout uses this)."""
-    try:
-        from vllm.sampling_params import StructuredOutputsParams  # vLLM >= 0.12
-
-        return {"structured_outputs": StructuredOutputsParams(json=schema)}
-    except (ImportError, TypeError):
-        pass
-    try:
-        from vllm.sampling_params import GuidedDecodingParams  # older vLLM
-
-        return {"guided_decoding": GuidedDecodingParams(json=schema)}
-    except (ImportError, TypeError):
-        pass
-    logger.warning(
-        "Guided JSON unavailable in this vLLM version; relying on the model."
-    )
-    return {}
-
-
-def _mean_token_prob(completion_output) -> Optional[float]:
-    """Mean exp(logprob) of the sampled tokens -> Surya's per-block `confidence`."""
-    lps = getattr(completion_output, "logprobs", None)
-    if not lps:
-        return None
-    probs: List[float] = []
-    for tid, lp_dict in zip(completion_output.token_ids, lps):
-        if not lp_dict:
-            continue
-        entry = lp_dict.get(tid)
-        if (
-            entry is None
-        ):  # sampled token not in the returned top-k; use the best we have
-            entry = max(lp_dict.values(), key=lambda e: e.logprob)
-        probs.append(math.exp(entry.logprob))
-    return sum(probs) / len(probs) if probs else None
-
-
-# ---------------------------------------------------------------------------
-# Offline vLLM backend + Surya manager (verbatim from surya-ocr.py)
-# ---------------------------------------------------------------------------
-
-
-class OfflineVLLMBackend:
-    """Surya `Backend` (duck-typed) that runs vLLM's offline `LLM().chat()` engine.
-
-    Surya's predictors call `manager.generate(batch)` -> `backend.generate(batch)`;
-    we satisfy that contract in-process (no server). Surya keeps ownership of the
-    prompts (`PROMPT_MAPPING`), image scaling (`scale_to_fit`), and output parsing.
-    """
-
-    name = "offline-vllm"
-
-    def __init__(
-        self,
-        model: str,
-        max_model_len: int,
-        gpu_memory_utilization: float,
-        dtype: str = "bfloat16",
-        max_tokens_default: int = 2048,
-        logprobs_default: bool = True,
-    ):
-        self.model = model
-        self.max_model_len = max_model_len
-        self.gpu_memory_utilization = gpu_memory_utilization
-        self.dtype = dtype
-        self.max_tokens_default = max_tokens_default
-        self.logprobs_default = logprobs_default
-        self.llm = None
-        self._build_messages = None
-        self._scale_to_fit = None
-        self._prompt_mapping = None
-
-    def start(self):
-        from vllm import LLM
-
-        logger.info(
-            f"Loading {self.model} into vLLM offline engine (dtype={self.dtype})..."
-        )
-        self.llm = LLM(
-            model=self.model,
-            dtype=self.dtype,
-            max_model_len=self.max_model_len,
-            gpu_memory_utilization=self.gpu_memory_utilization,
-            mm_processor_kwargs=MM_PROCESSOR_KWARGS,
-            limit_mm_per_prompt={"image": 1},
-        )
-        # Reuse Surya's exact request shaping so the offline path matches the server.
-        from surya.inference.backends.openai_client import _build_messages
-        from surya.inference.prompts import PROMPT_MAPPING
-        from surya.inference.util import scale_to_fit
-
-        self._build_messages = _build_messages
-        self._scale_to_fit = scale_to_fit
-        self._prompt_mapping = PROMPT_MAPPING
-        return None
-
-    def stop(self) -> None:
-        self.llm = None
-
-    def _sampling_params(self, item):
-        from vllm import SamplingParams
-
-        max_tokens = item.max_tokens or self.max_tokens_default
-        want_logprobs = item.request_logprobs or self.logprobs_default
-        kwargs: Dict[str, Any] = dict(temperature=0.0, top_p=0.1, max_tokens=max_tokens)
-        if want_logprobs:
-            kwargs["logprobs"] = 1
-        if item.guided_json is not None:
-            kwargs.update(build_structured_outputs(item.guided_json))
-        return SamplingParams(**kwargs)
-
-    def generate(self, batch):
-        from surya.inference.schema import BatchOutputItem
-
-        if self.llm is None:
-            self.start()
-        if not batch:
-            return []
-
-        conversations = []
-        sampling_params = []
-        for item in batch:
-            prompt = item.prompt or self._prompt_mapping[item.prompt_type]
-            image = self._scale_to_fit(item.image)
-            conversations.append(self._build_messages(image, prompt))
-            sampling_params.append(self._sampling_params(item))
-
-        outputs = self.llm.chat(
-            conversations,
-            sampling_params,
-            chat_template_content_format="openai",
-            use_tqdm=False,
-        )
-
-        results = []
-        for item, out in zip(batch, outputs):
-            comp = out.outputs[0]
-            results.append(
-                BatchOutputItem(
-                    raw=comp.text,
-                    token_count=len(comp.token_ids),
-                    error=False,
-                    mean_token_prob=_mean_token_prob(comp),
-                    logprobs=None,
-                    metadata=item.metadata,  # carries page_idx/block_idx — must round-trip
-                )
-            )
-        return results
-
-
-def make_manager(backend: OfflineVLLMBackend):
-    """A SuryaInferenceManager wired to our offline backend (bypassing autodetect)."""
-    from surya.inference import SuryaInferenceManager
-
-    manager = SuryaInferenceManager.__new__(SuryaInferenceManager)
-    manager.method = backend.name
-    manager.backend = backend
-    return manager
-
-
-# ---------------------------------------------------------------------------
-# Result serialization (verbatim from surya-ocr.py)
-# ---------------------------------------------------------------------------
-
-
-def _html_to_text(html: str) -> str:
-    from bs4 import BeautifulSoup
-
-    return BeautifulSoup(html, "html.parser").get_text(" ", strip=True)
-
-
-def serialize_pages(task: str, pages: List[Any]) -> Tuple[str, List[Dict[str, Any]]]:
-    """(text, structured-per-page) for one document's page results."""
-    structured = [p.model_dump(mode="json") for p in pages]
-    page_texts: List[str] = []
-    for page in pages:
-        if task == "ocr":
-            parts = []
-            for b in sorted(page.blocks, key=lambda b: b.reading_order):
-                if b.skipped or not b.html:
-                    continue
-                txt = _html_to_text(b.html)
-                if txt:
-                    parts.append(txt)
-            page_texts.append("\n".join(parts))
-        elif task == "layout":
-            # No OCR text in layout mode — emit a reading-order outline of labels.
-            page_texts.append(
-                "\n".join(
-                    f"{b.position}: {b.label}"
-                    for b in sorted(page.bboxes, key=lambda b: b.position)
-                )
-            )
-        else:  # table
-            if page.html:  # mode="full"
-                page_texts.append(page.html)
-            else:  # mode="simple"
-                page_texts.append(f"{len(page.rows)} rows x {len(page.cols)} cols")
-    return "\n\n".join(page_texts), structured
-
-
-def serialize_per_page(task: str, pages: List[Any]) -> List[Tuple[str, Dict[str, Any]]]:
-    """Per-page (text, structured-dict). Reuses `serialize_pages` one page at a time
-    so the per-file dataset row and the per-page bucket files share one code path."""
-    out: List[Tuple[str, Dict[str, Any]]] = []
-    for page in pages:
-        text, structured = serialize_pages(task, [page])
-        out.append((text, structured[0]))
-    return out
-
-
-# ---------------------------------------------------------------------------
-# Bucket-URL helpers (verbatim from pp-doclayout.py)
-# ---------------------------------------------------------------------------
-
-
-def is_bucket_url(s: str) -> bool:
-    return s.startswith(BUCKET_PREFIX)
-
-
-def parse_bucket_url(url: str) -> Tuple[str, str]:
-    """Split `hf://buckets/ns/bucket/path/in/bucket` into (`ns/bucket`, `path/in/bucket`)."""
-    if not is_bucket_url(url):
-        raise ValueError(f"Not a bucket URL: {url}")
-    rest = url[len(BUCKET_PREFIX) :].strip("/")
-    parts = rest.split("/", 2)
-    if len(parts) < 2:
-        raise ValueError(f"Bucket URL must include namespace and bucket name: {url}")
-    bucket_id = f"{parts[0]}/{parts[1]}"
-    prefix = parts[2] if len(parts) > 2 else ""
-    return bucket_id, prefix
-
-
-# ---------------------------------------------------------------------------
-# Image / PDF loading
-# ---------------------------------------------------------------------------
-
-
-def open_image(path: Path) -> Image.Image:
-    """Open one image as RGB. Falls back to imagecodecs for JPEG-2000, which the
-    image's bundled Pillow may not decode (no OpenJPEG)."""
-    try:
-        return Image.open(path).convert("RGB")
-    except (UnidentifiedImageError, OSError):
-        if path.suffix.lower() in JP2_EXTENSIONS:
-            import imagecodecs
-
-            arr = imagecodecs.imread(str(path))
-            logger.debug(f"Decoded {path.name} via imagecodecs (Pillow fallback)")
-            return Image.fromarray(arr).convert("RGB")
-        raise
-
-
-def load_pages(
-    kind: str,
-    local_path: Path,
-    load_pdf,
-    page_indices: Optional[List[int]],
-    pdf_dpi: int,
-) -> List[Image.Image]:
-    """A local document file -> list of RGB page images (1 for an image, N for a PDF)."""
-    if kind == "pdf":
-        images, _ = load_pdf(str(local_path), page_indices, dpi=pdf_dpi)
-        return [im.convert("RGB") for im in images]
-    return [open_image(local_path)]
-
-
-# ---------------------------------------------------------------------------
-# File listing + sources (mount vs copy)
-# ---------------------------------------------------------------------------
-
-
-@dataclass
-class FileRef:
-    """One input document. `key`/`rel` are the source-relative POSIX path (stable
-    across runs -> resume) and drive output mirroring. `local_path` is set in mount
-    mode; `bucket_file`/`bucket_path` in copy mode."""
-
-    key: str
-    rel: PurePosixPath
-    kind: str  # "image" | "pdf"
-    local_path: Optional[Path] = None
-    bucket_file: Any = None
-    bucket_path: Optional[str] = None
-
-
-def classify(path_str: str, exts: set) -> Optional[str]:
-    """Map a path to "pdf"/"image"/None using the allowed-extension set."""
-    ext = PurePosixPath(path_str).suffix.lower()
-    if ext == PDF_EXTENSION and PDF_EXTENSION in exts:
-        return "pdf"
-    if ext in exts:
-        return "image"
-    return None
-
-
-def _shuffle_slice(
-    refs: List[FileRef], shuffle: bool, seed: int, max_samples: Optional[int]
-) -> List[FileRef]:
-    refs.sort(key=lambda r: r.key)
-    if shuffle:
-        import random
-
-        random.Random(seed).shuffle(refs)
-    if max_samples:
-        refs = refs[:max_samples]
-    return refs
-
-
-class MountSource:
-    """Read files straight off a directory (a bucket mounted read-only at /in)."""
-
-    mode = "mount"
-
-    def __init__(self, root: Path, glob: str, exts: set):
-        self.root = root
-        self.glob = glob
-        self.exts = exts
-
-    def list_refs(
-        self, shuffle: bool, seed: int, max_samples: Optional[int]
-    ) -> List[FileRef]:
-        refs: List[FileRef] = []
-        for path in self.root.rglob("*"):
-            if not path.is_file():
-                continue
-            rel = path.relative_to(self.root)
-            rel_posix = rel.as_posix()
-            kind = classify(rel_posix, self.exts)
-            if kind is None or not fnmatch(rel_posix, self.glob):
-                continue
-            refs.append(
-                FileRef(
-                    key=rel_posix,
-                    rel=PurePosixPath(rel_posix),
-                    kind=kind,
-                    local_path=path,
-                )
-            )
-        return _shuffle_slice(refs, shuffle, seed, max_samples)
-
-    @contextmanager
-    def materialize(
-        self, chunk: List[FileRef], load_pdf, page_indices, pdf_dpi
-    ) -> Iterator[List[Tuple[FileRef, Optional[List[Image.Image]]]]]:
-        loaded: List[Tuple[FileRef, Optional[List[Image.Image]]]] = []
-        for ref in chunk:
-            loaded.append(
-                (
-                    ref,
-                    _safe_load(
-                        ref.kind, ref.local_path, load_pdf, page_indices, pdf_dpi
-                    ),
-                )
-            )
-        yield loaded  # nothing to clean up — reads are off the mount
-
-
-class CopySource:
-    """List + batch-download bucket files via huggingface_hub to local temp, then
-    delete the batch. The non-FUSE path (sidesteps the bulk-read stall)."""
-
-    mode = "copy"
-
-    def __init__(self, bucket_url: str, glob: str, exts: set, hf_token: Optional[str]):
-        from huggingface_hub import HfApi
-
-        self.bucket_id, self.prefix = parse_bucket_url(bucket_url)
-        self.glob = glob
-        self.exts = exts
-        self.hf_token = hf_token
-        self.api = HfApi(token=hf_token)
-
-    def list_refs(
-        self, shuffle: bool, seed: int, max_samples: Optional[int]
-    ) -> List[FileRef]:
-        logger.info(
-            f"Listing bucket {self.bucket_id}"
-            + (f"/{self.prefix}" if self.prefix else "")
-        )
-        refs: List[FileRef] = []
-        for item in self.api.list_bucket_tree(
-            self.bucket_id, prefix=self.prefix or None, recursive=True
-        ):
-            path = getattr(item, "path", None)
-            if not path:
-                continue
-            kind = classify(path, self.exts)
-            if kind is None:
-                continue
-            rel = path[len(self.prefix) :].lstrip("/") if self.prefix else path
-            if not fnmatch(rel, self.glob):
-                continue
-            refs.append(
-                FileRef(
-                    key=rel,
-                    rel=PurePosixPath(rel),
-                    kind=kind,
-                    bucket_file=item,
-                    bucket_path=path,
-                )
-            )
-        logger.info(f"Found {len(refs)} matching file(s) in bucket")
-        return _shuffle_slice(refs, shuffle, seed, max_samples)
-
-    @contextmanager
-    def materialize(
-        self, chunk: List[FileRef], load_pdf, page_indices, pdf_dpi
-    ) -> Iterator[List[Tuple[FileRef, Optional[List[Image.Image]]]]]:
-        tmp = Path(tempfile.mkdtemp(prefix="surya-copy-"))
-        try:
-            # Pass the BucketFile objects from list_bucket_tree so download skips the
-            # per-file metadata HEAD. Local names are index-keyed to avoid collisions.
-            files = []
-            locals_: List[Path] = []
-            for i, ref in enumerate(chunk):
-                local = tmp / f"{i:05d}{PurePosixPath(ref.bucket_path).suffix}"
-                files.append((ref.bucket_file, str(local)))
-                locals_.append(local)
-            self.api.download_bucket_files(
-                self.bucket_id, files=files, token=self.hf_token
-            )
-            loaded: List[Tuple[FileRef, Optional[List[Image.Image]]]] = []
-            for ref, local in zip(chunk, locals_):
-                if not local.exists():
-                    logger.warning(f"Download missing for {ref.key}; skipping")
-                    loaded.append((ref, None))
-                    continue
-                loaded.append(
-                    (ref, _safe_load(ref.kind, local, load_pdf, page_indices, pdf_dpi))
-                )
-            yield loaded
-        finally:
-            shutil.rmtree(tmp, ignore_errors=True)
-
-
-def _safe_load(
-    kind: str, local_path: Path, load_pdf, page_indices, pdf_dpi
-) -> Optional[List[Image.Image]]:
-    try:
-        return load_pages(kind, local_path, load_pdf, page_indices, pdf_dpi)
-    except Exception as e:  # noqa: BLE001 — a single bad file shouldn't kill the run
-        logger.warning(f"Failed to load {local_path.name}: {type(e).__name__}: {e}")
-        return None
-
-
-# ---------------------------------------------------------------------------
-# Sinks
-# ---------------------------------------------------------------------------
-
-
-class BucketFilesSink:
-    """Per page, write `<rel>.md` + `<rel>.json` (PDFs: `<stem>/page_NNN.{md,json}`),
-    mirroring the input structure, to a mounted dir OR an `hf://buckets/...` URL.
-    Streaming / O(1) memory. Resume-by-skip keys on the `.json` (written last)."""
-
-    def __init__(self, output_target: str, hf_token: Optional[str], resume: bool):
-        self.resume = resume
-        self.api_mode = is_bucket_url(output_target)
-        if self.api_mode:
-            from huggingface_hub import HfApi
-
-            self.bucket_id, self.prefix = parse_bucket_url(output_target)
-            self.api = HfApi(token=hf_token)
-            self.token = hf_token
-            self._buffer: List[Tuple[bytes, str]] = []
-            self._existing = self._load_existing() if resume else set()
-        else:
-            self.root = Path(output_target)
-            self.root.mkdir(parents=True, exist_ok=True)
-
-    @property
-    def label(self) -> str:
-        return (
-            f"hf://buckets/{self.bucket_id}/{self.prefix}".rstrip("/")
-            if self.api_mode
-            else str(self.root)
-        )
-
-    def _join(self, rel: str) -> str:
-        return f"{self.prefix}/{rel}".lstrip("/") if self.prefix else rel
-
-    def _load_existing(self) -> set:
-        existing = set()
-        try:
-            for item in self.api.list_bucket_tree(
-                self.bucket_id, prefix=self.prefix or None, recursive=True
-            ):
-                p = getattr(item, "path", None)
-                if p and p.endswith(".json"):
-                    existing.add(p)
-        except Exception as e:  # noqa: BLE001
-            logger.warning(f"Could not pre-list output bucket for resume: {e}")
-        if existing:
-            logger.info(f"Resume: {len(existing)} output file(s) already present")
-        return existing
-
-    def _page_targets(self, ref: FileRef, n_pages: int) -> List[Tuple[str, str]]:
-        if ref.kind == "pdf":
-            stem = ref.rel.with_suffix("")
-            return [
-                (
-                    str(stem / f"page_{i + 1:03d}.md"),
-                    str(stem / f"page_{i + 1:03d}.json"),
-                )
-                for i in range(n_pages)
-            ]
-        return [(str(ref.rel.with_suffix(".md")), str(ref.rel.with_suffix(".json")))]
-
-    def is_done(self, ref: FileRef) -> bool:
-        # Resume applies to single-image files only; PDFs are re-rendered (idempotent
-        # overwrite) since page count isn't known without opening them.
-        if not self.resume or ref.kind == "pdf":
-            return False
-        json_rel = str(ref.rel.with_suffix(".json"))
-        if self.api_mode:
-            return self._join(json_rel) in self._existing
-        return (self.root / json_rel).exists()
-
-    def write_pages(
-        self,
-        ref: FileRef,
-        per_page: List[Tuple[str, Dict[str, Any]]],
-        pages: Optional[List[Image.Image]],
-    ) -> None:
-        targets = self._page_targets(ref, len(per_page))
-        for (text, struct), (md_rel, json_rel) in zip(per_page, targets):
-            md_bytes = text.encode("utf-8")
-            json_bytes = json.dumps(struct, ensure_ascii=False).encode("utf-8")
-            if self.api_mode:
-                # .md first, .json last so a present .json marks the page complete.
-                self._buffer.append((md_bytes, self._join(md_rel)))
-                self._buffer.append((json_bytes, self._join(json_rel)))
-            else:
-                mp = self.root / md_rel
-                mp.parent.mkdir(parents=True, exist_ok=True)
-                mp.write_bytes(md_bytes)
-                (self.root / json_rel).write_bytes(json_bytes)
-
-    def write_error(self, ref: FileRef) -> None:
-        # Write nothing on error so the file is retried on the next (resumed) run.
-        pass
-
-    def flush(self) -> None:
-        if self.api_mode and self._buffer:
-            self.api.batch_bucket_files(
-                self.bucket_id, add=self._buffer, token=self.token
-            )
-            self._buffer = []
-
-    def finalize(self, summary: Dict[str, Any]) -> None:
-        self.flush()
-        logger.info(f"Bucket files written to {self.label}")
-
-
-class DatasetSink:
-    """Buffer one row per file, push a parquet dataset at the end (like surya-ocr.py)."""
-
-    def __init__(
-        self,
-        repo_id: str,
-        *,
-        hf_token: Optional[str],
-        private: bool,
-        config: Optional[str],
-        create_pr: bool,
-        include_images: bool,
-        output_column: str,
-        blocks_column: str,
-    ):
-        self.repo_id = repo_id
-        self.hf_token = hf_token
-        self.private = private
-        self.config = config
-        self.create_pr = create_pr
-        self.include_images = include_images
-        self.output_column = output_column
-        self.blocks_column = blocks_column
-        self._rows: List[Dict[str, Any]] = []
-
-    def is_done(self, ref: FileRef) -> bool:
-        return False  # single push at the end; no per-file resume
-
-    def write_pages(
-        self,
-        ref: FileRef,
-        per_page: List[Tuple[str, Dict[str, Any]]],
-        pages: Optional[List[Image.Image]],
-    ) -> None:
-        row = {
-            "file_name": ref.key,
-            "num_pages": len(per_page),
-            self.output_column: "\n\n".join(t for t, _ in per_page),
-            self.blocks_column: json.dumps(
-                [s for _, s in per_page], ensure_ascii=False
-            ),
-        }
-        if self.include_images and pages:
-            # First page only (keeps a single Image column); documented limitation.
-            row["image"] = pages[0]
-        self._rows.append(row)
-
-    def write_error(self, ref: FileRef) -> None:
-        self._rows.append(
-            {
-                "file_name": ref.key,
-                "num_pages": 0,
-                self.output_column: "[SURYA ERROR]",
-                self.blocks_column: None,
-            }
-        )
-
-    def flush(self) -> None:
-        pass  # single push at finalize
-
-    def finalize(self, summary: Dict[str, Any]) -> None:
-        from datasets import Dataset
-
-        if not self._rows:
-            logger.warning("No rows produced; nothing to push to the dataset.")
-            return
-
-        inference_entry = {
-            "model": summary["model"],
-            "model_name": "surya-ocr-2",
-            "column_name": self.output_column,
-            "blocks_column": self.blocks_column,
-            "task": summary["task"],
-            "table_mode": summary["table_mode"] if summary["task"] == "table" else None,
-            "backend": "vllm-offline",
-            "source": summary["source"],
-            "io_mode": summary["io_mode"],
-            "glob": summary["glob"],
-            "page_range": summary["page_range"],
-            "error_rate": summary["error_rate"],
-            "timestamp": datetime.now(timezone.utc).isoformat(),
-            "script": "surya-ocr-bucket.py",
-        }
-        for row in self._rows:
-            row["inference_info"] = json.dumps([inference_entry])
-
-        ds = Dataset.from_list(self._rows)
-        if self.include_images and "image" in ds.column_names:
-            try:
-                from datasets import Image as HFImage
-
-                ds = ds.cast_column("image", HFImage())
-            except Exception as e:  # noqa: BLE001
-                logger.warning(f"Could not cast image column: {e}")
-
-        logger.info(f"Pushing {len(ds)} rows to {self.repo_id}")
-        push_kwargs = {
-            "private": self.private,
-            "token": self.hf_token,
-            "max_shard_size": "500MB",
-            "create_pr": self.create_pr,
-            "commit_message": f"Add Surya OCR 2 {summary['task']} results ({len(ds)} files)"
-            + (f" [{self.config}]" if self.config else ""),
-        }
-        if self.config:
-            push_kwargs["config_name"] = self.config
-
-        for attempt in range(1, 4):
-            try:
-                if attempt > 1:
-                    logger.warning("Disabling XET (fallback to HTTP upload)")
-                    os.environ["HF_HUB_DISABLE_XET"] = "1"
-                ds.push_to_hub(self.repo_id, **push_kwargs)
-                break
-            except Exception as e:  # noqa: BLE001
-                logger.error(f"Upload attempt {attempt}/3 failed: {e}")
-                if attempt == 3:
-                    logger.error("All upload attempts failed.")
-                    raise
-                time.sleep(30 * (2 ** (attempt - 1)))
-
-        self._push_card(summary, len(ds))
-        logger.info(f"Dataset: https://huggingface.co/datasets/{self.repo_id}")
-
-    def _push_card(self, summary: Dict[str, Any], n_rows: int) -> None:
-        try:
-            from huggingface_hub import DatasetCard
-
-            card = DatasetCard(
-                _dataset_card(
-                    source=summary["source"],
-                    model=summary["model"],
-                    task=summary["task"],
-                    table_mode=summary["table_mode"],
-                    io_mode=summary["io_mode"],
-                    n_files=n_rows,
-                    n_ok=summary["n_ok"],
-                    output_column=self.output_column,
-                    blocks_column=self.blocks_column,
-                    processing_time=summary["processing_time"],
-                )
-            )
-            card.push_to_hub(self.repo_id, token=self.hf_token)
-        except Exception as e:  # noqa: BLE001
-            logger.warning(f"Could not push dataset card: {e}")
-
-
-def _dataset_card(
-    source: str,
-    model: str,
-    task: str,
-    table_mode: str,
-    io_mode: str,
-    n_files: int,
-    n_ok: int,
-    output_column: str,
-    blocks_column: str,
-    processing_time: str,
-) -> str:
-    task_desc = {
-        "ocr": "full-page OCR (structured HTML + bounding boxes)",
-        "layout": "layout analysis (labelled regions + reading order)",
-        "table": f"table recognition (mode `{table_mode}`)",
-    }[task]
-    return f"""---
-tags:
-- ocr
-- document-processing
-- surya
-- structured
-- uv-script
-- generated
----
-
-# Surya OCR 2 ({task}) on {source}
-
-{task_desc.capitalize()} over document files in the HF bucket
-`{source}`, using [Surya OCR 2](https://huggingface.co/{model}) (650M, Qwen3.5-based)
-by Datalab, via the [`surya-ocr`](https://github.com/datalab-to/surya) package, run
-as **offline vLLM batch inference** on Hugging Face Jobs (`surya-ocr-bucket.py`).
-
-## Processing Details
-
-- **Source bucket**: `{source}`
-- **Model**: [{model}](https://huggingface.co/{model})
-- **Task**: `{task}`{f" (table mode `{table_mode}`)" if task == "table" else ""}
-- **I/O mode**: `{io_mode}`
-- **Text column**: `{output_column}` (flattened, reading-order text per file)
-- **Structured column**: `{blocks_column}` (JSON: per-page blocks with bbox / polygon / label / reading_order / confidence / html)
-- **Files**: {n_files:,}
-- **Processed OK**: {n_ok:,} / {n_files:,}
-- **Processing time**: {processing_time}
-- **Date**: {datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC")}
-
-## License note
-
-Surya's code is Apache-2.0, but the model **weights** use a modified OpenRAIL-M
-license: free for research, personal use, and startups under $5M funding/revenue,
-restricted from competitive use against Datalab's API. See the
-[model card](https://huggingface.co/{model}).
-
-## Dataset Structure
-
-One row per source file:
-- `file_name`: source-relative path in the bucket
-- `num_pages`: pages OCR'd (1 for an image, N for a PDF)
-- `{output_column}`: flattened text (OCR), label outline (layout), or table HTML (table)
-- `{blocks_column}`: structured result as a JSON string (one entry per page)
-- `inference_info`: JSON list tracking models applied
-
-Generated with [UV Scripts](https://huggingface.co/uv-scripts).
-"""
-
-
-# ---------------------------------------------------------------------------
-# Predictor + processing loop
-# ---------------------------------------------------------------------------
-
-
-def build_predictor(task: str, table_mode: str, manager):
-    """Return a `run(images) -> page_results` closure (verbatim dispatch from parent)."""
-    if task == "ocr":
-        from surya.recognition import RecognitionPredictor
-
-        predictor = RecognitionPredictor(manager)
-
-        def run(images):
-            return predictor(images, full_page=True)
-    elif task == "layout":
-        from surya.layout import LayoutPredictor
-
-        predictor = LayoutPredictor(manager)
-
-        def run(images):
-            return predictor(images)
-    else:  # table
-        from surya.table_rec import TableRecPredictor
-
-        predictor = TableRecPredictor(manager)
-
-        def run(images):
-            return predictor(images, mode=table_mode)
-
-    return run
-
-
-def process(
-    refs: List[FileRef],
-    source,
-    run,
-    task: str,
-    sinks: List[Any],
-    batch_size: int,
-    load_pdf,
-    page_indices: Optional[List[int]],
-    pdf_dpi: int,
-) -> Tuple[int, int, int, float, float]:
-    """Resume-filter, then OCR file-by-file in batches.
-
-    Returns (processed, ok, errors, io_secs, inf_secs). `io_secs` is time spent
-    materializing batches (FUSE reads in mount mode; list-skip + batch download in
-    copy mode); `inf_secs` is engine time (incl. one-time model load on the first
-    batch). The split lets the mount-vs-copy benchmark isolate I/O from inference."""
-    pending = [r for r in refs if not all(s.is_done(r) for s in sinks)]
-    skipped = len(refs) - len(pending)
-    if skipped:
-        logger.info(f"Resume: skipping {skipped} already-complete file(s)")
-    logger.info(f"Processing {len(pending)} file(s)")
-
-    processed = ok = errors = 0
-    io_secs = inf_secs = 0.0
-    pbar = tqdm(total=len(pending), desc=f"Surya {task}")
-    for chunk in partition_all(batch_size, pending):
-        chunk = list(chunk)
-        t_io = time.monotonic()
-        with source.materialize(chunk, load_pdf, page_indices, pdf_dpi) as loaded:
-            io_secs += time.monotonic() - t_io
-            entries: List[Tuple[FileRef, List[Image.Image], int, int]] = []
-            flat: List[Image.Image] = []
-            for ref, pages in loaded:
-                if not pages:
-                    for s in sinks:
-                        s.write_error(ref)
-                    errors += 1
-                    processed += 1
-                    pbar.update(1)
-                    continue
-                entries.append((ref, pages, len(flat), len(pages)))
-                flat.extend(pages)
-
-            if flat:
-                t_inf = time.monotonic()
-                try:
-                    results = run(flat)
-                except Exception as e:  # noqa: BLE001
-                    logger.error(f"Batch generate failed: {e}")
-                    results = None
-                inf_secs += time.monotonic() - t_inf
-
-                if results is None:
-                    for ref, _pages, _start, _count in entries:
-                        for s in sinks:
-                            s.write_error(ref)
-                        errors += 1
-                        processed += 1
-                        pbar.update(1)
-                else:
-                    for ref, pages, start, count in entries:
-                        per_page = serialize_per_page(
-                            task, results[start : start + count]
-                        )
-                        for s in sinks:
-                            s.write_pages(ref, per_page, pages)
-                        ok += 1
-                        processed += 1
-                        pbar.update(1)
-
-        for s in sinks:
-            s.flush()
-    pbar.close()
-    return processed, ok, errors, io_secs, inf_secs
-
-
-# ---------------------------------------------------------------------------
-# Main
-# ---------------------------------------------------------------------------
-
-
-def resolve_io_mode(io_mode: str, input_source: str) -> str:
-    if io_mode == "auto":
-        return "copy" if is_bucket_url(input_source) else "mount"
-    return io_mode
-
-
-def main(args: argparse.Namespace) -> None:
-    # Unlock full Xet bandwidth for the model download (repo convention).
-    os.environ["HF_XET_HIGH_PERFORMANCE"] = "1"
-    # Surya reads settings from env at import; pin the checkpoint and forbid any
-    # server autostart (we inject our own offline backend instead).
-    os.environ["SURYA_MODEL_CHECKPOINT"] = args.model
-    os.environ["SURYA_INFERENCE_AUTOSTART"] = "False"
-
-    check_cuda_availability()
-    start_time = datetime.now(timezone.utc)
-
-    hf_token = args.hf_token or os.environ.get("HF_TOKEN")
-    if hf_token:
-        from huggingface_hub import login
-
-        login(token=hf_token)
-
-    exts = {e.strip().lower() for e in args.extensions.split(",") if e.strip()}
-    io_mode = resolve_io_mode(args.io_mode, args.input_source)
-
-    # ---------- source ----------
-    if io_mode == "copy":
-        if not is_bucket_url(args.input_source):
-            logger.error("--io-mode copy requires an hf://buckets/... input.")
-            sys.exit(1)
-        source = CopySource(args.input_source, args.glob, exts, hf_token)
-    else:
-        root = Path(args.input_source)
-        if not root.is_dir():
-            logger.error(
-                f"--io-mode mount requires an existing directory (got {root}). "
-                "Mount the bucket with -v hf://buckets/<id>:/in:ro and pass /in."
-            )
-            sys.exit(1)
-        source = MountSource(root, args.glob, exts)
-    logger.info(f"I/O mode: {io_mode}  Input: {args.input_source}")
-
-    # ---------- sinks ----------
-    sinks: List[Any] = []
-    if args.output_bucket:
-        sinks.append(
-            BucketFilesSink(args.output_bucket, hf_token, resume=not args.no_resume)
-        )
-    if args.output_dataset:
-        sinks.append(
-            DatasetSink(
-                args.output_dataset,
-                hf_token=hf_token,
-                private=args.private,
-                config=args.config,
-                create_pr=args.create_pr,
-                include_images=args.include_images,
-                output_column=args.output_column,
-                blocks_column=args.blocks_column,
-            )
-        )
-
-    # ---------- import Surya only after env is set ----------
-    from surya.input.load import load_pdf
-    from surya.settings import settings
-
-    page_indices = parse_page_range(args.page_range)
-    pdf_dpi = args.pdf_dpi if args.pdf_dpi else settings.IMAGE_DPI_HIGHRES
-
-    t_list = time.monotonic()
-    refs = source.list_refs(args.shuffle, args.seed, args.max_samples)
-    list_secs = time.monotonic() - t_list
-    if not refs:
-        logger.error("No matching files found. Check --glob / --extensions / input.")
-        sys.exit(1)
-    logger.info(
-        f"{len(refs)} file(s) listed in {list_secs:.1f}s | Model: {args.model}  "
-        f"Task: {args.task}"
-        + (f" (mode {args.table_mode})" if args.task == "table" else "")
-    )
-
-    # ---------- engine ----------
-    backend = OfflineVLLMBackend(
-        model=args.model,
-        max_model_len=args.max_model_len,
-        gpu_memory_utilization=args.gpu_memory_utilization,
-        dtype=args.dtype,
-    )
-    manager = make_manager(backend)
-    run = build_predictor(args.task, args.table_mode, manager)
-
-    processed, ok, errors, io_secs, inf_secs = process(
-        refs,
-        source,
-        run,
-        args.task,
-        sinks,
-        args.batch_size,
-        load_pdf,
-        page_indices,
-        pdf_dpi,
-    )
-
-    processing_time = (
-        f"{(datetime.now(timezone.utc) - start_time).total_seconds() / 60:.1f} min"
-    )
-    logger.info(
-        f"Processed {processed} (ok {ok}, errors {errors}) in {processing_time}"
-    )
-    # Benchmark breakdown: separate listing + per-batch I/O from engine time so the
-    # mount-vs-copy comparison isn't swamped by (identical) inference + model load.
-    pages_per_sec = ok / io_secs if io_secs else 0.0
-    logger.info(
-        f"[timing] io_mode={io_mode} list={list_secs:.1f}s io={io_secs:.1f}s "
-        f"inference={inf_secs:.1f}s files={ok} io_files_per_sec={pages_per_sec:.2f}"
-    )
-
-    summary = {
-        "model": args.model,
-        "task": args.task,
-        "table_mode": args.table_mode,
-        "source": args.input_source,
-        "io_mode": io_mode,
-        "glob": args.glob,
-        "page_range": args.page_range,
-        "n_ok": ok,
-        "error_rate": (processed - ok) / processed if processed else 0.0,
-        "processing_time": processing_time,
-    }
-    for s in sinks:
-        s.finalize(summary)
-
-    logger.info("Done! Surya OCR 2 (bucket) complete.")
-
-    if args.verbose:
-        import importlib.metadata
-
-        logger.info("--- Resolved package versions ---")
-        for pkg in [
-            "surya-ocr",
-            "vllm",
-            "transformers",
-            "torch",
-            "datasets",
-            "huggingface-hub",
-            "pillow",
-            "imagecodecs",
-        ]:
-            try:
-                logger.info(f"  {pkg}=={importlib.metadata.version(pkg)}")
-            except importlib.metadata.PackageNotFoundError:
-                logger.info(f"  {pkg}: not installed")
-
-
-# ---------------------------------------------------------------------------
-# CLI
-# ---------------------------------------------------------------------------
-
-
-def build_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser(
-        description="Surya OCR 2 (650M): structured OCR / layout / tables over a bucket of files",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog="""
-I/O modes (--io-mode):
-  auto    copy for an hf://buckets/... input, mount for a local dir (default)
-  mount   read off a bucket mounted read-only at /in (-v hf://buckets/<id>:/in:ro)
-  copy    list + batch-download via huggingface_hub to temp, OCR, delete the batch
-
-Outputs (at least one required):
-  --output-bucket   per-page .md + .json mirroring input structure (mounted dir or
-                    hf://buckets/... URL); resumable, O(1) memory
-  --output-dataset  parquet dataset push (one row per file)
-
-Run on the vllm/vllm-openai:v0.20.1 image (offline vLLM batch; qwen3_5 is
-version-sensitive — the site-packages python path is load-bearing):
-  --image vllm/vllm-openai:v0.20.1 --python /usr/local/bin/python3 \\
-    -e PYTHONPATH=/usr/local/lib/python3.12/site-packages
-""",
-    )
-    parser.add_argument(
-        "input_source",
-        help="Mounted dir (e.g. /in) OR hf://buckets/<ns>/<bucket>[/prefix]",
-    )
-    parser.add_argument(
-        "--io-mode",
-        choices=["auto", "mount", "copy"],
-        default="auto",
-        help="Input I/O strategy (default: auto)",
-    )
-    parser.add_argument(
-        "--glob",
-        default="*",
-        help="fnmatch pattern over the source-relative path (default: '*'; "
-        "e.g. '*.jp2'). Applied on top of --extensions.",
-    )
-    parser.add_argument(
-        "--extensions",
-        default=DEFAULT_EXTENSIONS,
-        help=f"Comma-separated file extensions to read (default: {DEFAULT_EXTENSIONS})",
-    )
-    parser.add_argument(
-        "--output-bucket",
-        default=None,
-        help="Per-file .md + .json output: a mounted dir OR hf://buckets/<id>[/prefix]",
-    )
-    parser.add_argument(
-        "--output-dataset",
-        default=None,
-        help="Output dataset repo ID (parquet, one row per file)",
-    )
-    parser.add_argument(
-        "--no-resume",
-        action="store_true",
-        help="Disable resume-by-skip for --output-bucket (re-OCR everything)",
-    )
-    parser.add_argument(
-        "--task", choices=TASKS, default="ocr", help="Task (default: ocr)"
-    )
-    parser.add_argument(
-        "--table-mode",
-        choices=["full", "simple"],
-        default="full",
-        help="Table task: 'full' = HTML, 'simple' = rows/cols/cells (default: full)",
-    )
-    parser.add_argument(
-        "--page-range",
-        default=None,
-        help="Pages from PDFs, e.g. '0-5,7' (PDFs only)",
-    )
-    parser.add_argument(
-        "--pdf-dpi",
-        type=int,
-        default=None,
-        help="DPI for PDF rendering (default: Surya's IMAGE_DPI_HIGHRES)",
-    )
-    parser.add_argument(
-        "--max-samples", type=int, help="Limit number of files (for testing)"
-    )
-    parser.add_argument(
-        "--shuffle", action="store_true", help="Shuffle before sampling"
-    )
-    parser.add_argument(
-        "--seed", type=int, default=42, help="Shuffle seed (default: 42)"
-    )
-    parser.add_argument(
-        "--batch-size",
-        type=int,
-        default=16,
-        help="Images per offline llm.chat batch AND per copy-mode download/cleanup unit (default: 16)",
-    )
-    parser.add_argument(
-        "--max-model-len",
-        type=int,
-        default=18000,
-        help="vLLM context length (default: 18000)",
-    )
-    parser.add_argument(
-        "--gpu-memory-utilization",
-        type=float,
-        default=0.85,
-        help="vLLM GPU memory fraction (default: 0.85)",
-    )
-    parser.add_argument(
-        "--dtype",
-        default="bfloat16",
-        help="vLLM dtype (default: bfloat16; use float16 on T4/Turing)",
-    )
-    parser.add_argument(
-        "--model", default=DEFAULT_MODEL, help=f"Model ID (default: {DEFAULT_MODEL})"
-    )
-    parser.add_argument(
-        "--output-column",
-        default="markdown",
-        help="Dataset text column (default: markdown)",
-    )
-    parser.add_argument(
-        "--blocks-column",
-        default="surya_blocks",
-        help="Dataset structured JSON column (default: surya_blocks)",
-    )
-    parser.add_argument(
-        "--include-images",
-        action="store_true",
-        help="Embed the first page image in --output-dataset (memory-heavy)",
-    )
-    parser.add_argument(
-        "--private", action="store_true", help="Make output dataset private"
-    )
-    parser.add_argument(
-        "--config",
-        default=None,
-        help="Config/subset name when pushing (for benchmarking in one repo)",
-    )
-    parser.add_argument(
-        "--create-pr",
-        action="store_true",
-        help="Push dataset as a pull request instead of directly",
-    )
-    parser.add_argument("--hf-token", help="Hugging Face API token (or set HF_TOKEN)")
-    parser.add_argument(
-        "--verbose",
-        action="store_true",
-        help="Log resolved package versions after processing",
-    )
-    return parser
-
-
-def _print_banner() -> None:
-    print(
-        "Surya OCR 2 (bucket) — structured OCR / layout / tables over a bucket of files (650M)"
-    )
-    print("\nUsage:")
-    print(
-        "  uv run surya-ocr-bucket.py INPUT [--output-bucket ... | --output-dataset ...] [options]"
-    )
-    print("\nExamples:")
-    print("  # copy a bucket of .jp2 -> a dataset")
-    print("  uv run surya-ocr-bucket.py hf://buckets/me/news --io-mode copy \\")
-    print("      --glob '*.jp2' --output-dataset me/news-ocr --private")
-    print("\n  # mount a bucket -> per-file .md + .json in an output bucket")
-    print("  uv run surya-ocr-bucket.py /in --io-mode mount --output-bucket /out")
-    print("\nRun on the vllm/vllm-openai:v0.20.1 image (offline vLLM batch):")
-    print("  hf jobs uv run --flavor l4x1 -s HF_TOKEN \\")
-    print("      --image vllm/vllm-openai:v0.20.1 --python /usr/local/bin/python3 \\")
-    print("      -e PYTHONPATH=/usr/local/lib/python3.12/site-packages \\")
-    print("      -v hf://buckets/me/news:/in:ro -v hf://buckets/me/news-ocr:/out \\")
-    print(
-        "      https://huggingface.co/datasets/uv-scripts/ocr/raw/main/surya-ocr-bucket.py \\"
-    )
-    print("      /in --io-mode mount --glob '*.jp2' --output-bucket /out")
-    print("\nFor full help: uv run surya-ocr-bucket.py --help")
-
-
-if __name__ == "__main__":
-    if len(sys.argv) == 1:
-        _print_banner()
-        sys.exit(0)
-
-    args = build_parser().parse_args()
-    if not args.output_bucket and not args.output_dataset:
-        build_parser().error(
-            "at least one of --output-bucket or --output-dataset is required"
-        )
-    if args.no_resume and not args.output_bucket:
-        logger.warning("--no-resume has no effect without --output-bucket")
-    if args.include_images and not args.output_dataset:
-        logger.warning("--include-images has no effect without --output-dataset")
-
-    main(args)

surya-ocr.py

@@ -1,915 +0,0 @@
-# /// script
-# requires-python = ">=3.10"
-# dependencies = [
-#     "surya-ocr",
-#     "datasets>=3.1.0",
-#     "huggingface-hub",
-#     "pillow",
-#     "toolz",
-#     "tqdm",
-# ]
-# ///
-"""
-Document intelligence on images OR multi-page PDFs with Datalab's **Surya OCR 2**
-(`datalab-to/surya-ocr-2`, 650M, Qwen3.5-style).
-
-Surya is *structured* OCR: instead of a flat markdown blob, it returns per-block
-HTML with bounding boxes, reading order, and labels (equations in `<math>`). This
-recipe writes **both**:
-
-  --output-column   (default `markdown`)  flattened, reading-order text per row
-  surya_blocks                            the full structured result as JSON
-                                          (bbox / polygon / label / reading_order /
-                                          confidence / html per block), one entry
-                                          per page.
-
-Three tasks via `--task`:
-  ocr     (default)  full-page OCR  -> text + per-block HTML/bboxes
-  layout             layout regions -> labelled boxes + reading order
-  table              table structure -> HTML (mode `full`) or rows/cols/cells
-                     (mode `simple`, via --table-mode)
-
-Input is one document per row:
-  --image-column COL   (default `image`)  one image per row
-  --pdf-column COL                        PDF bytes per row (multi-page; honors
-                                          --page-range). Pages are concatenated in
-                                          the text column and kept per-page in
-                                          `surya_blocks`.
-
-ENGINE: Surya normally spawns a vLLM **server** (Docker) — which can't run inside
-an HF Job. This script instead does **offline batch inference**: it injects a
-custom in-process backend into Surya's `SuryaInferenceManager` that runs vLLM's
-offline `LLM().chat()` engine (no server, no HTTP). Surya still owns all the
-prompting, image preprocessing, and HTML/bbox parsing — we only swap the
-transport. Run on the **`vllm/vllm-openai:v0.20.1`** image (Surya's known-good
-vLLM build; the model is the recent, version-sensitive `qwen3_5` architecture).
-
-LICENSE NOTE: Surya's *code* is Apache-2.0 but the *weights* are a modified
-OpenRAIL-M license — free for research, personal use, and startups under $5M
-funding/revenue, but restricted from competitive use against Datalab's API.
-Confirm you are within those terms. https://huggingface.co/datalab-to/surya-ocr-2
-
-HF Jobs (use the pinned vLLM image so vLLM + qwen3_5 support are present):
-
-    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 \\
-        https://huggingface.co/datasets/uv-scripts/ocr/raw/main/surya-ocr.py \\
-        INPUT_DATASET OUTPUT_DATASET \\
-        --max-samples 5 --shuffle --seed 42
-
-Model: datalab-to/surya-ocr-2  (package: surya-ocr, https://github.com/datalab-to/surya)
-"""
-
-import argparse
-import io
-import json
-import logging
-import math
-import os
-import sys
-import tempfile
-import time
-from datetime import datetime, timezone
-from typing import Any, Dict, List, Optional, Tuple
-from urllib.request import urlopen
-
-from datasets import load_dataset
-from huggingface_hub import DatasetCard, login
-from PIL import Image
-from toolz import partition_all
-from tqdm import tqdm
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-DEFAULT_MODEL = "datalab-to/surya-ocr-2"
-# Surya's own vision-tiling bounds (from its vLLM backend), applied to the
-# offline engine too so preprocessing matches the server path exactly.
-MM_PROCESSOR_KWARGS = {"min_pixels": 3136, "max_pixels": 6291456}
-TASKS = ("ocr", "layout", "table")
-
-
-def check_cuda_availability() -> None:
-    """Exit early with a clear message if there's no GPU."""
-    import torch
-
-    if not torch.cuda.is_available():
-        logger.error("CUDA is not available. This script requires a GPU.")
-        logger.error(
-            "Run on Hugging Face Jobs with: hf jobs uv run --flavor l4x1 "
-            "--image vllm/vllm-openai:v0.20.1 ..."
-        )
-        sys.exit(1)
-    logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
-
-
-def check_vllm_available() -> None:
-    """Fail fast (before loading 400 rows) if vLLM isn't importable.
-
-    Surya-2 runs its VLM through vLLM's offline engine, but `vllm` is deliberately
-    NOT a PEP723 dependency: the recent hybrid `qwen3_5` architecture is only in the
-    pinned `vllm/vllm-openai:v0.20.1` image, which also provides torch/transformers via
-    PYTHONPATH. Launched on the bare uv image (no `--image`), the import fails per-batch
-    and every row silently gets "[SURYA GENERATE ERROR]". Detect that up front instead.
-    """
-    import importlib.util
-
-    if importlib.util.find_spec("vllm") is None:
-        logger.error("vLLM is not importable — this recipe cannot run on the bare uv image.")
-        logger.error(
-            "Surya-2 needs the pinned vLLM build; re-run with the image + interpreter flags:"
-        )
-        logger.error(
-            "  hf jobs uv run --flavor l4x1 -s HF_TOKEN \\\n"
-            "      --image vllm/vllm-openai:v0.20.1 --python /usr/local/bin/python3 \\\n"
-            "      -e PYTHONPATH=/usr/local/lib/python3.12/site-packages \\\n"
-            "      <script_url> INPUT_DATASET OUTPUT_DATASET ..."
-        )
-        sys.exit(1)
-
-
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
-def parse_page_range(spec: Optional[str]) -> Optional[List[int]]:
-    """Turn '0-3,5' into [0,1,2,3,5]. None/empty -> None (all pages)."""
-    if not spec:
-        return None
-    pages: List[int] = []
-    for part in spec.split(","):
-        part = part.strip()
-        if not part:
-            continue
-        if "-" in part:
-            lo, hi = part.split("-", 1)
-            pages.extend(range(int(lo), int(hi) + 1))
-        else:
-            pages.append(int(part))
-    return pages or None
-
-
-def cell_to_bytes(cell: Any) -> bytes:
-    """Normalize an HF dataset cell (image or document) to raw file bytes."""
-    if isinstance(cell, Image.Image):
-        buf = io.BytesIO()
-        cell.convert("RGB").save(buf, format="PNG")
-        return buf.getvalue()
-    if isinstance(cell, dict):
-        if cell.get("bytes"):
-            return cell["bytes"]
-        if cell.get("path"):
-            with open(cell["path"], "rb") as f:
-                return f.read()
-        raise ValueError(
-            f"Unsupported image/document dict (no bytes/path): {list(cell)}"
-        )
-    if isinstance(cell, (bytes, bytearray)):
-        return bytes(cell)
-    if isinstance(cell, str):
-        if cell.startswith(("http://", "https://")):
-            return urlopen(cell).read()  # noqa: S310
-        with open(cell, "rb") as f:
-            return f.read()
-    raise ValueError(f"Unsupported cell type: {type(cell)}")
-
-
-def cell_to_pil(cell: Any) -> Image.Image:
-    """One image cell -> RGB PIL image."""
-    if isinstance(cell, Image.Image):
-        return cell.convert("RGB")
-    return Image.open(io.BytesIO(cell_to_bytes(cell))).convert("RGB")
-
-
-def load_pdf_images(
-    load_pdf, cell: Any, page_indices: Optional[List[int]], dpi: int
-) -> List[Image.Image]:
-    """Render one PDF cell into page images via Surya's own pypdfium2 loader."""
-    data = cell_to_bytes(cell)
-    with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp:
-        tmp.write(data)
-        path = tmp.name
-    try:
-        images, _ = load_pdf(path, page_indices, dpi=dpi)
-        return [im.convert("RGB") for im in images]
-    finally:
-        os.unlink(path)
-
-
-# --- structured-output shim (vLLM API moved between versions) ---
-def build_structured_outputs(schema: Dict[str, Any]) -> Dict[str, Any]:
-    """SamplingParams kwargs for guided JSON, across vLLM versions (layout uses this)."""
-    try:
-        from vllm.sampling_params import StructuredOutputsParams  # vLLM >= 0.12
-
-        return {"structured_outputs": StructuredOutputsParams(json=schema)}
-    except (ImportError, TypeError):
-        pass
-    try:
-        from vllm.sampling_params import GuidedDecodingParams  # older vLLM
-
-        return {"guided_decoding": GuidedDecodingParams(json=schema)}
-    except (ImportError, TypeError):
-        pass
-    logger.warning(
-        "Guided JSON unavailable in this vLLM version; relying on the model."
-    )
-    return {}
-
-
-def _mean_token_prob(completion_output) -> Optional[float]:
-    """Mean exp(logprob) of the sampled tokens -> Surya's per-block `confidence`."""
-    lps = getattr(completion_output, "logprobs", None)
-    if not lps:
-        return None
-    probs: List[float] = []
-    for tid, lp_dict in zip(completion_output.token_ids, lps):
-        if not lp_dict:
-            continue
-        entry = lp_dict.get(tid)
-        if (
-            entry is None
-        ):  # sampled token not in the returned top-k; use the best we have
-            entry = max(lp_dict.values(), key=lambda e: e.logprob)
-        probs.append(math.exp(entry.logprob))
-    return sum(probs) / len(probs) if probs else None
-
-
-class OfflineVLLMBackend:
-    """Surya `Backend` (duck-typed) that runs vLLM's offline `LLM().chat()` engine.
-
-    Surya's predictors call `manager.generate(batch)` -> `backend.generate(batch)`;
-    we satisfy that contract in-process (no server). Surya keeps ownership of the
-    prompts (`PROMPT_MAPPING`), image scaling (`scale_to_fit`), and output parsing.
-    """
-
-    name = "offline-vllm"
-
-    def __init__(
-        self,
-        model: str,
-        max_model_len: int,
-        gpu_memory_utilization: float,
-        dtype: str = "bfloat16",
-        max_tokens_default: int = 2048,
-        logprobs_default: bool = True,
-    ):
-        self.model = model
-        self.max_model_len = max_model_len
-        self.gpu_memory_utilization = gpu_memory_utilization
-        self.dtype = dtype
-        self.max_tokens_default = max_tokens_default
-        self.logprobs_default = logprobs_default
-        self.llm = None
-        self._build_messages = None
-        self._scale_to_fit = None
-        self._prompt_mapping = None
-
-    def start(self):
-        from vllm import LLM
-
-        logger.info(
-            f"Loading {self.model} into vLLM offline engine (dtype={self.dtype})..."
-        )
-        self.llm = LLM(
-            model=self.model,
-            dtype=self.dtype,
-            max_model_len=self.max_model_len,
-            gpu_memory_utilization=self.gpu_memory_utilization,
-            mm_processor_kwargs=MM_PROCESSOR_KWARGS,
-            limit_mm_per_prompt={"image": 1},
-        )
-        # Reuse Surya's exact request shaping so the offline path matches the server.
-        from surya.inference.backends.openai_client import _build_messages
-        from surya.inference.prompts import PROMPT_MAPPING
-        from surya.inference.util import scale_to_fit
-
-        self._build_messages = _build_messages
-        self._scale_to_fit = scale_to_fit
-        self._prompt_mapping = PROMPT_MAPPING
-        return None
-
-    def stop(self) -> None:
-        self.llm = None
-
-    def _sampling_params(self, item):
-        from vllm import SamplingParams
-
-        max_tokens = item.max_tokens or self.max_tokens_default
-        want_logprobs = item.request_logprobs or self.logprobs_default
-        kwargs: Dict[str, Any] = dict(temperature=0.0, top_p=0.1, max_tokens=max_tokens)
-        if want_logprobs:
-            kwargs["logprobs"] = 1
-        if item.guided_json is not None:
-            kwargs.update(build_structured_outputs(item.guided_json))
-        return SamplingParams(**kwargs)
-
-    def generate(self, batch):
-        from surya.inference.schema import BatchOutputItem
-
-        if self.llm is None:
-            self.start()
-        if not batch:
-            return []
-
-        conversations = []
-        sampling_params = []
-        for item in batch:
-            prompt = item.prompt or self._prompt_mapping[item.prompt_type]
-            image = self._scale_to_fit(item.image)
-            conversations.append(self._build_messages(image, prompt))
-            sampling_params.append(self._sampling_params(item))
-
-        outputs = self.llm.chat(
-            conversations,
-            sampling_params,
-            chat_template_content_format="openai",
-            use_tqdm=False,
-        )
-
-        results = []
-        for item, out in zip(batch, outputs):
-            comp = out.outputs[0]
-            results.append(
-                BatchOutputItem(
-                    raw=comp.text,
-                    token_count=len(comp.token_ids),
-                    error=False,
-                    mean_token_prob=_mean_token_prob(comp),
-                    logprobs=None,
-                    metadata=item.metadata,  # carries page_idx/block_idx — must round-trip
-                )
-            )
-        return results
-
-
-def make_manager(backend: OfflineVLLMBackend):
-    """A SuryaInferenceManager wired to our offline backend (bypassing autodetect)."""
-    from surya.inference import SuryaInferenceManager
-
-    manager = SuryaInferenceManager.__new__(SuryaInferenceManager)
-    manager.method = backend.name
-    manager.backend = backend
-    return manager
-
-
-# --- result serialization (text column + structured surya_blocks) ---
-def _html_to_text(html: str) -> str:
-    from bs4 import BeautifulSoup
-
-    return BeautifulSoup(html, "html.parser").get_text(" ", strip=True)
-
-
-def serialize_pages(task: str, pages: List[Any]) -> Tuple[str, List[Dict[str, Any]]]:
-    """(text, structured-per-page) for one row's page results."""
-    structured = [p.model_dump(mode="json") for p in pages]
-    page_texts: List[str] = []
-    for page in pages:
-        if task == "ocr":
-            parts = []
-            for b in sorted(page.blocks, key=lambda b: b.reading_order):
-                if b.skipped or not b.html:
-                    continue
-                txt = _html_to_text(b.html)
-                if txt:
-                    parts.append(txt)
-            page_texts.append("\n".join(parts))
-        elif task == "layout":
-            # No OCR text in layout mode — emit a reading-order outline of labels.
-            page_texts.append(
-                "\n".join(
-                    f"{b.position}: {b.label}"
-                    for b in sorted(page.bboxes, key=lambda b: b.position)
-                )
-            )
-        else:  # table
-            if page.html:  # mode="full"
-                page_texts.append(page.html)
-            else:  # mode="simple"
-                page_texts.append(f"{len(page.rows)} rows x {len(page.cols)} cols")
-    return "\n\n".join(page_texts), structured
-
-
-def create_dataset_card(
-    source_dataset: str,
-    model: str,
-    task: str,
-    table_mode: str,
-    num_samples: int,
-    n_ok: int,
-    source_column: str,
-    is_pdf: bool,
-    page_range: Optional[str],
-    output_column: str,
-    blocks_column: str,
-    split: str,
-    processing_time: str,
-) -> str:
-    input_kind = "PDF documents" if is_pdf else "images"
-    col_desc = "PDF" if is_pdf else "image"
-    if page_range:
-        col_desc += f", pages {page_range}"
-    task_desc = {
-        "ocr": "full-page OCR (structured HTML + bounding boxes)",
-        "layout": "layout analysis (labelled regions + reading order)",
-        "table": f"table recognition (mode `{table_mode}`)",
-    }[task]
-    return f"""---
-tags:
-- ocr
-- document-processing
-- surya
-- structured
-- uv-script
-- generated
----
-
-# Surya OCR 2 ({task}) on {source_dataset}
-
-{task_desc.capitalize()} over {input_kind} in
-[{source_dataset}](https://huggingface.co/datasets/{source_dataset}) using
-[Surya OCR 2](https://huggingface.co/{model}) (650M, Qwen3.5-based) by Datalab, via the
-[`surya-ocr`](https://github.com/datalab-to/surya) package, run as **offline vLLM batch
-inference** on Hugging Face Jobs.
-
-## Processing Details
-
-- **Source Dataset**: [{source_dataset}](https://huggingface.co/datasets/{source_dataset})
-- **Model**: [{model}](https://huggingface.co/{model})
-- **Task**: `{task}`{f" (table mode `{table_mode}`)" if task == "table" else ""}
-- **Input column**: `{source_column}` ({col_desc})
-- **Text column**: `{output_column}` (flattened, reading-order text per row)
-- **Structured column**: `{blocks_column}` (JSON: per-page blocks with bbox / polygon / label / reading_order / confidence / html)
-- **Split**: `{split}`
-- **Samples**: {num_samples:,}
-- **Processed OK**: {n_ok:,} / {num_samples:,}
-- **Processing time**: {processing_time}
-- **Date**: {datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC")}
-
-## License note
-
-Surya's code is Apache-2.0, but the model **weights** use a modified OpenRAIL-M
-license: free for research, personal use, and startups under $5M funding/revenue,
-restricted from competitive use against Datalab's API. See the
-[model card](https://huggingface.co/{model}).
-
-## Dataset Structure
-
-Original columns plus:
-- `{output_column}`: flattened text (OCR), label outline (layout), or table HTML (table)
-- `{blocks_column}`: structured result as a JSON string (one entry per page)
-- `inference_info`: JSON list tracking models applied to this dataset
-
-Generated with [UV Scripts](https://huggingface.co/uv-scripts).
-"""
-
-
-def main(
-    input_dataset: str,
-    output_dataset: str,
-    task: str = "ocr",
-    table_mode: str = "full",
-    image_column: str = "image",
-    pdf_column: Optional[str] = None,
-    output_column: str = "markdown",
-    overwrite: bool = False,
-    blocks_column: str = "surya_blocks",
-    page_range: Optional[str] = None,
-    split: str = "train",
-    max_samples: Optional[int] = None,
-    shuffle: bool = False,
-    seed: int = 42,
-    batch_size: int = 16,
-    max_model_len: int = 18000,
-    gpu_memory_utilization: float = 0.85,
-    dtype: str = "bfloat16",
-    model: str = DEFAULT_MODEL,
-    private: bool = False,
-    config: Optional[str] = None,
-    create_pr: bool = False,
-    hf_token: Optional[str] = None,
-    verbose: bool = False,
-) -> None:
-    # Unlock full Xet bandwidth for the model download (repo convention).
-    os.environ["HF_XET_HIGH_PERFORMANCE"] = "1"
-    # Surya reads settings from env at import; pin the checkpoint and forbid any
-    # server autostart (we inject our own offline backend instead).
-    os.environ["SURYA_MODEL_CHECKPOINT"] = model
-    os.environ["SURYA_INFERENCE_AUTOSTART"] = "False"
-
-    check_cuda_availability()
-    check_vllm_available()
-    start_time = datetime.now(timezone.utc)
-
-    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")
-    if HF_TOKEN:
-        login(token=HF_TOKEN)
-
-    # Import Surya only after env is set.
-    from surya.input.load import load_pdf
-    from surya.settings import settings
-
-    source_column = pdf_column or image_column
-    is_pdf = pdf_column is not None
-    page_indices = parse_page_range(page_range)
-    pdf_dpi = settings.IMAGE_DPI_HIGHRES
-
-    logger.info(
-        f"Model: {model}  Task: {task}"
-        + (f" (mode {table_mode})" if task == "table" else "")
-    )
-    logger.info(f"Loading dataset: {input_dataset} (split={split})")
-    dataset = load_dataset(input_dataset, split=split)
-    if source_column not in dataset.column_names:
-        logger.error(
-            f"Column '{source_column}' not found. Available: {dataset.column_names}"
-        )
-        sys.exit(1)
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(
-        dataset, [output_column, blocks_column], overwrite=overwrite
-    )
-    if shuffle:
-        dataset = dataset.shuffle(seed=seed)
-    if max_samples:
-        dataset = dataset.select(range(min(max_samples, len(dataset))))
-    n = len(dataset)
-    logger.info(f"Processing {n} documents from column '{source_column}'")
-
-    # Build the offline engine + inject it into a Surya manager, then pick the predictor.
-    backend = OfflineVLLMBackend(
-        model=model,
-        max_model_len=max_model_len,
-        gpu_memory_utilization=gpu_memory_utilization,
-        dtype=dtype,
-    )
-    manager = make_manager(backend)
-
-    if task == "ocr":
-        from surya.recognition import RecognitionPredictor
-
-        predictor = RecognitionPredictor(manager)
-
-        def run(images):
-            return predictor(images, full_page=True)
-    elif task == "layout":
-        from surya.layout import LayoutPredictor
-
-        predictor = LayoutPredictor(manager)
-
-        def run(images):
-            return predictor(images)
-    else:  # table
-        from surya.table_rec import TableRecPredictor
-
-        predictor = TableRecPredictor(manager)
-
-        def run(images):
-            return predictor(images, mode=table_mode)
-
-    texts: List[Optional[str]] = [None] * n
-    blocks: List[Optional[str]] = [None] * n
-    error_flags: List[bool] = [True] * n
-
-    for chunk in tqdm(list(partition_all(batch_size, range(n))), desc=f"Surya {task}"):
-        chunk = list(chunk)
-        flat_images: List[Image.Image] = []
-        spans: List[Tuple[int, int, int]] = []  # (row_idx, start, count)
-        for i in chunk:
-            try:
-                if is_pdf:
-                    imgs = load_pdf_images(
-                        load_pdf, dataset[i][source_column], page_indices, pdf_dpi
-                    )
-                else:
-                    imgs = [cell_to_pil(dataset[i][source_column])]
-            except Exception as e:
-                logger.warning(f"Row {i}: failed to load document: {e}")
-                texts[i] = f"[SURYA LOAD ERROR] {e}"
-                blocks[i] = None
-                continue
-            if not imgs:
-                texts[i] = "[SURYA EMPTY DOCUMENT]"
-                continue
-            spans.append((i, len(flat_images), len(imgs)))
-            flat_images.extend(imgs)
-
-        if not flat_images:
-            continue
-        try:
-            results = run(flat_images)
-        except Exception as e:
-            logger.error(f"Batch generate failed: {e}")
-            for i, _, _ in spans:
-                texts[i] = "[SURYA GENERATE ERROR]"
-                blocks[i] = None
-            continue
-
-        for i, start, count in spans:
-            page_results = results[start : start + count]
-            text, structured = serialize_pages(task, page_results)
-            texts[i] = text
-            blocks[i] = json.dumps(structured, ensure_ascii=False)
-            error_flags[i] = False
-
-    n_ok = sum(not f for f in error_flags)
-    logger.info(f"Processed OK: {n_ok}/{n}")
-
-    dataset = dataset.add_column(output_column, texts)
-    dataset = dataset.add_column(blocks_column, blocks)
-
-    inference_entry = {
-        "model": model,
-        "model_name": "surya-ocr-2",
-        "column_name": output_column,
-        "blocks_column": blocks_column,
-        "task": task,
-        "table_mode": table_mode if task == "table" else None,
-        "backend": "vllm-offline",
-        "page_range": page_range,
-        "error_rate": (n - n_ok) / n if n else 0.0,
-        "timestamp": datetime.now(timezone.utc).isoformat(),
-        "script": "surya-ocr.py",
-    }
-    if "inference_info" in dataset.column_names:
-
-        def update_info(example):
-            try:
-                existing = (
-                    json.loads(example["inference_info"])
-                    if example["inference_info"]
-                    else []
-                )
-            except (json.JSONDecodeError, TypeError):
-                existing = []
-            existing.append(inference_entry)
-            return {"inference_info": json.dumps(existing)}
-
-        dataset = dataset.map(update_info)
-    else:
-        dataset = dataset.add_column(
-            "inference_info", [json.dumps([inference_entry])] * n
-        )
-
-    processing_time = (
-        f"{(datetime.now(timezone.utc) - start_time).total_seconds() / 60:.1f} min"
-    )
-
-    logger.info(f"Pushing to {output_dataset}")
-    max_retries = 3
-    for attempt in range(1, max_retries + 1):
-        try:
-            if attempt > 1:
-                logger.warning("Disabling XET (fallback to HTTP upload)")
-                os.environ["HF_HUB_DISABLE_XET"] = "1"
-            dataset.push_to_hub(
-                output_dataset,
-                private=private,
-                token=HF_TOKEN,
-                max_shard_size="500MB",
-                create_pr=create_pr,
-                **({"config_name": config} if config else {}),
-                commit_message=f"Add Surya OCR 2 {task} results ({n} samples)"
-                + (f" [{config}]" if config else ""),
-            )
-            break
-        except Exception as e:
-            logger.error(f"Upload attempt {attempt}/{max_retries} failed: {e}")
-            if attempt < max_retries:
-                delay = 30 * (2 ** (attempt - 1))
-                logger.info(f"Retrying in {delay}s...")
-                time.sleep(delay)
-            else:
-                logger.error("All upload attempts failed. Results are lost.")
-                sys.exit(1)
-
-    try:
-        card = DatasetCard(
-            create_dataset_card(
-                source_dataset=input_dataset,
-                model=model,
-                task=task,
-                table_mode=table_mode,
-                num_samples=n,
-                n_ok=n_ok,
-                source_column=source_column,
-                is_pdf=is_pdf,
-                page_range=page_range,
-                output_column=output_column,
-                blocks_column=blocks_column,
-                split=split,
-                processing_time=processing_time,
-            )
-        )
-        card.push_to_hub(output_dataset, token=HF_TOKEN)
-    except Exception as e:
-        logger.warning(f"Could not push dataset card: {e}")
-
-    logger.info("Done! Surya OCR 2 complete.")
-    logger.info(f"Dataset: https://huggingface.co/datasets/{output_dataset}")
-    logger.info(f"Processing time: {processing_time}")
-
-    if verbose:
-        import importlib.metadata
-
-        logger.info("--- Resolved package versions ---")
-        for pkg in ["surya-ocr", "vllm", "transformers", "torch", "datasets", "pillow"]:
-            try:
-                logger.info(f"  {pkg}=={importlib.metadata.version(pkg)}")
-            except importlib.metadata.PackageNotFoundError:
-                logger.info(f"  {pkg}: not installed")
-
-
-if __name__ == "__main__":
-    if len(sys.argv) == 1:
-        print(
-            "Surya OCR 2 — structured OCR / layout / tables from images & PDFs (650M)"
-        )
-        print("\nUsage:")
-        print("  uv run surya-ocr.py INPUT OUTPUT [--task ocr|layout|table] [options]")
-        print("\nExamples:")
-        print("  # full-page OCR -> text + structured surya_blocks")
-        print("  uv run surya-ocr.py my-images my-ocr")
-        print("\n  # layout regions / table structure")
-        print("  uv run surya-ocr.py my-images my-layout --task layout")
-        print("  uv run surya-ocr.py my-tables my-tables-out --task table")
-        print("\n  # multi-page PDFs")
-        print("  uv run surya-ocr.py my-pdfs my-ocr --pdf-column pdf --page-range 0-5")
-        print("\nRun on the vllm/vllm-openai:v0.20.1 image (offline vLLM batch).")
-        print("For full help: uv run surya-ocr.py --help")
-        sys.exit(0)
-
-    parser = argparse.ArgumentParser(
-        description="Surya OCR 2 (650M): structured OCR / layout / tables, offline vLLM batch",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog="""
-Tasks (--task):
-  ocr     full-page OCR -> reading-order text + per-block HTML/bboxes (default)
-  layout  layout regions -> labelled boxes + reading order
-  table   table structure -> HTML (--table-mode full) or rows/cols/cells (simple)
-
-Output columns:
-  --output-column   flattened text per row (default: markdown)
-  surya_blocks      structured JSON per row (bbox/label/reading_order/confidence/html)
-
-Input (one document per row):
-  --image-column COL   one image per row   (default: image)
-  --pdf-column COL     PDF bytes per row    (multi-page; honors --page-range)
-
-Run on the vllm/vllm-openai:v0.20.1 image:
-  --image vllm/vllm-openai:v0.20.1 --python /usr/local/bin/python3 \\
-    -e PYTHONPATH=/usr/local/lib/python3.12/site-packages
-""",
-    )
-    parser.add_argument(
-        "input_dataset", help="Input dataset ID from the Hugging Face Hub"
-    )
-    parser.add_argument(
-        "output_dataset", help="Output dataset ID for the Hugging Face Hub"
-    )
-    parser.add_argument(
-        "--task", choices=TASKS, default="ocr", help="Task (default: ocr)"
-    )
-    parser.add_argument(
-        "--table-mode",
-        choices=["full", "simple"],
-        default="full",
-        help="Table task: 'full' = HTML, 'simple' = rows/cols/cells (default: full)",
-    )
-    parser.add_argument(
-        "--image-column", default="image", help="Image column (default: image)"
-    )
-    parser.add_argument(
-        "--pdf-column",
-        default=None,
-        help="PDF column (bytes/path). Mutually exclusive with --image-column.",
-    )
-    parser.add_argument(
-        "--output-column",
-        default="markdown",
-        help="Text output column (default: markdown)",
-    )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
-    parser.add_argument(
-        "--blocks-column",
-        default="surya_blocks",
-        help="Structured JSON output column (default: surya_blocks)",
-    )
-    parser.add_argument(
-        "--page-range",
-        default=None,
-        help="Pages from PDFs, e.g. '0-5,7' (PDF column only)",
-    )
-    parser.add_argument(
-        "--split", default="train", help="Dataset split (default: train)"
-    )
-    parser.add_argument(
-        "--max-samples", type=int, help="Limit number of documents (for testing)"
-    )
-    parser.add_argument(
-        "--shuffle", action="store_true", help="Shuffle before sampling"
-    )
-    parser.add_argument(
-        "--seed", type=int, default=42, help="Shuffle seed (default: 42)"
-    )
-    parser.add_argument(
-        "--batch-size",
-        type=int,
-        default=16,
-        help="Rows (images) per offline llm.chat batch (default: 16)",
-    )
-    parser.add_argument(
-        "--max-model-len",
-        type=int,
-        default=18000,
-        help="vLLM context length (default: 18000)",
-    )
-    parser.add_argument(
-        "--gpu-memory-utilization",
-        type=float,
-        default=0.85,
-        help="vLLM GPU memory fraction (default: 0.85)",
-    )
-    parser.add_argument(
-        "--dtype",
-        default="bfloat16",
-        help="vLLM dtype (default: bfloat16; use float16 on T4/Turing)",
-    )
-    parser.add_argument(
-        "--model", default=DEFAULT_MODEL, help=f"Model ID (default: {DEFAULT_MODEL})"
-    )
-    parser.add_argument(
-        "--private", action="store_true", help="Make output dataset private"
-    )
-    parser.add_argument(
-        "--config",
-        default=None,
-        help="Config/subset name when pushing (for benchmarking in one repo)",
-    )
-    parser.add_argument(
-        "--create-pr",
-        action="store_true",
-        help="Push as a pull request instead of directly",
-    )
-    parser.add_argument("--hf-token", help="Hugging Face API token (or set HF_TOKEN)")
-    parser.add_argument(
-        "--verbose",
-        action="store_true",
-        help="Log resolved package versions after processing",
-    )
-
-    args = parser.parse_args()
-
-    if args.pdf_column and args.image_column != "image":
-        parser.error("--image-column and --pdf-column are mutually exclusive.")
-
-    main(
-        input_dataset=args.input_dataset,
-        output_dataset=args.output_dataset,
-        task=args.task,
-        table_mode=args.table_mode,
-        image_column=args.image_column,
-        pdf_column=args.pdf_column,
-        output_column=args.output_column,
-        overwrite=args.overwrite,
-        blocks_column=args.blocks_column,
-        page_range=args.page_range,
-        split=args.split,
-        max_samples=args.max_samples,
-        shuffle=args.shuffle,
-        seed=args.seed,
-        batch_size=args.batch_size,
-        max_model_len=args.max_model_len,
-        gpu_memory_utilization=args.gpu_memory_utilization,
-        dtype=args.dtype,
-        model=args.model,
-        private=args.private,
-        config=args.config,
-        create_pr=args.create_pr,
-        hf_token=args.hf_token,
-        verbose=args.verbose,
-    )

tesseract-ocr.py

@@ -1,646 +0,0 @@
-# /// script
-# requires-python = ">=3.10"
-# dependencies = [
-#     "pytesseract",
-#     "datasets>=3.1.0",
-#     "huggingface-hub[hf_transfer]",
-#     "pillow",
-#     "tqdm",
-# ]
-# ///
-"""
-Plain-text OCR with **Tesseract** — the classical CPU OCR engine, as a baseline.
-
-This is the odd one out in this collection: no GPU, no VLM. It runs Google's
-Tesseract (v5, LSTM) over an image dataset and writes the recognised text to a
-column, so you can put a cheap, fast, no-GPU baseline next to the VLM recipes
-(e.g. in a per-collection leaderboard like ocr-bench). Output is plain text, not
-markdown — Tesseract has no notion of tables/formulas/layout-as-markdown.
-
-CPU-ONLY: this recipe deliberately does not require a GPU. Run it on a
-`cpu-basic` / `cpu-upgrade` flavor. `--num-workers` fans OCR out across cores.
-
-SYSTEM DEPENDENCY: the `tesseract` binary is NOT in the default Jobs image. On
-startup this script installs it via apt (`tesseract-ocr`; Jobs containers run as
-root). For non-English languages it also tries to install the matching
-`tesseract-ocr-<lang>` data pack. If your Jobs image blocks apt, bake Tesseract
-into a custom `--image` instead.
-
-HF Jobs (CPU — no GPU needed):
-
-    hf jobs uv run --flavor cpu-upgrade -s HF_TOKEN \\
-        https://huggingface.co/datasets/uv-scripts/ocr/raw/main/tesseract-ocr.py \\
-        input-dataset output-dataset \\
-        --max-samples 100 --shuffle
-
-Language packs (apt codes are 3-letter ISO 639-2, same as Tesseract's `--lang`):
-
-    --lang eng            English (default; ships with the base package)
-    --lang fra            French   (installs tesseract-ocr-fra)
-    --lang eng+fra        multiple languages, '+'-joined
-
-Engine: Tesseract OCR (https://github.com/tesseract-ocr/tesseract), Apache-2.0.
-"""
-
-import argparse
-import io
-import json
-import logging
-import os
-import shutil
-import subprocess
-import sys
-import time
-from concurrent.futures import ThreadPoolExecutor
-from datetime import datetime, timezone
-from typing import Any, Dict, List, Union
-
-from datasets import load_dataset
-from huggingface_hub import DatasetCard, login
-from PIL import Image
-from tqdm import tqdm
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# Stable leaderboard label. Tesseract is not on the Hub, so this is deliberately
-# NOT an org/name Hub id (faking one would produce broken links in dataset cards
-# and a misleading leaderboard row). The precise installed version is recorded
-# separately in inference_info as `tesseract_version`.
-MODEL_ID = "tesseract-5"
-MODEL_NAME = "Tesseract"
-PROJECT_URL = "https://github.com/tesseract-ocr/tesseract"
-
-# Error sentinel written to the output column when a single image fails, so one
-# bad page never sinks the whole run (matches the other recipes in this repo).
-OCR_ERROR = "[OCR ERROR]"
-
-
-def ensure_tesseract_installed(lang: str) -> None:
-    """Install the `tesseract` binary (and language data) if it's missing.
-
-    The default Jobs image has no `tesseract`. Containers run as root, so apt
-    works. The base `tesseract-ocr` package ships English; other languages need
-    their own `tesseract-ocr-<code>` data pack. Base install is fatal on
-    failure (nothing to OCR with); language-pack installs are best-effort (a
-    missing pack surfaces as a clear Tesseract error at OCR time anyway).
-    """
-    requested = [c.strip() for c in lang.split("+") if c.strip()]
-
-    if shutil.which("tesseract") is None:
-        logger.info("`tesseract` not found — installing via apt (Jobs runs as root)...")
-        try:
-            subprocess.run(["apt-get", "update", "-qq"], check=True)
-            subprocess.run(
-                ["apt-get", "install", "-y", "-qq", "tesseract-ocr"], check=True
-            )
-        except Exception as e:
-            logger.error(f"Failed to apt-install tesseract-ocr: {e}")
-            logger.error(
-                "If this Jobs image blocks apt, run with a custom --image that has "
-                "Tesseract preinstalled (apt package `tesseract-ocr`)."
-            )
-            sys.exit(1)
-        if shutil.which("tesseract") is None:
-            logger.error("tesseract still not on PATH after install — aborting.")
-            sys.exit(1)
-        logger.info("Installed tesseract.")
-
-    # Install any missing non-English language data packs (best-effort).
-    try:
-        import pytesseract
-
-        available = set(pytesseract.get_languages(config=""))
-    except Exception:
-        available = set()
-    missing = [c for c in requested if c not in available and c not in ("osd",)]
-    if missing:
-        pkgs = [f"tesseract-ocr-{c}" for c in missing]
-        logger.info(f"Installing language data pack(s): {pkgs}")
-        try:
-            subprocess.run(["apt-get", "update", "-qq"], check=True)
-            subprocess.run(["apt-get", "install", "-y", "-qq", *pkgs], check=True)
-        except Exception as e:
-            logger.warning(
-                f"Could not install language pack(s) {pkgs}: {e}. "
-                f"OCR will fail for those languages if the data is absent."
-            )
-
-
-def detect_tesseract_version() -> str:
-    """Return the installed Tesseract version string (best-effort)."""
-    try:
-        import pytesseract
-
-        return str(pytesseract.get_tesseract_version())
-    except Exception:
-        return "unknown"
-
-
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a
-    ground-truth `text`/`markdown` column) or crashes on push with a
-    duplicate-column error only *after* OCR has run. Catch it up front. With
-    overwrite=True, drop the clashing column(s) here instead (logged).
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error(
-        "Choose a different --output-column, or pass --overwrite to replace them."
-    )
-    sys.exit(1)
-
-
-def to_pil(image: Union[Image.Image, Dict[str, Any], str]) -> Image.Image:
-    """Coerce a datasets image cell to an RGB PIL image.
-
-    Handles the three shapes a HF image column yields: a decoded PIL image, a
-    `{"bytes": ...}` dict, or a file path string.
-    """
-    if isinstance(image, Image.Image):
-        pil_img = image
-    elif isinstance(image, dict) and "bytes" in image:
-        pil_img = Image.open(io.BytesIO(image["bytes"]))
-    elif isinstance(image, str):
-        pil_img = Image.open(image)
-    else:
-        raise ValueError(f"Unsupported image type: {type(image)}")
-    return pil_img.convert("RGB")
-
-
-def ocr_image(
-    image: Union[Image.Image, Dict[str, Any], str],
-    lang: str,
-    config: str,
-) -> str:
-    """Run Tesseract on a single image, returning stripped text (or the error sentinel)."""
-    import pytesseract
-
-    try:
-        pil_img = to_pil(image)
-        return pytesseract.image_to_string(pil_img, lang=lang, config=config).strip()
-    except Exception as e:  # noqa: BLE001 — one bad page shouldn't sink the run
-        logger.error(f"Error OCR'ing image: {e}")
-        return OCR_ERROR
-
-
-def run_ocr(
-    dataset,
-    image_column: str,
-    lang: str,
-    config: str,
-    num_workers: int,
-) -> List[str]:
-    """OCR every row's image, preserving dataset order.
-
-    Tesseract shells out to a subprocess, so ThreadPoolExecutor gives real
-    parallelism (the GIL is released during the call). We process in chunks so
-    only a chunk's worth of decoded images is held in memory at once — the full
-    dataset stays on disk.
-    """
-    n = len(dataset)
-    results: List[str] = []
-    chunk = max(num_workers * 4, 16)
-
-    with tqdm(total=n, desc="OCR", unit="img") as pbar:
-        if num_workers <= 1:
-            for i in range(n):
-                results.append(ocr_image(dataset[i][image_column], lang, config))
-                pbar.update(1)
-        else:
-            with ThreadPoolExecutor(max_workers=num_workers) as pool:
-                for start in range(0, n, chunk):
-                    stop = min(start + chunk, n)
-                    images = [dataset[i][image_column] for i in range(start, stop)]
-                    for text in pool.map(
-                        lambda img: ocr_image(img, lang, config), images
-                    ):
-                        results.append(text)
-                        pbar.update(1)
-    return results
-
-
-def create_dataset_card(
-    source_dataset: str,
-    num_samples: int,
-    processing_time: str,
-    tesseract_version: str,
-    lang: str,
-    psm: int,
-    oem: int,
-    num_workers: int,
-    output_column: str,
-    image_column: str = "image",
-    split: str = "train",
-) -> str:
-    """Create a dataset card documenting the Tesseract OCR run."""
-    return f"""---
-tags:
-- ocr
-- text-recognition
-- tesseract
-- uv-script
-- generated
----
-
-# Document OCR using Tesseract
-
-This dataset contains OCR results from images in [{source_dataset}](https://huggingface.co/datasets/{source_dataset}) using [Tesseract]({PROJECT_URL}), the classical open-source CPU OCR engine — a cheap, no-GPU baseline alongside the VLM OCR recipes.
-
-## Processing Details
-
-- **Source Dataset**: [{source_dataset}](https://huggingface.co/datasets/{source_dataset})
-- **Engine**: [Tesseract]({PROJECT_URL}) `{tesseract_version}`
-- **Language(s)**: `{lang}`
-- **Number of Samples**: {num_samples:,}
-- **Processing Time**: {processing_time}
-- **Processing Date**: {datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC")}
-
-### Configuration
-
-- **Image Column**: `{image_column}`
-- **Output Column**: `{output_column}`
-- **Dataset Split**: `{split}`
-- **Page Segmentation Mode (psm)**: {psm}
-- **OCR Engine Mode (oem)**: {oem}
-- **Workers**: {num_workers}
-
-## Model Information
-
-Tesseract is a classical (non-VLM) OCR engine:
-- Runs on CPU — no GPU required
-- v4+ uses an LSTM-based recognition engine
-- 100+ languages via installable data packs
-- Plain-text output (no markdown / table / formula structure)
-- Apache-2.0 licensed
-
-## Dataset Structure
-
-The dataset contains all original columns plus:
-- `{output_column}`: The recognised text (plain text)
-- `inference_info`: JSON list tracking all OCR models applied to this dataset
-
-## Reproduction
-
-```bash
-uv run https://huggingface.co/datasets/uv-scripts/ocr/raw/main/tesseract-ocr.py \\
-    {source_dataset} \\
-    <output-dataset> \\
-    --image-column {image_column} \\
-    --lang {lang} \\
-    --psm {psm}
-```
-
-Generated with [UV Scripts](https://huggingface.co/uv-scripts)
-"""
-
-
-def main(
-    input_dataset: str,
-    output_dataset: str,
-    image_column: str = "image",
-    lang: str = "eng",
-    psm: int = 3,
-    oem: int = 3,
-    num_workers: int = 0,
-    hf_token: str = None,
-    split: str = "train",
-    max_samples: int = None,
-    private: bool = False,
-    shuffle: bool = False,
-    seed: int = 42,
-    output_column: str = "markdown",
-    overwrite: bool = False,
-    dry_run: bool = False,
-    verbose: bool = False,
-    config: str = None,
-    create_pr: bool = False,
-):
-    """Process images from an HF dataset through Tesseract OCR."""
-
-    start_time = datetime.now(timezone.utc)
-
-    # Fan-out defaults to all cores. When running >1 worker, cap Tesseract's own
-    # OpenMP threads to 1 so N processes don't oversubscribe the CPU (per the
-    # Tesseract FAQ). Must be set before the binary is invoked.
-    if num_workers <= 0:
-        num_workers = os.cpu_count() or 1
-    if num_workers > 1:
-        os.environ.setdefault("OMP_THREAD_LIMIT", "1")
-
-    ensure_tesseract_installed(lang)
-    tesseract_version = detect_tesseract_version()
-    logger.info(
-        f"Using Tesseract {tesseract_version} (lang={lang}, psm={psm}, oem={oem})"
-    )
-    logger.info(f"CPU-only run with {num_workers} worker(s)")
-
-    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")
-    if HF_TOKEN and not dry_run:
-        login(token=HF_TOKEN)
-
-    # Load dataset
-    logger.info(f"Loading dataset: {input_dataset}")
-    dataset = load_dataset(input_dataset, split=split)
-
-    if image_column not in dataset.column_names:
-        raise ValueError(
-            f"Column '{image_column}' not found. Available: {dataset.column_names}"
-        )
-
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
-    if shuffle:
-        logger.info(f"Shuffling dataset with seed {seed}")
-        dataset = dataset.shuffle(seed=seed)
-
-    if max_samples:
-        dataset = dataset.select(range(min(max_samples, len(dataset))))
-        logger.info(f"Limited to {len(dataset)} samples")
-
-    tess_config = f"--psm {psm} --oem {oem}"
-    logger.info(f"Processing {len(dataset)} images -> column '{output_column}'")
-
-    all_outputs = run_ocr(dataset, image_column, lang, tess_config, num_workers)
-
-    n_errors = sum(1 for t in all_outputs if t == OCR_ERROR)
-    if n_errors:
-        logger.warning(f"{n_errors}/{len(all_outputs)} images failed OCR ({OCR_ERROR})")
-
-    processing_duration = datetime.now(timezone.utc) - start_time
-    processing_time_str = f"{processing_duration.total_seconds() / 60:.1f} min"
-
-    logger.info(f"Adding '{output_column}' column to dataset")
-    dataset = dataset.add_column(output_column, all_outputs)
-
-    # Inference info tracking. `model_id` + `column_name` are the fields consumers
-    # (e.g. ocr-bench) read to find the text column and the leaderboard label.
-    inference_entry = {
-        "model_id": MODEL_ID,
-        "model_name": MODEL_NAME,
-        "column_name": output_column,
-        "timestamp": datetime.now(timezone.utc).isoformat(),
-        "engine": "tesseract",
-        "tesseract_version": tesseract_version,
-        "lang": lang,
-        "psm": psm,
-        "oem": oem,
-    }
-
-    if "inference_info" in dataset.column_names:
-        logger.info("Updating existing inference_info column")
-
-        def update_inference_info(example):
-            try:
-                existing_info = (
-                    json.loads(example["inference_info"])
-                    if example["inference_info"]
-                    else []
-                )
-            except (json.JSONDecodeError, TypeError):
-                existing_info = []
-            existing_info.append(inference_entry)
-            return {"inference_info": json.dumps(existing_info)}
-
-        dataset = dataset.map(update_inference_info)
-    else:
-        logger.info("Creating new inference_info column")
-        inference_list = [json.dumps([inference_entry])] * len(dataset)
-        dataset = dataset.add_column("inference_info", inference_list)
-
-    if dry_run:
-        logger.info("--dry-run: skipping push to Hub.")
-        preview = next((t for t in all_outputs if t and t != OCR_ERROR), "")
-        logger.info(f"Sample OCR output (first non-empty, truncated):\n{preview[:500]}")
-        logger.info(
-            f"Done (dry run). {len(dataset)} rows OCR'd in {processing_time_str}."
-        )
-        return
-
-    # Push to hub with retry and XET fallback
-    logger.info(f"Pushing to {output_dataset}")
-    max_retries = 3
-    for attempt in range(1, max_retries + 1):
-        try:
-            if attempt > 1:
-                logger.warning("Disabling XET (fallback to HTTP upload)")
-                os.environ["HF_HUB_DISABLE_XET"] = "1"
-            dataset.push_to_hub(
-                output_dataset,
-                private=private,
-                token=HF_TOKEN,
-                max_shard_size="500MB",
-                **({"config_name": config} if config else {}),
-                create_pr=create_pr,
-                commit_message=f"Add {MODEL_ID} OCR results ({len(dataset)} samples)"
-                + (f" [{config}]" if config else ""),
-            )
-            break
-        except Exception as e:
-            logger.error(f"Upload attempt {attempt}/{max_retries} failed: {e}")
-            if attempt < max_retries:
-                delay = 30 * (2 ** (attempt - 1))
-                logger.info(f"Retrying in {delay}s...")
-                time.sleep(delay)
-            else:
-                logger.error("All upload attempts failed. OCR results are lost.")
-                sys.exit(1)
-
-    # Create and push dataset card
-    logger.info("Creating dataset card")
-    card_content = create_dataset_card(
-        source_dataset=input_dataset,
-        num_samples=len(dataset),
-        processing_time=processing_time_str,
-        tesseract_version=tesseract_version,
-        lang=lang,
-        psm=psm,
-        oem=oem,
-        num_workers=num_workers,
-        output_column=output_column,
-        image_column=image_column,
-        split=split,
-    )
-    card = DatasetCard(card_content)
-    card.push_to_hub(output_dataset, token=HF_TOKEN)
-
-    logger.info("Done! Tesseract OCR processing complete.")
-    logger.info(
-        f"Dataset available at: https://huggingface.co/datasets/{output_dataset}"
-    )
-    logger.info(f"Processing time: {processing_time_str}")
-    if processing_duration.total_seconds() > 0:
-        logger.info(
-            f"Processing speed: {len(dataset) / processing_duration.total_seconds():.2f} images/sec"
-        )
-
-    if verbose:
-        import importlib.metadata
-
-        logger.info("--- Resolved package versions ---")
-        for pkg in ["pytesseract", "datasets", "pyarrow", "pillow"]:
-            try:
-                logger.info(f"  {pkg}=={importlib.metadata.version(pkg)}")
-            except importlib.metadata.PackageNotFoundError:
-                logger.info(f"  {pkg}: not installed")
-        logger.info(f"  tesseract=={tesseract_version}")
-        logger.info("--- End versions ---")
-
-
-if __name__ == "__main__":
-    if len(sys.argv) == 1:
-        print("=" * 70)
-        print("Tesseract OCR — CPU baseline")
-        print("=" * 70)
-        print("\nClassical LSTM OCR engine. No GPU. Plain-text output.")
-        print("\nExamples:")
-        print("\n1. Basic OCR (English):")
-        print("   uv run tesseract-ocr.py input-dataset output-dataset")
-        print("\n2. Another language (installs the data pack on Jobs):")
-        print("   uv run tesseract-ocr.py docs results --lang fra")
-        print("\n3. Test with a small sample, no push:")
-        print("   uv run tesseract-ocr.py large-dataset out --max-samples 5 --dry-run")
-        print("\n4. Running on HF Jobs (CPU flavor):")
-        print("   hf jobs uv run --flavor cpu-upgrade \\")
-        print("     -s HF_TOKEN \\")
-        print(
-            "     https://huggingface.co/datasets/uv-scripts/ocr/raw/main/tesseract-ocr.py \\"
-        )
-        print("       input-dataset output-dataset --max-samples 100 --shuffle")
-        print("\nFor full help: uv run tesseract-ocr.py --help")
-        sys.exit(0)
-
-    parser = argparse.ArgumentParser(
-        description="CPU OCR baseline using Tesseract (classical LSTM engine, plain text)",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog="""
-Examples:
-  uv run tesseract-ocr.py my-docs analyzed-docs
-  uv run tesseract-ocr.py docs results --lang eng+fra --psm 4
-  uv run tesseract-ocr.py large-dataset test --max-samples 50 --shuffle --dry-run
-
-Page segmentation modes (--psm), common ones:
-  3  Fully automatic page segmentation, no OSD (default; good for full pages)
-  4  Assume a single column of text of variable sizes
-  6  Assume a single uniform block of text
-  1  Automatic page segmentation with OSD
-        """,
-    )
-
-    parser.add_argument("input_dataset", help="Input dataset ID from Hugging Face Hub")
-    parser.add_argument("output_dataset", help="Output dataset ID for Hugging Face Hub")
-    parser.add_argument(
-        "--image-column",
-        default="image",
-        help="Column containing images (default: image)",
-    )
-    parser.add_argument(
-        "--lang",
-        default="eng",
-        help="Tesseract language code(s), '+'-joined (default: eng). "
-        "Non-English packs are apt-installed on Jobs.",
-    )
-    parser.add_argument(
-        "--psm",
-        type=int,
-        default=3,
-        help="Page segmentation mode (default: 3, full-page auto)",
-    )
-    parser.add_argument(
-        "--oem",
-        type=int,
-        default=3,
-        help="OCR engine mode (default: 3, based on what's available; 1=LSTM only)",
-    )
-    parser.add_argument(
-        "--num-workers",
-        type=int,
-        default=0,
-        help="Parallel OCR workers (default: 0 = all CPU cores)",
-    )
-    parser.add_argument("--hf-token", help="Hugging Face API token")
-    parser.add_argument(
-        "--split", default="train", help="Dataset split to use (default: train)"
-    )
-    parser.add_argument(
-        "--max-samples",
-        type=int,
-        help="Maximum number of samples to process (for testing)",
-    )
-    parser.add_argument(
-        "--private", action="store_true", help="Make output dataset private"
-    )
-    parser.add_argument(
-        "--config",
-        help="Config/subset name when pushing to Hub (for benchmarking multiple models in one repo)",
-    )
-    parser.add_argument(
-        "--create-pr",
-        action="store_true",
-        help="Create a pull request instead of pushing directly (for parallel benchmarking)",
-    )
-    parser.add_argument(
-        "--shuffle", action="store_true", help="Shuffle dataset before processing"
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=42,
-        help="Random seed for shuffling (default: 42)",
-    )
-    parser.add_argument(
-        "--output-column",
-        default="markdown",
-        help="Column name for output text (default: markdown, for cross-model consistency)",
-    )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
-    parser.add_argument(
-        "--dry-run",
-        action="store_true",
-        help="Run OCR but do NOT push to the Hub (local smoke testing).",
-    )
-    parser.add_argument(
-        "--verbose",
-        action="store_true",
-        help="Log resolved package versions after processing (useful for pinning deps)",
-    )
-
-    args = parser.parse_args()
-
-    main(
-        input_dataset=args.input_dataset,
-        output_dataset=args.output_dataset,
-        image_column=args.image_column,
-        lang=args.lang,
-        psm=args.psm,
-        oem=args.oem,
-        num_workers=args.num_workers,
-        hf_token=args.hf_token,
-        split=args.split,
-        max_samples=args.max_samples,
-        private=args.private,
-        shuffle=args.shuffle,
-        seed=args.seed,
-        output_column=args.output_column,
-        overwrite=args.overwrite,
-        dry_run=args.dry_run,
-        verbose=args.verbose,
-        config=args.config,
-        create_pr=args.create_pr,
-    )

unlimited-ocr-vllm.py

@@ -1,577 +0,0 @@
-# /// script
-# requires-python = ">=3.11"
-# dependencies = [
-#     "datasets>=4.0.0",
-#     "huggingface-hub",
-#     "pillow",
-#     "tqdm",
-#     "toolz",
-# ]
-# ///
-
-"""
-Convert document images to markdown using Baidu Unlimited-OCR with vLLM.
-
-Unlimited-OCR (baidu/Unlimited-OCR, 3.3B, MIT) is a DeepSeek-OCR / DeepSeek-OCR-2 descendant. This
-recipe runs it as an offline vLLM batch job (dataset in -> markdown out), mirroring the proven
-deepseek-ocr-vllm.py pattern: llm.generate() with PIL images and the model's
-NGramPerReqLogitsProcessor to stop coordinate-token loops on long documents.
-
-One image per row -> one markdown. Output is layout-grounded markdown: text spans are tagged
-<|ref|>...<|/ref|> with <|det|>...<|/det|> coordinate boxes (coords normalized 0-1000); tables come
-back as HTML and equations as LaTeX. Pass --strip-grounding to drop the tags and keep clean text;
-add --grounding-column to keep the raw grounded output (with bboxes) in a second column too.
-
-Multi-page / "long-horizon" parsing (the model's headline feature) is not in this single-image batch
-recipe — for multi-page, serve the model and send all pages in one request (see serving-unlimited-ocr.md).
-Multi-page *does* work via vLLM serving: on a clean 2-page doc it returned both pages, <PAGE>-separated.
-But on hard/degraded scans (dense historical pages, newspaper clippings) vLLM multi-page degraded to
-hallucination in our tests, where the model's own SGLang build held up better — so SGLang is the more
-robust multi-page path. (vLLM's upstream PR, vllm-project/vllm#46564, benchmarks single-page only.)
-
-IMPORTANT: Unlimited-OCR's architecture is not in a stable vLLM pip wheel, so this script MUST run on
-Baidu's dedicated vLLM image (vllm and torch come from the image, not the PEP 723 deps):
-
-    hf jobs uv run --flavor l4x1 -s HF_TOKEN \\
-        --image vllm/vllm-openai:unlimited-ocr --python /usr/bin/python3 \\
-        -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages \\
-        https://huggingface.co/datasets/uv-scripts/ocr/raw/main/unlimited-ocr-vllm.py \\
-        your-input-dataset your-output-dataset --max-samples 10
-
-Use the vllm/vllm-openai:unlimited-ocr-cu129 tag on Hopper GPUs (h100/h200).
-
-Model card: https://huggingface.co/baidu/Unlimited-OCR
-vLLM recipe: https://recipes.vllm.ai/baidu/Unlimited-OCR
-"""
-
-import argparse
-import io
-import json
-import logging
-import os
-import re
-import sys
-import time
-from datetime import datetime
-from typing import Any, Dict, List, Optional, Union
-
-import torch
-from datasets import load_dataset
-from huggingface_hub import DatasetCard, login
-from PIL import Image
-from toolz import partition_all
-from tqdm.auto import tqdm
-
-# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc. Greedy OCR doesn't
-# use it; on the dedicated vllm/vllm-openai image it's a harmless no-op.
-os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
-from vllm import LLM, SamplingParams
-from vllm.model_executor.models.unlimited_ocr import NGramPerReqLogitsProcessor
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-MODEL = "baidu/Unlimited-OCR"
-
-# Prompt and no-repeat-ngram knobs straight from the model card / vLLM recipe (single image).
-PROMPT = "<image>document parsing."
-NGRAM_SIZE = 35
-WINDOW_SIZE = 128
-
-# Strip the model's grounding markup to recover clean text:
-# drop the <|det|>...<|/det|> coordinate boxes, then unwrap the <|ref|>...<|/ref|> spans.
-_DET_RE = re.compile(r"<\|det\|>.*?<\|/det\|>", re.DOTALL)
-_REF_RE = re.compile(r"<\|/?ref\|>")
-
-
-def strip_grounding(text: str) -> str:
-    """Remove <|det|> boxes and <|ref|> wrappers, keeping the inner text."""
-    text = _DET_RE.sub("", text)
-    text = _REF_RE.sub("", text)
-    # collapse the blank lines left behind by removed boxes
-    return re.sub(r"\n{3,}", "\n\n", text).strip()
-
-
-def check_cuda_availability():
-    """Check if CUDA is available and exit if not."""
-    if not torch.cuda.is_available():
-        logger.error("CUDA is not available. This script requires a GPU.")
-        sys.exit(1)
-    logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
-
-
-def ensure_output_columns_free(dataset, columns, overwrite=False):
-    """Fail fast if an output column would collide with an existing input column.
-
-    Adding a column that already exists silently overwrites it (e.g. a ground-truth
-    `text`/`markdown` column) or crashes on push with a duplicate-column error only
-    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing
-    column(s) here instead (logged) so the later add_column is clean.
-    """
-    clash = [c for c in columns if c in dataset.column_names]
-    if not clash:
-        return dataset
-    if overwrite:
-        logger.warning(f"--overwrite: replacing existing column(s) {clash}")
-        return dataset.remove_columns(clash)
-    logger.error(
-        f"Output column(s) {clash} already exist in the input dataset "
-        f"(columns: {dataset.column_names})."
-    )
-    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
-    sys.exit(1)
-
-
-def to_pil(image: Union[Image.Image, Dict[str, Any], str]) -> Image.Image:
-    """Convert various dataset image cell formats to an RGB PIL image."""
-    if isinstance(image, Image.Image):
-        return image.convert("RGB")
-    if isinstance(image, dict) and "bytes" in image:
-        return Image.open(io.BytesIO(image["bytes"])).convert("RGB")
-    if isinstance(image, str):
-        return Image.open(image).convert("RGB")
-    raise ValueError(f"Unsupported image type: {type(image)}")
-
-
-def create_dataset_card(
-    source_dataset: str,
-    output_dataset: str,
-    model: str,
-    num_samples: int,
-    processing_time: str,
-    output_column: str,
-    strip_grounding_enabled: bool,
-    split: str,
-) -> str:
-    """Create a dataset card documenting the OCR run."""
-    if strip_grounding_enabled:
-        grounding = "Grounding markup was stripped (`--strip-grounding`); the column holds clean text."
-    else:
-        grounding = (
-            "The column holds the model's raw layout-grounded markdown: text spans tagged "
-            "`<|ref|>...<|/ref|>` with `<|det|>...<|/det|>` coordinate boxes (coords 0-1000). "
-            "Strip them with "
-            "`re.sub(r'<\\|det\\|>.*?<\\|/det\\|>', '', t)` then `re.sub(r'<\\|/?ref\\|>', '', t)`."
-        )
-    return f"""---
-tags:
-- ocr
-- document-processing
-- unlimited-ocr
-- baidu
-- markdown
-- uv-script
-- generated
----
-
-# Document OCR using Unlimited-OCR
-
-This dataset contains OCR results for [{source_dataset}](https://huggingface.co/datasets/{source_dataset})
-produced by [{model}](https://huggingface.co/{model}) with vLLM.
-
-## Processing Details
-
-- **Source Dataset**: [{source_dataset}](https://huggingface.co/datasets/{source_dataset})
-- **Model**: [{model}](https://huggingface.co/{model})
-- **Number of Samples**: {num_samples:,}
-- **Processing Time**: {processing_time}
-- **Processing Date**: {datetime.now().strftime("%Y-%m-%d %H:%M UTC")}
-- **Output Column**: `{output_column}`
-- **Split**: `{split}`
-
-## Output
-
-{grounding}
-
-Tables are returned as HTML and equations as LaTeX.
-
-## Usage
-
-```python
-from datasets import load_dataset
-
-ds = load_dataset("{output_dataset}", split="{split}")
-print(ds[0]["{output_column}"])
-```
-
-## Reproduction
-
-Generated with the [uv-scripts/ocr](https://huggingface.co/datasets/uv-scripts/ocr) Unlimited-OCR
-vLLM recipe. Unlimited-OCR needs Baidu's dedicated vLLM image:
-
-```bash
-hf jobs uv run --flavor l4x1 -s HF_TOKEN \\
-    --image vllm/vllm-openai:unlimited-ocr --python /usr/bin/python3 \\
-    -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages \\
-    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/unlimited-ocr-vllm.py \\
-    {source_dataset} <output-dataset>
-```
-
-Generated with [UV Scripts](https://huggingface.co/uv-scripts)
-"""
-
-
-def main(
-    input_dataset: str,
-    output_dataset: str,
-    model: str = MODEL,
-    image_column: str = "image",
-    output_column: str = "markdown",
-    overwrite: bool = False,
-    grounding_column: Optional[str] = None,
-    batch_size: int = 8,
-    max_model_len: int = 32768,
-    max_tokens: int = 8192,
-    gpu_memory_utilization: float = 0.8,
-    strip_grounding_enabled: bool = False,
-    hf_token: Optional[str] = None,
-    split: str = "train",
-    max_samples: Optional[int] = None,
-    private: bool = False,
-    shuffle: bool = False,
-    seed: int = 42,
-    config: Optional[str] = None,
-    create_pr: bool = False,
-    verbose: bool = False,
-):
-    """Process images from an HF dataset through Unlimited-OCR with vLLM."""
-    if grounding_column and grounding_column == output_column:
-        raise ValueError("--grounding-column must differ from --output-column")
-    check_cuda_availability()
-    start_time = datetime.now()
-
-    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")
-    if HF_TOKEN:
-        login(token=HF_TOKEN)
-
-    logger.info(f"Loading dataset: {input_dataset}")
-    dataset = load_dataset(input_dataset, split=split)
-    if image_column not in dataset.column_names:
-        raise ValueError(
-            f"Column '{image_column}' not found. Available: {dataset.column_names}"
-        )
-
-    # Fail fast if the output column would collide with an existing input column
-    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
-
-    if shuffle:
-        logger.info(f"Shuffling dataset with seed {seed}")
-        dataset = dataset.shuffle(seed=seed)
-    if max_samples:
-        dataset = dataset.select(range(min(max_samples, len(dataset))))
-        logger.info(f"Limited to {len(dataset)} samples")
-
-    logger.info(f"Initializing vLLM with model: {model}")
-    logger.info("This may take a few minutes on first run...")
-
-    llm = LLM(
-        model=model,
-        trust_remote_code=True,
-        max_model_len=max_model_len,
-        gpu_memory_utilization=gpu_memory_utilization,
-        enable_prefix_caching=False,
-        mm_processor_cache_gb=0,
-        limit_mm_per_prompt={"image": 1},
-        logits_processors=[NGramPerReqLogitsProcessor],
-    )
-
-    sampling_params = SamplingParams(
-        temperature=0.0,
-        max_tokens=max_tokens,
-        skip_special_tokens=False,
-        extra_args=dict(ngram_size=NGRAM_SIZE, window_size=WINDOW_SIZE),
-    )
-
-    logger.info(f"Processing {len(dataset)} images in batches of {batch_size}")
-    all_outputs: List[str] = []
-    all_grounded: List[
-        str
-    ] = []  # raw grounded text, only when --grounding-column is set
-    for batch_indices in tqdm(
-        partition_all(batch_size, range(len(dataset))),
-        total=(len(dataset) + batch_size - 1) // batch_size,
-        desc="Unlimited-OCR",
-    ):
-        batch_indices = list(batch_indices)
-        try:
-            model_inputs = [
-                {
-                    "prompt": PROMPT,
-                    "multi_modal_data": {"image": to_pil(dataset[i][image_column])},
-                }
-                for i in batch_indices
-            ]
-            outputs = llm.generate(model_inputs, sampling_params)
-            for output in outputs:
-                raw = output.outputs[0].text.strip()
-                all_outputs.append(
-                    strip_grounding(raw) if strip_grounding_enabled else raw
-                )
-                if grounding_column:
-                    all_grounded.append(raw)
-        except Exception as e:
-            logger.error(f"Error processing batch: {e}")
-            all_outputs.extend(["[OCR FAILED]"] * len(batch_indices))
-            if grounding_column:
-                all_grounded.extend(["[OCR FAILED]"] * len(batch_indices))
-
-    processing_time_str = (
-        f"{(datetime.now() - start_time).total_seconds() / 60:.1f} min"
-    )
-
-    logger.info(f"Adding '{output_column}' column to dataset")
-    if output_column in dataset.column_names:
-        logger.warning(f"Column '{output_column}' already exists, replacing it")
-        dataset = dataset.remove_columns([output_column])
-    dataset = dataset.add_column(output_column, all_outputs)
-    if grounding_column:
-        logger.info(f"Adding '{grounding_column}' column (raw grounded output)")
-        if grounding_column in dataset.column_names:
-            logger.warning(f"Column '{grounding_column}' already exists, replacing it")
-            dataset = dataset.remove_columns([grounding_column])
-        dataset = dataset.add_column(grounding_column, all_grounded)
-
-    # inference_info: append-only log so several models can write into one dataset and be compared.
-    inference_entry = {
-        "model_id": model,
-        "model_name": "Unlimited-OCR",
-        "column_name": output_column,
-        "timestamp": datetime.now().isoformat(),
-        "batch_size": batch_size,
-        "max_tokens": max_tokens,
-        "max_model_len": max_model_len,
-        "gpu_memory_utilization": gpu_memory_utilization,
-        "strip_grounding": strip_grounding_enabled,
-        "grounding_column": grounding_column,
-        "script": "unlimited-ocr-vllm.py",
-        "script_url": "https://huggingface.co/datasets/uv-scripts/ocr/raw/main/unlimited-ocr-vllm.py",
-    }
-    if "inference_info" in dataset.column_names:
-        logger.info("Updating existing inference_info column")
-
-        def update_inference_info(example):
-            try:
-                existing = (
-                    json.loads(example["inference_info"])
-                    if example["inference_info"]
-                    else []
-                )
-            except (json.JSONDecodeError, TypeError):
-                existing = []
-            existing.append(inference_entry)
-            return {"inference_info": json.dumps(existing)}
-
-        dataset = dataset.map(update_inference_info)
-    else:
-        logger.info("Creating new inference_info column")
-        dataset = dataset.add_column(
-            "inference_info", [json.dumps([inference_entry])] * len(dataset)
-        )
-
-    logger.info(f"Pushing to {output_dataset}")
-    max_retries = 3
-    for attempt in range(1, max_retries + 1):
-        try:
-            if attempt > 1:
-                logger.warning("Disabling XET (fallback to HTTP upload)")
-                os.environ["HF_HUB_DISABLE_XET"] = "1"
-            dataset.push_to_hub(
-                output_dataset,
-                private=private,
-                token=HF_TOKEN,
-                max_shard_size="500MB",
-                **({"config_name": config} if config else {}),
-                create_pr=create_pr,
-                commit_message=f"Add {model} OCR results ({len(dataset)} samples)"
-                + (f" [{config}]" if config else ""),
-            )
-            break
-        except Exception as e:
-            logger.error(f"Upload attempt {attempt}/{max_retries} failed: {e}")
-            if attempt < max_retries:
-                delay = 30 * (2 ** (attempt - 1))
-                logger.info(f"Retrying in {delay}s...")
-                time.sleep(delay)
-            else:
-                logger.error("All upload attempts failed. OCR results are lost.")
-                sys.exit(1)
-
-    logger.info("Creating dataset card...")
-    card = DatasetCard(
-        create_dataset_card(
-            source_dataset=input_dataset,
-            output_dataset=output_dataset,
-            model=model,
-            num_samples=len(dataset),
-            processing_time=processing_time_str,
-            output_column=output_column,
-            strip_grounding_enabled=strip_grounding_enabled,
-            split=split,
-        )
-    )
-    card.push_to_hub(output_dataset, token=HF_TOKEN)
-
-    logger.info("✅ OCR conversion complete!")
-    logger.info(f"Dataset: https://huggingface.co/datasets/{output_dataset}")
-    logger.info(f"Processing time: {processing_time_str}")
-
-    if verbose:
-        import importlib.metadata
-
-        logger.info("--- Resolved package versions ---")
-        for pkg in ["vllm", "transformers", "torch", "datasets", "pillow"]:
-            try:
-                logger.info(f"  {pkg}=={importlib.metadata.version(pkg)}")
-            except importlib.metadata.PackageNotFoundError:
-                logger.info(f"  {pkg}: not installed")
-        logger.info("--- End versions ---")
-
-
-if __name__ == "__main__":
-    if len(sys.argv) == 1:
-        print("=" * 80)
-        print("Unlimited-OCR to Markdown Converter (vLLM)")
-        print("=" * 80)
-        print("\nBaidu Unlimited-OCR (3.3B, MIT) — one image per row -> markdown.")
-        print("\nMUST run on the dedicated image: vllm/vllm-openai:unlimited-ocr")
-        print("(use the -cu129 tag on Hopper GPUs).")
-        print("\nExample:")
-        print("   hf jobs uv run --flavor l4x1 -s HF_TOKEN \\")
-        print(
-            "     --image vllm/vllm-openai:unlimited-ocr --python /usr/bin/python3 \\"
-        )
-        print("     -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages \\")
-        print("     unlimited-ocr-vllm.py my-images my-markdown --max-samples 10")
-        print(
-            "\nMulti-page documents: serve the model instead (see serving-unlimited-ocr.md)."
-        )
-        print("\nFor full help, run with --help")
-        sys.exit(0)
-
-    parser = argparse.ArgumentParser(
-        description="OCR images to markdown using Unlimited-OCR (vLLM)",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog="""
-Examples:
-  # Basic usage
-  uv run unlimited-ocr-vllm.py my-images ocr-results
-
-  # Clean text (strip grounding tags)
-  uv run unlimited-ocr-vllm.py my-images ocr-results --strip-grounding
-
-  # On HF Jobs (dedicated image required)
-  hf jobs uv run --flavor l4x1 -s HF_TOKEN \\
-      --image vllm/vllm-openai:unlimited-ocr --python /usr/bin/python3 \\
-      -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages \\
-      https://huggingface.co/datasets/uv-scripts/ocr/raw/main/unlimited-ocr-vllm.py \\
-      my-dataset my-output --max-samples 10
-        """,
-    )
-    parser.add_argument("input_dataset", help="Input dataset ID from Hugging Face Hub")
-    parser.add_argument("output_dataset", help="Output dataset ID for Hugging Face Hub")
-    parser.add_argument(
-        "--model",
-        default=MODEL,
-        help=f"Model to use (default: {MODEL}). Override only for a same-architecture mirror.",
-    )
-    parser.add_argument(
-        "--image-column", default="image", help="Column with images (default: image)"
-    )
-    parser.add_argument(
-        "--output-column",
-        default="markdown",
-        help="Output column name (default: markdown)",
-    )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Replace the output column if it already exists in the input dataset "
-        "(default: error out to avoid clobbering an existing column).",
-    )
-    parser.add_argument(
-        "--strip-grounding",
-        action="store_true",
-        help="Drop <|det|>/<|ref|> grounding tags from the output column, keeping clean text",
-    )
-    parser.add_argument(
-        "--grounding-column",
-        help="Also store the RAW grounded output (boxes + tags) in this extra column "
-        "(pair with --strip-grounding to keep clean text AND the layout/bboxes)",
-    )
-    parser.add_argument(
-        "--batch-size", type=int, default=8, help="Images per batch (default: 8)"
-    )
-    parser.add_argument(
-        "--max-model-len",
-        type=int,
-        default=32768,
-        help="Max context length (default: 32768)",
-    )
-    parser.add_argument(
-        "--max-tokens",
-        type=int,
-        default=8192,
-        help="Max output tokens (default: 8192)",
-    )
-    parser.add_argument(
-        "--gpu-memory-utilization",
-        type=float,
-        default=0.8,
-        help="GPU memory fraction (default: 0.8)",
-    )
-    parser.add_argument("--hf-token", help="Hugging Face API token")
-    parser.add_argument(
-        "--split", default="train", help="Dataset split (default: train)"
-    )
-    parser.add_argument(
-        "--max-samples", type=int, help="Max samples to process (for testing)"
-    )
-    parser.add_argument(
-        "--private", action="store_true", help="Make output dataset private"
-    )
-    parser.add_argument(
-        "--shuffle", action="store_true", help="Shuffle before processing"
-    )
-    parser.add_argument(
-        "--seed", type=int, default=42, help="Shuffle seed (default: 42)"
-    )
-    parser.add_argument(
-        "--config",
-        help="Config/subset name when pushing (for benchmarking multiple models)",
-    )
-    parser.add_argument(
-        "--create-pr",
-        action="store_true",
-        help="Push as a PR instead of a direct commit",
-    )
-    parser.add_argument(
-        "--verbose",
-        action="store_true",
-        help="Log resolved package versions after the run",
-    )
-
-    args = parser.parse_args()
-
-    main(
-        input_dataset=args.input_dataset,
-        output_dataset=args.output_dataset,
-        model=args.model,
-        image_column=args.image_column,
-        output_column=args.output_column,
-        overwrite=args.overwrite,
-        grounding_column=args.grounding_column,
-        batch_size=args.batch_size,
-        max_model_len=args.max_model_len,
-        max_tokens=args.max_tokens,
-        gpu_memory_utilization=args.gpu_memory_utilization,
-        strip_grounding_enabled=args.strip_grounding,
-        hf_token=args.hf_token,
-        split=args.split,
-        max_samples=args.max_samples,
-        private=args.private,
-        shuffle=args.shuffle,
-        seed=args.seed,
-        config=args.config,
-        create_pr=args.create_pr,
-        verbose=args.verbose,
-    )