kurenai 紅

Small, fast page-routing classifiers distilled from VLM judgments.

A vision-language model is good at judging document pages — and expensive at scale. kurenai distills those judgments into small image models: feed in a page image, get back an 8-class layout-score vector and a binary routing decision (agentic / non_agentic) in milliseconds instead of a VLM call.

Typical use: in a document-processing pipeline, route easy pages (clean text) to cheap OCR/text extraction and hard pages (tables, forms, charts, complex or degraded layouts) to an agentic VLM pipeline.

Models

Three ONNX checkpoints, all trained on the same corpus and evaluated on the same 152,182-page held-out test split (documents never straddle splits):

file	backbone	params	route acc	route F1	score MAE
finetune-convnext.onnx	ConvNeXt-Tiny	28M	0.951	0.914	0.028
finetune-mobilenetv3.onnx	MobileNetV3-Large	5M	0.952	0.916	0.029
probe-dinov2s.onnx	DINOv2-S (frozen) + MLP	22M + 0.2M	0.934	0.886	0.039

Majority-class baseline: 0.714 accuracy. Teacher-rule ceiling: 0.987.

Which one?

• finetune-convnext.onnx — best soft-score fidelity (lowest score MAE) and best indirect routing; the default in the kurenai SDK.
• finetune-mobilenetv3.onnx — leads the binary route at 5.6× fewer parameters; the pick for routing-only or latency-sensitive use.
• probe-dinov2s.onnx — cheapest to retrain (only the MLP head); ~3 F1 points behind the fine-tunes.

What it predicts

Eight layout classes, each scored independently with a sigmoid:

clean_single_column, multi_column_text, table_heavy, form_or_key_value, figure_chart_heavy, complex_mixed, scanned_or_degraded, title_blank_other

plus a dedicated binary route head. The routing rule marks a page agentic when it belongs to table_heavy, form_or_key_value, figure_chart_heavy, complex_mixed, or scanned_or_degraded; everything else is non_agentic.

Usage (onnxruntime)

import numpy as np
import onnxruntime as ort
from PIL import Image
from huggingface_hub import hf_hub_download

CLASSES = ["clean_single_column", "multi_column_text", "table_heavy",
           "form_or_key_value", "figure_chart_heavy", "complex_mixed",
           "scanned_or_degraded", "title_blank_other"]

def preprocess(img: Image.Image, size: int = 224) -> np.ndarray:
    if img.mode != "RGB":                      # composite RGBA onto white
        rgba = img.convert("RGBA")
        img = Image.new("RGB", rgba.size, (255, 255, 255))
        img.paste(rgba, mask=rgba.getchannel("A"))
    w, h = img.size                            # pad to centered square on white
    side = max(w, h)
    canvas = Image.new("RGB", (side, side), (255, 255, 255))
    canvas.paste(img, ((side - w) // 2, (side - h) // 2))
    arr = np.asarray(canvas.resize((size, size), Image.Resampling.BILINEAR),
                     dtype=np.float32) / 255.0
    arr = arr.transpose(2, 0, 1)               # CHW
    mean = np.array([0.485, 0.456, 0.406], np.float32).reshape(3, 1, 1)
    std = np.array([0.229, 0.224, 0.225], np.float32).reshape(3, 1, 1)
    return (arr - mean) / std

path = hf_hub_download("4thel00z/kurenai", "finetune-mobilenetv3.onnx")
session = ort.InferenceSession(path, providers=["CPUExecutionProvider"])

batch = np.stack([preprocess(Image.open("page.png"))]).astype(np.float32)
score_logits, route_logit = session.run(None, {"pixel_values": batch})

sigmoid = lambda x: 1.0 / (1.0 + np.exp(-x))
route_prob = float(sigmoid(route_logit)[0])
scores = dict(zip(CLASSES, sigmoid(score_logits)[0].tolist()))

print("agentic" if route_prob >= 0.5 else "non_agentic", route_prob)
print(max(scores, key=scores.get), scores)

There is also a small torch-free Python SDK (Classifier(), batch inference, CLI, automatic CUDA/CoreML provider selection) in the kurenai repository.

Model I/O

tensor	shape	dtype	notes
pixel_values (input)	[batch, 3, 224, 224]	float32	ImageNet-normalized, see preprocessing above
score_logits (output)	[batch, 8]	float32	logits — apply sigmoid per class
route_logit (output)	[batch, 1]	float32	logit — sigmoid ≥ 0.5 ⇒ agentic

All files are single self-contained ONNX graphs (opset 18, no external-data sidecars), verified against the torch checkpoints at export time (max deviation < 1e-3).

Training

• Data: 1,015,861 rendered PDF pages from 151,915 documents drawn from the FinePDFs corpus, 100% teacher-labeled by the kimi-k2.6 vision-language model (eight soft layout scores + a routing decision per page).
• Splits: frozen document-level stratified splits — 709,234 train / 154,445 val / 152,182 test pages. Pages of one document never straddle splits (page-level splits would leak).
• Objective: route_BCE(pos_weighted) + λ · score_BCE — the 8-class score head doubles as an auxiliary signal that regularizes the route head.
• Augmentation: small rotation and mild color jitter only. No horizontal flips — they destroy reading-order layout cues.
• Input: aspect-preserving — pages are padded to square on white, never squashed.
• Recipe: AdamW, cosine schedule, early stopping on validation route F1.

Limitations

• Trained exclusively on rendered PDF page thumbnails; behavior on photos, screenshots, or non-document images is undefined.
• Labels are distilled from a VLM teacher, not human annotations — the models inherit the teacher's biases, and the deterministic label→route rule caps achievable accuracy at 0.987 on this corpus.
• scanned_or_degraded never occurs as a hard label in the training data; it is routed agentic by rule, so its score is the least calibrated of the eight.
• The route threshold (0.5) was not tuned per deployment; shift it to trade precision against recall for your pipeline.

License

MIT