"""HF-Hub-compatible model for Stoicheia-meter (macronization + metrical scansion). Self-contained: vendors the same transformer primitives as modeling_char_bert.py, plus the fine-tune-only additions meter/model.py::MeterModel and meter/backbone.py:: CharBertWithHidden make on top of the plain backbone: - a zero-init `cap_emb` capitalization input embedding (fine-tune-only; base pretraining treats capitalization as output-only) - an ELMo-style learned scalar mix over every block's output (+ the final normed hidden state) instead of using only the last layer - two extra per-letter heads: `head_mac` (2-way: long/short vowel quantity) and `head_scan` (4-way: none/heavy/light/verse-final syllable weight) The submodule layout (`self.encoder.*` for the frozen backbone, `head_mac`/ `head_scan`/`mix_w` at the top level) matches meter.model.MeterModel's real state dict exactly -- converted checkpoints load with strict=True and no key remapping. """ from __future__ import annotations from dataclasses import dataclass from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel from transformers.modeling_outputs import ModelOutput from .configuration_char_bert_meter import CharBertMeterConfig class RMSNorm(nn.Module): def __init__(self, d, eps=1e-6): super().__init__() self.w = nn.Parameter(torch.ones(d)) self.eps = eps def forward(self, x): x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) return x * self.w class RoPE(nn.Module): def __init__(self, dim, base=10000.0): super().__init__() self.dim = dim self.base = base def cos_sin(self, pos): # Recomputed on every call rather than cached in a registered buffer: a # persistent=False buffer is never covered by the checkpoint's state dict, # so it depends entirely on __init__-time materialization -- which some # transformers versions' meta-device/low_cpu_mem_usage loading path can # skip, silently leaving this tensor uninitialized. Recomputing here is # immune to that regardless of how the model was constructed/loaded. inv = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=pos.device).float() / self.dim)) f = torch.outer(pos.float(), inv) emb = torch.cat([f, f], -1) return emb.cos(), emb.sin() def _rotate_half(x): d = x.shape[-1] // 2 return torch.cat([-x[..., d:], x[..., :d]], -1) def apply_rope(q, k, cos, sin): cos = cos[None, None] sin = sin[None, None] return q * cos + _rotate_half(q) * sin, k * cos + _rotate_half(k) * sin class Attention(nn.Module): def __init__(self, d, n_heads, rope: RoPE, qk_norm=False): super().__init__() self.h = n_heads self.dh = d // n_heads self.qkv = nn.Linear(d, 3 * d, bias=False) self.o = nn.Linear(d, d, bias=False) self.rope = rope self.qk_norm = qk_norm if qk_norm: self.q_norm = RMSNorm(self.dh) self.k_norm = RMSNorm(self.dh) def forward(self, x, pos, attn_mask): B, T, D = x.shape qkv = self.qkv(x).view(B, T, 3, self.h, self.dh).permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] if self.qk_norm: q, k = self.q_norm(q), self.k_norm(k) cos, sin = self.rope.cos_sin(pos) cos, sin = cos.to(x.dtype), sin.to(x.dtype) q, k = apply_rope(q, k, cos, sin) out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) out = out.transpose(1, 2).reshape(B, T, D) return self.o(out) class GeGLU(nn.Module): def __init__(self, d, mult=8 / 3): super().__init__() hidden = int(d * mult) hidden = (hidden + 63) // 64 * 64 self.wi = nn.Linear(d, 2 * hidden, bias=False) self.wo = nn.Linear(hidden, d, bias=False) def forward(self, x): a, b = self.wi(x).chunk(2, -1) return self.wo(F.gelu(a) * b) class Block(nn.Module): def __init__(self, d, n_heads, rope, window=0, qk_norm=False): super().__init__() self.n1 = RMSNorm(d) self.attn = Attention(d, n_heads, rope, qk_norm=qk_norm) self.n2 = RMSNorm(d) self.mlp = GeGLU(d) self.window = window # 0 = global; >0 = local sliding window (characters) def forward(self, x, pos, base_mask): x = x + self.attn(self.n1(x), pos, base_mask) x = x + self.mlp(self.n2(x)) return x def build_attn_mask(seg_id, window, device, dtype): """Additive mask (B,1,T,T): same-segment AND (window==0 or |i-j| 0: idx = torch.arange(T, device=device) near = (idx[None, :] - idx[:, None]).abs() < window same = same & near[None] mask = torch.zeros(B, 1, T, T, dtype=dtype, device=device) mask.masked_fill_(~same[:, None], float("-inf")) return mask class _MeterEncoder(nn.Module): """Same submodule names/shapes as a plain CharBertEncoder (so a pretraining backbone loads into it with no remapping), plus an optional zero-init cap_emb and per-layer output collection for the scalar mix -- mirrors meter.backbone.CharBertWithHidden exactly.""" def __init__(self, config: CharBertMeterConfig): super().__init__() self.e_char = nn.Embedding(config.n_char_ids, config.d_model) self.e_bnd = nn.Embedding(config.n_boundary, config.d_model) self.e_dia = nn.Embedding(config.n_dia, config.d_model) self.e_punct = nn.Embedding(config.n_punct, config.d_model) if config.use_cap: self.cap_emb = nn.Embedding(2, config.d_model) rope = RoPE(config.d_model // config.n_heads) blocks = [] for i in range(config.depth): win = 0 if i % 4 == 3 else config.char_window # 3 local : 1 global blocks.append(Block(config.d_model, config.n_heads, rope, window=win, qk_norm=config.qk_norm)) self.blocks = nn.ModuleList(blocks) self.norm_out = RMSNorm(config.d_model) # frozen pretraining output heads: not used by the meter heads, but part of # the backbone's real state dict (kept so a pretraining checkpoint -- or this # converted meter checkpoint -- loads with strict=True) self.head_char = nn.Linear(config.d_model, config.n_char_ids, bias=False) self.head_bnd = nn.Linear(config.d_model, 3, bias=False) self.head_dia = nn.Linear(config.d_model, 48, bias=False) self.head_cap = nn.Linear(config.d_model, 2, bias=False) self.head_punct = nn.Linear(config.d_model, 6, bias=False) self.cfg = config def forward(self, input_ids, boundary, dia, punct, cap=None, seg_id=None, collect_layers=False): cfg = self.cfg B, T = input_ids.shape pos = torch.arange(T, device=input_ids.device) seg = seg_id if seg_id is not None else torch.zeros(B, T, dtype=torch.long, device=input_ids.device) x = self.e_char(input_ids) + self.e_bnd(boundary) + self.e_dia(dia) + self.e_punct(punct) cap_emb = getattr(self, "cap_emb", None) if cap_emb is not None and cap is not None: x = x + cap_emb(cap) attn_mask = build_attn_mask(seg, cfg.char_window, input_ids.device, x.dtype) glob_mask = build_attn_mask(seg, 0, input_ids.device, x.dtype) layers = [] for blk in self.blocks: m = glob_mask if blk.window == 0 else attn_mask x = blk(x, pos, m) if collect_layers: layers.append(x) x = self.norm_out(x) return layers, x @dataclass class CharBertMeterOutput(ModelOutput): mac: torch.FloatTensor = None scan: torch.FloatTensor = None class CharBertMeterModel(PreTrainedModel): config_class = CharBertMeterConfig def __init__(self, config: CharBertMeterConfig): super().__init__(config) self.encoder = _MeterEncoder(config) self.head_mac = nn.Linear(config.d_model, 2, bias=False) # 0=long, 1=short self.head_scan = nn.Linear(config.d_model, 4, bias=False) # 0=none,1=heavy,2=light,3=verse-final if config.scalar_mix: self.mix_w = nn.Parameter(torch.zeros(config.depth + 1)) self.post_init() def _init_weights(self, module): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, std=0.02) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, std=0.02) def forward( self, input_ids: torch.LongTensor, boundary: torch.LongTensor, dia: torch.LongTensor, punct: torch.LongTensor, cap: Optional[torch.LongTensor] = None, seg_id: Optional[torch.LongTensor] = None, return_dict: bool = True, **kwargs, ): collect = bool(self.config.scalar_mix) layers, x = self.encoder(input_ids, boundary, dia, punct, cap=cap, seg_id=seg_id, collect_layers=collect) if self.config.scalar_mix: h = torch.stack(layers + [x]) # (L+1, B, T, D) mix = torch.softmax(self.mix_w, 0) h = torch.einsum("l,lbtd->btd", mix.to(h.dtype), h) else: h = x mac = self.head_mac(h) scan = self.head_scan(h) if not return_dict: return (mac, scan) return CharBertMeterOutput(mac=mac, scan=scan)