feat: interscript.ml — unified maps+models transliterate

pyproject.toml

@@ -28,7 +28,8 @@ Homepage = "https://www.interscript.org"
 Issues = "https://github.com/interscript/interscript-py/issues"
 
 [project.optional-dependencies]
-dev = ["pytest>=7"]
+ml = ["numpy>=1.26", "pyyaml>=6.0", "onnxruntime>=1.17"]
+dev = ["pytest>=7", "onnx>=1.16"]
 
 [tool.setuptools.packages.find]
 where = ["src"]

src/interscript/interscript.py

@@ -79,5 +79,26 @@ def load_map(map_name: str, on_unsupported: str = "raise") -> Engine:
     return engine
 
 
-def transliterate(map_name: str, text: str) -> str:
-    return load_map(map_name).transliterate(text)
+def transliterate(id: str, text: str, index_url: str | None = None) -> str:
+    """Unified conversion surface: map ids run the rule engine; model
+    ids resolve through the interscript-ml index and run the neural
+    runtime (lazy — requires the interscript[ml] extra)."""
+    if map_exist(id):
+        return load_map(id).transliterate(text)
+    try:
+        from interscript.ml import Model, PlaneModel
+        from interscript.ml.registry import RegistryError, load_index, resolve
+    except ImportError as e:  # pragma: no cover - exercised when extra missing
+        raise FileNotFoundError(
+            f"id '{id}' is not a map, and the ML layer is not installed (pip install interscript[ml])"
+        ) from e
+    try:
+        entry = load_index(index_url)[id]
+        path = resolve(id, index_url=index_url)
+    except (RegistryError, KeyError) as e:
+        raise FileNotFoundError(
+            f"id '{id}' is neither a map in the load paths nor a model in the interscript-ml index"
+        ) from e
+    if entry.kind == "plane":
+        return PlaneModel.from_zip(Path(path).read_bytes()).translate(text)
+    return Model.load(path).translate(text)

src/interscript/ml/__init__.py

@@ -0,0 +1,33 @@
+"""interscript.ml — the neural layer: model index resolution and the
+IMF v1 / plane zip runtimes (seq2seq byte models + plane-factorized).
+
+Installed as an optional extra:  pip install interscript[ml].
+Maps transliteration lives in the core package; this module resolves
+model ids from the interscript-ml index and runs their ONNX graphs.
+
+    from interscript.ml import Model, PlaneModel
+    Model.load("khm-latn-1.0").translate("ភាសា")
+"""
+
+from interscript.ml.loader import Manifest, ModelFormatError
+from interscript.ml.model import Model
+from interscript.ml.plane import PlaneModel
+from interscript.ml.registry import RegistryError, cache_dir, load_index, resolve
+from interscript.ml.tokens import BYTE_OFFSET, EOS_ID, PAD_ID, UNK_ID, decode, encode
+
+__all__ = [
+    "BYTE_OFFSET",
+    "EOS_ID",
+    "Manifest",
+    "Model",
+    "ModelFormatError",
+    "PAD_ID",
+    "PlaneModel",
+    "RegistryError",
+    "UNK_ID",
+    "cache_dir",
+    "decode",
+    "encode",
+    "load_index",
+    "resolve",
+]

src/interscript/ml/loader.py

@@ -0,0 +1,68 @@
+"""IMF v1 zip loading: sha256 verification + extraction."""
+
+from __future__ import annotations
+
+import hashlib
+import zipfile
+from dataclasses import dataclass
+from pathlib import Path
+
+import yaml
+
+
+class ModelFormatError(ValueError):
+    """The zip is not a valid IMF v1 artifact (or fails integrity)."""
+
+
+@dataclass(frozen=True)
+class Manifest:
+    id: str
+    task: str
+    decoder: str
+    precision: str
+    opset: int
+    sha256: dict[str, str]
+
+
+def load_manifest(zip_path: Path | str) -> Manifest:
+    with zipfile.ZipFile(zip_path) as zf:
+        names = zf.namelist()
+        for required in ("metadata.yaml", "encoder.onnx", "decoder.onnx"):
+            if required not in names:
+                raise ModelFormatError(f"missing required file: {required}")
+        raw = yaml.safe_load(zf.read("metadata.yaml"))
+        if raw.get("format") != "imf-v1":
+            raise ModelFormatError(f"unsupported format: {raw.get('format')!r}")
+        if raw.get("tokenizer") != "bytes":
+            raise ModelFormatError(
+                f"tokenizer {raw.get('tokenizer')!r}: this runtime is byte-level only"
+            )
+        return Manifest(
+            id=raw["id"],
+            task=raw["task"],
+            decoder=raw.get("decoder", "plain"),
+            precision=raw.get("precision", "fp32"),
+            opset=int(raw.get("opset", 14)),
+            sha256=dict(raw.get("sha256", {})),
+        )
+
+
+def verify_and_read(zip_path: Path | str) -> dict[str, bytes]:
+    """Read .onnx members after verifying each sha256 against the
+    manifest — the corrupt-download failure mode fails loudly here."""
+    manifest = load_manifest(zip_path)
+    graphs: dict[str, bytes] = {}
+    with zipfile.ZipFile(zip_path) as zf:
+        for name in [n for n in zf.namelist() if n.endswith(".onnx")]:
+            member = zf.read(name)
+            recorded = manifest.sha256.get(name)
+            if recorded is None:
+                raise ModelFormatError(f"{name} is not covered by metadata sha256")
+            actual = hashlib.sha256(member).hexdigest()
+            if actual != recorded:
+                raise ModelFormatError(
+                    f"{name} sha256 mismatch: zip has {actual}, "
+                    f"metadata says {recorded}"
+                )
+            graphs[name] = member
+    return graphs

