| """Multimodal prototype: typed questions about IMAGES (optionally with a text state). |
| |
| from jevlike.mm import MultimodalSystemOne |
| m = MultimodalSystemOne.load("checkpoints/jevlike-mm-proto") |
| m.predict(None, {"kind": Choice("What type of document is this?", ["invoice", "letter", "resume"]), |
| "signed": Noul("The page is handwritten")}, image="scan.png") |
| |
| Design (research prototype, see reports/mm_proto.md): |
| |
| * **Vision**: a frozen SigLIP (default google/siglip-base-patch16-224). Its 14x14 patch grid is |
| average-pooled to 7x7 and SigLIP's pooled (attention-pooling head) embedding is prepended, giving |
| ``n_img_tokens`` = 50 vectors. A trainable 2-layer MLP projector maps them into ModernBERT's |
| token-embedding space (plus a learned per-slot position vector). |
| * **Input**: image tokens are spliced into the sequence right after [CLS] through |
| ``inputs_embeds``, framed by learned [IMG_START] / [IMG_END] vectors:: |
| |
| [CLS] [IMG_START] <50 image tokens> [IMG_END] <state text> [SEP] <type tag> <instr> [SEP] [MASK] key: text ... |
| |
| The question part, the option markers, the decision head, the marker scorer and the escalate |
| head are the text model's (jevlike/model.py), unchanged. |
| * **Similarity prior** (``sim_prior``): SigLIP's own image-text score for every option (option text |
| for choice/score, the statement for noul, scored at the ``true`` marker) is turned into a few |
| features (raw logit, sigmoid, z-score / softmax / gap-to-max within the question, question type) |
| and added, through a zero-initialized MLP, to that option's [MASK] embedding. This hands the |
| language side SigLIP's zero-shot knowledge of label names it has never been trained on; the |
| encoder still reads the instructions and options and decides. Training drops the prior for a |
| fraction of questions (``sim_dropout``), so ``use_sim=False`` is a meaningful ablation. |
| * **Text-only inputs** go through exactly the text path (no image tokens, no prior), so the |
| checkpoint directory is also a valid text ``SystemOne`` checkpoint: ``model.safetensors`` holds |
| the DecisionModel, ``mm.safetensors`` + ``mm_config.json`` the multimodal parts; SigLIP weights |
| come from ``mm_config.vision`` (HF id / path) or ``vision.safetensors`` if present. |
| |
| Limitations: 224x224 input (small document text is unreadable), one image per question, long |
| states are truncated (no chunking) when an image is present. |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import math |
| import os |
| import re |
| import time |
| from dataclasses import asdict, dataclass, field |
| from typing import Callable, Mapping, Sequence |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| from safetensors.torch import load_file, save_file |
| from torch import nn |
|
|
| from jevlike.calibrate import temperature |
| from jevlike.model import DecisionModel |
| from jevlike.predict import SystemOne |
| from jevlike.serialize import TYPE_IDS, Encoded, Serializer, TokenBudgetBatchSampler, collate |
| from jevlike.types import QTYPES, PublicQuestion, Question, Record |
|
|
| MM_CONFIG, MM_WEIGHTS, VISION_WEIGHTS = "mm_config.json", "mm.safetensors", "vision.safetensors" |
| N_SIM_FEATS = 6 + len(QTYPES) |
| IMG_OFFSET = 1 |
|
|
| Questions = Mapping[str, "PublicQuestion | Question"] |
| ImageLike = "PIL.Image.Image | str | os.PathLike | np.ndarray | None" |
|
|
|
|
| @dataclass |
| class MMConfig: |
| vision: str = "google/siglip-base-patch16-224" |
| vision_config: dict = field(default_factory=dict) |
| image_size: int = 224 |
| pool: int = 2 |
| use_pooled: bool = True |
| proj_hidden: int = 2048 |
| sim_prior: bool = True |
| sim_dropout: float = 0.3 |
| sim_template: str = "{}" |
| text_max_len: int = 64 |
|
|
| @property |
| def grid(self) -> int: |
| vc = self.vision_config.get("vision_config", {}) |
| return (vc.get("image_size", self.image_size) // vc.get("patch_size", 16)) // self.pool |
|
|
| @property |
| def n_img_tokens(self) -> int: |
| return self.grid ** 2 + int(self.use_pooled) |
|
|
| @property |
| def n_slots(self) -> int: |
| return self.n_img_tokens + 2 |
|
|
| @classmethod |
| def load(cls, path: str) -> "MMConfig": |
| with open(os.path.join(path, MM_CONFIG)) as f: |
| d = json.load(f) |
| return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__}) |
|
|
|
|
| |
|
|
| def load_pixels(image, size: int = 224) -> torch.Tensor: |
| """PIL image / path / HxWx3 uint8 array -> [3, size, size] float tensor, SigLIP preprocessing |
| (squash-resize bicubic, scale to [-1, 1]).""" |
| from PIL import Image |
| if isinstance(image, (str, os.PathLike)): |
| image = Image.open(image) |
| elif isinstance(image, np.ndarray): |
| image = Image.fromarray(image) |
| image = image.convert("RGB") |
| if image.size != (size, size): |
| image = image.resize((size, size), Image.BICUBIC) |
| x = torch.from_numpy(np.asarray(image, dtype=np.float32).copy()).permute(2, 0, 1) / 255.0 |
| return (x - 0.5) / 0.5 |
|
|
|
|
| _OPAQUE = re.compile(r"^([A-Za-z]|\d+|(opt_|option_|class_|label_|c|cat_|choice_|id_)\d+|[A-Z]{2,4})$") |
|
|
|
|
| def option_prompts(q: Question) -> list[str | None]: |
| """The text SigLIP scores for each marker: option text (readable key + text) for choice/score, |
| the statement at the ``true`` marker for noul (``false`` gets none).""" |
| if q.type == "noul": |
| return [None, q.instructions.strip()] |
| out = [] |
| for o in q.options: |
| key, text = o.key.strip(), o.text.strip() |
| if q.type == "score" or _OPAQUE.match(key): |
| out.append(text or key) |
| else: |
| k = key.replace("_", " ") |
| out.append(f"{k}: {text}" if text and text.lower() != k.lower() else k) |
| return out |
|
|
|
|
| class MMSerializer: |
| """Text serializer + image slots: image items get ``n_slots`` placeholder positions after [CLS] |
| (their text is serialized with ``max_len - n_slots``); markers are shifted accordingly.""" |
|
|
| def __init__(self, tokenizer, max_len: int, head_max_len: int, n_slots: int): |
| self.text = Serializer(tokenizer, max_len, head_max_len) |
| img_len = max_len - n_slots |
| self.img = Serializer(tokenizer, img_len, min(head_max_len, img_len - 64)) |
| self.n_slots, self.pad = n_slots, self.text.pad |
|
|
| def encode_many(self, items: Sequence[tuple[str, Question, bool]]) -> list[Encoded]: |
| out: list[Encoded | None] = [None] * len(items) |
| for has_img, ser in ((False, self.text), (True, self.img)): |
| idx = [i for i, it in enumerate(items) if bool(it[2]) == has_img] |
| if not idx: |
| continue |
| for i, e in zip(idx, ser.encode_many([(items[i][0] or "", items[i][1]) for i in idx])): |
| if has_img: |
| e = Encoded([e.input_ids[0]] + [self.pad] * self.n_slots + e.input_ids[1:], |
| [m + self.n_slots for m in e.markers], e.qtype, e.trunc) |
| out[i] = e |
| return out |
|
|
|
|
| def mm_collate(encs: Sequence[Encoded], labels, pad_id: int, image_of: Sequence[int], pixels: Sequence[torch.Tensor], |
| prompts: Sequence[Sequence[str | None]] | None, use_sim: Sequence[bool] | None = None) -> dict: |
| """``serialize.collate`` + image fields. |
| |
| image_of[i]: index into ``pixels`` for row i, or -1 (text-only row). prompts[i]: per-marker SigLIP |
| prompt (or None); use_sim[i]: False withholds the prior for row i. Adds ``pixel_values`` [N, 3, S, S], |
| ``image_rows`` / ``image_index`` [R], ``sim_strings`` (list, unique) and ``sim_marker_text`` [M] |
| (index into sim_strings, -1 = none).""" |
| b = collate(encs, labels, pad_id) |
| used = sorted({j for j in image_of if j >= 0}) |
| remap = {j: k for k, j in enumerate(used)} |
| rows = [i for i, j in enumerate(image_of) if j >= 0] |
| b["image_rows"] = torch.tensor(rows, dtype=torch.long) |
| b["image_index"] = torch.tensor([remap[image_of[i]] for i in rows], dtype=torch.long) |
| b["pixel_values"] = torch.stack([pixels[j] for j in used]) if used else torch.zeros(0, 3, 1, 1) |
| strings: dict[str, int] = {} |
| mt: list[int] = [] |
| for i, e in enumerate(encs): |
| ps = prompts[i] if prompts is not None and image_of[i] >= 0 and (use_sim is None or use_sim[i]) else None |
| for r in range(len(e.markers)): |
| s = ps[r] if ps is not None else None |
| mt.append(strings.setdefault(s, len(strings)) if s else -1) |
| b["sim_strings"] = list(strings) |
| b["sim_marker_text"] = torch.tensor(mt, dtype=torch.long) |
| return b |
|
|
|
|
| |
|
|
| class MultimodalDecisionModel(nn.Module): |
| def __init__(self, dm: DecisionModel, mcfg: MMConfig, siglip: nn.Module | None = None, |
| text_tokenize: Callable[[list[str]], torch.Tensor] | None = None): |
| super().__init__() |
| from transformers import SiglipConfig, SiglipModel |
| if siglip is None: |
| siglip = SiglipModel(SiglipConfig.from_dict(mcfg.vision_config)) |
| self.dm, self.mcfg = dm, mcfg |
| mcfg.vision_config = siglip.config.to_dict() |
| self.siglip = siglip.eval().requires_grad_(False) |
| self.text_tokenize = text_tokenize |
| Dv, D = siglip.config.vision_config.hidden_size, dm.backbone.config.hidden_size |
| n = mcfg.n_img_tokens |
| self.proj = nn.Sequential(nn.LayerNorm(Dv), nn.Linear(Dv, mcfg.proj_hidden), nn.GELU(), |
| nn.Linear(mcfg.proj_hidden, D)) |
| self.img_pos = nn.Parameter(torch.randn(n, D) * 0.02) |
| self.img_start = nn.Parameter(torch.randn(D) * 0.02) |
| self.img_end = nn.Parameter(torch.randn(D) * 0.02) |
| self.sim_mlp = nn.Sequential(nn.Linear(N_SIM_FEATS, 256), nn.GELU(), nn.Linear(256, D)) |
| nn.init.zeros_(self.sim_mlp[-1].weight) |
| nn.init.zeros_(self.sim_mlp[-1].bias) |
| self._text_cache: dict[str, torch.Tensor] = {} |
|
|
| |
| def train(self, mode: bool = True): |
| super().train(mode) |
| self.siglip.eval() |
| return self |
|
|
| def mm_parameters(self) -> list[tuple[str, nn.Parameter]]: |
| return [(n, p) for n, p in self.named_parameters() if not n.startswith(("dm.", "siglip."))] |
|
|
| @property |
| def device(self) -> torch.device: |
| return self.img_pos.device |
|
|
| |
| @torch.no_grad() |
| def image_features(self, pixel_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| """-> (tokens [N, n_img_tokens, Dv], normalized image embedding [N, Dv]).""" |
| vm = self.siglip.vision_model |
| dtype = next(vm.parameters()).dtype |
| out = vm(pixel_values=pixel_values.to(self.device, dtype)) |
| h = out.last_hidden_state |
| N, P, Dv = h.shape |
| g = int(math.isqrt(P)) |
| grid = F.avg_pool2d(h.transpose(1, 2).reshape(N, Dv, g, g).float(), self.mcfg.pool) |
| toks = grid.flatten(2).transpose(1, 2) |
| pooled = out.pooler_output.float() |
| if self.mcfg.use_pooled: |
| toks = torch.cat([pooled[:, None], toks], 1) |
| return toks, F.normalize(pooled, dim=-1) |
|
|
| @torch.no_grad() |
| def text_embeds(self, strings: Sequence[str]) -> torch.Tensor: |
| """Normalized SigLIP text embeddings [U, Dv] (cached per string).""" |
| if not strings: |
| return torch.zeros(0, self.siglip.config.text_config.hidden_size, device=self.device) |
| todo = [s for s in dict.fromkeys(strings) if s not in self._text_cache] |
| for i in range(0, len(todo), 256): |
| chunk = todo[i:i + 256] |
| ids = self.text_tokenize([self.mcfg.sim_template.format(s) for s in chunk]).to(self.device) |
| emb = F.normalize(self.siglip.text_model(input_ids=ids).pooler_output.float(), dim=-1) |
| for s, e in zip(chunk, emb): |
| self._text_cache[s] = e |
| if len(self._text_cache) > 200_000: |
| self._text_cache.clear() |
| return torch.stack([self._text_cache[s] for s in strings]) |
|
|
| def sim_features(self, b: dict, img_emb: torch.Tensor, B: int, K: int) -> tuple[torch.Tensor, torch.Tensor]: |
| """Per-marker prior features. Returns (marker positions into marker_idx, features [m, N_SIM_FEATS]).""" |
| mt = b["sim_marker_text"] |
| sel = (mt >= 0).nonzero().flatten() |
| if len(sel) == 0: |
| return sel, torch.zeros(0, N_SIM_FEATS, device=self.device) |
| txt = self.text_embeds(b["sim_strings"]) |
| row_img = torch.full((B,), -1, dtype=torch.long, device=self.device) |
| row_img[b["image_rows"]] = b["image_index"] |
| seg, rank = b["marker_seg"][sel], b["marker_rank"][sel] |
| cos = (img_emb[row_img[seg]] * txt[mt[sel]]).sum(-1) |
| sl = self.siglip |
| logit = cos * sl.logit_scale.float().exp() + sl.logit_bias.float() |
| |
| M = torch.full((B, K), float("nan"), device=self.device).index_put((seg, rank), logit) |
| valid = ~M.isnan() |
| cnt = valid.sum(-1).clamp(min=1) |
| Mz = M.nan_to_num(0.0) |
| mean = Mz.sum(-1) / cnt |
| std = (((Mz - mean[:, None]) ** 2) * valid).sum(-1).div(cnt).sqrt().clamp(min=1e-3) |
| mx = M.masked_fill(~valid, -1e9).max(-1).values |
| soft = M.masked_fill(~valid, -1e9).softmax(-1) |
| qt = b["qtype"][seg] |
| multi = (cnt[seg] > 1).float() |
| feats = torch.stack([ |
| logit / 10, torch.sigmoid(logit), |
| multi * (logit - mean[seg]) / std[seg], multi * soft[seg, rank], multi * (logit - mx[seg]) / 10, |
| torch.ones_like(logit)], -1) |
| feats = torch.cat([feats, F.one_hot(qt, len(QTYPES)).float()], -1) |
| return sel, feats |
|
|
| |
| def embed(self, b: dict) -> torch.Tensor: |
| ids = b["input_ids"] |
| emb = self.dm.backbone.get_input_embeddings()(ids) |
| B, T, D = emb.shape |
| rows = b.get("image_rows") |
| if rows is None or len(rows) == 0: |
| return emb |
| toks, img_emb = self.image_features(b["pixel_values"]) |
| v = (self.proj(toks) + self.img_pos).to(emb.dtype) |
| n = self.mcfg.n_img_tokens |
| emb = emb.clone() |
| emb[rows, IMG_OFFSET] = self.img_start.to(emb.dtype) |
| emb[rows, IMG_OFFSET + 1:IMG_OFFSET + 1 + n] = v[b["image_index"]] |
| emb[rows, IMG_OFFSET + 1 + n] = self.img_end.to(emb.dtype) |
| if self.mcfg.sim_prior and "sim_marker_text" in b: |
| K = int(b["n_options"].max()) |
| sel, feats = self.sim_features(b, img_emb, B, K) |
| if len(sel): |
| add = self.sim_mlp(feats).to(emb.dtype) |
| emb = emb.reshape(B * T, D).index_add(0, b["marker_idx"][sel], add).reshape(B, T, D) |
| return emb |
|
|
| def forward(self, **b) -> dict[str, torch.Tensor]: |
| rows = b.get("image_rows") |
| if rows is None or len(rows) == 0: |
| return self.dm(**b) |
| return self.dm(**b, inputs_embeds=self.embed(b)) |
|
|
| |
| def save_pretrained(self, path: str, tokenizer=None, save_vision: bool = False) -> None: |
| """model.safetensors / config.json = the DecisionModel (a valid text checkpoint); |
| mm.safetensors / mm_config.json = projector, image slots, prior MLP.""" |
| self.dm.save_pretrained(path, tokenizer) |
| with open(os.path.join(path, MM_CONFIG), "w") as f: |
| json.dump(asdict(self.mcfg), f, indent=1, default=str) |
| save_file({n: p.detach().contiguous().cpu() for n, p in self.mm_parameters()}, os.path.join(path, MM_WEIGHTS)) |
| if save_vision: |
| save_file({k: v.detach().contiguous().cpu() for k, v in self.siglip.state_dict().items()}, |
| os.path.join(path, VISION_WEIGHTS)) |
|
|
| @classmethod |
| def from_pretrained(cls, path: str, device: str | torch.device = "cpu", vision: str | None = None, |
| text_tokenize=None) -> "MultimodalDecisionModel": |
| device = torch.device(device) |
| dm = DecisionModel.from_pretrained(path, device) |
| mcfg = MMConfig.load(path) |
| siglip, tt = load_siglip(path if os.path.exists(os.path.join(path, VISION_WEIGHTS)) else (vision or mcfg.vision), |
| mcfg, device) |
| model = cls(dm, mcfg, siglip, text_tokenize or tt) |
| sd = load_file(os.path.join(path, MM_WEIGHTS)) |
| missing = model.load_state_dict(sd, strict=False) |
| bad = [k for k in missing.missing_keys if not k.startswith(("dm.", "siglip."))] |
| assert not bad and not missing.unexpected_keys, (bad, missing.unexpected_keys) |
| return model.to(device) |
|
|
| @classmethod |
| def from_text_checkpoint(cls, path: str, mcfg: MMConfig | None = None, device="cpu") -> "MultimodalDecisionModel": |
| """New multimodal model around a trained text checkpoint (fresh projector / prior MLP).""" |
| device = torch.device(device) |
| mcfg = mcfg or MMConfig() |
| siglip, tt = load_siglip(mcfg.vision, mcfg, device) |
| return cls(DecisionModel.from_pretrained(path, device), mcfg, siglip, tt).to(device) |
|
|
|
|
| def load_siglip(src: str, mcfg: MMConfig, device: torch.device): |
| """SigLIP model (frozen; bf16 on CUDA) + a text tokenize function. ``src`` is an HF id / dir, or a |
| checkpoint dir holding vision.safetensors (+ mm_config.json with the architecture).""" |
| from transformers import AutoTokenizer, SiglipConfig, SiglipModel |
| if os.path.exists(os.path.join(src, VISION_WEIGHTS)): |
| model = SiglipModel(SiglipConfig.from_dict(mcfg.vision_config)) |
| model.load_state_dict(load_file(os.path.join(src, VISION_WEIGHTS))) |
| tok_src = mcfg.vision |
| else: |
| model = SiglipModel.from_pretrained(src) |
| tok_src = src |
| if device.type == "cuda": |
| model = model.to(torch.bfloat16) |
| model = model.to(device).eval().requires_grad_(False) |
| try: |
| tok = AutoTokenizer.from_pretrained(tok_src) |
| except Exception: |
| tok = None |
|
|
| def tokenize(strings: list[str]) -> torch.Tensor: |
| return tok(strings, padding="max_length", truncation=True, max_length=mcfg.text_max_len, |
| return_tensors="pt")["input_ids"] |
| return model, (tokenize if tok is not None else None) |
|
|
|
|
| |
|
|
| def read_mm_jsonl(path: str, root: str = ".") -> list[tuple[Record, str | None, dict]]: |
| """Record JSONL with an optional extra "image" (path, relative to ``root``) -> (record, image path, raw dict).""" |
| out = [] |
| with open(path, encoding="utf-8") as f: |
| for line in f: |
| if line.strip(): |
| d = json.loads(line) |
| img = d.get("image") |
| out.append((Record.from_json(line), os.path.join(root, img) if img else None, d)) |
| return out |
|
|
|
|
| class MMDataset(torch.utils.data.Dataset): |
| """Pre-encoded (record, image) items; images are decoded in ``collate`` (DataLoader workers).""" |
|
|
| def __init__(self, items: Sequence[tuple[Record, str | None]], ser: MMSerializer, image_size: int = 224, |
| sim_dropout: float = 0.0, seed: int = 0): |
| self.records = [r for r, _ in items] |
| self.images = [img for _, img in items] |
| self.encs = ser.encode_many([(r.state, r.question, img is not None) for r, img in items]) |
| self.prompts = [option_prompts(r.question) for r in self.records] |
| self.labels = [r.label for r in self.records] |
| self.lengths = np.array([len(e) for e in self.encs]) |
| self.pad_id, self.size, self.sim_dropout = ser.pad, image_size, sim_dropout |
| self.rng = np.random.default_rng(seed) |
|
|
| def __len__(self) -> int: |
| return len(self.encs) |
|
|
| def __getitem__(self, i: int) -> int: |
| return i |
|
|
| def collate(self, idx: Sequence[int]) -> dict: |
| paths = list(dict.fromkeys(self.images[i] for i in idx if self.images[i])) |
| pix = [load_pixels(p, self.size) for p in paths] |
| where = {p: k for k, p in enumerate(paths)} |
| image_of = [where[self.images[i]] if self.images[i] else -1 for i in idx] |
| use_sim = (self.rng.random(len(idx)) >= self.sim_dropout).tolist() if self.sim_dropout else None |
| b = mm_collate([self.encs[i] for i in idx], [self.labels[i] for i in idx], self.pad_id, image_of, pix, |
| [self.prompts[i] for i in idx], use_sim) |
| b["index"] = torch.tensor(list(idx)) |
| return b |
|
|
|
|
| def to_device(b: dict, device) -> dict: |
| out = {k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v) for k, v in b.items()} |
| out["k_max"] = int(b["n_options"].max()) |
| return out |
|
|
|
|
| @torch.no_grad() |
| def collect(model: MultimodalDecisionModel, ds: MMDataset, device, max_tokens: int = 8192) -> dict: |
| """Raw per-question outputs over a dataset (the train.collect contract, so calibrate.fit and |
| train.summarize work on it).""" |
| model.eval() |
| sampler = TokenBudgetBatchSampler(ds.lengths, max_tokens=max_tokens, max_batch=128, shuffle=False) |
| logits: list = [None] * len(ds) |
| esc = torch.zeros(len(ds)) |
| for idx in sampler: |
| b = to_device(ds.collate(idx), device) |
| with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"): |
| out = model(**b) |
| lg, e = out["logits"].float().cpu(), out["escalate_logit"].float().cpu() |
| for j, i in enumerate(idx): |
| logits[i] = lg[j, :len(ds.encs[i].markers)] |
| esc[idx] = e |
| return {"logits": logits, "escalate_logit": esc, "labels": torch.tensor(ds.labels), |
| "qtype": torch.tensor([e.qtype for e in ds.encs]), |
| "n_options": torch.tensor([len(e.markers) for e in ds.encs]), |
| "source": [r.source for r in ds.records]} |
|
|
|
|
| |
|
|
| class MultimodalSystemOne(SystemOne): |
| """``SystemOne`` plus images. ``predict(state_or_None, questions, image=...)`` returns the same |
| answer dicts as ``SystemOne.predict``. Calls without any image take the text path unchanged |
| (including ``long_state="chunk"``).""" |
|
|
| def __init__(self, model: MultimodalDecisionModel, tokenizer, calibration: dict | None, device: torch.device, |
| **kw): |
| super().__init__(model.dm, tokenizer, calibration, device, **kw) |
| self.mm = model.eval() |
| cfg = model.dm.cfg |
| self.mser = MMSerializer(tokenizer, cfg.max_len, cfg.head_max_len, model.mcfg.n_slots) |
|
|
| @classmethod |
| def load(cls, path: str, device: str | torch.device | None = None, vision: str | None = None, |
| **kw) -> "MultimodalSystemOne": |
| from transformers import AutoTokenizer |
| device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu")) |
| cal_path = os.path.join(path, "calibration.json") |
| cal = json.load(open(cal_path)) if os.path.exists(cal_path) else None |
| model = MultimodalDecisionModel.from_pretrained(path, device, vision) |
| return cls(model, AutoTokenizer.from_pretrained(path), cal, device, **kw) |
|
|
| def predict(self, state: str | None, questions: Questions, image: ImageLike = None, |
| calibrated: bool = True, use_sim: bool = True, **kw) -> dict: |
| return self.predict_batch([(state, questions, image)], calibrated=calibrated, use_sim=use_sim, **kw)[0] |
|
|
| @torch.inference_mode() |
| def predict_batch(self, items: Sequence[tuple], calibrated: bool = True, use_sim: bool = True, |
| **kw) -> list[dict]: |
| """items: (state, questions) or (state, questions, image); state may be None/"" with an image.""" |
| items = [tuple(it) + (None,) * (3 - len(it)) for it in items] |
| if all(img is None for _, _, img in items): |
| return super().predict_batch([(s or "", qs) for s, qs, _ in items], calibrated=calibrated, **kw) |
| t0 = time.perf_counter() |
| pixels, image_of_item = [], [] |
| for _, _, img in items: |
| if img is None: |
| image_of_item.append(-1) |
| else: |
| image_of_item.append(len(pixels)) |
| pixels.append(img if isinstance(img, torch.Tensor) else load_pixels(img, self.mm.mcfg.image_size)) |
| flat: list[tuple[int, str, Question]] = [] |
| for i, (_, qs, _) in enumerate(items): |
| for name, q in qs.items(): |
| q = q if isinstance(q, Question) else Question.from_public(q) |
| try: |
| q.validate() |
| except AssertionError as e: |
| raise ValueError(f"question {name!r}: {e}") from None |
| flat.append((i, name, q)) |
| encs = self.mser.encode_many([(items[i][0] or "", q, image_of_item[i] >= 0) for i, _, q in flat]) |
| prompts = [option_prompts(q) for _, _, q in flat] |
| logits: list[torch.Tensor] = [None] * len(encs) |
| esc = torch.zeros(len(encs)) |
| for idx in self.batches([len(e) for e in encs]) if encs else []: |
| b = mm_collate([encs[j] for j in idx], None, self.ser.pad, [image_of_item[flat[j][0]] for j in idx], |
| pixels, [prompts[j] for j in idx], [use_sim] * len(idx)) |
| b = to_device(b, self.device) |
| with torch.autocast(self.device.type, dtype=torch.bfloat16, enabled=self.device.type == "cuda"): |
| out = self.mm(**b) |
| lg, e = out["logits"].float().cpu(), out["escalate_logit"].float().cpu() |
| for r, j in enumerate(idx): |
| logits[j] = lg[r, :len(encs[j].markers)] |
| esc[idx] = e |
| cal = self.cal if calibrated else None |
| if cal: |
| esc = cal["escalate"]["a"] * esc + cal["escalate"]["b"] |
| p_wrong = 1 - torch.sigmoid(esc) |
| results = [{"answers": {}, "model": self.name, "latency_ms": 0.0} for _ in items] |
| for j, (i, name, q) in enumerate(flat): |
| p = (logits[j] / temperature(cal, q.type, q.n_options)).softmax(-1) |
| results[i]["answers"][name] = self._answer(q, p, float(p_wrong[j])) |
| ms = (time.perf_counter() - t0) * 1000 |
| for r in results: |
| r["latency_ms"] = ms |
| return results |
|
|
