logogram 0.1.0__py3-none-any.whl

logogram/runner.py

@@ -0,0 +1,396 @@
+"""Run a spec end to end and write its experiment folder. The GUI and the CLI both call this."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import platform
+import random
+import threading
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+from importlib.metadata import PackageNotFoundError, version
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+
+from logogram import __version__
+from logogram.backends.base import ModelBackend
+from logogram.datasets import parse_jsonl
+from logogram.engine import Cancelled, run_experiment
+from logogram.fileio import atomic_output, write_text_atomic
+from logogram.project import Project, now_iso
+from logogram.prompts import prepare_prompts
+from logogram.results import (
+    build_summary,
+    compute_stats,
+    results_table,
+    site_payload,
+    write_results,
+)
+from logogram.schema import Manifest
+from logogram.spec import ModelRef, SAERef, Spec
+from logogram.stats import resample_counts
+
+EventFn = Callable[[str, dict[str, Any]], None]
+ModelProvider = Callable[[ModelRef], ModelBackend]
+SAEProvider = Callable[[SAERef, ModelBackend], Any]
+
+
+class RunError(RuntimeError):
+    pass
+
+
+def configure_determinism(seed: int = 0) -> None:
+    """Make repeated runs on the same machine bit-identical."""
+    os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
+    torch.use_deterministic_algorithms(True)
+    random.seed(seed)
+    np.random.seed(seed % (2**32))
+    torch.manual_seed(seed)
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+
+
+def _pkg_version(name: str) -> str | None:
+    try:
+        return version(name)
+    except PackageNotFoundError:
+        return None
+
+
+def environment_versions() -> dict[str, Any]:
+    return {
+        "logogram": __version__,
+        "python": platform.python_version(),
+        "torch": torch.__version__,
+        "transformer_lens": _pkg_version("transformer-lens"),
+        "transformers": _pkg_version("transformers"),
+        "numpy": _pkg_version("numpy"),
+        "pyarrow": _pkg_version("pyarrow"),
+    }
+
+
+def model_matches(backend: ModelBackend | None, ref: ModelRef) -> bool:
+    if backend is None:
+        return False
+    from logogram.backends.transformer_lens import resolve_device
+
+    info = backend.info
+    try:
+        device = resolve_device(ref.device)
+    except Exception:  # noqa: BLE001
+        return False
+    return (
+        info.id == ref.id
+        and (ref.revision is None or ref.revision == info.revision)
+        and info.dtype == ref.dtype
+        and info.device == device
+        and info.process_weights == ref.process_weights
+    )
+
+
+def default_provider(on_event: EventFn | None = None) -> ModelProvider:
+    def provide(ref: ModelRef) -> ModelBackend:
+        from logogram.backends.transformer_lens import load_model
+
+        return load_model(
+            ref.id,
+            revision=ref.revision,
+            dtype=ref.dtype,
+            device=ref.device,
+            process_weights=ref.process_weights,
+            on_progress=(lambda p: on_event("model", p)) if on_event else None,
+        )
+
+    return provide
+
+
+@dataclass
+class RunOutcome:
+    run_id: str
+    folder: Path
+    status: str
+    summary: dict[str, Any] | None
+    manifest: dict[str, Any]
+    backend: ModelBackend | None
+
+
+def write_manifest(path: Path, manifest: dict[str, Any]) -> None:
+    write_json(path, Manifest.model_validate(manifest).model_dump(mode="json"))
+
+
+def write_json(path: Path, data: Any) -> None:
+    """Write atomically, so readers (the history panel) never see a half-written file."""
+    write_text_atomic(path, json.dumps(data, indent=2, ensure_ascii=False) + "\n")
+
+
+def run_spec(
+    spec: Spec,
+    project: Project,
+    *,
+    run_id: str | None = None,
+    backend: ModelBackend | None = None,
+    provider: ModelProvider | None = None,
+    on_event: EventFn | None = None,
+    cancel: threading.Event | None = None,
+    derived_from: dict[str, Any] | None = None,
+    sae: Any = None,
+    sae_provider: SAEProvider | None = None,
+) -> RunOutcome:
+    """Execute ``spec`` and write ``experiments/<run_id>/`` in ``project``.
+
+    ``sae`` is an already loaded SAE, used when it is the one the spec names; otherwise
+    ``sae_provider`` (by default, a download from Hugging Face) loads it."""
+    emit = on_event or (lambda kind, data: None)
+    configure_determinism(spec.statistics.seed)
+    run_id = run_id or project.new_run_id(spec.name)
+    folder = project.prepare_run_dir(run_id)
+    started_at = now_iso()
+    t0 = time.perf_counter()
+
+    manifest: dict[str, Any] = {
+        "logogram_manifest": 1,
+        "run_id": run_id,
+        "status": "running",
+        "started_at": started_at,
+        "versions": environment_versions(),
+        "derived_from": derived_from,
+    }
+    # The spec as executed: written first so a failed run still records what was attempted.
+    write_json(folder / "spec.json", spec.model_dump(mode="json"))
+    write_manifest(folder / "manifest.json", manifest)
+
+    def fail(status: str, message: str | None) -> RunOutcome:
+        manifest.update(
+            {
+                "status": status,
+                "error": message,
+                "finished_at": now_iso(),
+                "wall_time_s": round(time.perf_counter() - t0, 3),
+            }
+        )
+        for name in ("results.parquet", "summary.json", "predictions.json"):
+            (folder / name).unlink(missing_ok=True)
+        write_manifest(folder / "manifest.json", manifest)
+        return RunOutcome(run_id, folder, status, None, manifest, backend)
+
+    try:
+        dataset_path = project.resolve_dataset(spec.dataset.path)
+        # Hash, parse and snapshot the same bytes, even if the source is subsequently edited.
+        content = dataset_path.read_bytes()
+        sha = hashlib.sha256(content).hexdigest()
+        if spec.dataset.sha256 and sha != spec.dataset.sha256:
+            raise RunError(
+                f"{spec.dataset.path} has changed since this spec was written (its hash no longer "
+                "matches). Restore the original file, or remove dataset.sha256 to accept the new one."
+            )
+        records = parse_jsonl(content.decode("utf-8"), source=spec.dataset.path)
+        if spec.dataset.limit is not None:
+            records = records[: spec.dataset.limit]
+        source_path = spec.dataset.path
+        snapshots = project.writable(project.datasets_dir / "snapshots")
+        snapshots.mkdir(exist_ok=True)
+        snapshot = project.writable(snapshots / f"{sha}.jsonl")
+        if os.path.lexists(snapshot):
+            if project.readable(snapshot).read_bytes() != content:
+                raise RunError(
+                    "The pinned dataset snapshot has changed. Restore it from a backup before rerunning."
+                )
+        else:
+            with atomic_output(snapshot) as output:
+                output.write(content)
+        spec = spec.model_copy(
+            update={
+                "dataset": spec.dataset.model_copy(
+                    update={
+                        "path": snapshot.relative_to(project.root).as_posix(),
+                        "sha256": sha,
+                    }
+                )
+            }
+        )
+        write_json(folder / "spec.json", spec.model_dump(mode="json"))
+        if cancel is not None and cancel.is_set():
+            raise Cancelled()
+
+        if not model_matches(backend, spec.model):
+            emit("status", {"message": f"Loading {spec.model.id}"})
+            backend = (provider or default_provider(on_event))(spec.model)
+        assert backend is not None
+        info = backend.info
+
+        # Pin exactly what ran, so `logogram run` on this file reproduces it.
+        spec = spec.model_copy(
+            update={
+                "model": spec.model.model_copy(update={"revision": info.revision}),
+                "dataset": spec.dataset.model_copy(update={"sha256": sha}),
+            }
+        )
+        write_json(folder / "spec.json", spec.model_dump(mode="json"))
+        manifest.update(
+            {
+                "model": {
+                    "id": info.id,
+                    "revision": info.revision,
+                    "architecture": info.architecture,
+                    "process_weights": info.process_weights,
+                    "backend": info.backend,
+                    "backend_version": info.backend_version,
+                },
+                "device": {
+                    "type": info.device,
+                    "name": info.device_name,
+                    "cuda": torch.version.cuda if info.device == "cuda" else None,
+                    "threads": torch.get_num_threads(),
+                    "deterministic_algorithms": torch.are_deterministic_algorithms_enabled(),
+                },
+                "dtype": info.dtype,
+                "dataset": {
+                    "path": spec.dataset.path,
+                    "sha256": sha,
+                    "n": len(records),
+                    "source_path": source_path,
+                },
+            }
+        )
+        write_manifest(folder / "manifest.json", manifest)
+
+        if spec.sae is not None:
+            from logogram.sae import load_sae, sae_matches
+
+            if not sae_matches(sae, spec.sae):
+                emit("status", {"message": f"Loading the SAE {spec.sae.repo}"})
+                sae = (
+                    sae_provider(spec.sae, backend)
+                    if sae_provider is not None
+                    else load_sae(spec.sae, info.device, cancel=cancel)
+                )
+            spec = spec.model_copy(
+                update={"sae": spec.sae.model_copy(update={"revision": sae.revision})}
+            )
+            write_json(folder / "spec.json", spec.model_dump(mode="json"))
+            manifest["sae"] = {
+                "repo": sae.repo,
+                "path": sae.path,
+                "revision": sae.revision,
+                "format": sae.params.format,
+                "site": sae.site,
+                "layer": sae.layer,
+                "d_sae": sae.d_sae,
+            }
+            write_manifest(folder / "manifest.json", manifest)
+
+        prompts = prepare_prompts(backend, records, spec.tokenization.prepend_bos)
+        # Bootstrap resamples over the prompts a method measures (steering measures held-out
+        # prompts only), drawn once per run so every site shares them.
+        resamples: dict[int, np.ndarray] = {}
+
+        def counts_for(result: Any) -> np.ndarray:
+            n = len(result.prompts)
+            if n not in resamples:
+                resamples[n] = resample_counts(n, spec.statistics.bootstrap, spec.statistics.seed)
+            return resamples[n]
+
+        model_shape = {
+            "id": info.id,
+            "n_layers": info.n_layers,
+            "n_heads": info.n_heads,
+            "d_model": info.d_model,
+            "d_head": info.d_head,
+            "d_mlp": info.d_mlp,
+            "site_kinds": list(info.site_kinds),
+            "block_structure": info.extra.get("block_structure"),
+            "normalization": info.extra.get("normalization"),
+            "activation": info.extra.get("activation"),
+        }
+
+        if spec.predictions is not None:
+            from logogram.analysis import prediction_report
+
+            emit("status", {"message": "Measuring the configured layer predictions"})
+            with backend.lock:
+                predictions = prediction_report(
+                    backend,
+                    records,
+                    index=spec.predictions.prompt_index,
+                    settings=spec.predictions,
+                    prepend_bos=spec.tokenization.prepend_bos,
+                    batch_size=spec.execution.batch_size,
+                )
+            if cancel is not None and cancel.is_set():
+                raise Cancelled()
+            write_json(folder / "predictions.json", predictions)
+
+        def on_start(sites: list[Any], layout: dict[str, Any]) -> None:
+            emit(
+                "started",
+                {
+                    "run_id": run_id,
+                    "n_prompts": len(prompts),
+                    "model": model_shape,
+                    "layout": layout,
+                    "sites": [s.to_dict() for s in sites],
+                },
+            )
+
+        def on_layer(layer: int, indices: list[int], partial: Any) -> None:
+            # A finished site's statistics depend only on its own rows, so they are final.
+            stats = compute_stats(spec, partial, counts_for(partial), indices)
+            emit(
+                "layer",
+                {"run_id": run_id, "layer": layer, "sites": site_payload(partial, stats, indices)},
+            )
+
+        def on_progress(done: int, total: int, layer: int) -> None:
+            emit(
+                "progress",
+                {
+                    "run_id": run_id,
+                    "done": done,
+                    "total": total,
+                    "layer": layer,
+                    "memory": backend.memory_in_use() if backend else None,
+                    "elapsed_s": round(time.perf_counter() - t0, 2),
+                },
+            )
+
+        result = run_experiment(
+            spec,
+            backend,
+            prompts,
+            on_progress=on_progress,
+            on_layer=on_layer,
+            on_start=on_start,
+            cancel=cancel,
+            sae=sae,
+        )
+        stats = compute_stats(spec, result, counts_for(result))
+        summary = build_summary(spec, result, stats, run_id, model_shape)
+        write_results(folder / "results.parquet", results_table(result, stats))
+        write_json(folder / "summary.json", summary)
+        manifest.update(
+            {
+                "status": "finished",
+                "finished_at": now_iso(),
+                "wall_time_s": round(time.perf_counter() - t0, 3),
+            }
+        )
+        write_manifest(folder / "manifest.json", manifest)
+        emit("finished", {"run_id": run_id})
+        return RunOutcome(run_id, folder, "finished", summary, manifest, backend)
+    except Cancelled:
+        outcome = fail("cancelled", None)
+        emit("cancelled", {"run_id": run_id})
+        return outcome
+    except Exception as exc:  # noqa: BLE001 - recorded in the manifest and reported
+        message = str(exc) or type(exc).__name__
+        outcome = fail("failed", message)
+        emit("failed", {"run_id": run_id, "error": message})
+        return outcome

logogram/runs.py

@@ -0,0 +1,98 @@
+"""Read finished runs: the per-prompt evidence behind one site's number."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+from logogram.datasets import DatasetError, file_sha256, load_dataset
+from logogram.project import Project
+from logogram.results import read_site_rows
+from logogram.spec import Spec
+
+
+class RunNotReady(ValueError):
+    pass
+
+
+def load_summary(project: Project, run_id: str) -> dict[str, Any]:
+    path = project.run_dir(run_id) / "summary.json"
+    if not path.is_file():
+        raise RunNotReady("This run has no results yet.")
+    return project.read_json(path) or {}
+
+
+def site_detail(
+    project: Project, run_id: str, site_index: int, extremes: int = 3
+) -> dict[str, Any]:
+    folder = project.run_dir(run_id)
+    summary = load_summary(project, run_id)
+    sites = summary["sites"]
+    if not 0 <= site_index < len(sites):
+        raise RunNotReady(f"Site {site_index} isn't part of this run.")
+    site = sites[site_index]
+    spec = Spec.from_path(project.readable(folder / "spec.json"))
+    rows = read_site_rows(project.readable(folder / "results.parquet"), site_index)
+
+    records = None
+    dataset_changed = False
+    try:
+        path = project.resolve_dataset(spec.dataset.path)
+        records = load_dataset(path)
+        if spec.dataset.limit:
+            records = records[: spec.dataset.limit]
+        if spec.dataset.sha256 and file_sha256(path) != spec.dataset.sha256:
+            dataset_changed = True
+    except (DatasetError, ValueError):
+        records = None
+
+    effects = np.asarray(rows["effect"], dtype=np.float64)
+    prompts = []
+    for i, prompt_index in enumerate(rows["prompt"]):
+        record = records[prompt_index] if records and prompt_index < len(records) else None
+        receiver = rows["receiver_logit_diff"][i]
+        patched = rows["patched_logit_diff"][i]
+        prompts.append(
+            {
+                "index": prompt_index,
+                "clean": record.clean if record else None,
+                "corrupt": record.corrupt if record else None,
+                "answer": record.answer if record else None,
+                "distractor": record.distractor if record else None,
+                "effect": _f(rows["effect"][i]),
+                "delta": _f(rows["delta"][i]),
+                "patched_logit_diff": _f(patched),
+                "receiver_logit_diff": _f(receiver),
+                "reference_logit_diff": _f(rows["reference_logit_diff"][i]),
+                "patched_answer_prob": _f(rows["patched_answer_prob"][i]),
+                "receiver_answer_prob": _f(rows["receiver_answer_prob"][i]),
+                "flipped": bool((receiver > 0 > patched) or (receiver < 0 < patched)),
+            }
+        )
+    mean = site["effect"]["mean"] or 0.0
+    direction = 1.0 if mean >= 0 else -1.0
+    finite = np.where(np.isfinite(effects), effects, 0.0)
+    order = np.argsort(-direction * finite, kind="mergesort")
+    k = min(extremes, len(order))
+    return {
+        "run_id": run_id,
+        "site": site,
+        "receiver": summary["receiver"],
+        "reference": summary["reference"],
+        "metric": summary["metric"],
+        "statistics": summary["statistics"],
+        "prompts": prompts,
+        "strongest": [int(rows["prompt"][j]) for j in order[:k]],
+        "weakest": [int(rows["prompt"][j]) for j in order[::-1][:k]],
+        "dataset_changed": dataset_changed,
+        "dataset_available": records is not None,
+    }
+
+
+def _f(x: Any) -> float | None:
+    try:
+        value = float(x)
+    except (TypeError, ValueError):
+        return None
+    return value if np.isfinite(value) else None

