planetbridging 0.7.0__py3-none-any.whl

planetbridging/__init__.py

@@ -0,0 +1,91 @@
+"""Planet Bridging — stream AI engine weights into Loom .entity checkpoints."""
+
+from .bedrocks import BEDROCK_IDS, BEDROCK_PLANETS, BEDROCKS, BedrockInfo, STREAM_MODEL_IDS
+from .bedrock_ladder import LadderResult, run_all_bedrock_ladders, run_bedrock_ladder
+from .bedrock_smoke import SmokeResult, run_all_bedrock_smokes, run_bedrock_smoke
+from .compare import compare_outputs, diff_label
+from . import absorb
+from . import engines
+from .engines import EngineStreamResult
+from .stream import (
+    StreamResult,
+    stream_bedrock,
+    stream_cnn1,
+    stream_cnn2,
+    stream_cnn3,
+    stream_dense,
+    stream_embedding,
+    stream_layernorm,
+    stream_lstm,
+    stream_mha,
+    stream_mixer,
+    stream_residual,
+    stream_rmsnorm,
+    stream_rnn,
+    stream_swiglu,
+)
+
+try:
+    from .welvet_infer import (
+        import_welvet,
+        infer_bedrock_entity,
+        infer_dense_entity,
+        infer_mixer_v1_entity,
+        print_compare_ladder,
+        try_infer_bedrock_entity,
+        WELVET_RELOAD_BEDROCKS,
+        WELVET_RELOAD_SKIP,
+    )
+except ImportError:  # pragma: no cover
+    import_welvet = None  # type: ignore
+    infer_bedrock_entity = None  # type: ignore
+    infer_dense_entity = None  # type: ignore
+    infer_mixer_v1_entity = None  # type: ignore
+    print_compare_ladder = None  # type: ignore
+    try_infer_bedrock_entity = None  # type: ignore
+    WELVET_RELOAD_BEDROCKS = frozenset()  # type: ignore
+    WELVET_RELOAD_SKIP = frozenset()  # type: ignore
+
+__version__ = "0.7.0"
+__all__ = [
+    "__version__",
+    "BEDROCK_IDS",
+    "BEDROCK_PLANETS",
+    "BEDROCKS",
+    "BedrockInfo",
+    "STREAM_MODEL_IDS",
+    "StreamResult",
+    "EngineStreamResult",
+    "SmokeResult",
+    "LadderResult",
+    "engines",
+    "absorb",
+    "run_bedrock_ladder",
+    "run_all_bedrock_ladders",
+    "stream_bedrock",
+    "stream_dense",
+    "stream_cnn1",
+    "stream_cnn2",
+    "stream_cnn3",
+    "stream_mha",
+    "stream_lstm",
+    "stream_rnn",
+    "stream_layernorm",
+    "stream_embedding",
+    "stream_rmsnorm",
+    "stream_swiglu",
+    "stream_residual",
+    "stream_mixer",
+    "run_bedrock_smoke",
+    "run_all_bedrock_smokes",
+    "compare_outputs",
+    "diff_label",
+    "import_welvet",
+    "infer_dense_entity",
+    "infer_bedrock_entity",
+    "try_infer_bedrock_entity",
+    "infer_mixer_v1_entity",
+    "print_compare_ladder",
+    "WELVET_RELOAD_BEDROCKS",
+    "WELVET_RELOAD_SKIP",
+]

planetbridging/_binary.py

