invarlock 0.3.5__py3-none-any.whl → 0.3.7__py3-none-any.whl

invarlock/adapters/hf_causal.py

@@ -0,0 +1,418 @@
+"""
+HuggingFace causal LM adapter (decoder-only).
+=============================================
+
+Role-based adapter for HuggingFace decoder-only causal language models.
+
+This adapter intentionally avoids model-family naming. It selects a structural
+spec at runtime (dense FFN vs MoE vs GPT-2-like blocks) and exposes a stable
+`describe()` contract for InvarLock gates and reporting.
+"""
+
+from __future__ import annotations
+
+import os
+from types import SimpleNamespace
+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
+
+LIGHT_IMPORT = os.getenv("INVARLOCK_LIGHT_IMPORT", "").strip().lower() in {
+    "1",
+    "true",
+    "yes",
+}
+
+
+def _first_item(seq: Any) -> Any | None:
+    try:
+        if hasattr(seq, "__len__") and len(seq) > 0:  # type: ignore[arg-type]
+            return seq[0]  # type: ignore[index]
+    except Exception:
+        pass
+    try:
+        return next(iter(seq))
+    except Exception:
+        return None
+
+
+def _has_set_attr(obj: Any, name: str) -> bool:
+    d = getattr(obj, "__dict__", None)
+    if isinstance(d, dict) and name in d:
+        return True
+    if isinstance(obj, nn.Module):
+        if hasattr(obj, "_modules") and name in obj._modules:
+            return True
+        if hasattr(obj, "_parameters") and name in obj._parameters:
+            return True
+        if hasattr(obj, "_buffers") and name in obj._buffers:
+            return True
+    return False
+
+
+class _CausalSpec:
+    spec_name = "base"
+
+    def matches(self, model: Any, base: Any, layers: Any) -> bool:
+        raise NotImplementedError
+
+    def infer_mlp_dim(self, layer: Any, config: Any, hidden_size: int) -> int:
+        raise NotImplementedError
+
+    def layer_modules(self, model: Any, layer: Any) -> dict[str, Any]:
+        raise NotImplementedError
+
+    def tying_map(self, model: Any, base: Any) -> dict[str, str]:
+        return {}
+
+
+class _DenseDecoderSpec(_CausalSpec):
+    spec_name = "dense_decoder"
+
+    def matches(self, model: Any, base: Any, layers: Any) -> bool:
+        layer = _first_item(layers)
+        if layer is None:
+            return False
+        has_attn = (
+            hasattr(layer, "self_attn")
+            and _has_set_attr(layer.self_attn, "q_proj")
+            and _has_set_attr(layer.self_attn, "k_proj")
+            and _has_set_attr(layer.self_attn, "v_proj")
+            and _has_set_attr(layer.self_attn, "o_proj")
+        )
+        has_mlp = (
+            hasattr(layer, "mlp")
+            and _has_set_attr(layer.mlp, "gate_proj")
+            and _has_set_attr(layer.mlp, "up_proj")
+            and _has_set_attr(layer.mlp, "down_proj")
+        )
+        has_norms = _has_set_attr(layer, "input_layernorm") and _has_set_attr(
+            layer, "post_attention_layernorm"
+        )
+        return bool(has_attn and has_mlp and has_norms)
+
+    def infer_mlp_dim(self, layer: Any, config: Any, hidden_size: int) -> int:
+        mlp_dim = int(getattr(config, "intermediate_size", hidden_size * 4) or 0)
+        try:
+            gate_proj = getattr(getattr(layer, "mlp", None), "gate_proj", None)
+            if gate_proj is not None and hasattr(gate_proj, "weight"):
+                mlp_dim = int(gate_proj.weight.shape[0])
+        except Exception:
+            pass
+        return int(mlp_dim)
+
+    def layer_modules(self, model: Any, layer: Any) -> dict[str, Any]:
+        mlp = layer.mlp
+        return {
+            "self_attn.q_proj": layer.self_attn.q_proj,
+            "self_attn.k_proj": layer.self_attn.k_proj,
+            "self_attn.v_proj": layer.self_attn.v_proj,
+            "self_attn.o_proj": layer.self_attn.o_proj,
+            "input_layernorm": layer.input_layernorm,
+            "post_attention_layernorm": layer.post_attention_layernorm,
+            "mlp.gate_proj": mlp.gate_proj,
+            "mlp.up_proj": mlp.up_proj,
+            "mlp.down_proj": mlp.down_proj,
+        }
+
+    def tying_map(self, model: Any, base: Any) -> dict[str, str]:
+        tying: dict[str, str] = {}
+        try:
+            if hasattr(model, "lm_head") and hasattr(base, "embed_tokens"):
+                if model.lm_head.weight is base.embed_tokens.weight:
+                    tying["lm_head.weight"] = "model.embed_tokens.weight"
+        except Exception:
+            pass
+        return tying
+
+
+class _MoEDecoderSpec(_CausalSpec):
+    spec_name = "moe_decoder"
+
+    def matches(self, model: Any, base: Any, layers: Any) -> bool:
+        layer = _first_item(layers)
+        if layer is None:
+            return False
+        has_attn = (
+            hasattr(layer, "self_attn")
+            and _has_set_attr(layer.self_attn, "q_proj")
+            and _has_set_attr(layer.self_attn, "k_proj")
+            and _has_set_attr(layer.self_attn, "v_proj")
+            and _has_set_attr(layer.self_attn, "o_proj")
+        )
+        moe = getattr(layer, "block_sparse_moe", None)
+        experts = getattr(moe, "experts", None) if moe is not None else None
+        expert0 = _first_item(experts) if experts is not None else None
+        has_moe = bool(
+            expert0 is not None
+            and _has_set_attr(expert0, "w1")
+            and _has_set_attr(expert0, "w2")
+        )
+        has_norms = _has_set_attr(layer, "input_layernorm") and _has_set_attr(
+            layer, "post_attention_layernorm"
+        )
+        return bool(has_attn and has_moe and has_norms)
+
+    def infer_mlp_dim(self, layer: Any, config: Any, hidden_size: int) -> int:
+        mlp_dim = int(getattr(config, "intermediate_size", hidden_size * 4) or 0)
+        try:
+            moe = getattr(layer, "block_sparse_moe", None)
+            experts = getattr(moe, "experts", None) if moe is not None else None
+            expert0 = _first_item(experts) if experts is not None else None
+            if expert0 is not None:
+                w1 = getattr(expert0, "w1", None)
+                if w1 is not None and hasattr(w1, "weight"):
+                    mlp_dim = int(w1.weight.shape[0])
+        except Exception:
+            pass
+        return int(mlp_dim)
+
+    def layer_modules(self, model: Any, layer: Any) -> dict[str, Any]:
+        moe = layer.block_sparse_moe
+        expert0 = _first_item(moe.experts)
+        if expert0 is None:
+            raise AdapterError(
+                code="E202",
+                message="ADAPTER-STRUCTURE-INVALID: MoE layer missing experts",
+                details={"layer_class": layer.__class__.__name__},
+            )
+        return {
+            "self_attn.q_proj": layer.self_attn.q_proj,
+            "self_attn.k_proj": layer.self_attn.k_proj,
+            "self_attn.v_proj": layer.self_attn.v_proj,
+            "self_attn.o_proj": layer.self_attn.o_proj,
+            "input_layernorm": layer.input_layernorm,
+            "post_attention_layernorm": layer.post_attention_layernorm,
+            # Best-effort mapping to dense naming used elsewhere in the stack.
+            "mlp.gate_proj": expert0.w1,
+            "mlp.up_proj": getattr(expert0, "w3", expert0.w1),
+            "mlp.down_proj": expert0.w2,
+        }
+
+    def tying_map(self, model: Any, base: Any) -> dict[str, str]:
+        return _DenseDecoderSpec().tying_map(model, base)
+
+
+class _GPT2LikeDecoderSpec(_CausalSpec):
+    spec_name = "gpt2_like"
+
+    def matches(self, model: Any, base: Any, layers: Any) -> bool:
+        layer = _first_item(layers)
+        if layer is None:
+            return False
+        return bool(
+            hasattr(layer, "attn")
+            and hasattr(layer.attn, "c_proj")
+            and hasattr(layer, "mlp")
+            and hasattr(layer.mlp, "c_proj")
+        )
+
+    def infer_mlp_dim(self, layer: Any, config: Any, hidden_size: int) -> int:
+        try:
+            c_fc = getattr(getattr(layer, "mlp", None), "c_fc", None)
+            if c_fc is not None and hasattr(c_fc, "weight"):
+                # HF GPT-style uses Conv1D where nf is out_features.
+                if hasattr(c_fc, "nf"):
+                    return int(c_fc.nf)
+                return int(c_fc.weight.shape[0])
+        except Exception:
+            pass
+        return int(getattr(config, "n_inner", hidden_size * 4) or 0)
+
+    def layer_modules(self, model: Any, layer: Any) -> dict[str, Any]:
+        return {
+            "attn.c_attn": layer.attn.c_attn,
+            "attn.c_proj": layer.attn.c_proj,
+            "mlp.c_fc": layer.mlp.c_fc,
+            "mlp.c_proj": layer.mlp.c_proj,
+            "ln_1": layer.ln_1,
+            "ln_2": layer.ln_2,
+        }
+
+    def tying_map(self, model: Any, base: Any) -> dict[str, str]:
+        tying: dict[str, str] = {}
+        try:
+            if hasattr(model, "lm_head") and hasattr(base, "wte"):
+                if model.lm_head.weight is base.wte.weight:
+                    tying["lm_head.weight"] = "transformer.wte.weight"
+        except Exception:
+            pass
+        return tying
+
+
+_SPECS: list[_CausalSpec] = [
+    _MoEDecoderSpec(),
+    _DenseDecoderSpec(),
+    _GPT2LikeDecoderSpec(),
+]
+
+
+class HF_Causal_Adapter(HFAdapterMixin, ModelAdapter):
+    """Spec-driven adapter for decoder-only causal LMs."""
+
+    name = "hf_causal"
+
+    def load_model(
+        self, model_id: str, device: str = "auto", **kwargs: Any
+    ) -> ModuleType | Any:
+        try:
+            with wrap_errors(
+                DependencyError,
+                "E203",
+                "DEPENDENCY-MISSING: transformers",
+                lambda e: {"dependency": "transformers"},
+            ):
+                from transformers import AutoModelForCausalLM  # type: ignore
+
+            with wrap_errors(
+                ModelLoadError,
+                "E201",
+                "MODEL-LOAD-FAILED: transformers AutoModelForCausalLM",
+                lambda e: {"model_id": model_id},
+            ):
+                model = AutoModelForCausalLM.from_pretrained(model_id, **kwargs)
+
+            return self._safe_to_device(model, device)
+        except DependencyError:
+            if LIGHT_IMPORT:
+                stub = SimpleNamespace(name="hf_causal_stub")
+                stub.to = lambda *_a, **_k: stub  # type: ignore[attr-defined]
+                return stub
+            raise
+
+    def _unwrap(self, model: Any) -> tuple[Any, Any, Any]:
+        config = getattr(model, "config", None)
+        if hasattr(model, "model") and hasattr(model.model, "layers"):
+            return model.model, model.model.layers, config
+        if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
+            return model.transformer, model.transformer.h, config
+        if hasattr(model, "layers"):
+            return model, model.layers, config
+        if hasattr(model, "h"):
+            return model, model.h, config
+        raise AdapterError(
+            code="E202",
+            message="ADAPTER-STRUCTURE-INVALID: unrecognized HF causal LM structure",
+            details={"model_class": model.__class__.__name__},
+        )
+
+    def _select_spec(self, model: Any, base: Any, layers: Any) -> _CausalSpec:
+        for spec in _SPECS:
+            try:
+                if spec.matches(model, base, layers):
+                    return spec
+            except Exception:
+                continue
+        return _DenseDecoderSpec()
+
+    def can_handle(self, model: ModuleType | Any) -> bool:
+        try:
+            base, layers, _cfg = self._unwrap(model)
+        except Exception:
+            return False
+        return any(spec.matches(model, base, layers) for spec in _SPECS)
+
+    def describe(self, model: ModuleType | Any) -> dict[str, Any]:
+        base, layers, config = self._unwrap(model)
+        if config is None:
+            raise AdapterError(
+                code="E202",
+                message="ADAPTER-STRUCTURE-INVALID: missing HuggingFace config on model",
+                details={"model_class": model.__class__.__name__},
+            )
+
+        try:
+            n_layers = len(layers)
+        except Exception:
+            n_layers = sum(1 for _ in iter(layers))
+
+        def _coerce_int(value: Any) -> int | None:
+            try:
+                if isinstance(value, bool):
+                    return None
+                if isinstance(value, int):
+                    return int(value)
+                if isinstance(value, float):
+                    return int(value)
+                if isinstance(value, str):
+                    stripped = value.strip()
+                    if stripped and stripped.isdigit():
+                        return int(stripped)
+            except Exception:
+                return None
+            return None
+
+        n_heads = _coerce_int(getattr(config, "num_attention_heads", None))
+        if n_heads is None:
+            n_heads = _coerce_int(getattr(config, "n_head", None))
+
+        hidden_size = _coerce_int(getattr(config, "hidden_size", None))
+        if hidden_size is None:
+            hidden_size = _coerce_int(getattr(config, "n_embd", None))
+
+        vocab_size = _coerce_int(getattr(config, "vocab_size", None))
+
+        if n_heads is None or hidden_size is None:
+            raise AdapterError(
+                code="E202",
+                message="ADAPTER-STRUCTURE-INVALID: missing head/hidden size metadata",
+                details={"model_class": model.__class__.__name__},
+            )
+
+        spec = self._select_spec(model, base, layers)
+
+        heads_per_layer = [int(n_heads)] * int(n_layers)
+        mlp_dims: list[int] = []
+        for idx in range(int(n_layers)):
+            layer = layers[idx]
+            mlp_dims.append(spec.infer_mlp_dim(layer, config, int(hidden_size)))
+
+        tying = spec.tying_map(model, base)
+
+        total_params = 0
+        try:
+            total_params = sum(p.numel() for p in model.parameters())
+        except Exception:
+            total_params = 0
+
+        try:
+            device = next(model.parameters()).device
+        except Exception:
+            device = torch.device("cpu")
+
+        return {
+            "n_layer": int(n_layers),
+            "heads_per_layer": heads_per_layer,
+            "mlp_dims": mlp_dims,
+            "tying": tying,
+            "model_type": str(getattr(config, "model_type", "") or "causal"),
+            "model_class": model.__class__.__name__,
+            "hf_model_type": str(getattr(config, "model_type", "") or ""),
+            "hf_config_class": config.__class__.__name__
+            if hasattr(config, "__class__")
+            else "unknown",
+            "n_heads": int(n_heads),
+            "hidden_size": int(hidden_size),
+            "vocab_size": int(vocab_size) if vocab_size is not None else None,
+            "total_params": int(total_params),
+            "device": str(device),
+            "spec": spec.spec_name,
+        }
+
+    def get_layer_modules(
+        self, model: ModuleType | Any, layer_idx: int
+    ) -> dict[str, Any]:
+        base, layers, _cfg = self._unwrap(model)
+        spec = self._select_spec(model, base, layers)
+        layer = layers[layer_idx]
+        return spec.layer_modules(model, layer)

