document-classification-demo — open-vocab, zero-shot Gradio demo

.gitattributes

@@ -0,0 +1,38 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/financial-reports_2.png filter=lfs diff=lfs merge=lfs -text
+examples/manuals_1.png filter=lfs diff=lfs merge=lfs -text
+examples/scientific-articles_2.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,33 @@
+---
+title: Document Classification Demo
+emoji: 🖼️
+colorFrom: blue
+colorTo: green
+sdk: gradio
+sdk_version: 6.20.0
+app_file: app.py
+pinned: false
+license: other
+short_description: Classify any document image against labels you choose
+---
+
+# Nutrient document classification demo
+
+Open-vocabulary, zero-shot: add candidate classes (a **label** + optional **description** per row), upload a
+document page, and the model ranks them — no fixed class list, no per-class training. Two models, one tab each:
+
+- **v2 flagship** — the commercial model (best accuracy).
+- **v1 (open-weight)** — a downloadable open-vocab classifier.
+
+→ [model](https://huggingface.co/nutrientdocs/document-classification-v2) ·
+[leaderboard](https://huggingface.co/spaces/nutrientdocs/document-classification-leaderboard) ·
+[benchmark](https://huggingface.co/datasets/nutrientdocs/document-classification-benchmark)
+
+## Secret required
+
+The **v2** model is commercial and loaded from a **private** repo, so this Space needs an `HF_TOKEN`
+**secret** (Settings → Variables and secrets) with read access to it. The weights are used server-side only
+and are never downloadable. Without the secret, the demo marks the v2 tab _unavailable_ and the open **v1**
+model still runs. Inference runs on **ZeroGPU**.
+
+Sample document images for quick testing are in `examples/`.

app.py

@@ -0,0 +1,235 @@
+"""Nutrient Document Classification — open-vocab, zero-shot demo (public-ready).
+
+Censored like the sibling grounding demo: NO proprietary weights and NO custom architecture code ship in this
+Space. Everything trained runs as ONNX loaded at RUNTIME:
+  * v2 flagship (commercial) — the full model runs as ONNX fetched from the PRIVATE repo via the
+    HF_TOKEN secret (server-side only, never committed here). Marked unavailable if the secret is absent.
+  * v1 (open-weight)         — ONNX from the public v1 repo.
+The preprocessor + tokenizer are bundled alongside the ONNX in the source repos, so no model identifiers
+appear in this app. Runs on ZeroGPU; ONNX inference on CPU. Results are a ranked Markdown table.
+"""
+import glob as _glob
+import os
+
+import gradio as gr
+import numpy as np
+import spaces
+from huggingface_hub import hf_hub_download
+from transformers import AutoImageProcessor, AutoTokenizer
+
+V2_PRIV = os.environ.get("V2_REPO", "nutrientdocs/document-classification-v2-private")
+V1_REPO = os.environ.get("V1_REPO", "nutrientdocs/document-classification-v1")
+TOKEN = os.environ.get("HF_TOKEN")
+
+_V2, _V1 = {}, {}   # lazy caches: v2 (private) ONNX + processors / v1 (public) ONNX
+
+DEFAULT_ROWS = [
+    ["invoice", "an itemized bill listing goods or services and a total due"],
+    ["letter", "correspondence with a salutation, body, and signature"],
+    ["memo", "an internal memorandum headed To, From, Date, Subject"],
+    ["form", "a structured template with labeled fields to fill in"],
+    ["scientific article", "a research paper with an abstract, methods, and references"],
+    ["resume", "a summary of education, work experience, and skills"],
+]
+
+
+def _rows_to_text(rows):
+    return "\n".join(f"{lab} | {desc}" if desc else lab for lab, desc in rows)
+
+
+DEFAULT_TEXT = _rows_to_text(DEFAULT_ROWS)
+
+# Clickable examples — each loads the page image AND its candidate-label table. Only pages the model
+# classifies correctly are included (verified on the benchmark fixtures).
+_DESC = {
+    "financial reports": "a corporate financial statement with balance sheets, income tables, and figures",
+    "scientific articles": "a research paper with an abstract, methods, results, and references",
+    "laws and regulations": "statutory legal text with numbered sections and articles",
+    "government tenders": "a public procurement notice or call for bids",
+    "manuals": "a product or technical manual with instructions and labeled diagrams",
+    "patents": "a patent document with claims, drawings, and an application number",
+    "invoice": "an itemized bill listing goods or services, quantities, prices, and a total due",
+    "letter": "correspondence with a salutation, body, and signature",
+    "presentation": "a slide from a slideshow presentation", "receipt": "a store receipt with a total",
+    "bar chart": "a chart comparing values with rectangular bars",
+    "line chart": "a chart showing a trend as a line over an axis",
+    "pie chart": "a circular chart divided into proportional slices",
+    "qr code": "a square two-dimensional matrix barcode", "bar code": "a striped one-dimensional barcode",
+    "signature": "a handwritten signature mark", "screenshot": "a screenshot of a software user interface",
+    "table": "a grid of rows and columns of data",
+}
+_DOC = ["financial reports", "scientific articles", "laws and regulations", "government tenders", "manuals",
+        "patents", "invoice", "letter", "presentation", "receipt"]
+_ELEM = ["bar chart", "line chart", "pie chart", "qr code", "bar code", "signature", "screenshot", "table",
+         "invoice", "financial reports"]
+_ELEM_GT = {"bar chart", "pie chart", "line chart", "signature", "screenshot", "qr code", "bar code", "table"}
+
+
+def _example_set():
+    out = []
+    for p in sorted(_glob.glob(os.path.join(os.path.dirname(__file__), "examples", "*.png"))):
+        gt = os.path.basename(p).rsplit("_", 1)[0].replace("-", " ")
+        group = _ELEM if gt in _ELEM_GT else _DOC
+        if gt not in group:
+            group = group + [gt]
+        out.append((gt, p, [[lab, _DESC.get(lab, "")] for lab in group]))
+    return out
+
+
+EXAMPLES = _example_set()
+# one clickable example per document type (keeps the row short + each pick unambiguous)
+_EX_BY_NAME = {}
+for _gt, _p, _t in EXAMPLES:
+    _EX_BY_NAME.setdefault(_gt, (_p, _t))
+EXAMPLE_NAMES = list(_EX_BY_NAME)
+
+
+def _load_example(name):
+    p, t = _EX_BY_NAME.get(name, (None, DEFAULT_ROWS))
+    return p, _rows_to_text(t)
+
+
+def _parse(text):
+    labels, queries = [], []
+    for line in (text or "").splitlines():
+        lab, _, desc = line.partition("|")
+        lab, desc = lab.strip(), desc.strip()
+        if not lab:
+            continue
+        labels.append(lab); queries.append(f"{lab}. {desc}" if desc else lab)
+    return labels, queries
+
+
+def _table(labels, probs):
+    ranked = sorted(zip(labels, probs), key=lambda x: -x[1])
+    rows = []
+    for i, (lab, p) in enumerate(ranked, 1):
+        pct = max(0.0, min(1.0, p)) * 100
+        star = " ✅" if i == 1 else ""
+        rows.append(
+            f"<tr><td style='text-align:right;padding:4px 10px;color:#5c6773'>{i}</td>"
+            f"<td style='padding:4px 10px'><b>{lab}</b>{star}</td>"
+            f"<td style='padding:4px 10px;width:180px'>"
+            f"<div style='background:#e6ecf5;border-radius:4px;height:14px'>"
+            f"<div style='background:#2f52d0;height:14px;border-radius:4px;width:{pct:.0f}%'></div></div></td>"
+            f"<td style='text-align:right;padding:4px 10px;font-variant-numeric:tabular-nums'>{p:.3f}</td></tr>")
+    return ("<table style='border-collapse:collapse;font:14px system-ui'>"
+            "<thead><tr>"
+            "<th style='text-align:right;padding:4px 10px'>#</th>"
+            "<th style='text-align:left;padding:4px 10px'>Label</th>"
+            "<th style='padding:4px 10px'></th>"
+            "<th style='text-align:right;padding:4px 10px'>Score</th></tr></thead>"
+            "<tbody>" + "".join(rows) + "</tbody></table>")
+
+
+def _ort(path):
+    import onnxruntime as ort
+    return ort.InferenceSession(path, providers=["CPUExecutionProvider"])
+
+
+# ---------- v2 flagship: full model as ONNX from the PRIVATE repo (censored) ----------
+def _ensure_v2():
+    if _V2:
+        return
+    if not TOKEN:
+        _V2["error"] = "no HF_TOKEN secret"; return
+    try:
+        dl = lambda f: hf_hub_download(V2_PRIV, f, repo_type="model", token=TOKEN)
+        _V2["img"] = _ort(dl("weights/image_encoder.onnx"))    # pixel_values -> pooled patches [1,256,1152]
+        _V2["txt"] = _ort(dl("weights/text_encoder.onnx"))     # input_ids/mask -> L2 label embeds [N,1024]
+        _V2["head"] = _ort(dl("weights/head-v3-ep6.onnx"))     # patches + label_embeds -> probs
+        _V2["ip"] = AutoImageProcessor.from_pretrained(V2_PRIV, subfolder="encoder/image", token=TOKEN)
+        _V2["tk"] = AutoTokenizer.from_pretrained(V2_PRIV, subfolder="encoder/text", token=TOKEN)
+    except Exception as e:  # noqa: BLE001
+        _V2["error"] = f"{type(e).__name__}"
+
+
+def _run_v2(image, queries):
+    _ensure_v2()
+    if "error" in _V2:
+        return None, f"⚠️ Flagship unavailable ({_V2['error']}). This model needs the commercial weights."
+    pv = _V2["ip"](images=[image.convert("RGB")], return_tensors="np")["pixel_values"].astype(np.float32)
+    patches = _V2["img"].run(["patches"], {"pixel_values": pv})[0]
+    enc = _V2["tk"](queries, padding=True, truncation=True, max_length=512, return_tensors="np")
+    lab = _V2["txt"].run(["embeds"], {"input_ids": enc["input_ids"].astype(np.int64),
+                                      "attention_mask": enc["attention_mask"].astype(np.int64)})[0]
+    probs = _V2["head"].run(["probs"], {"patches": patches.astype(np.float32),
+                                        "label_embeds": lab.astype(np.float32)})[0][0].tolist()
+    return probs, None
+
+
+# ---------- v1 open-weight: public ONNX (image + text towers) ----------
+def _ensure_v1():
+    if _V1:
+        return
+    try:
+        img = hf_hub_download(V1_REPO, "modules/omni-image/image_model.onnx", repo_type="model", token=TOKEN)
+        txt = hf_hub_download(V1_REPO, "modules/omni-image/text_model.onnx", repo_type="model", token=TOKEN)
+        import json
+        cfg = hf_hub_download(V1_REPO, "modules/omni-image/config.json", repo_type="model", token=TOKEN)
+        cal = json.load(open(cfg)).get("calibration") or {"scale": 1.0, "bias": 0.0}
+        _V1["img"] = _ort(img); _V1["txt"] = _ort(txt)
+        _V1["tok"] = AutoTokenizer.from_pretrained(V1_REPO, subfolder="modules/omni-image", token=TOKEN)
+        _V1["proc"] = AutoImageProcessor.from_pretrained(V1_REPO, subfolder="modules/omni-image", token=TOKEN)
+        _V1["scale"], _V1["bias"] = float(cal["scale"]), float(cal["bias"])
+    except Exception as e:  # noqa: BLE001
+        _V1["error"] = f"{type(e).__name__}"
+
+
+def _run_v1(image, queries):
+    _ensure_v1()
+    if "error" in _V1:
+        return None, f"⚠️ v1 unavailable ({_V1['error']})."
+    pix = _V1["proc"](images=[image.convert("RGB")], return_tensors="np")["pixel_values"].astype(np.float16)
+    ie = _V1["img"].run(["image_emb"], {"pixel_values": pix})[0]                 # [1,D]
+    enc = _V1["tok"](queries, padding=True, truncation=True, max_length=64, return_tensors="np")
+    te = _V1["txt"].run(["text_emb"], {"input_ids": enc["input_ids"].astype(np.int64),
+                                       "attention_mask": enc["attention_mask"].astype(np.int64)})[0]  # [N,D]
+    cos = (ie @ te.T)[0]
+    probs = (1.0 / (1.0 + np.exp(-(_V1["scale"] * cos + _V1["bias"])))).tolist()
+    return probs, None
+
+
+LINKS = ("→ [model](https://huggingface.co/nutrientdocs/document-classification-v2) · "
+         "[leaderboard](https://huggingface.co/spaces/nutrientdocs/document-classification-leaderboard) · "
+         "[benchmark](https://huggingface.co/datasets/nutrientdocs/document-classification-benchmark)")
+
+MODELS = {"v2 flagship — best accuracy (commercial)": "v2",
+          "v1 — open-weight (downloadable)": "v1"}
+
+
+@spaces.GPU(duration=120)
+def classify(image, text, model_name):
+    if image is None:
+        return "Upload a document image first."
+    labels, queries = _parse(text)
+    if not labels:
+        return "Add at least one class (one <b>label</b> per line)."
+    probs, err = (_run_v2 if MODELS.get(model_name) == "v2" else _run_v1)(image, queries)
+    return err if err else _table(labels, probs)
+
+
+with gr.Blocks(title="Nutrient Document Classification", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# Classify any document against labels you choose\n"
+                "Open-vocabulary, zero-shot — add candidate classes (label + optional description), upload a "
+                "page, and the model ranks them. No fixed class list.\n\n" + LINKS)
+    with gr.Row():
+        img = gr.Image(type="pil", label="Document image", height=360)
+        with gr.Column():
+            model = gr.Radio(list(MODELS), value=list(MODELS)[0], label="Model")
+            tbl = gr.Textbox(value=DEFAULT_TEXT, lines=10, label="Candidate classes",
+                             info="One class per line — `label | description` (the description is optional).")
+    btn = gr.Button("Classify", variant="primary")
+    out = gr.HTML()
+    sel = gr.Textbox(visible=False)
+    gr.Examples(examples=[[n] for n in EXAMPLE_NAMES], inputs=[sel], outputs=[img, tbl],
+                fn=_load_example, run_on_click=True, cache_examples=False,
+                label="Examples the model classifies correctly — click to load the page + its labels")
+    btn.click(classify, [img, tbl, model], out)
+
+    def _q(request: gr.Request):
+        m = (request.query_params or {}).get("model", "")
+        return list(MODELS)[1] if "v1" in m.lower() else list(MODELS)[0]
+    demo.load(_q, None, model)
+
+demo.queue().launch(ssr_mode=False)

requirements.txt

@@ -0,0 +1,9 @@
+torch
+torchvision
+transformers>=4.49
+onnxruntime
+sentencepiece
+protobuf
+pillow
+accelerate
+spaces