processing_char_bert_joint.py

"""HF-Hub-compatible processor for CharDiff-grc-tagger-parser.

Turns one or more already word-tokenized sentences into the model's input tensors (the same
four character planes as the base CharBertProcessor, plus a `word_id` array pooling characters
into words, a `cap` plane -- an actual model input for this fine-tuned checkpoint, unlike the
base model where capitalization is output-only -- and a `seg_id` plane), and decodes the
model's raw logits back into predictions. The model's NATIVE tag/relation inventory is the
Perseus/AGDT scheme (9-position morphological tags; Prague-style dependency relations like
SBJ/OBJ/ATR/COORD/AuxP -- the scheme its OGA training treebank uses), not Universal
Dependencies. Pass `ud=True` to `decode()` for Universal Dependencies-style output instead
(UPOS/FEATS/deprel) -- see `to_ud()` below for the conversion and its scope/limitations.

Interface: call the processor with `List[str]` (one already-tokenized sentence, e.g.
`["ὁ", "ἄνθρωπος", "τρέχει", "."]`) or `List[List[str]]` (a batch of such sentences) --
i.e. tokens are CoNLL-U-style FORMs, punctuation already split into its own element. This
mirrors tagger/dataset.py's `sent.tokens` convention directly (each token's FORM is encoded
independently) without pulling in a real Greek word tokenizer as a dependency; a convenience
`.tokenize(text)` staticmethod is provided for plain-text input (naive whitespace split +
peeling off trailing sentence punctuation), which is enough for a usage-example sentence but
is not a substitute for real tokenization on running text.

Not a `PreTrainedTokenizer` subclass, for the same reason as the base CharBertProcessor: the
underlying encoding is parallel per-letter planes, not a single token-id stream. Follows the
same `register_for_auto_class` mechanism, so `AutoProcessor.from_pretrained(repo_id,
trust_remote_code=True)` works the same way.
"""
from __future__ import annotations

import json
import re
import unicodedata
from pathlib import Path

import numpy as np
import torch

PAD_ID = 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}
_PUNCT_CLASS = {",": 1, "·": 2, "·": 2, ":": 3}   # small punct-plane LUT (comma/mid-dot/colon);
                                                    # sentence-final .;!? are captured via the
                                                    # boundary plane instead, same convention as
                                                    # processing_char_bert.py's base processor.


def _pack_dia(acc, br, iota, diaer):
    return ((acc * 3 + br) * 2 + iota) * 2 + diaer


# ==== Universal Dependencies conversion ========================================================
#
# The model's native output is the Perseus/AGDT annotation scheme its OGA training treebank
# uses: a 9-position morphological tag (pos.person.number.tense.mood.voice.gender.case.degree,
# one letter per position, "-" for n/a) and Prague-style dependency relations (SBJ, OBJ, ATR,
# ADV, PNOM, PRED, APOS, COORD, Aux*, ATV/AtvV, ExD, MWE, OCOMP) -- NOT Universal Dependencies.
#
# This is a solved, documented conversion problem, not something designed from scratch here:
# both UD_Ancient_Greek-Perseus and UD_Ancient_Greek-PROIEL are themselves automatic
# conversions of exactly this annotation style, via Zeman & Ramasamy's HamleDT pipeline
# (Zeman et al. 2014, "HamleDT: Harmonized Multi-Language Dependency Treebank", LRE) --
# specifically Treex::Block::HamleDT::{GRC::Harmonize,HarmonizePerseus,Udep} and
# Treex::Tool::PhraseBuilder::PragueToUD (github.com/ufal/treex), plus the morphology decoder
# Lingua::Interset::Tagset::GRC::Conll (github.com/dan-zeman/interset). The tables and
# restructuring rules below port that design to our exact tag/relation inventory.
#
# Two relations require real TREE restructuring, not just a label rename, because UD attaches
# function words/conjuncts differently than Prague style:
#   - COORD: Prague style makes the coordinating conjunction the head, conjuncts its children.
#     UD makes the FIRST conjunct the head, the conjunction becomes a `cc` dependent of it, and
#     other conjuncts become `conj` dependents.
#   - AuxP: Prague style makes the preposition the head of its complement. UD reverses this:
#     the nominal complement becomes the head, the preposition becomes a `case` dependent of it.
#   - Pnom (predicate nominal / copula): Prague style makes the copula verb the head with the
#     nominal attached as Pnom. UD makes the NOMINAL the head, the copula a `cop` dependent,
#     and reattaches the subject to the nominal instead of the verb.
#
# What relation does the PROMOTED node (first conjunct / nominal complement) get in UD, given
# it used to be a plain dependent of COORD/AuxP? In full AGDT/HamleDT, the coordinator's or
# preposition's own outbound relation carries the phrase's true external function via a
# _CO/_AP label suffix (e.g. "Sb_Co"), which HamleDT strips off and hands to the promoted node.
# Our 24-label inventory has no such suffix mechanism -- verified directly against real training
# examples in oga_sota.conllu (e.g. "ἀπὸ" and "ἕως" each attach to their external governor with
# deprel=AuxP verbatim, and "καὶ" attaches with deprel=COORD verbatim, regardless of whether the
# phrase functions as a locative adjunct, a coordinated subject, or anything else). So the label
# itself never carries a recoverable function here, and the promoted node gets UD's generic
# fallback instead of a falsely precise guess: "obl" for AuxP nominals, "dep" for COORD first
# conjuncts (see Stage 1/2 below).
#
# SCOPE: this ports the core, well-documented HamleDT/UD restructuring rules and a straight
# per-label relabeling for everything else. It does NOT replicate every Ancient-Greek-specific
# refinement in the full pipeline (e.g. GRC::Harmonize's regex-based negation-particle
# splitting, or the is_member/shared-modifier _CO/_AP distinction discussed above) -- treat this
# as a faithful, reasonably-scoped port, not a byte-exact reproduction of the official
# UD_Ancient_Greek-Perseus release.

