archscope 0.2.3__tar.gz → 0.2.4__tar.gz

{archscope-0.2.3/src/archscope.egg-info → archscope-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: archscope
-Version: 0.2.3
+Version: 0.2.4
 Summary: Lightweight workbench for cross-architecture mechanistic interpretability experiments on small models
 Author: Juan Cruz Dovzak
 License: Apache-2.0

{archscope-0.2.3 → archscope-0.2.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "archscope"
-version = "0.2.3"
+version = "0.2.4"
 description = "Lightweight workbench for cross-architecture mechanistic interpretability experiments on small models"
 readme = "README.md"
 authors = [{name = "Juan Cruz Dovzak"}]

{archscope-0.2.3 → archscope-0.2.4}/src/archscope/__init__.py

@@ -25,7 +25,7 @@ Quick start::
     print(result.to_markdown())
 """
 
-__version__ = "0.2.3"
+__version__ = "0.2.4"
 
 from . import probes, sae, neurons, attribute, backends, circuits, transfer, bench, lens, diff
 from .loader import load_model, make_tokenize_fn

archscope-0.2.4/src/archscope/kazdov_backend.py

@@ -0,0 +1,99 @@
+"""Backend for custom architectures that expose layers via ``model.blocks``.
+
+Originally written for kazdov-α (a transformer-style decoder LM with hybrid
+MoBE-BCN + MHA attention) — but the backend is generic. It works for ANY
+PyTorch model where:
+
+- residual blocks are exposed as ``model.blocks`` (a ``nn.ModuleList``)
+- ``model.d_model`` (or ``model.hidden_size``) is set on the model
+- forward signature is ``model(input_ids, attention_mask=None, ...)``
+
+This is the simplest pattern for registering a custom architecture with
+archscope. If your model uses a different convention (e.g., ``model.layers``
+under another parent), subclass ``Backend`` directly — this module is a
+working example.
+
+The backend registers under the name ``"kazdov"`` for historical reasons.
+It used to be coupled to a private model-loading function; that function
+was moved out of the shipped package since it depended on a private
+repository. To load your own custom model, do it yourself and then call
+``Backend.for_model(model, hint="kazdov")``.
+"""
+from __future__ import annotations
+import torch
+
+from .backends import Backend, ActivationRecord
+
+
+@Backend.register("kazdov")
+class KazdovBackend(Backend):
+    """Generic backend for models exposing layers via ``model.blocks``.
+
+    Captures the output of each block via forward hooks (the model is
+    expected to not implement ``output_hidden_states=True`` natively).
+
+    Requirements on the model:
+    - ``model.blocks`` is a ``nn.ModuleList`` of residual blocks.
+    - ``model.d_model`` or ``model.hidden_size`` is set.
+    - ``model(input_ids, attention_mask=...)`` is the forward signature.
+    """
+
+    def layer_names(self) -> list[str]:
+        n_layers = len(self.model.blocks)
+        return [f"layer_{i}.residual" for i in range(n_layers)]
+
+    def extract(self, inputs, layers=None):
+        layers = layers or self.layer_names()
+        self._validate_layers(layers)
+        captures: dict[str, torch.Tensor] = {}
+
+        # Register a forward hook on each requested block.
+        hooks = []
+        for layer_name in layers:
+            idx = int(layer_name.split("_")[1].split(".")[0])
+            if idx >= len(self.model.blocks):
+                continue
+            block = self.model.blocks[idx]
+
+            def make_hook(name):
+                def hook(module, inp, out):
+                    tensor = out if isinstance(out, torch.Tensor) else out[0]
+                    captures[name] = tensor.detach()
+                return hook
+            hooks.append(block.register_forward_hook(make_hook(layer_name)))
+
+        try:
+            with torch.no_grad():
+                if isinstance(inputs, dict):
+                    input_ids = inputs["input_ids"]
+                    attn = inputs.get("attention_mask")
+                else:
+                    input_ids = inputs
+                    attn = None
+                self.model(input_ids, attention_mask=attn)
+        finally:
+            for h in hooks:
+                h.remove()
+
+        records = []
+        for layer_name in layers:
+            if layer_name not in captures:
+                continue
+            records.append(ActivationRecord(
+                layer_name=layer_name,
+                activations=captures[layer_name],
+                meta={"kind": "residual", "arch": "kazdov-blocks"},
+            ))
+        return records
+
+    def hidden_dim(self, layer_name: str) -> int:
+        # Some custom models expose this as `d_model`, others as `hidden_size`.
+        for attr in ("d_model", "hidden_size"):
+            v = getattr(self.model, attr, None)
+            if v is not None:
+                return v
+        raise ValueError(
+            f"Cannot infer hidden_dim for {type(self.model).__name__}: "
+            f"set model.d_model or model.hidden_size, or subclass KazdovBackend "
+            f"and override hidden_dim()."
+        )

{archscope-0.2.3 → archscope-0.2.4/src/archscope.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: archscope
-Version: 0.2.3
+Version: 0.2.4
 Summary: Lightweight workbench for cross-architecture mechanistic interpretability experiments on small models
 Author: Juan Cruz Dovzak
 License: Apache-2.0

{archscope-0.2.3 → archscope-0.2.4}/tests/test_circuits_3arch.py

@@ -11,13 +11,14 @@ import os
 
 import torch
 
-sys.path.insert(0, "/Users/kazdov/code/OriginalKazdov/archscope/src")
+sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent.parent / "src"))
 
 from archscope import circuits
-from archscope.kazdov_backend import load_kazdov_checkpoint
+import sys as _sys; _sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent.parent / "scripts"))
+from _kazdov_loader import load_kazdov_checkpoint
 
 
-CHECKPOINT_KAZDOV = "/Users/kazdov/code/OriginalKazdov/_models/kazdov-98m-alpha"
+CHECKPOINT_KAZDOV = __import__("os").environ.get("KAZDOV_CHECKPOINT", "/Users/kazdov/code/OriginalKazdov/_models/kazdov-98m-alpha")
 PYTHIA_NAME = "EleutherAI/pythia-160m"
 MAMBA_NAME = "state-spaces/mamba-130m-hf"
 
@@ -97,7 +98,7 @@ def main():
     print("  • concentration relative ≈ 0 → highly confident predictions (concentrated)")
 
     # Save
-    out_path = "/Users/kazdov/code/OriginalKazdov/archscope/_research/circuits_3arch.json"
+    out_path = "str(__import__("pathlib").Path(__file__).parent.parent / "_research")/circuits_3arch.json"
     os.makedirs(os.path.dirname(out_path), exist_ok=True)
     with open(out_path, "w") as f:
         json.dump(all_results, f, indent=2, default=str)

{archscope-0.2.3 → archscope-0.2.4}/tests/test_diff.py

@@ -11,7 +11,7 @@ import time
 import copy
 import torch
 
-sys.path.insert(0, "/Users/kazdov/code/OriginalKazdov/archscope/src")
+sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent.parent / "src"))
 
 from archscope import diff
 

{archscope-0.2.3 → archscope-0.2.4}/tests/test_kazdov_integration.py

@@ -9,14 +9,15 @@ import sys
 import time
 import torch
 
-sys.path.insert(0, "/Users/kazdov/code/OriginalKazdov/archscope/src")
+sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent.parent / "src"))
 
 from archscope import probes, sae, neurons
 from archscope.backends import Backend
-from archscope.kazdov_backend import load_kazdov_checkpoint
+import sys as _sys; _sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent.parent / "scripts"))
+from _kazdov_loader import load_kazdov_checkpoint
 
 
-CHECKPOINT = "/Users/kazdov/code/OriginalKazdov/_models/kazdov-98m-alpha"
+CHECKPOINT = __import__("os").environ.get("KAZDOV_CHECKPOINT", "/Users/kazdov/code/OriginalKazdov/_models/kazdov-98m-alpha")
 
 
 def tokenize(tokenizer, texts: list[str]) -> dict:

{archscope-0.2.3 → archscope-0.2.4}/tests/test_lens.py

@@ -9,7 +9,7 @@ import sys
 import time
 import torch
 
-sys.path.insert(0, "/Users/kazdov/code/OriginalKazdov/archscope/src")
+sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent.parent / "src"))
 
 from archscope import lens
 

{archscope-0.2.3 → archscope-0.2.4}/tests/test_mamba_integration.py

@@ -10,7 +10,7 @@ import sys
 import time
 import torch
 
-sys.path.insert(0, "/Users/kazdov/code/OriginalKazdov/archscope/src")
+sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent.parent / "src"))
 
 from archscope import probes, sae, neurons
 from archscope.backends import Backend

{archscope-0.2.3 → archscope-0.2.4}/tests/test_mamba_ssm_state.py

@@ -11,7 +11,7 @@ import sys
 import time
 import torch
 
-sys.path.insert(0, "/Users/kazdov/code/OriginalKazdov/archscope/src")
+sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent.parent / "src"))
 
 from archscope import sae
 from archscope.backends import Backend

{archscope-0.2.3 → archscope-0.2.4}/tests/test_probe_transfer.py

@@ -14,14 +14,15 @@ import sys
 import time
 import torch
 
-sys.path.insert(0, "/Users/kazdov/code/OriginalKazdov/archscope/src")
+sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent.parent / "src"))
 
 from archscope import transfer
 from archscope.backends import Backend
-from archscope.kazdov_backend import load_kazdov_checkpoint
+import sys as _sys; _sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent.parent / "scripts"))
+from _kazdov_loader import load_kazdov_checkpoint
 
 
-CHECKPOINT_KAZDOV = "/Users/kazdov/code/OriginalKazdov/_models/kazdov-98m-alpha"
+CHECKPOINT_KAZDOV = __import__("os").environ.get("KAZDOV_CHECKPOINT", "/Users/kazdov/code/OriginalKazdov/_models/kazdov-98m-alpha")
 PYTHIA_NAME = "EleutherAI/pythia-160m"
 
 

{archscope-0.2.3 → archscope-0.2.4}/tests/test_pythia_end_to_end.py

@@ -13,7 +13,7 @@ import sys
 import time
 import torch
 
-sys.path.insert(0, "/Users/kazdov/code/OriginalKazdov/archscope/src")
+sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent.parent / "src"))
 
 from archscope import probes, sae, neurons, attribute
 from archscope.backends import Backend

{archscope-0.2.3 → archscope-0.2.4}/tests/test_unit.py

@@ -22,7 +22,7 @@ def test_imports():
     import archscope
     from archscope import (probes, sae, neurons, attribute, backends,
                             circuits, transfer, bench, lens, diff)
-    assert archscope.__version__ == "0.2.3"
+    assert archscope.__version__ == "0.2.4"
 
 
 def test_loader_exports():

archscope-0.2.3/src/archscope/kazdov_backend.py

@@ -1,142 +0,0 @@
-"""Backend for kazdov-α (and related Kazdov family models).
-
-Kazdov-α is a transformer-style decoder LM with hybrid attention (MoBE-BCN
-mixture of bilinear experts + standard MHA in parallel). Architecturally
-closer to standard transformer than to pure RNN/SSM — but the BCN attention
-branch makes it a distinct architecture family for cross-arch interp.
-
-Differences from HF transformer:
-- No HF AutoModelForCausalLM interface (custom forward signature)
-- Layers exposed as `model.blocks` (ModuleList)
-- No `output_hidden_states=True` argument — we capture via forward hooks
-- Forward signature: (input_ids, attention_mask=None, labels=None)
-"""
-from __future__ import annotations
-import sys
-from pathlib import Path
-import torch
-
-from .backends import Backend, ActivationRecord
-
-
-KAZDOV_REPO = Path.home() / "code" / "OriginalKazdov" / "kazdov"
-
-
-def _ensure_kazdov_importable():
-    """Add kazdov repo to sys.path so we can import KazdovLM."""
-    p = str(KAZDOV_REPO)
-    if p not in sys.path:
-        sys.path.insert(0, p)
-
-
-def load_kazdov_checkpoint(checkpoint_path: str | Path, device: str = "cpu"):
-    """Load kazdov-α from a checkpoint directory.
-
-    Expects: config.json + final.pt (or latest.pt) in the directory.
-    Returns: (model in eval mode, tokenizer wrapper).
-    """
-    _ensure_kazdov_importable()
-    from kazdov.kazdov_lm import KazdovLM
-    import json
-
-    ckpt_dir = Path(checkpoint_path)
-    config = json.loads((ckpt_dir / "config.json").read_text())
-    model_cfg = config["model_cfg"]
-
-    model = KazdovLM(
-        vocab_size=model_cfg["vocab_size"],
-        d_model=model_cfg["d_model"],
-        n_layers=model_cfg["n_layers"],
-        n_heads=model_cfg["n_heads"],
-        rank=model_cfg["rank"],
-        mlp_dim=model_cfg.get("mlp_dim"),
-        max_len=model_cfg.get("max_len", 256),
-        use_trilinear=model_cfg.get("use_trilinear", False),
-        use_bi_bcn=model_cfg.get("use_bi_bcn", False),
-        use_hybrid_mha=model_cfg.get("use_hybrid_mha", True),
-        use_mobe=model_cfg.get("use_mobe", False),
-        n_experts=model_cfg.get("n_experts", 1),
-    )
-
-    # Try final.pt then latest.pt
-    for fname in ("final.pt", "latest.pt"):
-        f = ckpt_dir / fname
-        if f.exists():
-            state = torch.load(f, map_location=device, weights_only=False)
-            if isinstance(state, dict) and "model" in state:
-                state = state["model"]
-            model.load_state_dict(state, strict=False)
-            break
-    else:
-        raise FileNotFoundError(f"No final.pt or latest.pt in {ckpt_dir}")
-
-    model.to(device).eval()
-
-    # Tokenizer: kazdov used GPT-2 tokenizer per memory
-    from transformers import GPT2Tokenizer
-    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-
-    return model, tokenizer
-
-
-@Backend.register("kazdov")
-class KazdovBackend(Backend):
-    """Backend for kazdov-family models (KazdovLM, MoBE-BCN variants).
-
-    Uses forward hooks to capture residual stream after each KazdovBlock,
-    since the model doesn't expose output_hidden_states.
-    """
-
-    def layer_names(self) -> list[str]:
-        n_layers = len(self.model.blocks)
-        return [f"layer_{i}.residual" for i in range(n_layers)]
-
-    def extract(self, inputs, layers=None):
-        layers = layers or self.layer_names()
-        self._validate_layers(layers)
-        captures: dict[str, torch.Tensor] = {}
-
-        # Register a forward hook on each requested block.
-        hooks = []
-        for layer_name in layers:
-            idx = int(layer_name.split("_")[1].split(".")[0])
-            if idx >= len(self.model.blocks):
-                continue
-            block = self.model.blocks[idx]
-
-            def make_hook(name):
-                def hook(module, inp, out):
-                    tensor = out if isinstance(out, torch.Tensor) else out[0]
-                    captures[name] = tensor.detach()
-                return hook
-            hooks.append(block.register_forward_hook(make_hook(layer_name)))
-
-        try:
-            # Kazdov forward signature: model(input_ids, attention_mask=None)
-            with torch.no_grad():
-                if isinstance(inputs, dict):
-                    input_ids = inputs["input_ids"]
-                    attn = inputs.get("attention_mask")
-                else:
-                    input_ids = inputs
-                    attn = None
-                self.model(input_ids, attention_mask=attn)
-        finally:
-            for h in hooks:
-                h.remove()
-
-        records = []
-        for layer_name in layers:
-            if layer_name not in captures:
-                continue
-            records.append(ActivationRecord(
-                layer_name=layer_name,
-                activations=captures[layer_name],
-                meta={"kind": "residual", "arch": "kazdov-mobe-bcn"},
-            ))
-        return records
-
-    def hidden_dim(self, layer_name: str) -> int:
-        return self.model.d_model