nutrientdocs/document-classification-v1

1---2license: apache-2.03pipeline_tag: zero-shot-image-classification4language:5  - en6tags:7  - zero-shot-image-classification8  - image-classification9  - document-ai10  - open-vocabulary11  - open-weights12datasets:13  - nutrientdocs/document-classification-benchmark14---15 16# document-classification-v1 — open-weight17 18**An open-weight, open-vocabulary document classifier you can download and run.** Supply any set of text19labels at inference; the model scores a document image against them by calibrated cosine and returns a20per-label match probability. No fixed class list, no per-class training.21 22The **open-weight** sibling of the commercial flagship23[`document-classification-v2`](https://huggingface.co/nutrientdocs/document-classification-v2). It ships as24two self-contained **ONNX** graphs — an image tower and a text tower — that you run with `onnxruntime`.25`embed_dim: 1024`; classification `p = sigmoid(scale·cos + bias)` (calibration in26`modules/omni-image/config.json`).27 28- 🎯 **Try it:** [document-classification-demo](https://huggingface.co/spaces/nutrientdocs/document-classification-demo)29- 🏆 **Leaderboard:** [document-classification-leaderboard](https://huggingface.co/spaces/nutrientdocs/document-classification-leaderboard)30- 📊 **Benchmark:** [document-classification-benchmark](https://huggingface.co/datasets/nutrientdocs/document-classification-benchmark)31- 🏵️ **Flagship (commercial):** [document-classification-v2](https://huggingface.co/nutrientdocs/document-classification-v2)32 33## Results (macro-F1, zero-shot)34 35| Benchmark | **v1 (open)** | v2 (commercial) | best cloud VLM |36| --- | ---: | ---: | ---: |37| DocLayNet | 0.75 | **0.97** | 0.83 |38| Forms | 0.80 | **1.00** | 1.00 |39| Tobacco | 0.61 | 0.74 | **0.85** |40| OOD (unseen types) | 0.86 | **0.95** | — |41| OOV (synonym wording) | 0.73 | **0.83** | — |42 43Every entry is scored by the same open scorer — full ranking, plus a **generalist zero-shot baseline** and44each cloud model, on the45[leaderboard](https://huggingface.co/spaces/nutrientdocs/document-classification-leaderboard). v1 is the46free, open-weight sibling: it trails the commercial [`v2`](https://huggingface.co/nutrientdocs/document-classification-v2)47and the large cloud VLMs on accuracy, but it's Apache-2.0 and downloadable. Like all embedding models it48trails VLMs most on Tobacco (a read-the-header task). ~**5.7 pages/s on an A40** (fused image+text).49 50## Usage (ONNX)51 52```python53import numpy as np, onnxruntime as ort, json54from transformers import AutoImageProcessor, AutoTokenizer55from huggingface_hub import hf_hub_download56from PIL import Image57 58R = "nutrientdocs/document-classification-v1"59img_sess = ort.InferenceSession(hf_hub_download(R, "modules/omni-image/image_model.onnx"))   # SigLIP image tower60txt_sess = ort.InferenceSession(hf_hub_download(R, "modules/omni-image/text_model.onnx"))     # Qwen text tower61cal = json.load(open(hf_hub_download(R, "modules/omni-image/config.json")))["calibration"]62proc = AutoImageProcessor.from_pretrained(R, subfolder="modules/omni-image")   # SigLIP image processor63tok  = AutoTokenizer.from_pretrained(R, subfolder="modules/omni-image")         # Qwen tokenizer64 65labels = ["invoice", "letter", "memo", "form", "scientific article", "resume"]66calib = lambda cos: 1 / (1 + np.exp(-(cal["scale"] * cos + cal["bias"])))67 68def embed_text(texts, maxlen):69    e = tok(texts, padding=True, truncation=True, max_length=maxlen, return_tensors="np")70    return txt_sess.run(["text_emb"], {"input_ids": e["input_ids"].astype(np.int64),71                                       "attention_mask": e["attention_mask"].astype(np.int64)})[0]  # [.,1024] L272 73lab = embed_text(labels, 64)                                              # label embeds, once74 75# --- image branch: page image vs labels (image ONNX has batch=1; loop+pool for multi-page) ---76pix = proc(images=[Image.open("doc.png").convert("RGB")], return_tensors="np")["pixel_values"].astype(np.float16)77ie  = img_sess.run(["image_emb"], {"pixel_values": pix})[0]              # [1,1024] L278image_probs = calib((ie @ lab.T)[0])                                     # [N]79 80# --- text branch: the page's OCR text vs labels (up to ~2048 tokens) ---81doc_text = open("doc.txt").read()82text_probs = calib((embed_text([doc_text], 2048) @ lab.T)[0])            # [N]83 84# --- reliability fusion: weight each branch by how DECISIVE it is (top1-top2 margin) ---85margin = lambda p: float(np.partition(p, -2)[-1] - np.partition(p, -2)[-2])86wi, wt = margin(image_probs), margin(text_probs); s = wi + wt + 1e-987fused = (wi / s) * image_probs + (wt / s) * text_probs88print(dict(zip(labels, fused.round(3).tolist())))89```90 91## What's in this repo92- `modules/omni-image/{image_model.onnx, text_model.onnx}` — the image + text towers (fp16, `onnxruntime`).93- `modules/omni-image/{config.json, preprocessor_config.json, tokenizer.json}` — calibration + the94  preprocessor and tokenizer needed to run them. That's it — nothing else required.95 96Open weights under **Apache-2.0** — free to download and run. For the higher-accuracy commercial flagship97(on-prem, calibrated), see [`document-classification-v2`](https://huggingface.co/nutrientdocs/document-classification-v2).98 99## About the author100 101<a href="https://nutrient.io/">102  <img src="https://avatars2.githubusercontent.com/u/1527679?v=3&s=200" height="80" />103</a>104 105This 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.106