# position 1 (POS) -> UD UPOS. 'c' defaults to CCONJ, refined to SCONJ if its converted deprel
# is "mark" (i.e. it was AuxC); 'v' stays VERB regardless of mood (participles etc. get
# VerbForm=Part as a feature, not a UPOS change, matching standard UD_Ancient_Greek practice).
_POS_TO_UPOS = {
    "a": "ADJ", "c": "CCONJ", "d": "ADV", "g": "PART", "i": "INTJ", "l": "DET",
    "m": "NUM", "n": "NOUN", "p": "PRON", "r": "ADP", "u": "PUNCT", "v": "VERB", "x": "X",
}
_NUMBER = {"d": "Dual", "p": "Plur", "s": "Sing"}
_TENSE_ASPECT = {   # (Tense, Aspect-or-None); aorist/perfect/future-perfect analysis per
    "a": ("Past", "Perf"),   # aorist = perfective past
    "f": ("Fut", None),
    "i": ("Past", "Imp"),    # imperfect
    "l": ("Pqp", None),      # pluperfect
    "p": ("Pres", None),
    "r": ("Pres", "Perf"),   # perfect = present relevance of a completed action
    "t": ("Fut", "Perf"),    # future perfect
}
_MOOD_VERBFORM = {   # (Mood-or-None, VerbForm)
    "i": ("Ind", "Fin"), "s": ("Sub", "Fin"), "o": ("Opt", "Fin"), "m": ("Imp", "Fin"),
    "n": (None, "Inf"), "p": (None, "Part"), "g": (None, "Gdv"),
}
_VOICE = {"a": "Act", "m": "Mid", "p": "Pass", "e": "Mid"}   # e = mediopassive, approximated as Mid
_GENDER = {"c": "Com", "f": "Fem", "m": "Masc", "n": "Neut"}
_CASE = {"a": "Acc", "d": "Dat", "g": "Gen", "n": "Nom", "v": "Voc"}
_DEGREE = {"c": "Cmp", "s": "Sup"}


