| --- |
| license: apache-2.0 |
| pipeline_tag: zero-shot-image-classification |
| language: |
| - en |
| tags: |
| - zero-shot-image-classification |
| - image-classification |
| - document-ai |
| - open-vocabulary |
| - open-weights |
| datasets: |
| - nutrientdocs/document-classification-benchmark |
| --- |
| |
| # document-classification-v1 β open-weight |
|
|
| **An open-weight, open-vocabulary document classifier you can download and run.** Supply any set of text |
| labels at inference; the model scores a document image against them by calibrated cosine and returns a |
| per-label match probability. No fixed class list, no per-class training. |
|
|
| The **open-weight** sibling of the commercial flagship |
| [`document-classification-v2`](https://huggingface.co/nutrientdocs/document-classification-v2). It ships as |
| two self-contained **ONNX** graphs β an image tower and a text tower β that you run with `onnxruntime`. |
| `embed_dim: 1024`; classification `p = sigmoid(scaleΒ·cos + bias)` (calibration in |
| `modules/omni-image/config.json`). |
|
|
| - π― **Try it:** [document-classification-demo](https://huggingface.co/spaces/nutrientdocs/document-classification-demo) |
| - π **Leaderboard:** [document-classification-leaderboard](https://huggingface.co/spaces/nutrientdocs/document-classification-leaderboard) |
| - π **Benchmark:** [document-classification-benchmark](https://huggingface.co/datasets/nutrientdocs/document-classification-benchmark) |
| - π΅οΈ **Flagship (commercial):** [document-classification-v2](https://huggingface.co/nutrientdocs/document-classification-v2) |
|
|
| ## Results (macro-F1, zero-shot) |
|
|
| | Benchmark | **v1 (open)** | v2 (commercial) | best cloud VLM | |
| | --- | ---: | ---: | ---: | |
| | DocLayNet | 0.75 | **0.97** | 0.83 | |
| | Forms | 0.80 | **1.00** | 1.00 | |
| | Tobacco | 0.61 | 0.74 | **0.85** | |
| | OOD (unseen types) | 0.86 | **0.95** | β | |
| | OOV (synonym wording) | 0.73 | **0.83** | β | |
|
|
| Every entry is scored by the same open scorer β full ranking, plus a **generalist zero-shot baseline** and |
| each cloud model, on the |
| [leaderboard](https://huggingface.co/spaces/nutrientdocs/document-classification-leaderboard). v1 is the |
| free, open-weight sibling: it trails the commercial [`v2`](https://huggingface.co/nutrientdocs/document-classification-v2) |
| and the large cloud VLMs on accuracy, but it's Apache-2.0 and downloadable. Like all embedding models it |
| trails VLMs most on Tobacco (a read-the-header task). ~**5.7 pages/s on an A40** (fused image+text). |
|
|
| ## Usage (ONNX) |
|
|
| ```python |
| import numpy as np, onnxruntime as ort, json |
| from transformers import AutoImageProcessor, AutoTokenizer |
| from huggingface_hub import hf_hub_download |
| from PIL import Image |
| |
| R = "nutrientdocs/document-classification-v1" |
| img_sess = ort.InferenceSession(hf_hub_download(R, "modules/omni-image/image_model.onnx")) # SigLIP image tower |
| txt_sess = ort.InferenceSession(hf_hub_download(R, "modules/omni-image/text_model.onnx")) # Qwen text tower |
| cal = json.load(open(hf_hub_download(R, "modules/omni-image/config.json")))["calibration"] |
| proc = AutoImageProcessor.from_pretrained(R, subfolder="modules/omni-image") # SigLIP image processor |
| tok = AutoTokenizer.from_pretrained(R, subfolder="modules/omni-image") # Qwen tokenizer |
| |
| labels = ["invoice", "letter", "memo", "form", "scientific article", "resume"] |
| calib = lambda cos: 1 / (1 + np.exp(-(cal["scale"] * cos + cal["bias"]))) |
| |
| def embed_text(texts, maxlen): |
| e = tok(texts, padding=True, truncation=True, max_length=maxlen, return_tensors="np") |
| return txt_sess.run(["text_emb"], {"input_ids": e["input_ids"].astype(np.int64), |
| "attention_mask": e["attention_mask"].astype(np.int64)})[0] # [.,1024] L2 |
| |
| lab = embed_text(labels, 64) # label embeds, once |
| |
| # --- image branch: page image vs labels (image ONNX has batch=1; loop+pool for multi-page) --- |
| pix = proc(images=[Image.open("doc.png").convert("RGB")], return_tensors="np")["pixel_values"].astype(np.float16) |
| ie = img_sess.run(["image_emb"], {"pixel_values": pix})[0] # [1,1024] L2 |
| image_probs = calib((ie @ lab.T)[0]) # [N] |
| |
| # --- text branch: the page's OCR text vs labels (up to ~2048 tokens) --- |
| doc_text = open("doc.txt").read() |
| text_probs = calib((embed_text([doc_text], 2048) @ lab.T)[0]) # [N] |
| |
| # --- reliability fusion: weight each branch by how DECISIVE it is (top1-top2 margin) --- |
| margin = lambda p: float(np.partition(p, -2)[-1] - np.partition(p, -2)[-2]) |
| wi, wt = margin(image_probs), margin(text_probs); s = wi + wt + 1e-9 |
| fused = (wi / s) * image_probs + (wt / s) * text_probs |
| print(dict(zip(labels, fused.round(3).tolist()))) |
| ``` |
|
|
| ## What's in this repo |
| - `modules/omni-image/{image_model.onnx, text_model.onnx}` β the image + text towers (fp16, `onnxruntime`). |
| - `modules/omni-image/{config.json, preprocessor_config.json, tokenizer.json}` β calibration + the |
| preprocessor and tokenizer needed to run them. That's it β nothing else required. |
|
|
| Open weights under **Apache-2.0** β free to download and run. For the higher-accuracy commercial flagship |
| (on-prem, calibrated), see [`document-classification-v2`](https://huggingface.co/nutrientdocs/document-classification-v2). |
|
|
| ## 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. |
|
|