Pluto-1 private release candidate: dual-encoder ensemble, calibrated Typed Decisions specialist
2778af7 verified | """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)) | |
| 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 | |
| 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] | |