"""Doc-Split demo — upload a merged PDF, watch it split into its constituent documents.
Each model predicts, per page, whether it starts a new document. Text comes from the PDF's embedded text
layer (free); scanned pages fall back to vision-only (OCR gate). Every model runs as portable ONNX graphs
under onnxruntime, with the per-page confidence smoothing done in numpy. Commercial weights load at runtime
from a private repo via the HF_TOKEN secret (server-side only, never downloadable).
ZeroGPU: models load lazily on first call and are cached across calls."""
import os, io, base64, json
import numpy as np
import gradio as gr
import spaces
import onnxruntime as ort
import fitz # pymupdf
from PIL import Image
from transformers import AutoTokenizer
from huggingface_hub import snapshot_download
TOKEN = os.environ.get("HF_TOKEN")
_PROV = ["CUDAExecutionProvider", "CPUExecutionProvider"] # onnxruntime-gpu if present, else CPU fallback
MODELS = {
"doc-split-v2 (flagship · commercial)": "nutrientdocs/doc-split-v2-private",
"doc-split-v1 (open-weight)": "nutrientdocs/doc-split-v1",
}
_CACHE = {}
def _lse(x, axis): # numerically-stable log-sum-exp along one axis (keeps shape reduced)
m = x.max(axis, keepdims=True)
return (m + np.log(np.exp(x - m).sum(axis, keepdims=True))).squeeze(axis)
def _crf_marginals(bl, crf):
"""Per-page P(boundary) via forward-backward over a 2-tag linear chain. bl: raw boundary logits [N].
Emission per page t is [0, bl[t]] (tag 0 = interior, tag 1 = boundary). Public CRF math, no arch."""
trans = np.asarray(crf["trans"], np.float64) # [2,2] i->j
start = np.asarray(crf["start"], np.float64); end = np.asarray(crf["end"], np.float64)
N = len(bl); e = np.stack([np.zeros(N), np.asarray(bl, np.float64)], 1) # [N,2]
a = np.zeros((N, 2)); a[0] = start + e[0]
for t in range(1, N):
a[t] = _lse(a[t - 1][:, None] + trans, 0) + e[t]
bta = np.zeros((N, 2)); bta[N - 1] = end
for t in range(N - 2, -1, -1):
bta[t] = _lse(trans + (e[t + 1] + bta[t + 1])[None, :], 1)
m = a + bta; m = m - m.max(1, keepdims=True); p = np.exp(m)
return (p / p.sum(1, keepdims=True))[:, 1] # posterior P(tag=boundary) per page
# Beta calibration (fit on our-domain val): smooth, monotonic map raw CRF marginal -> honest P(boundary).
# The raw model is over-confident (e.g. flagship raw 0.99 is really ~0.86 likely a boundary); this corrects it.
# p_cal = sigmoid(a*ln(p) + b*ln(1-p) + c). Fit values are hard-coded per model.
_BETA = {
"v2": (0.4898, -0.5100, -0.5071), # flagship; ECE 0.045 -> 0.013
"v1": (0.5156, -0.4023, -0.1545), # open; ECE 0.044 -> 0.012
}
def _calibrate(p, model_name):
a, b, c = _BETA["v1" if "v1" in model_name.lower() else "v2"]
p = min(max(float(p), 1e-6), 1 - 1e-6)
return float(1.0 / (1.0 + np.exp(-(a * np.log(p) + b * np.log(1 - p) + c))))
def _load(name):
if name in _CACHE:
return _CACHE[name]
d = snapshot_download(MODELS[name], token=TOKEN, # commercial weights via HF_TOKEN; v1 is open
allow_patterns=["*.onnx", "*.onnx.data", "crf.json", "tokenizer*", "special_tokens*"])
obj = dict(tok=AutoTokenizer.from_pretrained(d), crf=json.load(open(os.path.join(d, "crf.json"))),
img=ort.InferenceSession(os.path.join(d, "image_model.onnx"), providers=_PROV),
txt=ort.InferenceSession(os.path.join(d, "text_model.onnx"), providers=_PROV),
head=ort.InferenceSession(os.path.join(d, "head.onnx"), providers=_PROV))
_CACHE[name] = obj
return obj
def _pdf_pages(path, max_pages=40):
doc = fitz.open(path); out = []
for i, pg in enumerate(doc):
if i >= max_pages:
break
pix = pg.get_pixmap(matrix=fitz.Matrix(150 / 72, 150 / 72))
im = Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
out.append((im, pg.get_text("text") or ""))
doc.close()
return out
def _thumb(im, w=150):
t = im.copy(); t.thumbnail((w, w * 2)); b = io.BytesIO(); t.save(b, "PNG")
return "data:image/png;base64," + base64.b64encode(b.getvalue()).decode()
def _thumbstrip(pages, label):
h = [f"{label}
"
""]
for i, (im, _) in enumerate(pages):
h.append(f"
})
"
f"
p{i+1}
")
h.append("
Upload a PDF (a few concatenated documents) and press Split.
"
obj = _load(model_name)
pages = _pdf_pages(pdf)
if not pages:
return "No pages found.
"
N = len(pages)
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) # [N,3,512,512]
prompts = ["query: " + (t or " ") for _, t in pages]
gate = np.array([1. if (t and t.strip()) else 0. for _, t in pages], np.float32)
vi = obj["img"].run(["image_embed"], {"pixel_values": arr})[0] # image tower -> per-page embedding
b = obj["tok"](prompts, padding=True, truncation=True, max_length=512, return_tensors="np")
vt = obj["txt"].run(["text_embed"], {"input_ids": b["input_ids"].astype(np.int64),
"attention_mask": b["attention_mask"].astype(np.int64)})[0]
vi = vi.astype(np.float32)[None] # [1,N,d_img]
vt = (vt.astype(np.float32) * gate[:, None])[None] # [1,N,d_txt], OCR-gated
bl = obj["head"].run(["boundary_logit"], {"v_img": vi, "v_txt": vt, # ONNX boundary head
"gate": gate[None], "mask": np.ones((1, N), np.float32)})[0][0]
bl[0] = 30.0 # keep page 0 forced
raw = _crf_marginals(bl, obj["crf"]).tolist() # per-page P(boundary), uncalibrated
conf = [_calibrate(c, model_name) for c in raw] # honest, calibrated confidence (beta)
# start a new document only where the CALIBRATED confidence >= the chosen threshold.
tau = float(min_conf) / 100.0
pred = [1 if (i == 0 or conf[i] >= tau) else 0 for i in range(len(conf))]
# group pages into documents at each boundary
docs, cur = [], []
for i, p in enumerate(pred):
if p and cur:
docs.append(cur); cur = []
cur.append(i)
if cur:
docs.append(cur)
txt_pages = sum(1 for _, t in pages if t and t.strip())
html = [f"Split into {len(docs)} "
f"document(s) across {N} pages · {model_name} · text layer on {txt_pages}/{N} pages
"]
for k, grp in enumerate(docs, 1):
c = conf[grp[0]]
split = "forced start (p1)" if grp[0] == 0 else f"split confidence {c:.0%}"
dot = "#0f7a58" if c >= .8 else ("#9a6a12" if c >= .5 else "#a83a3a")
html.append(f""
f"
Document {k} · pages "
f"{grp[0]+1}–{grp[-1]+1} ({len(grp)}p) · "
f"● {split}
"
f"
")
for i in grp:
html.append(f"
"
f"
})
"
f"
p{i+1}
")
html.append("