laya-coreml 0.1.0__py3-none-any.whl

laya_coreml/__init__.py

@@ -0,0 +1,6 @@
+"""Core ML runtime for Laya. PyTorch is only needed when converting weights."""
+
+from .agent import Agent, load
+
+__all__ = ["Agent", "load"]
+__version__ = "0.1.0"

laya_coreml/__main__.py

@@ -0,0 +1,3 @@
+from .cli import main
+
+main()

laya_coreml/agent.py

@@ -0,0 +1,91 @@
+"""Core ML inference. No MLX, Transformers, or PyTorch dependency at runtime."""
+
+import json
+import math
+
+import numpy as np
+
+from .artifacts import package_for_coreml
+from .hub import DEFAULT_MODEL, resolve_checkpoint
+from .prompt import PromptMixin
+from .result import ResultMixin
+from .tokenizer import Tokenizer
+
+COMPUTE_UNITS = {"all": "ALL", "cpu": "CPU_ONLY", "cpu_gpu": "CPU_AND_GPU", "cpu_ne": "CPU_AND_NE"}
+
+
+class Agent(PromptMixin, ResultMixin):
+    def __init__(
+        self,
+        model_dir,
+        *,
+        compute_units="cpu_gpu",
+        allow_unvalidated_gpu=False,
+        revision=None,
+        local_files_only=False,
+    ):
+        if compute_units not in COMPUTE_UNITS:
+            raise ValueError(f"compute_units must be one of {list(COMPUTE_UNITS)}")
+        self.model_dir = resolve_checkpoint(
+            model_dir, revision=revision, local_files_only=local_files_only
+        )
+        self.manifest = json.loads((self.model_dir / "coreml_config.json").read_text())
+        if self.manifest.get("format") != "laya-coreml" or self.manifest.get("format_version") != 1:
+            raise ValueError("Unsupported Core ML export format")
+        self.shape = self.manifest["shape"]
+        if (
+            compute_units == "cpu_gpu"
+            and self.shape["flexible"]
+            and not self.shape.get("lengths")
+            and not allow_unvalidated_gpu
+        ):
+            raise ValueError(
+                "RangeDim + CPU_AND_GPU failed local fidelity and repeatability checks. "
+                "Re-export with the default enumerated shapes, or use compute_units='cpu'. "
+                "allow_unvalidated_gpu=True is for reproducing the failure only."
+            )
+        self.cfg = json.loads((self.model_dir / "rl_agent_config.json").read_text())
+        self.temperature = self.cfg.get("temperature", [1.0, 1.0, 1.0])
+        self.temperature_by_options = self.cfg.get("temperature_by_options", {})
+        if len(self.temperature) != 3 or any(
+            not math.isfinite(float(t)) or float(t) <= 0
+            for t in [*self.temperature, *self.temperature_by_options.values()]
+        ):
+            raise ValueError("Calibration temperatures must be finite and positive")
+        self.tok = Tokenizer(self.model_dir / "tokenizer")
+        self.batch_size = self.shape["batch_size"]
+        self.pad_to_multiple = 16
+        import coremltools as ct
+
+        self.compute_units = compute_units
+        self.model = ct.models.MLModel(
+            str(package_for_coreml(self.model_dir / "model.mlpackage")),
+            compute_units=getattr(ct.ComputeUnit, COMPUTE_UNITS[compute_units]),
+        )
+
+    def forward(self, batch):
+        outputs = self.model.predict(batch)
+        return np.asarray(outputs["logits"], np.float32), np.asarray(
+            outputs["action_logits"], np.float32
+        )
+
+
+def load(
+    model_dir=DEFAULT_MODEL,
+    *,
+    revision=None,
+    local_files_only=False,
+    compute_units=None,
+    allow_unvalidated_gpu=False,
+):
+    directory = resolve_checkpoint(model_dir, revision=revision, local_files_only=local_files_only)
+    manifest = json.loads((directory / "coreml_config.json").read_text())
+    if manifest.get("format") == "laya-coreml-ane":
+        from .ane import ANEAgent
+
+        return ANEAgent(directory, compute_units=compute_units or "cpu_ne")
+    return Agent(
+        directory,
+        compute_units=compute_units or "cpu_gpu",
+        allow_unvalidated_gpu=allow_unvalidated_gpu,
+    )

