Pluto-1 private release candidate: dual-encoder ensemble, calibrated Typed Decisions specialist
2778af7 verified | """Pluto: a compact typed-decision model. | |
| A bidirectional encoder reads the state, the typed question and every candidate | |
| option. Each option is scored at its own marker token. A small permutation- | |
| equivariant interaction stack compares options globally, then a scalar head | |
| returns one logit per option. A second head predicts the teacher's confidence so | |
| the runtime can abstain under a calibrated risk target. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import math | |
| import os | |
| from dataclasses import asdict, dataclass | |
| 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 | |
| from .schema import KIND_TO_ID | |
| IGNORE_INDEX = -1 | |
| class PlutoArch: | |
| hidden: int = 768 | |
| interaction_layers: int = 2 | |
| interaction_heads: int = 8 | |
| dropout: float = 0.05 | |
| max_options: int = 192 | |
| max_score_levels: int = 12 | |
| class PlutoForDecisions(nn.Module): | |
| def __init__( | |
| self, | |
| backbone_name: str | None = None, | |
| backbone_config: AutoConfig | None = None, | |
| arch: PlutoArch | None = None, | |
| vocab_size: int | None = None, | |
| freeze_layers: int = 0, | |
| ): | |
| super().__init__() | |
| self.arch = arch or PlutoArch() | |
| if backbone_config is None: | |
| self.backbone = AutoModel.from_pretrained( | |
| backbone_name, attn_implementation="sdpa" | |
| ) | |
| else: | |
| self.backbone = AutoModel.from_config( | |
| backbone_config, attn_implementation="sdpa" | |
| ) | |
| if vocab_size is not None and vocab_size != self.backbone.config.vocab_size: | |
| self.backbone.resize_token_embeddings(vocab_size) | |
| self.arch.hidden = self.backbone.config.hidden_size | |
| h = self.arch.hidden | |
| self.kind_embedding = nn.Embedding(len(KIND_TO_ID), h) | |
| self.score_position = nn.Embedding(self.arch.max_score_levels, h) | |
| self.option_in = nn.Sequential(nn.Linear(h, h), nn.LayerNorm(h)) | |
| self.global_in = nn.Sequential(nn.Linear(h, h), nn.LayerNorm(h)) | |
| layer = nn.TransformerEncoderLayer( | |
| d_model=h, | |
| nhead=self.arch.interaction_heads, | |
| dim_feedforward=2 * h, | |
| dropout=self.arch.dropout, | |
| activation="gelu", | |
| batch_first=True, | |
| norm_first=True, | |
| ) | |
| self.interaction = nn.TransformerEncoder( | |
| layer, num_layers=self.arch.interaction_layers, enable_nested_tensor=False | |
| ) | |
| self.option_norm = nn.LayerNorm(h) | |
| self.scorer = nn.Sequential( | |
| nn.Linear(h, h), | |
| nn.GELU(), | |
| nn.Dropout(self.arch.dropout), | |
| nn.LayerNorm(h), | |
| nn.Linear(h, 1), | |
| ) | |
| self.confidence_head = nn.Sequential( | |
| nn.Linear(h, h // 2), nn.GELU(), nn.Linear(h // 2, 1) | |
| ) | |
| if freeze_layers > 0: | |
| for p in self.backbone.embeddings.parameters(): | |
| p.requires_grad = False | |
| for layer_mod in self.backbone.layers[:freeze_layers]: | |
| for p in layer_mod.parameters(): | |
| p.requires_grad = False | |
| def forward( | |
| self, | |
| input_ids: torch.Tensor, | |
| attention_mask: torch.Tensor, | |
| marker_positions: torch.Tensor, | |
| n_options: torch.Tensor, | |
| kind_ids: torch.Tensor, | |
| ) -> dict[str, torch.Tensor]: | |
| hidden = self.backbone( | |
| input_ids=input_ids, attention_mask=attention_mask | |
| ).last_hidden_state | |
| b, nmax = marker_positions.shape | |
| h = hidden.shape[-1] | |
| safe_positions = marker_positions.clamp(min=0).unsqueeze(-1).expand(-1, -1, h) | |
| gathered = torch.gather(hidden, 1, safe_positions) | |
| option_ids = torch.arange(nmax, device=input_ids.device).unsqueeze(0) | |
| valid = option_ids < n_options.unsqueeze(1) | |
| kind_vec = self.kind_embedding(kind_ids) | |
| options = self.option_in(gathered) + kind_vec.unsqueeze(1) | |
| score_rows = kind_ids == KIND_TO_ID["score"] | |
| if score_rows.any(): | |
| pos = torch.arange(nmax, device=input_ids.device).clamp( | |
| max=self.arch.max_score_levels - 1 | |
| ) | |
| pos_vec = self.score_position(pos).unsqueeze(0) | |
| options = options + pos_vec * score_rows.view(-1, 1, 1).to(options.dtype) | |
| global_vec = self.global_in(hidden[:, 0]) + kind_vec | |
| sequence = torch.cat([global_vec.unsqueeze(1), options], dim=1) | |
| pad_mask = torch.cat( | |
| [ | |
| torch.zeros(b, 1, dtype=torch.bool, device=input_ids.device), | |
| ~valid, | |
| ], | |
| dim=1, | |
| ) | |
| sequence = self.interaction(sequence, src_key_padding_mask=pad_mask) | |
| logits = self.scorer(sequence[:, 1:]).squeeze(-1).float() | |
| logits = logits.masked_fill(~valid, -1e9) | |
| correct_logit = self.confidence_head(sequence[:, 0]).squeeze(-1).float() | |
| return { | |
| "logits": logits, | |
| "correct_logit": correct_logit, | |
| "valid": valid, | |
| } | |
| def decision_loss( | |
| outputs: dict[str, torch.Tensor], | |
| targets: torch.Tensor, | |
| kind_ids: torch.Tensor, | |
| weights: dict[str, float] | None = None, | |
| ) -> tuple[torch.Tensor, dict[str, float]]: | |
| weights = weights or {"kl": 1.0, "ce": 0.5, "rps": 0.5, "conf": 0.1} | |
| logits = outputs["logits"] | |
| valid = outputs["valid"] | |
| target = targets.clamp_min(0.0) | |
| target = target / target.sum(-1, keepdim=True).clamp_min(1e-9) | |
| log_probs = F.log_softmax(logits, dim=-1) | |
| row_valid = valid.any(-1) | |
| kind_w = torch.tensor([1.0, 0.8, 1.2], device=logits.device)[kind_ids] | |
| kl = -(target * log_probs).sum(-1) | |
| kl = (kl * kind_w)[row_valid].mean() | |
| argmax = target.argmax(-1) | |
| ce = F.nll_loss(log_probs, argmax, reduction="none") | |
| ce = (ce * kind_w)[row_valid].mean() | |
| score_rows = (kind_ids == KIND_TO_ID["score"]) & row_valid | |
| rps = torch.zeros((), device=logits.device) | |
| if score_rows.any(): | |
| probs = log_probs.exp() | |
| cs = (probs.cumsum(-1) - target.cumsum(-1))[:, :-1] | |
| row_rps = cs.pow(2).mean(-1) | |
| rps = row_rps[score_rows].mean() | |
| correct_target = (logits.argmax(-1) == argmax).float() | |
| conf = F.binary_cross_entropy_with_logits( | |
| outputs["correct_logit"], correct_target, reduction="none" | |
| ) | |
| conf = conf[row_valid].mean() | |
| total = ( | |
| weights["kl"] * kl | |
| + weights["ce"] * ce | |
| + weights["rps"] * rps | |
| + weights["conf"] * conf | |
| ) | |
| parts = { | |
| "kl": float(kl.detach()), | |
| "ce": float(ce.detach()), | |
| "rps": float(rps.detach()), | |
| "correct": float(conf.detach()), | |
| } | |
| return total, parts | |
| def save_pluto( | |
| model: PlutoForDecisions, | |
| tokenizer, | |
| out_dir: str, | |
| extra: dict, | |
| ) -> None: | |
| os.makedirs(out_dir, exist_ok=True) | |
| state = {k: v.detach().cpu() for k, v in model.state_dict().items()} | |
| save_file(state, 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")) | |
| config = { | |
| "format": "pluto-1", | |
| "arch": asdict(model.arch), | |
| "backbone_type": model.backbone.config.model_type, | |
| "vocab_size": model.backbone.config.vocab_size, | |
| **extra, | |
| } | |
| with open(os.path.join(out_dir, "pluto_config.json"), "w", encoding="utf-8") as f: | |
| json.dump(config, f, indent=2) | |
| def load_pluto(path: str, device: str | torch.device = "cpu") -> tuple[PlutoForDecisions, dict]: | |
| with open(os.path.join(path, "pluto_config.json"), "r", encoding="utf-8") as f: | |
| config = json.load(f) | |
| backbone_config = AutoConfig.for_model( | |
| **json.load(open(os.path.join(path, "backbone_config.json"), encoding="utf-8")) | |
| ) | |
| arch = PlutoArch(**config["arch"]) | |
| model = PlutoForDecisions( | |
| backbone_config=backbone_config, | |
| arch=arch, | |
| vocab_size=config["vocab_size"], | |
| ) | |
| state = load_file(os.path.join(path, "model.safetensors")) | |
| model.load_state_dict(state, strict=True) | |
| model.to(device) | |
| model.eval() | |
| return model, config | |