invarlock/adapters/{hf_onnx.py → hf_causal_onnx.py}

@@ -22,7 +22,7 @@ from invarlock.core.error_utils import wrap_errors
 from invarlock.core.exceptions import DependencyError, ModelLoadError
 
 
-class HF_ORT_CausalLM_Adapter(ModelAdapter):
+class HF_Causal_ONNX_Adapter(ModelAdapter):
     """Optimum/ONNXRuntime causal LM adapter.
 
     Provides a lightweight bridge that loads an ORTModelForCausalLM and
@@ -31,7 +31,7 @@ class HF_ORT_CausalLM_Adapter(ModelAdapter):
     back to reload in the CLI runner.
     """
 
-    name = "hf_onnx"
+    name = "hf_causal_onnx"
 
     # --- Loading ---------------------------------------------------------
     def load_model(self, model_id: str, device: str = "cpu", **kwargs: Any):  # type: ignore[override]
@@ -109,4 +109,4 @@ class HF_ORT_CausalLM_Adapter(ModelAdapter):
         raise NotImplementedError("restore not supported for ONNXRuntime models")
 
 
-__all__ = ["HF_ORT_CausalLM_Adapter"]
+__all__ = ["HF_Causal_ONNX_Adapter"]

invarlock/adapters/hf_mixin.py

