plane: preserve_diacritics mode (WO02)

src/interscript/ml/plane.py

@@ -27,10 +27,43 @@ def __init__(self, graph: bytes, classes: list[str], k_passes: int = 2):
         self.mask_id = self.n_classes
         self.k = k_passes
         self.sess = ort.InferenceSession(graph, providers=["CPUExecutionProvider"])
+        # a character is a diacritic iff it occurs inside some class the
+        # artifact knows — no per-language mark tables needed
+        self._mark_chars = {ch for cls in classes for ch in cls}
+        self._class_to_id = {cls: i for i, cls in enumerate(classes)}
 
     def encode(self, text: str) -> list[int]:
         return [b + 3 for b in text.encode("utf-8")]
 
+    def _split_planes(self, text: str) -> tuple[str, list[str]]:
+        """Text -> (skeleton, per-base classes). Marks FOLLOW their base
+        letter; leading marks bind to a '\\x00' anchor (render drops it).
+        Cluster order is preserved as written — render-exactness beats
+        normalization (the Hebrew canon lesson)."""
+        skeleton: list[str] = []
+        classes: list[str] = []
+        current: list[str] = []
+
+        def close() -> None:
+            if not current:
+                return
+            if skeleton:
+                classes[-1] = classes[-1] + "".join(current)
+            else:
+                skeleton.append("\x00")
+                classes.append("".join(current))
+            current.clear()
+
+        for ch in text:
+            if ch in self._mark_chars:
+                current.append(ch)
+                continue
+            close()
+            skeleton.append(ch)
+            classes.append("")
+        close()
+        return "".join(skeleton), classes
+
     def predict_token_classes(self, text: str) -> list[int]:
         ids = np.array([self.encode(text)], dtype=np.int64)
         plane = np.full_like(ids, self.mask_id)
@@ -39,18 +72,62 @@ def predict_token_classes(self, text: str) -> list[int]:
             plane = logits.argmax(-1)
         return plane[0].tolist()
 
-    def translate(self, text: str) -> str:
-        tok_preds = self.predict_token_classes(text)
-        # majority vote of a char's byte-token predictions
+    def _decode_pinned(self, skeleton: str, char_classes: list[str]) -> list[int]:
+        """K-pass decode with user classes pinned: pinned positions are
+        conditioned AND never overwritten."""
+        ids = np.array([self.encode(skeleton)], dtype=np.int64)
+        per_token_pin: list[int] = []
+        for ch, cls in zip(skeleton, char_classes):
+            pid = self._class_to_id.get(cls, self.mask_id)
+            per_token_pin.extend([pid] * len(ch.encode("utf-8")))
+        pin_arr = np.array(per_token_pin, dtype=np.int64)
+        pinned = pin_arr != self.mask_id
+        plane = np.full_like(ids, self.mask_id)
+        if pinned.any():
+            plane[0, pinned] = pin_arr[pinned]
+        for _ in range(self.k):
+            logits = self.sess.run(None, {"input_ids": ids, "plane_ids": plane})[0]
+            pred = logits.argmax(-1)
+            if pinned.any():
+                pred[0, pinned] = pin_arr[pinned]
+            plane = pred
+        return plane[0].tolist()
+
+    def translate(self, text: str, preserve_diacritics: bool = False) -> str:
+        if not preserve_diacritics:
+            tok_preds = self.predict_token_classes(text)
+            pos = 0
+            char_classes: list[str] = []
+            for ch in text:
+                n = len(ch.encode("utf-8"))
+                votes = tok_preds[pos : pos + n]
+                cid = max(set(votes), key=votes.count)
+                char_classes.append(self.classes[cid] if cid < self.n_classes else "")
+                pos += n
+            return render_plane(text, char_classes)
+
+        # preserve path: user diacritics are pinned inside the decode and
+        # round-trip byte-exactly (leading marks included)
+        skeleton, char_classes = self._split_planes(text)
+        tok_preds = self._decode_pinned(skeleton, char_classes)
         pos = 0
-        char_classes: list[str] = []
-        for ch in text:
+        final: list[str] = []
+        for ch, user_cls in zip(skeleton, char_classes):
             n = len(ch.encode("utf-8"))
-            votes = tok_preds[pos : pos + n]
-            cid = max(set(votes), key=votes.count)
-            char_classes.append(self.classes[cid] if cid < self.n_classes else "")
+            if user_cls:
+                final.append(user_cls)
+            else:
+                votes = tok_preds[pos : pos + n]
+                cid = max(set(votes), key=votes.count)
+                final.append(self.classes[cid] if cid < self.n_classes else "")
             pos += n
