logogram 0.1.0__py3-none-any.whl

logogram/__init__.py

@@ -0,0 +1,6 @@
+"""Logogram: a local-first workbench for causal experiments inside language models."""
+
+__version__ = "0.1.0"
+# The release date of this version: after a few months the app suggests looking for a newer one
+# (no network needed). Set it with every release.
+__released__ = "2026-10-08"

logogram/__main__.py

@@ -0,0 +1,5 @@
+"""python -m logogram"""
+
+from logogram.cli import main
+
+main()

logogram/analysis.py

@@ -0,0 +1,419 @@
+"""Interactive analyses: the token strip, the baseline check and attention patterns."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import torch
+
+from logogram.backends.base import ModelBackend
+from logogram.datasets import PromptRecord
+from logogram.engine import compute_baselines
+from logogram.prompts import PreparedPrompt, PromptIssue, group_by_length, prepare_prompt
+from logogram.spec import IndexPosition, PredictionSettings
+
+
+def _f(x: float) -> float | None:
+    x = float(x)
+    return x if np.isfinite(x) else None
+
+
+def tokenize_pair(backend: ModelBackend, record: PromptRecord, prepend_bos: bool) -> dict[str, Any]:
+    """Everything the token strip shows for one prompt pair, including what's wrong with it."""
+    prepared, issues, clean, corrupt = prepare_prompt(backend, record, 0, prepend_bos)
+    n = min(len(clean.ids), len(corrupt.ids))
+    differs = [i for i in range(n) if clean.ids[i] != corrupt.ids[i]]
+    differs += list(range(n, max(len(clean.ids), len(corrupt.ids))))
+
+    def answer_info(text: str) -> dict[str, Any]:
+        pieces = backend.tokenize(text, prepend_bos=False)
+        return {
+            "text": text,
+            "tokens": pieces.tokens,
+            "id": pieces.ids[0] if len(pieces.ids) == 1 else None,
+        }
+
+    return {
+        "clean": {"tokens": clean.tokens, "ids": clean.ids},
+        "corrupt": {"tokens": corrupt.tokens, "ids": corrupt.ids},
+        "aligned": len(clean.ids) == len(corrupt.ids),
+        "differs": differs,
+        "labels": prepared.labels if prepared else {},
+        "answer": answer_info(record.answer),
+        "distractor": answer_info(record.distractor),
+        "issues": [i.to_dict() for i in issues],
+    }
+
+
+def prepare_with_issues(
+    backend: ModelBackend, records: list[PromptRecord], prepend_bos: bool
+) -> tuple[list[PreparedPrompt], list[PromptIssue]]:
+    prepared, issues = [], []
+    for i, record in enumerate(records):
+        p, prompt_issues, _, _ = prepare_prompt(backend, record, i, prepend_bos)
+        issues.extend(prompt_issues)
+        if p is not None:
+            prepared.append(p)
+    return prepared, issues
+
+
+def baseline_report(
+    backend: ModelBackend,
+    records: list[PromptRecord],
+    *,
+    prepend_bos: bool = True,
+    batch_size: int = 64,
+    top_k: int = 5,
+) -> dict[str, Any]:
+    """Run clean and corrupt prompts unpatched: logit differences, probabilities, top tokens."""
+    prepared, issues = prepare_with_issues(backend, records, prepend_bos)
+    if not prepared:
+        return {"n": 0, "prompts": [], "issues": [i.to_dict() for i in issues], "summary": None}
+    groups = group_by_length(prepared)
+    base = compute_baselines(backend, prepared, groups, batch_size)
+    top: dict[str, list[list[dict[str, Any]]]] = {
+        "clean": [[] for _ in prepared],
+        "corrupt": [[] for _ in prepared],
+    }
+    for which in ("clean", "corrupt"):
+        for group in groups:
+            tokens = group.clean if which == "clean" else group.corrupt
+            for start in range(0, len(group.members), batch_size):
+                idx = group.members[start : start + batch_size]
+                logits = backend.final_logits(tokens[start : start + batch_size])
+                probs = torch.softmax(logits, dim=-1)
+                values, ids = probs.topk(top_k, dim=-1)
+                for row, p in enumerate(idx):
+                    top[which][p] = [
+                        {"token": backend.token_str(int(t)), "id": int(t), "prob": float(v)}
+                        for v, t in zip(values[row].tolist(), ids[row].tolist(), strict=True)
+                    ]
+    rows = []
+    for i, p in enumerate(prepared):
+        rows.append(
+            {
+                "index": p.index,
+                "clean": p.record.clean,
+                "corrupt": p.record.corrupt,
+                "answer": p.record.answer,
+                "distractor": p.record.distractor,
+                "clean_logit_diff": _f(base.clean_ld[i]),
+                "corrupt_logit_diff": _f(base.corrupt_ld[i]),
+                "clean_answer_prob": _f(base.clean_prob[i]),
+                "corrupt_answer_prob": _f(base.corrupt_prob[i]),
+                "clean_top": top["clean"][i],
+                "corrupt_top": top["corrupt"][i],
+            }
+        )
+    gap = base.clean_ld - base.corrupt_ld
+    summary = {
+        "clean_logit_diff": _f(base.clean_ld.mean()),
+        "corrupt_logit_diff": _f(base.corrupt_ld.mean()),
+        "gap": _f(gap.mean()),
+        "clean_prefers_answer": int((base.clean_ld > 0).sum()),
+        "corrupt_prefers_answer": int((base.corrupt_ld > 0).sum()),
+        "clean_answer_prob": _f(base.clean_prob.mean()),
+        "corrupt_answer_prob": _f(base.corrupt_prob.mean()),
+    }
+    return {
+        "n": len(prepared),
+        "prompts": rows,
+        "issues": [i.to_dict() for i in issues],
+        "summary": summary,
+    }
+
+
+def attention_report(
+    backend: ModelBackend,
+    records: list[PromptRecord],
+    *,
+    index: int,
+    layer: int,
+    head: int,
+    which: str = "clean",
+    prepend_bos: bool = True,
+    batch_size: int = 64,
+) -> dict[str, Any]:
+    """The head's attention on one prompt, and averaged over prompts of the same token length."""
+    info = backend.info
+    if not 0 <= layer < info.n_layers or not 0 <= head < info.n_heads:
+        raise ValueError(f"L{layer} H{head} doesn't exist in this model.")
+    if not 0 <= index < len(records):
+        raise ValueError(f"There is no prompt {index}.")
+    prepared, _ = prepare_with_issues(backend, records, prepend_bos)
+    target = next((p for p in prepared if p.index == index), None)
+    if target is None:
+        raise ValueError(f"Prompt {index} can't be used; fix it in the dataset first.")
+    same = [p for p in prepared if p.length == target.length]
+    tokens = torch.tensor(
+        [(p.clean if which == "clean" else p.corrupt).ids for p in same], dtype=torch.long
+    )
+    total = None
+    chosen = None
+    for start in range(0, len(same), batch_size):
+        pattern = backend.attention_pattern(tokens[start : start + batch_size], layer)[:, head]
+        pattern = pattern.double().cpu()
+        total = pattern.sum(0) if total is None else total + pattern.sum(0)
+        for row, p in enumerate(same[start : start + batch_size]):
+            if p.index == index:
+                chosen = pattern[row]
+    assert total is not None and chosen is not None
+    average = total / len(same)
+    tok = (target.clean if which == "clean" else target.corrupt).tokens
+    # The average mixes prompts, so only labels at the same index in all of them, and tokens
+    # identical in all of them, describe its positions.
+    seqs = [(p.clean if which == "clean" else p.corrupt).tokens for p in same]
+    common_tokens = [t if all(s[i] == t for s in seqs) else None for i, t in enumerate(tok)]
+    common_labels = {
+        label: pos
+        for label, pos in target.labels.items()
+        if all(p.labels.get(label) == pos for p in same)
+    }
+    return {
+        "layer": layer,
+        "head": head,
+        "which": which,
+        "index": index,
+        "tokens": tok,
+        "labels": target.labels,
+        "pattern": chosen.tolist(),
+        "average": average.tolist(),
+        "average_tokens": common_tokens,
+        "average_labels": common_labels,
+        "n_average": len(same),
+        "n_total": len(prepared),
+        "length": target.length,
+    }
+
+
+def prediction_report(
+    backend: ModelBackend,
+    records: list[PromptRecord],
+    *,
+    index: int,
+    settings: PredictionSettings,
+    prepend_bos: bool,
+    batch_size: int,
+) -> dict[str, Any]:
+    """A descriptive logit lens, preserving the baseline's length groups and batch shape."""
+    if index != settings.prompt_index:
+        raise ValueError("The prompt index must match the prediction settings in the spec.")
+    if not 0 <= index < len(records):
+        raise ValueError(f"There is no prompt {index}.")
+    prepared, issues = prepare_with_issues(backend, records, prepend_bos)
+    target_idx = next((i for i, p in enumerate(prepared) if p.index == index), None)
+    if target_idx is None:
+        raise ValueError(f"Prompt {index} can't be used; fix it in the dataset first.")
+    target = prepared[target_idx]
+    pos = settings.position.index if isinstance(settings.position, IndexPosition) else -1
+    pos = target.length + pos if pos < 0 else pos
+    if not 0 <= pos < target.length:
+        raise ValueError(f"Token position {pos} is outside this prompt ({target.length} tokens).")
+    group = next(g for g in group_by_length(prepared) if target_idx in g.members)
+    within = group.members.index(target_idx)
+    start = (within // batch_size) * batch_size
+    tokens = group.clean if settings.which == "clean" else group.corrupt
+    batch = tokens[start : start + batch_size]
+    logits = backend.layer_logits(batch, pos, within - start)
+    if not torch.isfinite(logits).all():
+        raise ValueError("The prediction diagnostic produced non-finite logits. Try float32.")
+    probs = torch.softmax(logits, dim=-1)
+    # Stable ordering makes ties deterministic, including toy or degenerate models.
+    top = torch.argsort(logits, dim=-1, descending=True, stable=True)[:, : settings.top_k]
+    answer = target.answer_id
+    distractor = target.distractor_id
+    rows = []
+    for layer in range(backend.info.n_layers):
+        rows.append(
+            {
+                "layer": layer,
+                "top": [
+                    {
+                        "id": int(t),
+                        "token": backend.token_str(int(t)),
+                        "prob": float(probs[layer, t]),
+                    }
+                    for t in top[layer]
+                ],
+                "answer_prob": float(probs[layer, answer]),
+                "distractor_prob": float(probs[layer, distractor]),
+                "logit_diff": float(logits[layer, answer] - logits[layer, distractor]),
+            }
+        )
+    return {
+        "index": index,
+        "position": pos,
+        "tokens": (target.clean if settings.which == "clean" else target.corrupt).tokens,
+        "settings": settings.model_dump(mode="json"),
+        "layers": rows,
+        "batch_members": [prepared[i].index for i in group.members[start : start + batch_size]],
+        "answer": target.record.answer,
+        "distractor": target.record.distractor,
+        "issues": [issue.to_dict() for issue in issues],
+    }
+
+
+# -- sparse autoencoder features -------------------------------------------------------------
+
+
+def _first_real_token(prepend_bos: bool) -> int:
+    """Where real tokens start: SAEs aren't trained on the beginning-of-sequence token, whose
+    activations are unlike any other, so fits and splices leave it alone."""
+    return 1 if prepend_bos else 0
+
+
+def sae_fit_report(
+    backend: ModelBackend,
+    sae: Any,
+    records: list[PromptRecord],
+    *,
+    prepend_bos: bool,
+    batch_size: int,
+) -> dict[str, Any]:
+    """How well the SAE fits these prompts: the variance of the activations it explains, how many
+    features fire per token, and what the logit difference becomes when the model runs on the
+    SAE's reconstruction instead of the activation (its error removed)."""
+    from logogram.sae import fit_on
+
+    prepared, issues = prepare_with_issues(backend, records, prepend_bos)
+    if not prepared:
+        raise ValueError("None of these prompts can be used with the loaded model.")
+    start = _first_real_token(prepend_bos)
+    key = (sae.site, sae.layer)
+    acts, clean, spliced = [], [], []
+
+    def splice(x: torch.Tensor) -> torch.Tensor:
+        out = x.float().clone()
+        f, stats = sae.encode(out[:, start:])
+        out[:, start:] = sae.decode(f, stats)
+        return out
+
+    for group in group_by_length(prepared):
+        for begin in range(0, len(group.members), batch_size):
+            idx = group.members[begin : begin + batch_size]
+            tokens = group.clean[begin : begin + batch_size]
+            acts.append(backend.capture(tokens, [key])[key][:, start:].float())
+            answers = torch.tensor([prepared[i].answer_id for i in idx])
+            distractors = torch.tensor([prepared[i].distractor_id for i in idx])
+            rows = torch.arange(len(idx))
+            plain = backend.final_logits(tokens)
+            edited = backend.edit_logits(tokens, sae.site, sae.layer, splice)
+            clean.append((plain[rows, answers] - plain[rows, distractors]).double())
+            spliced.append((edited[rows, answers] - edited[rows, distractors]).double())
+    flat = torch.cat([a.reshape(-1, a.shape[-1]) for a in acts])
+    fit = fit_on(sae, flat)
+    ld, ld_spliced = torch.cat(clean), torch.cat(spliced)
+    fit.update(
+        {
+            "logit_diff": _f(ld.mean()),
+            "spliced_logit_diff": _f(ld_spliced.mean()),
+            "n": len(prepared),
+            "skipped": len(issues),
+        }
+    )
+    sae.fit = fit
+    return fit
+
+
+def _single(backend: ModelBackend, records: list[PromptRecord], index: int, prepend_bos: bool):  # type: ignore[no-untyped-def]
+    prepared, _ = prepare_with_issues(backend, records, prepend_bos)
+    target = next((p for p in prepared if p.index == index), None)
+    if target is None:
+        raise ValueError(f"Prompt {index} can't be used; fix it in the dataset first.")
+    return prepared, target
+
+
+def token_features_report(
+    backend: ModelBackend,
+    sae: Any,
+    records: list[PromptRecord],
+    *,
+    index: int,
+    which: str,
+    prepend_bos: bool,
+    top_k: int = 8,
+) -> dict[str, Any]:
+    """The features that fire most on each token of one prompt."""
+    _, target = _single(backend, records, index, prepend_bos)
+    tokenized = target.clean if which == "clean" else target.corrupt
+    tokens = torch.tensor([tokenized.ids], dtype=torch.long)
+    key = (sae.site, sae.layer)
+    f, _ = sae.encode(backend.capture(tokens, [key])[key][0].float())
+    values, ids = f.topk(min(top_k, f.shape[-1]), dim=-1)
+    per_token = [
+        [
+            {"feature": int(i), "activation": float(v)}
+            for v, i in zip(values[p].tolist(), ids[p].tolist(), strict=True)
+            if v > 0
+        ]
+        for p in range(f.shape[0])
+    ]
+    return {
+        "index": index,
+        "which": which,
+        "tokens": tokenized.tokens,
+        "labels": target.labels,
+        "features": per_token,
+        "active": [int((f[p] > 0).sum()) for p in range(f.shape[0])],
+        "first_real_token": _first_real_token(prepend_bos),
+    }
+
+
+def feature_report(
+    backend: ModelBackend,
+    sae: Any,
+    records: list[PromptRecord],
+    *,
+    feature: int,
+    index: int,
+    which: str,
+    prepend_bos: bool,
+    batch_size: int,
+    top: int = 10,
+) -> dict[str, Any]:
+    """Where one feature fires: along one prompt's tokens, and on which prompts of the dataset
+    most strongly (all computed here, from the project's own prompts)."""
+    if not 0 <= feature < sae.d_sae:
+        raise ValueError(f"The SAE has {sae.d_sae} features; there is no feature {feature}.")
+    prepared, target = _single(backend, records, index, prepend_bos)
+    key = (sae.site, sae.layer)
+    start = _first_real_token(prepend_bos)
+    strongest: list[dict[str, Any]] = []
+    along: list[float] = []
+    for group in group_by_length(prepared):
+        tokens = group.clean if which == "clean" else group.corrupt
+        for begin in range(0, len(group.members), batch_size):
+            idx = group.members[begin : begin + batch_size]
+            f, _ = sae.encode(
+                backend.capture(tokens[begin : begin + batch_size], [key])[key].float()
+            )
+            column = f[..., feature].double().cpu()  # [B, pos]
+            for row, p in enumerate(idx):
+                prompt = prepared[p]
+                values = column[row]
+                if prompt.index == index:
+                    along = [float(v) for v in values]
+                best = int(values[start:].argmax()) + start if values.shape[0] > start else 0
+                seq = prompt.clean if which == "clean" else prompt.corrupt
+                strongest.append(
+                    {
+                        "index": prompt.index,
+                        "max": float(values[best]),
+                        "position": best,
+                        "token": seq.tokens[best],
+                        "text": prompt.record.clean if which == "clean" else prompt.record.corrupt,
+                    }
+                )
+    strongest.sort(key=lambda r: (-r["max"], r["index"]))
+    tokenized = target.clean if which == "clean" else target.corrupt
+    return {
+        "feature": feature,
+        "which": which,
+        "index": index,
+        "tokens": tokenized.tokens,
+        "activations": along,
+        "top": [r for r in strongest[:top] if r["max"] > 0],
+        "active_prompts": sum(1 for r in strongest if r["max"] > 0),
+        "n": len(strongest),
+    }

logogram/atp.py

@@ -0,0 +1,120 @@
+"""Attribution patching: estimate activation patching at every site from one gradient.
+
+Patching a site changes the logit difference by f(receiver with the source activation) - f(receiver).
+To first order that is (source activation - receiver activation) · the gradient of the logit
+difference with respect to the activation, at the receiver run. One forward pass on the source
+prompts and one forward and backward pass on the receiver prompts give the estimate for every site
+at once, so a sweep costs a few passes per batch instead of one patched run per site and prompt.
+
+The estimate is only first order. It misses saturation (in attention patterns, normalization and
+the final softmax) and can miss or even invert an effect, which is why the results say "estimated"
+and offer to verify the strongest sites with real patching.
+"""
+
+from __future__ import annotations
+
+import threading
+from collections.abc import Callable
+from typing import Any
+
+import numpy as np
+import torch
+
+from logogram.backends.base import Cancelled, ModelBackend
+from logogram.engine import (
+    EngineResult,
+    LayerFn,
+    ProgressFn,
+    _answer_tensors,
+    _chunks,
+    check_gap,
+    compute_baselines,
+)
+from logogram.prompts import PreparedPrompt, group_by_length
+from logogram.sites import ResolvedSite, expand_scope, resolve_position
+from logogram.spec import AllPositions, AttributionPatching, Spec
+
+
+def run_attribution_patching(
+    spec: Spec,
+    backend: ModelBackend,
+    prompts: list[PreparedPrompt],
+    *,
+    on_progress: ProgressFn | None = None,
+    on_layer: LayerFn | None = None,
+    cancel: threading.Event | None = None,
+    on_start: Callable[[list[ResolvedSite], dict[str, Any]], None] | None = None,
+    receiver_override: str | None = None,
+    source_override: str | None = None,
+) -> EngineResult:
+    exp = spec.experiment
+    assert isinstance(exp, AttributionPatching)
+    info = backend.info
+    batch_size = spec.execution.batch_size
+    sites, layout = expand_scope(spec, info, prompts)
+    if on_start is not None:
+        on_start(sites, layout)
+    groups = group_by_length(prompts)
+    baselines = compute_baselines(backend, prompts, groups, batch_size, cancel)
+    receiver, source = (
+        ("corrupt", "clean") if exp.direction == "clean_to_corrupt" else ("clean", "corrupt")
+    )
+    reference = source
+    receiver = receiver_override or receiver
+    source = source_override or source
+    warnings = check_gap(spec, baselines, prompts, receiver, reference)
+
+    n = len(prompts)
+    needed = list(dict.fromkeys((rs.kind, rs.layer) for rs in sites))
+    estimate = np.zeros((len(sites), n))
+    done = 0
+    for group in groups:
+        receiver_tokens = group.clean if receiver == "clean" else group.corrupt
+        source_tokens = group.clean if source == "clean" else group.corrupt
+        for sl in _chunks(len(group.members), batch_size):
+            if cancel is not None and cancel.is_set():
+                raise Cancelled()
+            idx = group.members[sl]
+            source_acts = backend.capture(source_tokens[sl], needed)
+            receiver_acts, grads = backend.gradients(
+                receiver_tokens[sl], *_answer_tensors(prompts, idx), needed
+            )
+            rows = torch.arange(len(idx))
+            for rs in sites:
+                key = (rs.kind, rs.layer)
+                diff = source_acts[key].double() - receiver_acts[key].double()
+                grad = grads[key].double()
+                if rs.kind == "head":
+                    diff, grad = diff[:, :, rs.head], grad[:, :, rs.head]
+                per_position = (diff * grad).sum(-1)  # [B, pos]
+                if isinstance(rs.site.position, AllPositions):
+                    values = per_position.sum(1)
+                else:
+                    positions = torch.tensor(
+                        [resolve_position(rs.site.position, prompts[p]) for p in idx]
+                    )
+                    values = per_position[rows, positions]
+                estimate[rs.index, idx] = values.cpu().numpy()
+            done += len(idx)
+            if on_progress is not None:
+                on_progress(done, n, info.n_layers - 1)
+
+    receiver_ld = baselines.ld(receiver)
+    result = EngineResult(
+        sites=sites,
+        layout=layout,
+        prompts=prompts,
+        baselines=baselines,
+        receiver=receiver,
+        reference=reference,
+        # The logit difference patching would give, to first order. Probabilities aren't estimated.
+        patched_ld=receiver_ld[None, :] + estimate,
+        patched_prob=np.full((len(sites), n), np.nan),
+        warnings=warnings,
+        measure="estimate",
+        delta=estimate,
+    )
+    if on_layer is not None:
+        for layer in sorted({rs.layer for rs in sites}):
+            on_layer(layer, [rs.index for rs in sites if rs.layer == layer], result)
+    return result

logogram/backends/__init__.py

@@ -0,0 +1,5 @@
+"""Model backends. Experiments talk to models only through ModelBackend."""
+
+from logogram.backends.base import BackendError, ModelBackend, ModelInfo, Patch, Tokenized
+
+__all__ = ["BackendError", "ModelBackend", "ModelInfo", "Patch", "Tokenized"]