nneditor 1.0.0__py3-none-any.whl

nneditor/__init__.py

@@ -0,0 +1,5 @@
+"""NNEditor package."""
+
+__all__ = ["__version__"]
+
+__version__ = "1.0.0"

nneditor/__main__.py

@@ -0,0 +1,12 @@
+"""Run NNEditor with ``python -m nneditor``."""
+
+from nneditor.ui.app import cli
+
+
+def main() -> None:
+    """Launch the Flet application."""
+    cli()
+
+
+if __name__ == "__main__":  # pragma: no cover - interpreter entry point
+    main()

nneditor/adapters/__init__.py

@@ -0,0 +1,6 @@
+"""Source-format adapters.
+
+Each adapter maps one artifact contract from :mod:`nneditor.ir.capabilities`
+into the canonical IR, preserving source-specific data in extension namespaces
+rather than discarding it.
+"""

nneditor/adapters/detect.py

@@ -0,0 +1,91 @@
+"""Artifact-kind detection by content, not by extension (P6.5/P7.5).
+
+Architectural rule 2 makes format support an artifact-and-capability matrix
+rather than an extension list, so opening dispatches on what a file *is*:
+its magic bytes, and for zip containers the members it holds. Extensions are
+only a tie-breaker for formats whose content is genuinely ambiguous (textual
+MLIR versus arbitrary text).
+"""
+
+from __future__ import annotations
+
+import struct
+import zipfile
+from pathlib import Path
+
+from nneditor.ir.capabilities import ArtifactKind
+
+__all__ = ["DetectionError", "detect_artifact_kind"]
+
+_ZIP_MAGIC = b"PK\x03\x04"
+_ONNX_PROTO_HINT = b"\x08"  # ModelProto.ir_version, field 1 varint
+
+
+class DetectionError(ValueError):
+    """Raised when a file matches no supported artifact contract."""
+
+
+def _zip_kind(path: Path) -> ArtifactKind:
+    try:
+        with zipfile.ZipFile(path) as archive:
+            names = archive.namelist()
+    except zipfile.BadZipFile as error:
+        raise DetectionError(f"{path.name} is not a readable archive") from error
+    # data.pkl wins over archive_format when both are present: a torch
+    # checkpoint always holds data.pkl, while anyone can append a spare
+    # archive_format member to reroute it into the PT2 parser.
+    if any(name.endswith("data.pkl") for name in names):
+        # Both plain state dicts and pickled FX modules live here; the FX
+        # reader is the one that can say which, so the caller tries the
+        # checkpoint path first and falls back on a scan refusal.
+        return ArtifactKind.PYTORCH_STATE_DICT
+    if any(name.endswith("archive_format") for name in names):
+        return ArtifactKind.PYTORCH_EXPORTED_PROGRAM
+    raise DetectionError(
+        f"{path.name} is a zip archive but holds no recognized model container"
+    )
+
+
+def _leading_content(prefix: bytes) -> bytes:
+    """The first content-bearing line, with ``//`` comment lines skipped.
+
+    StableHLO exporters routinely emit header comments before ``module``, so
+    the probe must look at the first non-comment line.
+    """
+    for line in prefix.split(b"\n"):
+        stripped = line.lstrip()
+        if not stripped or stripped.startswith(b"//"):
+            continue
+        return stripped
+    return b""
+
+
+def detect_artifact_kind(path: Path | str) -> ArtifactKind:
+    """The artifact contract a file satisfies, decided from its content."""
+    target = Path(path)
+    with open(target, "rb") as handle:
+        prefix = handle.read(4096)
+    if not prefix:
+        raise DetectionError(f"{target.name} is empty")
+    if prefix.startswith(_ZIP_MAGIC):
+        return _zip_kind(target)
+    # safetensors: an 8-byte little-endian header length that fits inside
+    # the file, followed immediately by '{'. Without the length check any
+    # file with a '{' at offset 8 would claim the format.
+    if len(prefix) > 8 and prefix[8:9] == b"{":
+        (header_length,) = struct.unpack("<Q", prefix[:8])
+        if 8 + header_length <= target.stat().st_size:
+            return ArtifactKind.SAFETENSORS
+    content = _leading_content(prefix)
+    if content.startswith(b"module") or content.startswith(b'"builtin.module"'):
+        return ArtifactKind.JAX_STABLEHLO
+    if prefix.startswith(_ONNX_PROTO_HINT):
+        return ArtifactKind.ONNX_MODEL
+    if target.suffix.lower() in {".onnx", ".pb"}:
+        return ArtifactKind.ONNX_MODEL
+    if target.suffix.lower() in {".mlir", ".stablehlo"}:
+        return ArtifactKind.JAX_STABLEHLO
+    raise DetectionError(
+        f"{target.name} matches no supported artifact contract; see "
+        "docs/artifact-capabilities.md for the accepted inputs"
+    )

