archscope 0.2.2__tar.gz → 0.2.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: archscope
-Version: 0.2.2
+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.2 → archscope-0.2.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "archscope"
-version = "0.2.2"
+version = "0.2.4"
 description = "Lightweight workbench for cross-architecture mechanistic interpretability experiments on small models"
 readme = "README.md"
 authors = [{name = "Juan Cruz Dovzak"}]
@@ -43,3 +43,6 @@ build-backend = "setuptools.build_meta"
 
 [tool.setuptools.packages.find]
 where = ["src"]
+
+[tool.setuptools.package-data]
+"archscope" = ["py.typed"]

archscope-0.2.4/src/archscope/__init__.py

@@ -0,0 +1,44 @@
+"""archscope — cross-architecture mechanistic interpretability workbench.
+
+Core methods (architecture-agnostic):
+- probes:    linear/MLP probes on hidden states (Drop the Act-style)
+- sae:       Dense + Rank-1 sparse autoencoders (WriteSAE-style)
+- neurons:   contrastive neuron modulation
+- attribute: activation patching + DIM decomposition
+- circuits:  induction head, copy, attention-concentration detectors
+- lens:      logit lens + tuned lens (Belrose et al 2023)
+- diff:      base vs fine-tuned model comparison
+
+Experiment infrastructure:
+- backends:  unified extraction API across architectures
+- transfer:  cross-arch probe transfer via paired-activation alignment
+- bench:     InterpProfile standardized benchmark
+- loader:    one-call HuggingFace model + tokenizer + backend loader
+
+Backends: ``transformer``, ``mamba`` (incl. ssm_state), ``kazdov``, ``recurrent``.
+
+Quick start::
+
+    import archscope as ai
+    model, tok, backend = ai.load_model("EleutherAI/pythia-160m", arch="transformer")
+    result = ai.lens.logit_lens(model, tok, "The capital of France is", target_token=" Paris")
+    print(result.to_markdown())
+"""
+
+__version__ = "0.2.4"
+
+from . import probes, sae, neurons, attribute, backends, circuits, transfer, bench, lens, diff
+from .loader import load_model, make_tokenize_fn
+
+# Kazdov backend registers itself on import — optional, only if kazdov repo present
+try:
+    from . import kazdov_backend  # noqa: F401
+except ImportError:
+    pass
+
+__all__ = [
+    "probes", "sae", "neurons", "attribute", "backends",
+    "circuits", "transfer", "bench", "lens", "diff",
+    "load_model", "make_tokenize_fn",
+    "__version__",
+]

{archscope-0.2.2 → archscope-0.2.4}/src/archscope/backends.py

@@ -77,6 +77,21 @@ class Backend(abc.ABC):
         """Dimensionality of activations at a given layer."""
         ...
 
+    def _validate_layers(self, layers: list[str]) -> None:
+        """Raise a clear error if any requested layer name isn't valid."""
+        valid = set(self.layer_names())
+        bad = [ln for ln in layers if ln not in valid]
+        if bad:
+            # Show first few valid examples so users see the format.
+            sample = ", ".join(self.layer_names()[:3])
+            n_total = len(valid)
+            raise ValueError(
+                f"Unknown layer name(s) for {type(self).__name__}: {bad}. "
+                f"Valid layer names look like: {sample}{', ...' if n_total > 3 else ''} "
+                f"(total: {n_total} layer names). Call `backend.layer_names()` "
+                f"to see all valid options."
+            )
+
 
 @Backend.register("transformer")
 class TransformerBackend(Backend):
@@ -89,6 +104,7 @@ class TransformerBackend(Backend):
 
     def extract(self, inputs, layers=None):
         layers = layers or self.layer_names()
+        self._validate_layers(layers)
         # Use HF's output_hidden_states=True for clean extraction.
         # Wrap in no_grad: extraction shouldn't build a backward graph.
         with torch.no_grad():
@@ -134,6 +150,7 @@ class MambaBackend(Backend):
 
     def extract(self, inputs, layers=None):
         layers = layers or self.layer_names()
+        self._validate_layers(layers)
 
         need_residual = any(".residual" in ln for ln in layers)
         need_ssm = any(".ssm_state" in ln for ln in layers)

{archscope-0.2.2 → archscope-0.2.4}/src/archscope/cli.py

@@ -13,9 +13,13 @@ from . import __version__
 console = Console()
 
 
-@click.group()
-def cli() -> None:
-    """archscope — cross-architecture mechanistic interpretability toolkit."""
+@click.group(invoke_without_command=True)
+@click.version_option(__version__, "-V", "--version", prog_name="archscope")
+@click.pass_context
+def cli(ctx: click.Context) -> None:
+    """archscope — cross-architecture mechanistic interpretability workbench."""
+    if ctx.invoked_subcommand is None:
+        click.echo(ctx.get_help())
 
 
 @cli.command()

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.4/src/archscope/loader.py