src/interscript/ml/model.py

@@ -0,0 +1,127 @@
+"""Model.load(zip) + translate(text): greedy KV decode with plain fallback."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+
+from interscript.ml.tokens import EOS_ID, PAD_ID, decode, encode
+from interscript.ml.loader import load_manifest, verify_and_read
+
+
+class Model:
+    """A loaded, checksum-verified IMF v1 model.
+
+    >>> model = Model.load("khm-latn-1.0.zip")
+    >>> model.translate("ភាសា")
+    """
+
+    def __init__(self, zip_path: Path | str):
+        self.zip_path = Path(zip_path)
+        self.manifest = load_manifest(self.zip_path)
+        import onnxruntime as ort
+
+        options = ort.SessionOptions()
+        options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+        graphs = verify_and_read(self.zip_path)
+        self._encoder = ort.InferenceSession(
+            graphs["encoder.onnx"], options, providers=_providers()
+        )
+        decoder_name = (
+            "decoder-kv.onnx"
+            if self.manifest.decoder == "kv" and "decoder-kv.onnx" in graphs
+            else "decoder.onnx"
+        )
+        self._kv_session = decoder_name == "decoder-kv.onnx"
+        self._decoder = ort.InferenceSession(
+            graphs[decoder_name], options, providers=_providers()
+        )
+        self._pasts = {
+            meta.name: _zero_past(meta)
+            for meta in self._decoder.get_inputs()
+            if meta.name.startswith("past_")
+        }
+        self._output_names = [o.name for o in self._decoder.get_outputs()]
+
+    @classmethod
+    def load(cls, path_or_id: Path | str, index_url: str | None = None) -> "Model":
+        """Accepts a zip path OR a model id from models.yaml (dynamic
+        fetch: download -> verify -> cache)."""
+        candidate = str(path_or_id)
+        if candidate.endswith(".zip") or Path(candidate).exists():
+            return cls(candidate)
+        from interscript.ml.registry import resolve
+
+        return cls(resolve(candidate, index_url))
+
+    @property
+    def id(self) -> str:
+        return self.manifest.id
+
+    def translate(self, text: str, max_len: int = 256) -> str:
+        token_ids = self.generate(text, max_len=max_len)
+        return decode(token_ids)
+
+    def generate(self, text: str, max_len: int = 256) -> list[int]:
+        ids = np.array([encode(text)], dtype=np.int64)
+        if ids.shape[1] == 1:  # only the trailing EOS: empty input
+            return []
+        hidden = self._encoder.run(None, {"input_ids": ids})[0]
+        if self._kv_session:
+            return self._greedy_kv(hidden, max_len)
+        return self._greedy_plain(hidden, max_len)
+
+    def _greedy_kv(self, hidden, max_len: int) -> list[int]:
+        pasts = dict(self._pasts)
+        current = np.array([[PAD_ID]], dtype=np.int64)
+        generated: list[int] = []
+        for _ in range(max_len):
+            outputs = self._decoder.run(
+                None,
+                {"input_ids": current, "encoder_hidden_states": hidden, **pasts},
+            )
+            results = dict(zip(self._output_names, outputs, strict=True))
+            token = int(np.argmax(results["logits"][0, -1]))
+            if token == EOS_ID:
+                break
+            generated.append(token)
+            pasts = {
+                name: results[name.replace("past_", "present_", 1)]
+                for name in pasts
+            }
+            current = np.array([[token]], dtype=np.int64)
+        return generated
+
+    def _greedy_plain(self, hidden, max_len: int) -> list[int]:
+        decoder_ids = np.array([[PAD_ID]], dtype=np.int64)
+        generated: list[int] = []
+        for _ in range(max_len):
+            logits = self._decoder.run(
+                None,
+                {"input_ids": decoder_ids, "encoder_hidden_states": hidden},
+            )[0]
+            token = int(np.argmax(logits[0, -1]))
+            if token == EOS_ID:
+                break
+            generated.append(token)
+            decoder_ids = np.concatenate(
+                [decoder_ids, np.array([[token]], dtype=np.int64)], axis=1
+            )
+        return generated
+
+
+def _providers() -> list[str]:
+    import onnxruntime as ort
+
+    available = ort.get_available_providers()
+    preferred = [p for p in ("CPUExecutionProvider",) if p in available]
+    return preferred or available
+
+
+def _zero_past(meta) -> object:
+    shape = meta.shape  # [batch, heads, past_seq, d_kv], dynamic dims are str
+    heads = shape[1] if isinstance(shape[1], int) else 4
+    d_kv = shape[3] if isinstance(shape[3], int) else 8
+    dtype = np.float16 if meta.type == "tensor(float16)" else np.float32
+    return np.zeros((1, heads, 0, d_kv), dtype=dtype)

src/interscript/ml/plane.py

@@ -0,0 +1,92 @@
+"""Plane model kind: the run-018 on-device artifact family.
+
+One ONNX graph maps (input_ids, plane_ids) -> per-position class
+logits; decode is a host-side K-pass loop (pass k>1 conditions on
+pass k-1 argmax); the haraqat plane renders onto the skeleton.
+Byte table: id = byte + 3 (the canonical ByT5 table); MASK id =
+len(classes).
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import onnxruntime as ort
+
+MASK_OFFSET = 0  # plane vocab = classes + MASK at index len(classes)
+
+
+def render_plane(skeleton: str, char_classes: list[str]) -> str:
+    """One class per CHARACTER (multi-byte chars consume one class)."""
+    return "".join(ch + cls for ch, cls in zip(skeleton, char_classes))
+
+
+class PlaneModel:
+    def __init__(self, graph: bytes, classes: list[str], k_passes: int = 2):
+        self.classes = classes
+        self.n_classes = len(classes)
+        self.mask_id = self.n_classes
+        self.k = k_passes
+        self.sess = ort.InferenceSession(graph, providers=["CPUExecutionProvider"])
+
+    def encode(self, text: str) -> list[int]:
+        return [b + 3 for b in text.encode("utf-8")]
+
+    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)
+        for _ in range(self.k):
+            logits = self.sess.run(None, {"input_ids": ids, "plane_ids": plane})[0]
+            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
+        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)
+
+
+def from_zip(data: bytes) -> PlaneModel:
+    """Load a plane artifact zip: metadata.yaml (id, precision, kind,
+    k_passes, classes, member sha256s) + plane.onnx + classes.json.
+    Every member is sha256-verified before the session is created."""
+    import hashlib
+    import io
+    import json as _json
+    import zipfile
+
+    import yaml
+
+    zf = zipfile.ZipFile(io.BytesIO(data))
+    names = set(zf.namelist())
+    if "metadata.yaml" not in names:
+        raise ValueError("plane zip: metadata.yaml missing")
+    meta = yaml.safe_load(zf.read("metadata.yaml"))
+    if meta.get("kind") != "plane":
+        raise ValueError(f"plane zip: kind={meta.get('kind')!r}, expected 'plane'")
+    members = meta.get("members") or {}
+    for member in ("plane.onnx", "classes.json"):
+        if member not in names:
+            raise ValueError(f"plane zip: {member} missing")
+        want = members.get(member)
+        got = hashlib.sha256(zf.read(member)).hexdigest()
+        if want is not None and want != got:
+            raise ValueError(f"plane zip: {member} sha256 mismatch ({got})")
+        if want is None:
+            raise ValueError(f"plane zip: {member} has no sha256 in metadata")
+    classes = _json.loads(zf.read("classes.json"))
+    return PlaneModel(
+        graph=zf.read("plane.onnx"),
+        classes=classes,
+        k_passes=int(meta.get("k_passes", 2)),
+    )
+
+
+PlaneModel.from_zip = staticmethod(from_zip)