interscript-ml 0.1.0__py3-none-any.whl

interscript_ml/__init__.py

@@ -0,0 +1,33 @@
+"""interscript-ml — the Python runtime for Interscript Model Format (IMF v1).
+
+The reference implementation: the Ruby and TypeScript runtimes are
+diffed against this one on shared golden sets.
+
+    from interscript_ml import Model
+    model = Model.load("khm-latn-1.0.zip")
+    model.translate("ភាសា")        # -> "pheasaea"
+
+Byte-level only: the tokenizer is the canonical ByT5 table (byte b ->
+token id b+3, trailing EOS), fixed and documented — no vocab files.
+"""
+
+from __future__ import annotations
+
+from interscript_ml.loader import Manifest, ModelFormatError
+from interscript_ml.model import Model
+from interscript_ml.registry import RegistryError, 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",
+    "RegistryError",
+    "UNK_ID",
+    "decode",
+    "encode",
+    "resolve",
+]

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

interscript_ml/model.py

@@ -0,0 +1,211 @@
+"""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.loader import load_manifest, verify_and_read
+from interscript_ml.tokens import EOS_ID, PAD_ID, decode, encode
+
+
+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, num_beams: int = 1) -> str:
+        token_ids = self.generate(text, max_len=max_len, num_beams=num_beams)
+        return decode(token_ids)
+
+    def generate(self, text: str, max_len: int = 256, num_beams: int = 1) -> 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:
+            if num_beams > 1:
+                return self._beam_kv(hidden, max_len, num_beams)
+            return self._greedy_kv(hidden, max_len)
+        return self._greedy_plain(hidden, max_len)
+
+    def _beam_kv(self, hidden, max_len: int, num_beams: int) -> list[int]:
+        """Batched beam search over the KV graph: the export's batch axis
+        carries the beams; per-step presents are gathered on beam reorder.
+        Canonical semantics: EOS hypotheses are recorded but never shrink
+        the live set (candidates come from a 2K window); the search runs
+        to max_len or exhaustion, and the winner is picked by raw
+        cumulative logprob — length normalization measurably rewards
+        long garbage on low-confidence byte models."""
+        beams = num_beams
+        enc = np.repeat(hidden, beams, axis=0)
+        pasts = {
+            name: np.repeat(zero, beams, axis=0)
+            for name, zero in self._pasts.items()
+        }
+        current = np.full((beams, 1), PAD_ID, dtype=np.int64)
+        scores = np.full((beams,), -np.inf, dtype=np.float32)
+        scores[0] = 0.0  # only beam 0 is live at step 0
+        sequences: list[list[int]] = [[] for _ in range(beams)]
+        finished: list[tuple[float, list[int]]] = []
+
+        for _ in range(max_len):
+            outputs = self._decoder.run(
+                None,
+                {"input_ids": current, "encoder_hidden_states": enc, **pasts},
+            )
+            results = dict(zip(self._output_names, outputs, strict=True))
+            logits = results["logits"][:, -1, :].astype(np.float32)
+            logprobs = logits - np.log(
+                np.exp(logits - logits.max(axis=1, keepdims=True)).sum(
+                    axis=1, keepdims=True
+                )
+            )  # stable log-softmax
+            cand = np.where(
+                (scores > -np.inf)[:, None],
+                scores[:, None] + logprobs,
+                -np.inf,
+            )
+            flat = cand.reshape(-1)
+            window = min(2 * beams, flat.size)
+            order = np.argpartition(flat, -window)[-window:]
+            order = order[np.argsort(-flat[order])]
+            new_pasts = {
+                name: results[name.replace("past_", "present_", 1)] for name in pasts
+            }
+            next_pasts: dict[str, np.ndarray] = {}
+            next_sequences: list[list[int]] = []
+            next_scores = np.full((beams,), -np.inf, dtype=np.float32)
+            next_current = np.full((beams, 1), PAD_ID, dtype=np.int64)
+            slot = 0
+            for idx in order:
+                src = int(idx // cand.shape[1])
+                token = int(idx % cand.shape[1])
+                score = float(flat[idx])
+                if token == EOS_ID:
+                    finished.append((score, sequences[src]))
+                    continue
+                if slot >= beams:
+                    continue
+                next_scores[slot] = score
+                next_sequences.append(sequences[src] + [token])
+                next_current[slot, 0] = token
+                for name, tensor in new_pasts.items():
+                    holder = next_pasts.setdefault(
+                        name, np.empty((beams,) + tensor.shape[1:], tensor.dtype)
+                    )
+                    holder[slot : slot + 1] = tensor[src : src + 1]
+                slot += 1
+            if slot == 0:
+                break
+            pasts = next_pasts
+            sequences = next_sequences + [
+                [] for _ in range(beams - len(next_sequences))
+            ]
+            scores = next_scores
+            current = next_current
+
+        pool = finished + [
+            (float(scores[i]), sequences[i]) for i in range(beams) if scores[i] > -np.inf
+        ]
+        best = max(pool, key=lambda s: s[0])
+        return best[1]
+
+    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)

interscript_ml/registry.py

@@ -0,0 +1,159 @@
+"""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 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://raw.githubusercontent.com/interscript/interscript-ml/main/models.yaml"
+)
+ENV_INDEX = "SECRYST_INDEX"
+ENV_CACHE = "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
+    parts: tuple[Part, ...] = ()
+
+
+def cache_dir() -> Path:
+    if os.environ.get(ENV_CACHE):
+        return Path(os.environ[ENV_CACHE])
+    return Path.home() / ".cache" / "secryst"
+
+
+def load_index(index_url: str | None = None) -> dict[str, IndexEntry]:
+    source = index_url or os.environ.get(ENV_INDEX) or DEFAULT_INDEX_URL
+    if source.startswith(("http://", "https://")):
+        with urllib.request.urlopen(source) as response:
+            text = response.read().decode("utf-8")
+    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", ""),
+            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

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")

