Publish safetensors weights, config and model card

README.md

@@ -0,0 +1,55 @@
+---
+license: apache-2.0
+language:
+- grc
+library_name: transformers
+tags:
+- ancient-greek
+- classical-philology
+- character-level
+- masked-diffusion
+- macronization
+- metrical-scansion
+pipeline_tag: token-classification
+---
+
+# Stoicheia -- macronization and metrical scansion (joint)
+
+**Stoicheia** is a 405M-parameter character-level masked-diffusion encoder for Ancient Greek
+(`d_model` 1024, depth 32, banded attention: three of every four blocks attend within a
+256-character window, the fourth globally). Its input is factored into five aligned planes --
+letters, word/sentence boundaries, diacritics, capitalization, punctuation -- each of which can
+be masked independently to an explicit *unknown* state at inference. That is what lets one model
+read an edited text, *scriptio continua*, and a lacuna of unknown length without changing
+anything but its input.
+
+Anonymous release accompanying a paper under review.
+
+Two per-letter heads on `Stoicheia-doc_clean`, trained jointly: **macronization** (long versus
+short at ambiguous bare α/ι/υ -- Greek orthography never marks vowel length) and **scansion**
+(none / heavy-end / light-end / verse-end). An optional Viterbi decoder constrains the scansion
+output to valid paths through a set of metre automata.
+
+Use this checkpoint when you want both tasks from one model; use `Stoicheia-macronizer` when you
+want vowel length alone, where a dedicated head does better.
+
+## Usage
+
+```python
+import torch
+from transformers import AutoModel
+from huggingface_hub import hf_hub_download
+
+REPO = "Ericu950/Stoicheia-meter"
+model = AutoModel.from_pretrained(REPO, trust_remote_code=True).eval()
+
+hf_hub_download(repo_id=REPO, filename="processing_char_bert_meter.py", local_dir=".")
+from processing_char_bert_meter import CharBertMeterProcessor
+
+proc = CharBertMeterProcessor()
+batch = proc("ἄνδρα μοι ἔννεπε, μοῦσα, πολύτροπον, ὃς μάλα πολλὰ")
+with torch.no_grad():
+    out = model(**{k: v for k, v in batch.items() if not k.startswith("_")})
+print(proc.decode_macronization(out, batch))   # _ long, ^ short
+print(proc.decode_scansion(out, batch))        # [heavy] {light}
+```

config.json

@@ -0,0 +1,23 @@
+{
+  "n_alpha": 24,
+  "mask_id": 24,
+  "blank_id": 25,
+  "pad_id": 26,
+  "n_char_ids": 27,
+  "n_boundary": 4,
+  "n_dia": 49,
+  "n_punct": 7,
+  "d_model": 1024,
+  "n_heads": 16,
+  "depth": 32,
+  "char_window": 256,
+  "attn_impl": "sdpa",
+  "qk_norm": true,
+  "use_cap": true,
+  "scalar_mix": true,
+  "model_type": "char_bert_meter",
+  "auto_map": {
+    "AutoConfig": "configuration_char_bert_meter.CharBertMeterConfig",
+    "AutoModel": "modeling_char_bert_meter.CharBertMeterModel"
+  }
+}

configuration_char_bert_meter.py

@@ -0,0 +1,50 @@
+"""HF-Hub-compatible config for Stoicheia-meter (macronization + metrical scansion).
+
+Same backbone hyperparameters as CharBertConfig (this wraps a Stoicheia backbone
+fine-tuned with two extra per-letter heads), plus the two fields that change the
+model's *shape* (use_cap, scalar_mix) -- head_dropout/w_mac/w_scan/class weights are
+training-only and irrelevant to inference, so they aren't part of this config.
+"""
+from transformers import PretrainedConfig
+
+
+class CharBertMeterConfig(PretrainedConfig):
+    model_type = "char_bert_meter"
+
+    def __init__(
+        self,
+        n_alpha: int = 24,
+        mask_id: int = 24,
+        blank_id: int = 25,
+        pad_id: int = 26,
+        n_char_ids: int = 27,
+        n_boundary: int = 4,
+        n_dia: int = 49,
+        n_punct: int = 7,
+        d_model: int = 1024,
+        n_heads: int = 16,
+        depth: int = 32,
+        char_window: int = 256,
+        attn_impl: str = "sdpa",
+        qk_norm: bool = True,
+        use_cap: bool = True,
+        scalar_mix: bool = True,
+        **kwargs,
+    ):
+        self.n_alpha = n_alpha
+        self.mask_id = mask_id
+        self.blank_id = blank_id
+        self.pad_id = pad_id
+        self.n_char_ids = n_char_ids
+        self.n_boundary = n_boundary
+        self.n_dia = n_dia
+        self.n_punct = n_punct
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.depth = depth
+        self.char_window = char_window
+        self.attn_impl = attn_impl
+        self.qk_norm = qk_norm
+        self.use_cap = use_cap
+        self.scalar_mix = scalar_mix
+        super().__init__(**kwargs)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d17f0518e64281d3b8ae18f48922ade29aabd7581a860a3a24909f6a7149eac1
+size 1620052996

