hung-k-nguyen's picture
v1 re-scored heuristic + fused usage; v2 col
6c930d5 verified
|
Raw
History Blame Contribute Delete
5.66 kB
metadata
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. 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).

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. v1 is the free, open-weight sibling: it trails the commercial 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)

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.

About the author

This project is maintained and funded by Nutrient - 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.