Upload modeling_char_bert_meter.py with huggingface_hub

modeling_char_bert_meter.py

@@ -0,0 +1,243 @@
+"""HF-Hub-compatible model for CharDiff-grc-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|<window)."""
+    B, T = seg_id.shape
+    same = seg_id[:, None, :] == seg_id[:, :, None]
+    if window and window > 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)