archscope 0.2.2__tar.gz → 0.2.3__tar.gz

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

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

@@ -1,6 +1,6 @@
 [project]
 name = "archscope"
-version = "0.2.2"
+version = "0.2.3"
 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.3/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.3"
+
+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.3}/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.3}/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.2 → archscope-0.2.3}/src/archscope/kazdov_backend.py

@@ -95,6 +95,7 @@ class KazdovBackend(Backend):
 
     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.

archscope-0.2.3/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.3}/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.3/src/archscope.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: archscope
-Version: 0.2.2
+Version: 0.2.3
 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.3}/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.3}/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.3"
+
+
+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__",
-]