eka-doc-classifier: base1024 hier (flat L2 default + scope/l1 select, quality 1-100 score)

README.md

@@ -0,0 +1,157 @@
+---
+license: apache-2.0
+pipeline_tag: image-classification
+tags:
+  - siglip2
+  - document-classification
+  - hierarchical
+  - multi-task
+library_name: transformers
+---
+
+# Hierarchical SigLIP2 Document Classifier (~97M)
+
+A **frozen SigLIP2 base** vision tower (`google/siglip2-base-patch16-naflex`, single
+high-res forward at `max_num_patches=1024`) plus lightweight heads that classify a
+document image across several tasks from one shared pooled embedding:
+
+| Task | Output |
+|------|--------|
+| `l2` (flat, **default**) | leaf class — flat 27-way head; **L1 inferred** from it |
+| `l1` | document group — one of 5 (Document, Miscellaneous, Diagnostic imaging, Card / credential, Body images) |
+| `l2` (leaf) | leaf class via a per-L1 head, when the group is fixed/predicted |
+| `medical` | sigmoid — `medical` / `non_medical` (+ `p_positive`) |
+| `handwritten` | sigmoid — `handwritten` / `printed` (+ `p_positive`); tuned on documents |
+| `quality` | sigmoid — a single **`score` in 1–100** (P(good) mapped onto 1–100; higher = better), no label |
+
+## L2 has two paths, selected by `scope`
+
+There is **both** a flat 27-way L2 head (the backbone's own end-to-end classifier,
+run on the raw pooled feature) **and** a per-L1 leaf head per group (run on the
+standardized embedding). You choose how L2 / L1 are produced:
+
+| `scope` | what runs | L1 | L2 |
+|---------|-----------|----|----|
+| `"flat"` *(default)* | flat 27-way head | **inferred** from the L2 prediction | flat |
+| `"l1"` | L1 group head only | predicted | — (not produced) |
+| `"hierarchical"` | L1 head → that group's leaf head | predicted | leaf |
+
+**Fix the group yourself** by passing `l1="Document"` (any group name): L2 is then
+read from that group's leaf head, overriding the scope's routing.
+
+```
+image → frozen SigLIP2 tower → pooled embedding ──► flat L2 head ─► L2 (+ inferred L1)   [scope="flat"]
+                                     │
+                                     ├─ standardize (mu/sd) → L1 head ─► L1               [scope="l1"]
+                                     │                          └─► leaf L2 head[L1] ─► L2 [scope="hierarchical" / fixed l1]
+                                     └─ medical · handwritten · quality  (sigmoid heads, independent)
+```
+
+### Taxonomy (5 L1 groups → 27 L2 leaves)
+
+```
+Document  (10 leaves)
+    ├─ Letter                 (key: letter)
+    ├─ Other document         (key: other_document)
+    ├─ Scan interpretation    (key: scan_interpretation)
+    ├─ Lab / diagnostic report (key: lab_diagnostic_report)
+    ├─ OPD Consultation Record (key: opd_consultation_record)
+    ├─ Discharge summary      (key: discharge_summary)
+    ├─ Insurance document     (key: insurance_document)
+    ├─ Form                   (key: form)
+    ├─ Invoice / bill         (key: invoice_bill)
+    ├─ Certificate            (key: certificate)
+Miscellaneous  (2 leaves)
+    ├─ Other (Miscellaneous)  (key: other_miscellaneous)
+    ├─ Medication image       (key: medication_image)
+Diagnostic imaging  (5 leaves)
+    ├─ Ultrasound             (key: ultrasound)
+    ├─ X-ray                  (key: x_ray)
+    ├─ MRI                    (key: mri)
+    ├─ CT                     (key: ct)
+    ├─ Other (Diagnostic imaging) (key: other_diagnostic_imaging)
+Card / credential  (7 leaves)
+    ├─ ABHA card              (key: abha_card)
+    ├─ Voter ID               (key: voter_id)
+    ├─ Aadhaar card           (key: aadhaar_card)
+    ├─ Insurance e-card       (key: insurance_e_card)
+    ├─ PAN card               (key: pan_card)
+    ├─ PMJAY / Ayushman       (key: pmjay_ayushman)
+    ├─ Other card             (key: other_card)
+Body images  (3 leaves)
+    ├─ Headshot               (key: headshot)
+    ├─ Skin / wound           (key: skin_wound)
+    ├─ Other (Body images)    (key: other_body_images)
+```
+
+## Usage
+
+```python
+from transformers import AutoModel
+from PIL import Image
+
+model = AutoModel.from_pretrained("ekacare/eka-doc-classifier", trust_remote_code=True).eval()
+img = Image.open("doc.jpg")
+
+model.classify(img)                       # default: flat L2 + inferred L1 + binaries
+```
+
+```jsonc
+// example return value (scope="flat")
+{
+  "l2":          { "source": "flat", "key": "lab_diagnostic_report",
+                    "value": "Lab / diagnostic report", "confidence": 0.97 },
+  "l1":          { "value": "Document", "source": "inferred_from_flat_l2" },
+  "medical":     { "value": "medical", "confidence": 0.99, "p_positive": 0.99 },
+  "handwritten": { "value": "printed", "confidence": 0.98, "p_positive": 0.02 },
+  "quality":     { "score": 94, "p_good": 0.94 }
+}
+```
+
+### Choosing the L2 scope
+
+```python
+model.classify(img)                       # scope="flat" (default): flat L2, L1 inferred
+model.classify(img, scope="l1")           # L1 group only (no L2)
+model.classify(img, scope="hierarchical") # L1 head → leaf L2
+model.classify(img, l1="Document")         # FIX L1 = Document → that group's leaf L2
+model.classify(img, medical=False, handwritten=False, quality=False)  # L2/L1 only
+```
+
+`medical`, `handwritten`, `quality` are independent sigmoid heads — toggle them
+freely. `quality` reports a single **`score` in 1–100** (the sigmoid P(good) mapped
+onto the range; higher = better quality) plus the raw `p_good` — there is no
+poor/good label. `handwritten` is trained on document images.
+
+### Top-k
+
+```python
+model.classify(img, top_k=3)              # top-3 candidates for the multi-class heads
+```
+
+## Quantized loading (optimum-quanto — CPU & GPU)
+
+int8 **and int4** run on CPU as well as GPU. Only the vision tower is quantized;
+`patch_embedding` and the heads stay full precision.
+
+```python
+# pip install optimum-quanto
+from modeling_siglip2_hier import load_classifier   # downloaded by trust_remote_code
+
+model = load_classifier("ekacare/eka-doc-classifier", quantization="int4")  # None | "int8" | "int4"
+model.classify(img)
+```
+
+Flat-L2 accuracy by precision (vision tower quantized; full internal test = 16,204
+images, external OOD = 953; see `benchmark_results.json`):
+
+| precision | internal acc / mF1 (cluster) | external acc / mF1 (cluster) |
+|-----------|------------------------------|------------------------------|
+| bf16 | 96.45 / 0.965 (97.36) | 80.90 / 0.702 (90.95) |
+| int8 | 96.43 / 0.965 (97.36) | 80.90 / 0.653 (91.15) |
+| int4 | 95.22 / 0.955 (94.23) | 77.23 / 0.630 (85.60) |
+
+**int8 is effectively lossless on accuracy** (internal −0.02, external acc identical;
+only the external macro-F1 dips as rare classes lose a little balance). **int4** trades
+~1.2pt internal / ~3.7pt external for the smallest footprint. Base weights ship in bf16
+(~193 MB); int8/int4 quantize the vision tower only (heads + `patch_embedding` stay fp).