nneditor/adapters/jax/__init__.py

@@ -0,0 +1,23 @@
+"""JAX / StableHLO artifact adapters (Phase 7).
+
+Textual StableHLO modules — what ``jax.export`` and ``jax.jit(...).lower()``
+emit — are parsed natively: MLIR's textual form is regular enough that a
+focused reader recovers functions, operations, values, and types without
+linking MLIR or executing any JAX code. Orbax/Flax checkpoint trees open
+weights-only. ``jax`` itself is a dev-only dependency used to generate
+fixtures and validate this reader.
+"""
+
+from nneditor.adapters.jax.checkpoint import open_orbax_checkpoint
+from nneditor.adapters.jax.stablehlo import (
+    StableHloError,
+    open_stablehlo,
+    parse_stablehlo,
+)
+
+__all__ = [
+    "StableHloError",
+    "open_orbax_checkpoint",
+    "open_stablehlo",
+    "parse_stablehlo",
+]

nneditor/adapters/jax/checkpoint.py

@@ -0,0 +1,158 @@
+"""Weights-only JAX/Flax/Orbax checkpoint sessions (P7.3).
+
+An Orbax checkpoint directory holds one msgpack/JSON tree plus array data;
+the portable, universally readable member is the **safetensors** or ``.npy``
+payload many JAX pipelines emit alongside. This adapter opens the directory's
+readable tensor files as a weights-only document with the *checkpoint tree*
+preserved as a path-keyed structure — it never claims topology, because a
+parameter tree describes no computation (the Orbax contract says exactly
+that).
+"""
+
+from __future__ import annotations
+
+import json
+from itertools import islice
+from pathlib import Path
+from typing import Final
+
+from nneditor import __version__
+from nneditor.adapters.pytorch.safetensors import open_safetensors
+from nneditor.diagnostics import DiagnosticLog, Severity
+from nneditor.ir.capabilities import ArtifactKind, contract_for
+from nneditor.ir.core import (
+    ArtifactRef,
+    Document,
+    Graph,
+    JsonValue,
+    ProvenanceEntry,
+    TensorRef,
+)
+from nneditor.ir.identity import ROOT_GRAPH_ID
+from nneditor.storage.reader import hash_file
+
+__all__ = ["OrbaxError", "open_orbax_checkpoint"]
+
+_MAX_TREE_FILES: Final = 256
+_MAX_TREE_FILE_BYTES: Final = 8 * 1024 * 1024
+
+
+class OrbaxError(ValueError):
+    """Raised when a checkpoint directory yields no readable tensors."""
+
+
+def open_orbax_checkpoint(path: Path | str) -> Document:
+    """Open a checkpoint directory (or single safetensors file) weights-only.
+
+    A directory containing ``*.safetensors`` reads through the native
+    safetensors path; ``.npy`` members are read directly. The parameter tree
+    (``*.json`` metadata) is carried as an extension without interpretation.
+
+    Every failure — including unexpected ones from hostile input — surfaces
+    as :class:`OrbaxError`; the layered error contract in the session depends
+    on adapters never leaking raw exceptions.
+    """
+    target = Path(path)
+    try:
+        return _open_orbax_checkpoint(target)
+    except OrbaxError:
+        raise
+    except Exception as error:
+        raise OrbaxError(
+            f"{target.name} could not be opened as a checkpoint: "
+            f"{type(error).__name__}: {error}"
+        ) from error
+
+
+def _open_orbax_checkpoint(target: Path) -> Document:
+    if target.is_file():
+        document = open_safetensors(target)
+        return document
+
+    if not target.is_dir():
+        raise OrbaxError(f"{target} is neither a file nor a directory")
+
+    log = DiagnosticLog()
+    safetensor_files = sorted(target.rglob("*.safetensors"))
+    tensors: list[TensorRef] = []
+    tree: dict[str, JsonValue] = {}
+
+    # The directory contents are untrusted: both the number of metadata
+    # files and each file's size are capped so a hostile checkpoint cannot
+    # turn the tree walk into unbounded reads.
+    metadata_files = sorted(islice(target.rglob("*.json"), _MAX_TREE_FILES + 1))
+    if len(metadata_files) > _MAX_TREE_FILES:
+        log.add(
+            "jax.tree-metadata-truncated",
+            Severity.WARNING,
+            f"the checkpoint holds more than {_MAX_TREE_FILES} metadata "
+            "files; only the first ones are carried",
+            target.name,
+        )
+        metadata_files = metadata_files[:_MAX_TREE_FILES]
+    for metadata in metadata_files:
+        try:
+            if metadata.stat().st_size > _MAX_TREE_FILE_BYTES:
+                log.add(
+                    "jax.tree-metadata-too-large",
+                    Severity.WARNING,
+                    f"{metadata.name} exceeds the {_MAX_TREE_FILE_BYTES}-byte "
+                    "metadata limit and is skipped",
+                    metadata.name,
+                )
+                continue
+            tree[metadata.relative_to(target).as_posix()] = json.loads(
+                metadata.read_text(encoding="utf-8")
+            )
+        except (OSError, ValueError, RecursionError):
+            log.add(
+                "jax.unreadable-tree-metadata",
+                Severity.WARNING,
+                f"{metadata.name} is not readable JSON and is skipped",
+                metadata.name,
+            )
+
+    if not safetensor_files:
+        raise OrbaxError(
+            f"{target} holds no safetensors payload; export the checkpoint "
+            "with a safetensors writer to inspect its weights"
+        )
+
+    primary = safetensor_files[0]
+    if len(safetensor_files) > 1:
+        log.add(
+            "jax.sharded-checkpoint",
+            Severity.WARNING,
+            f"{len(safetensor_files)} shards found; only {primary.name} is "
+            "opened in this build",
+            target.name,
+        )
+    document = open_safetensors(primary)
+    tensors = list(document.tensors.values())
+
+    extensions: list[tuple[str, JsonValue]] = list(document.extensions)
+    if tree:
+        extensions.append(("x-orbax.tree", tree))
+
+    contract = contract_for(ArtifactKind.ORBAX_CHECKPOINT)
+    return Document(
+        source=ArtifactRef(
+            path=str(primary),
+            content_hash=hash_file(primary),
+            byte_size=primary.stat().st_size,
+        ),
+        artifact_kind=ArtifactKind.ORBAX_CHECKPOINT,
+        capabilities=contract.statuses,
+        graphs=[Graph(id=ROOT_GRAPH_ID, name="checkpoint")],
+        tensors=tensors,
+        provenance=[
+            ProvenanceEntry(
+                operation="import",
+                tool_version=f"nneditor {__version__}",
+                target=ROOT_GRAPH_ID,
+                parameters=(("loading_mode", "safe artifact"),),
+            )
+        ],
+        diagnostics=(*document.diagnostics, *tuple(log)),
+        extensions=extensions,
+    )