@@ -0,0 +1,86 @@
+"""Locate the loom-stream Go binary (stdlib subprocess backend)."""
+
+from __future__ import annotations
+
+import json
+import os
+import platform
+import shutil
+import stat
+import subprocess
+from pathlib import Path
+
+
+def repo_root() -> Path:
+    """Planetbridging repo root (../../ from src/planetbridging/)."""
+    return Path(__file__).resolve().parents[2]
+
+
+def default_binary_path() -> Path:
+    name = "loom-stream.exe" if platform.system() == "Windows" else "loom-stream"
+    return repo_root() / "bin" / name
+
+
+def find_loom_stream(explicit: str | Path | None = None) -> Path:
+    """Resolve loom-stream: explicit → env → PATH → repo bin/."""
+    if explicit is not None:
+        path = Path(explicit).expanduser().resolve()
+        if not path.is_file():
+            raise FileNotFoundError(f"loom-stream not found: {path}")
+        return path
+
+    env = os.environ.get("PLANETBRIDGING_LOOM_STREAM")
+    if env:
+        path = Path(env).expanduser().resolve()
+        if path.is_file():
+            return path
+
+    which = shutil.which("loom-stream")
+    if which:
+        return Path(which)
+
+    candidate = default_binary_path()
+    if candidate.is_file():
+        return candidate
+
+    raise FileNotFoundError(
+        "loom-stream binary not found. Build from the planetbridging repo:\n"
+        "  go build -o bin/loom-stream ./cmd/loom-stream/\n"
+        "Or set PLANETBRIDGING_LOOM_STREAM=/path/to/loom-stream"
+    )
+
+
+def ensure_executable(path: Path) -> None:
+    if not os.access(path, os.X_OK):
+        path.chmod(path.stat().st_mode | stat.S_IXUSR)
+
+
+def run_loom_stream(
+    *,
+    binary: Path | None,
+    envelope: dict,
+    payload: dict,
+    timeout: float = 120.0,
+) -> dict:
+    """Invoke loom-stream with JSON on stdin, parse JSON stdout."""
+    exe = find_loom_stream(binary)
+    ensure_executable(exe)
+
+    body = {**envelope, "payload": payload}
+    proc = subprocess.run(
+        [str(exe)],
+        input=json.dumps(body).encode("utf-8"),
+        capture_output=True,
+        timeout=timeout,
+        check=False,
+    )
+    if proc.returncode != 0:
+        stderr = proc.stderr.decode("utf-8", errors="replace").strip()
+        stdout = proc.stdout.decode("utf-8", errors="replace").strip()
+        detail = stderr or stdout or f"exit {proc.returncode}"
+        raise RuntimeError(f"loom-stream failed: {detail}")
+
+    try:
+        return json.loads(proc.stdout.decode("utf-8"))
+    except json.JSONDecodeError as exc:
+        raise RuntimeError(f"loom-stream returned invalid JSON: {proc.stdout!r}") from exc

planetbridging/absorb.py

