Hub-native packaging: namespaced encoder shards, complete safetensors index, loaders
d9a201b verified | """Per-question specialist heads on top of the Pluto backbone.""" | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from safetensors.torch import load_file, save_file | |
| from transformers import AutoConfig, AutoModel, AutoTokenizer | |
| from .prompt import option_texts | |
| from .schema import KIND_TO_ID, Question, normalize_state, parse_questions | |
| class TaskSpecialist(nn.Module): | |
| def __init__( | |
| self, | |
| qids: dict[str, int], | |
| n_classes: dict[str, int], | |
| kinds: dict[str, str], | |
| backbone_name: str | None = None, | |
| backbone_config=None, | |
| vocab_size: int | None = None, | |
| ): | |
| super().__init__() | |
| if backbone_config is None: | |
| attn = "eager" if backbone_name and "deberta" in backbone_name.lower() else "sdpa" | |
| self.backbone = AutoModel.from_pretrained(backbone_name, attn_implementation=attn) | |
| else: | |
| attn = "eager" if getattr(backbone_config, "model_type", "") in ("deberta", "deberta-v2") else "sdpa" | |
| self.backbone = AutoModel.from_config(backbone_config, attn_implementation=attn) | |
| if vocab_size is not None: | |
| self.backbone.resize_token_embeddings(vocab_size) | |
| if next(self.backbone.parameters()).dtype == torch.float16: | |
| self.backbone = self.backbone.float() | |
| h = self.backbone.config.hidden_size | |
| self.qids = dict(qids) | |
| self.n_classes = dict(n_classes) | |
| self.kinds = dict(kinds) | |
| self.heads = nn.ModuleDict({q: nn.Linear(h, n) for q, n in n_classes.items()}) | |
| self.register_buffer("_qid_index", torch.tensor([qids[q] for q in n_classes], dtype=torch.long), persistent=False) | |
| def forward(self, input_ids, attention_mask, qid_ids: torch.Tensor) -> torch.Tensor: | |
| h = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state[:, 0] | |
| maxc = max(self.n_classes.values()) | |
| out = torch.full((h.shape[0], maxc), -1e9, device=h.device) | |
| for q, idx in self.qids.items(): | |
| m = qid_ids == idx | |
| if m.any(): | |
| out[m, : self.n_classes[q]] = self.heads[q](h[m]).float() | |
| return out.float() | |
| def context_text(state: str, q: Question) -> str: | |
| instruction = q.instructions | |
| if q.kind == "score": | |
| instruction += " Options are ordered from lowest to highest." | |
| opts = "\n".join(f"{i}. {t}" for i, t in enumerate(option_texts(q))) | |
| return ( | |
| f"State:\n{normalize_state(state)}\n\n" | |
| f"Question:\n{instruction}\n\n" | |
| f"Options:\n{opts}" | |
| ) | |
| def save_specialist(model: TaskSpecialist, tokenizer, out_dir: str, extra: dict) -> None: | |
| os.makedirs(out_dir, exist_ok=True) | |
| save_file({k: v.detach().cpu() for k, v in model.state_dict().items()}, os.path.join(out_dir, "model.safetensors")) | |
| tokenizer.save_pretrained(out_dir) | |
| model.backbone.config.to_json_file(os.path.join(out_dir, "backbone_config.json")) | |
| cfg = { | |
| "format": "pluto-task-1", | |
| "qids": model.qids, | |
| "n_classes": model.n_classes, | |
| "kinds": model.kinds, | |
| "vocab_size": model.backbone.config.vocab_size, | |
| **extra, | |
| } | |
| with open(os.path.join(out_dir, "pluto_task_config.json"), "w", encoding="utf-8") as f: | |
| json.dump(cfg, f, indent=2) | |
| def load_specialist(path: str, device: str | torch.device = "cpu") -> tuple[TaskSpecialist, dict]: | |
| cfg = json.load(open(os.path.join(path, "pluto_task_config.json"), encoding="utf-8")) | |
| backbone_cfg = AutoConfig.for_model( | |
| **json.load(open(os.path.join(path, "backbone_config.json"), encoding="utf-8")) | |
| ) | |
| model = TaskSpecialist( | |
| cfg["qids"], cfg["n_classes"], cfg["kinds"], backbone_config=backbone_cfg, vocab_size=cfg["vocab_size"] | |
| ) | |
| state = load_file(os.path.join(path, "model.safetensors")) | |
| prefix = cfg.get("tensor_prefix") | |
| if prefix and state and all(k.startswith(prefix) for k in state): | |
| state = {k[len(prefix):]: v for k, v in state.items()} | |
| elif not prefix and state: | |
| base_prefix = os.path.basename(os.path.normpath(path)) + "." | |
| if all(k.startswith(base_prefix) for k in state): | |
| state = {k[len(base_prefix):]: v for k, v in state.items()} | |
| missing, unexpected = model.load_state_dict(state, strict=False) | |
| if missing or unexpected: | |
| raise RuntimeError( | |
| f"checkpoint {path!r} did not load cleanly: missing={list(missing)[:4]} unexpected={list(unexpected)[:4]}" | |
| ) | |
| model.to(device).eval() | |
| return model, cfg | |
| class TaskSpecialistEngine: | |
| def from_pretrained(cls, pretrained_model_name_or_path: str, device: str | None = None, **kwargs) -> "TaskSpecialistEngine": | |
| return cls(pretrained_model_name_or_path, device=device) | |
| 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_specialist(path, device=self.device) | |
| self.tokenizer = AutoTokenizer.from_pretrained(path) | |
| self.model_name = self.config.get("model_name", "pluto-typed") | |
| self.temperatures = self.config.get("calibration", {}).get("temperature", {}) | |
| self.max_len = int(self.config.get("training", {}).get("max_len", 512)) | |
| def info(self) -> dict: | |
| return { | |
| "name": self.model_name, | |
| "path": self.path, | |
| "device": str(self.device), | |
| "qids": sorted(self.config["qids"]), | |
| "temperatures": self.temperatures, | |
| "params": int(sum(p.numel() for p in self.model.parameters())), | |
| } | |
| def _resolve_qid(self, name: str) -> str: | |
| if name in self.config["qids"]: | |
| return name | |
| for cq in self.config["qids"]: | |
| if name.endswith("_" + cq) or cq.endswith("_" + name): | |
| return cq | |
| raise ValueError(f"question {name!r} is not supported by {self.model_name}") | |
| def _forward_strings(self, texts: list[str], qids: list[str]) -> np.ndarray: | |
| enc = self.tokenizer( | |
| texts, | |
| padding=True, | |
| truncation=True, | |
| max_length=self.max_len, | |
| return_tensors="pt", | |
| ) | |
| input_ids = enc["input_ids"].to(self.device) | |
| mask = enc["attention_mask"].to(self.device) | |
| resolved = [self._resolve_qid(q) for q in qids] | |
| qid_ids = torch.tensor([self.config["qids"][q] for q in resolved], dtype=torch.long, device=self.device) | |
| ctx = ( | |
| torch.autocast(device_type="cuda", dtype=torch.bfloat16) | |
| if self.device.type == "cuda" | |
| else __import__("contextlib").nullcontext() | |
| ) | |
| with ctx: | |
| logits = self.model(input_ids, mask, qid_ids) | |
| logits = logits.float().cpu() | |
| temps = torch.tensor( | |
| [self.temperatures.get(k, 1.0) for k in ("choice", "noul", "score")], dtype=torch.float32 | |
| ) | |
| kinds = torch.tensor([KIND_TO_ID[self.config["kinds"][q]] for q in resolved]) | |
| probs = F.softmax(logits / temps[kinds].unsqueeze(-1), dim=-1) | |
| ns = torch.tensor([self.config["n_classes"][q] for q in resolved]) | |
| probs = probs.masked_fill(torch.arange(probs.shape[1]).unsqueeze(0) >= ns.unsqueeze(1), 0.0) | |
| probs = probs / probs.sum(-1, keepdim=True).clamp_min(1e-9) | |
| return probs.numpy() | |
| def predict(self, state, questions, max_len: int | None = None, alpha: float | None = None) -> dict: | |
| qs = parse_questions(questions) if not isinstance(questions, list) else questions | |
| texts, qids = [], [] | |
| for q in qs: | |
| if q.name not in self.config["qids"]: | |
| self._resolve_qid(q.name) | |
| texts.append(context_text(state, q)) | |
| qids.append(q.name) | |
| probs = self._forward_strings(texts, qids) | |
| answers = {} | |
| escalate = False | |
| for i, q in enumerate(qs): | |
| p = probs[i, : q.n_options] | |
| top = int(np.argmax(p)) | |
| conf = float(p[top]) | |
| item = { | |
| "type": q.kind, | |
| "confidence": conf, | |
| "accept_probability": conf, | |
| "escalate": conf < 0.5, | |
| "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: | |
| item["score"] = float(np.dot(np.arange(q.n_options), p)) | |
| item["score_normalized"] = item["score"] / max(1, q.n_options - 1) | |
| answers[q.name] = item | |
| escalate = escalate or item["escalate"] | |
| return { | |
| "answers": answers, | |
| "escalate": escalate, | |
| "routing": {"model": self.model_name, "device": str(self.device)}, | |
| } | |
| def predict_many(self, items, batch_size: int = 64, max_len: int | None = None) -> list[dict]: | |
| items = list(items) | |
| texts = [context_text(s, q) for s, q in items] | |
| qids = [q.name for _, q in items] | |
| out = [] | |
| for start in range(0, len(items), batch_size): | |
| probs = self._forward_strings(texts[start : start + batch_size], qids[start : start + batch_size]) | |
| for row, (_, q) in enumerate(items[start : start + batch_size]): | |
| p = probs[row, : q.n_options] | |
| top = int(np.argmax(p)) | |
| value = q.keys[top] if q.kind == "choice" else (float(p[1]) if q.kind == "noul" else float(np.dot(np.arange(q.n_options), p))) | |
| out.append( | |
| { | |
| "index": start + row, | |
| "top": top, | |
| "value": value, | |
| "probs": p.astype(np.float32), | |
| "confidence": float(p[top]), | |
| "accept_probability": float(p[top]), | |
| } | |
| ) | |
| return out | |