archscope 0.2.2__py3-none-any.whl

archscope/bench.py

@@ -0,0 +1,262 @@
+"""InterpBench — standardized mechanistic interpretability benchmark.
+
+Run a fixed test suite on any model to get a comparable "interp profile".
+Designed to support cross-architecture comparisons (transformer, hybrid, SSM).
+
+Tests:
+1. probe_sentiment_auroc   — Can we linearly probe pos/neg sentiment?
+2. probe_math_auroc        — Can we linearly probe math vs non-math?
+3. induction_head          — Does the model copy A->B in-context?
+4. copy_circuit            — Can it copy a word list verbatim?
+5. concentration           — How peaked are next-token predictions?
+6. sae_dictionary_quality  — Dense SAE reconstruction at mid-layer
+7. ssm_state_info          — (Mamba only) variance ratio of SSM state across inputs
+
+Output: dataclass `InterpProfile` with all scores, JSON-serializable.
+A model's "InterpProfile" is its interp signature.
+"""
+from __future__ import annotations
+from dataclasses import dataclass, field
+import time
+import torch
+
+from . import probes, sae, circuits
+from .backends import Backend
+
+
+@dataclass
+class InterpProfile:
+    """Standardized interp test results for one model."""
+    model_name: str
+    arch_family: str            # "transformer" | "hybrid" | "ssm" | "custom"
+    n_params: int
+    n_layers: int
+    hidden_dim: int
+
+    # Probe tasks (in-arch AUROC at representative mid-layer)
+    probe_sentiment_auroc: float = float("nan")
+    probe_math_auroc: float = float("nan")
+
+    # Circuit detection (relative to baseline)
+    induction_head_relative: float = float("nan")
+    copy_accuracy: float = float("nan")
+    concentration_relative: float = float("nan")
+
+    # SAE quality (dense SAE recon at mid-layer)
+    sae_dense_recon: float = float("nan")
+    sae_rank1_recon: float = float("nan")
+    sae_better: str = "—"   # "dense" or "rank1"
+
+    # SSM-specific (NaN for non-SSM)
+    ssm_state_variance_ratio: float = float("nan")
+
+    runtime_seconds: float = 0.0
+    notes: list = field(default_factory=list)
+
+
+# ---------- Probe tasks ----------
+
+SENTIMENT_POS = [
+    "I love this movie, it's amazing!", "What a wonderful day today.",
+    "Best book I have ever read.", "She is so kind and thoughtful.",
+    "The food was absolutely delicious.", "I'm thrilled about the great news.",
+    "Such a beautiful sunset tonight.", "He always makes me laugh.",
+    "Amazing performance overall.", "Truly a delightful experience.",
+    "Incredible work, congratulations!", "I admire her dedication.",
+    "The concert was breathtaking.", "He is the kindest person I know.",
+    "This was pure joy from start.", "I had a wonderful evening.",
+]
+SENTIMENT_NEG = [
+    "I hate this place absolutely.", "What a terrible movie that was.",
+    "This is the worst day ever.", "She is mean and selfish.",
+    "The food was awful and cold.", "I'm so disappointed with everything.",
+    "Such a horrible meeting today.", "He always annoys me deeply.",
+    "What an awful experience overall.", "I despise everything about this.",
+    "This restaurant is dreadful.", "She is the rudest person here.",
+    "The concert was a disaster.", "He never listens to anyone properly.",
+    "Truly a miserable evening.", "I regret coming here entirely.",
+]
+MATH = [
+    "Solve for x: 2x + 3 = 11.", "Compute the derivative of x cubed.",
+    "The integral from 0 to 1 of x dx.", "Triangle with angles 30 60 90.",
+    "The eigenvalue of matrix M is.", "By chain rule d/dx of sin x squared.",
+    "Find roots of x squared minus 5x plus 6.", "The Taylor series of e to x.",
+    "Prove sum of two evens is even.", "The Cauchy-Schwarz inequality.",
+    "Let f be continuous on closed interval.", "Define the limit as n grows large.",
+    "By induction on the natural numbers.", "The dot product of u and v.",
+    "A group is abelian if commutative.", "Fundamental theorem of calculus.",
+]
+NONMATH = [
+    "The cat sat on the mat softly.", "Music has the power to move us.",
+    "Mountains stretch to the horizon.", "She whispered her secret carefully.",
+    "Birds sing at dawn every morning.", "The chef prepared dinner slowly.",
+    "Children laughed in the city park.", "Rain pattered against the window.",
+    "Coffee aroma filled the kitchen.", "Time passes faster when busy.",
+    "Stars appeared in the dark sky.", "The river flowed gently downstream.",
+    "Snow blanketed the entire valley.", "Dance is a universal language.",
+    "Books contain endless worlds.", "Travel broadens one's perspective.",
+]
+
+
+def _run_probe(model, backend_hint, tokenize_fn, pos_texts, neg_texts, layer_name):
+    inputs_pos = tokenize_fn(pos_texts)
+    inputs_neg = tokenize_fn(neg_texts)
+    pf = probes.fit_probe(
+        model, inputs_pos=inputs_pos, inputs_neg=inputs_neg,
+        layer_name=layer_name, backend_hint=backend_hint,
+    )
+    # Use train AUROC (val with tiny n is too noisy)
+    return float(pf.metrics["train_auroc"])
+
+
+def _run_sae(model, backend, layer_name, hidden_dim, tokenize_fn):
+    """Train both Dense and Rank-1 SAE, return both recon errors."""
+    corpus = SENTIMENT_POS + SENTIMENT_NEG + MATH + NONMATH
+    inputs = tokenize_fn(corpus)
+    with torch.no_grad():
+        rec = backend.extract(inputs, layers=[layer_name])[0]
+    acts = rec.activations.reshape(-1, hidden_dim).detach()
+    cfg_d = sae.SAEConfig(input_dim=hidden_dim, n_features=512, sae_type="dense",
+                           sparsity=1e-4, learning_rate=3e-3)
+    cfg_r = sae.SAEConfig(input_dim=hidden_dim, n_features=512, sae_type="rank1",
+                           sparsity=1e-4, learning_rate=3e-3)
+    sae_d = sae.fit_sae(acts, cfg_d, epochs=60)
+    sae_r = sae.fit_sae(acts, cfg_r, epochs=60)
+    return sae_d.last_metrics["recon"], sae_r.last_metrics["recon"]
+
+
+def _ssm_state_variance_ratio(model, backend, layer_name, tokenize_fn) -> float:
+    """For Mamba: how much variance does the SSM state have across diverse inputs?
+
+    Ratio = variance across inputs / total variance. High = state encodes input strongly.
+    """
+    corpus = SENTIMENT_POS[:8] + MATH[:8]   # diverse inputs
+    inputs = tokenize_fn(corpus)
+    rec = backend.extract(inputs, layers=[layer_name])[0]
+    # rec.activations shape: (B, intermediate, ssm_state)
+    flat = rec.activations.reshape(rec.activations.shape[0], -1)
+    var_across = flat.var(dim=0).mean().item()
+    var_total = flat.var().item()
+    return float(var_across / (var_total + 1e-9))
+
+
+# ---------- Main runner ----------
+
+def benchmark(
+    model_name: str,
+    model,
+    tokenizer,
+    backend_hint: str,
+    arch_family: str = "transformer",
+    tokenize_fn=None,
+    sentiment_layer: int | None = None,
+    math_layer: int | None = None,
+    sae_layer: int | None = None,
+    ssm_layer: int | None = None,
+) -> InterpProfile:
+    """Run InterpBench on a model. Returns InterpProfile.
+
+    `tokenize_fn(texts: list[str]) -> dict` is the model-specific tokenizer wrapper.
+    """
+    backend = Backend.for_model(model, hint=backend_hint)
+    t_start = time.time()
+
+    len(backend.layer_names())
+    # Filter to only residual layers if mixed (e.g., Mamba has both .residual and .ssm_state)
+    residual_layers = [n for n in backend.layer_names() if ".residual" in n]
+    n_blocks = len(residual_layers)
+    hidden_dim = backend.hidden_dim(residual_layers[0])
+    n_params = sum(p.numel() for p in model.parameters())
+
+    profile = InterpProfile(
+        model_name=model_name,
+        arch_family=arch_family,
+        n_params=n_params,
+        n_layers=n_blocks,
+        hidden_dim=hidden_dim,
+    )
+
+    # Choose representative layers: sentiment probes at shallow depth, math
+    # and SAE at mid-depth. Each can be overridden by the caller.
+    if sentiment_layer is None:
+        sentiment_layer = n_blocks // 4
+    if math_layer is None:
+        math_layer = n_blocks // 2
+    if sae_layer is None:
+        sae_layer = n_blocks // 2
+
+    try:
+        profile.probe_sentiment_auroc = _run_probe(
+            model, backend_hint, tokenize_fn,
+            SENTIMENT_POS, SENTIMENT_NEG,
+            f"layer_{sentiment_layer}.residual",
+        )
+    except Exception as e:
+        profile.notes.append(f"probe_sentiment error: {str(e)[:60]}")
+
+    try:
+        profile.probe_math_auroc = _run_probe(
+            model, backend_hint, tokenize_fn,
+            MATH, NONMATH,
+            f"layer_{math_layer}.residual",
+        )
+    except Exception as e:
+        profile.notes.append(f"probe_math error: {str(e)[:60]}")
+
+    # Circuit detection
+    try:
+        circs = circuits.run_all_circuits(model, tokenizer=tokenizer, device="cpu")
+        profile.induction_head_relative = circs["induction_head"].relative
+        if "copy_circuit" in circs:
+            profile.copy_accuracy = circs["copy_circuit"].score
+        if "early_token_concentration" in circs:
+            profile.concentration_relative = circs["early_token_concentration"].relative
+    except Exception as e:
+        profile.notes.append(f"circuits error: {str(e)[:60]}")
+
+    # SAE quality
+    try:
+        d_recon, r_recon = _run_sae(
+            model, backend, f"layer_{sae_layer}.residual", hidden_dim, tokenize_fn,
+        )
+        profile.sae_dense_recon = d_recon
+        profile.sae_rank1_recon = r_recon
+        profile.sae_better = "rank1" if r_recon < d_recon else "dense"
+    except Exception as e:
+        profile.notes.append(f"sae error: {str(e)[:60]}")
+
+    # SSM state (mamba/ssm only)
+    if arch_family == "ssm" and ssm_layer is not None:
+        try:
+            profile.ssm_state_variance_ratio = _ssm_state_variance_ratio(
+                model, backend, f"layer_{ssm_layer}.ssm_state", tokenize_fn,
+            )
+        except Exception as e:
+            profile.notes.append(f"ssm_state error: {str(e)[:60]}")
+
+    profile.runtime_seconds = time.time() - t_start
+    return profile
+
+
+def profile_to_markdown(profile: InterpProfile) -> str:
+    """Format profile as markdown row for the leaderboard."""
+    lines = [
+        f"### {profile.model_name}",
+        f"  Arch: {profile.arch_family} | Params: {profile.n_params/1e6:.1f}M | "
+        f"Layers: {profile.n_layers} | Hidden: {profile.hidden_dim}",
+        "  | Test | Score |",
+        "  |------|-------|",
+        f"  | Sentiment probe AUROC    | {profile.probe_sentiment_auroc:.3f} |",
+        f"  | Math probe AUROC          | {profile.probe_math_auroc:.3f} |",
+        f"  | Induction head (×chance)  | {profile.induction_head_relative:>.1f} |",
+        f"  | Copy accuracy             | {profile.copy_accuracy:.2%} |",
+        f"  | Concentration (rel)       | {profile.concentration_relative:.3f} |",
+        f"  | SAE Dense recon           | {profile.sae_dense_recon:.4f} |",
+        f"  | SAE Rank-1 recon          | {profile.sae_rank1_recon:.4f} |",
+        f"  | SAE better                | {profile.sae_better} |",
+        f"  | SSM state var ratio       | {profile.ssm_state_variance_ratio:.3f} |",
+        f"  | Runtime                   | {profile.runtime_seconds:.1f}s |",
+    ]
+    if profile.notes:
+        lines.append(f"  Notes: {'; '.join(profile.notes)}")
+    return "\n".join(lines)