@@ -0,0 +1,188 @@
+"""High-level absorb API — stream live models into Loom .entity for any bedrock."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Sequence
+
+import numpy as np
+
+from . import engines
+from .engines import EngineStreamResult
+from .layers.dense import LayerSpec, layers_from_jax, layers_from_keras, layers_from_pytorch, layers_from_sklearn
+from .stream import StreamResult, stream_dense
+
+
+def stream(
+    bedrock: str,
+    planet: str,
+    *,
+    model_id: str | None = None,
+    **kwargs: Any,
+) -> EngineStreamResult:
+    """Stream any bedrock layer type from a live AI engine (pytorch/tensorflow/jax/sklearn)."""
+    return engines.stream(bedrock, planet, model_id=model_id, **kwargs)
+
+
+def stream_all_bedrocks(planet: str, **kwargs: Any) -> list[EngineStreamResult]:
+    return engines.stream_all_bedrocks(planet, **kwargs)
+
+
+def stream_all_planets(bedrock: str, **kwargs: Any) -> list[EngineStreamResult]:
+    return engines.stream_all_planets(bedrock, **kwargs)
+
+
+def _default_specs(units: Sequence[int], *, activation: str = "relu") -> tuple[LayerSpec, ...]:
+    return tuple(LayerSpec(u, activation) for u in units)
+
+
+def pytorch(
+    net: Any,
+    *,
+    planet: str = "pytorch",
+    model_id: str = "model",
+    input_dim: int,
+    layer_specs: Sequence[LayerSpec] | None = None,
+    layer_units: Sequence[int] | None = None,
+    activation: str = "relu",
+    inputs: np.ndarray | None = None,
+    native_outputs: np.ndarray | None = None,
+    output_path: str | Path | None = None,
+    models_dir: str | Path | None = None,
+    fixture_version: str = "",
+    root: str | Path | None = None,
+    **kwargs: Any,
+) -> StreamResult:
+    """Stream a PyTorch nn.Module (sequential Dense stack) into Loom .entity."""
+    specs = _resolve_specs(layer_specs, layer_units, activation)
+    layers = layers_from_pytorch(net, specs, input_dim=input_dim)
+    return stream_dense(
+        planet=planet,
+        model_id=model_id,
+        layers=layers,
+        input_dim=input_dim,
+        inputs=inputs,
+        native_outputs=native_outputs,
+        output_path=output_path,
+        models_dir=models_dir,
+        fixture_version=fixture_version,
+        root=root,
+        **kwargs,
+    )
+
+
+def keras(
+    net: Any,
+    *,
+    planet: str = "tensorflow",
+    model_id: str = "model",
+    input_dim: int,
+    layer_specs: Sequence[LayerSpec] | None = None,
+    layer_units: Sequence[int] | None = None,
+    activation: str = "relu",
+    inputs: np.ndarray | None = None,
+    native_outputs: np.ndarray | None = None,
+    output_path: str | Path | None = None,
+    models_dir: str | Path | None = None,
+    fixture_version: str = "",
+    root: str | Path | None = None,
+    **kwargs: Any,
+) -> StreamResult:
+    """Stream a Keras / TensorFlow model into Loom .entity."""
+    specs = _resolve_specs(layer_specs, layer_units, activation)
+    layers = layers_from_keras(net, specs, input_dim=input_dim)
+    return stream_dense(
+        planet=planet,
+        model_id=model_id,
+        layers=layers,
+        input_dim=input_dim,
+        inputs=inputs,
+        native_outputs=native_outputs,
+        output_path=output_path,
+        models_dir=models_dir,
+        fixture_version=fixture_version,
+        root=root,
+        **kwargs,
+    )
+
+
+def jax(
+    params: Any,
+    *,
+    planet: str = "jax",
+    model_id: str = "model",
+    input_dim: int,
+    layer_specs: Sequence[LayerSpec] | None = None,
+    layer_units: Sequence[int] | None = None,
+    activation: str = "relu",
+    inputs: np.ndarray | None = None,
+    native_outputs: np.ndarray | None = None,
+    output_path: str | Path | None = None,
+    models_dir: str | Path | None = None,
+    fixture_version: str = "",
+    root: str | Path | None = None,
+    **kwargs: Any,
+) -> StreamResult:
+    """Stream Flax/JAX params dict into Loom .entity."""
+    specs = _resolve_specs(layer_specs, layer_units, activation)
+    layers = layers_from_jax(params, specs, input_dim=input_dim)
+    return stream_dense(
+        planet=planet,
+        model_id=model_id,
+        layers=layers,
+        input_dim=input_dim,
+        inputs=inputs,
+        native_outputs=native_outputs,
+        output_path=output_path,
+        models_dir=models_dir,
+        fixture_version=fixture_version,
+        root=root,
+        **kwargs,
+    )
+
+
+def sklearn(
+    reg: Any,
+    *,
+    planet: str = "sklearn",
+    model_id: str = "model",
+    input_dim: int,
+    layer_specs: Sequence[LayerSpec] | None = None,
+    layer_units: Sequence[int] | None = None,
+    activation: str = "linear",
+    inputs: np.ndarray | None = None,
+    native_outputs: np.ndarray | None = None,
+    output_path: str | Path | None = None,
+    models_dir: str | Path | None = None,
+    fixture_version: str = "",
+    root: str | Path | None = None,
+    **kwargs: Any,
+) -> StreamResult:
+    """Stream sklearn MLPRegressor into Loom .entity."""
+    specs = _resolve_specs(layer_specs, layer_units, activation)
+    layers = layers_from_sklearn(reg, specs, input_dim=input_dim)
+    return stream_dense(
+        planet=planet,
+        model_id=model_id,
+        layers=layers,
+        input_dim=input_dim,
+        inputs=inputs,
+        native_outputs=native_outputs,
+        output_path=output_path,
+        models_dir=models_dir,
+        fixture_version=fixture_version,
+        root=root,
+        **kwargs,
+    )
+
+
+def _resolve_specs(
+    layer_specs: Sequence[LayerSpec] | None,
+    layer_units: Sequence[int] | None,
+    activation: str,
+) -> tuple[LayerSpec, ...]:
+    if layer_specs is not None:
+        return tuple(layer_specs)
+    if layer_units is None:
+        raise ValueError("layer_specs or layer_units required")
+    return _default_specs(layer_units, activation=activation)

planetbridging/bedrock_ladder.py

@@ -0,0 +1,173 @@
+"""Full compare ladder per bedrock: native → loom-stream → welvet reload."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from .bedrock_smoke import (
+    _SMOKE_RUNNERS,
+    _load_bedrock_module,
+    _load_fixture_npz,
+    _repo,
+)
+from .bedrocks import BEDROCK_IDS, BEDROCKS, BedrockInfo
+from .compare import compare_outputs, diff_label
+from .stream import StreamResult
+from .welvet_infer import (
+    WELVET_RELOAD_SKIP,
+    try_infer_bedrock_entity,
+)
+
+
+@dataclass(frozen=True)
+class LadderResult:
+    bedrock: str
+    model_id: str
+    planet: str
+    layer_count: int
+    entity_path: str
+    native: np.ndarray
+    loom_stream: np.ndarray
+    welvet: np.ndarray | None
+    native_vs_loom: str
+    native_vs_welvet: str | None
+    loom_vs_welvet: str | None
+    welvet_note: str | None = None
+
+
+def _fixture_inputs(
+    bedrock: str,
+    root: Path,
+    model: Any,
+    manifest: Any,
+    fixtures_mod: Any,
+    data: dict[str, np.ndarray],
+    out_dim: int,
+) -> tuple[np.ndarray, int | None]:
+    """Return (x_test, input_dim_for_dense)."""
+    if bedrock == "dense":
+        x = fixtures_mod.ensure_fixtures(manifest)["x_test"][:, : model.input_dim]
+        return x, int(model.input_dim)
+    if bedrock == "mixer":
+        return data["x_test"], None
+    if bedrock == "residual":
+        _, _, _, xm, xs, _ = fixtures_mod.slice_model_inputs(data, model, out_dim)
+        # Residual infer uses paired tensors; welvet skip uses main only for display.
+        return xm, None
+    sliced = fixtures_mod.slice_model_inputs(data, model, out_dim)
+    return sliced[2], None
+
+
+def run_bedrock_ladder(
+    bedrock: str,
+    *,
+    root: Path | None = None,
+    model_id: str | None = None,
+    binary: Path | None = None,
+    out_dir: Path | None = None,
+    max_samples: int = 5,
+    try_welvet: bool = True,
+    planet: str = "numpy",
+) -> LadderResult:
+    """Stream one bedrock model and compare native / loom-stream / welvet reload."""
+    bedrock = bedrock.lower()
+    if bedrock not in BEDROCKS:
+        raise ValueError(f"unknown bedrock {bedrock!r}")
+
+    info = BEDROCKS[bedrock]
+    mid = model_id or info.smoke_model_id
+    root = _repo(root)
+    runner = _SMOKE_RUNNERS.get(bedrock)
+    if runner is None:
+        raise ValueError(f"no smoke runner for {bedrock}")
+
+    stream: StreamResult = runner(root, mid, info, binary=binary, out_dir=out_dir)
+    native = stream.native_reference
+    if native is None:
+        raise RuntimeError(f"{bedrock}: stream result missing native_reference")
+
+    manifest_mod = _load_bedrock_module(root, bedrock, "manifest")
+    fixtures_mod = _load_bedrock_module(root, bedrock, "fixtures")
+    manifest = manifest_mod.load_manifest()
+    model = next(m for m in manifest.models if m.id == mid)
+    data = _load_fixture_npz(root, bedrock, info.fixture_version)
+    out_dim = manifest_mod.model_output_dim(model)
+    x_test, input_dim = _fixture_inputs(bedrock, root, model, manifest, fixtures_mod, data, out_dim)
+
+    n = min(max_samples, len(native), len(stream.outputs))
+    native_s = native[:n]
+    loom_s = stream.outputs[:n]
+    x_s = x_test[:n]
+
+    n_max, _, n_exact = compare_outputs(native_s, loom_s)
+    native_vs_loom = diff_label(n_max, exact=n_exact)
+
+    welvet_out: np.ndarray | None = None
+    native_vs_welvet: str | None = None
+    loom_vs_welvet: str | None = None
+    welvet_note: str | None = None
+
+    if bedrock in WELVET_RELOAD_SKIP:
+        welvet_note = "welvet reload skipped (C ABI gap on this bedrock)"
+    elif not try_welvet:
+        welvet_note = "welvet not requested"
+    else:
+        welvet_out = try_infer_bedrock_entity(
+            bedrock,
+            stream.entity_path,
+            x_s,
+            output_dim=out_dim if bedrock != "dense" else native_s.shape[-1],
+            input_dim=input_dim,
+        )
+        if welvet_out is None:
+            welvet_note = "welvet reload failed or unavailable"
+        else:
+            w_max, _, w_exact = compare_outputs(native_s, welvet_out)
+            lw_max, _, lw_exact = compare_outputs(loom_s, welvet_out)
+            native_vs_welvet = diff_label(w_max, exact=w_exact)
+            loom_vs_welvet = diff_label(lw_max, exact=lw_exact)
+
+    return LadderResult(
+        bedrock=bedrock,
+        model_id=mid,
+        planet=planet,
+        layer_count=stream.layer_count,
+        entity_path=stream.entity_path,
+        native=native_s,
+        loom_stream=loom_s,
+        welvet=welvet_out,
+        native_vs_loom=native_vs_loom,
+        native_vs_welvet=native_vs_welvet,
+        loom_vs_welvet=loom_vs_welvet,
+        welvet_note=welvet_note,
+    )
+
+
+def run_all_bedrock_ladders(**kwargs: Any) -> list[LadderResult]:
+    results: list[LadderResult] = []
+    for bid in BEDROCK_IDS:
+        try:
+            results.append(run_bedrock_ladder(bid, **kwargs))
+        except Exception as exc:
+            info = BEDROCKS[bid]
+            results.append(
+                LadderResult(
+                    bedrock=bid,
+                    model_id=info.smoke_model_id,
+                    planet=str(kwargs.get("planet", "numpy")),
+                    layer_count=0,
+                    entity_path=str(exc),
+                    native=np.zeros((0, 0)),
+                    loom_stream=np.zeros((0, 0)),
+                    welvet=None,
+                    native_vs_loom="ERROR",
+                    native_vs_welvet=None,
+                    loom_vs_welvet=None,
+                    welvet_note=str(exc),
+                )
+            )
+    return results