archscope 0.2.2__py3-none-any.whl

archscope/diff.py

@@ -0,0 +1,212 @@
+"""Model Diff — compare base model vs fine-tuned version, find what changed.
+
+This is the practical question every fine-tuner asks:
+  "I fine-tuned X on my data. What did it change in the model's brain?"
+
+Given (base_model, fine_tuned_model) with the same architecture, archscope.diff
+returns a structured ModelDiff with:
+- Per-layer activation drift (how much each layer's residual stream moved)
+- Top-K shifted neurons per layer
+- Circuit score deltas (induction, copy, etc.)
+- Optional: probe direction drift (do trained probes still work?)
+
+Requirements:
+- base and fine_tuned must share architecture + tokenizer
+- A small calibration_texts list (16-100 short texts) to measure drift on
+
+Usage:
+    from archscope.diff import compare
+    result = compare(base, finetuned, tokenizer, calibration_texts, backend_hint="transformer")
+    print(result.to_markdown())
+"""
+from __future__ import annotations
+from dataclasses import dataclass, field
+import torch
+
+from .backends import Backend
+from . import circuits
+
+
+@dataclass
+class LayerDrift:
+    """How much one layer's residual stream changed under fine-tuning."""
+    layer: int
+    layer_name: str
+    mean_l2_delta: float           # mean L2 norm of (a_ft - a_base) per token
+    relative_drift: float           # mean_l2_delta / mean ||a_base||
+    cosine_similarity: float        # mean cos(a_base, a_ft) — closer to 1 = less change
+    top_shifted_neurons: list[tuple[int, float]] = field(default_factory=list)
+
+
+@dataclass
+class CircuitDelta:
+    """How a single circuit score changed."""
+    name: str
+    base_score: float
+    fine_tuned_score: float
+    delta: float                   # ft - base
+    relative_change: float          # delta / |base| (capped at large values)
+
+
+@dataclass
+class ModelDiff:
+    """Full diff between base and fine-tuned model."""
+    arch_family: str
+    n_layers: int
+    n_calibration_texts: int
+    layer_drift: list[LayerDrift] = field(default_factory=list)
+    circuit_deltas: list[CircuitDelta] = field(default_factory=list)
+    notes: list[str] = field(default_factory=list)
+
+    def top_changed_layers(self, k: int = 3) -> list[LayerDrift]:
+        """Return the k layers with highest relative drift."""
+        return sorted(self.layer_drift, key=lambda d: -d.relative_drift)[:k]
+
+    def to_markdown(self) -> str:
+        lines = [f"# Model Diff — {self.arch_family} ({self.n_layers} layers)\n"]
+        lines.append(f"Calibration texts: {self.n_calibration_texts}\n")
+
+        # Layer drift table
+        lines.append("## Per-layer residual drift\n")
+        lines.append("| Layer | mean ‖Δa‖ | relative | cosine sim |")
+        lines.append("|------:|---------:|---------:|----------:|")
+        for d in self.layer_drift:
+            lines.append(
+                f"| {d.layer:>5} | {d.mean_l2_delta:>8.3f} | {d.relative_drift:>7.3%} | {d.cosine_similarity:>9.3f} |"
+            )
+
+        # Top changed layers
+        lines.append("\n## Top-3 most-changed layers\n")
+        top = self.top_changed_layers(3)
+        for d in top:
+            lines.append(f"- **layer {d.layer}**: relative drift {d.relative_drift:.1%}, cosine {d.cosine_similarity:.3f}")
+            if d.top_shifted_neurons:
+                idxs = ", ".join(f"#{i}({delta:+.2f})" for i, delta in d.top_shifted_neurons[:5])
+                lines.append(f"  - top-5 neurons shifted: {idxs}")
+
+        # Circuit deltas
+        if self.circuit_deltas:
+            lines.append("\n## Circuit deltas (induction, copy, concentration)\n")
+            lines.append("| Circuit | base | fine-tuned | Δ | Δ% |")
+            lines.append("|---------|-----:|-----------:|--:|---:|")
+            for c in self.circuit_deltas:
+                pct = f"{c.relative_change:+.1%}" if abs(c.base_score) > 1e-9 else "—"
+                lines.append(
+                    f"| {c.name} | {c.base_score:.3f} | {c.fine_tuned_score:.3f} | {c.delta:+.3f} | {pct} |"
+                )
+
+        if self.notes:
+            lines.append("\n## Notes\n")
+            for n in self.notes:
+                lines.append(f"- {n}")
+        return "\n".join(lines)
+
+
+def _activation_drift(
+    base_acts: torch.Tensor,
+    ft_acts: torch.Tensor,
+    top_k_neurons: int = 10,
+) -> tuple[float, float, float, list[tuple[int, float]]]:
+    """Compute drift stats between two activation tensors of shape (B, T, H) or (B, H).
+
+    Returns (mean_l2_delta, relative_drift, cosine_sim, top_k_shifted_neurons).
+    """
+    if base_acts.shape != ft_acts.shape:
+        raise ValueError(f"Shape mismatch: {tuple(base_acts.shape)} vs {tuple(ft_acts.shape)}")
+    # Flatten leading dims; keep hidden dim
+    H = base_acts.shape[-1]
+    a = base_acts.reshape(-1, H)
+    b = ft_acts.reshape(-1, H)
+    diff = (b - a).float()
+    l2_per_token = diff.norm(dim=-1)
+    mean_l2 = float(l2_per_token.mean().item())
+    base_l2 = float(a.float().norm(dim=-1).mean().item()) + 1e-9
+    relative = mean_l2 / base_l2
+    # Cosine sim (per-token, then averaged)
+    cos = torch.nn.functional.cosine_similarity(a.float(), b.float(), dim=-1)
+    cosine_sim = float(cos.mean().item())
+    # Per-neuron drift: mean absolute delta per channel
+    per_neuron_delta = diff.abs().mean(dim=0)   # (H,)
+    topk = torch.topk(per_neuron_delta, k=min(top_k_neurons, H))
+    top_shifted = [(int(i), float(v)) for i, v in zip(topk.indices.tolist(), topk.values.tolist())]
+    return mean_l2, relative, cosine_sim, top_shifted
+
+
+def compare(
+    base_model,
+    fine_tuned_model,
+    tokenizer,
+    calibration_texts: list[str],
+    backend_hint: str | None = None,
+    run_circuits: bool = True,
+    max_length: int = 32,
+    device: str = "cpu",
+) -> ModelDiff:
+    """Compare base vs fine-tuned model on calibration texts.
+
+    Both models must share the same architecture and tokenizer.
+    """
+    base_backend = Backend.for_model(base_model, hint=backend_hint)
+    ft_backend = Backend.for_model(fine_tuned_model, hint=backend_hint)
+
+    layer_names = [n for n in base_backend.layer_names() if ".residual" in n]
+    ft_layer_names = [n for n in ft_backend.layer_names() if ".residual" in n]
+    if layer_names != ft_layer_names:
+        raise ValueError("base and fine_tuned have different layer structure — "
+                         "they must share architecture")
+
+    # Tokenize calibration
+    enc = tokenizer(calibration_texts, return_tensors="pt", padding=True,
+                     truncation=True, max_length=max_length)
+    inputs = {"input_ids": enc["input_ids"].to(device)}
+    if "attention_mask" in enc:
+        inputs["attention_mask"] = enc["attention_mask"].to(device)
+
+    # Need attention_mask as bool for kazdov backend
+    if backend_hint == "kazdov" and "attention_mask" in inputs:
+        inputs["attention_mask"] = inputs["attention_mask"].bool()
+
+    with torch.no_grad():
+        base_records = base_backend.extract(inputs, layers=layer_names)
+        ft_records = ft_backend.extract(inputs, layers=layer_names)
+
+    diff = ModelDiff(
+        arch_family=backend_hint or "auto",
+        n_layers=len(layer_names),
+        n_calibration_texts=len(calibration_texts),
+    )
+
+    for base_rec, ft_rec in zip(base_records, ft_records):
+        mean_l2, rel, cos_sim, top_neurons = _activation_drift(
+            base_rec.activations, ft_rec.activations,
+        )
+        idx = int(base_rec.layer_name.split("_")[1].split(".")[0])
+        diff.layer_drift.append(LayerDrift(
+            layer=idx, layer_name=base_rec.layer_name,
+            mean_l2_delta=mean_l2,
+            relative_drift=rel,
+            cosine_similarity=cos_sim,
+            top_shifted_neurons=top_neurons,
+        ))
+
+    # Circuit deltas
+    if run_circuits:
+        try:
+            base_circs = circuits.run_all_circuits(base_model, tokenizer=tokenizer, device=device)
+            ft_circs = circuits.run_all_circuits(fine_tuned_model, tokenizer=tokenizer, device=device)
+            for name in base_circs:
+                if name not in ft_circs:
+                    continue
+                bs = base_circs[name].score
+                fs = ft_circs[name].score
+                delta = fs - bs
+                rel = delta / (abs(bs) + 1e-9)
+                diff.circuit_deltas.append(CircuitDelta(
+                    name=name,
+                    base_score=bs, fine_tuned_score=fs,
+                    delta=delta, relative_change=rel,
+                ))
+        except Exception as e:
+            diff.notes.append(f"circuit comparison error: {str(e)[:80]}")
+
+    return diff

