File size: 8,373 Bytes
c7893fa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 | """DecisionModel: ModernBERT backbone -> 2-layer decision head -> per-marker scorer -> softmax
per question, plus an escalate head predicting whether the argmax answer is correct."""
from __future__ import annotations
import json
import math
import os
from dataclasses import asdict, dataclass, field
import torch
import torch.nn.functional as F
from safetensors.torch import load_file, save_file
from torch import nn
from transformers import AutoConfig, AutoModel
from jevlike.serialize import TYPE_IDS
N_TYPES = len(TYPE_IDS)
N_ESC_FEATS = 3 + N_TYPES # max prob, normalized entropy, log n_options, one-hot type
@dataclass
class DecisionConfig:
backbone: str = "answerdotai/ModernBERT-base"
backbone_config: dict = field(default_factory=dict) # filled from the backbone; makes ckpts self-contained
attn_implementation: str = "sdpa"
head_layers: int = 2
head_heads: int = 12
head_ff: int = 2048
dropout: float = 0.1
esc_hidden: int = 256
max_len: int = 512
head_max_len: int = 192
escalate_threshold: float = 0.5 # escalate when escalate_prob = P(argmax wrong) > threshold
name: str = "jevlike-base"
@classmethod
def load(cls, path: str) -> "DecisionConfig":
with open(os.path.join(path, "config.json")) as f:
d = json.load(f)
return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})
class DecisionModel(nn.Module):
def __init__(self, cfg: DecisionConfig, backbone: nn.Module | None = None):
super().__init__()
if backbone is None: # architecture only (weights come from a checkpoint)
bc = dict(cfg.backbone_config)
backbone = AutoModel.from_config(AutoConfig.for_model(bc.pop("model_type"), **bc),
attn_implementation=cfg.attn_implementation)
self.cfg, self.backbone = cfg, backbone
cfg.backbone_config = backbone.config.to_dict()
D = backbone.config.hidden_size
if D % cfg.head_heads: # e.g. ModernBERT-large (1024): use 64-dim heads
cfg.head_heads = D // 64
self.head = nn.ModuleList(nn.TransformerEncoderLayer(
D, cfg.head_heads, cfg.head_ff, cfg.dropout, activation="gelu", batch_first=True, norm_first=True)
for _ in range(cfg.head_layers))
self.head_norm = nn.LayerNorm(D)
self.scorer = nn.Sequential(nn.Linear(D, D), nn.GELU(), nn.Dropout(cfg.dropout), nn.Linear(D, 1))
self.escalate = nn.Sequential(nn.Linear(D + N_ESC_FEATS, cfg.esc_hidden), nn.GELU(),
nn.Dropout(cfg.dropout), nn.Linear(cfg.esc_hidden, 1))
# skip path: escalate logit starts at logit(max prob), i.e. "P(correct) = own confidence"
self.esc_skip = nn.Parameter(torch.tensor([1.0, 0.0]))
for lin in (self.scorer[-1], self.escalate[-1]): # start near-uniform / uninformed
nn.init.normal_(lin.weight, std=0.02)
nn.init.zeros_(lin.bias)
@classmethod
def from_backbone(cls, cfg: DecisionConfig) -> "DecisionModel":
"""New model with pretrained backbone weights and freshly initialized heads."""
bb = AutoModel.from_pretrained(cfg.backbone, attn_implementation=cfg.attn_implementation)
return cls(cfg, bb)
def head_parameters(self):
return [p for n, p in self.named_parameters() if not n.startswith("backbone.")]
def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, marker_idx: torch.Tensor,
marker_seg: torch.Tensor, marker_rank: torch.Tensor, n_options: torch.Tensor,
qtype: torch.Tensor, k_max: int | None = None, inputs_embeds: torch.Tensor | None = None,
**_) -> dict[str, torch.Tensor]:
"""Returns ``logits`` [B, K] (-inf beyond each question's n_options; one question per
row) and ``escalate_logit`` [B] (logit of P(argmax correct)). The escalate head sees the
[CLS] state plus detached answer-distribution features, so it cannot move the answers.
``inputs_embeds`` [B, T, D] (optional) replaces the token-embedding lookup of ``input_ids``
(used by the multimodal prototype, jevlike/mm.py, to splice in image tokens)."""
if inputs_embeds is not None:
h = self.backbone(inputs_embeds=inputs_embeds, attention_mask=attention_mask).last_hidden_state
else:
h = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
pad = attention_mask == 0
for layer in self.head:
h = layer(h, src_key_padding_mask=pad)
h = self.head_norm(h)
B, T, D = h.shape
marker_logit = self.scorer(h.reshape(B * T, D)[marker_idx]).squeeze(-1).float()
K = k_max or int(n_options.max())
logits = torch.full((B, K), -math.inf, device=h.device).index_put((marker_seg, marker_rank), marker_logit)
with torch.no_grad(): # answer-distribution features, detached
logp = logits.log_softmax(-1)
p = logp.exp()
ent = -(p * logp.masked_fill(p == 0, 0)).sum(-1) / n_options.float().log()
maxp = p.max(-1).values
feats = torch.cat([maxp[:, None], ent[:, None], (n_options.float().log() / math.log(256))[:, None],
F.one_hot(qtype, N_TYPES).float()], -1)
conf_logit = torch.logit(maxp, eps=1e-6)
esc = self.escalate(torch.cat([h[:, 0].float(), feats], -1)).squeeze(-1).float()
esc = esc + self.esc_skip[0] * conf_logit + self.esc_skip[1]
return {"logits": logits, "escalate_logit": esc}
# ------------------------------------------------------------ io
def save_pretrained(self, path: str, tokenizer=None) -> None:
os.makedirs(path, exist_ok=True)
with open(os.path.join(path, "config.json"), "w") as f:
json.dump(asdict(self.cfg), f, indent=1, default=str)
save_file({k: v.detach().contiguous().cpu() for k, v in self.state_dict().items()},
os.path.join(path, "model.safetensors"))
if tokenizer is not None:
tokenizer.save_pretrained(path)
@classmethod
def from_pretrained(cls, path: str, device: str | torch.device = "cpu") -> "DecisionModel":
model = cls(DecisionConfig.load(path))
model.load_state_dict(load_file(os.path.join(path, "model.safetensors"), device=str(device)))
return model.to(device)
# ---------------------------------------------------------------- losses & metrics
def compute_loss(out: dict[str, torch.Tensor], batch: dict[str, torch.Tensor], rps_weight: float = 1.0,
esc_weight: float = 0.2) -> tuple[torch.Tensor, dict[str, float]]:
"""Per-question loss = CE + rps_weight * RPS (score questions) + esc_weight * BCE(escalate).
Returns (sum over questions, stats); callers normalize by the number of questions per step.
No label smoothing: all terms are strictly proper scoring rules.
"""
logits, labels, n = out["logits"], batch["labels"], batch["n_options"]
logp = logits.log_softmax(-1)
nll = -logp.gather(1, labels[:, None]).squeeze(1)
cdf_err = logp.exp().cumsum(-1) - F.one_hot(labels, logits.shape[1]).float().cumsum(-1)
rps = cdf_err.pow(2).sum(-1) / (n - 1).clamp(min=1)
is_score = (batch["qtype"] == TYPE_IDS["score"]).float()
correct = (logits.detach().argmax(-1) == labels).float()
bce = F.binary_cross_entropy_with_logits(out["escalate_logit"], correct, reduction="none")
per = nll + rps_weight * rps * is_score + esc_weight * bce
stats = {"nll": nll.sum().item(), "rps": (rps * is_score).sum().item(), "n_score": is_score.sum().item(),
"bce": bce.sum().item(), "correct": correct.sum().item(), "n": float(len(labels))}
return per.sum(), stats
def ece(conf: torch.Tensor, correct: torch.Tensor, n_bins: int = 15) -> float:
"""Expected calibration error of top-1 confidence (equal-width bins)."""
if len(conf) == 0:
return float("nan")
bins = (conf.float() * n_bins).long().clamp(0, n_bins - 1)
tot = torch.zeros(n_bins).index_add_(0, bins, torch.ones_like(conf))
c = torch.zeros(n_bins).index_add_(0, bins, conf.float())
a = torch.zeros(n_bins).index_add_(0, bins, correct.float())
return ((c - a).abs().sum() / len(conf)).item()
|