benchmark_results.json

@@ -0,0 +1,51 @@
+{
+  "flat_l2": {
+    "bf16": {
+      "internal": {
+        "acc": 96.45,
+        "macro_f1": 0.9646,
+        "cluster_acc": 97.36,
+        "n": 16204
+      },
+      "external": {
+        "acc": 80.9,
+        "macro_f1": 0.7019,
+        "cluster_acc": 90.95,
+        "n": 953
+      },
+      "img_per_s": 78.4
+    },
+    "int8": {
+      "internal": {
+        "acc": 96.43,
+        "macro_f1": 0.9645,
+        "cluster_acc": 97.36,
+        "n": 16204
+      },
+      "external": {
+        "acc": 80.9,
+        "macro_f1": 0.6527,
+        "cluster_acc": 91.15,
+        "n": 953
+      },
+      "img_per_s": 69.3
+    },
+    "int4": {
+      "internal": {
+        "acc": 95.22,
+        "macro_f1": 0.9553,
+        "cluster_acc": 94.23,
+        "n": 16204
+      },
+      "external": {
+        "acc": 77.23,
+        "macro_f1": 0.6304,
+        "cluster_acc": 85.6,
+        "n": 953
+      },
+      "img_per_s": 24.8
+    }
+  },
+  "internal_n": 16204,
+  "external_n": 953
+}

config.json

