"""Inference engine for Pluto checkpoints.""" from __future__ import annotations import json import os import time from contextlib import nullcontext from typing import Any, Iterable import numpy as np import torch import torch.nn.functional as F from .calibration import ConformalGate from .model import load_pluto from .prompt import OPTION_MARKER, SPECIAL_TOKENS, option_texts, render_prompt from .schema import KIND_TO_ID, Question, normalize_state, parse_questions class PlutoEngine: def __init__(self, path: str, device: str | None = None): if device is None: device = "cuda" if torch.cuda.is_available() and torch.cuda.device_count() > 0 else "cpu" self.path = path self.device = torch.device(device) self.model, self.config = load_pluto(path, device=self.device) from transformers import AutoTokenizer self.tokenizer = AutoTokenizer.from_pretrained(path) missing = [t for t in SPECIAL_TOKENS if t not in self.tokenizer.get_vocab()] if missing: self.tokenizer.add_special_tokens({"additional_special_tokens": missing}) self.marker_id = self.tokenizer.convert_tokens_to_ids(OPTION_MARKER) self.model_name = self.config.get("model_name", "pluto") cal = self.config.get("calibration", {}) self.temperatures = { k: float(cal.get("temperature", {}).get(k, 1.0)) for k in KIND_TO_ID } self.gates = {} for k, v in cal.get("gates", {}).items(): try: self.gates[k] = ConformalGate.from_dict(v) except Exception: pass self.train_max_len = int(self.config.get("training", {}).get("max_len", 640)) @property def info(self) -> dict: return { "name": self.model_name, "path": self.path, "device": str(self.device), "train_max_len": self.train_max_len, "temperatures": self.temperatures, "gates": {k: v.to_dict() for k, v in self.gates.items()}, "params": int(sum(p.numel() for p in self.model.parameters())), } def _tokenize_one(self, state: str, q: Question, max_len: int) -> tuple[list[int], list[int]]: prompt = render_prompt(state, q) ids = self.tokenizer(prompt, add_special_tokens=True)["input_ids"] if len(ids) > max_len: s = normalize_state(state) best = None for frac in (0.75, 0.55, 0.38, 0.25, 0.12): keep = max(48, int(len(s) * frac)) head = keep // 2 tail = keep - head clipped = s[:head] + "\n...\n" + (s[-tail:] if tail else "") cand = render_prompt(clipped, q) cand_ids = self.tokenizer(cand, add_special_tokens=True)["input_ids"] best = cand_ids if len(cand_ids) <= max_len: ids = cand_ids break else: ids = best positions = [i for i, t in enumerate(ids) if t == self.marker_id] if len(positions) != q.n_options: raise ValueError( f"prompt for question {q.name!r} has {len(positions)} option markers, " f"expected {q.n_options}; shorten the state or the option descriptions" ) return ids, positions @torch.no_grad() def _forward(self, encoded: list[tuple[list[int], list[int], int]]) -> tuple[np.ndarray, np.ndarray]: max_len = max(len(ids) for ids, _, _ in encoded) max_opt = max(len(pos) for _, pos, _ in encoded) b = len(encoded) input_ids = torch.zeros((b, max_len), dtype=torch.long) attention = torch.zeros((b, max_len), dtype=torch.long) marker = torch.zeros((b, max_opt), dtype=torch.long) n_options = torch.zeros((b,), dtype=torch.long) kinds = torch.zeros((b,), dtype=torch.long) for i, (ids, pos, kind) in enumerate(encoded): input_ids[i, : len(ids)] = torch.tensor(ids, dtype=torch.long) attention[i, : len(ids)] = 1 marker[i, : len(pos)] = torch.tensor(pos, dtype=torch.long) n_options[i] = len(pos) kinds[i] = kind input_ids = input_ids.to(self.device) attention = attention.to(self.device) marker = marker.to(self.device) n_options = n_options.to(self.device) kinds = kinds.to(self.device) ctx = ( torch.autocast(device_type="cuda", dtype=torch.bfloat16) if self.device.type == "cuda" else nullcontext() ) with ctx: outputs = self.model(input_ids, attention, marker, n_options, kinds) logits = outputs["logits"].float().cpu() correct_prob = torch.sigmoid(outputs["correct_logit"].float()).cpu() temps = torch.tensor( [self.temperatures[k] for k in ("choice", "noul", "score")], dtype=torch.float32 )[kinds.cpu()] probs = F.softmax(logits / temps.unsqueeze(-1), dim=-1) opt_ids = torch.arange(max_opt).unsqueeze(0) probs = probs.masked_fill(opt_ids >= n_options.cpu().unsqueeze(1), 0.0) probs = probs / probs.sum(-1, keepdim=True).clamp_min(1e-9) return probs.numpy(), correct_prob.numpy() def predict( self, state: Any, questions: dict[str, Any] | list[Question], max_len: int | None = None, alpha: float | None = None, ) -> dict: qs = questions if isinstance(questions, list) else parse_questions(questions) state_s = normalize_state(state) max_len = max_len or self.train_max_len encoded = [] for q in qs: ids, pos = self._tokenize_one(state_s, q, max_len) encoded.append((ids, pos, KIND_TO_ID[q.kind])) t0 = time.perf_counter() probs, correct_prob = self._forward(encoded) latency = (time.perf_counter() - t0) * 1000.0 answers = {} escalate = False for i, q in enumerate(qs): p = probs[i, : q.n_options] top = int(np.argmax(p)) confidence = float(p[top]) gate = self.gates.get(q.kind) threshold = gate.threshold if gate else 0.5 accept_p = float(correct_prob[i]) item = { "type": q.kind, "confidence": confidence, "accept_probability": accept_p, "escalate": bool(accept_p < threshold), "probabilities": {k: float(v) for k, v in zip(q.keys, p)}, } if q.kind == "choice": item["choice"] = q.keys[top] elif q.kind == "noul": item["noul"] = float(p[1]) else: labels = option_texts(q) expected = float(np.dot(np.arange(q.n_options), p)) item["score"] = expected item["score_normalized"] = expected / max(1, q.n_options - 1) item["labels"] = labels answers[q.name] = item escalate = escalate or item["escalate"] return { "answers": answers, "escalate": escalate, "routing": { "model": self.model_name, "device": str(self.device), "latency_ms": round(latency, 3), }, } def predict_many( self, items: Iterable[tuple[Any, Question]], batch_size: int = 64, max_len: int | None = None, ) -> list[dict]: items = list(items) max_len = max_len or self.train_max_len enc = [] for state, q in items: ids, pos = self._tokenize_one(normalize_state(state), q, max_len) enc.append((ids, pos, KIND_TO_ID[q.kind])) order = sorted(range(len(enc)), key=lambda i: len(enc[i][0])) out: list[dict | None] = [None] * len(enc) for start in range(0, len(order), batch_size): chunk_idx = order[start : start + batch_size] chunk = [enc[i] for i in chunk_idx] probs, correct_prob = self._forward(chunk) for row, i in enumerate(chunk_idx): q = items[i][1] p = probs[row, : q.n_options] top = int(np.argmax(p)) if q.kind == "choice": value = q.keys[top] elif q.kind == "noul": value = float(p[1]) else: value = float(np.dot(np.arange(q.n_options), p)) out[i] = { "index": i, "top": top, "value": value, "probs": p.astype(np.float32), "confidence": float(p[top]), "accept_probability": float(correct_prob[row]), } return [o for o in out if o is not None]