archscope/kazdov_backend.py

@@ -0,0 +1,141 @@
+"""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

archscope/lens.py

@@ -0,0 +1,304 @@
+"""Logit lens and Tuned lens — project intermediate residual stream into vocab space.
+
+The logit lens (Nostalgebraist, 2020) applies the model's own final norm + unembedding
+to every layer's residual stream, revealing "what the model would predict if forced
+to commit at this layer." It tells you which layers do the work of forming the final
+prediction.
+
+The tuned lens (Belrose et al, 2023) learns per-layer affine transformations that
+correct for representation drift, producing a more faithful intermediate decoding.
+
+Both work on any architecture exposing residual stream via Backend.extract — i.e.
+transformer, mamba, kazdov, custom recurrent.
+"""
+from __future__ import annotations
+from dataclasses import dataclass, field
+import torch
+import torch.nn as nn
+
+from .backends import Backend
+from ._utils import resolve_unembedding, resolve_final_norm
+
+
+@dataclass
+class LayerPrediction:
+    """What the lens says at one layer."""
+    layer: int
+    layer_name: str
+    top_tokens: list[tuple[int, str, float]]    # (token_id, token_str, prob)
+    target_prob: float | None = None             # if target_token provided, prob at this layer
+    target_rank: int | None = None               # rank of target token in this layer's distribution
+    entropy: float = 0.0                          # confidence (low = peaked, high = uncertain)
+
+
+@dataclass
+class LensResult:
+    """Output of a lens pass over multiple layers."""
+    prompt: str
+    target_token: str | None = None
+    target_token_id: int | None = None
+    layers: list[LayerPrediction] = field(default_factory=list)
+    method: str = "logit_lens"
+
+    def to_markdown(self) -> str:
+        """Quick formatted display."""
+        lines = [f"### {self.method} on `{self.prompt}`"]
+        if self.target_token:
+            lines.append(f"Target: `{self.target_token}` (id={self.target_token_id})")
+        lines.append("")
+        lines.append("| Layer | top-1 token | top-1 prob | target prob | rank | entropy |")
+        lines.append("|-------|-------------|-----------:|------------:|-----:|--------:|")
+        for lp in self.layers:
+            if lp.top_tokens:
+                top_id, top_str, top_p = lp.top_tokens[0]
+            else:
+                top_str, top_p = "?", float("nan")
+            tgt_p = f"{lp.target_prob:.3f}" if lp.target_prob is not None else "—"
+            tgt_r = f"{lp.target_rank}" if lp.target_rank is not None else "—"
+            lines.append(
+                f"| {lp.layer:>2} | `{top_str[:20]}` | {top_p:.3f} | {tgt_p} | {tgt_r} | {lp.entropy:.3f} |"
+            )
+        return "\n".join(lines)
+
+
+def _logits_to_pred(
+    logits: torch.Tensor,
+    tokenizer,
+    target_id: int | None = None,
+    top_k: int = 5,
+) -> tuple[list[tuple[int, str, float]], float, int | None, float]:
+    """Convert logit vector → (top_k_tokens, target_prob, target_rank, entropy)."""
+    probs = torch.softmax(logits.float(), dim=-1)
+    top_p, top_i = probs.topk(top_k)
+    top_tokens = []
+    for p, i in zip(top_p.tolist(), top_i.tolist()):
+        try:
+            tok_str = tokenizer.decode([i])
+        except Exception:
+            tok_str = f"<id={i}>"
+        top_tokens.append((i, tok_str, p))
+    if target_id is not None:
+        target_prob = float(probs[target_id].item())
+        target_rank = int((probs > probs[target_id]).sum().item())
+    else:
+        target_prob = None
+        target_rank = None
+    # Shannon entropy in nats
+    entropy = float(-(probs.clamp(min=1e-12) * probs.clamp(min=1e-12).log()).sum().item())
+    return top_tokens, target_prob, target_rank, entropy
+
+
+def logit_lens(
+    model,
+    tokenizer,
+    prompt: str,
+    target_token: str | None = None,
+    layers: list[int] | None = None,
+    backend_hint: str | None = None,
+    top_k: int = 5,
+    device: str = "cpu",
+) -> LensResult:
+    """Apply model's own final-norm + unembedding to each layer's residual stream.
+
+    Returns per-layer predictions: top-k tokens + entropy + (optional) target rank/prob.
+
+    The classic question this answers: "at which layer does the model commit to the
+    final token?" Often you'll see entropy decrease and target rank climb across layers.
+    """
+    backend = Backend.for_model(model, hint=backend_hint)
+    norm = resolve_final_norm(model)
+    unembed = resolve_unembedding(model)
+    if unembed is None:
+        raise ValueError("Could not locate model's lm_head / unembedding module.")
+
+    # Tokenize prompt (handle both HF tokenizer + kazdov-style inputs)
+    enc = tokenizer(prompt, return_tensors="pt")
+    if hasattr(enc, "input_ids"):
+        inputs = {"input_ids": enc.input_ids.to(device)}
+    else:
+        inputs = {"input_ids": enc["input_ids"].to(device)}
+
+    target_id = None
+    if target_token is not None:
+        target_ids = tokenizer(target_token, add_special_tokens=False).input_ids
+        target_id = target_ids[0] if target_ids else None
+
+    # Pick layers
+    all_layer_names = [n for n in backend.layer_names() if ".residual" in n]
+    if layers is None:
+        layer_names = all_layer_names
+    else:
+        layer_names = [f"layer_{i}.residual" for i in layers if f"layer_{i}.residual" in all_layer_names]
+
+    # Extract residual streams
+    with torch.no_grad():
+        records = backend.extract(inputs, layers=layer_names)
+
+    result = LensResult(
+        prompt=prompt,
+        target_token=target_token,
+        target_token_id=target_id,
+        method="logit_lens",
+    )
+
+    for rec in records:
+        # rec.activations shape: (B, T, hidden) — take last-token residual
+        last = rec.activations[:, -1, :]   # (B, hidden)
+        # Apply final norm if present (transformer family). For kazdov it's `model.ln_f`.
+        if norm is not None:
+            try:
+                last = norm(last)
+            except Exception:
+                pass
+        # Project to vocab
+        logits = unembed(last)[0]   # (vocab,)
+        top_tokens, tgt_p, tgt_r, ent = _logits_to_pred(logits, tokenizer, target_id, top_k)
+        idx = int(rec.layer_name.split("_")[1].split(".")[0])
+        result.layers.append(LayerPrediction(
+            layer=idx, layer_name=rec.layer_name,
+            top_tokens=top_tokens,
+            target_prob=tgt_p, target_rank=tgt_r,
+            entropy=ent,
+        ))
+
+    return result
+
+
+# -------------------- TUNED LENS --------------------
+
+class TunedLens(nn.Module):
+    """Learned per-layer affine transformations into the unembedding space.
+
+    For each layer, learn a transformation: residual → adjusted_residual that, when
+    fed through the model's final_norm + unembedding, better matches the final-layer
+    distribution. Following Belrose et al 2023.
+
+    Usage:
+        tl = TunedLens.fit(model, tokenizer, calibration_texts, backend_hint="transformer")
+        result = tl.predict(model, tokenizer, prompt="...", backend_hint="transformer")
+    """
+
+    def __init__(self, model, n_layers: int, hidden_dim: int, backend_hint: str | None = None):
+        super().__init__()
+        self.n_layers = n_layers
+        self.hidden_dim = hidden_dim
+        self.backend_hint = backend_hint
+        # One Linear per layer; identity initialization (so untrained tuned lens ≈ logit lens)
+        self.translators = nn.ModuleList([
+            nn.Linear(hidden_dim, hidden_dim, bias=True)
+            for _ in range(n_layers)
+        ])
+        for t in self.translators:
+            nn.init.eye_(t.weight)
+            nn.init.zeros_(t.bias)
+        self._model_ref = [model]   # avoid registering as submodule
+
+    @classmethod
+    def fit(
+        cls,
+        model,
+        tokenizer,
+        calibration_texts: list[str],
+        backend_hint: str | None = None,
+        epochs: int = 30,
+        lr: float = 1e-3,
+        max_len: int = 64,
+        device: str = "cpu",
+    ) -> "TunedLens":
+        """Train per-layer translators to match the model's own final distribution."""
+        backend = Backend.for_model(model, hint=backend_hint)
+        layer_names = [n for n in backend.layer_names() if ".residual" in n]
+        n_layers = len(layer_names)
+        # Get hidden dim from first layer
+        hidden_dim = backend.hidden_dim(layer_names[0])
+
+        tl = cls(model, n_layers=n_layers, hidden_dim=hidden_dim, backend_hint=backend_hint).to(device)
+        norm = resolve_final_norm(model)
+        unembed = resolve_unembedding(model)
+
+        opt = torch.optim.AdamW(tl.translators.parameters(), lr=lr)
+
+        # Pre-extract all activations + target logits once
+        enc = tokenizer(calibration_texts, return_tensors="pt", padding=True,
+                         truncation=True, max_length=max_len)
+        inputs = {"input_ids": enc["input_ids"].to(device)}
+        with torch.no_grad():
+            records = backend.extract(inputs, layers=layer_names)
+            # Target: model's actual final logits at last position
+            final_residual = records[-1].activations[:, -1, :]
+            if norm is not None:
+                final_residual = norm(final_residual)
+            target_logits = unembed(final_residual).detach()    # (B, vocab)
+            target_log_probs = torch.log_softmax(target_logits.float(), dim=-1)
+
+        for epoch in range(epochs):
+            opt.zero_grad()
+            total_loss = 0.0
+            for i, rec in enumerate(records):
+                last = rec.activations[:, -1, :].detach()
+                translated = tl.translators[i](last)
+                if norm is not None:
+                    translated = norm(translated)
+                pred_logits = unembed(translated)
+                pred_log_probs = torch.log_softmax(pred_logits.float(), dim=-1)
+                # KL divergence target_log_probs || pred_log_probs
+                loss = torch.nn.functional.kl_div(pred_log_probs, target_log_probs,
+                                                   log_target=True, reduction="batchmean")
+                total_loss = total_loss + loss
+            total_loss.backward()
+            opt.step()
+
+        tl.last_loss = float(total_loss.item() / max(n_layers, 1))
+        return tl
+
+    def predict(
+        self,
+        model,
+        tokenizer,
+        prompt: str,
+        target_token: str | None = None,
+        backend_hint: str | None = None,
+        layers: list[int] | None = None,
+        top_k: int = 5,
+        device: str = "cpu",
+    ) -> LensResult:
+        """Apply the trained tuned lens to a new prompt."""
+        backend = Backend.for_model(model, hint=backend_hint or self.backend_hint)
+        norm = resolve_final_norm(model)
+        unembed = resolve_unembedding(model)
+
+        enc = tokenizer(prompt, return_tensors="pt")
+        inputs = {"input_ids": enc["input_ids"].to(device)}
+        target_id = None
+        if target_token is not None:
+            ids = tokenizer(target_token, add_special_tokens=False).input_ids
+            target_id = ids[0] if ids else None
+
+        all_layer_names = [n for n in backend.layer_names() if ".residual" in n]
+        if layers is None:
+            layer_names = all_layer_names
+        else:
+            layer_names = [f"layer_{i}.residual" for i in layers if f"layer_{i}.residual" in all_layer_names]
+
+        with torch.no_grad():
+            records = backend.extract(inputs, layers=layer_names)
+
+        result = LensResult(prompt=prompt, target_token=target_token,
+                            target_token_id=target_id, method="tuned_lens")
+
+        for rec in records:
+            idx = int(rec.layer_name.split("_")[1].split(".")[0])
+            last = rec.activations[:, -1, :]
+            with torch.no_grad():
+                translated = self.translators[idx](last)
+                if norm is not None:
+                    translated = norm(translated)
+                logits = unembed(translated)[0]
+            top_tokens, tgt_p, tgt_r, ent = _logits_to_pred(logits, tokenizer, target_id, top_k)
+            result.layers.append(LayerPrediction(
+                layer=idx, layer_name=rec.layer_name,
+                top_tokens=top_tokens,
+                target_prob=tgt_p, target_rank=tgt_r,
+                entropy=ent,
+            ))
+        return result