diff --git a/src/interscript/ml/plane.py b/src/interscript/ml/plane.py index 3df7f2b..cd24ac7 100644 --- a/src/interscript/ml/plane.py +++ b/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: diff --git a/tests/test_plane_preserve.py b/tests/test_plane_preserve.py new file mode 100644 index 0000000..7805378 --- /dev/null +++ b/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