Publish safetensors weights, config and model card

README.md

@@ -1,17 +1,50 @@
 ---
 license: apache-2.0
-language: [grc]
-tags: [ancient-greek, masked-diffusion, fine-tuned]
+language:
+- grc
+library_name: transformers
+tags:
+- ancient-greek
+- classical-philology
+- character-level
+- masked-diffusion
+- text-restoration
+- epigraphy
+- papyrology
+pipeline_tag: fill-mask
 ---
-# Stoicheia-restoration-test4
 
-Documentary restoration, PHI/TM final digit 4 held out. Part of the Stoicheia release: a 405M-parameter character-level
-masked-diffusion Transformer for Ancient Greek whose input factors into five
-independently maskable planes (letters, word/sentence boundaries, diacritics,
-capitalization, punctuation).
+# Stoicheia -- documentary restoration (digit 4 held out)
 
-Anonymous release accompanying a paper under review. Code, usage examples, and
-reproduction instructions: https://huggingface.co/anonymous-stoicheia/Stoicheia-code
+**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.
 
-Checkpoint format: PyTorch state dict (`torch.load`), architecture and loading
-code in the accompanying code repository (`model/char_bert.py`).
+Anonymous release accompanying a paper under review.
+
+Fine-tuned from `Stoicheia-doc_clean` to restore damaged inscriptions and papyri, on every
+document **except** those whose PHI/TM identifier ends in **4**.
+
+Ten such checkpoints are released, one per digit, because a fixed split makes a model useless
+for exactly the documents an editor cares about: whatever inscription or papyrus you are working
+on, one of the ten has provably never read it. Pick that one and the reading it proposes cannot
+be a memory of the edition you are trying to check.
+
+## Usage
+
+```python
+import torch
+from transformers import AutoModel
+model = AutoModel.from_pretrained("anonymous-stoicheia/Stoicheia-restoration-test4", trust_remote_code=True).eval()
+
+from processing_char_bert import CharBertProcessor
+proc = CharBertProcessor()
+# "[7]" is a lacuna of known width; "[7±2]" if the width itself is uncertain
+text = "αγαθηιτυχηιεδοξεντ[7]βουληικαιτωιδημωι"
+best, width, cands = proc.restore_elastic(model, text, mask_dia_boundary=True)
+print(best)
+```

config.json

@@ -0,0 +1,21 @@
+{
+  "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,
+  "n_heads": 16,
+  "char_window": 256,
+  "attn_impl": "sdpa",
+  "qk_norm": true,
+  "d_model": 1024,
+  "depth": 32,
+  "model_type": "char_bert",
+  "auto_map": {
+    "AutoConfig": "configuration_char_bert.CharBertConfig",
+    "AutoModel": "modeling_char_bert.CharBertModel"
+  }
+}

configuration_char_bert.py

@@ -0,0 +1,40 @@
+"""HF-Hub-compatible config for Stoicheia (CharBertEncoder)."""
+from transformers import PretrainedConfig
+
+
+class CharBertConfig(PretrainedConfig):
+    model_type = "char_bert"
+
+    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,
+        **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
+        super().__init__(**kwargs)

model.safetensors

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

modeling_char_bert.py

@@ -0,0 +1,210 @@
+"""HF-Hub-compatible model for Stoicheia (CharBertEncoder).
+
+Self-contained: vendors the transformer primitives (RMSNorm/RoPE/Attention/GeGLU/Block)
+so this file has no dependency on the original research repo. Uses the SDPA attention
+path only (portable to CPU and any CUDA GPU) -- the original training code also supports
+a compiled FlexAttention block-sparse path for long packed sequences on GPU, which is not
+needed for standalone inference on single passages and is omitted here for portability.
+"""
+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 import CharBertConfig
+
+
+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
+
+
+@dataclass
+class CharBertOutput(ModelOutput):
+    char: torch.FloatTensor = None
+    boundary: torch.FloatTensor = None
+    dia: torch.FloatTensor = None
+    cap: torch.FloatTensor = None
+    punct: torch.FloatTensor = None
+    hidden_states: Optional[tuple] = None
+
+
+class CharBertModel(PreTrainedModel):
+    config_class = CharBertConfig
+
+    def __init__(self, config: CharBertConfig):
+        super().__init__(config)
+        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)
+        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)
+        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.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,
+        seg_id: Optional[torch.LongTensor] = None,
+        output_hidden_states: bool = False,
+        return_dict: bool = True,
+        **kwargs,
+    ):
+        cfg = self.config
+        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)
+        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)
+
+        hidden_states = [] if output_hidden_states else None
+        for blk in self.blocks:
+            m = glob_mask if blk.window == 0 else attn_mask
+            x = blk(x, pos, m)
+            if output_hidden_states:
+                hidden_states.append(x)
+
+        x = self.norm_out(x)
+        out = dict(
+            char=self.head_char(x),
+            boundary=self.head_bnd(x),
+            dia=self.head_dia(x),
+            cap=self.head_cap(x),
+            punct=self.head_punct(x),
+        )
+        if output_hidden_states:
+            out["hidden_states"] = tuple(hidden_states) + (x,)
+        if not return_dict:
+            return tuple(v for v in out.values() if v is not None)
+        return CharBertOutput(**out)

