--- 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 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.