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

invarlock/guards/invariants.py

@@ -5,6 +5,7 @@ InvarLock Guards - Invariants
 Invariant checking for model edits to ensure structural integrity.
 """
 
+import hashlib
 from typing import Any
 
 import torch
@@ -33,6 +34,7 @@ class InvariantsGuard(Guard):
         self.on_fail = on_fail
         self.prepared = False
         self.baseline_checks: dict[str, Any] = {}
+        self.last_current_checks: dict[str, Any] = {}
         self.profile_checks: tuple[str, ...] = ()
 
     def prepare(
@@ -102,6 +104,10 @@ class InvariantsGuard(Guard):
             "action": outcome.action,
             "violations": outcome.violations,
             "metrics": outcome.metrics,
+            "details": {
+                "baseline_checks": self.baseline_checks,
+                "current_checks": self.last_current_checks,
+            },
         }
 
     def finalize(self, model: Any) -> GuardOutcome:
@@ -125,6 +131,7 @@ class InvariantsGuard(Guard):
 
         # Check current invariants
         current_checks = self._capture_invariants(model, None)
+        self.last_current_checks = current_checks
         violations: list[dict[str, Any]] = []
         tokenizer_mismatches: list[dict[str, Any]] = []
 
@@ -354,14 +361,14 @@ class InvariantsGuard(Guard):
         except Exception:
             pass
 
-        # LLaMA style (model.embed_tokens <-> lm_head)
+        # Decoder embed_tokens style (model.embed_tokens <-> lm_head)
         try:
-            llama_model = getattr(model, "model", None)
-            embed_tokens = getattr(llama_model, "embed_tokens", None)
+            decoder_model = getattr(model, "model", None)
+            embed_tokens = getattr(decoder_model, "embed_tokens", None)
             embed_weight = getattr(embed_tokens, "weight", None)
-            llama_head_weight = getattr(getattr(model, "lm_head", None), "weight", None)
-            if embed_weight is not None and llama_head_weight is not None:
-                weight_tying_flags["llama"] = _is_tied(embed_weight, llama_head_weight)
+            head_weight = getattr(getattr(model, "lm_head", None), "weight", None)
+            if embed_weight is not None and head_weight is not None:
+                weight_tying_flags["embed_tokens"] = _is_tied(embed_weight, head_weight)
         except Exception:
             pass
 
@@ -376,8 +383,10 @@ class InvariantsGuard(Guard):
             structure_items = []
             for name, module in model.named_modules():
                 structure_items.append(f"{name}:{type(module).__name__}")
-            structure_hash = hash(tuple(structure_items))
-            checks["structure_hash"] = structure_hash
+            canonical = "\n".join(sorted(structure_items))
+            checks["structure_hash"] = hashlib.sha256(
+                canonical.encode("utf-8")
+            ).hexdigest()[:16]
         except Exception:
             checks["structure_hash"] = 0
 
@@ -424,7 +433,7 @@ class InvariantsGuard(Guard):
             return "bert" in model_type or has_cls_decoder
 
         if name in {"rope_rotary_embedding", "rotary_embedding"}:
-            # Detect rotary embeddings used by LLaMA-style models
+            # Detect rotary embeddings used by RoPE-style models
             if hasattr(model, "model") and hasattr(model.model, "layers"):
                 first_layer = model.model.layers[0] if model.model.layers else None
             else:
@@ -443,7 +452,7 @@ class InvariantsGuard(Guard):
             model_type = getattr(config, "model_type", "") if config else ""
             return any(
                 keyword in model_type
-                for keyword in ("gpt", "llama", "mistral", "opt", "phi")
+                for keyword in ("gpt", "mistral", "mixtral", "qwen", "opt", "phi")
             )
 
         return True

invarlock/guards/policies.py

@@ -15,7 +15,7 @@ from typing import Any, Literal
 
 try:  # Python 3.12+
     from typing import NotRequired, TypedDict
-except ImportError:  # Legacy fallback
+except ImportError:  # Python <3.12 fallback
     from typing import NotRequired
 
     from typing_extensions import TypedDict
@@ -40,6 +40,7 @@ SPECTRAL_CONSERVATIVE: SpectralPolicy = {
     "scope": "ffn",  # FFN layers only (safest)
     "correction_enabled": True,
     "max_caps": 3,
+    "max_spectral_norm": None,
     "multiple_testing": {"method": "bonferroni", "alpha": 0.02, "m": 4},
 }
 
@@ -50,6 +51,7 @@ SPECTRAL_BALANCED: SpectralPolicy = {
     "scope": "ffn",  # FFN layers only
     "correction_enabled": False,
     "max_caps": 5,
+    "max_spectral_norm": None,
     "multiple_testing": {"method": "bh", "alpha": 0.05, "m": 4},
 }
 
@@ -60,6 +62,7 @@ SPECTRAL_AGGRESSIVE: SpectralPolicy = {
     "scope": "all",  # All layers including attention
     "correction_enabled": True,
     "max_caps": 8,
+    "max_spectral_norm": None,
     "multiple_testing": {"method": "bh", "alpha": 0.1, "m": 4},
 }
 
@@ -70,6 +73,7 @@ SPECTRAL_ATTN_AWARE: SpectralPolicy = {
     "scope": "attn",  # Attention layers only
     "correction_enabled": False,
     "max_caps": 5,
+    "max_spectral_norm": None,
     "multiple_testing": {"method": "bh", "alpha": 0.05, "m": 4},
 }
 
@@ -81,7 +85,8 @@ RMT_CONSERVATIVE: RMTPolicyDict = {
     "deadband": 0.05,  # 5% deadband - strict threshold
     "margin": 1.3,  # Lower margin for conservative detection
     "correct": True,  # Enable automatic correction
-    "epsilon": {"attn": 0.05, "ffn": 0.06, "embed": 0.07, "other": 0.07},
+    "epsilon_default": 0.06,
+    "epsilon_by_family": {"attn": 0.05, "ffn": 0.06, "embed": 0.07, "other": 0.07},
 }
 
 # Balanced RMT policy - good for most use cases
@@ -90,7 +95,8 @@ RMT_BALANCED: RMTPolicyDict = {
     "deadband": 0.10,  # 10% deadband - reasonable tolerance
     "margin": 1.5,  # Standard margin for outlier detection
     "correct": False,  # Monitor-only by default
-    "epsilon": {"attn": 0.08, "ffn": 0.10, "embed": 0.12, "other": 0.12},
+    "epsilon_default": 0.10,
+    "epsilon_by_family": {"attn": 0.08, "ffn": 0.10, "embed": 0.12, "other": 0.12},
 }
 
 # Aggressive RMT policy - for research/experimental use
@@ -99,7 +105,8 @@ RMT_AGGRESSIVE: RMTPolicyDict = {
     "deadband": 0.15,  # 15% deadband - more permissive
     "margin": 1.8,  # Higher margin allows more deviation
     "correct": True,  # Enable automatic correction
-    "epsilon": {"attn": 0.15, "ffn": 0.15, "embed": 0.15, "other": 0.15},
+    "epsilon_default": 0.15,
+    "epsilon_by_family": {"attn": 0.15, "ffn": 0.15, "embed": 0.15, "other": 0.15},
 }
 
 # === Variance Guard Policies ===
@@ -276,6 +283,8 @@ def get_spectral_policy(
                     policy["scope"] = tier_config["scope"]
                 if "max_caps" in tier_config:
                     policy["max_caps"] = tier_config["max_caps"]
+                if "max_spectral_norm" in tier_config:
+                    policy["max_spectral_norm"] = tier_config["max_spectral_norm"]
                 if "family_caps" in tier_config:
                     policy["family_caps"] = tier_config["family_caps"]
                 if "multiple_testing" in tier_config:
@@ -390,9 +399,10 @@ def get_rmt_policy(name: str = "balanced", *, use_yaml: bool = True) -> RMTPolic
                     policy["deadband"] = tier_config["deadband"]
                 if "margin" in tier_config:
                     policy["margin"] = tier_config["margin"]
-                # Use epsilon_by_family as the epsilon dict
+                if "epsilon_default" in tier_config:
+                    policy["epsilon_default"] = tier_config["epsilon_default"]
                 if "epsilon_by_family" in tier_config:
-                    policy["epsilon"] = tier_config["epsilon_by_family"]
+                    policy["epsilon_by_family"] = tier_config["epsilon_by_family"]
         except Exception:
             # Fallback to hardcoded values on any error
             pass