processing_char_bert.py

@@ -0,0 +1,346 @@
+"""HF-Hub-compatible processor for Stoicheia: text <-> the model's four input planes.
+
+Wraps the reference normalization/denormalization logic (character classification,
+diacritic packing, word/sentence-boundary detection) into a single callable that
+produces model-ready tensors, plus decode helpers for the three masking use cases
+described in the model card (restoration, accent recovery, re-segmentation).
+
+This is intentionally NOT a `PreTrainedTokenizer` subclass: the underlying encoding is
+a row-per-letter, four-parallel-plane structure (not a single token-id stream), which
+doesn't fit that base class's assumptions. It follows the same `register_for_auto_class`
+mechanism transformers uses for tokenizers/feature extractors, 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 dataclasses import dataclass, field
+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}
+_MARK_CHARS = {"acute": "́", "grave": "̀", "circ": "͂",
+               "smooth": "̓", "rough": "̔", "iota": "ͅ", "diaer": "̈"}
+# punct plane classes (matches data/normalize.py's canonical 6-way scheme):
+# 0 none, 1 comma, 2 high-dot(·), 3 colon, 4 period, 5 question/exclamation
+_PUNCT_CHARS = {1: ",", 2: "·", 3: ":", 4: ".", 5: ";"}
+
+# "λόγ[5±3]καὶ" -- a lacuna of uncertain width: best guess 5 letters, plausible
+# range 5-3..5+3. See CharBertProcessor.restore_elastic.
+_ELASTIC_RE = re.compile(r"\[(\d+)±(\d+)\]")
+
+
+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
+
+
+@dataclass
+class _Encoded:
+    chars: list  # int, may include MASK
+    boundary: list
+    dia: list
+    punct: list
+    cap: list  # original capitalization, for round-tripping non-masked positions
+
+
+class CharBertProcessor:
+    """`processor(text)` -> dict of batched tensors ready for `CharBertModel(**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": "CharBertProcessor"}))
+
+    # ---------------------------------------------------------------- encode
+
+    def _classify(self, text: str) -> _Encoded:
+        """Turn raw NFC/NFD polytonic text into per-letter plane lists, damage ('-' runs)
+        preserved as MASK/UNK positions, gold values kept for every other position."""
+        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)  # filled in by trailing combining marks below
+                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, mask_planes: list[str] | None = None, has_boundaries: bool = True):
+        """Encode `text` into model-ready tensors.
+
+        mask_planes: any subset of {"chars", "boundary", "dia", "punct"} to force to
+            UNKNOWN at every position (in addition to any '-' runs, which are always
+            treated as a damaged/masked span regardless of mask_planes).
+        has_boundaries: set False for scriptio continua input (no real spaces) so the
+            boundary plane starts fully UNKNOWN rather than "no boundaries found".
+        """
+        mask_planes = set(mask_planes or [])
+        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)
+
+        if "chars" in mask_planes:
+            chars[:] = MASK
+        if "boundary" in mask_planes or not has_boundaries:
+            boundary[:] = UNK_BND
+        if "dia" in mask_planes:
+            dia[:] = UNK_DIA
+        if "punct" in mask_planes:
+            punct[:] = UNK_PUNCT
+
+        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],
+            seg_id=torch.zeros(1, n, dtype=torch.long),
+        )
+        batch["_cap"] = enc.cap  # kept out-of-band; not a model input (cap is output-only)
+        return batch
+
+    # ---------------------------------------------------------------- decode
+
+    @staticmethod
+    def _restore_polytonic(chars, dia, cap, boundary, punct=None) -> str:
+        words, cur = [], []
+        for i in range(len(chars)):
+            ch = ID2LETTER[chars[i]] if chars[i] < 24 else "?"
+            a, b, io, dd = _unpack_dia(int(dia[i]))
+            if cap[i]:
+                ch = ch.upper()
+            s = ch
+            if b:
+                s += _MARK_CHARS[{1: "smooth", 2: "rough"}[b]]
+            if dd:
+                s += _MARK_CHARS["diaer"]
+            if a:
+                s += _MARK_CHARS[{1: "acute", 2: "grave", 3: "circ"}[a]]
+            if io:
+                s += _MARK_CHARS["iota"]
+            cur.append(s)
+            p = int(punct[i]) if punct is not None else 0
+            if boundary[i] >= 1:
+                w = "".join(cur)
+                if w and w[-1] == "σ":
+                    w = w[:-1] + "ς"
+                w = unicodedata.normalize("NFC", w)
+                if p in _PUNCT_CHARS:
+                    w += _PUNCT_CHARS[p]
+                elif boundary[i] == 2:
+                    w += "."
+                words.append(w)
+                cur = []
+        if cur:
+            w = unicodedata.normalize("NFC", "".join(cur))
+            p = int(punct[-1]) if punct is not None else 0
+            if p in _PUNCT_CHARS:
+                w += _PUNCT_CHARS[p]
+            words.append(w)
+        return " ".join(words)
+
+    def decode_restoration(self, model_out, batch) -> str:
+        """Fill masked positions with the model's argmax predictions; keep every
+        other position exactly as given. Each plane is filled independently
+        wherever IT is unknown -- chars/cap only inside a '-' gap (chars==MASK),
+        but boundary/dia/punct wherever THAT plane is UNK, which may be the whole
+        sequence if mask_planes was also used for joint gap+accent+boundary
+        restoration (not just the '-' gap itself)."""
+        pred_char = model_out.char.argmax(-1)[0].tolist()
+        pred_bnd = model_out.boundary.argmax(-1)[0].tolist()
+        pred_dia = model_out.dia.argmax(-1)[0].tolist()
+        pred_cap = model_out.cap.argmax(-1)[0].tolist()
+        pred_punct = model_out.punct.argmax(-1)[0].tolist()
+        chars = batch["input_ids"][0].tolist()
+        boundary = batch["boundary"][0].tolist()
+        dia = batch["dia"][0].tolist()
+        punct = batch["punct"][0].tolist()
+        cap = batch["_cap"]
+        for i in range(len(chars)):
+            if chars[i] == MASK:
+                chars[i] = pred_char[i] if pred_char[i] < 24 else 0
+                cap[i] = pred_cap[i]
+            if boundary[i] == UNK_BND:
+                boundary[i] = pred_bnd[i]
+            if dia[i] == UNK_DIA:
+                dia[i] = pred_dia[i]
+            if punct[i] == UNK_PUNCT:
+                punct[i] = pred_punct[i]
+        return self._restore_polytonic(chars, dia, cap, boundary, punct)
+
+    def decode_diacritics(self, model_out, batch) -> str:
+        """Replace the diacritic plane with the model's predictions; letters/boundaries/
+        capitalization/punctuation are taken from the input as given."""
+        pred_dia = model_out.dia.argmax(-1)[0].tolist()
+        chars = batch["input_ids"][0].tolist()
+        boundary = batch["boundary"][0].tolist()
+        punct = batch["punct"][0].tolist()
+        cap = batch["_cap"]
+        return self._restore_polytonic(chars, pred_dia, cap, boundary, punct)
+
+    def decode_boundaries(self, model_out, batch) -> str:
+        """Replace the boundary plane with the model's predictions (0/1/2); letters/
+        diacritics/capitalization/punctuation are taken from the input as given.
+        Only useful when the input truly has no accents either (a spaced-out or
+        scriptio-continua text that already carries accents gives the boundary head
+        a strong shortcut -- each word carries exactly one accent -- so this isn't a
+        meaningful standalone test of the boundary head specifically; see
+        decode_restoration/restore_elastic for the realistic joint case)."""
+        pred_bnd = model_out.boundary.argmax(-1)[0].tolist()
+        chars = batch["input_ids"][0].tolist()
+        dia = batch["dia"][0].tolist()
+        punct = batch["punct"][0].tolist()
+        cap = batch["_cap"]
+        return self._restore_polytonic(chars, dia, cap, pred_bnd, punct)
+
+    def restore_elastic(self, model, text: str, min_width: int = 1,
+                         mask_dia_boundary: bool = False):
+        """Restore a lacuna of *uncertain* width -- the realistic editorial case,
+        since editors estimate a lacuna's length, they rarely know it exactly.
+
+        `text` must contain exactly one `[N±M]` marker (best-guess width N,
+        plausible range N-M..N+M), e.g. `"λόγ[5±3]καὶ ὁ λόγος ἦν πρὸς τὸν θεόν"`.
+        For every candidate width in that range, this fills the gap, then scores
+        each candidate by the mean log-probability of the model's own letter
+        predictions inside the gap specifically (that's what distinguishes widths).
+
+        `mask_dia_boundary` controls what happens OUTSIDE the gap:
+          - False (default): real accents/word-boundaries already present in
+            `text` are kept as given -- only the gap itself is filled. Use this
+            for text where the surrounding context is already known/accented (the
+            common editorial case: a lacuna in an otherwise-legible inscription).
+          - True: accents and word-boundaries are masked and reconstructed
+            everywhere, not just inside the gap -- for fully bare scriptio
+            continua surrounding the lacuna too (no spaces, no accents at all).
+
+        Returns `(best_text, best_width, candidates)`, where `candidates` is every
+        `(width, filled_text, mean_logp)` tried, sorted best-first. Needs the model
+        (not just its output), since it runs one forward pass per candidate width.
+        """
+        m = _ELASTIC_RE.search(text)
+        if not m:
+            raise ValueError("text must contain one '[N±M]' marker, e.g. 'λόγ[5±3]καὶ'")
+        n, spread = int(m.group(1)), int(m.group(2))
+        prefix, suffix = text[:m.start()], text[m.end():]
+        gap_start = len(self._classify(prefix).chars)
+
+        candidates = []
+        for L in range(max(min_width, n - spread), n + spread + 1):
+            probe = prefix + ("-" * L) + suffix
+            if mask_dia_boundary:
+                batch = self(probe, mask_planes=["dia", "boundary"], has_boundaries=False)
+            else:
+                batch = self(probe)
+            with torch.no_grad():
+                out = model(**{k: v for k, v in batch.items() if not k.startswith("_")})
+
+            logp = torch.log_softmax(out.char, dim=-1)[0]
+            pred_char = out.char.argmax(-1)[0].tolist()
+            gap_logp = sum(logp[gap_start + i, pred_char[gap_start + i]].item()
+                            for i in range(L)) / L
+
+            pred_bnd = out.boundary.argmax(-1)[0].tolist()
+            pred_dia = out.dia.argmax(-1)[0].tolist()
+            pred_cap = out.cap.argmax(-1)[0].tolist()
+            pred_punct = out.punct.argmax(-1)[0].tolist()
+            chars = batch["input_ids"][0].tolist()
+            boundary = batch["boundary"][0].tolist()
+            dia = batch["dia"][0].tolist()
+            punct = batch["punct"][0].tolist()
+            cap = batch["_cap"]
+            for i in range(len(chars)):
+                if chars[i] == MASK:
+                    chars[i] = pred_char[i] if pred_char[i] < 24 else 0
+                    cap[i] = pred_cap[i]
+                if boundary[i] == UNK_BND:
+                    boundary[i] = pred_bnd[i]
+                if dia[i] == UNK_DIA:
+                    dia[i] = pred_dia[i]
+                if punct[i] == UNK_PUNCT:
+                    punct[i] = pred_punct[i]
+            filled = self._restore_polytonic(chars, dia, cap, boundary, punct)
+            candidates.append((L, filled, gap_logp))
+
+        candidates.sort(key=lambda c: -c[2])
+        best_L, best_text, _ = candidates[0]
+        return best_text, best_L, candidates

training_metadata.json

@@ -0,0 +1,54 @@
+{
+  "name": "whole_v4_t4v5",
+  "test_digit": "4",
+  "val_digit": "5",
+  "seed": 0,
+  "attn": "flex",
+  "compile": true,
+  "seq_len": 8192,
+  "micro_batch": 4,
+  "grad_accum": 2,
+  "lr": 0.0001,
+  "wd": 0.1,
+  "warmup": 200,
+  "total_steps": 6000,
+  "decay_len": 1500,
+  "eval_every": 1000000000,
+  "eval_n": 256,
+  "dev_cer_n_per_L": 60,
+  "dev_cer_L_max": 20,
+  "stall_window": 8,
+  "stall_eps": 0.002,
+  "anneal_patience": 3,
+  "log_every": 20,
+  "lam": 0.1,
+  "meta_condition": true,
+  "min_len": 20,
+  "torso": "best.pt",
+  "mix": {
+    "iphi": 1.0,
+    "iphi_syn": 0.5,
+    "papyri": 1.0,
+    "iphi_seg": 0.5,
+    "pap_seg": 0.3,
+    "literary": 0.3,
+    "literary_n": 60000
+  },
+  "noise": {
+    "w_span": 0.45,
+    "w_word": 0.15,
+    "w_elastic": 0.05,
+    "w_iid": 0.15,
+    "w_halfword": 0.15,
+    "w_substitute": 0.05,
+    "span_mean": 4.0,
+    "span_max": 12,
+    "p_region_none": 0.3,
+    "p_century_none": 0.3
+  },
+  "dev_exclude": "",
+  "out_dir": "whole_v4_t4v5",
+  "save_every": 0,
+  "_source_checkpoint": "final.pt",
+  "_source_step": 6000
+}