Pluto-1 private release candidate: dual-encoder ensemble, calibrated Typed Decisions specialist
2778af7 verified | """Cross-encoder specialist: state + question + one option in a single encoder pass.""" | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import torch | |
| import torch.nn as nn | |
| from safetensors.torch import load_file, save_file | |
| from transformers import AutoConfig, AutoModel | |
| class PlutoCrossEncoder(nn.Module): | |
| def __init__(self, backbone_name: str | None = None, backbone_config=None, vocab_size: int | None = None): | |
| super().__init__() | |
| 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: | |
| self.backbone.resize_token_embeddings(vocab_size) | |
| h = self.backbone.config.hidden_size | |
| self.head = nn.Sequential( | |
| nn.Linear(h, h), | |
| nn.GELU(), | |
| nn.Dropout(0.1), | |
| nn.LayerNorm(h), | |
| nn.Linear(h, 1), | |
| ) | |
| def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: | |
| h = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state[:, 0] | |
| return self.head(h).squeeze(-1).float() | |
| def save_cross_encoder(model: PlutoCrossEncoder, 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-ce-1", | |
| "vocab_size": model.backbone.config.vocab_size, | |
| **extra, | |
| } | |
| with open(os.path.join(out_dir, "pluto_ce_config.json"), "w", encoding="utf-8") as f: | |
| json.dump(cfg, f, indent=2) | |
| def load_cross_encoder(path: str, device: str | torch.device = "cpu") -> tuple[PlutoCrossEncoder, dict]: | |
| cfg = json.load(open(os.path.join(path, "pluto_ce_config.json"), encoding="utf-8")) | |
| backbone_cfg = AutoConfig.for_model( | |
| **json.load(open(os.path.join(path, "backbone_config.json"), encoding="utf-8")) | |
| ) | |
| model = PlutoCrossEncoder(backbone_config=backbone_cfg, vocab_size=cfg["vocab_size"]) | |
| model.load_state_dict(load_file(os.path.join(path, "model.safetensors")), strict=True) | |
| model.to(device).eval() | |
| return model, cfg | |