archscope/circuits.py

@@ -0,0 +1,255 @@
+"""Cross-architecture circuit detection.
+
+Classical mech-interp circuits were discovered in transformers (induction heads,
+attention sinks, copy circuits). It's an OPEN QUESTION whether they exist in
+SSMs (Mamba) or hybrid attention (Kazdov-α).
+
+Each detector is BEHAVIORAL — measures the model's outputs given crafted
+inputs — so it works on ANY architecture without per-arch internals.
+
+Detectors implemented:
+- induction_head_score: does the model copy A->B associations? (Olsson et al 2022)
+- copy_score: can the model copy a sequence verbatim?
+- early_token_attention: proxy for attention-sink behavior, measures prob mass
+  concentrated on early tokens (BOS-like).
+"""
+from __future__ import annotations
+import random
+import torch
+import torch.nn.functional as F
+from dataclasses import dataclass
+
+
+@dataclass
+class CircuitScore:
+    """Score from one circuit test."""
+    name: str
+    score: float                # primary metric, higher = circuit more present
+    baseline: float             # random-model baseline
+    relative: float             # score / baseline (>1 means circuit present)
+    raw: dict                   # detailed numbers
+
+
+# ---------- Helper: model-agnostic forward ----------
+
+def _model_logits(model, input_ids: torch.Tensor) -> torch.Tensor:
+    """Call model and extract final-step logits. Works for HF + kazdov."""
+    with torch.no_grad():
+        out = model(input_ids)
+    if isinstance(out, dict):
+        logits = out.get("logits", out.get("last_hidden_state"))
+    elif hasattr(out, "logits"):
+        logits = out.logits
+    else:
+        logits = out
+    return logits
+
+
+# ---------- Induction head detection ----------
+
+def induction_head_score(
+    model,
+    n_pairs: int = 20,
+    seq_len: int = 6,
+    n_trials: int = 50,
+    vocab_size: int | None = None,
+    device: str = "cpu",
+    seed: int = 0,
+) -> CircuitScore:
+    """Olsson-style induction head test.
+
+    Construct a sequence: [A1 B1] [A2 B2] ... [A_k B_k] [A_1]
+    Score = probability that next token is B_1 (the original pair completion).
+
+    If induction heads exist, the model should learn to copy A_1 → B_1 in-context.
+    Compared to random-token baseline (1/vocab_size).
+    """
+    rng = random.Random(seed)
+    if vocab_size is None:
+        # Try to infer (HF transformers + custom models like kazdov)
+        if hasattr(model, "config") and hasattr(model.config, "vocab_size"):
+            vocab_size = model.config.vocab_size
+        elif hasattr(model, "vocab_size"):
+            vocab_size = model.vocab_size
+        else:
+            vocab_size = 50257   # GPT-2 default
+
+    successes = 0
+    rank_sum = 0.0
+    prob_target_sum = 0.0
+    for trial in range(n_trials):
+        # Pick n_pairs random token pairs
+        tokens = rng.sample(range(100, min(vocab_size, 40000)), 2 * n_pairs)
+        seq = []
+        pairs = []
+        for i in range(n_pairs):
+            a, b = tokens[2*i], tokens[2*i+1]
+            seq.extend([a, b])
+            pairs.append((a, b))
+        # Append cue: A1 — model should predict B1
+        a1, b1 = pairs[0]
+        seq.append(a1)
+
+        ids = torch.tensor([seq], dtype=torch.long, device=device)
+        logits = _model_logits(model, ids)
+        # last position prediction
+        last_logits = logits[0, -1, :]
+        probs = F.softmax(last_logits, dim=-1)
+
+        # Did the model predict b1 as top-1?
+        pred = int(torch.argmax(last_logits).item())
+        if pred == b1:
+            successes += 1
+        # Rank of b1 in the logit distribution (0 = predicted top-1).
+        rank = int((last_logits > last_logits[b1]).sum().item())
+        rank_sum += rank
+        prob_target_sum += float(probs[b1].item())
+
+    accuracy = successes / n_trials
+    avg_rank = rank_sum / n_trials
+    avg_prob = prob_target_sum / n_trials
+    chance = 1.0 / vocab_size
+    return CircuitScore(
+        name="induction_head",
+        score=avg_prob,
+        baseline=chance,
+        relative=avg_prob / chance,
+        raw={
+            "accuracy_top1": accuracy,
+            "avg_rank_target": avg_rank,
+            "avg_prob_target": avg_prob,
+            "n_trials": n_trials,
+            "seq_len_pairs": n_pairs,
+        },
+    )
+
+
+# ---------- Copy circuit detection ----------
+
+def copy_score(
+    model,
+    tokenizer,
+    n_trials: int = 30,
+    n_words: int = 5,
+    device: str = "cpu",
+    seed: int = 0,
+) -> CircuitScore:
+    """Test whether model can copy a short word-list verbatim after a separator.
+
+    Input pattern: "list: A B C D E. list: " → predict "A"
+                   then after predicting A, predict B, etc.
+
+    Score = fraction of correctly-copied tokens.
+    """
+    rng = random.Random(seed)
+    word_pool = [
+        "cat", "dog", "tree", "moon", "star", "fish", "bird", "rock",
+        "leaf", "wave", "sky", "rain", "snow", "wind", "fire", "ice",
+        "book", "lamp", "key", "door", "wall", "road", "hill", "sand",
+    ]
+
+    correct = 0
+    total = 0
+    for trial in range(n_trials):
+        words = rng.sample(word_pool, n_words)
+        prompt = f"list: {' '.join(words)}. list: "
+        ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+        ids.shape[1]
+
+        # Token IDs for target words (first token of each)
+        target_tokens = []
+        for w in words:
+            target_tokens.append(tokenizer(" " + w, add_special_tokens=False).input_ids[0])
+
+        # Autoregressively predict n_words tokens, chaining the model's own
+        # predictions (not teacher-forcing) — measures cumulative copy ability.
+        cur = ids.clone()
+        for tgt in target_tokens:
+            logits = _model_logits(model, cur)
+            next_tok = int(torch.argmax(logits[0, -1, :]).item())
+            if next_tok == tgt:
+                correct += 1
+            total += 1
+            cur = torch.cat([cur, torch.tensor([[next_tok]], device=device)], dim=1)
+
+    acc = correct / total if total > 0 else 0.0
+    # Random baseline for copying is ~1/vocab_size; effectively ~0
+    vocab_size = (getattr(model, "vocab_size", None)
+                  or getattr(getattr(model, "config", None), "vocab_size", 50257))
+    chance = 1.0 / vocab_size
+    return CircuitScore(
+        name="copy_circuit",
+        score=acc,
+        baseline=chance,
+        relative=acc / chance if chance > 0 else float("inf"),
+        raw={"n_trials": n_trials, "n_words": n_words, "correct": correct, "total": total},
+    )
+
+
+# ---------- Early-token attention proxy (attention sink) ----------
+
+def early_token_attention(
+    model,
+    tokenizer,
+    texts: list[str] | None = None,
+    device: str = "cpu",
+) -> CircuitScore:
+    """Behavioral proxy for "attention sink" — measures how much logit mass
+    concentrates on the first token's vocab prediction when given a continuation cue.
+
+    Construct: "[BOS] X1 X2 ... X_n [predict ?]"
+    Measure: entropy of next-token distribution.
+
+    Low entropy → model is highly concentrated (possible sink-like behavior).
+    High entropy → no concentration.
+
+    NOTE: this is a coarse behavioral proxy. True attention-sink analysis
+    requires architecture-specific attention weights. But the proxy works
+    on ANY architecture (including SSMs that don't have attention).
+    """
+    if texts is None:
+        texts = [
+            "The cat sat on the", "Music has the power to",
+            "Mountains stretch to the", "She wrote a letter to",
+            "The sun set behind the", "Children laughed in the",
+            "Solve for x in the equation", "Compute the derivative of the",
+        ]
+    entropies = []
+    for txt in texts:
+        ids = tokenizer(txt, return_tensors="pt").input_ids.to(device)
+        logits = _model_logits(model, ids)
+        last_logits = logits[0, -1, :]
+        probs = F.softmax(last_logits, dim=-1)
+        # Shannon entropy (in nats)
+        ent = -(probs * (probs.clamp(min=1e-12).log())).sum().item()
+        entropies.append(ent)
+
+    avg_ent = sum(entropies) / len(entropies)
+    # Reference: log(vocab_size) is max entropy
+    vocab_size = (getattr(model, "vocab_size", None)
+                  or getattr(getattr(model, "config", None), "vocab_size", 50257))
+    max_ent = torch.log(torch.tensor(float(vocab_size))).item()
+    return CircuitScore(
+        name="early_token_concentration",
+        score=avg_ent,             # in nats
+        baseline=max_ent,
+        relative=avg_ent / max_ent,  # 0 = full concentration, 1 = uniform
+        raw={"per_text_entropy": entropies, "max_entropy": max_ent},
+    )
+
+
+# ---------- Runner ----------
+
+def run_all_circuits(model, tokenizer=None, device: str = "cpu") -> dict[str, CircuitScore]:
+    """Run all available circuit tests on a model.
+
+    Returns a dict keyed by circuit name. Some tests require tokenizer; if not
+    provided, those are skipped.
+    """
+    results = {}
+    # Induction head only needs model + vocab_size
+    results["induction_head"] = induction_head_score(model, device=device)
+    if tokenizer is not None:
+        results["copy_circuit"] = copy_score(model, tokenizer, device=device)
+        results["early_token_concentration"] = early_token_attention(model, tokenizer, device=device)
+    return results