laya_coreml/ane.py

@@ -0,0 +1,167 @@
+"""Host embedding lookup + one ANE graph + CPU action-head runtime."""
+
+import json
+import math
+from pathlib import Path
+
+import coremltools as ct
+import numpy as np
+from safetensors import safe_open
+
+from laya_coreml.prompt import PromptMixin
+from laya_coreml.result import ResultMixin
+from laya_coreml.tokenizer import Tokenizer
+
+from .artifacts import package_for_coreml, verify_files, verify_research_manifest
+
+
+class ANEAgent(PromptMixin, ResultMixin):
+    def __init__(self, source, package=None, *, length=None, compute_units="cpu_ne"):
+        self.source = Path(source)
+        if package is None:
+            self.manifest = json.loads((self.source / "coreml_config.json").read_text())
+            if (self.manifest.get("format"), self.manifest.get("format_version")) != (
+                "laya-coreml-ane",
+                1,
+            ):
+                raise ValueError("Unsupported ANE bundle format")
+            shape = self.manifest["shape"]
+            if shape["batch_size"] != 1 or shape["max_options"] != 32 or shape["flexible"]:
+                raise ValueError("ANE runtime requires a fixed B1/K32 bundle")
+            if length is not None and length != shape["max_length"]:
+                raise ValueError("Requested length does not match the ANE bundle")
+            length = shape["max_length"]
+            verify_files(self.source, self.manifest["files"])
+            package = self.source / "model.mlpackage"
+            host_weights = self.source / "host_weights.safetensors"
+        else:
+            length = 96 if length is None else length
+            self.manifest = verify_research_manifest(self.source, package, length=length)
+            host_weights = self.source / "model.safetensors"
+        self.model_dir = self.source
+        self.cfg = json.loads((self.source / "rl_agent_config.json").read_text())
+        self.temperature = self.cfg.get("temperature", [1.0, 1.0, 1.0])
+        self.temperature_by_options = self.cfg.get("temperature_by_options", {})
+        self.tok = Tokenizer(self.source / "tokenizer")
+        self.batch_size, self.pad_to_multiple = 1, 16
+        self.shape = {
+            "batch_size": 1,
+            "max_length": length,
+            "min_length": length,
+            "max_options": 32,
+            "flexible": False,
+            "lengths": None,
+        }
+        self.compute_units = compute_units
+        units = {
+            "cpu_ne": ct.ComputeUnit.CPU_AND_NE,
+            "cpu_gpu": ct.ComputeUnit.CPU_AND_GPU,
+            "all": ct.ComputeUnit.ALL,
+            "cpu": ct.ComputeUnit.CPU_ONLY,
+        }
+        if compute_units not in units:
+            raise ValueError(f"compute_units must be one of {list(units)}")
+        self.model = ct.models.MLModel(
+            str(package_for_coreml(package)), compute_units=units[compute_units]
+        )
+        self.encoder_cfg = json.loads((self.source / "encoder/config.json").read_text())
+        width = int(self.encoder_cfg["hidden_size"])
+        expected_shapes = {
+            "embeddings": (1, width, 1, length),
+            "full_mask": (1, length, 1, length),
+            "local_mask": (1, length, 1, length),
+            "type_vectors": (1, width, 1, 1),
+            "marker_map": (1, length, 1, 32),
+        }
+        actual_shapes = {
+            feature.name: tuple(feature.type.multiArrayType.shape)
+            for feature in self.model.get_spec().description.input
+        }
+        if actual_shapes != expected_shapes:
+            raise ValueError(
+                f"Package signature mismatch: expected {expected_shapes}, got {actual_shapes}"
+            )
+        with safe_open(str(host_weights), framework="numpy") as weights:
+            self.embedding = weights.get_tensor("encoder.embeddings.tok_embeddings.weight")
+            self.type_embedding = weights.get_tensor("type_emb.weight")
+            self.action = {
+                key: weights.get_tensor("act_head." + key).astype(np.float32)
+                for key in ("0.weight", "0.bias", "2.weight", "2.bias")
+            }
+        spec = self.model.get_spec()
+        self.output_names = [output.name for output in spec.description.output]
+        self._erf = np.frompyfunc(math.erf, 1, 1)
+        positions = np.arange(length)
+        self.window = (
+            np.abs(positions[:, None] - positions[None, :])
+            <= int(self.encoder_cfg.get("local_attention", 128)) // 2
+        )
+
+    def model_inputs(self, batch):
+        expected_shapes = {
+            "input_ids": (1, self.shape["max_length"]),
+            "attention_mask": (1, self.shape["max_length"]),
+            "marker_pos": (1, 32),
+            "marker_mask": (1, 32),
+            "qtype": (1,),
+        }
+        if set(batch) != set(expected_shapes):
+            raise ValueError("Prepared batch fields do not match the fixed ANE signature")
+        for name, shape in expected_shapes.items():
+            if batch[name].shape != shape or not np.issubdtype(batch[name].dtype, np.integer):
+                raise ValueError(f"{name} must have integer dtype and shape {shape}")
+        if np.any(batch["input_ids"] < 0) or np.any(batch["input_ids"] >= self.embedding.shape[0]):
+            raise ValueError("Token id outside checkpoint vocabulary")
+        for name in ("attention_mask", "marker_mask"):
+            if not np.isin(batch[name], (0, 1)).all():
+                raise ValueError(f"{name} must contain only zero or one")
+        if not batch["attention_mask"].any(axis=-1).all():
+            raise ValueError("Every batch row needs at least one valid attention key")
+        if np.any(batch["qtype"] < 0) or np.any(batch["qtype"] > 2):
+            raise ValueError("Question type must be 0, 1 or 2")
+        if np.any(batch["marker_pos"] < 0) or np.any(
+            batch["marker_pos"] >= self.shape["max_length"]
+        ):
+            raise ValueError("Marker position outside exported sequence")
+        ids, valid = batch["input_ids"], batch["attention_mask"].astype(bool)
+        embeddings = self.embedding[ids].transpose(0, 2, 1)[:, :, None, :]
+        full = np.broadcast_to(valid[:, None, :], (ids.shape[0], ids.shape[1], ids.shape[1]))
+        local = (self.window[None] | ~valid[:, :, None]) & full
+        # Core ML BC1S attention scores are [B,key,1,query].
+        masks = {"full_mask": full, "local_mask": local}
+        result = {
+            name: np.where(value.transpose(0, 2, 1)[:, :, None, :], 0, -1e4).astype(np.float16)
+            for name, value in masks.items()
+        }
+        result["embeddings"] = np.ascontiguousarray(embeddings, dtype=np.float16)
+        result["type_vectors"] = np.ascontiguousarray(
+            self.type_embedding[batch["qtype"]][:, :, None, None], dtype=np.float16
+        )
+        marker_map = np.zeros((ids.shape[0], ids.shape[1], 1, 32), np.float16)
+        for row in range(ids.shape[0]):
+            marker_map[row, batch["marker_pos"][row], 0, np.arange(32)] = 1
+        result["marker_map"] = marker_map
+        return result
+
+    def forward(self, batch):
+        outputs = self.model.predict(self.model_inputs(batch))
+        # Output names are traced identifiers; shapes uniquely identify these outputs.
+        logits = (
+            next(v for v in outputs.values() if v.shape[1] == 1).reshape(1, 32).astype(np.float32)
+        )
+        pooled = (
+            next(v for v in outputs.values() if v.shape[1] != 1).reshape(1, -1).astype(np.float32)
+        )
+        logits = np.where(batch["marker_mask"].astype(bool), logits, -1e4)
+        p = np.exp(logits - logits.max(axis=-1, keepdims=True))
+        p /= p.sum(axis=-1, keepdims=True)
+        k = np.maximum(batch["marker_mask"].sum(axis=-1), 2).astype(np.float32)
+        entropy = -(p * np.log(np.maximum(p, 1e-9))).sum(axis=-1) / np.log(k)
+        top = np.sort(p, axis=-1)[:, -2:]
+        features = np.stack((top[:, 1], top[:, 1] - top[:, 0], entropy, k / 255.0), axis=-1)
+        action_input = np.concatenate((pooled, features), axis=-1)
+        hidden = action_input @ self.action["0.weight"].T + self.action["0.bias"]
+        # Only 256 host elements: exact erf GELU, no tanh/sigmoid approximation.
+        hidden = hidden * (1 + self._erf(hidden / np.sqrt(2)).astype(np.float32)) / 2
+        action = hidden @ self.action["2.weight"].T + self.action["2.bias"]
+        return logits, action

