Ericu950/Stoicheia-macronizer

1"""HF-Hub-compatible processor for Stoicheia-meter: text <-> the model's five2input planes (chars/boundary/dia/punct/cap -- capitalization is a real input here,3unlike the base pretrained model, where it is output-only), plus decode helpers for4the two production tasks: macronization (vowel-length marks on ambiguous dichrona)5and metrical scansion (heavy/light/verse-final syllable bracketing).6 7Mark conventions and the ambiguous-dichrona rule are ported verbatim from8meter/marks.py (the training-side projection code) so a converted checkpoint's9predictions decode identically to the original project's `meter.predict` CLI.10"""11from __future__ import annotations12 13import json14import unicodedata15from dataclasses import dataclass16from pathlib import Path17 18import numpy as np19import torch20 21MASK, BLANK, PAD = 24, 25, 2622UNK_BND, UNK_DIA, UNK_PUNCT = 3, 48, 623 24ALPHABET = "αβγδεζηθικλμνξοπρστυφχψω"25LETTER_IDS = {c: i for i, c in enumerate(ALPHABET)}26ID2LETTER = np.array(list(ALPHABET))27 28_EXTRA_BASE = {29    "ς": "σ", "ϲ": "σ", "Ϲ": "σ", "ϐ": "β", "ϑ": "θ", "ϕ": "φ", "ϰ": "κ", "ϱ": "ρ", "ϖ": "π",30}31_MARK_MAP = {32    0x0301: "acute", 0x0341: "acute", 0x0300: "grave", 0x0340: "grave",33    0x0342: "circ", 0x0302: "circ", 0x0313: "smooth", 0x0343: "smooth",34    0x0314: "rough", 0x0345: "iota", 0x0308: "diaer",35}36_ACC = {"acute": 1, "grave": 2, "circ": 3}37_BR = {"smooth": 1, "rough": 2}38 39 40def _pack_dia(acc, br, iota, diaer):41    return ((acc * 3 + br) * 2 + iota) * 2 + diaer42 43 44def _unpack_dia(d):45    diaer = d % 2; d //= 246    iota = d % 2; d //= 247    br = d % 3; acc = d // 348    return acc, br, iota, diaer49 50 51# ---- macron/scansion label conventions (meter/marks.py, reproduced verbatim) ----52MAC_LONG, MAC_SHORT = 0, 153SCAN_O, SCAN_HEAVY, SCAN_LIGHT, SCAN_VERSE = 0, 1, 2, 354 55_A, _E, _H, _I, _O, _Y, _W = (LETTER_IDS[c] for c in "αεηιουω")56DICHRONA_IDS = np.array([_A, _I, _Y])57VOWEL_IDS = np.array([_A, _E, _H, _I, _O, _Y, _W])58DIPHTHONGS = {(_A, _I), (_A, _Y), (_E, _I), (_E, _Y), (_H, _Y),59              (_O, _I), (_O, _Y), (_Y, _I), (_W, _Y)}60 61 62def ambiguous_mask(chars: np.ndarray, boundary: np.ndarray, dia: np.ndarray) -> np.ndarray:63    """Which plane positions are ambiguous dichrona (the macronizer's domain)? A64    position is ambiguous iff it's a base alpha/iota/upsilon without circumflex or65    iota subscript, and not part of a diphthong (diaeresis on the second vowel66    breaks the diphthong; pairs never span a word boundary)."""67    n = len(chars)68    d = np.asarray(dia, dtype=np.int64)69    acc, _br, iota, diaer = _unpack_dia(d.copy())70    is_dich = np.isin(chars, DICHRONA_IDS)71    out = is_dich & (acc != 3) & (iota == 0)72    if n > 1:73        chars = np.asarray(chars)74        pair = np.zeros(n - 1, dtype=bool)75        for f, s in DIPHTHONGS:76            pair |= (chars[:-1] == f) & (chars[1:] == s)77        pair &= np.asarray(boundary[:-1]) == 078        out[1:] &= ~(pair & (diaer[1:] == 0))79        out[:-1] &= ~(pair & (diaer[1:] == 0))80    return out81 82 83def merge_vowelless_syllables(chars: np.ndarray, scan_labels: np.ndarray) -> np.ndarray:84    """A predicted syllable span with no vowel isn't a real syllable -- it's a85    boundary placed one letter early, typically at the first letter of a geminate86    consonant pair (e.g. "{λε}[ν]" for what should be one closed syllable87    "[λεν]"). Merge any such span into the preceding one, keeping its own weight88    label: that label (usually already correct, since it's typically a closing89    consonant) is normally right for the merged syllable -- only the boundary was90    misplaced. A vowel-less span at the very start of the line is left as-is."""91    out = np.asarray(scan_labels).copy()92    is_vowel = np.isin(np.asarray(chars), VOWEL_IDS)93    kept = []94    start = 095    for i in range(len(out)):96        if out[i] == SCAN_O:97            continue98        if not is_vowel[start:i + 1].any() and kept:99            out[kept.pop()] = SCAN_O100        kept.append(i)101        start = i + 1102    return out103 104 105def enforce_circumflex_heavy(dia: np.ndarray, scan_labels: np.ndarray) -> np.ndarray:106    """A syllable containing a circumflexed vowel is always heavy -- a fixed rule107    of Greek prosody, not something the per-letter classifier can get wrong in108    principle, only in practice. Flip any SCAN_LIGHT span containing a circumflex109    to SCAN_HEAVY; SCAN_VERSE is left alone (it already renders as a heavy-looking110    bracket). Boundary placement itself is untouched -- this only corrects weight."""111    d = np.asarray(dia, dtype=np.int64)112    acc, _br, _iota, _diaer = _unpack_dia(d.copy())113    has_circ = acc == 3114    out = np.asarray(scan_labels).copy()115    start = 0116    for i in range(len(out)):117        if out[i] != SCAN_O:118            if out[i] == SCAN_LIGHT and has_circ[start:i + 1].any():119                out[i] = SCAN_HEAVY120            start = i + 1121    return out122 123 124def _is_letter(ch: str) -> bool:125    low = ch.lower()126    if low in LETTER_IDS or low in _EXTRA_BASE:127        return True128    dec = unicodedata.normalize("NFD", low)129    return bool(dec) and (dec[0] in LETTER_IDS or dec[0] in _EXTRA_BASE)130 131 132def insert_marks(plain: str, labels: dict) -> str:133    """Write `_` (long) / `^` (short) after the letters given by134    {letter_ordinal: MAC_LONG|MAC_SHORT} -- production macron output format."""135    nfc = unicodedata.normalize("NFC", plain)136    out, ordinal, pending = [], -1, None137    for ch in nfc:138        if pending is not None and not unicodedata.category(ch).startswith("M"):139            out.append(pending)140            pending = None141        out.append(ch)142        if _is_letter(ch):143            ordinal += 1144            if ordinal in labels:145                pending = "_" if labels[ordinal] == MAC_LONG else "^"146    if pending is not None:147        out.append(pending)148    return "".join(out)149 150 151def bracketize(plain: str, labels: dict) -> str:152    """Render per-letter scan labels back into [heavy]{light} syllable spans153    (verse-final span rendered as heavy, matching brevis-in-longo display)."""154    nfc = unicodedata.normalize("NFC", plain)155    out, cur, ordinal = [], [], -1156    for ch in nfc:157        cur.append(ch)158        if _is_letter(ch):159            ordinal += 1160            lab = labels.get(ordinal, SCAN_O)161            if lab != SCAN_O:162                o, c = ("{", "}") if lab == SCAN_LIGHT else ("[", "]")163                out.append(o + "".join(cur) + c)164                cur = []165    if cur:166        out.append("".join(cur))167    return "".join(out)168 169 170@dataclass171class _Encoded:172    chars: list173    boundary: list174    dia: list175    punct: list176    cap: list177 178 179class CharBertMeterProcessor:180    """`processor(text)` -> dict of batched tensors ready for `CharBertMeterModel(**batch)`."""181 182    def __init__(self):183        pass184 185    @classmethod186    def from_pretrained(cls, *_args, **_kwargs):187        return cls()188 189    def save_pretrained(self, save_directory, **_kwargs):190        Path(save_directory).mkdir(parents=True, exist_ok=True)191        (Path(save_directory) / "processor_config.json").write_text(192            json.dumps({"processor_class": "CharBertMeterProcessor"}))193 194    def _classify(self, text: str) -> _Encoded:195        nfd = unicodedata.normalize("NFD", text)196        chars, boundary, dia, punct, cap = [], [], [], [], []197        acc = br = iota = diaer = 0198        pending_bnd = 0199        i = 0200        while i < len(nfd):201            ch = nfd[i]202            if ch == "-":203                run = 0204                while i < len(nfd) and nfd[i] == "-":205                    run += 1206                    i += 1207                for _ in range(run):208                    chars.append(MASK); boundary.append(UNK_BND)209                    dia.append(UNK_DIA); punct.append(UNK_PUNCT); cap.append(0)210                continue211            low = ch.lower()212            base = low if low in LETTER_IDS else _EXTRA_BASE.get(low)213            if base is not None:214                if chars and pending_bnd:215                    boundary[-1] = pending_bnd216                    pending_bnd = 0217                chars.append(LETTER_IDS[base])218                cap.append(1 if ch != low else 0)219                boundary.append(0)220                dia.append(0)221                punct.append(0)222                acc = br = iota = diaer = 0223            elif unicodedata.combining(ch) or ord(ch) in _MARK_MAP:224                kind = _MARK_MAP.get(ord(ch))225                if kind in _ACC:226                    acc = _ACC[kind]227                elif kind in _BR:228                    br = _BR[kind]229                elif kind == "iota":230                    iota = 1231                elif kind == "diaer":232                    diaer = 1233                if dia:234                    dia[-1] = _pack_dia(acc, br, iota, diaer)235            elif ch.isspace():236                pending_bnd = max(pending_bnd, 1)237            elif ch in ".;!?":238                pending_bnd = max(pending_bnd, 2)239                if punct:240                    punct[-1] = 4 if ch == "." else 5241            elif ch in ",:··":242                if punct:243                    punct[-1] = {",": 1, "·": 2, "·": 2, ":": 3}.get(ch, 0)244            i += 1245        if boundary:246            boundary[-1] = max(boundary[-1], 2)247        return _Encoded(chars, boundary, dia, punct, cap)248 249    def __call__(self, text: str, has_boundaries: bool = True):250        """Encode `text` into model-ready tensors. Unlike the base pretraining251        processor, `cap` is a real model input here (fine-tune-only channel)."""252        enc = self._classify(text)253        n = len(enc.chars)254        chars = np.array(enc.chars, dtype=np.int64)255        boundary = np.array(enc.boundary, dtype=np.int64)256        dia = np.array(enc.dia, dtype=np.int64)257        punct = np.array(enc.punct, dtype=np.int64)258        cap = np.array(enc.cap, dtype=np.int64)259        if not has_boundaries:260            boundary[:] = UNK_BND261 262        batch = dict(263            input_ids=torch.from_numpy(chars)[None],264            boundary=torch.from_numpy(boundary)[None],265            dia=torch.from_numpy(dia)[None],266            punct=torch.from_numpy(punct)[None],267            cap=torch.from_numpy(cap)[None],268            seg_id=torch.zeros(1, n, dtype=torch.long),269        )270        batch["_text"] = text  # kept out-of-band for decode (marks splice into the original string)271        batch["_chars"] = chars272        batch["_boundary"] = boundary273        batch["_dia"] = dia274        return batch275 276    # ---------------------------------------------------------------- decode277 278    def decode_macronization(self, model_out, batch) -> str:279        """Insert `_`/`^` (long/short) after every ambiguous alpha/iota/upsilon --280        matches `meter.predict --macronize` exactly (only ambiguous dichrona get a281        mark; unambiguous positions -- eta, omega, diphthongs, iota subscript,282        circumflexed vowels -- are left bare, since their length isn't in doubt)."""283        pred_mac = model_out.mac.argmax(-1)[0].numpy()284        amb = ambiguous_mask(batch["_chars"], batch["_boundary"], batch["_dia"])285        labels = {int(i): int(pred_mac[i]) for i in np.flatnonzero(amb)}286        return insert_marks(batch["_text"], labels)287 288    def decode_scansion(self, model_out, batch) -> str:289        """Bracket every syllable the model assigns a non-trivial weight to:290        [heavy], {light}, with the line-final syllable (brevis in longo) shown as291        heavy -- matches `meter.predict --scan` exactly, including its two292        deterministic corrections: merge_vowelless_syllables (a vowel-less293        predicted span gets folded into the preceding syllable) and294        enforce_circumflex_heavy (a circumflexed syllable is always heavy)."""295        pred_scan = model_out.scan.argmax(-1)[0].numpy()296        pred_scan = merge_vowelless_syllables(batch["_chars"], pred_scan)297        pred_scan = enforce_circumflex_heavy(batch["_dia"], pred_scan)298        labels = {i: int(c) for i, c in enumerate(pred_scan) if c > 0}299        return bracketize(batch["_text"], labels)300