doc-split-v1 / README.md
hung-k-nguyen's picture
Release snapshot
dccadc6
|
Raw
History Blame Contribute Delete
6.3 kB
---
license: apache-2.0
pipeline_tag: image-text-to-text
language: [multilingual]
tags:
- page-stream-segmentation
- document-boundary-detection
- document-ai
- document-splitting
- open-weights
datasets:
- nutrientdocs/doc-split-benchmark
---
# doc-split-v1 — open-weight
**Where does one document end and the next begin?** An open-weight page-stream-segmentation model you can
download and run: it splits a stream of pages (a scanned batch / merged PDF) back into its constituent
documents.
The lightweight, open sibling of the commercial flagship
[`doc-split-v2`](https://huggingface.co/nutrientdocs/doc-split-v2) — compact, ~4.5× faster, near-flagship
accuracy on our data and multilingual out of the box. Shipped as **ONNX** — runs with `onnxruntime`, no
framework or modelling code to install.
- 🎯 **Try it:** [doc-split-demo](https://huggingface.co/spaces/nutrientdocs/doc-split-demo?model=v1)
- 🏆 **Leaderboard:** [doc-split-leaderboard](https://huggingface.co/spaces/nutrientdocs/doc-split-leaderboard)
- 📊 **Benchmark:** [doc-split-benchmark](https://huggingface.co/datasets/nutrientdocs/doc-split-benchmark)
- 🔒 **Higher accuracy?** [doc-split-v2](https://huggingface.co/nutrientdocs/doc-split-v2) (commercial)
## Results — boundary F1 (κ)
Per-page boundary detection, page 0 forced. **This model** vs the private doc-split-v2, the strongest cloud VLM,
and prior work.
| Cut | **doc-split-v1** | doc-split-v2 | best cloud VLM | OpenPSS specialist |
|---|---|---|---|---|
| OpenPSS-short (sparse) | **0.585** (.53) | 0.619 | 0.598 (gemini-flash) | 0.76 |
| OpenPSS-long | **0.859** (.82) | 0.886 | 0.244 (gemini-flash) | 0.83 |
| our-200 (synthetic) | **0.936** (.78) | 0.934 | 0.942 (gpt-sol) | — |
| TABME++ test | **0.704** (.56) | 0.901 | — | — |
| Tobacco800 test | **0.820** (.60) | 0.957 | — | — |
| val (real-doc) | **0.918** (.86) | 0.908 | — | — |
Beats every evaluated cloud VLM on OpenPSS-**long** (0.859 vs 0.244) at a fraction of the cost, and holds up
on our data. TABME++/Tobacco800 are zero-shot for this model (in-domain for doc-split-v2).
## What's in this repo
Runs entirely under `onnxruntime` — nothing else to install.
- `image_model.onnx`, `text_model.onnx` — the image and text towers (per-page embeddings).
- `head.onnx` — the boundary head (per-page boundary score).
- `crf.json` — smoothing parameters for the per-page confidence.
- `tokenizer.json` (+ config) — the bundled text tokenizer.
## Usage (ONNX)
```python
# pip install onnxruntime transformers numpy huggingface_hub
import numpy as np, onnxruntime as ort, json
from transformers import AutoTokenizer
from huggingface_hub import snapshot_download
d = snapshot_download("nutrientdocs/doc-split-v1")
img = ort.InferenceSession(f"{d}/image_model.onnx", providers=["CPUExecutionProvider"])
text = ort.InferenceSession(f"{d}/text_model.onnx", providers=["CPUExecutionProvider"])
head = ort.InferenceSession(f"{d}/head.onnx", providers=["CPUExecutionProvider"])
tok = AutoTokenizer.from_pretrained(d); crf = json.load(open(f"{d}/crf.json"))
def _lse(x, ax):
m = x.max(ax, keepdims=True); return (m + np.log(np.exp(x - m).sum(ax, keepdims=True))).squeeze(ax)
def marginals(bl, crf): # per-page confidence via forward-backward over a 2-tag chain
T = np.asarray(crf["trans"]); s = np.asarray(crf["start"]); e_ = np.asarray(crf["end"])
N = len(bl); e = np.stack([np.zeros(N), bl], 1); a = np.zeros((N, 2)); a[0] = s + e[0]
for t in range(1, N): a[t] = _lse(a[t-1][:, None] + T, 0) + e[t]
b = np.zeros((N, 2)); b[N-1] = e_
for t in range(N-2, -1, -1): b[t] = _lse(T + (e[t+1] + b[t+1])[None, :], 1)
m = a + b; m -= m.max(1, keepdims=True); p = np.exp(m); return (p / p.sum(1, keepdims=True))[:, 1]
def split(pages, tau=0.5): # pages: list of (PIL image, ocr_text or "")
arr = np.stack([(np.asarray(im.convert("RGB").resize((512, 512)), np.float32)/255 - .5)/.5
for im, _ in pages]).transpose(0, 3, 1, 2).astype(np.float32)
vi = img.run(["image_embed"], {"pixel_values": arr})[0]
b = tok(["query: "+(t or " ") for _, t in pages], padding=True, truncation=True,
max_length=512, return_tensors="np")
vt = text.run(["text_embed"], {"input_ids": b["input_ids"].astype(np.int64),
"attention_mask": b["attention_mask"].astype(np.int64)})[0]
g = np.array([1. if (t and t.strip()) else 0. for _, t in pages], np.float32); N = len(pages)
vt = vt * g[:, None] # OCR gate: text ignored on pages with no text layer
bl = head.run(["boundary_logit"], {"v_img": vi[None], "v_txt": vt[None],
"gate": g[None], "mask": np.ones((1, N), np.float32)})[0][0]
bl[0] = 30.0 # force page 0 to start a document
conf = marginals(bl, crf) # per-page confidence in [0,1]
return [1 if (i == 0 or conf[i] >= tau) else 0 for i in range(N)] # 1 = this page starts a new document
```
## Intended use & limits
**Use it for:** splitting merged/batch-scanned PDFs into documents; routing; pre-processing for
classification/extraction. **Limits:** boundary detection only (does not classify document *type*); the
sparse low-boundary regime (OpenPSS-short) is hardest; OCR text helps on text-heavy pages.
## License
Apache-2.0.
## Calibrated confidence
The raw boundary score is over-confident (a raw 0.85 is really ~63% likely a true boundary). We ship a
**beta calibration** (fit on held-out data) so the reported confidence is honest and usable as a threshold:
```
p_calibrated = sigmoid(a·ln(p) + b·ln(1-p) + c), (a, b, c) = (0.516, -0.402, -0.155)
```
ECE 0.044 → 0.012. The demo applies this and lets you set a minimum-confidence threshold on the calibrated value.
## About the author
<a href="https://nutrient.io/">
<img src="https://avatars2.githubusercontent.com/u/1527679?v=3&s=200" height="80" />
</a>
This project is maintained and funded by [Nutrient](https://nutrient.io/) - The deterministic document infrastructure enterprises run their highest-stakes workflows on: replayable output, clear exceptions, and full audit trails on the messy, regulated documents where AI alone breaks.