def xpos_to_upos_feats(xpos: str) -> tuple[str, dict]:
    """9-position Perseus/AGDT tag (e.g. "v3spia---") -> (UD UPOS, UD FEATS dict).
    "-" in any position means not applicable/unspecified and is simply omitted from FEATS."""
    pos = xpos[0] if xpos else "x"
    upos = _POS_TO_UPOS.get(pos, "X")
    feats = {}
    rest = xpos[1:].ljust(8, "-")
    person, number, tense, mood, voice, gender, case, degree = rest[:8]
    if person in "123":
        feats["Person"] = person
    if number in _NUMBER:
        feats["Number"] = _NUMBER[number]
    if tense in _TENSE_ASPECT:
        t, a = _TENSE_ASPECT[tense]
        feats["Tense"] = t
        if a:
            feats["Aspect"] = a
    if mood in _MOOD_VERBFORM:
        m, vf = _MOOD_VERBFORM[mood]
        if m:
            feats["Mood"] = m
        feats["VerbForm"] = vf
    if voice in _VOICE:
        feats["Voice"] = _VOICE[voice]
    if gender in _GENDER:
        feats["Gender"] = _GENDER[gender]
    if case in _CASE:
        feats["Case"] = _CASE[case]
    if degree in _DEGREE:
        feats["Degree"] = _DEGREE[degree]
    return upos, feats


# Straight per-label relabeling for everything that ISN'T COORD/AuxP/Pnom (those need the tree
# restructuring in to_ud() below). Homoglyph artifacts in the training treebank (Greek "Ζ"/"Κ"
# instead of Latin "Z"/"K" in a couple of AuxZ/AuxK instances) are normalized first.
_HOMOGLYPH_FIX = {"AuxΖ": "AuxZ", "AuxΚ": "AuxK"}
_SIMPLE_DEPREL = {
    "SBJ": "nsubj", "OBJ": "obj", "OCOMP": "obj", "APOS": "appos",
    "AuxC": "mark", "AuxV": "aux", "AuxG": "punct", "AuxK": "punct", "AuxX": "punct",
    "AuxR": "expl", "AuxY": "cc", "AuxZ": "advmod", "AuxΖ": "advmod", "AuxΚ": "punct",
    "ATV": "xcomp", "AtvV": "xcomp", "ExD": "orphan", "MWE": "fixed",
    "PRED": "root",
}


def _normalize_deprel(label: str | None) -> str | None:
    return _HOMOGLYPH_FIX.get(label, label)


def to_ud(sent: list[dict]) -> list[dict]:
    """Convert one sentence's native decode() output (list of {form, xpos, lemma, upos, head,
    deprel} dicts, 1-indexed word order, head=0 meaning root) into UD-style output: each dict
    becomes {form, lemma, upos (UD), xpos (native tag, kept as-is per UD convention that XPOS
    is language-specific and optional), feats (dict), head, deprel (UD)}.

    See the module-level comment above for the restructuring rules and scope/limitations.
    """
    n = len(sent)
    if n == 0:
        return []
    heads = [0] * (n + 1)
    deprels = [None] * (n + 1)
    native_upos = [None] * (n + 1)
    for i, w in enumerate(sent, start=1):
        heads[i] = w["head"] if w["head"] is not None else 0
        deprels[i] = _normalize_deprel(w["deprel"])
        native_upos[i] = w["upos"]

    def children_of():
        c = {i: [] for i in range(0, n + 1)}
        for i in range(1, n + 1):
            c.setdefault(heads[i], []).append(i)
        return c

    # Stage 1: COORD -> first-conjunct-head restructuring.
    for c in [i for i in range(1, n + 1) if deprels[i] == "COORD"]:
        kids = sorted(children_of().get(c, []))
        if not kids:
            continue
        first, rest = kids[0], kids[1:]
        heads[first] = heads[c]
        deprels[first] = "dep"   # COORD's own outbound label carries no recoverable function; see above
        heads[c] = first
        deprels[c] = "cc"
        for k in rest:
            heads[k] = first
            if deprels[k] not in ("AuxG", "AuxK", "AuxX", "AuxΖ", "AuxΚ"):  # not a punctuation child
                deprels[k] = "conj"
            # punctuation children just reattach to `first`, keeping their own (later-relabeled) deprel

    # Stage 2: AuxP -> case restructuring (nominal complement becomes the head).
    for p in [i for i in range(1, n + 1) if deprels[i] == "AuxP"]:
        kids = sorted(children_of().get(p, []))
        if not kids:
            continue
        nominal, others = kids[0], kids[1:]
        heads[nominal] = heads[p]
        deprels[nominal] = "obl"   # AuxP's own outbound label carries no recoverable function; see above
        heads[p] = nominal
        deprels[p] = "case"
        for o in others:
            heads[o] = nominal

    # Stage 3: Pnom (predicate nominal) -> copula reversal. The nominal becomes the head in
    # place of the copula verb; the verb becomes `cop`; anything else attached to the verb
    # (chiefly the subject) is reattached to the nominal.
    for nom in [i for i in range(1, n + 1) if deprels[i] == "PNOM"]:
        verb = heads[nom]
        if verb == 0:
            continue
        heads[nom] = heads[verb]
        deprels[nom] = deprels[verb] if deprels[verb] != "PRED" else "root"
        heads[verb] = nom
        deprels[verb] = "cop"
        for i in range(1, n + 1):
            if i != nom and heads[i] == verb:
                heads[i] = nom

    # Stage 4: everything else -- straight relabel, with a few UPOS/context-sensitive refinements.
    for i in range(1, n + 1):
        d = deprels[i]
        if d in ("cc", "conj", "case", "cop", "dep", "obl"):     # already converted above
            continue
        if d == "ATR":
            deprels[i] = ("det" if native_upos[i] == "l"
                          else "amod" if native_upos[i] == "a"
                          else "nmod")
        elif d == "ADV":
            deprels[i] = ("advmod" if native_upos[i] == "d"
                          else "advcl" if native_upos[i] == "v" else "obl")
        elif d in _SIMPLE_DEPREL:
            deprels[i] = _SIMPLE_DEPREL[d]
        # else: already UD-style (e.g. a second pass over an already-converted label) or
        # unrecognized -- left as-is rather than silently dropped.

    out = []
    for i, w in enumerate(sent, start=1):
        upos, feats = xpos_to_upos_feats(w["xpos"])
        if deprels[i] == "mark" and upos == "CCONJ":
            upos = "SCONJ"   # refine c (CCONJ default) once we know it introduced a subordinate clause
        out.append(dict(form=w["form"], lemma=w["lemma"], upos=upos, xpos=w["xpos"],
                         feats=feats, head=heads[i], deprel=deprels[i]))
    return out


