File size: 7,816 Bytes
0415692 | 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 | """ModernBERT-architecture encoder (Ettin weights), written from scratch.
Attention is controlled by per-token vectors instead of a precomputed mask:
position_ids: RoPE positions (branches restart after the state).
segment_ids: 0 = state, k > 0 = question branch k, -1 = padding.
option_ids: 0 = not an option, m > 0 = option m of its question.
A token attends to state tokens, its question's own tokens, and (for option tokens)
its own option. Local layers additionally restrict attention to
|pos_i - pos_j| <= window // 2, measured in position ids (not sequence index) so
packed and separate requests match.
Comparison mode lets each question and its options attend to the whole question branch.
"""
import json
from dataclasses import dataclass
import torch
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from torch import nn
@dataclass
class Config:
vocab_size: int
hidden_size: int
intermediate_size: int
num_hidden_layers: int
num_attention_heads: int
global_attn_every_n_layers: int = 3
local_attention: int = 128
global_rope_theta: float = 160000.0
local_rope_theta: float = 10000.0
norm_eps: float = 1e-5
pad_token_id: int = 50283
cls_token_id: int = 50281
sep_token_id: int = 50282
model_type: str = "modernbert"
compare_options: bool = False
date_facts: bool = False # an input-format setting: facts.date_facts runs on every state (model.pack_request)
@classmethod
def from_hf(cls, d: dict) -> "Config":
"""Reads a Hugging Face config or a run's saved config, ignoring fields this class does not define."""
return cls(**{k: d[k] for k in cls.__dataclass_fields__ if k in d})
def rope(x: torch.Tensor, pos: torch.Tensor, theta: float) -> torch.Tensor:
"""x: [B, H, L, D], pos: [B, L]. Half-split rotation (rotate_half convention)."""
d = x.shape[-1]
inv_freq = 1.0 / (theta ** (torch.arange(0, d, 2, device=x.device, dtype=torch.float32) / d))
freqs = pos[:, None, :, None].float() * inv_freq # [B, 1, L, D/2]
cos, sin = freqs.cos().to(x.dtype), freqs.sin().to(x.dtype)
x1, x2 = x[..., : d // 2], x[..., d // 2 :]
return torch.cat((x1 * cos - x2 * sin, x2 * cos + x1 * sin), dim=-1)
class Layer(nn.Module):
def __init__(self, c: Config, idx: int):
super().__init__()
self.is_global = idx % c.global_attn_every_n_layers == 0
self.theta = c.global_rope_theta if self.is_global else c.local_rope_theta
self.heads = c.num_attention_heads
self.attn_norm = nn.LayerNorm(c.hidden_size, eps=c.norm_eps, bias=False) if idx else nn.Identity()
self.Wqkv = nn.Linear(c.hidden_size, 3 * c.hidden_size, bias=False)
self.Wo = nn.Linear(c.hidden_size, c.hidden_size, bias=False)
self.mlp_norm = nn.LayerNorm(c.hidden_size, eps=c.norm_eps, bias=False)
self.Wi = nn.Linear(c.hidden_size, 2 * c.intermediate_size, bias=False)
self.Wo_mlp = nn.Linear(c.intermediate_size, c.hidden_size, bias=False)
def forward(self, h, pos, global_mask, local_mask, prefix=None):
"""prefix: cached (k, v) of earlier tokens, prepended to this call's keys and values."""
B, L, _ = h.shape
q, k, v = self.Wqkv(self.attn_norm(h)).view(B, L, 3, self.heads, -1).permute(2, 0, 3, 1, 4)
q, k = rope(q, pos, self.theta), rope(k, pos, self.theta)
kv = (k, v)
if prefix is not None:
k, v = torch.cat([prefix[0], k], 2), torch.cat([prefix[1], v], 2)
mask = global_mask if self.is_global else local_mask
a = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
h = h + self.Wo(a.transpose(1, 2).reshape(B, L, -1))
x, gate = self.Wi(self.mlp_norm(h)).chunk(2, dim=-1)
return h + self.Wo_mlp(F.gelu(x) * gate), kv
class Encoder(nn.Module):
def __init__(self, c: Config):
super().__init__()
self.config = c
self.tok_embeddings = nn.Embedding(c.vocab_size, c.hidden_size)
self.emb_norm = nn.LayerNorm(c.hidden_size, eps=c.norm_eps, bias=False)
self.layers = nn.ModuleList(Layer(c, i) for i in range(c.num_hidden_layers))
self.final_norm = nn.LayerNorm(c.hidden_size, eps=c.norm_eps, bias=False)
@staticmethod
def masks(pos, seg, opt, window, pos_k=None, seg_k=None, opt_k=None, *, compare_options=False):
"""Boolean [B, 1, Lq, Lk] masks, True = may attend. Global and local variants.
Keys default to the queries; pass *_k to attend over a longer key sequence (cached state + queries).
The masks are dense, so their memory grows with Lq * Lk.
"""
pos_k, seg_k, opt_k = (x if x is not None else y for x, y in ((pos_k, pos), (seg_k, seg), (opt_k, opt)))
si, sj, oi, oj = seg[:, :, None], seg_k[:, None, :], opt[:, :, None], opt_k[:, None, :]
within = (sj == si) & ((si > 0) if compare_options else ((oj == 0) | (oj == oi)))
allowed = (sj != -1) & ((sj == 0) | within)
local = allowed & ((pos[:, :, None] - pos_k[:, None, :]).abs() <= window // 2)
return allowed[:, None], local[:, None]
def forward(self, input_ids, position_ids=None, segment_ids=None, option_ids=None, inputs_embeds=None):
B, L = input_ids.shape
if position_ids is None:
position_ids = torch.arange(L, device=input_ids.device).expand(B, L)
if segment_ids is None:
segment_ids = torch.zeros_like(input_ids)
if option_ids is None:
option_ids = torch.zeros_like(input_ids)
emb = self.tok_embeddings(input_ids) if inputs_embeds is None else inputs_embeds
return self.run(emb, position_ids, segment_ids, option_ids)[0]
def run(self, emb, pos, seg, opt, cache=None):
"""Encode embeddings, optionally on top of a cached prefix. Returns (hidden, cache for these tokens).
A cache holds each layer's rotated keys and values plus the pos/seg/opt ids of its tokens. The state
never attends to questions, so a state encoded once can serve any number of later questions.
"""
if cache is None:
g, l = self.masks(pos, seg, opt, self.config.local_attention, compare_options=self.config.compare_options)
else:
g, l = self.masks(pos, seg, opt, self.config.local_attention,
torch.cat([cache["pos"], pos], 1), torch.cat([cache["seg"], seg], 1), torch.cat([cache["opt"], opt], 1),
compare_options=self.config.compare_options)
h, kvs = self.emb_norm(emb), []
for i, layer in enumerate(self.layers):
h, kv = layer(h, pos, g, l, None if cache is None else cache["kv"][i])
kvs.append(kv)
return self.final_norm(h), {"kv": kvs, "pos": pos, "seg": seg, "opt": opt}
@classmethod
def from_pretrained(cls, repo: str) -> "Encoder":
with open(hf_hub_download(repo, "config.json")) as f:
model = cls(Config.from_hf(json.load(f)))
sd = torch.load(hf_hub_download(repo, "pytorch_model.bin"), map_location="cpu", weights_only=True)
rename = {
"model.embeddings.tok_embeddings.": "tok_embeddings.",
"model.embeddings.norm.": "emb_norm.",
"model.final_norm.": "final_norm.",
".attn.Wqkv.": ".Wqkv.",
".attn.Wo.": ".Wo.",
".mlp.Wi.": ".Wi.",
".mlp.Wo.": ".Wo_mlp.",
"model.layers.": "layers.",
}
out = {}
for k, v in sd.items():
if not k.startswith("model."): # MLM head, unused
continue
for a, b in rename.items():
k = k.replace(a, b)
out[k] = v
model.load_state_dict(out, strict=True)
return model.eval()
|