invarlock 0.2.0__py3-none-any.whl

invarlock/adapters/hf_mixin.py

@@ -0,0 +1,383 @@
+"""
+Shared HuggingFace adapter mixin.
+=================================
+
+Provides reusable functionality for InvarLock's HuggingFace adapters:
+- Device resolution helpers
+- Snapshot/restore with device awareness
+- Chunked snapshot helpers to reduce peak memory usage
+- Lightweight config serialization
+- Weight-tying detection plumbing
+"""
+
+from __future__ import annotations
+
+import io
+import json
+import os
+import tempfile
+from pathlib import Path
+from typing import Any
+
+import torch
+
+from invarlock.security import is_secure_path
+
+SCALAR_TYPES = (int, float, str, bool)
+
+
+def _sanitize_param_name(name: str) -> str:
+    """Return a filesystem-safe parameter name."""
+    return name.replace(".", "__").replace("/", "_")
+
+
+def _ensure_secure_dir(path: Path) -> None:
+    """Ensure snapshot directory uses 0o700 permissions."""
+    path.mkdir(parents=True, exist_ok=True)
+    os.chmod(path, 0o700)
+    if not is_secure_path(path):
+        raise RuntimeError(
+            f"Snapshot directory {path} must have permissions 0o700 for security."
+        )
+
+
+def _resolve_named_parameter(
+    module: torch.nn.Module, path: str
+) -> torch.nn.Parameter | None:
+    """Resolve a parameter by dotted path, returning None if missing."""
+    current: Any = module
+    parts = path.split(".")
+    for name in parts[:-1]:
+        current = getattr(current, name, None)
+        if current is None:
+            return None
+    leaf = getattr(current, parts[-1], None)
+    if isinstance(leaf, torch.nn.Parameter):
+        return leaf
+    return None
+
+
+class HFAdapterMixin:
+    """Reusable utilities for HuggingFace-backed adapters."""
+
+    # ------------------------------------------------------------------
+    # Device helpers
+    # ------------------------------------------------------------------
+    def _resolve_device(
+        self, device: str | torch.device | None = "auto"
+    ) -> torch.device:
+        """
+        Resolve a target torch.device for model placement.
+
+        Args:
+            device: Requested device ("auto" selects CUDA→MPS→CPU).
+
+        Returns:
+            torch.device for placement.
+        """
+
+        if isinstance(device, torch.device):
+            return device
+
+        device_str = "auto" if device is None else str(device)
+        device_str = device_str.lower()
+
+        if device_str == "auto":
+            if torch.cuda.is_available():
+                return torch.device("cuda")
+            if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+                return torch.device("mps")
+            return torch.device("cpu")
+
+        return torch.device(device_str)
+
+    # ------------------------------------------------------------------
+    # HF save/export helpers
+    # ------------------------------------------------------------------
+    def save_pretrained(self, model: torch.nn.Module, path: str | Path) -> bool:
+        """
+        Save a HuggingFace model in a HF-loadable directory.
+
+        Args:
+            model: HF Transformers model implementing save_pretrained
+            path: Target directory path
+
+        Returns:
+            True on success, False otherwise
+        """
+        try:
+            p = Path(path)
+            p.mkdir(parents=True, exist_ok=True)
+            # Most HF models implement save_pretrained
+            save = getattr(model, "save_pretrained", None)
+            if callable(save):
+                save(str(p))
+                return True
+        except Exception:
+            return False
+        return False
+
+    # ------------------------------------------------------------------
+    # Snapshot / restore
+    # ------------------------------------------------------------------
+    def snapshot(self, model: torch.nn.Module) -> bytes:
+        """
+        Serialize model state with device awareness and weight-tying metadata.
+
+        Args:
+            model: HuggingFace model instance.
+
+        Returns:
+            Bytes payload produced by torch.save.
+        """
+
+        state_dict: dict[str, Any] = {}
+        device_map: dict[str, str] = {}
+
+        for name, param in model.named_parameters():
+            state_key = f"params.{name}"
+            state_dict[state_key] = param.detach().cpu().clone()
+            device_map[state_key] = str(param.device)
+
+        for name, buffer in model.named_buffers():
+            state_key = f"buffers.{name}"
+            state_dict[state_key] = buffer.detach().cpu().clone()
+            device_map[state_key] = str(buffer.device)
+
+        if hasattr(model, "config"):
+            state_dict["config"] = self._serialize_config(model.config)
+
+        state_dict["device_map"] = device_map
+        state_dict["model_class"] = model.__class__.__name__
+        state_dict["weight_tying"] = self._extract_weight_tying_info(model)
+
+        buffer = io.BytesIO()
+        torch.save(state_dict, buffer)
+        return buffer.getvalue()
+
+    def restore(self, model: torch.nn.Module, blob: bytes) -> None:
+        """
+        Restore model state produced by `snapshot`.
+
+        Args:
+            model: Model to restore in-place.
+            blob: Bytes payload from snapshot.
+        """
+
+        buffer = io.BytesIO(blob)
+        state_dict = torch.load(buffer, map_location="cpu", weights_only=False)
+
+        device_map: dict[str, str] = state_dict.get("device_map", {})
+
+        for name, param in model.named_parameters():
+            state_key = f"params.{name}"
+            if state_key not in state_dict:
+                continue
+            target_device = torch.device(device_map.get(state_key, "cpu"))
+            with torch.no_grad():
+                param.copy_(state_dict[state_key].to(target_device))
+
+        for name, buffer_param in model.named_buffers():
+            state_key = f"buffers.{name}"
+            if state_key not in state_dict:
+                continue
+            target_device = torch.device(device_map.get(state_key, "cpu"))
+            buffer_param.copy_(state_dict[state_key].to(target_device))
+
+        original_tying = state_dict.get("weight_tying", {})
+        if isinstance(original_tying, dict) and original_tying:
+            current_tying = self._extract_weight_tying_info(model)
+            for tied_param, source_param in original_tying.items():
+                if current_tying.get(tied_param) != source_param:
+                    self._restore_weight_tying(model, tied_param, source_param)
+
+    # ------------------------------------------------------------------
+    # Chunked snapshot helpers
+    # ------------------------------------------------------------------
+    def snapshot_chunked(
+        self, model: torch.nn.Module, *, prefix: str = "invarlock-snap-"
+    ) -> str:
+        """
+        Create a chunked snapshot on disk to minimise in-memory footprint.
+
+        Each parameter and buffer is serialized individually so only a single
+        tensor resides in memory at a time. Metadata is recorded in manifest.json.
+        """
+
+        snapshot_dir = Path(tempfile.mkdtemp(prefix=prefix))
+        _ensure_secure_dir(snapshot_dir)
+
+        manifest: dict[str, Any] = {
+            "model_class": model.__class__.__name__,
+            "config": self._serialize_config(model.config)
+            if hasattr(model, "config")
+            else {},
+            "params": {},
+            "buffers": {},
+            "device_map": {},
+            "weight_tying": self._extract_weight_tying_info(model),
+        }
+
+        for name, param in model.named_parameters():
+            filename = f"param__{_sanitize_param_name(name)}.pt"
+            file_path = snapshot_dir / filename
+            torch.save(param.detach().cpu(), file_path)
+            manifest["params"][name] = filename
+            manifest["device_map"][name] = str(param.device)
+
+        for name, buffer in model.named_buffers():
+            filename = f"buffer__{_sanitize_param_name(name)}.pt"
+            file_path = snapshot_dir / filename
+            torch.save(buffer.detach().cpu(), file_path)
+            manifest["buffers"][name] = filename
+            manifest["device_map"][f"buffer::{name}"] = str(buffer.device)
+
+        manifest_path = snapshot_dir / "manifest.json"
+        manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8")
+        return str(snapshot_dir)
+
+    def restore_chunked(self, model: torch.nn.Module, snapshot_path: str) -> None:
+        """
+        Restore a chunked snapshot produced by `snapshot_chunked`.
+
+        Args:
+            model: Model to restore in-place.
+            snapshot_path: Directory path created by `snapshot_chunked`.
+        """
+
+        snapshot_dir = Path(snapshot_path)
+        manifest_path = snapshot_dir / "manifest.json"
+        if not manifest_path.exists():
+            raise FileNotFoundError(f"Missing manifest for snapshot at {snapshot_path}")
+
+        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
+        param_map = dict(model.named_parameters())
+        buffer_map = dict(model.named_buffers())
+
+        device_map = manifest.get("device_map", {})
+
+        for name, filename in manifest.get("params", {}).items():
+            file_path = snapshot_dir / filename
+            if name not in param_map or not file_path.exists():
+                continue
+            target = param_map[name]
+            target_device = torch.device(device_map.get(name, str(target.device)))
+            tensor = torch.load(file_path, map_location="cpu")
+            with torch.no_grad():
+                target.copy_(tensor.to(target_device))
+
+        for name, filename in manifest.get("buffers", {}).items():
+            file_path = snapshot_dir / filename
+            if name not in buffer_map or not file_path.exists():
+                continue
+            target = buffer_map[name]
+            key = f"buffer::{name}"
+            target_device = torch.device(device_map.get(key, str(target.device)))
+            tensor = torch.load(file_path, map_location="cpu")
+            target.copy_(tensor.to(target_device))
+
+        original_tying = manifest.get("weight_tying", {})
+        if isinstance(original_tying, dict) and original_tying:
+            current_tying = self._extract_weight_tying_info(model)
+            for tied_param, source_param in original_tying.items():
+                if current_tying.get(tied_param) != source_param:
+                    self._restore_weight_tying(model, tied_param, source_param)
+
+    # ------------------------------------------------------------------
+    # Weight-tying hooks (overridden by concrete adapters)
+    # ------------------------------------------------------------------
+    def _extract_weight_tying_info(self, model: torch.nn.Module) -> dict[str, str]:
+        """Return mapping of tied parameter names to source parameter names."""
+
+        tying: dict[str, str] = {}
+        param_names = set(dict(model.named_parameters()).keys())
+
+        if "lm_head.weight" in param_names and "transformer.wte.weight" in param_names:
+            tying["lm_head.weight"] = "transformer.wte.weight"
+
+        decoder_name = "cls.predictions.decoder.weight"
+        if decoder_name in param_names:
+            for candidate in (
+                "bert.embeddings.word_embeddings.weight",
+                "embeddings.word_embeddings.weight",
+            ):
+                if candidate in param_names:
+                    tying[decoder_name] = candidate
+                    break
+
+        return tying
+
+    def _restore_weight_tying(
+        self, model: torch.nn.Module, tied_param: str, source_param: str
+    ) -> None:
+        """Restore a weight-tying relationship (no-op by default)."""
+        model_params = dict(model.named_parameters())
+        tied = model_params.get(tied_param)
+        source = model_params.get(source_param)
+        if tied is None or source is None:
+            return
+        with torch.no_grad():
+            tied.copy_(source)
+
+    def validate_weight_tying(self, model: torch.nn.Module) -> None:
+        """Raise if a known weight-tying relationship has been broken."""
+        tying = self._extract_weight_tying_info(model)
+        if not tying:
+            return
+
+        model_params = dict(model.named_parameters())
+        for tied_param, source_param in tying.items():
+            tied = model_params.get(tied_param)
+            source = model_params.get(source_param)
+            if tied is None:
+                tied = _resolve_named_parameter(model, tied_param)
+            if source is None:
+                source = _resolve_named_parameter(model, source_param)
+            if tied is None or source is None:
+                from invarlock.core.exceptions import AdapterError
+
+                raise AdapterError(
+                    code="E202",
+                    message="ADAPTER-STRUCTURE-INVALID: missing tied/source parameter",
+                    details={
+                        "tied_param": tied_param,
+                        "source_param": source_param,
+                    },
+                )
+            if not torch.allclose(tied, source):
+                from invarlock.core.exceptions import AdapterError
+
+                raise AdapterError(
+                    code="E202",
+                    message="ADAPTER-STRUCTURE-INVALID: weight-tying invariant violated",
+                    details={
+                        "tied_param": tied_param,
+                        "source_param": source_param,
+                    },
+                )
+
+    # ------------------------------------------------------------------
+    # Helper utilities
+    # ------------------------------------------------------------------
+    def _serialize_config(self, config: Any) -> dict[str, Any]:
+        """Serialize HuggingFace config fields into simple Python types."""
+
+        result: dict[str, Any] = {}
+        for key in dir(config):
+            if key.startswith("_"):
+                continue
+
+            try:
+                value = getattr(config, key)
+            except AttributeError:
+                continue
+
+            if callable(value):
+                continue
+
+            if value is None or isinstance(value, SCALAR_TYPES):
+                result[key] = value
+            elif isinstance(value, list | dict):
+                result[key] = value
+
+        return result