@@ -0,0 +1,76 @@
+"""High-level model loading helper.
+
+Eliminates ~5 lines of HuggingFace boilerplate per example:
+
+    # Before
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    tok = AutoTokenizer.from_pretrained(name)
+    if tok.pad_token is None: tok.pad_token = tok.eos_token
+    model = AutoModelForCausalLM.from_pretrained(name, dtype=torch.float32)
+    model.eval()
+    backend = Backend.for_model(model, hint=arch)
+
+    # After
+    model, tok, backend = archscope.load_model(name, arch="transformer")
+"""
+from __future__ import annotations
+from typing import Any
+
+from .backends import Backend
+
+
+def load_model(
+    name: str,
+    arch: str | None = None,
+    dtype: Any = None,
+    device: str = "cpu",
+) -> tuple[Any, Any, Backend]:
+    """Load a HuggingFace model + tokenizer + matching backend in one call.
+
+    Args:
+        name: HuggingFace model id (e.g., "EleutherAI/pythia-160m") OR a path
+            to a local archscope-registered model (e.g., kazdov checkpoint).
+        arch: Backend hint — one of "transformer", "mamba", "kazdov",
+            "recurrent". If None, auto-detect from HF config.model_type.
+        dtype: torch dtype. Defaults to torch.float32.
+        device: device string. Default "cpu".
+
+    Returns:
+        (model, tokenizer, backend) ready for use in probes/sae/lens/etc.
+    """
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    if dtype is None:
+        dtype = torch.float32
+
+    tokenizer = AutoTokenizer.from_pretrained(name)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    model = AutoModelForCausalLM.from_pretrained(name, dtype=dtype)
+    model = model.to(device).eval()
+
+    backend = Backend.for_model(model, hint=arch)
+    return model, tokenizer, backend
+
+
+def make_tokenize_fn(tokenizer, max_length: int = 32, attention_mask_bool: bool = False):
+    """Return a tokenize function suitable for ``Backend.extract`` and ``probes.fit_probe``.
+
+    Args:
+        tokenizer: a HuggingFace tokenizer.
+        max_length: max sequence length to truncate to.
+        attention_mask_bool: if True, returns attention_mask as a bool tensor
+            (required for ``kazdov`` backend); HF default is int64.
+
+    Returns:
+        A callable ``texts -> dict`` that matches the ``inputs`` format used
+        across archscope.
+    """
+    def fn(texts):
+        out = tokenizer(texts, return_tensors="pt", padding=True,
+                          truncation=True, max_length=max_length)
+        if attention_mask_bool and "attention_mask" in out:
+            out["attention_mask"] = out["attention_mask"].bool()
+        return out
+    return fn

{archscope-0.2.2 → archscope-0.2.4}/src/archscope/probes.py

@@ -110,28 +110,72 @@ class ProbeFit:
 
 
 def _auroc(logits: torch.Tensor, labels: torch.Tensor) -> float:
-    """Simple AUROC from logits + binary labels."""
+    """AUROC from logits + binary labels.
+
+    Returns 0.5 (chance) when only one class is present in `labels`
+    (the typical small-split case) — this is more informative than NaN
+    and avoids sklearn's UndefinedMetricWarning leaking to user code.
+    """
+    import warnings
     from sklearn.metrics import roc_auc_score
     scores = torch.sigmoid(logits).cpu().numpy()
     y = labels.cpu().numpy()
-    try:
-        return float(roc_auc_score(y, scores))
-    except ValueError:
-        return float("nan")   # happens when only one class present
+    if len(set(y.tolist())) < 2:
+        warnings.warn(
+            "Only one class present in this split — AUROC undefined; "
+            "returning 0.5 (chance). Increase val_split or dataset size.",
+            stacklevel=2,
+        )
+        return 0.5
+    return float(roc_auc_score(y, scores))
 
 
 # High-level API matching paper
 
 def fit_probe(
     model,
-    inputs_pos: list,    # examples where target=1 (e.g., faithful)
-    inputs_neg: list,    # examples where target=0 (e.g., reasoning theater)
-    layer_name: str,
+    inputs_pos=None,
+    inputs_neg=None,
+    layer_name: str = "",
     backend_hint: str | None = None,
     config: ProbeConfig | None = None,
     device: str = "cpu",
+    *,
+    tokenizer=None,
+    pos_texts: list[str] | None = None,
+    neg_texts: list[str] | None = None,
+    max_length: int = 32,
 ) -> ProbeFit:
