logogram 0.1.0__py3-none-any.whl

logogram/backends/base.py

@@ -0,0 +1,202 @@
+"""The interface between experiments and a model library.
+
+Experiments speak only in abstract sites (``resid_pre``, ``head`` ...) and tensors. A backend
+maps those sites onto its own hooks. TransformerLens is the only backend in v0.1; others (for
+example remote execution) can be added by implementing :class:`ModelBackend`.
+"""
+
+from __future__ import annotations
+
+import threading
+from abc import ABC, abstractmethod
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+import torch
+
+STREAM_KINDS = ("resid_pre", "resid_mid", "resid_post", "attn_out", "mlp_out")
+ALL_KINDS = (*STREAM_KINDS, "head")
+
+
+@dataclass(frozen=True)
+class ModelInfo:
+    id: str
+    revision: str | None
+    architecture: str
+    n_layers: int
+    n_heads: int
+    d_model: int
+    d_head: int
+    d_mlp: int | None
+    d_vocab: int
+    n_ctx: int
+    n_params: int | None
+    dtype: str
+    device: str  # "cpu", "cuda" or "mps"
+    device_name: str
+    process_weights: bool
+    site_kinds: tuple[str, ...]
+    backend: str
+    backend_version: str
+    extra: dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        data = asdict(self)
+        data["site_kinds"] = list(self.site_kinds)
+        return data
+
+
+@dataclass
+class Tokenized:
+    ids: list[int]
+    tokens: list[str]
+    offsets: list[tuple[int, int]]  # character span of each token; (0, 0) for added BOS
+
+
+@dataclass
+class Patch:
+    """Replace one site's activation, row by row, during a forward pass.
+
+    ``heads`` gives each row's head for head sites. ``positions`` gives each row's single
+    position, or is ``None`` to replace every position. ``values`` is ``[B, pos, d]`` when every
+    position is replaced and ``[B, d]`` otherwise (``d`` is ``d_head`` for heads).
+    """
+
+    kind: str
+    layer: int
+    values: torch.Tensor
+    heads: torch.Tensor | None = None
+    positions: torch.Tensor | None = None
+
+
+class BackendError(RuntimeError):
+    """A model can't be loaded or run. The message says what went wrong and how to fix it."""
+
+
+class Cancelled(Exception):
+    """The user cancelled the job (a run between batches, or a model load between steps)."""
+
+
+class ModelBackend(ABC):
+    info: ModelInfo
+
+    def __init__(self) -> None:
+        # All model use goes through this lock, one forward pass at a time, so interactive
+        # requests can interleave with a running sweep between batches.
+        self.lock = threading.RLock()
+
+    @abstractmethod
+    def tokenize(self, text: str, prepend_bos: bool) -> Tokenized: ...
+
+    @abstractmethod
+    def single_token_id(self, text: str) -> int | None:
+        """The token id if ``text`` is exactly one token, else None."""
+
+    @abstractmethod
+    def token_str(self, token_id: int) -> str: ...
+
+    @abstractmethod
+    def final_logits(self, tokens: torch.Tensor, patch: Patch | None = None) -> torch.Tensor:
+        """Logits at the last position, ``[B, vocab]`` in float32."""
+
+    @abstractmethod
+    def capture(
+        self, tokens: torch.Tensor, sites: list[tuple[str, int]]
+    ) -> dict[tuple[str, int], torch.Tensor]:
+        """Activations for ``(kind, layer)`` sites: ``[B, pos, d]``, or ``[B, pos, H, d_head]``."""
+
+    @abstractmethod
+    def attention_pattern(self, tokens: torch.Tensor, layer: int) -> torch.Tensor:
+        """Attention probabilities ``[B, H, query, key]`` in float32."""
+
+    def edit_logits(
+        self,
+        tokens: torch.Tensor,
+        kind: str,
+        layer: int,
+        edit: Any,
+    ) -> torch.Tensor:
+        """Logits at the last position, ``[B, vocab]``, with the activation at ``(kind, layer)``
+        replaced by ``edit(activation)``: ``[B, pos, d]`` in, the same shape out. For
+        residual-stream sites the edit changes the stream itself, as patching does."""
+        raise BackendError("Editing activations isn't supported by this model backend.")
+
+    def path_patch(
+        self,
+        tokens: torch.Tensor,
+        sender: Patch,
+        frozen_heads: dict[int, torch.Tensor],
+        frozen_mlps: dict[int, torch.Tensor] | None,
+        receivers: list[tuple[str, int, int, str]],
+    ) -> torch.Tensor:
+        """Logits at the last position, ``[B, vocab]``, after patching the sender's effect into
+        the receivers' inputs only.
+
+        First pass: the sender is patched, every attention head's output ``z`` is held at
+        ``frozen_heads`` (the receiver run's own values, ``[B, pos, H, d_head]`` per layer) except
+        the sender's, and so are MLP outputs when ``frozen_mlps`` is given; the receivers' inputs
+        are recorded. Second pass: only those inputs are patched in. Receivers are
+        ``("head", layer, head, "q" | "k" | "v")`` or ``("logits", -1, -1, "")``.
+        """
+        raise BackendError("Path patching isn't supported by this model backend.")
+
+    def gradients(
+        self,
+        tokens: torch.Tensor,
+        answers: torch.Tensor,
+        distractors: torch.Tensor,
+        sites: list[tuple[str, int]],
+    ) -> tuple[dict[tuple[str, int], torch.Tensor], dict[tuple[str, int], torch.Tensor]]:
+        """Activations at ``(kind, layer)`` sites and the gradient of logit(answer) -
+        logit(distractor) at the last position with respect to each, in one forward and backward
+        pass. Shapes as in :meth:`capture`. The gradient of a residual-stream site is with respect
+        to the residual stream itself, not only to what the next component reads."""
+        raise BackendError("Gradients aren't supported by this model backend.")
+
+    def direct_effects(
+        self,
+        tokens: torch.Tensor,
+        answers: torch.Tensor,
+        distractors: torch.Tensor,
+        heads: bool,
+    ) -> dict[str, torch.Tensor]:
+        """Direct contributions to logit(answer) - logit(distractor) at the last position.
+
+        Returns float64 tensors on the CPU, one row per prompt: ``embed`` ``[B]``, ``attn_out``
+        and ``mlp_out`` ``[B, layers]``, ``head`` ``[B, layers, heads]`` when ``heads`` is true,
+        ``logit_diff`` ``[B]`` and ``remainder`` ``[B]`` (what biases add: the logit difference
+        minus every component's term).
+        """
+        raise BackendError("Direct logit attribution isn't supported by this model backend.")
+
+    def layer_logits(self, tokens: torch.Tensor, position: int, row: int) -> torch.Tensor:
+        """Final-norm logit lens at resid_post, ``[layer, vocab]`` for one row, on CPU.
+
+        Preserve the supplied batch composition. Backends must opt in rather than assuming
+        that every model has the same normalization and vocabulary projection.
+        """
+        raise BackendError("Per-layer predictions aren't supported by this model backend.")
+
+    @property
+    def device(self) -> torch.device:
+        return torch.device(self.info.device)
+
+    def memory_in_use(self) -> int | None:
+        """Bytes in use: accelerator memory on a GPU, this process's resident memory on a CPU."""
+        dev = self.info.device
+        if dev == "cuda" and torch.cuda.is_available():
+            return int(torch.cuda.memory_allocated())
+        if dev == "mps" and hasattr(torch, "mps"):
+            try:
+                return int(torch.mps.current_allocated_memory())
+            except Exception:
+                return None
+        try:
+            import psutil
+
+            return int(psutil.Process().memory_info().rss)
+        except Exception:
+            return None
+
+    def close(self) -> None:  # noqa: B027 - optional hook
+        """Release memory held by the model."""