modeling_char_bert_meter.py

@@ -0,0 +1,243 @@
+"""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|<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)

processing_char_bert_meter.py

@@ -0,0 +1,299 @@
+"""HF-Hub-compatible processor for Stoicheia-meter: text <-> the model's five
+input planes (chars/boundary/dia/punct/cap -- capitalization is a real input here,
+unlike the base pretrained model, where it is output-only), plus decode helpers for
+the two production tasks: macronization (vowel-length marks on ambiguous dichrona)
+and metrical scansion (heavy/light/verse-final syllable bracketing).
+
+Mark conventions and the ambiguous-dichrona rule are ported verbatim from
+meter/marks.py (the training-side projection code) so a converted checkpoint's
+predictions decode identically to the original project's `meter.predict` CLI.
+"""
+from __future__ import annotations
+
+import json
+import unicodedata
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import torch
+
+MASK, BLANK, PAD = 24, 25, 26
+UNK_BND, UNK_DIA, UNK_PUNCT = 3, 48, 6
+
+ALPHABET = "αβγδεζηθικλμνξοπρστυφχψω"
+LETTER_IDS = {c: i for i, c in enumerate(ALPHABET)}
+ID2LETTER = np.array(list(ALPHABET))
+
+_EXTRA_BASE = {
+    "ς": "σ", "ϲ": "σ", "Ϲ": "σ", "ϐ": "β", "ϑ": "θ", "ϕ": "φ", "ϰ": "κ", "ϱ": "ρ", "ϖ": "π",
+}
+_MARK_MAP = {
+    0x0301: "acute", 0x0341: "acute", 0x0300: "grave", 0x0340: "grave",
+    0x0342: "circ", 0x0302: "circ", 0x0313: "smooth", 0x0343: "smooth",
+    0x0314: "rough", 0x0345: "iota", 0x0308: "diaer",
+}
+_ACC = {"acute": 1, "grave": 2, "circ": 3}
+_BR = {"smooth": 1, "rough": 2}
+
+
+def _pack_dia(acc, br, iota, diaer):
+    return ((acc * 3 + br) * 2 + iota) * 2 + diaer
+
+
+def _unpack_dia(d):
+    diaer = d % 2; d //= 2
+    iota = d % 2; d //= 2
+    br = d % 3; acc = d // 3
+    return acc, br, iota, diaer
+
+
+# ---- macron/scansion label conventions (meter/marks.py, reproduced verbatim) ----
+MAC_LONG, MAC_SHORT = 0, 1
+SCAN_O, SCAN_HEAVY, SCAN_LIGHT, SCAN_VERSE = 0, 1, 2, 3
+
+_A, _E, _H, _I, _O, _Y, _W = (LETTER_IDS[c] for c in "αεηιουω")
+DICHRONA_IDS = np.array([_A, _I, _Y])
+VOWEL_IDS = np.array([_A, _E, _H, _I, _O, _Y, _W])
+DIPHTHONGS = {(_A, _I), (_A, _Y), (_E, _I), (_E, _Y), (_H, _Y),
+              (_O, _I), (_O, _Y), (_Y, _I), (_W, _Y)}
+
+
+def ambiguous_mask(chars: np.ndarray, boundary: np.ndarray, dia: np.ndarray) -> np.ndarray:
+    """Which plane positions are ambiguous dichrona (the macronizer's domain)? A
+    position is ambiguous iff it's a base alpha/iota/upsilon without circumflex or
+    iota subscript, and not part of a diphthong (diaeresis on the second vowel
+    breaks the diphthong; pairs never span a word boundary)."""
+    n = len(chars)
+    d = np.asarray(dia, dtype=np.int64)
+    acc, _br, iota, diaer = _unpack_dia(d.copy())
+    is_dich = np.isin(chars, DICHRONA_IDS)
+    out = is_dich & (acc != 3) & (iota == 0)
+    if n > 1:
+        chars = np.asarray(chars)
+        pair = np.zeros(n - 1, dtype=bool)
+        for f, s in DIPHTHONGS:
+            pair |= (chars[:-1] == f) & (chars[1:] == s)
+        pair &= np.asarray(boundary[:-1]) == 0
+        out[1:] &= ~(pair & (diaer[1:] == 0))
+        out[:-1] &= ~(pair & (diaer[1:] == 0))
+    return out
+
+
+def merge_vowelless_syllables(chars: np.ndarray, scan_labels: np.ndarray) -> np.ndarray:
+    """A predicted syllable span with no vowel isn't a real syllable -- it's a
+    boundary placed one letter early, typically at the first letter of a geminate
+    consonant pair (e.g. "{λε}[ν]" for what should be one closed syllable
+    "[λεν]"). Merge any such span into the preceding one, keeping its own weight
+    label: that label (usually already correct, since it's typically a closing
+    consonant) is normally right for the merged syllable -- only the boundary was
+    misplaced. A vowel-less span at the very start of the line is left as-is."""
+    out = np.asarray(scan_labels).copy()
+    is_vowel = np.isin(np.asarray(chars), VOWEL_IDS)
+    kept = []
+    start = 0
+    for i in range(len(out)):
+        if out[i] == SCAN_O:
+            continue
+        if not is_vowel[start:i + 1].any() and kept:
+            out[kept.pop()] = SCAN_O
+        kept.append(i)
+        start = i + 1
+    return out
+
+
+def enforce_circumflex_heavy(dia: np.ndarray, scan_labels: np.ndarray) -> np.ndarray:
+    """A syllable containing a circumflexed vowel is always heavy -- a fixed rule
+    of Greek prosody, not something the per-letter classifier can get wrong in
+    principle, only in practice. Flip any SCAN_LIGHT span containing a circumflex
+    to SCAN_HEAVY; SCAN_VERSE is left alone (it already renders as a heavy-looking
+    bracket). Boundary placement itself is untouched -- this only corrects weight."""
+    d = np.asarray(dia, dtype=np.int64)
+    acc, _br, _iota, _diaer = _unpack_dia(d.copy())
+    has_circ = acc == 3
+    out = np.asarray(scan_labels).copy()
+    start = 0
+    for i in range(len(out)):
+        if out[i] != SCAN_O:
+            if out[i] == SCAN_LIGHT and has_circ[start:i + 1].any():
+                out[i] = SCAN_HEAVY
+            start = i + 1
+    return out
+
+
+def _is_letter(ch: str) -> bool:
+    low = ch.lower()
+    if low in LETTER_IDS or low in _EXTRA_BASE:
+        return True
+    dec = unicodedata.normalize("NFD", low)
+    return bool(dec) and (dec[0] in LETTER_IDS or dec[0] in _EXTRA_BASE)
+
+
+def insert_marks(plain: str, labels: dict) -> str:
+    """Write `_` (long) / `^` (short) after the letters given by
+    {letter_ordinal: MAC_LONG|MAC_SHORT} -- production macron output format."""
+    nfc = unicodedata.normalize("NFC", plain)
+    out, ordinal, pending = [], -1, None
+    for ch in nfc:
+        if pending is not None and not unicodedata.category(ch).startswith("M"):
+            out.append(pending)
+            pending = None
+        out.append(ch)
+        if _is_letter(ch):
+            ordinal += 1
+            if ordinal in labels:
+                pending = "_" if labels[ordinal] == MAC_LONG else "^"
+    if pending is not None:
+        out.append(pending)
+    return "".join(out)
+
+
+def bracketize(plain: str, labels: dict) -> str:
+    """Render per-letter scan labels back into [heavy]{light} syllable spans
+    (verse-final span rendered as heavy, matching brevis-in-longo display)."""
+    nfc = unicodedata.normalize("NFC", plain)
+    out, cur, ordinal = [], [], -1
+    for ch in nfc:
+        cur.append(ch)
+        if _is_letter(ch):
+            ordinal += 1
+            lab = labels.get(ordinal, SCAN_O)
+            if lab != SCAN_O:
+                o, c = ("{", "}") if lab == SCAN_LIGHT else ("[", "]")
+                out.append(o + "".join(cur) + c)
+                cur = []
+    if cur:
+        out.append("".join(cur))
+    return "".join(out)
+
+
+@dataclass
+class _Encoded:
+    chars: list
+    boundary: list
+    dia: list
+    punct: list
+    cap: list
+
+
+class CharBertMeterProcessor:
+    """`processor(text)` -> dict of batched tensors ready for `CharBertMeterModel(**batch)`."""
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def from_pretrained(cls, *_args, **_kwargs):
+        return cls()
+
+    def save_pretrained(self, save_directory, **_kwargs):
+        Path(save_directory).mkdir(parents=True, exist_ok=True)
+        (Path(save_directory) / "processor_config.json").write_text(
+            json.dumps({"processor_class": "CharBertMeterProcessor"}))
+
+    def _classify(self, text: str) -> _Encoded:
+        nfd = unicodedata.normalize("NFD", text)
+        chars, boundary, dia, punct, cap = [], [], [], [], []
+        acc = br = iota = diaer = 0
+        pending_bnd = 0
+        i = 0
+        while i < len(nfd):
+            ch = nfd[i]
+            if ch == "-":
+                run = 0
+                while i < len(nfd) and nfd[i] == "-":
+                    run += 1
+                    i += 1
+                for _ in range(run):
+                    chars.append(MASK); boundary.append(UNK_BND)
+                    dia.append(UNK_DIA); punct.append(UNK_PUNCT); cap.append(0)
+                continue
+            low = ch.lower()
+            base = low if low in LETTER_IDS else _EXTRA_BASE.get(low)
+            if base is not None:
+                if chars and pending_bnd:
+                    boundary[-1] = pending_bnd
+                    pending_bnd = 0
+                chars.append(LETTER_IDS[base])
+                cap.append(1 if ch != low else 0)
+                boundary.append(0)
+                dia.append(0)
+                punct.append(0)
+                acc = br = iota = diaer = 0
+            elif unicodedata.combining(ch) or ord(ch) in _MARK_MAP:
+                kind = _MARK_MAP.get(ord(ch))
+                if kind in _ACC:
+                    acc = _ACC[kind]
+                elif kind in _BR:
+                    br = _BR[kind]
+                elif kind == "iota":
+                    iota = 1
+                elif kind == "diaer":
+                    diaer = 1
+                if dia:
+                    dia[-1] = _pack_dia(acc, br, iota, diaer)
+            elif ch.isspace():
+                pending_bnd = max(pending_bnd, 1)
+            elif ch in ".;!?":
+                pending_bnd = max(pending_bnd, 2)
+                if punct:
+                    punct[-1] = 4 if ch == "." else 5
+            elif ch in ",:··":
+                if punct:
+                    punct[-1] = {",": 1, "·": 2, "·": 2, ":": 3}.get(ch, 0)
+            i += 1
+        if boundary:
+            boundary[-1] = max(boundary[-1], 2)
+        return _Encoded(chars, boundary, dia, punct, cap)
+
+    def __call__(self, text: str, has_boundaries: bool = True):
+        """Encode `text` into model-ready tensors. Unlike the base pretraining
+        processor, `cap` is a real model input here (fine-tune-only channel)."""
+        enc = self._classify(text)
+        n = len(enc.chars)
+        chars = np.array(enc.chars, dtype=np.int64)
+        boundary = np.array(enc.boundary, dtype=np.int64)
+        dia = np.array(enc.dia, dtype=np.int64)
+        punct = np.array(enc.punct, dtype=np.int64)
+        cap = np.array(enc.cap, dtype=np.int64)
+        if not has_boundaries:
+            boundary[:] = UNK_BND
+
+        batch = dict(
+            input_ids=torch.from_numpy(chars)[None],
+            boundary=torch.from_numpy(boundary)[None],
+            dia=torch.from_numpy(dia)[None],
+            punct=torch.from_numpy(punct)[None],
+            cap=torch.from_numpy(cap)[None],
+            seg_id=torch.zeros(1, n, dtype=torch.long),
+        )
+        batch["_text"] = text  # kept out-of-band for decode (marks splice into the original string)
+        batch["_chars"] = chars
+        batch["_boundary"] = boundary
+        batch["_dia"] = dia
+        return batch
+
+    # ---------------------------------------------------------------- decode
+
+    def decode_macronization(self, model_out, batch) -> str:
+        """Insert `_`/`^` (long/short) after every ambiguous alpha/iota/upsilon --
+        matches `meter.predict --macronize` exactly (only ambiguous dichrona get a
+        mark; unambiguous positions -- eta, omega, diphthongs, iota subscript,
+        circumflexed vowels -- are left bare, since their length isn't in doubt)."""
+        pred_mac = model_out.mac.argmax(-1)[0].numpy()
+        amb = ambiguous_mask(batch["_chars"], batch["_boundary"], batch["_dia"])
+        labels = {int(i): int(pred_mac[i]) for i in np.flatnonzero(amb)}
+        return insert_marks(batch["_text"], labels)
+
+    def decode_scansion(self, model_out, batch) -> str:
+        """Bracket every syllable the model assigns a non-trivial weight to:
+        [heavy], {light}, with the line-final syllable (brevis in longo) shown as
+        heavy -- matches `meter.predict --scan` exactly, including its two
+        deterministic corrections: merge_vowelless_syllables (a vowel-less
+        predicted span gets folded into the preceding syllable) and
+        enforce_circumflex_heavy (a circumflexed syllable is always heavy)."""
+        pred_scan = model_out.scan.argmax(-1)[0].numpy()
+        pred_scan = merge_vowelless_syllables(batch["_chars"], pred_scan)
+        pred_scan = enforce_circumflex_heavy(batch["_dia"], pred_scan)
+        labels = {i: int(c) for i, c in enumerate(pred_scan) if c > 0}
+        return bracketize(batch["_text"], labels)

training_metadata.json

@@ -0,0 +1,119 @@
+{
+  "cfg": {
+    "name": "meter_joint_docclean",
+    "seed": 0,
+    "ckpt": "best.pt",
+    "out_dir": "meter_joint_docclean",
+    "encoded": "encoded",
+    "attn": "sdpa",
+    "T": 2048,
+    "micro_batch": 16,
+    "eval_micro": 16,
+    "lr_enc": 0.0001,
+    "lr_head": 0.0008,
+    "wd": 0.01,
+    "clip": 1.0,
+    "warmup_frac": 0.1,
+    "epochs": 25,
+    "patience": 6,
+    "head_dropout": 0.33,
+    "use_cap": true,
+    "scalar_mix": true,
+    "w_mac": 1.0,
+    "w_scan": 1.0,
+    "select": "mean",
+    "prose_per_epoch": 150000,
+    "scan_passes": 2,
+    "scan_max_verses": 8,
+    "mac_class_w": "auto",
+    "scan_class_w": "auto",
+    "log_every": 50
+  },
+  "mcfg": {
+    "head_dropout": 0.33,
+    "w_mac": 1.0,
+    "w_scan": 1.0,
+    "use_cap": true,
+    "scalar_mix": true,
+    "mac_class_w": [
+      3.935332098208277,
+      0.5727731680236648
+    ],
+    "scan_class_w": [
+      0.4482191825355474,
+      1.2170555073099973,
+      1.2004924282537328,
+      8.748756624191806
+    ]
+  },
+  "pretrain_cfg": {
+    "name": "gcb_doc_clean",
+    "seed": 0,
+    "d_model": 1024,
+    "depth": 32,
+    "char_window": 256,
+    "attn": "flex",
+    "compile": true,
+    "qk_norm": true,
+    "seq_len": 8192,
+    "micro_batch": 4,
+    "grad_accum": 2,
+    "lr": 0.0003,
+    "wd": 0.1,
+    "clip": 1.0,
+    "warmup_frac": 0.04,
+    "decay_frac": 0.4,
+    "tier_weights": [
+      1.0,
+      1.0,
+      0.3
+    ],
+    "lam": 0.1,
+    "total_steps": 120000,
+    "hard_max_steps": 1000000,
+    "num_workers": 16,
+    "log_every": 20,
+    "ckpt_every": 2000,
+    "eval_every_ckpt": true,
+    "eval_n": 256,
+    "out_dir": "gcb_doc_clean",
+    "anneal_phases": [
+      [
+        0.5,
+        [
+          1.0,
+          1.0,
+          0.3
+        ]
+      ],
+      [
+        0.9,
+        [
+          1.0,
+          1.0,
+          0.0
+        ]
+      ],
+      [
+        1.0,
+        [
+          1.0,
+          0.0,
+          0.0
+        ]
+      ]
+    ]
+  },
+  "epoch": 12,
+  "dev": {
+    "mac_bal": 0.8927,
+    "mac_acc": 0.9195,
+    "mac_n": 2498,
+    "scan_bal": 0.9811,
+    "scan_acc": 0.9787,
+    "end_f1": 0.9791,
+    "scan_n": 68004
+  },
+  "T": 2048,
+  "_source_checkpoint": "best.pt"
+}