archscope/cli.py

@@ -0,0 +1,120 @@
+"""Command-line interface for archscope."""
+from __future__ import annotations
+
+import json
+import os
+
+import click
+from rich.console import Console
+from rich.table import Table
+
+from . import __version__
+
+console = Console()
+
+
+@click.group()
+def cli() -> None:
+    """archscope — cross-architecture mechanistic interpretability toolkit."""
+
+
+@cli.command()
+def info() -> None:
+    """Show available methods + supported backends."""
+    methods = Table(title=f"archscope v{__version__}")
+    methods.add_column("Method")
+    methods.add_column("Module")
+    methods.add_column("Source paper")
+    for row in (
+        ("Probes",              "probes.fit_probe",           "Drop the Act (2605.11467)"),
+        ("SAE",                 "sae.fit_sae",                "WriteSAE (2605.12770)"),
+        ("Neuron mod",          "neurons.find_neurons",       "Targeted Neuron Mod (2605.12290)"),
+        ("Activation patch",    "attribute.activation_patch", "Multi-Agent Sycophancy (2605.12991)"),
+        ("Cross-arch transfer", "transfer.evaluate_transfer", "this library"),
+        ("Circuit detection",   "circuits.run_all_circuits",  "this library"),
+        ("Logit/tuned lens",    "lens.logit_lens",            "Belrose et al 2023"),
+        ("Model diff",          "diff.compare",               "this library"),
+        ("InterpBench",         "bench.benchmark",            "this library"),
+    ):
+        methods.add_row(*row)
+    console.print(methods)
+
+    backends = Table(title="Backends")
+    backends.add_column("Name")
+    backends.add_column("Architecture family")
+    for row in (
+        ("transformer", "HuggingFace decoder LMs (Llama, GPT, Qwen, Pythia, ...)"),
+        ("mamba",       "Mamba / Mamba-2 SSM — exposes .ssm_state (recurrent h_t)"),
+        ("kazdov",      "Kazdov-α hybrid MoBE-BCN+MHA"),
+        ("recurrent",   "Generic RNN (subclass per model)"),
+    ):
+        backends.add_row(*row)
+    console.print(backends)
+
+
+@cli.command()
+@click.argument("model_name")
+@click.option("--arch", default="transformer",
+              type=click.Choice(["transformer", "mamba", "kazdov"]),
+              help="Architecture family.")
+@click.option("--out", default=None,
+              help="Output file. Format inferred from extension: .json or .md. "
+                   "Without --out, prints markdown to stdout.")
+def bench(model_name: str, arch: str, out: str | None) -> None:
+    """Run archscope InterpBench on a HuggingFace model.
+
+    Examples:
+      archscope bench EleutherAI/pythia-160m --arch transformer
+      archscope bench EleutherAI/pythia-160m --arch transformer --out pythia.md
+      archscope bench state-spaces/mamba-130m-hf --arch mamba --out mamba.json
+    """
+    # Lazy imports keep `archscope info` fast (no torch/transformers).
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    from . import bench as bench_mod
+
+    console.print(f"[cyan]→ Loading {model_name}…[/cyan]")
+    tok = AutoTokenizer.from_pretrained(model_name)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+    model = AutoModelForCausalLM.from_pretrained(model_name, dtype=torch.float32)
+    model.eval()
+
+    def tokenize_fn(texts):
+        return tok(texts, return_tensors="pt", padding=True, truncation=True, max_length=32)
+
+    arch_family = {"transformer": "transformer", "mamba": "ssm", "kazdov": "hybrid"}[arch]
+    profile = bench_mod.benchmark(
+        model_name=model_name,
+        model=model,
+        tokenizer=tok,
+        backend_hint=arch,
+        arch_family=arch_family,
+        tokenize_fn=tokenize_fn,
+    )
+
+    markdown = bench_mod.profile_to_markdown(profile)
+    console.print()
+    console.print(markdown)
+
+    if out:
+        from dataclasses import asdict
+        os.makedirs(os.path.dirname(out) or ".", exist_ok=True)
+        ext = os.path.splitext(out)[1].lower()
+        if ext == ".md":
+            with open(out, "w") as f:
+                f.write(markdown + "\n")
+            console.print(f"\n[green]Saved markdown report to {out}[/green]")
+        elif ext == ".json" or ext == "":
+            with open(out, "w") as f:
+                json.dump(asdict(profile), f, indent=2, default=str)
+            console.print(f"\n[green]Saved JSON profile to {out}[/green]")
+        else:
+            raise click.UsageError(
+                f"Unsupported --out extension '{ext}'. Use .md or .json."
+            )
+
+
+if __name__ == "__main__":
+    cli()