logogram/sae.py

@@ -0,0 +1,161 @@
+"""Sparse autoencoders on a model's activations: encode, decode, edit one feature, and check fit.
+
+An SAE writes an activation x as a sparse sum of decoder directions: x = f · W_dec + b_dec + e,
+where f are the feature activations and e is what the SAE misses (its error). Editing a feature
+changes x by (new - old) · W_dec[i] and keeps the error as it was, so only that feature changes.
+Whether an SAE fits a loaded model at all depends on how it was trained (on which model, and with
+or without TransformerLens's weight processing), so Logogram measures the fit on the project's
+prompts instead of assuming it.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+from logogram.backends.saes import SAEParams
+from logogram.spec import SAERef
+
+LN_EPS = 1e-5
+
+
+@dataclass
+class SAE:
+    """A loaded SAE, its parameters on the model's device, and where it came from."""
+
+    repo: str
+    path: str
+    revision: str
+    params: SAEParams
+    device: str = "cpu"
+    fit: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def site(self) -> str:
+        return self.params.site
+
+    @property
+    def layer(self) -> int:
+        return self.params.layer
+
+    @property
+    def d_sae(self) -> int:
+        return self.params.d_sae
+
+    def to(self, device: str) -> SAE:
+        p = self.params
+        for name in ("W_enc", "b_enc", "W_dec", "b_dec", "threshold"):
+            value = getattr(p, name)
+            if value is not None:
+                setattr(p, name, value.to(device))
+        self.device = device
+        return self
+
+    def describe(self) -> dict[str, Any]:
+        p = self.params
+        return {
+            "repo": self.repo,
+            "path": self.path,
+            "revision": self.revision,
+            "site": p.site,
+            "layer": p.layer,
+            "d_in": p.d_in,
+            "d_sae": p.d_sae,
+            "activation": p.activation,
+            "k": p.k,
+            "normalize": p.normalize,
+            "format": p.format,
+            "note": p.note,
+            "fit": self.fit or None,
+        }
+
+    # -- the map --------------------------------------------------------------------------------
+
+    def _standardize(
+        self, x: torch.Tensor
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
+        if self.params.normalize != "layer_norm":
+            return x, None
+        mu = x.mean(dim=-1, keepdim=True)
+        centered = x - mu
+        std = centered.std(dim=-1, keepdim=True)
+        return centered / (std + LN_EPS), (mu, std)
+
+    def encode(self, x: torch.Tensor) -> tuple[torch.Tensor, Any]:
+        """Feature activations ``[..., d_sae]`` of activations ``[..., d_in]``, and what
+        decoding needs to restore a standardized input (or None)."""
+        p = self.params
+        x = x.float()
+        sae_in, stats = self._standardize(x)
+        if p.subtract_b_dec:
+            sae_in = sae_in - p.b_dec
+        pre = sae_in @ p.W_enc + p.b_enc
+        if p.activation == "relu":
+            f = torch.relu(pre)
+        elif p.activation == "jumprelu":
+            assert p.threshold is not None
+            f = torch.relu(pre) * (pre > p.threshold)
+        else:
+            assert p.k is not None
+            top = pre.topk(min(p.k, pre.shape[-1]), dim=-1, sorted=False)
+            f = torch.zeros_like(pre).scatter(-1, top.indices, torch.relu(top.values))
+        return f, stats
+
+    def decode(self, f: torch.Tensor, stats: Any) -> torch.Tensor:
+        p = self.params
+        out = f @ p.W_dec + p.b_dec
+        if stats is not None:
+            mu, std = stats
+            out = out * std + mu
+        return out
+
+    def feature_direction(self, feature: int, stats: Any = None) -> torch.Tensor:
+        """What one unit of the feature adds to the activation: its decoder row, scaled back by
+        each input's standard deviation for SAEs that standardize their inputs."""
+        row = self.params.W_dec[feature]
+        if stats is None:
+            return row
+        _, std = stats
+        return row * std
+
+    # -- fit ------------------------------------------------------------------------------------
+
+
+def fit_on(sae: SAE, activations: torch.Tensor) -> dict[str, float]:
+    """How well the SAE reconstructs ``activations`` ``[N, d_in]``: the fraction of variance it
+    explains (1 is perfect, 0 is no better than the mean), and how many features fire per token."""
+    x = activations.float()
+    f, stats = sae.encode(x)
+    x_hat = sae.decode(f, stats)
+    residual = ((x - x_hat) ** 2).sum()
+    total = ((x - x.mean(dim=0, keepdim=True)) ** 2).sum().clamp_min(1e-12)
+    return {
+        "variance_explained": float(1.0 - residual / total),
+        "l0": float((f > 0).float().sum(-1).mean()),
+        "tokens": int(x.shape[0]),
+    }
+
+
+def load_sae(
+    ref: SAERef,
+    device: str,
+    progress: Any = None,
+    cancel: Any = None,
+) -> SAE:
+    """Download (once) and read a published SAE at an exact revision, on ``device``."""
+    from logogram.backends.saes import download_sae, read_sae
+
+    revision, folder = download_sae(ref.repo, ref.path, ref.revision, progress, cancel)
+    params = read_sae(folder, ref.path)
+    return SAE(repo=ref.repo, path=ref.path, revision=revision, params=params).to(device)
+
+
+def sae_matches(sae: SAE | None, ref: SAERef) -> bool:
+    return (
+        sae is not None
+        and sae.repo == ref.repo
+        and sae.path.strip("/") == ref.path.strip("/")
+        and (ref.revision is None or ref.revision == sae.revision)
+    )