interscript_ml-0.1.0.dist-info/METADATA

@@ -0,0 +1,55 @@
+Metadata-Version: 2.4
+Name: interscript-ml
+Version: 0.1.0
+Summary: Python runtime for Interscript Model Format (IMF v1) — the phonological layer of Interscript
+Author: Interscript Project
+License: BSD-3-Clause
+Keywords: transliteration,diacritization,g2p,onnx,byt5
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=1.26
+Requires-Dist: pyyaml>=6.0
+Requires-Dist: onnxruntime>=1.17
+Provides-Extra: dev
+Requires-Dist: pytest>=8.0; extra == "dev"
+Requires-Dist: onnx>=1.16; extra == "dev"
+
+# interscript-ml (Python runtime)
+
+The reference Python runtime for **IMF v1** model zips — the phonological
+layer of Interscript. The Ruby (secryst gem) and TypeScript
+(@interscript/ml) runtimes are diffed against this one on shared golden
+sets.
+
+```python
+from interscript_ml import Model
+
+model = Model.load("khm-latn-1.0")        # id: index resolve -> download
+                                          # -> sha256-verify -> cache -> load
+model.translate("ភាសា")                   # -> "pheasaea"
+model.id                                   # "khm-latn-1.0"
+
+model = Model.load("khm-latn-1.0.zip")    # or: a local zip path directly
+```
+
+- Byte-level only: the canonical ByT5 table (byte `b` → id `b+3`,
+  trailing EOS) — no vocab files, no per-model tokenization code.
+- Greedy KV-cache decode when the zip ships `decoder-kv.onnx`
+  (default), plain full-recompute fallback otherwise.
+- Every `.onnx` member is sha256-verified against `metadata.yaml`
+  before the session is created; corrupt downloads fail loudly.
+- Dynamic fetch per the `models.yaml` contract (shared with the Ruby and
+  TypeScript runtimes): resolve id -> channel URL, download to temp,
+  verify whole-file sha256 against the index, atomically install into
+  `~/.cache/interscript/models/<id>/`. Overrides:
+  `SECRYST_INDEX` (URL or path), `SECRYST_CACHE`.
+
+Install: `pip install ./runtime` (from the interscript-ml checkout) or
+`pip install -e "./runtime[dev]"` for development.
+
+Tests: `python -m pytest runtime/tests` — tiny-graph zips, no torch
+needed. The end-to-end golden test runs when `SECRYST_E2E_ZIP`
+points at a real zip (e.g. `models/khm-latn/khm-latn-1.0-fp32.zip`)
+and asserts byte-identical outputs against `golden/khm-latn-100.jsonl`.
+
+License: BSD-3-Clause.

interscript_ml-0.1.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+interscript_ml/__init__.py,sha256=qnoVrOcnMimA93q9PyGxtMj8A9SWjjxJjif-IwB-e7g,955
+interscript_ml/loader.py,sha256=lau1Ed9EXmhRdGiwKt_BuYApZrfAO1bXvo6rUZ9_NgU,2321
+interscript_ml/model.py,sha256=ZmPtzuhezZ91ecw0bgb9wPhl122-HS6GMjZBUO4Uo7o,8432
+interscript_ml/registry.py,sha256=qmy7NTgXXHZaVzCt6h4VC1SGImIbn57KQLPFwbVzYh4,5303
+interscript_ml/tokens.py,sha256=UV59DlvAte1VDZRB_rW7hQ2dOCgIiMRWkvcoJ3N9nTk,884
+interscript_ml-0.1.0.dist-info/METADATA,sha256=7zBvwM1uIaUtO9dF3Dt_PC_YVCGj39YKPHaJ8POKMEU,2291
+interscript_ml-0.1.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+interscript_ml-0.1.0.dist-info/top_level.txt,sha256=YBwZ6Bpz1TTlnsES0VcaIGF0aG4azG1RTeY_kmRFQkI,15
+interscript_ml-0.1.0.dist-info/RECORD,,

interscript_ml-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

interscript_ml-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+interscript_ml