logogram/backends/hub.py

@@ -0,0 +1,375 @@
+"""Hugging Face Hub access: resolve an exact revision, download with progress, read metadata.
+
+This is the only module that talks to the network, and only when the user starts it (loading a
+model or asking for a memory estimate). Weights are loaded only from safetensors files, which
+cannot execute code; pickle-based ``.bin`` checkpoints are refused.
+"""
+
+from __future__ import annotations
+
+import fnmatch
+import json
+import logging
+import sys
+import threading
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, TypeVar
+
+from logogram.backends.base import BackendError, Cancelled
+
+log = logging.getLogger(__name__)
+
+# Top-level files needed to build the model and tokenizer. Weights are added separately.
+_SUPPORT_PATTERNS = ("*.json", "*.txt", "*.model", "*.tiktoken")
+_SKIP_FILES = {"README.md", ".gitattributes"}
+
+
+@dataclass
+class RepoFiles:
+    revision: str
+    files: list[tuple[str, int]]  # (path, size in bytes) to download
+    n_params: int | None
+    gated: bool
+
+
+ProgressFn = Callable[[int, int, str], None]  # (bytes done, bytes total, current file)
+T = TypeVar("T")
+
+
+def _is_offline_error(exc: Exception) -> bool:
+    from huggingface_hub.errors import OfflineModeIsEnabled
+
+    if isinstance(exc, OfflineModeIsEnabled):
+        return True
+    name = type(exc).__name__
+    return (
+        any(k in name for k in ("Connect", "Timeout", "NameResolution"))
+        or "connect" in str(exc).lower()
+    )
+
+
+def _friendly_hub_error(model_id: str, exc: Exception) -> BackendError:
+    from huggingface_hub.errors import (
+        GatedRepoError,
+        RepositoryNotFoundError,
+        RevisionNotFoundError,
+    )
+
+    if isinstance(exc, GatedRepoError):
+        return BackendError(
+            f"{model_id} is gated. Accept its license on huggingface.co, then log in on this "
+            "machine with `hf auth login` and try again."
+        )
+    if isinstance(exc, RepositoryNotFoundError):
+        return BackendError(
+            f"No model called {model_id} was found on Hugging Face. Check the spelling (ids look "
+            "like owner/name). Private models need `hf auth login` first."
+        )
+    if isinstance(exc, RevisionNotFoundError):
+        return BackendError(
+            f"That revision of {model_id} doesn't exist. Check the commit or branch."
+        )
+    return BackendError(f"Couldn't reach Hugging Face for {model_id}: {exc}")
+
+
+def _select_files(siblings: list[tuple[str, int]]) -> list[tuple[str, int]]:
+    top_level = [(f, s) for f, s in siblings if "/" not in f and f not in _SKIP_FILES]
+    weights = [(f, s) for f, s in top_level if f.endswith(".safetensors")]
+    if not weights:
+        has_pickle = any(f.endswith((".bin", ".pt", ".pth", ".ckpt")) for f, _ in top_level)
+        if has_pickle:
+            raise BackendError(
+                "This model only publishes pickle weights (.bin). Logogram loads safetensors "
+                "weights only, because pickle files can run code when opened. Choose a model "
+                "with model.safetensors, or convert the weights yourself."
+            )
+        raise BackendError("This repository has no safetensors weights at its top level.")
+    support = [
+        (f, s)
+        for f, s in top_level
+        if any(fnmatch.fnmatch(f, p) for p in _SUPPORT_PATTERNS)
+        and not f.startswith(("onnx", "flax", "tf_"))
+    ]
+    return sorted(support) + sorted(weights)
+
+
+def resolve(model_id: str, revision: str | None) -> RepoFiles:
+    """Ask the Hub for the exact commit and the files to fetch. Falls back to the local cache."""
+    from huggingface_hub import HfApi
+
+    try:
+        info = HfApi().model_info(model_id, revision=revision, files_metadata=True)
+    except Exception as exc:  # noqa: BLE001 - classified below
+        if _is_offline_error(exc):
+            cached = cached_snapshot(model_id, revision)
+            if cached is not None:
+                sha, path = cached
+                files = [(p.name, p.stat().st_size) for p in path.iterdir() if p.is_file()]
+                return RepoFiles(
+                    revision=sha, files=_select_files(files), n_params=None, gated=False
+                )
+            raise BackendError(
+                f"Couldn't reach Hugging Face, and {model_id} isn't in your local cache. "
+                "Connect to the internet to download it once."
+            ) from exc
+        raise _friendly_hub_error(model_id, exc) from exc
+    siblings = [(s.rfilename, int(s.size or 0)) for s in (info.siblings or [])]
+    n_params = None
+    if getattr(info, "safetensors", None) is not None:
+        n_params = int(info.safetensors.total)
+    return RepoFiles(
+        revision=info.sha,
+        files=_select_files(siblings),
+        n_params=n_params,
+        gated=bool(info.gated),
+    )
+
+
+def cached_snapshot(model_id: str, revision: str | None) -> tuple[str, Path] | None:
+    """The (commit, folder) of a cached snapshot containing config.json, if any."""
+    from huggingface_hub import try_to_load_from_cache
+
+    found = try_to_load_from_cache(model_id, "config.json", revision=revision or "main")
+    if isinstance(found, str):
+        path = Path(found).parent
+        return path.name, path
+    return None
+
+
+def validate_local_weights(folder: Path) -> None:
+    """Refuse pickle-only caches and unsafe or incomplete safetensors shard indexes."""
+    if (folder / "adapter_config.json").exists():
+        raise BackendError(
+            "Adapter checkpoints are not supported. Choose a complete base or merged model published as safetensors, without a separate adapter configuration."
+        )
+    single = folder / "model.safetensors"
+    index = folder / "model.safetensors.index.json"
+    if single.is_file():
+        return
+    if not index.is_file():
+        raise BackendError(
+            "No model.safetensors weights are available. Download a model with safetensors weights."
+        )
+    try:
+        data = json.loads(index.read_text(encoding="utf-8"))
+        mapping = data["weight_map"]
+        if not isinstance(mapping, dict) or not mapping:
+            raise ValueError("Empty weight map")
+        for name in mapping.values():
+            if (
+                not isinstance(name, str)
+                or "/" in name
+                or "\\" in name
+                or ":" in name
+                or not name.endswith(".safetensors")
+            ):
+                raise ValueError("Unsafe shard name")
+            if not (folder / name).is_file():
+                raise ValueError("Missing shard")
+    except (OSError, ValueError, KeyError, TypeError) as exc:
+        raise BackendError(
+            "The safetensors index is invalid or has missing shards. Download the complete safetensors checkpoint again."
+        ) from exc
+
+
+def _make_bar_class(on_value: Callable[[int], None], size: int) -> type:
+    """A silent progress bar that reports how much of one file has arrived.
+
+    Plain HTTP downloads report bytes as they arrive. Xet downloads report network bytes often
+    and bytes written to disk rarely; the larger of the two, capped at the file size, gives
+    smooth progress that still ends exactly at the file size.
+    """
+    from tqdm.std import tqdm
+
+    class _Bar(tqdm):  # type: ignore[misc]
+        monitor_interval = 0  # no display, so no tqdm monitor thread per file
+
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            kwargs["disable"] = True
+            super().__init__(*args, **kwargs)
+            self._written = 0
+            self._received = 0
+
+        def _report(self) -> None:
+            on_value(min(size, max(self._written, self._received)))
+
+        def update(self, n: float | None = 1) -> bool | None:
+            if n:
+                self._written += int(n)
+                self._report()
+            return None
+
+        def update_transfer(self, n: float | None = 1) -> None:
+            if n:
+                self._received += int(n)
+                self._report()
+
+        def set_transfer_postfix_str(self, *args: Any, **kwargs: Any) -> None:
+            return None
+
+        def __getattr__(self, name: str) -> Any:
+            # Display hooks that newer huggingface_hub versions may call: nothing to display.
+            if name.startswith(("set_", "update_")):
+                return lambda *args, **kwargs: None
+            raise AttributeError(name)
+
+    return _Bar
+
+
+def download(
+    model_id: str,
+    repo: RepoFiles,
+    progress: ProgressFn | None = None,
+    *,
+    cancel: threading.Event | None = None,
+) -> Path:
+    """Download the selected files at the pinned revision and return the snapshot folder.
+
+    Setting ``cancel`` raises ``Cancelled`` within a fraction of a second and stops the transfer;
+    a partial file stays in the cache, where the next attempt resumes it.
+    """
+    from huggingface_hub import hf_hub_download, try_to_load_from_cache
+
+    total = sum(size for _, size in repo.files)
+    finished = 0  # bytes of files already complete
+    folder: Path | None = None
+    for filename, size in repo.files:
+        if cancel is not None and cancel.is_set():
+            raise Cancelled()
+        cached = try_to_load_from_cache(model_id, filename, revision=repo.revision)
+        if isinstance(cached, str):
+            finished += size
+            folder = Path(cached).parent
+            if progress:
+                progress(finished, total, filename)
+            continue
+
+        def on_value(value: int, _name: str = filename, _base: int = finished) -> None:
+            if cancel is not None and cancel.is_set():
+                # Raising stops a plain HTTP transfer. Xet transfers are aborted from outside
+                # (see _abort_xet); an exception raised in their callback would only be printed.
+                if not _called_from_xet():
+                    raise Cancelled()
+                return
+            if progress:
+                progress(_base + value, total, _name)
+
+        def fetch(_name: str = filename, _size: int = size, _on: Any = on_value) -> str:
+            return hf_hub_download(
+                model_id,
+                _name,
+                revision=repo.revision,
+                tqdm_class=_make_bar_class(_on, _size),
+            )
+
+        try:
+            path = _cancellable(fetch, cancel)
+        except Cancelled:
+            raise
+        except Exception as exc:  # noqa: BLE001
+            raise _friendly_hub_error(model_id, exc) from exc
+        finished += size
+        folder = Path(path).parent
+        if progress:
+            progress(finished, total, filename)
+    if folder is None:
+        raise BackendError(f"Nothing to download for {model_id}.")
+    return folder
+
+
+def _cancellable(fn: Callable[[], T], cancel: threading.Event | None) -> T:
+    """Run ``fn`` on a helper thread, so waiting for it can stop as soon as ``cancel`` is set."""
+    if cancel is None:
+        return fn()
+    box: dict[str, Any] = {}
+    done = threading.Event()
+
+    def target() -> None:
+        try:
+            box["value"] = fn()
+        except BaseException as exc:  # noqa: BLE001 - handed to the waiting thread
+            box["error"] = exc
+        finally:
+            done.set()
+
+    threading.Thread(target=target, name="logogram-download", daemon=True).start()
+    while not done.wait(0.1):
+        if cancel.is_set():
+            _abort_xet()
+            raise Cancelled()
+    if "error" in box:
+        raise box["error"]
+    return box["value"]
+
+
+def _called_from_xet() -> bool:
+    frame = sys._getframe(1)
+    while frame is not None:
+        if "xet" in frame.f_globals.get("__name__", ""):
+            return True
+        frame = frame.f_back
+    return False
+
+
+def _abort_xet() -> None:
+    """Stop Xet transfers in flight, as huggingface_hub does on Ctrl+C. Without this (if the
+    private helper moves), a cancelled Xet download finishes in the background, into the cache."""
+    try:
+        from huggingface_hub.utils._xet import abort_xet_session
+
+        abort_xet_session()
+    except Exception:  # noqa: BLE001
+        log.debug("couldn't abort Xet transfers", exc_info=True)
+
+
+@dataclass
+class ArchitectureSummary:
+    n_layers: int
+    n_heads: int
+    d_model: int
+    d_mlp: int
+    d_vocab: int
+    n_ctx: int
+    architecture: str
+
+
+def read_architecture(config: dict[str, Any]) -> ArchitectureSummary:
+    """Read the shape of a transformer from a Hugging Face config.json (best effort)."""
+
+    def first(*keys: str, default: int | None = None) -> int:
+        for key in keys:
+            value = config.get(key)
+            if isinstance(value, int) and value > 0:
+                return value
+        text = config.get("text_config")
+        if isinstance(text, dict):
+            for key in keys:
+                value = text.get(key)
+                if isinstance(value, int) and value > 0:
+                    return value
+        if default is None:
+            raise BackendError(f"config.json has none of {', '.join(keys)}.")
+        return default
+
+    d_model = first("hidden_size", "n_embd", "d_model")
+    return ArchitectureSummary(
+        n_layers=first("num_hidden_layers", "n_layer", "num_layers"),
+        n_heads=first("num_attention_heads", "n_head", "num_heads"),
+        d_model=d_model,
+        d_mlp=first("intermediate_size", "n_inner", "ffn_dim", default=4 * d_model),
+        d_vocab=first("vocab_size"),
+        n_ctx=first("max_position_embeddings", "n_positions", "n_ctx", default=2048),
+        architecture=(config.get("architectures") or ["unknown"])[0],
+    )
+
+
+def fetch_config(model_id: str, revision: str) -> dict[str, Any]:
+    from huggingface_hub import hf_hub_download
+
+    try:
+        path = hf_hub_download(model_id, "config.json", revision=revision)
+    except Exception as exc:  # noqa: BLE001
+        raise _friendly_hub_error(model_id, exc) from exc
+    return json.loads(Path(path).read_text(encoding="utf-8"))