@@ -0,0 +1,201 @@
+{
+  "architectures": [
+    "Siglip2HierForDocClassification"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_siglip2_hier.Siglip2HierConfig",
+    "AutoModel": "modeling_siglip2_hier.Siglip2HierForDocClassification"
+  },
+  "base_model_id": "google/siglip2-base-patch16-naflex",
+  "default_scope": "flat",
+  "dtype": "bfloat16",
+  "flat_l2_classes": [
+    "letter",
+    "other_document",
+    "scan_interpretation",
+    "lab_diagnostic_report",
+    "opd_consultation_record",
+    "discharge_summary",
+    "insurance_document",
+    "form",
+    "invoice_bill",
+    "certificate",
+    "other_miscellaneous",
+    "medication_image",
+    "ultrasound",
+    "x_ray",
+    "mri",
+    "ct",
+    "other_diagnostic_imaging",
+    "abha_card",
+    "voter_id",
+    "aadhaar_card",
+    "insurance_e_card",
+    "pan_card",
+    "pmjay_ayushman",
+    "other_card",
+    "headshot",
+    "skin_wound",
+    "other_body_images"
+  ],
+  "flat_l2_to_l1": {
+    "aadhaar_card": "Card / credential",
+    "abha_card": "Card / credential",
+    "certificate": "Document",
+    "ct": "Diagnostic imaging",
+    "discharge_summary": "Document",
+    "form": "Document",
+    "headshot": "Body images",
+    "insurance_document": "Document",
+    "insurance_e_card": "Card / credential",
+    "invoice_bill": "Document",
+    "lab_diagnostic_report": "Document",
+    "letter": "Document",
+    "medication_image": "Miscellaneous",
+    "mri": "Diagnostic imaging",
+    "opd_consultation_record": "Document",
+    "other_body_images": "Body images",
+    "other_card": "Card / credential",
+    "other_diagnostic_imaging": "Diagnostic imaging",
+    "other_document": "Document",
+    "other_miscellaneous": "Miscellaneous",
+    "pan_card": "Card / credential",
+    "pmjay_ayushman": "Card / credential",
+    "scan_interpretation": "Document",
+    "skin_wound": "Body images",
+    "ultrasound": "Diagnostic imaging",
+    "voter_id": "Card / credential",
+    "x_ray": "Diagnostic imaging"
+  },
+  "hand_classes": [
+    "printed",
+    "handwritten"
+  ],
+  "head_hidden": 512,
+  "hidden_size": 768,
+  "l1_classes": [
+    "Document",
+    "Miscellaneous",
+    "Diagnostic imaging",
+    "Card / credential",
+    "Body images"
+  ],
+  "l2_by_l1": {
+    "Body images": [
+      "headshot",
+      "skin_wound",
+      "other_body_images"
+    ],
+    "Card / credential": [
+      "abha_card",
+      "voter_id",
+      "aadhaar_card",
+      "insurance_e_card",
+      "pan_card",
+      "pmjay_ayushman",
+      "other_card"
+    ],
+    "Diagnostic imaging": [
+      "ultrasound",
+      "x_ray",
+      "mri",
+      "ct",
+      "other_diagnostic_imaging"
+    ],
+    "Document": [
+      "letter",
+      "other_document",
+      "scan_interpretation",
+      "lab_diagnostic_report",
+      "opd_consultation_record",
+      "discharge_summary",
+      "insurance_document",
+      "form",
+      "invoice_bill",
+      "certificate"
+    ],
+    "Miscellaneous": [
+      "other_miscellaneous",
+      "medication_image"
+    ]
+  },
+  "l2_display": {
+    "aadhaar_card": "Aadhaar card",
+    "abha_card": "ABHA card",
+    "certificate": "Certificate",
+    "ct": "CT",
+    "discharge_summary": "Discharge summary",
+    "form": "Form",
+    "headshot": "Headshot",
+    "insurance_document": "Insurance document",
+    "insurance_e_card": "Insurance e-card",
+    "invoice_bill": "Invoice / bill",
+    "lab_diagnostic_report": "Lab / diagnostic report",
+    "letter": "Letter",
+    "medication_image": "Medication image",
+    "mri": "MRI",
+    "opd_consultation_record": "OPD Consultation Record",
+    "other_body_images": "Other (Body images)",
+    "other_card": "Other card",
+    "other_diagnostic_imaging": "Other (Diagnostic imaging)",
+    "other_document": "Other document",
+    "other_miscellaneous": "Other (Miscellaneous)",
+    "pan_card": "PAN card",
+    "pmjay_ayushman": "PMJAY / Ayushman",
+    "scan_interpretation": "Scan interpretation",
+    "skin_wound": "Skin / wound",
+    "ultrasound": "Ultrasound",
+    "voter_id": "Voter ID",
+    "x_ray": "X-ray"
+  },
+  "max_num_patches": 1024,
+  "med_classes": [
+    "non_medical",
+    "medical"
+  ],
+  "model_type": "siglip2_hier_doc",
+  "qual_classes": [
+    "poor",
+    "good"
+  ],
+  "quality_score_label": "good",
+  "quality_score_max": 100,
+  "quality_score_min": 1,
+  "sigmoid_heads": [
+    "med",
+    "hand",
+    "qual"
+  ],
+  "transformers_version": "5.10.2",
+  "vision_config": {
+    "_name_or_path": "",
+    "architectures": null,
+    "attention_dropout": 0.0,
+    "chunk_size_feed_forward": 0,
+    "dtype": null,
+    "hidden_act": "gelu_pytorch_tanh",
+    "hidden_size": 768,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "intermediate_size": 3072,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-06,
+    "model_type": "siglip2_vision_model",
+    "num_attention_heads": 12,
+    "num_channels": 3,
+    "num_hidden_layers": 12,
+    "num_patches": 256,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "patch_size": 16,
+    "problem_type": null,
+    "return_dict": true,
+    "transformers_version": "5.10.2"
+  }
+}