# ---- edit-script machinery (tagger/edits.py, ported verbatim on strings) ---------------------

GRAVE, ACUTE = "̀", "́"
_STRIP_MARKS = {"̄", "̆", "̣"}   # macron, breve, dot-below


def form_key(s: str) -> str:
    """Fold a surface form for script/lexicon keys: lowercase, grave->acute, strip
    macron/breve/underdot. NFC output. (tagger.edits.form_key, verbatim.)"""
    s = unicodedata.normalize("NFD", s.lower())
    s = s.replace(GRAVE, ACUTE)
    s = "".join(ch for ch in s if ch not in _STRIP_MARKS)
    return unicodedata.normalize("NFC", s)


def apply_script(form: str, sc) -> str | None:
    """tagger.edits.apply_script, verbatim. sc = (p_cut, p_add, s_cut, s_add, cap)."""
    p_cut, p_add, s_cut, s_add, _cap = sc
    if len(form) < p_cut + s_cut:
        return None
    return p_add + form[p_cut:len(form) - s_cut or None] + s_add


# ---- per-word character classification (subset of processing_char_bert.py's _classify) -------


def classify_word(word: str):
    """NFD-normalize `word`; return (chars, dia, cap) int lists, or None if it has no Greek
    letters (e.g. pure punctuation / digits / Latin -- these get no word slot, matching
    tagger/dataset.py's encode_word/encode_sentence convention)."""
    nfd = unicodedata.normalize("NFD", word)
    chars, dia, cap = [], [], []
    acc = br = iota = diaer = 0
    for ch in nfd:
        low = ch.lower()
        base = low if low in LETTER_IDS else _EXTRA_BASE.get(low)
        if base is not None:
            chars.append(LETTER_IDS[base])
            cap.append(1 if ch != low else 0)
            dia.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)
        # anything else inside a single pre-tokenized word (stray punctuation) is dropped
    if not chars:
        return None
    return chars, dia, cap


def punct_class(tok: str) -> int:
    """Punctuation class a non-Greek token contributes to the preceding word's punct plane."""
    return max((_PUNCT_CLASS.get(c, 0) for c in tok), default=0)


