"""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 try: # package import (trust_remote_code) — falls back for direct use in a clone from .configuration_siglip2_hier import Siglip2HierConfig except ImportError: 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 = [] @classmethod def from_pretrained(cls, *args, **kwargs): """Load, then REFUSE to return a model with a random vision tower. transformers < 5.6 expects the pre-rename tensor names (vision.vision_model.*) and silently leaves every vision.* tensor newly initialized — the model then classifies with high confidence from random weights. Missing vision keys are an error, not a warning.""" wants_info = kwargs.pop("output_loading_info", False) model, info = super().from_pretrained(*args, output_loading_info=True, **kwargs) missing = [k for k in info.get("missing_keys", []) if k.startswith("vision.")] if missing: import transformers raise RuntimeError( f"{len(missing)} vision-tower tensors were missing from the " f"checkpoint and would be randomly initialized (e.g. " f"{missing[0]!r}). This happens on transformers < 5.6 " f"(installed: {transformers.__version__}), which expects " f"different tensor names. Install transformers>=5.6.0." ) return (model, info) if wants_info else model 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).""" # importlib rather than a plain import: transformers' dynamic-module # loader scans import statements statically and would demand `optimum` # for every user, even though only quantized loading needs it. import importlib _quanto = importlib.import_module("optimum.quanto") freeze, qint4, qint8, quantize = _quanto.freeze, _quanto.qint4, _quanto.qint8, _quanto.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()