invarlock/adapters/hf_onnx.py

@@ -0,0 +1,112 @@
+"""
+HuggingFace Optimum ONNX Runtime Adapter
+=======================================
+
+Minimal adapter to load CPU-friendly ONNX Runtime models via Optimum for
+causal language modeling. This enables certification of pre-quantized int8
+models published for ONNX/Optimum, when used with `edit: noop`.
+
+Notes
+- This adapter targets inference (perplexity/log-loss) and does not expose
+  parameter/module traversal; guards that require PyTorch module access may
+  be inapplicable for ONNX models.
+- Optional dependency: `optimum[onnxruntime]` and `onnxruntime`.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from invarlock.core.api import ModelAdapter
+from invarlock.core.error_utils import wrap_errors
+from invarlock.core.exceptions import DependencyError, ModelLoadError
+
+
+class HF_ORT_CausalLM_Adapter(ModelAdapter):
+    """Optimum/ONNXRuntime causal LM adapter.
+
+    Provides a lightweight bridge that loads an ORTModelForCausalLM and
+    presents enough interface compatibility for evaluation paths that operate
+    on logits (perplexity). Snapshot/restore are not supported and will fall
+    back to reload in the CLI runner.
+    """
+
+    name = "hf_onnx"
+
+    # --- Loading ---------------------------------------------------------
+    def load_model(self, model_id: str, device: str = "cpu", **kwargs: Any):  # type: ignore[override]
+        with wrap_errors(
+            DependencyError,
+            "E203",
+            "DEPENDENCY-MISSING: optimum/onnxruntime",
+            lambda e: {"dependency": "optimum[onnxruntime]"},
+        ):
+            from optimum.onnxruntime import ORTModelForCausalLM  # type: ignore
+
+        # Prefer CPUExecutionProvider by default; callers may override via kwargs
+        providers = kwargs.pop("providers", None)
+        if providers is None:
+            providers = ["CPUExecutionProvider"]
+
+        # Trust remote code where necessary; users can set to False via kwargs
+        trust_remote_code = kwargs.pop("trust_remote_code", True)
+
+        # Some repos use non-default file names; accept overrides but default to standard
+        file_name = kwargs.pop("file_name", None)
+
+        with wrap_errors(
+            ModelLoadError,
+            "E201",
+            "MODEL-LOAD-FAILED: optimum ORTModelForCausalLM",
+            lambda e: {"model_id": model_id},
+        ):
+            model = ORTModelForCausalLM.from_pretrained(
+                model_id,
+                file_name=file_name,
+                provider="CPUExecutionProvider"
+                if providers == ["CPUExecutionProvider"]
+                else None,
+                providers=providers,
+                trust_remote_code=trust_remote_code,
+                **kwargs,
+            )
+        # ORT models manage device internally; return as-is
+        return model
+
+    # --- Capability probes ----------------------------------------------
+    def can_handle(self, model: Any) -> bool:  # type: ignore[override]
+        cls_name = model.__class__.__name__
+        mod_name = getattr(model.__class__, "__module__", "")
+        if "optimum.onnxruntime" in mod_name and "ORTModel" in cls_name:
+            return True
+        # Heuristic: object exposes generate and is directly callable
+        return hasattr(model, "generate") and callable(model)
+
+    # --- Structure description (best-effort) -----------------------------
+    def describe(self, model: Any) -> dict[str, Any]:  # type: ignore[override]
+        # Attempt to read from HF config when present
+        cfg = getattr(model, "config", None)
+        n_layer = int(
+            getattr(cfg, "n_layer", getattr(cfg, "num_hidden_layers", 0)) or 0
+        )
+        n_head = int(
+            getattr(cfg, "n_head", getattr(cfg, "num_attention_heads", 0)) or 0
+        )
+        heads = [n_head] * n_layer if n_layer and n_head else []
+        return {
+            "n_layer": n_layer,
+            "heads_per_layer": heads,
+            "mlp_dims": [],
+            "tying": {},
+            "backend": "onnxruntime",
+        }
+
+    # --- Snapshot/restore (unsupported; runner falls back to reload) -----
+    def snapshot(self, model: Any) -> bytes:  # type: ignore[override]
+        raise NotImplementedError("snapshot not supported for ONNXRuntime models")
+
+    def restore(self, model: Any, blob: bytes) -> None:  # type: ignore[override]
+        raise NotImplementedError("restore not supported for ONNXRuntime models")
+
+
+__all__ = ["HF_ORT_CausalLM_Adapter"]

invarlock/adapters/hf_t5.py

@@ -0,0 +1,137 @@
+"""
+HuggingFace T5 Model Adapter
+============================
+
+ModelAdapter implementation for HuggingFace T5 encoder-decoder models.
+
+Loads AutoModelForSeq2SeqLM (e.g., t5-small/base/large) and exposes a minimal
+describe() sufficient for guard policies and reporting.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from invarlock.core.api import ModelAdapter
+from invarlock.core.error_utils import wrap_errors
+from invarlock.core.exceptions import AdapterError, DependencyError, ModelLoadError
+
+from .hf_mixin import HFAdapterMixin
+
+TensorType = torch.Tensor
+ModuleType = nn.Module
+
+
+class HF_T5_Adapter(HFAdapterMixin, ModelAdapter):
+    """HuggingFace T5 adapter using AutoModelForSeq2SeqLM."""
+
+    name = "hf_t5"
+
+    def load_model(self, model_id: str, device: str = "auto") -> ModuleType | Any:  # type: ignore[override]
+        with wrap_errors(
+            DependencyError,
+            "E203",
+            "DEPENDENCY-MISSING: transformers",
+            lambda e: {"dependency": "transformers"},
+        ):
+            from transformers import AutoModelForSeq2SeqLM  # type: ignore
+
+        with wrap_errors(
+            ModelLoadError,
+            "E201",
+            "MODEL-LOAD-FAILED: transformers AutoModelForSeq2SeqLM",
+            lambda e: {"model_id": model_id},
+        ):
+            model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
+        return model.to(self._resolve_device(device))
+
+    def can_handle(self, model: ModuleType | Any) -> bool:  # type: ignore[override]
+        cfg = getattr(model, "config", None)
+        if cfg is None:
+            return False
+        try:
+            mt = str(getattr(cfg, "model_type", "")).lower()
+        except Exception:
+            mt = ""
+        if mt == "t5":
+            return True
+        # Heuristic: encoder-decoder with lm_head and shared embedding
+        return bool(
+            getattr(cfg, "is_encoder_decoder", False)
+            and hasattr(model, "lm_head")
+            and hasattr(model, "shared")
+        )
+
+    def describe(self, model: ModuleType | Any) -> dict[str, Any]:  # type: ignore[override]
+        cfg = getattr(model, "config", None)
+        if cfg is None:
+            raise AdapterError(
+                code="E202",
+                message="ADAPTER-STRUCTURE-INVALID: missing HuggingFace config on model",
+                details={"model_class": model.__class__.__name__},
+            )
+
+        # Extract key dimensions with safe fallbacks
+        enc_layers = int(
+            getattr(cfg, "num_layers", getattr(cfg, "num_encoder_layers", 0)) or 0
+        )
+        dec_layers = int(getattr(cfg, "num_decoder_layers", enc_layers) or 0)
+        n_layer = int(enc_layers + dec_layers or max(enc_layers, dec_layers))
+        n_heads = int(
+            getattr(cfg, "num_heads", getattr(cfg, "num_attention_heads", 0)) or 0
+        )
+        d_model = int(getattr(cfg, "d_model", getattr(cfg, "hidden_size", 0)) or 0)
+        d_ff = int(getattr(cfg, "d_ff", (d_model * 4) if d_model else 0) or 0)
+        vocab_size = int(getattr(cfg, "vocab_size", 0) or 0)
+
+        try:
+            device = next(model.parameters()).device
+        except Exception:
+            device = torch.device("cpu")
+
+        heads_per_layer = [n_heads] * max(1, n_layer)
+        mlp_dims = [d_ff] * max(1, n_layer)
+
+        # Weight tying: T5 ties lm_head ↔ shared embedding
+        tying_map: dict[str, str] = {}
+        try:
+            if hasattr(model, "lm_head") and hasattr(model, "shared"):
+                if getattr(model.lm_head, "weight", None) is getattr(
+                    model.shared, "weight", object()
+                ):
+                    tying_map["lm_head.weight"] = "shared.weight"
+        except Exception:
+            pass
+
+        total_params = 0
+        try:
+            total_params = sum(p.numel() for p in model.parameters())
+        except Exception:
+            total_params = 0
+
+        return {
+            "n_layer": int(max(1, n_layer)),
+            "heads_per_layer": heads_per_layer,
+            "mlp_dims": mlp_dims,
+            "tying": tying_map,
+            "model_type": "t5",
+            "model_class": model.__class__.__name__,
+            "n_heads": n_heads,
+            "hidden_size": d_model,
+            "vocab_size": vocab_size,
+            "total_params": total_params,
+            "device": str(device),
+        }
+
+    # snapshot/restore provided by HFAdapterMixin
+    def snapshot(self, model: ModuleType) -> bytes:  # type: ignore[override]
+        return super().snapshot(model)
+
+    def restore(self, model: ModuleType, blob: bytes) -> None:  # type: ignore[override]
+        return super().restore(model, blob)
+
+
+__all__ = ["HF_T5_Adapter"]