-    """End-to-end: extract activations from model, fit probe."""
+    """End-to-end: extract activations from a model and fit a probe.
+
+    Two calling conventions:
+
+    1. **Pre-tokenized**: pass ``inputs_pos`` and ``inputs_neg`` as already-tokenized
+       dicts (with ``input_ids``, optional ``attention_mask``).
+
+    2. **Texts + tokenizer**: pass ``tokenizer=…``, ``pos_texts=[…]``, ``neg_texts=[…]``
+       and archscope tokenizes for you. The kazdov backend requires a bool
+       attention_mask; we auto-handle that.
+
+    Returns:
+        A ``ProbeFit`` with trained probe and ``.metrics`` (train/val AUROC, loss).
+    """
+    if pos_texts is not None or neg_texts is not None:
+        if tokenizer is None:
+            raise ValueError("pos_texts/neg_texts require a tokenizer= argument")
+        if pos_texts is None or neg_texts is None:
+            raise ValueError("provide both pos_texts and neg_texts")
+        from .loader import make_tokenize_fn
+        tk = make_tokenize_fn(
+            tokenizer, max_length=max_length,
+            attention_mask_bool=(backend_hint == "kazdov"),
+        )
+        inputs_pos = tk(pos_texts)
+        inputs_neg = tk(neg_texts)
+    elif inputs_pos is None or inputs_neg is None:
+        raise ValueError("provide either (inputs_pos, inputs_neg) or "
+                          "(tokenizer, pos_texts, neg_texts)")
+
     backend = Backend.for_model(model, hint=backend_hint)
     config = config or ProbeConfig(layer_name=layer_name)
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: archscope
-Version: 0.2.2
+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.2 → archscope-0.2.4}/src/archscope.egg-info/SOURCES.txt

@@ -11,8 +11,10 @@ src/archscope/cli.py
 src/archscope/diff.py
 src/archscope/kazdov_backend.py
 src/archscope/lens.py
+src/archscope/loader.py
 src/archscope/neurons.py
 src/archscope/probes.py
+src/archscope/py.typed
 src/archscope/sae.py
 src/archscope/transfer.py
 src/archscope.egg-info/PKG-INFO

{archscope-0.2.2 → 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.2 → 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.2 → 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.2 → 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.2 → 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.2 → 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.2 → 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.2 → 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.2 → archscope-0.2.4}/tests/test_unit.py

@@ -22,7 +22,56 @@ def test_imports():
     import archscope
     from archscope import (probes, sae, neurons, attribute, backends,
                             circuits, transfer, bench, lens, diff)
-    assert archscope.__version__ == "0.2.2"
+    assert archscope.__version__ == "0.2.4"
+
+
+def test_loader_exports():
+    """load_model and make_tokenize_fn are exported at top level."""
+    import archscope
+    assert hasattr(archscope, "load_model")
+    assert hasattr(archscope, "make_tokenize_fn")
+    assert callable(archscope.load_model)
+    assert callable(archscope.make_tokenize_fn)
+
+
+def test_layer_name_validation_clear_error():
+    """Backend validates layer names with an informative error."""
+    from archscope.backends import Backend, ActivationRecord
+
+    # Build a minimal mock backend
+    class _MockBackend(Backend):
+        def layer_names(self): return ["layer_0.residual", "layer_1.residual"]
+        def extract(self, inputs, layers=None):
+            layers = layers or self.layer_names()
+            self._validate_layers(layers)
+            return []
+        def hidden_dim(self, layer_name): return 8
+
+    b = _MockBackend(model=None)
+    # Valid layer → no error
+    b.extract({}, layers=["layer_0.residual"])
+    # Invalid layer → clear error message
+    try:
+        b.extract({}, layers=["layer_99.residual"])
+    except ValueError as e:
+        assert "Unknown layer name" in str(e)
+        assert "layer_0.residual" in str(e)  # shows valid example
+        return
+    raise AssertionError("Expected ValueError for invalid layer name")
+
+
+def test_auroc_returns_chance_on_single_class():
+    """_auroc returns 0.5 instead of NaN when only one class is present."""
+    import warnings
+    import torch
+    from archscope.probes import _auroc
+
+    logits = torch.tensor([0.5, 0.2, -0.1])
+    labels = torch.tensor([1.0, 1.0, 1.0])   # only one class
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        result = _auroc(logits, labels)
+    assert result == 0.5
 
 
 def test_diff_dataclasses():

archscope-0.2.2/src/archscope/__init__.py

@@ -1,30 +0,0 @@
-"""archscope: unified mech interp toolkit across small + RNN + transformer.
-
-Four core methods unified under a single API:
-- probes: linear/MLP probes over hidden states (Drop the Act inspired)
-- sae: sparse autoencoders for residual + recurrent state (WriteSAE)
-- neurons: targeted neuron modulation via contrastive search (Nous Research)
-- attribute: activation patching + DIM decomposition (Multi-Agent Sycophancy)
-
-Each method exposes the same architecture-agnostic API:
-- .extract(model, inputs) -> hidden states / activations
-- .fit(activations, labels) -> learned tool
-- .apply(model, inputs) -> modified outputs / scores / explanations
-
-Designed for cross-architecture comparison: transformer, Mamba/SSM, custom RNN.
-"""
-
-__version__ = "0.2.2"
-
-from . import probes, sae, neurons, attribute, backends, circuits, transfer, bench, lens, diff
-
-# Kazdov backend registers itself on import — optional, only if kazdov repo present
-try:
-    from . import kazdov_backend  # noqa: F401
-except ImportError:
-    pass
-
-__all__ = [
-    "probes", "sae", "neurons", "attribute", "backends",
-    "circuits", "transfer", "bench", "lens", "diff", "__version__",
-]

archscope-0.2.2/src/archscope/kazdov_backend.py

@@ -1,141 +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()
-        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