@@ -490,18 +490,39 @@ class HFAdapterMixin:
         """Return mapping of tied parameter names to source parameter names."""
 
         tying: dict[str, str] = {}
-        param_names = set(dict(model.named_parameters()).keys())
+        try:
+            named = model.named_parameters(remove_duplicate=False)  # type: ignore[call-arg]
+        except TypeError:  # pragma: no cover - torch version dependent
+            named = model.named_parameters()
+        params = dict(named)
+
+        def _is_tied(name_a: str, name_b: str) -> bool:
+            a = params.get(name_a)
+            b = params.get(name_b)
+            if a is None or b is None:
+                return False
+            try:
+                if a is b:
+                    return True
+                if hasattr(a, "data_ptr") and hasattr(b, "data_ptr"):
+                    return int(a.data_ptr()) == int(b.data_ptr())
+            except Exception:
+                return False
+            return False
 
-        if "lm_head.weight" in param_names and "transformer.wte.weight" in param_names:
+        if _is_tied("lm_head.weight", "transformer.wte.weight"):
             tying["lm_head.weight"] = "transformer.wte.weight"
 
+        if _is_tied("lm_head.weight", "model.embed_tokens.weight"):
+            tying["lm_head.weight"] = "model.embed_tokens.weight"
+
         decoder_name = "cls.predictions.decoder.weight"