configuration_siglip2_hier.py

@@ -0,0 +1,83 @@
+"""Config for the hierarchical SigLIP2 document classifier.
+
+Self-contained on purpose: every piece of taxonomy/metadata needed at inference
+(L1 groups, the L2 leaves under each L1, the flat-L2 label space, the binary/quality
+class orders, and the human-readable display names) is stored on the config, so the
+published model repo does not depend on the parrotlet package.
+
+Two L2 paths are shipped:
+  * a **flat** 27-way L2 classifier (the backbone's own end-to-end head, run on the
+    RAW pooled feature) — this is the DEFAULT path; L1 is *inferred* from the flat
+    prediction via ``flat_l2_to_l1``.
+  * per-L1 **leaf** L2 heads (run on the standardized embedding) — used when the
+    caller fixes the L1 group, or asks for ``scope="hierarchical"``.
+
+The medical / handwritten heads are single-logit **sigmoid** heads; ``*_classes`` is
+ordered ``[negative, positive]`` and the sigmoid value is P(positive). The ``quality``
+head is also a sigmoid, but is reported **only** as a single continuous ``score`` in
+``[quality_score_min, quality_score_max]`` (default 1–100) — no poor/good label —
+where ``score = round(min + P(good) * (max - min))``.
+"""
+from __future__ import annotations
+
+from typing import Dict, List, Optional
+
+from transformers import PretrainedConfig
+
+
+class Siglip2HierConfig(PretrainedConfig):
+    model_type = "siglip2_hier_doc"
+
+    def __init__(
+        self,
+        base_model_id: str = "google/siglip2-base-patch16-naflex",
+        vision_config: Optional[dict] = None,
+        hidden_size: int = 768,
+        head_hidden: int = 512,
+        max_num_patches: int = 1024,
+        l1_classes: Optional[List[str]] = None,
+        l2_by_l1: Optional[Dict[str, List[str]]] = None,
+        flat_l2_classes: Optional[List[str]] = None,
+        flat_l2_to_l1: Optional[Dict[str, str]] = None,
+        med_classes: Optional[List[str]] = None,
+        hand_classes: Optional[List[str]] = None,
+        qual_classes: Optional[List[str]] = None,
+        sigmoid_heads: Optional[List[str]] = None,
+        quality_score_label: str = "good",
+        quality_score_min: int = 1,
+        quality_score_max: int = 100,
+        default_scope: str = "flat",
+        l2_display: Optional[Dict[str, str]] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.base_model_id = base_model_id
+        self.vision_config = vision_config
+        self.hidden_size = hidden_size
+        self.head_hidden = head_hidden
+        self.max_num_patches = max_num_patches
+        # Taxonomy. l1_classes is the ordered L1 label space; l2_by_l1 maps each L1
+        # to its ordered list of leaf L2 keys (the per-L1 leaf-head output space).
+        self.l1_classes = list(l1_classes or [])
+        self.l2_by_l1 = {k: list(v) for k, v in (l2_by_l1 or {}).items()}
+        # Flat 27-way L2 head label space + key -> L1 group (for inferring L1 from
+        # the flat prediction in the default scope).
+        self.flat_l2_classes = list(flat_l2_classes or [])
+        self.flat_l2_to_l1 = dict(flat_l2_to_l1 or {})
+        # Binary heads: ordered [negative, positive]; sigmoid value is P(positive).
+        self.med_classes = list(med_classes or ["no", "yes"])
+        self.hand_classes = list(hand_classes or ["no", "yes"])
+        self.qual_classes = list(qual_classes or ["poor", "good"])
+        # Names of the heads that are single-logit sigmoid (vs softmax).
+        self.sigmoid_heads = list(sigmoid_heads or ["med", "hand", "qual"])
+        # The quality class whose sigmoid probability is reported as the score.
+        self.quality_score_label = quality_score_label
+        # quality is reported as a single number in [min, max] (no poor/good label):
+        # score = round(min + P(good) * (max - min)).
+        self.quality_score_min = int(quality_score_min)
+        self.quality_score_max = int(quality_score_max)
+        # Default L2 path: "flat" (flat head + inferred L1), "l1" (L1 head only),
+        # or "hierarchical" (L1 head -> leaf head).
+        self.default_scope = default_scope
+        # Optional pretty names for L2 keys (e.g. "x_ray" -> "X-ray").
+        self.l2_display = dict(l2_display or {})

example.py

@@ -0,0 +1,10 @@
+"""Minimal usage example. Run: python example.py path/to/image.jpg"""
+import sys
+from transformers import AutoModel
+from PIL import Image
+
+model = AutoModel.from_pretrained(".", trust_remote_code=True).eval()
+img = Image.open(sys.argv[1] if len(sys.argv) > 1 else "doc.jpg")
+
+import json
+print(json.dumps(model.classify(img), indent=2))

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:737a15c2d5173dd41233ec742aedee843ac714a1f0cdb2f82a8f6b62e17cb228
+size 193058476

modeling_siglip2_hier.py

@@ -0,0 +1,348 @@
+"""Hierarchical SigLIP2 document classifier as a self-contained HF model.
+
+One pooled embedding of the (frozen) SigLIP2 tower drives several heads:
+
+    image -> SigLIP2 vision tower -> pooled [d]
+        flat L2 head      (on the RAW pooled feature)        -> 27-way leaf, L1 inferred
+        standardize (mu/sd) -> z, then:
+            L1 head                  -> 5-way document group
+            leaf L2 head[L1]         -> leaf within a (predicted or fixed) L1 group
+            medical / handwritten    -> sigmoid (P positive)
+            quality                  -> sigmoid (P good) reported as a 1-100 score
+
+Three L2 *scopes* (``classify(..., scope=...)``):
+  * ``"flat"`` (DEFAULT): run the flat 27-way head; **L1 is inferred** from the
+    predicted L2 class. This is the backbone's own strong end-to-end path.
+  * ``"l1"``: run the L1 group head only (no L2).
+  * ``"hierarchical"``: L1 head picks the group, then that group's leaf L2 head.
+
+Fixing the group: pass ``l1="Document"`` (any L1 name) to **fix the L1** and read L2
+from that group's leaf head — regardless of scope.
+
+Medical / handwritten are single-logit sigmoids (yes/no + ``p_positive``);
+``quality`` returns only a single ``score`` in 1-100 (P(good) mapped onto the range),
+with no poor/good label.
+
+Public API
+----------
+    model = AutoModel.from_pretrained(repo, trust_remote_code=True)
+    model.classify(image)                       # flat L2 + inferred L1 + binaries
+    model.classify(image, scope="l1")           # L1 group only
+    model.classify(image, l1="Document")         # fixed L1 -> leaf L2
+
+Quantized loading (int8/int4) uses optimum-quanto; see ``load_classifier`` /
+``quantize_in_place`` at the bottom.
+"""
+from __future__ import annotations
+
+import os
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel
+
+from .configuration_siglip2_hier import Siglip2HierConfig
+
+
+class _MLPHead(nn.Module):
+    """Linear -> GELU -> Dropout -> Linear. Must match the training-time head so
+    saved weights load 1:1. ``n_out=1`` for the sigmoid binary heads."""
+
+    def __init__(self, d_in: int, n_out: int, hidden: int = 512, dropout: float = 0.0):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(d_in, hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden, n_out)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class Siglip2HierForDocClassification(PreTrainedModel):
+    config_class = Siglip2HierConfig
+    main_input_name = "pixel_values"
+    _tied_weights_keys = []
+    _keys_to_ignore_on_load_missing = []
+
+    def __init__(self, config: Siglip2HierConfig):
+        super().__init__(config)
+        from transformers import Siglip2VisionConfig, Siglip2VisionModel
+
+        vc = config.vision_config
+        if isinstance(vc, dict):
+            vc = Siglip2VisionConfig(**vc)
+        # Only the vision tower of SigLIP2 — the text encoder is never instantiated.
+        self.vision = Siglip2VisionModel(vc) if vc is not None else None
+
+        d, hid = config.hidden_size, config.head_hidden
+        sig = set(config.sigmoid_heads)
+        self.l1_head = _MLPHead(d, len(config.l1_classes), hid)
+        # ModuleList aligned to config.l1_classes order (L1 names contain spaces /
+        # slashes, not valid ModuleDict keys).
+        self.leaf_heads = nn.ModuleList(
+            [_MLPHead(d, len(config.l2_by_l1[l1]), hid) for l1 in config.l1_classes]
+        )
+        # Flat 27-way L2 head = the backbone's own end-to-end classifier (a bare
+        # Linear), applied to the RAW pooled feature (NOT standardized).
+        self.flat_head = nn.Linear(d, len(config.flat_l2_classes)) if config.flat_l2_classes else None
+        # Binary heads: 1 logit when sigmoid, else len(classes).
+        self.med_head = _MLPHead(d, 1 if "med" in sig else len(config.med_classes), hid)
+        self.hand_head = _MLPHead(d, 1 if "hand" in sig else len(config.hand_classes), hid)
+        self.qual_head = _MLPHead(d, 1 if "qual" in sig else len(config.qual_classes), hid)
+
+        # Feature standardization fit on the training embeddings (fp32, always).
+        self.register_buffer("mu", torch.zeros(1, d, dtype=torch.float32))
+        self.register_buffer("sd", torch.ones(1, d, dtype=torch.float32))
+
+        self._processor = None
+        self.post_init()
+
+    # ------------------------------------------------------------------ utils
+    def gradient_checkpointing_enable(self, **kw):
+        if self.vision is not None:
+            self.vision.gradient_checkpointing_enable(**kw)
+
+    def set_processor(self, processor):
+        self._processor = processor
+        return self
+
+    def _get_processor(self):
+        if self._processor is None:
+            from transformers import AutoImageProcessor
+
+            src = self.config._name_or_path or self.name_or_path
+            if not src:
+                raise RuntimeError(
+                    "No image processor available. Load from a repo/dir that ships a "
+                    "preprocessor, or call model.set_processor(proc)."
+                )
+            proc = AutoImageProcessor.from_pretrained(src)
+            proc.max_num_patches = self.config.max_num_patches
+            self._processor = proc
+        return self._processor
+
+    @property
+    def _compute_dtype(self) -> torch.dtype:
+        dt = self.dtype  # PreTrainedModel.dtype skips int (quantized) params
+        return dt if dt.is_floating_point else torch.float16
+
+    # --------------------------------------------------------------- forward
+    @torch.no_grad()
+    def _embed_raw(self, image) -> torch.Tensor:
+        """Image -> RAW pooled embedding [1, d] (fp32). One vision forward."""
+        from PIL import Image
+
+        if isinstance(image, (str, os.PathLike)):
+            image = Image.open(image)
+        elif hasattr(image, "read"):  # file-like / bytes buffer
+            image = Image.open(image)
+        image = image.convert("RGB")
+
+        proc = self._get_processor()
+        enc = proc(images=[image], return_tensors="pt")
+        device, cdtype = self.device, self._compute_dtype
+        enc = {
+            k: (v.to(device=device, dtype=cdtype) if torch.is_floating_point(v) else v.to(device))
+            for k, v in enc.items()
+        }
+        return self.vision(**enc).pooler_output.float()
+
+    def _standardize(self, feat: torch.Tensor) -> torch.Tensor:
+        return (feat - self.mu.float()) / self.sd.float()
+
+    @torch.no_grad()
+    def embed(self, image) -> torch.Tensor:
+        """Back-compat: standardized pooled embedding [1, d] (fp32)."""
+        return self._standardize(self._embed_raw(image))
+
+    # --------------------------------------------------------- head helpers
+    def _softmax_probs(self, head: nn.Module, x: torch.Tensor) -> torch.Tensor:
+        logits = head(x.to(self._compute_dtype))
+        return torch.softmax(logits.float()[0], dim=-1)
+
+    def _sigmoid_pos(self, head: nn.Module, z: torch.Tensor) -> float:
+        """P(positive) from a single-logit sigmoid head."""
+        logit = head(z.to(self._compute_dtype)).float().reshape(-1)[0]
+        return float(torch.sigmoid(logit))
+
+    def _top(self, probs: torch.Tensor, classes, top_k: Optional[int]):
+        if top_k and top_k > 1:
+            k = min(top_k, probs.numel())
+            vals, idx = torch.topk(probs, k)
+            return [{"value": classes[int(i)], "confidence": round(float(v), 4)}
+                    for v, i in zip(vals, idx)]
+        i = int(probs.argmax())
+        return {"value": classes[i], "confidence": round(float(probs[i]), 4)}
+
+    def _flat_l2(self, raw: torch.Tensor, top_k):
+        """Flat 27-way L2 on the RAW feature; infer L1 from the top prediction."""
+        probs = self._softmax_probs(self.flat_head, raw)
+        disp = self.config.l2_display
+        res = self._top(probs, self.config.flat_l2_classes, top_k)
+        if isinstance(res, list):  # top_k
+            for r in res:
+                r["key"] = r["value"]
+                r["value"] = disp.get(r["key"], r["key"])
+            top_key = res[0]["key"]
+            l2 = {"source": "flat", "candidates": res}
+        else:
+            top_key = res["value"]
+            l2 = {"source": "flat", "key": top_key, "value": disp.get(top_key, top_key),
+                  "confidence": res["confidence"]}
+        l1_name = self.config.flat_l2_to_l1.get(top_key)
+        return {"l2": l2, "l1": {"value": l1_name, "source": "inferred_from_flat_l2"}}
+
+    def _leaf_l2(self, z: torch.Tensor, l1_name: str, top_k):
+        """Leaf L2 within a given L1 group."""
+        idx = self.config.l1_classes.index(l1_name)
+        probs = self._softmax_probs(self.leaf_heads[idx], z)
+        leaf_classes = self.config.l2_by_l1[l1_name]
+        disp = self.config.l2_display
+        res = self._top(probs, leaf_classes, top_k)
+        if isinstance(res, list):
+            for r in res:
+                r["key"] = r["value"]
+                r["value"] = disp.get(r["key"], r["key"])
+            return {"source": "leaf", "l1": l1_name, "candidates": res}
+        key = res["value"]
+        return {"source": "leaf", "l1": l1_name, "key": key,
+                "value": disp.get(key, key), "confidence": res["confidence"]}
+
+    def _binary(self, name: str, z: torch.Tensor, head, classes):
+        """yes/no from a sigmoid head; classes = [negative, positive]."""
+        if name in self.config.sigmoid_heads:
+            p = self._sigmoid_pos(head, z)               # P(positive)
+        else:  # softmax fallback
+            probs = self._softmax_probs(head, z)
+            p = float(probs[-1])
+        pos, neg = classes[1], classes[0]
+        is_pos = p >= 0.5
+        return {"value": pos if is_pos else neg,
+                "confidence": round(p if is_pos else 1 - p, 4),
+                "p_positive": round(p, 4)}
+
+    def _quality(self, z: torch.Tensor) -> dict:
+        """Quality as a single number in [min, max] (default 1-100) — no poor/good label.
+
+        The sigmoid head's P(good) is linearly mapped onto the configured range:
+        ``score = round(min + P(good) * (max - min))`` (P=0 -> min, P=1 -> max)."""
+        p = self._sigmoid_pos(self.qual_head, z)         # P(good)
+        if self.config.quality_score_label != self.config.qual_classes[1]:
+            p = 1.0 - p                                  # score tracks the configured label
+        lo, hi = self.config.quality_score_min, self.config.quality_score_max
+        return {"score": int(round(lo + p * (hi - lo))), "p_good": round(p, 4)}
+
+    @torch.no_grad()
+    def classify(
+        self,
+        image,
+        scope: Optional[str] = None,
+        l1: Optional[str] = None,
+        medical: bool = True,
+        quality: bool = True,
+        handwritten: bool = True,
+        top_k: Optional[int] = None,
+    ) -> dict:
+        """Classify one image.
+
+        Args:
+            image: PIL.Image, path, or file-like.
+            scope: L2 path — ``"flat"`` (default; flat head, L1 inferred), ``"l1"``
+                (L1 group head only), or ``"hierarchical"`` (L1 head -> leaf head).
+                Defaults to ``config.default_scope``.
+            l1: fix the L1 group (any name in ``config.l1_classes``) and read L2 from
+                that group's leaf head — overrides scope's L1/L2 routing.
+            medical / handwritten: sigmoid binaries (also report ``p_positive``).
+            quality: sigmoid head; reports a single ``score`` in 1-100 (mapped from
+                P(good)) plus ``p_good`` — no poor/good label.
+            top_k: if >1, return the top-k candidates for the multi-class heads.
+
+        Returns:
+            dict keyed by task. ``l2``/``l1`` entries carry a ``source`` field
+            indicating how they were produced.
+        """
+        scope = (scope or self.config.default_scope).lower()
+        if l1 is not None and l1 not in self.config.l1_classes:
+            raise ValueError(f"unknown l1 group {l1!r}; choose from {self.config.l1_classes}")
+        if scope not in ("flat", "l1", "hierarchical"):
+            raise ValueError(f"unknown scope {scope!r}; use 'flat', 'l1', or 'hierarchical'.")
+        if scope == "flat" and self.flat_head is None and l1 is None:
+            raise ValueError("scope='flat' requires a flat L2 head; this model has none.")
+
+        raw = self._embed_raw(image)        # single vision forward
+        z = self._standardize(raw)
+        out: dict = {}
+
+        if l1 is not None:
+            # Fixed L1 group -> leaf L2.
+            out["l1"] = {"value": l1, "source": "fixed"}
+            out["l2"] = self._leaf_l2(z, l1, top_k)
+        elif scope == "l1":
+            out["l1"] = {**self._top(self._softmax_probs(self.l1_head, z),
+                                     self.config.l1_classes, top_k), "source": "l1_head"}
+        elif scope == "flat":
+            out.update(self._flat_l2(raw, top_k))
+        else:  # hierarchical
+            l1_probs = self._softmax_probs(self.l1_head, z)
+            l1_name = self.config.l1_classes[int(l1_probs.argmax())]
+            out["l1"] = {**self._top(l1_probs, self.config.l1_classes, top_k), "source": "l1_head"}
+            out["l2"] = self._leaf_l2(z, l1_name, top_k)
+
+        if medical:
+            out["medical"] = self._binary("med", z, self.med_head, self.config.med_classes)
+        if handwritten:
+            out["handwritten"] = self._binary("hand", z, self.hand_head, self.config.hand_classes)
+        if quality:
+            out["quality"] = self._quality(z)
+        return out
+
+
+# --------------------------------------------------------------------------- #
+# optimum-quanto quantization shortcuts (int8/int4, CPU & GPU).
+# --------------------------------------------------------------------------- #
+# patch_embedding is excluded — SigLIP2 casts the pixel input to its weight dtype,
+# which an int weight breaks. Heads stay full precision (a few MB, decisive for
+# accuracy); pass quantize_heads=True to override.
+_EXCLUDE_FROM_QUANT = ["*patch_embedding*"]
+
+
+def quantize_in_place(model, bits: int = 4, quantize_heads: bool = False):
+    """Quantize a loaded model's vision tower with optimum-quanto (int8 or int4)."""
+    from optimum.quanto import freeze, qint4, qint8, quantize
+
+    if bits not in (4, 8):
+        raise ValueError(f"bits must be 4 or 8, got {bits}")
+    weights = qint8 if bits == 8 else qint4
+    target = model if quantize_heads else model.vision
+    quantize(target, weights=weights, exclude=list(_EXCLUDE_FROM_QUANT))
+    freeze(target)
+    return model.eval()
+
+
+def load_classifier(
+    model_id: str,
+    quantization: Optional[str] = None,
+    device: Optional[str] = None,
+    dtype: Optional[torch.dtype] = torch.bfloat16,
+    trust_remote_code: bool = True,
+    quantize_heads: bool = False,
+    **kwargs,
+):
+    """Load the hierarchical classifier, optionally quanto-quantized (int8/int4)."""
+    from transformers import AutoModel
+
+    model = AutoModel.from_pretrained(
+        model_id, trust_remote_code=trust_remote_code, dtype=dtype, **kwargs
+    ).eval()
+    if device:
+        model = model.to(device)
+
+    q = (quantization or "").lower()
+    if q in ("int8", "8bit", "8"):
+        quantize_in_place(model, bits=8, quantize_heads=quantize_heads)
+    elif q in ("int4", "4bit", "4"):
+        quantize_in_place(model, bits=4, quantize_heads=quantize_heads)
+    elif q not in ("", "none", "fp16", "bf16", "fp32"):
+        raise ValueError(f"Unknown quantization {quantization!r}; use None/int8/int4.")
+    return model.eval()

preprocessor_config.json

@@ -0,0 +1,20 @@
+{
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_processor_type": "Siglip2ImageProcessor",
+  "image_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "max_num_patches": 1024,
+  "patch_size": 16,
+  "resample": 2,
+  "rescale_factor": 0.00392156862745098
+}

requirements.txt

@@ -0,0 +1,5 @@
+torch>=2.7.0
+transformers>=4.52.0
+pillow
+# Optional — only needed for int8/int4 quantized loading (runs on CPU and GPU):
+optimum-quanto>=0.2.6