def encode_sentence(tokens: list[str]):
    """tokens: pre-tokenized FORMs (CoNLL-U-style; punctuation already its own element).
    Returns (chars, boundary, dia, punct, cap, word_id, forms) for one packed sentence-row.
    `forms` holds the surface FORM of each *encodable* (has-Greek-letters) token, in word-slot
    order -- needed later to apply the lemma edit-script back onto real text."""
    chars, boundary, dia, punct, cap, word_id = [], [], [], [], [], []
    forms = []
    w = 0
    for tok in tokens:
        enc = classify_word(tok)
        if enc is None:
            pc = punct_class(tok)
            if pc and punct:
                punct[-1] = max(punct[-1], pc)
            continue
        c, d, cp = enc
        n = len(c)
        chars.extend(c)
        dia.extend(d)
        cap.extend(cp)
        punct.extend([0] * n)
        boundary.extend([0] * (n - 1) + [1])
        word_id.extend([w] * n)
        forms.append(tok)
        w += 1
    if boundary:
        boundary[-1] = 2   # sentence end
    return chars, boundary, dia, punct, cap, word_id, forms


_TRAIL_PUNCT_RE = re.compile(r"^(.*?)([.;!?,·:]+)$")


class CharBertJointProcessor:
    """`processor(sentences)` -> dict of batched tensors ready for
    `CharBertForTaggingAndParsing(**batch)`; `processor.decode(model_out, batch)` -> per-sentence
    list of word-level (form, xpos, lemma, upos, head, deprel) records."""

    def __init__(self, vocab: dict | None = None, deprel_vocab: list | None = None):
        self.vocab = vocab or {}
        self.deprel_vocab = deprel_vocab or []
        self.scripts = [tuple(s[:4]) + (bool(s[4]),) for s in self.vocab.get("scripts", [])]
        self.xpos_alpha = self.vocab.get("xpos_alpha", [])
        self.tags = self.vocab.get("tags", [])
        self.upos = self.vocab.get("upos", [])
        self.lex_ft = self.vocab.get("lex_ft", {})
        self.lex_f = self.vocab.get("lex_f", {})
        self.nongreek = self.vocab.get("nongreek", {})
        self.fallback_xpos = self.vocab.get("fallback_xpos", "u--------")
        self.fallback_upos = self.vocab.get("fallback_upos", "u")

        if self.xpos_alpha and self.tags:
            xid = [{c: i for i, c in enumerate(a)} for a in self.xpos_alpha]
            xpos_len = len(self.xpos_alpha)
            self.tag_idx = torch.tensor(
                [[xid[p].get(t[p] if p < len(t) else "-", 0) for p in range(xpos_len)]
                 for t in self.tags],
                dtype=torch.long,
            )
        else:
            self.tag_idx = None
        self._script_pc = np.array([s[0] + s[2] for s in self.scripts]) if self.scripts else None
        self._app_cache = {}

    # ---------------------------------------------------------------- construction

    @classmethod
    def from_pretrained(cls, path, **_kwargs):
        path = Path(path)
        vocab = json.loads((path / "vocab.json").read_text(encoding="utf-8"))
        deprel_vocab = json.loads((path / "deprel_vocab.json").read_text(encoding="utf-8"))
        return cls(vocab, deprel_vocab)

    def save_pretrained(self, save_directory, **_kwargs):
        save_directory = Path(save_directory)
        save_directory.mkdir(parents=True, exist_ok=True)
        (save_directory / "vocab.json").write_text(json.dumps(self.vocab, ensure_ascii=False))
        (save_directory / "deprel_vocab.json").write_text(json.dumps(self.deprel_vocab, ensure_ascii=False))
        (save_directory / "processor_config.json").write_text(
            json.dumps({"processor_class": "CharBertJointProcessor"}))

    @staticmethod
    def tokenize(text: str) -> list[str]:
        """Naive plain-text -> token-list convenience: whitespace split + peel a run of
        trailing sentence/clause punctuation off each whitespace-chunk into its own token.
        Good enough for a short usage-example sentence; real corpora should be tokenized
        properly (this is not a substitute for that)."""
        out = []
        for chunk in text.strip().split():
            m = _TRAIL_PUNCT_RE.match(chunk)
            if m and m.group(1):
                out.append(m.group(1))
                out.extend(list(m.group(2)))
            else:
                out.append(chunk)
        return out

    # ---------------------------------------------------------------- encode

    def __call__(self, sentences: list[str] | list[list[str]]):
        if not sentences:
            raise ValueError("sentences must be non-empty")
        if isinstance(sentences[0], str):
            sentences = [sentences]   # single sentence -> batch of 1

        encs = [encode_sentence(s) for s in sentences]
        B = len(encs)
        Tmax = max(len(e[0]) for e in encs)
        Tmax = max(Tmax, 1)

        ids = np.full((B, Tmax), PAD_ID, dtype=np.int64)
        bnd = np.zeros((B, Tmax), dtype=np.int64)
        dia = np.zeros((B, Tmax), dtype=np.int64)
        pct = np.zeros((B, Tmax), dtype=np.int64)
        cp = np.zeros((B, Tmax), dtype=np.int64)
        seg = np.zeros((B, Tmax), dtype=np.int64)
        wid = np.full((B, Tmax), -1, dtype=np.int64)
        all_forms = []
        for b, (c, bd, d, p, ca, wi, forms) in enumerate(encs):
            n = len(c)
            if n:
                ids[b, :n] = c
                bnd[b, :n] = bd
                dia[b, :n] = d
                pct[b, :n] = p
                cp[b, :n] = ca
                seg[b, :n] = 1
                wid[b, :n] = wi
            all_forms.append(forms)

        batch = dict(
            input_ids=torch.from_numpy(ids),
            boundary=torch.from_numpy(bnd),
            dia=torch.from_numpy(dia),
            punct=torch.from_numpy(pct),
            cap=torch.from_numpy(cp),
            seg_id=torch.from_numpy(seg),
            word_id=torch.from_numpy(wid),
        )
        batch["_forms"] = all_forms   # kept out-of-band: not a model input, needed for decoding
        return batch

    # ---------------------------------------------------------------- decode: xpos / upos

    def decode_xpos(self, xpos_logits, word_mask, flat_logits=None):
        """Constrained decode: combine the 9 factored heads' log-probs (+ the flat full-tag
        head's log-probs, if present) over the attested-tag inventory, argmax. Mirrors
        tagger.decode.TagDecoder.xpos exactly."""
        assert self.tag_idx is not None, "vocab.json (xpos_alpha/tags) required for xpos decode"
        ti = self.tag_idx.to(xpos_logits[0].device)
        score = 0
        for p, lg in enumerate(xpos_logits):
            lp = torch.log_softmax(lg.float(), -1)
            score = score + lp[:, :, ti[:, p]]
        if flat_logits is not None:
            score = score + torch.log_softmax(flat_logits.float(), -1)
        best = score.argmax(-1).cpu()
        return [[self.tags[int(best[b, w])]
                 for w in range(word_mask.shape[1]) if word_mask[b, w]]
                for b in range(word_mask.shape[0])]

    def decode_upos(self, upos_logits, word_mask):
        pred = upos_logits.argmax(-1).cpu()
        return [[self.upos[int(pred[b, w])]
                 for w in range(word_mask.shape[1]) if word_mask[b, w]]
                for b in range(word_mask.shape[0])]

    # ---------------------------------------------------------------- decode: lemma

    def _applicable(self, L):
        if L not in self._app_cache:
            self._app_cache[L] = self._script_pc <= L
        return self._app_cache[L]

    def decode_lemma(self, form: str, script_logprobs: torch.Tensor, xpos: str | None = None,
                      use_lexicon: bool = True, topk: int = 64) -> str:
        """script_logprobs: (n_scripts,) log-softmax for this word. Mirrors
        tagger.decode.LemmaDecoder.__call__ exactly (lexicon-rescored edit script, with an
        open-script fallback for OOV forms)."""
        key = form_key(form)
        copy_cap = form[:1] != form[:1].lower()

        app = torch.from_numpy(self._applicable(len(key))).to(script_logprobs.device)
        masked = script_logprobs.masked_fill(~app, -1e30)
        k = min(topk, masked.shape[-1])
        topv, topi = masked.topk(k)
        topv, topi = topv.tolist(), topi.tolist()

        if use_lexicon:
            cands = None
            if xpos is not None:
                cands = self.lex_ft.get(key + "\t" + xpos)
            if not cands:
                cands = self.lex_f.get(key, {})
            cands = dict(cands)
            if cands:
                lower = {}
                for lemma in cands:
                    lower.setdefault(lemma.lower(), lemma)
                best, best_s = None, -1e30
                for s, i in zip(topv, topi):
                    if s <= -1e29:
                        break
                    out = apply_script(key, self.scripts[i])
                    lemma = lower.get(out)
                    if lemma is not None:
                        sc = s + 1e-3 * np.log1p(cands[lemma])
                        if sc > best_s:
                            best, best_s = lemma, sc
                return best if best is not None else max(cands, key=cands.get)

        sc = self.scripts[topi[0]]
        lemma = apply_script(key, sc)
        if lemma is None:
            return form
        if sc[4] or copy_cap:
            lemma = lemma[:1].upper() + lemma[1:]
        return lemma

    # ---------------------------------------------------------------- decode: dependency arcs

    @staticmethod
    def decode_arcs(arc_scores, rel_scores):
        """Greedy per-token argmax head (0=root) + label argmax at the chosen head. Mirrors
        parser.biaffine.BiaffineHead.decode exactly."""
        heads_out = arc_scores.argmax(-1)
        B, W = heads_out.shape
        bi = torch.arange(B, device=arc_scores.device)[:, None].expand(B, W)
        wi = torch.arange(W, device=arc_scores.device)[None, :].expand(B, W)
        labels_out = rel_scores[bi, wi, heads_out].argmax(-1)
        return heads_out.cpu(), labels_out.cpu()

    # ---------------------------------------------------------------- full decode

    def decode(self, model_out, batch, use_lexicon: bool = True, ud: bool = False) -> list[list[dict]]:
        """model_out: a CharBertJointOutput (or the equivalent return_dict=False tuple).
        batch: the dict returned by __call__ (needs `_forms`).
        -> per sentence, a list of dicts, word order.

        Default (ud=False): the model's NATIVE Perseus/AGDT-style output --
        {form, xpos, lemma, upos, head, deprel}, where `upos` is the Perseus single-letter
        code (not UD) and `deprel` is a Prague-style relation (SBJ/OBJ/ATR/COORD/AuxP/...).

        ud=True: Universal Dependencies-style output instead -- {form, lemma, upos (UD),
        xpos (native tag, kept as-is per UD convention), feats (dict), head, deprel (UD)}.
        See `to_ud()` for the conversion (a documented, but not 100%-official-treebank-exact,
        port of the HamleDT/UD_Ancient_Greek-Perseus conversion methodology)."""
        word_mask = model_out.word_mask if hasattr(model_out, "word_mask") else model_out[-1]
        xpos_logits = model_out.xpos_logits if hasattr(model_out, "xpos_logits") else model_out[0]
        script_logits = model_out.script_logits if hasattr(model_out, "script_logits") else model_out[1]
        upos_logits = model_out.upos_logits if hasattr(model_out, "upos_logits") else model_out[2]
        flat_logits = model_out.flat_logits if hasattr(model_out, "flat_logits") else model_out[3]
        arc_scores = model_out.arc_scores if hasattr(model_out, "arc_scores") else model_out[4]
        rel_scores = model_out.rel_scores if hasattr(model_out, "rel_scores") else model_out[5]

        xpos_pred = self.decode_xpos(xpos_logits, word_mask, flat_logits=flat_logits)
        upos_pred = self.decode_upos(upos_logits, word_mask)
        script_logprobs = torch.log_softmax(script_logits.float(), -1)   # (B,W,n_script)

        heads_out, labels_out = (None, None)
        if arc_scores is not None:
            heads_out, labels_out = self.decode_arcs(arc_scores, rel_scores)

        forms_batch = batch["_forms"]
        results = []
        for b, forms in enumerate(forms_batch):
            n = len(forms)
            sent_out = []
            for w in range(n):
                form = forms[w]
                xpos = xpos_pred[b][w]
                upos = upos_pred[b][w]
                lemma = self.decode_lemma(form, script_logprobs[b, w], xpos=xpos, use_lexicon=use_lexicon)
                head = int(heads_out[b, w]) if heads_out is not None else None
                deprel = (self.deprel_vocab[int(labels_out[b, w])]
                          if labels_out is not None and self.deprel_vocab else None)
                sent_out.append(dict(form=form, xpos=xpos, lemma=lemma, upos=upos,
                                      head=head, deprel=deprel))
            results.append(to_ud(sent_out) if ud else sent_out)
        return results