invarlock/adapters/py.typed

@@ -0,0 +1 @@
+# Marker file for PEP 561 type hints

invarlock/assurance/__init__.py

@@ -0,0 +1,43 @@
+"""Assurance namespace (`invarlock.assurance`).
+
+This namespace groups safety-certificate related surfaces. For now it forwards
+to `invarlock.eval` and guard modules; future work may move implementations here.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from invarlock.reporting.report_types import RunReport
+
+try:  # pragma: no cover - shim to reporting modules
+    from invarlock.reporting.certificate import (
+        CERTIFICATE_SCHEMA_VERSION,
+        make_certificate,
+        validate_certificate,
+    )
+
+    # Prefer direct import from render for rendering APIs
+    from invarlock.reporting.render import render_certificate_markdown
+except Exception:  # pragma: no cover - provide soft stubs
+    CERTIFICATE_SCHEMA_VERSION = "v1"
+
+    def make_certificate(
+        report: RunReport,
+        baseline: RunReport | dict[str, Any],
+    ) -> dict[str, Any]:
+        raise ImportError("invarlock.reporting.certificate not available")
+
+    def render_certificate_markdown(certificate: dict[str, Any]) -> str:
+        raise ImportError("invarlock.reporting.certificate not available")
+
+    def validate_certificate(certificate: dict[str, Any]) -> bool:
+        raise ImportError("invarlock.reporting.certificate not available")
+
+
+__all__ = [
+    "CERTIFICATE_SCHEMA_VERSION",
+    "make_certificate",
+    "render_certificate_markdown",
+    "validate_certificate",
+]

invarlock/cli/__init__.py

@@ -0,0 +1,8 @@
+"""CLI namespace wrapper for unified import path (`invarlock.cli`)."""
+
+from __future__ import annotations
+
+# Re-export common entry points for convenience
+from .app import app
+
+__all__ = ["app"]