-        if decoder_name in param_names:
+        if decoder_name in params:
             for candidate in (
                 "bert.embeddings.word_embeddings.weight",
                 "embeddings.word_embeddings.weight",
             ):
-                if candidate in param_names:
+                if _is_tied(decoder_name, candidate):
                     tying[decoder_name] = candidate
                     break
 

invarlock/adapters/{hf_bert.py → hf_mlm.py}

@@ -1,15 +1,8 @@
 """
-HuggingFace BERT Model Adapter
+HuggingFace masked LM adapter.
 ==============================
 
-ModelAdapter implementation for HuggingFace BERT architecture models.
-
-This adapter provides BERT-specific integration including:
-- Support for BERT, RoBERTa, DistilBERT, and other BERT variants
-- Proper handling of bidirectional attention layers
-- Support for classification heads and pooling layers
-- Token type embeddings and position embeddings handling
-- Proper device-aware state serialization
+ModelAdapter implementation for HuggingFace masked language models.
 """
 
 from typing import Any
@@ -27,7 +20,7 @@ TensorType = torch.Tensor
 ModuleType = nn.Module
 
 
-class HF_BERT_Adapter(HFAdapterMixin, ModelAdapter):
+class HF_MLM_Adapter(HFAdapterMixin, ModelAdapter):
     """
     HuggingFace-specific ModelAdapter implementation for BERT models.
 
@@ -39,7 +32,7 @@ class HF_BERT_Adapter(HFAdapterMixin, ModelAdapter):
     - Device-aware state serialization
     """
 