laya_coreml/artifacts.py

@@ -0,0 +1,90 @@
+"""Integrity checks for portable Core ML bundles."""
+
+import errno
+import hashlib
+import os
+import shutil
+import tempfile
+from pathlib import Path, PurePosixPath
+
+
+def file_digest(path):
+    digest = hashlib.sha256()
+    with Path(path).open("rb") as stream:
+        for block in iter(lambda: stream.read(8 * 1024**2), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def tree_digest(path):
+    path = Path(path)
+    digest = hashlib.sha256()
+    for file in sorted(path.rglob("*")):
+        if file.is_file():
+            digest.update(str(file.relative_to(path)).encode())
+            digest.update(bytes.fromhex(file_digest(file)))
+    return digest.hexdigest()
+
+
+def package_for_coreml(package):
+    """Materialize Hub symlinks: the native compiler can copy them into broken paths."""
+    package = Path(package)
+    if not package.is_symlink() and not any(path.is_symlink() for path in package.rglob("*")):
+        return package
+    digest = tree_digest(package)
+    root = (
+        Path(os.environ.get("LAYA_COREML_CACHE", Path.home() / ".cache/laya-coreml")) / "packages"
+    )
+    root.mkdir(parents=True, exist_ok=True)
+    destination = root / digest
+    target = destination / "model.mlpackage"
+    if destination.exists():
+        if not target.is_dir() or tree_digest(target) != digest:
+            raise ValueError(f"Core ML cache integrity failure; remove {destination} and retry")
+        return target
+    temporary = Path(tempfile.mkdtemp(prefix=".preparing-", dir=root))
+    try:
+        copied = temporary / "model.mlpackage"
+        shutil.copytree(package, copied, symlinks=False)
+        if tree_digest(copied) != digest:
+            raise ValueError("Core ML package changed while materializing cached weights")
+        try:
+            temporary.rename(destination)
+        except OSError as error:
+            if error.errno not in (errno.EEXIST, errno.ENOTEMPTY) or not target.is_dir():
+                raise
+        if tree_digest(target) != digest:
+            raise ValueError("Materialized Core ML package failed its integrity check")
+    finally:
+        shutil.rmtree(temporary, ignore_errors=True)
+    return target
+
+
+def verify_files(directory, files):
+    for name, expected in files.items():
+        relative = PurePosixPath(name)
+        if relative.is_absolute() or ".." in relative.parts or "\\" in name:
+            raise ValueError("Bundle file names must be relative paths inside the model directory")
+        if file_digest(Path(directory) / name) != expected["sha256"]:
+            raise ValueError(f"Bundle file does not match its manifest: {name}")
+
+
+def verify_research_manifest(source, package, *, length):
+    import json
+
+    source, package = Path(source), Path(package)
+    manifest = json.loads((package.parent / "manifest.json").read_text())
+    if manifest.get("format") != "laya-ane-research" or manifest.get("version") != 1:
+        raise ValueError("Unsupported ANE research artifact manifest")
+    if manifest.get("kind") != "body" or manifest.get("shape") != {
+        "batch": 1,
+        "length": length,
+        "options": 32,
+    }:
+        raise ValueError("Requested runtime shape does not match artifact manifest")
+    verify_files(
+        source, {name: {"sha256": sha} for name, sha in manifest["source_files_sha256"].items()}
+    )
+    if tree_digest(package) != manifest["package_sha256"]:
+        raise ValueError("ANE package content does not match its manifest")
+    return manifest

laya_coreml/cli.py

@@ -0,0 +1,73 @@
+"""Convert and run Laya Core ML packages."""
+
+import argparse
+import json
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    commands = parser.add_subparsers(dest="command", required=True)
+    export = commands.add_parser("convert", help="Export original Laya weights to Core ML")
+    export.add_argument("source", help="Local original checkpoint, or pinned Laya model name")
+    export.add_argument("output")
+    export.add_argument("--max-length", type=int)
+    export.add_argument("--batch-size", type=int, default=1)
+    export.add_argument("--max-options", type=int, default=32)
+    export.add_argument("--fixed", action="store_true")
+    export.add_argument("--precision", choices=["float16", "float32"], default="float16")
+    export.add_argument("--revision")
+    export.add_argument("--attention", choices=["explicit", "sdpa"], default="sdpa")
+    export.add_argument("--shape-mode", choices=["enumerated", "range"], default="enumerated")
+    predict = commands.add_parser("predict")
+    predict.add_argument("model_dir")
+    predict.add_argument("--state", required=True, help="Literal state text")
+    predict.add_argument(
+        "--questions", required=True, help="JSON file containing question definitions"
+    )
+    predict.add_argument(
+        "--compute-units",
+        choices=["all", "cpu", "cpu_gpu", "cpu_ne"],
+        help="Default: cpu_ne for ANE bundles; cpu_gpu for ordinary exports",
+    )
+    predict.add_argument(
+        "--offline", action="store_true", help="Use local files or cached Hub snapshots only"
+    )
+    predict.add_argument("--revision", help="Pinned Hugging Face commit or revision")
+    args = parser.parse_args()
+    if args.command == "convert":
+        from .convert import convert
+
+        convert(
+            args.source,
+            args.output,
+            max_length=args.max_length,
+            flexible=not args.fixed,
+            batch_size=args.batch_size,
+            max_options=args.max_options,
+            precision=args.precision,
+            revision=args.revision,
+            attention=args.attention,
+            shape_mode=args.shape_mode,
+        )
+    else:
+        from pathlib import Path
+
+        from .agent import load
+
+        agent = load(
+            args.model_dir,
+            compute_units=args.compute_units,
+            local_files_only=args.offline,
+            revision=args.revision,
+        )
+        print(
+            json.dumps(
+                agent.predict(args.state, json.loads(Path(args.questions).read_text())),
+                ensure_ascii=False,
+                indent=2,
+            )
+        )
+
+
+if __name__ == "__main__":
+    main()

laya_coreml/common.py

@@ -0,0 +1,119 @@
+"""Laya prompt construction and calibration, adapted from upstream (see NOTICE)."""
+
+import json
+import math
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+
+QTYPES = {"choice": 0, "score": 1, "noul": 2}
+QTYPE_NAMES = {v: k for k, v in QTYPES.items()}
+
+
+def serialize_state(state: Union[str, dict, list]) -> str:
+    if isinstance(state, str):
+        return state
+    return json.dumps(state, ensure_ascii=False)
+
+
+def render_criterion(value) -> str:
+    """Render one criterion value as text.
+
+    Strings pass through; anything structured (dict, list, number) becomes compact JSON, so a
+    rubric reads as JSON rather than a Python repr. Without this a dict-valued criterion
+    crashed `noul` outright and leaked `{'desc': ...}` into `choice` and `score` prompts.
+    """
+    if isinstance(value, str):
+        return value
+    return json.dumps(value, ensure_ascii=False, separators=(", ", ": "), default=str)
+
+
+def render_options(q: Dict) -> List[str]:
+    """Render option texts in label-index order. Noul is always [false, true]."""
+    t, crit = q["t"], q.get("crit")
+    if t == "choice":
+        # only None/"" mean "no description"; 0 and False are legitimate criterion values
+        return [
+            k if v is None or v == "" else "%s: %s" % (k, render_criterion(v))
+            for k, v in crit.items()
+        ]
+    if t == "score":
+        return ["level %d: %s" % (i, render_criterion(c)) for i, c in enumerate(crit)]
+    crit = crit or {}
+    false_crit, true_crit = crit.get("false"), crit.get("true")
+    return [
+        "false: "
+        + (
+            render_criterion(false_crit)
+            if false_crit not in (None, "")
+            else "no, the statement does not hold"
+        ),
+        "true: "
+        + (
+            render_criterion(true_crit)
+            if true_crit not in (None, "")
+            else "yes, the statement holds"
+        ),
+    ]
+
+
+def build_prefix(tok, q: Dict, head_max_len: int = 192, option_order=None):
+    """Build the question-only prefix, before state tokens and final truncation."""
+    mask_tok = tok.mask_token
+    opts = render_options(q)
+    order = option_order if option_order is not None else list(range(len(opts)))
+    ins = str(q["ins"]).replace(mask_tok, " ")
+    head_ids = tok("%s question: %s" % (q["t"], ins), add_special_tokens=False)["input_ids"]
+    opt_ids = []
+    for i in order:
+        opt_ids.append(
+            [tok.mask_token_id]
+            + tok(" " + opts[i].replace(mask_tok, " "), add_special_tokens=False)["input_ids"][:48]
+        )
+    opt_budget = head_max_len - sum(len(o) for o in opt_ids)
+    if opt_budget < 16:
+        per = max(4, (head_max_len - 16) // max(1, len(opt_ids)))
+        opt_ids = [o[:per] for o in opt_ids]
+        opt_budget = head_max_len - sum(len(o) for o in opt_ids)
+    head_ids = head_ids[: max(8, opt_budget)]
+    ids = [tok.cls_token_id] + head_ids + [tok.sep_token_id]
+    markers = []
+    for o in opt_ids:
+        markers.append(len(ids))
+        ids.extend(o)
+    ids.append(tok.sep_token_id)
+    return ids, markers
+
+
+def build_sequence(
+    tok,
+    state: Union[str, dict, list],
+    q: Dict,
+    max_len: int = 512,
+    head_max_len: int = 192,
+    option_order: Optional[List[int]] = None,
+    truncate_left: bool = False,
+):
+    """Format: [CLS] <type> instructions [SEP] [MASK] opt0 [MASK] opt1 ... [SEP] state [SEP]."""
+    ids, markers = build_prefix(tok, q, head_max_len, option_order)
+    room = max(0, max_len - len(ids) - 1)
+    st = tok(serialize_state(state).replace(tok.mask_token, " "), add_special_tokens=False)[
+        "input_ids"
+    ]
+    st = st[-room:] if truncate_left else st[:room]
+    ids = ids + st + [tok.sep_token_id]
+    return ids[:max_len], [m for m in markers if m < max_len]
+
+
+def confidence_from_probs(p: np.ndarray, k: int) -> float:
+    """Normalized Shannon entropy confidence: 1 - H(p) / log(k)."""
+    if k < 2:
+        return 1.0
+    p = p[:k]
+    ent = -(p * np.log(np.clip(p, 1e-12, 1.0))).sum()
+    return float(np.clip(1.0 - ent / math.log(k), 0.0, 1.0))
+
+
+def temp_bucket(qtype: int, k: int) -> str:
+    size = "2" if k <= 2 else "3-5" if k <= 5 else "6-10" if k <= 10 else "11+"
+    return "%s:%s" % (QTYPE_NAMES[int(qtype)], size)