-        return render_plane(text, char_classes)
+        out: list[str] = []
+        for ch, cls in zip(skeleton, final):
+            if ch == "\x00":
+                out.append(cls)  # anchor for leading marks: emit class only
+            else:
+                out.append(ch + cls)
+        return "".join(out)
 
 
 def from_zip(data: bytes) -> PlaneModel:

tests/test_plane_preserve.py

@@ -0,0 +1,94 @@
+"""WO02: user-diacritics preservation in the plane runtime.
+
+Builds a REAL deterministic ONNX graph (no doubles): class_logits =
+one_hot[(input_id + 7*plane_id) % TABLE], so K-pass cycling is fully
+predictable and every assertion is on actual model output.
+"""
+
+from __future__ import annotations
+
+import io
+
+import numpy as np
+import onnx
+from onnx import TensorProto, helper, numpy_helper
+
+TABLE = 512
+N_CLASSES = 3
+CLASSES = ["́", "̀", "́̀"]  # synthetic combining-mark classes
+
+
+def _graph_bytes() -> bytes:
+    table = np.zeros((TABLE, N_CLASSES), dtype=np.float32)
+    table[np.arange(TABLE), np.arange(TABLE) % N_CLASSES] = 1.0
+    init = numpy_helper.from_array(table, "table")
+    ids = helper.make_tensor_value_info("input_ids", TensorProto.INT64, ["B", "S"])
+    plane = helper.make_tensor_value_info("plane_ids", TensorProto.INT64, ["B", "S"])
+    out = helper.make_tensor_value_info("class_logits", TensorProto.FLOAT, ["B", "S", str(N_CLASSES)])
+    idx = helper.make_node("Add", ["input_ids", "plane7"], ["sum"])
+    mul = helper.make_node("Mul", ["sum", "plane7"], ["scaled"])
+    mod = helper.make_node("Mod", ["scaled", "tablen"], ["idx"], fmod=0)
+    gather = helper.make_node("Gather", ["table", "idx"], ["class_logits"], axis=0)
+    graph = helper.make_graph(
+        [idx, mul, mod, gather], "tiny_plane",
+        [ids, plane], [out],
+        initializer=[init,
+                     numpy_helper.from_array(np.array(7, dtype=np.int64), "plane7"),
+                     numpy_helper.from_array(np.array(TABLE, dtype=np.int64), "tablen")],
+    )
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    model.ir_version = 8
+    onnx.checker.check_model(model)
+    buf = io.BytesIO()
+    onnx.save(model, buf)
+    return buf.getvalue()
+
+
+def make_model():
+    from interscript.ml.plane import PlaneModel
+
+    return PlaneModel(graph=_graph_bytes(), classes=CLASSES, k_passes=2)
+
+
+def test_default_path_unchanged_by_flag():
+    m = make_model()
+    assert m.translate("abc") == m.translate("abc", preserve_diacritics=False)
+
+
+def test_preserve_identity_fully_labeled():
+    m = make_model()
+    labeled = "áb̀"
+    assert m.translate(labeled, preserve_diacritics=True) == labeled
+
+
+def test_preserve_partial_keeps_user_classes_and_fills_rest():
+    m = make_model()
+    out = m.translate("a\u0301bc", preserve_diacritics=True)
+    assert out.startswith("a\u0301")  # user class kept verbatim regardless of model cycle
+    marks = "".join(CLASSES)
+    assert "".join(ch for ch in out if ch not in marks) == "abc"  # bases in order
+
+
+def test_preserve_leading_marks_roundtrip():
+    m = make_model()
+    labeled = "\u0301ab"
+    out = m.translate(labeled, preserve_diacritics=True)
+    assert out.startswith(labeled[0])  # leading mark kept verbatim
+    assert "\x00" not in out  # anchor never leaks into output
+    marks = "".join(CLASSES)
+    assert "".join(ch for ch in out if ch not in marks) == "ab"
+
+
+def test_preserve_multibyte_base_chars():
+    m = make_model()
+    labeled = "\u00e9\u0301"  # precomposed two-byte base + mark
+    assert m.translate(labeled, preserve_diacritics=True) == labeled
+
+
+def test_plain_text_still_translates():
+    m = make_model()
+    out = m.translate("abc")
+    marks = "".join(CLASSES)
+    assert set(out) <= set("abc") | set(marks)
+    assert "".join(ch for ch in out if ch not in marks) == "abc"
+    assert len(m.predict_token_classes("abc")) == 3