-    name = "hf_bert"
+    name = "hf_mlm"
 
     def load_model(
         self, model_id: str, device: str = "auto", **kwargs: Any

invarlock/adapters/{hf_t5.py → hf_seq2seq.py}

@@ -1,11 +1,11 @@
 """
-HuggingFace T5 Model Adapter
-============================
+HuggingFace encoder-decoder adapter.
+===================================
 
-ModelAdapter implementation for HuggingFace T5 encoder-decoder models.
+ModelAdapter implementation for HuggingFace encoder-decoder (seq2seq) models.
 
-Loads AutoModelForSeq2SeqLM (e.g., t5-small/base/large) and exposes a minimal
-describe() sufficient for guard policies and reporting.
+Loads AutoModelForSeq2SeqLM and exposes a minimal describe() sufficient for
+guard policies and reporting.
 """
 
 from __future__ import annotations
@@ -25,10 +25,10 @@ TensorType = torch.Tensor
 ModuleType = nn.Module
 
 
-class HF_T5_Adapter(HFAdapterMixin, ModelAdapter):
-    """HuggingFace T5 adapter using AutoModelForSeq2SeqLM."""
+class HF_Seq2Seq_Adapter(HFAdapterMixin, ModelAdapter):
+    """HuggingFace encoder-decoder adapter using AutoModelForSeq2SeqLM."""
 
-    name = "hf_t5"
+    name = "hf_seq2seq"
 
     def load_model(  # type: ignore[override]
         self, model_id: str, device: str = "auto", **kwargs: Any
@@ -136,4 +136,4 @@ class HF_T5_Adapter(HFAdapterMixin, ModelAdapter):
         return super().restore(model, blob)
 
 
-__all__ = ["HF_T5_Adapter"]
+__all__ = ["HF_Seq2Seq_Adapter"]

invarlock/calibration/spectral_null.py

@@ -148,7 +148,7 @@ def _selected_families_for_alpha(
 
 
 def summarize_null_sweep_reports(
-    reports: list[dict[str, Any]],
+    reports: list[object],
     *,
     tier: str,
     safety_margin: float = 0.05,
@@ -186,20 +186,25 @@ def summarize_null_sweep_reports(
         mt = _extract_multiple_testing(metrics)
         if mt:
             mt_method = str(mt.get("method", mt_method))
-            if mt.get("alpha") is not None:
-                mt_alpha = float(mt.get("alpha"))
-            if mt.get("m") is not None:
-                mt_m = int(mt.get("m"))
+            alpha_value = mt.get("alpha")
+            if alpha_value is not None:
+                try:
+                    mt_alpha = float(alpha_value)
+                except Exception:
+                    pass
+            m_value = mt.get("m")
+            if m_value is not None:
+                try:
+                    mt_m = int(m_value)
+                except Exception:
+                    pass
 
         fam_z = _extract_family_max_z(metrics)
         for fam, z in fam_z.items():
             family_max_z[fam] = max(family_max_z[fam], float(z))
 
-        selection = (
-            metrics.get("multiple_testing_selection")
-            if isinstance(metrics.get("multiple_testing_selection"), dict)
-            else {}
-        )
+        raw_selection = metrics.get("multiple_testing_selection")
+        selection = raw_selection if isinstance(raw_selection, dict) else {}
         pvals = selection.get("family_pvalues")
         if not isinstance(pvals, dict):
             pvals = {}

invarlock/calibration/variance_ve.py

@@ -107,8 +107,6 @@ def summarize_ve_sweep_reports(
     evaluated = 0
 
     for report in reports:
-        if not isinstance(report, dict):
-            continue
         g = _extract_guard(report, "variance") or {}
         metrics = g.get("metrics", {}) if isinstance(g.get("metrics"), dict) else {}
         pg = metrics.get("predictive_gate")