From 5c97856c0adbb70abd90cb72aff5fc00a5928535 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Tue, 6 Oct 2026 11:24:37 +0800 Subject: [PATCH] feat: neural layer as interscript.ml - unified transliterate dispatch over maps + models Ports the secryst crystal (loader/model/plane/registry/tokens, 20 ported tests) into interscript as an optional [ml] extra. transliterate dispatches map ids to the rule engine and model ids to the ML runtime by index kind (plane vs seq2seq); index pinned at index-v7; env vars INTERSCRIPT_ML_INDEX/CACHE with SECRYST_* legacy aliases. --- pyproject.toml | 3 +- src/interscript/interscript.py | 25 +++- src/interscript/ml/__init__.py | 33 +++++ src/interscript/ml/loader.py | 68 ++++++++++ src/interscript/ml/model.py | 127 +++++++++++++++++ src/interscript/ml/plane.py | 92 +++++++++++++ src/interscript/ml/registry.py | 186 +++++++++++++++++++++++++ src/interscript/ml/tokens.py | 30 +++++ tests_ml/__init__.py | 0 tests_ml/helpers.py | 68 ++++++++++ tests_ml/test_dispatch.py | 68 ++++++++++ tests_ml/test_model.py | 159 ++++++++++++++++++++++ tests_ml/test_plane.py | 144 ++++++++++++++++++++ tests_ml/test_registry.py | 240 +++++++++++++++++++++++++++++++++ 14 files changed, 1240 insertions(+), 3 deletions(-) create mode 100644 src/interscript/ml/__init__.py create mode 100644 src/interscript/ml/loader.py create mode 100644 src/interscript/ml/model.py create mode 100644 src/interscript/ml/plane.py create mode 100644 src/interscript/ml/registry.py create mode 100644 src/interscript/ml/tokens.py create mode 100644 tests_ml/__init__.py create mode 100644 tests_ml/helpers.py create mode 100644 tests_ml/test_dispatch.py create mode 100644 tests_ml/test_model.py create mode 100644 tests_ml/test_plane.py create mode 100644 tests_ml/test_registry.py diff --git a/pyproject.toml b/pyproject.toml index 2609e84..147bc9d 100644 --- a/pyproject.toml +++ b/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"] diff --git a/src/interscript/interscript.py b/src/interscript/interscript.py index 02bed86..f50a14e 100644 --- a/src/interscript/interscript.py +++ b/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) diff --git a/src/interscript/ml/__init__.py b/src/interscript/ml/__init__.py new file mode 100644 index 0000000..ff67512 --- /dev/null +++ b/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", +] diff --git a/src/interscript/ml/loader.py b/src/interscript/ml/loader.py new file mode 100644 index 0000000..2c08b17 --- /dev/null +++ b/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 diff --git a/src/interscript/ml/model.py b/src/interscript/ml/model.py new file mode 100644 index 0000000..3a4ad93 --- /dev/null +++ b/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) diff --git a/src/interscript/ml/plane.py b/src/interscript/ml/plane.py new file mode 100644 index 0000000..3df7f2b --- /dev/null +++ b/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) diff --git a/src/interscript/ml/registry.py b/src/interscript/ml/registry.py new file mode 100644 index 0000000..fc103bf --- /dev/null +++ b/src/interscript/ml/registry.py @@ -0,0 +1,186 @@ +"""Model index resolution + cached downloads (the dynamic-fetch layer). + +Implements the models.yaml contract shared by the Ruby and TypeScript +runtimes: resolve an id, reuse a verified cache copy, or download + +sha256-verify + atomically install into the cache. +""" + +from __future__ import annotations + +import hashlib +import re +import urllib.error +import os +import shutil +import tempfile +import urllib.request +from dataclasses import dataclass +from pathlib import Path +from urllib.parse import urlparse + +import yaml + +DEFAULT_INDEX_URL = ( + "https://github.com/interscript/interscript-ml/releases/download/index-v7/models-index.yaml" +) +# primary env vars; the SECRYST_* names are deprecated aliases (kept one transition period) +ENV_INDEX = "INTERSCRIPT_ML_INDEX" +ENV_CACHE = "INTERSCRIPT_ML_CACHE" +ENV_INDEX_LEGACY = "SECRYST_INDEX" +ENV_CACHE_LEGACY = "SECRYST_CACHE" + + +class RegistryError(ValueError): + """The index cannot be fetched/parsed, or the id is unknown.""" + + +@dataclass(frozen=True) +class Part: + url: str + sha256: str + size: int + + +@dataclass(frozen=True) +class IndexEntry: + id: str + filename: str + url: str + sha256: str + size: int + precision: str + task: str + kind: str = "" + parts: tuple[Part, ...] = () + + +def cache_dir() -> Path: + env_cache = os.environ.get(ENV_CACHE) or os.environ.get(ENV_CACHE_LEGACY) + if env_cache: + return Path(env_cache) + return Path.home() / ".cache" / "interscript-ml" + + +def _fetch_http_with_sidecar(url: str) -> str: + """HTTP sources must ship a sibling .sha256 sidecar; verify before parsing.""" + with urllib.request.urlopen(url) as response: + body = response.read() + try: + with urllib.request.urlopen(f"{url}.sha256") as response: + sidecar = response.read().decode("utf-8").strip() + except urllib.error.HTTPError as exc: + raise RegistryError(f"index sha256 sidecar missing: {url}.sha256 -> {exc.code}") from exc + expected = sidecar.split()[0] if sidecar else "" + if not re.fullmatch(r"[0-9a-fA-F]{64}", expected): + raise RegistryError(f"index sha256 sidecar malformed: {url}.sha256") + actual = hashlib.sha256(body).hexdigest() + if actual != expected.lower(): + raise RegistryError( + f"index sha256 mismatch: got {actual}, sidecar says {expected.lower()}" + ) + return body.decode("utf-8") + + +def load_index(index_url: str | None = None) -> dict[str, IndexEntry]: + source = index_url or os.environ.get(ENV_INDEX) or os.environ.get(ENV_INDEX_LEGACY) or DEFAULT_INDEX_URL + if source.startswith(("http://", "https://")): + text = _fetch_http_with_sidecar(source) + else: + text = Path(source).read_text(encoding="utf-8") + raw = yaml.safe_load(text) + if not isinstance(raw, dict) or raw.get("version") != 1: + raise RegistryError("index must be a mapping with version: 1") + entries: dict[str, IndexEntry] = {} + for model_id, spec in raw.get("models", {}).items(): + parts = tuple( + Part(url=part["url"], sha256=part["sha256"], size=int(part.get("size", 0))) + for part in spec.get("parts", []) + ) + entries[model_id] = IndexEntry( + id=model_id, + filename=spec["filename"], + url=spec.get("url", ""), + sha256=spec["sha256"], + size=int(spec.get("size", 0)), + precision=spec.get("precision", "fp32"), + task=spec.get("task", ""), + kind=spec.get("kind", ""), + parts=parts, + ) + return entries + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as fh: + while chunk := fh.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def _open_channel(url: str): + if url.startswith("file://"): + return open(urlparse(url).path, "rb") + return urllib.request.urlopen(url) + + +def _download_parts(entry: IndexEntry, downloaded: Path) -> None: + """Stream parts into `downloaded` in index order, verifying each part's + sha256 as it lands. Used when the artifact exceeds GitHub's 2 GiB + per-asset cap; the assembled file is checked against entry.sha256 by + the caller, so the cache contract is identical to single-file models.""" + with downloaded.open("ab") as out: + for index, part in enumerate(entry.parts): + digest = hashlib.sha256() + with _open_channel(part.url) as remote: + while chunk := remote.read(1024 * 1024): + out.write(chunk) + digest.update(chunk) + actual = digest.hexdigest() + if actual != part.sha256: + raise RegistryError( + f"part {index} of {entry.filename} sha256 mismatch: " + f"got {actual}, index says {part.sha256}" + ) + + +def resolve(model_id: str, index_url: str | None = None) -> Path: + """Return a verified local zip path for `model_id`, downloading and + installing into the cache when needed. Never returns an unverified + file: cache hits are re-verified against the index sha256.""" + entries = load_index(index_url) + if model_id not in entries: + raise RegistryError( + f"unknown model id {model_id!r} (known: {sorted(entries)})" + ) + entry = entries[model_id] + target = cache_dir() / "models" / model_id / entry.filename + if target.is_file() and _sha256_file(target) == entry.sha256: + return target + + target.parent.mkdir(parents=True, exist_ok=True) + fd, tmp_name = tempfile.mkstemp(dir=target.parent, suffix=".part") + os.close(fd) + downloaded = Path(tmp_name) + if entry.parts: + _download_parts(entry, downloaded) + elif entry.url.startswith("file://"): + source = Path(urlparse(entry.url).path) + if not source.is_file(): + raise RegistryError(f"channel file missing: {source}") + shutil.copyfile(source, downloaded) # file:// is a mirror, not a move + else: + urllib.request.urlretrieve(entry.url, downloaded) + try: + actual = _sha256_file(downloaded) + if actual != entry.sha256: + raise RegistryError( + f"downloaded {entry.filename} sha256 mismatch: got {actual}, " + f"index says {entry.sha256}" + ) + if downloaded != target: + os.replace(downloaded, target) + finally: + if downloaded != target and downloaded.exists(): + downloaded.unlink() + return target diff --git a/src/interscript/ml/tokens.py b/src/interscript/ml/tokens.py new file mode 100644 index 0000000..b3f6af3 --- /dev/null +++ b/src/interscript/ml/tokens.py @@ -0,0 +1,30 @@ +"""The canonical ByT5 byte table (fixed, no vocab files). + +Stock google/byt5 tokenizers map UTF-8 byte b to id b+3 and append +EOS(1); pad=0, unk=2. Ids are NOT raw byte values — feeding text.bytes +directly produces silent garbage on real models. +""" + +from __future__ import annotations + +BYTE_OFFSET = 3 +PAD_ID = 0 +EOS_ID = 1 +UNK_ID = 2 + + +def encode(text: str) -> list[int]: + """Canonical byte-level tokenization (byte+3 table, trailing EOS).""" + return [b + BYTE_OFFSET for b in text.encode("utf-8")] + [EOS_ID] + + +def decode(token_ids: list[int]) -> str: + """Token ids -> text; stops at EOS, maps id-3 back to a byte.""" + out = bytearray() + for token in token_ids: + if token == EOS_ID: + break + if token in (PAD_ID, UNK_ID): + continue + out.append((token - BYTE_OFFSET) % 256) + return out.decode("utf-8", errors="replace") diff --git a/tests_ml/__init__.py b/tests_ml/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests_ml/helpers.py b/tests_ml/helpers.py new file mode 100644 index 0000000..fb5a175 --- /dev/null +++ b/tests_ml/helpers.py @@ -0,0 +1,68 @@ +"""Shared tiny-graph zip builder for runtime tests.""" + +from __future__ import annotations + +import hashlib +import zipfile +from pathlib import Path + +import numpy as np +import yaml +from onnx import TensorProto, helper, numpy_helper + +MANIFEST = { + "format": "imf-v1", + "id": "tiny-1.0", + "task": "translit", + "source_script": "Latn", + "target": "Latn", + "tokenizer": "bytes", + "opset": 14, + "decoder": "plain", + "precision": "fp32", + "license": "BSD-3-Clause", + "trained_from": "runtime test fixture", +} + + +def _add_graph(name: str, inputs: list[str], output: str) -> bytes: + graph = helper.make_graph( + nodes=[helper.make_node("Add", [inputs[0], "bias"], [output])], + name=name, + inputs=[ + helper.make_tensor_value_info(n, TensorProto.INT64, ["batch", "seq"]) + for n in inputs + ], + outputs=[ + helper.make_tensor_value_info(output, TensorProto.INT64, ["batch", "seq"]) + ], + initializer=[numpy_helper.from_array(np.zeros(1, dtype=np.int64), "bias")], + ) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("", 14)], ir_version=7 + ) + return model.SerializeToString() + + +def build_tiny_zip( + path: Path, tamper: bool = False, manifest: dict | None = None +) -> Path: + encoder = _add_graph("tiny-enc", ["input_ids"], "last_hidden_state") + decoder = _add_graph( + "tiny-dec", ["input_ids", "encoder_hidden_states"], "logits" + ) + sha = { + "encoder.onnx": hashlib.sha256(encoder).hexdigest(), + "decoder.onnx": hashlib.sha256(decoder).hexdigest(), + } + if tamper: + sha["encoder.onnx"] = "0" * 64 + meta = dict(manifest if manifest is not None else MANIFEST) + meta["sha256"] = sha + path.parent.mkdir(parents=True, exist_ok=True) + with zipfile.ZipFile(path, "w") as zf: + zf.writestr("metadata.yaml", yaml.safe_dump(meta)) + zf.writestr("encoder.onnx", encoder) + zf.writestr("decoder.onnx", decoder) + zf.writestr("README.md", "# tiny\n") + return path diff --git a/tests_ml/test_dispatch.py b/tests_ml/test_dispatch.py new file mode 100644 index 0000000..96f7847 --- /dev/null +++ b/tests_ml/test_dispatch.py @@ -0,0 +1,68 @@ +"""Unified dispatch: transliterate(id, text) resolves map ids via the +engine and model ids via the ML layer (lazy, optional extra).""" + +from __future__ import annotations + +from pathlib import Path + +import yaml + +from interscript import transliterate + + +def _index_file(tmp_path: Path, zip_path: Path) -> str: + import hashlib + + index = { + "version": 1, + "models": { + "tiny-1.0": { + "task": "translit", + "precision": "fp32", + "filename": zip_path.name, + "url": f"file://{zip_path}", + "sha256": hashlib.sha256(zip_path.read_bytes()).hexdigest(), + "size": zip_path.stat().st_size, + } + }, + } + path = tmp_path / "models.yaml" + path.write_text(yaml.safe_dump(index), encoding="utf-8") + return str(path) + + +def test_transliterate_routes_model_id_through_ml_layer(tmp_path: Path) -> None: + from interscript.ml import Model + from tests_ml.helpers import build_tiny_zip + + zip_path = build_tiny_zip(tmp_path / "channel" / "tiny.zip") + index = _index_file(tmp_path, zip_path) + + expected = Model.load(zip_path).translate("he", max_len=4) + got = transliterate("tiny-1.0", "he", index_url=index) + assert isinstance(got, str) + assert got == expected + + +def test_transliterate_unknown_id_still_raises_file_not_found(tmp_path: Path) -> None: + from tests_ml.helpers import build_tiny_zip + + index = _index_file(tmp_path, build_tiny_zip(tmp_path / "channel" / "tiny.zip")) + try: + transliterate("no-such-map-9.9", "he", index_url=index) + raise AssertionError("expected FileNotFoundError") + except FileNotFoundError as e: + assert "no-such-map-9.9" in str(e) + + +def test_transliterate_map_id_still_uses_engine(tmp_path: Path) -> None: + import pytest + + from interscript import add_load_path + + maps = Path(__file__).resolve().parent.parent / "maps" + if not maps.exists(): + pytest.skip("maps corpus not checked out") + add_load_path(str(maps)) + out = transliterate("bgnpcgn-rus-cyrl-latn-2016", "Москва") + assert isinstance(out, str) and out diff --git a/tests_ml/test_model.py b/tests_ml/test_model.py new file mode 100644 index 0000000..221431c --- /dev/null +++ b/tests_ml/test_model.py @@ -0,0 +1,159 @@ +"""Tests for the interscript-ml runtime. + +Tiny-graph zips built with the onnx package (no torch, no training +repo). The end-to-end golden test runs only when a real zip is provided +via SECRYST_E2E_ZIP. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import zipfile +from pathlib import Path + +import pytest +import yaml + +ort = pytest.importorskip("onnxruntime") +onnx = pytest.importorskip("onnx") + +from interscript.ml import Model, ModelFormatError, decode, encode # noqa: E402 +from onnx import TensorProto, helper, numpy_helper # noqa: E402 + +import numpy as np # noqa: E402 + + +def _graph(opset: int = 14) -> bytes: + graph = helper.make_graph( + nodes=[helper.make_node("Add", ["input_ids", "bias"], ["last_hidden_state"])], + name="tiny-enc", + inputs=[ + helper.make_tensor_value_info("input_ids", TensorProto.INT64, ["batch", "seq"]) + ], + outputs=[ + helper.make_tensor_value_info( + "last_hidden_state", TensorProto.INT64, ["batch", "seq"] + ) + ], + initializer=[numpy_helper.from_array(np.zeros(1, dtype=np.int64), "bias")], + ) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("", opset)], ir_version=7 + ) + return model.SerializeToString() + + +def _decoder_graph() -> bytes: + graph = helper.make_graph( + nodes=[ + helper.make_node("Add", ["input_ids", "bias"], ["logits"]) + ], + name="tiny-dec", + inputs=[ + helper.make_tensor_value_info("input_ids", TensorProto.INT64, ["batch", "seq"]), + helper.make_tensor_value_info( + "encoder_hidden_states", TensorProto.INT64, ["batch", "seq"] + ), + ], + outputs=[ + helper.make_tensor_value_info("logits", TensorProto.INT64, ["batch", "seq"]) + ], + initializer=[numpy_helper.from_array(np.zeros(1, dtype=np.int64), "bias")], + ) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("", 14)], ir_version=7 + ) + return model.SerializeToString() + + +MANIFEST = { + "format": "imf-v1", + "id": "tiny-1.0", + "task": "translit", + "source_script": "Latn", + "target": "Latn", + "tokenizer": "bytes", + "opset": 14, + "decoder": "plain", + "precision": "fp32", + "license": "BSD-3-Clause", + "trained_from": "runtime test fixture", +} + + +def _tiny_zip(path: Path, tamper: bool = False, manifest: dict | None = None) -> Path: + encoder, decoder = _graph(), _decoder_graph() + sha = { + "encoder.onnx": hashlib.sha256(encoder).hexdigest(), + "decoder.onnx": hashlib.sha256(decoder).hexdigest(), + } + if tamper: + sha["encoder.onnx"] = "0" * 64 + meta = dict(manifest if manifest is not None else MANIFEST) + meta["sha256"] = sha + with zipfile.ZipFile(path, "w") as zf: + zf.writestr("metadata.yaml", yaml.safe_dump(meta)) + zf.writestr("encoder.onnx", encoder) + zf.writestr("decoder.onnx", decoder) + zf.writestr("README.md", "# tiny\n") + return path + + +def test_token_table() -> None: + assert encode("rok") == [117, 114, 110, 1] + assert decode([117, 114, 110]) == "rok" + assert decode([117, 1, 114]) == "r" + assert decode([]) == "" + + +def test_load_and_decode_tiny(tmp_path: Path) -> None: + z = _tiny_zip(tmp_path / "tiny.zip") + model = Model.load(z) + assert model.id == "tiny-1.0" + # tiny graphs are identity Adds: logits echo the decoder prefix, so + # greedy emits encode(PAD-prefix input)+... — deterministic, not + # meaningful; what matters is that the loop runs and decodes. + tokens = model.generate("he", max_len=4) + assert isinstance(tokens, list) + text = model.translate("he", max_len=4) + assert isinstance(text, str) + + +def test_sha256_mismatch_rejected(tmp_path: Path) -> None: + z = _tiny_zip(tmp_path / "bad.zip", tamper=True) + with pytest.raises(ModelFormatError, match="sha256 mismatch"): + Model.load(z) + + +def test_non_bytes_tokenizer_rejected(tmp_path: Path) -> None: + manifest = dict(MANIFEST, tokenizer="sentencepiece") + z = _tiny_zip(tmp_path / "spm.zip", manifest=manifest) + with pytest.raises(ModelFormatError, match="byte-level only"): + Model.load(z) + + +def test_missing_graph_rejected(tmp_path: Path) -> None: + z = _tiny_zip(tmp_path / "m.zip") + truncated = tmp_path / "trunc.zip" + with zipfile.ZipFile(z) as src, zipfile.ZipFile(truncated, "w") as dst: + for name in src.namelist(): + if name != "decoder.onnx": + dst.writestr(name, src.read(name)) + with pytest.raises(ModelFormatError, match="decoder.onnx"): + Model.load(truncated) + + +def test_golden_set_e2e() -> None: + """Run against a real zip: byte-identical outputs on the golden set.""" + zip_path = os.environ.get("SECRYST_E2E_ZIP") + if not zip_path: + pytest.skip("set SECRYST_E2E_ZIP to a real IMF zip") + golden = Path(__file__).resolve().parent.parent.parent / "golden" / "khm-latn-100.jsonl" + if "khm" not in Path(zip_path).name: + pytest.skip("golden file is khm-latn specific") + model = Model.load(zip_path) + rows = [json.loads(line) for line in golden.read_text(encoding="utf-8").splitlines()] + for row in rows: + assert model.translate(row["input"], max_len=128) == row["output"], row["input"] diff --git a/tests_ml/test_plane.py b/tests_ml/test_plane.py new file mode 100644 index 0000000..d0cd79a --- /dev/null +++ b/tests_ml/test_plane.py @@ -0,0 +1,144 @@ +"""Tests for the plane model kind (run-018 on-device artifact). + +The plane artifact ships one ONNX graph taking (input_ids, plane_ids) +to per-position class logits; decode is a host-side K-pass loop; the +haraqat plane renders onto the skeleton. Fixture graph: logits[i] +depends only on plane_ids[i] via class = (plane + 1) % n_classes, so +predictions cycle deterministically - enough to pin the loop, the +rendering, and the byte table. +""" + +from __future__ import annotations + +import sys +import unittest +from pathlib import Path + +import pytest + +pytest.importorskip("onnx") +pytest.importorskip("onnxruntime") +import numpy as np +from onnx import TensorProto, helper, numpy_helper + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) + +from interscript.ml.plane import PlaneModel, render_plane + +N_CLASSES = 4 +CLASSES = ["", "َ", "ُ", "ْ"] + + +def tiny_plane_graph() -> bytes: + """class_logits[i, c] = 1 iff c == (plane_ids[i] + 1) % N_CLASSES + (bias on the desired class so argmax is exact).""" + init = numpy_helper.from_array( + np.array( + [[0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1], [1, 0, 0, 0], [0, 1, 0, 0]], + dtype=np.float32, + ), + "map", + ) + node = helper.make_node("Gather", ["map", "plane_ids"], ["class_logits"], axis=0) + graph = helper.make_graph( + [node], + "tiny-plane", + inputs=[ + helper.make_tensor_value_info("input_ids", TensorProto.INT64, ["b", "t"]), + helper.make_tensor_value_info("plane_ids", TensorProto.INT64, ["b", "t"]), + ], + outputs=[ + helper.make_tensor_value_info("class_logits", TensorProto.FLOAT, ["b", "t", "c"]) + ], + initializer=[init], + ) + m = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 16)]) + m.ir_version = 9 + return m.SerializeToString() + + +class TestRenderPlane(unittest.TestCase): + def test_renders_classes_onto_skeleton(self): + skel = "كتب" + classes = ["", "َ", "ُ"] + self.assertEqual(render_plane(skel, classes), "كتَبُ") + + def test_multibyte_chars_consume_one_class(self): + self.assertEqual(render_plane("كتب", ["َ", "", ""]), "كَتب") + + +class TestPlaneModel(unittest.TestCase): + def _model(self, k=2): + return PlaneModel(graph=tiny_plane_graph(), classes=list(CLASSES), k_passes=k) + + def test_k_pass_loop_cycles_predictions(self): + m = self._model(k=1) + # one pass from MASK(=4): class = (4+1) % 4 = 1 -> fatha everywhere + out = m.translate("كتب") + self.assertEqual(out, "كَتَبَ") + m2 = self._model(k=2) + # second pass: 1 -> 2 -> damma + self.assertEqual(m2.translate("كتب"), "كُتُبُ") + + def test_input_ids_follow_byte_table(self): + m = self._model(k=1) + ids = m.encode("ab") + self.assertEqual(ids, [ord("a") + 3, ord("b") + 3]) + + +if __name__ == "__main__": + unittest.main() + + +class TestPlaneZipRoundTrip(unittest.TestCase): + def _zip_bytes(self, tamper: bool = False) -> bytes: + import hashlib + import io + import json + import zipfile + + graph = tiny_plane_graph() + classes_json = json.dumps(CLASSES, ensure_ascii=False).encode("utf-8") + # metadata shas describe the AUTHENTIC members + sha_g = hashlib.sha256(graph).hexdigest() + sha_c = hashlib.sha256(classes_json).hexdigest() + meta = ( + f"id: plane-fixture-1.0\nprecision: fp32\nkind: plane\nk_passes: 2\n" + f"members:\n plane.onnx: {sha_g}\n classes.json: {sha_c}\n" + ) + stored = bytearray(graph) + if tamper: + # valid-CRC corruption: stored content differs, metadata sha stale + idx = stored.find(b"tiny-plane") + stored[idx] ^= 0xFF + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w") as z: + z.writestr("metadata.yaml", meta) + z.writestr("plane.onnx", bytes(stored)) + z.writestr("classes.json", classes_json) + return buf.getvalue() + + def test_round_trip_load(self): + from interscript.ml.plane import PlaneModel + + m = PlaneModel.from_zip(self._zip_bytes()) + self.assertEqual(m.k, 2) + self.assertEqual(m.translate("كتب"), "كُتُبُ") + + def test_tampered_graph_rejected(self): + from interscript.ml.plane import PlaneModel + + with self.assertRaisesRegex(ValueError, "sha256"): + PlaneModel.from_zip(self._zip_bytes(tamper=True)) + + def test_missing_member_rejected(self): + import io + import zipfile + + from interscript.ml.plane import PlaneModel + + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w") as z: + z.writestr("metadata.yaml", "kind: plane\nmembers:\n plane.onnx: " + "0" * 64 + "\n") + with self.assertRaisesRegex(ValueError, "plane.onnx"): + PlaneModel.from_zip(buf.getvalue()) diff --git a/tests_ml/test_registry.py b/tests_ml/test_registry.py new file mode 100644 index 0000000..8f69776 --- /dev/null +++ b/tests_ml/test_registry.py @@ -0,0 +1,240 @@ +"""Tests for the dynamic-fetch layer (models.yaml resolution + cache).""" + +from __future__ import annotations + +import hashlib +import zipfile +from pathlib import Path + +import pytest +import yaml + +from interscript.ml.registry import RegistryError, resolve +from tests_ml.helpers import build_tiny_zip + +import os # noqa: E402 + + +def _index_file(tmp_path: Path, zip_path: Path, sha256: str | None = None) -> Path: + index = { + "version": 1, + "models": { + "tiny-1.0": { + "task": "translit", + "precision": "fp32", + "filename": zip_path.name, + "url": f"file://{zip_path}", + "sha256": sha256 or hashlib.sha256(zip_path.read_bytes()).hexdigest(), + "size": zip_path.stat().st_size, + } + }, + } + path = tmp_path / "models.yaml" + path.write_text(yaml.safe_dump(index), encoding="utf-8") + return path + + +def test_resolve_downloads_verifies_and_caches(tmp_path: Path) -> None: + zip_path = build_tiny_zip(tmp_path / "channel" / "tiny.zip") + index = _index_file(tmp_path, zip_path) + cache = tmp_path / "cache" + os.environ["SECRYST_CACHE"] = str(cache) + try: + local = resolve("tiny-1.0", index_url=str(index)) + assert local == cache / "models" / "tiny-1.0" / "tiny.zip" + assert local.is_file() + # second resolve is a verified cache hit (channel dir removed) + zip_path.unlink() + assert resolve("tiny-1.0", index_url=str(index)) == local + finally: + os.environ.pop("SECRYST_CACHE", None) + + +def test_resolve_rejects_bad_download(tmp_path: Path) -> None: + zip_path = build_tiny_zip(tmp_path / "channel" / "tiny.zip") + index = _index_file(tmp_path, zip_path, sha256="0" * 64) + os.environ["SECRYST_CACHE"] = str(tmp_path / "cache") + try: + with pytest.raises(RegistryError, match="sha256 mismatch"): + resolve("tiny-1.0", index_url=str(index)) + finally: + os.environ.pop("SECRYST_CACHE", None) + + +def test_resolve_unknown_id(tmp_path: Path) -> None: + index = tmp_path / "models.yaml" + index.write_text(yaml.safe_dump({"version": 1, "models": {}}), encoding="utf-8") + with pytest.raises(RegistryError, match="unknown model id"): + resolve("nope-1.0", index_url=str(index)) + + +def test_model_load_by_id(tmp_path: Path) -> None: + zip_path = build_tiny_zip(tmp_path / "channel" / "tiny.zip") + index = _index_file(tmp_path, zip_path) + os.environ["SECRYST_CACHE"] = str(tmp_path / "cache") + try: + from interscript.ml import Model + + model = Model.load("tiny-1.0", index_url=str(index)) + assert model.id == "tiny-1.0" + assert isinstance(model.translate("he", max_len=4), str) + finally: + os.environ.pop("SECRYST_CACHE", None) + + +def test_resolve_parts_assembles_and_verifies(tmp_path: Path) -> None: + import hashlib + + zip_path = build_tiny_zip(tmp_path / "channel" / "tiny.zip") + blob = zip_path.read_bytes() + part_a, part_b = blob[: len(blob) // 2 + 3], blob[len(blob) // 2 + 3 :] + channel = tmp_path / "channel" + (channel / "tiny.zip.part-00").write_bytes(part_a) + (channel / "tiny.zip.part-01").write_bytes(part_b) + index = { + "version": 1, + "models": { + "tiny-1.0": { + "task": "translit", + "precision": "fp32", + "filename": "tiny.zip", + "sha256": hashlib.sha256(blob).hexdigest(), + "size": len(blob), + "parts": [ + { + "url": f"file://{channel / 'tiny.zip.part-00'}", + "sha256": hashlib.sha256(part_a).hexdigest(), + "size": len(part_a), + }, + { + "url": f"file://{channel / 'tiny.zip.part-01'}", + "sha256": hashlib.sha256(part_b).hexdigest(), + "size": len(part_b), + }, + ], + }, + }, + } + index_path = tmp_path / "models.yaml" + index_path.write_text(yaml.safe_dump(index), encoding="utf-8") + cache = tmp_path / "cache" + os.environ["SECRYST_CACHE"] = str(cache) + try: + local = resolve("tiny-1.0", index_url=str(index_path)) + assert local == cache / "models" / "tiny-1.0" / "tiny.zip" + assert local.read_bytes() == blob + zip_path.unlink() + (channel / "tiny.zip.part-00").unlink() + assert resolve("tiny-1.0", index_url=str(index_path)) == local + finally: + os.environ.pop("SECRYST_CACHE", None) + + +def test_resolve_parts_rejects_corrupt_part(tmp_path: Path) -> None: + import hashlib + + zip_path = build_tiny_zip(tmp_path / "channel" / "tiny.zip") + blob = zip_path.read_bytes() + part_a, part_b = blob[:7], blob[7:] + channel = tmp_path / "channel" + (channel / "tiny.zip.part-00").write_bytes(part_a) + (channel / "tiny.zip.part-01").write_bytes(part_b) + index = { + "version": 1, + "models": { + "tiny-1.0": { + "filename": "tiny.zip", + "sha256": hashlib.sha256(blob).hexdigest(), + "parts": [ + { + "url": f"file://{channel / 'tiny.zip.part-00'}", + "sha256": "0" * 64, + "size": len(part_a), + }, + { + "url": f"file://{channel / 'tiny.zip.part-01'}", + "sha256": hashlib.sha256(part_b).hexdigest(), + "size": len(part_b), + }, + ], + }, + }, + } + index_path = tmp_path / "models.yaml" + index_path.write_text(yaml.safe_dump(index), encoding="utf-8") + os.environ["SECRYST_CACHE"] = str(tmp_path / "cache") + try: + with pytest.raises(RegistryError, match="part 0"): + resolve("tiny-1.0", index_url=str(index_path)) + finally: + os.environ.pop("SECRYST_CACHE", None) + + +def test_default_index_url_pins_github_release() -> None: + # exact pin (not index-v\d+): the default must resolve every model + # the distribution contract has shipped, incl. static-int8 (v6). + # Bump this test WITH the registry constant, never ahead of it. + from interscript.ml.registry import DEFAULT_INDEX_URL + + assert DEFAULT_INDEX_URL == ( + "https://github.com/interscript/interscript-ml" + "/releases/download/index-v7/models-index.yaml" + ) + assert "raw.githubusercontent" not in DEFAULT_INDEX_URL + + +def test_http_index_requires_valid_sha256_sidecar(monkeypatch, tmp_path) -> None: + import threading + from http.server import BaseHTTPRequestHandler, HTTPServer + + from interscript.ml.registry import load_index + + body = b"version: 1\nmodels: {}\n" + good = hashlib.sha256(body).hexdigest() + bad = "0" * 64 + + class Handler(BaseHTTPRequestHandler): + sidecar: str | None = None + + def do_GET(self): # noqa: N802 + if self.path == "/models-index.yaml": + self.send_response(200) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + return + if self.path == "/models-index.yaml.sha256": + if Handler.sidecar is None: + self.send_response(404) + self.send_header("Content-Length", "0") + self.end_headers() + return + payload = f"{Handler.sidecar} models-index.yaml\n".encode() + self.send_response(200) + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + return + self.send_response(404) + self.send_header("Content-Length", "0") + self.end_headers() + + def log_message(self, *args): # silence + pass + + server = HTTPServer(("127.0.0.1", 0), Handler) + port = server.server_address[1] + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + url = f"http://127.0.0.1:{port}/models-index.yaml" + try: + Handler.sidecar = good + assert load_index(url) == {} + Handler.sidecar = bad + with pytest.raises(RegistryError, match="index sha256 mismatch"): + load_index(url) + Handler.sidecar = None + with pytest.raises(RegistryError, match="index sha256"): + load_index(url) + finally: + server.shutdown()