open-interpretability 0.0.20__py3-none-any.whl

openinterp/superposition_finders.py

@@ -0,0 +1,225 @@
+"""Ways to find superpositions: sparse directions over one layer's MLP neurons (issue #3686).
+
+Two finders give definitions that ``open-interpretability audit --superposition`` reads.
+
+- :func:`from_circuit` groups a traced circuit's neurons by layer. The weights are the
+  neurons' attributions.
+- :func:`grammatical_number` fits a sparse probe for the number of the noun that the
+  verb agrees with, at one layer's last token.
+
+``open-interpretability superpositions`` runs them (``openinterp/cli_superpositions.py``).
+"""
+
+from __future__ import annotations
+
+import math
+import random
+from collections.abc import Iterable, Sequence
+
+import torch
+
+from openinterp.agreement import TEMPLATES, pairs
+from openinterp.core.profiles import SuperpositionMember, SuperpositionProfile
+
+NUMBER_SOURCE = "grammatical number from the agreement templates (#3686)"
+TRAIN_SHARE = 0.8  # the share of examples the probe trains on. The rest is held out.
+L1_DECAY = 0.7  # each step of the L1 path multiplies the strength by this
+ISTA_STEPS = 200  # proximal-gradient steps at each L1 strength
+L2_STRENGTH = 1e-2  # the ridge penalty of the refit on the kept neurons
+BATCH = 64  # prompts in one forward pass
+
+
+def from_circuit(
+    rows: Iterable[dict], *, top: int = 100, min_members: int = 2, source: str
+) -> list[tuple[int, SuperpositionProfile]]:
+    """One (layer, profile) for each layer with ``min_members`` of the circuit's ``top`` neurons.
+
+    ``rows`` are a circuit's rows, as ``scripts/sva_circuit.py`` writes them. Only rows
+    with ``"table": "neuron"`` count, in ``rank`` order. A member's weight is its
+    attribution. Members with a zero attribution are not used.
+    """
+    if min_members < 2:
+        raise ValueError("a superposition has at least 2 members, so min_members must be ≥ 2")
+    neurons = sorted((r for r in rows if r.get("table") == "neuron"), key=lambda r: r["rank"])
+    by_layer: dict[int, list[SuperpositionMember]] = {}
+    for row in neurons[:top]:
+        if row["attribution"] != 0:
+            member = SuperpositionMember(neuron_idx=row["neuron_idx"], weight=row["attribution"])
+            by_layer.setdefault(row["layer"], []).append(member)
+    return [
+        (layer, _profile(members, "circuit", source))
+        for layer, members in sorted(by_layer.items())
+        if len(members) >= min_members
+    ]
+
+
+def _profile(members: list[SuperpositionMember], method: str, source: str):
+    ordered = sorted(members, key=lambda m: m.neuron_idx)
+    return SuperpositionProfile(members=ordered, method=method, source=source)
+
+
+def split_by_group(groups: torch.Tensor, seed: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """Train and held-out row indices. All rows of one group are on the same side.
+
+    A fixed shuffle of the groups puts the first 80% of them in training.
+    """
+    ids = groups.unique()
+    order = ids[torch.randperm(len(ids), generator=torch.Generator().manual_seed(seed))]
+    in_train = torch.isin(groups, order[: int(TRAIN_SHARE * len(ids))])
+    return in_train.nonzero().flatten(), (~in_train).nonzero().flatten()
+
+
+def sparse_probe(
+    acts: torch.Tensor,
+    labels: torch.Tensor,
+    *,
+    k: int,
+    seed: int = 0,
+    groups: torch.Tensor | None = None,
+) -> tuple[list[int], list[float], float]:
+    """The ``k`` neurons that best predict ``labels``, their weights and the held-out accuracy.
+
+    ``acts`` is ``[n, width]`` and ``labels`` is ``[n]`` of 0 and 1. ``groups`` gives
+    each row a group id, and :func:`split_by_group` holds out 20% of the groups. With no
+    ``groups``, each row is a group. Each column is standardised on the training part.
+    An L1 logistic regression gets weaker until at least ``k`` weights are not zero. The
+    ``k`` largest weights choose the neurons. Then an L2 logistic regression is fitted
+    again on those neurons only. The weights are in the raw activation space.
+    """
+    x, y = acts.float().cpu(), labels.float().cpu()
+    groups = torch.arange(len(y)) if groups is None else groups.cpu()
+    train, test = split_by_group(groups, seed)
+    mean, std = x[train].mean(0), x[train].std(0)
+    std = torch.where(std > 1e-6, std, torch.ones_like(std))  # a constant column stays 0
+    z = (x - mean) / std
+    weights = _l1_path(z[train], y[train], k)
+    keep = weights.abs().topk(k).indices.sort().values
+    w, b = _ridge_logistic(z[train][:, keep], y[train])
+    predicted = (z[test][:, keep] @ w + b > 0).float()
+    accuracy = float((predicted == y[test]).float().mean()) if len(test) else math.nan
+    return keep.tolist(), (w / std[keep]).tolist(), accuracy
+
+
+def _l1_path(z: torch.Tensor, y: torch.Tensor, k: int) -> torch.Tensor:
+    """L1 logistic weights at the strongest strength on the path with at least ``k`` non-zero."""
+    n = len(y)
+    strength = float((z.T @ (y - y.mean())).abs().max()) / n  # all weights are 0 above this
+    step = 1.0 / (0.25 * _top_eigenvalue(z) / n)  # 1 / the gradient's Lipschitz constant
+    w, b = torch.zeros(z.shape[1]), torch.zeros(())
+    while int((w != 0).sum()) < k and strength > 1e-8:
+        strength *= L1_DECAY
+        for _ in range(ISTA_STEPS):
+            residual = torch.sigmoid(z @ w + b) - y
+            w = _shrink(w - step * (z.T @ residual) / n, step * strength)
+            b = b - step * residual.mean()
+    return w
+
+
+def _shrink(w: torch.Tensor, by: float) -> torch.Tensor:
+    """Soft thresholding: the proximal step of the L1 penalty."""
+    return torch.sign(w) * torch.clamp(w.abs() - by, min=0.0)
+
+
+def _top_eigenvalue(z: torch.Tensor, steps: int = 50) -> float:
+    """The largest eigenvalue of ``zᵀz``, by power iteration."""
+    v = torch.ones(z.shape[1]) / math.sqrt(z.shape[1])
+    for _ in range(steps):
+        v = z.T @ (z @ v)
+        v = v / v.norm().clamp_min(1e-12)
+    return float((z @ v).pow(2).sum()) or 1.0
+
+
+def _ridge_logistic(z: torch.Tensor, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Weights and bias of an L2 logistic regression, fitted with L-BFGS."""
+    w = torch.zeros(z.shape[1], requires_grad=True)
+    b = torch.zeros((), requires_grad=True)
+    optimiser = torch.optim.LBFGS([w, b], max_iter=200, line_search_fn="strong_wolfe")
+
+    def closure():
+        optimiser.zero_grad()
+        loss = torch.nn.functional.binary_cross_entropy_with_logits(z @ w + b, y)
+        loss = loss + L2_STRENGTH * w.pow(2).sum() / 2
+        loss.backward()
+        return loss
+
+    with torch.enable_grad():
+        optimiser.step(closure)
+    return w.detach(), b.detach()
+
+
+def agreement_examples(n: int, seed: int) -> list[tuple[str, int]]:
+    """``2n`` (prompt, label) examples from ``n`` agreement pairs, spread over the templates.
+
+    The label is 1 when the noun that the verb agrees with is plural. Each pair gives
+    both of its prompts, so it gives one singular and one plural example.
+    """
+    rng = random.Random(seed)
+    out = []
+    for i, template in enumerate(TEMPLATES):
+        count = n // len(TEMPLATES) + (i < n % len(TEMPLATES))
+        for row in pairs(template, count, rng):
+            clean_plural = int(row["sign"] < 0)
+            out += [(row["clean"], clean_plural), (row["corrupted"], 1 - clean_plural)]
+    return out
+
+
+def last_token_activations(runner, texts: Sequence[str], layers: Sequence[int]):
+    """``{layer: [len(texts), width]}``: each layer's MLP activations at each text's last token.
+
+    Texts with the same number of tokens run together, so no text needs padding.
+    """
+    ids = [runner.tokenizer.encode(t) for t in texts]
+    by_length: dict[int, list[int]] = {}
+    for i, row in enumerate(ids):
+        by_length.setdefault(len(row), []).append(i)
+    out = {layer: torch.zeros(len(texts), runner.mlp_width(layer)) for layer in layers}
+    for group in by_length.values():
+        for start in range(0, len(group), BATCH):
+            part = group[start : start + BATCH]
+            batch = torch.tensor([ids[i] for i in part])
+            cache = runner.layer_activations({"input_ids": batch}, list(layers))
+            for layer in layers:
+                out[layer][part] = cache[layer][:, -1].float().cpu()
+    return out
+
+
+def grammatical_number(
+    runner, layers: Sequence[int], *, k: int, n: int, seed: int
+) -> list[tuple[int, SuperpositionProfile, float]]:
+    """One (layer, profile, held-out accuracy) per layer: a ``k``-neuron probe for noun number.
+
+    It uses ``n`` agreement pairs (:func:`agreement_examples`). The probe reads the
+    layer's MLP activations at the last token, where the verb comes next. The two
+    prompts of a pair are always on the same side of the held-out split.
+    """
+    examples = agreement_examples(n, seed)
+    labels = torch.tensor([label for _, label in examples])
+    pair_ids = torch.arange(len(examples)) // 2  # agreement_examples puts a pair's prompts together
+    acts = last_token_activations(runner, [text for text, _ in examples], layers)
+    out = []
+    for layer in layers:
+        neurons, weights, accuracy = sparse_probe(
+            acts[layer], labels, k=k, seed=seed, groups=pair_ids
+        )
+        members = [
+            SuperpositionMember(neuron_idx=i, weight=w)
+            for i, w in zip(neurons, weights, strict=True)
+            if w != 0
+        ]
+        out.append((layer, _profile(members, "sparse-probe", NUMBER_SOURCE), accuracy))
+    return out
+
+
+def definition(model: str, layer: int, profile: SuperpositionProfile) -> dict:
+    """The definition that ``open-interpretability audit --superposition`` reads."""
+    return {"model": model, "layer": layer, **profile.model_dump()}
+
+
+def compare(a: SuperpositionProfile, b: SuperpositionProfile) -> tuple[list[int], float]:
+    """The neurons that ``a`` and ``b`` share, and the cosine of their two directions."""
+    wa = {m.neuron_idx: m.weight for m in a.members}
+    wb = {m.neuron_idx: m.weight for m in b.members}
+    shared = sorted(wa.keys() & wb.keys())
+    dot = sum(wa[i] * wb[i] for i in shared)
+    norms = math.sqrt(sum(v * v for v in wa.values())) * math.sqrt(sum(v * v for v in wb.values()))
+    return shared, dot / norms

openinterp/sync/__init__.py

@@ -0,0 +1,5 @@
+"""Deploy the built site to Cloudflare R2 (``open-interpretability deploy``)."""
+
+from openinterp.sync.r2 import R2Config, SyncPlan, cache_control_for, make_client, sync_directory
+
+__all__ = ["R2Config", "SyncPlan", "cache_control_for", "make_client", "sync_directory"]

openinterp/sync/r2.py

@@ -0,0 +1,280 @@
+"""Delta-sync a built ``dist/`` directory to a Cloudflare R2 bucket over the S3 API.
+
+Only files whose MD5 differs from the remote ETag are uploaded (``put_object``
+always produces single-part uploads, whose ETag is the content MD5). Because the
+builder is deterministic, an unchanged registry syncs as a no-op: nothing is
+uploaded, deleted or purged.
+
+* Upload order: content-addressed files, then other assets, then HTML, so a page
+  never references an object that has not landed yet.
+* Derived prefixes (``pagefind/``) are only synced when a source file changed. The
+  Pagefind index is not byte-stable across runs, so syncing it alone would churn.
+* Protected prefixes (``data/snapshots/``) are never deleted: snapshot URLs are
+  permanent.
+* With a Cloudflare zone configured, only changed short-cache URLs are purged.
+
+Credentials come from ``CF_ACCOUNT_ID``, ``R2_ACCESS_KEY_ID`` and
+``R2_SECRET_ACCESS_KEY``; purging uses ``CF_ZONE_ID`` and ``CF_API_TOKEN``.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import mimetypes
+import os
+import re
+import urllib.request
+from collections.abc import Callable, Mapping, Sequence
+from dataclasses import dataclass, field
+from pathlib import Path, PurePosixPath
+from typing import Any
+
+from openinterp.core.exceptions import SyncError
+
+SHORT_CACHE = "public, max-age=3600"
+IMMUTABLE_CACHE = "public, max-age=31536000, immutable"
+DERIVED_PREFIXES = ("pagefind/",)
+PROTECTED_PREFIXES = ("data/snapshots/",)
+PURGE_BATCH = 30  # Cloudflare's per-request limit for purge-by-URL
+
+CONTENT_TYPES = {
+    ".html": "text/html; charset=utf-8",
+    ".json": "application/json",
+    ".parquet": "application/vnd.apache.parquet",
+    ".wasm": "application/wasm",
+    ".js": "text/javascript",
+    ".css": "text/css",
+    ".svg": "image/svg+xml",
+}
+# Pagefind's index shards are content-addressed. Its entry files are not, so they must
+# stay on the short cache: pagefind.js, pagefind-entry.json and pagefind-ui.*.
+HASHED_NAME = re.compile(r"[_.-][0-9a-f]{7,}\.[A-Za-z0-9_]+$")
+PAGEFIND_SHARD_SUFFIXES = {".pf_fragment", ".pf_index", ".pf_filter", ".pf_meta", ".pagefind"}
+
+
+def cache_control_for(key: str) -> str:
+    """Cache policy: immutable only for content-addressed files, one hour for everything else."""
+    path = PurePosixPath(key)
+    if key.startswith("data/snapshots/") or path.suffix in PAGEFIND_SHARD_SUFFIXES:
+        return IMMUTABLE_CACHE
+    if path.suffix not in {".html", ".json"} and HASHED_NAME.search(path.name):
+        return IMMUTABLE_CACHE
+    return SHORT_CACHE
+
+
+def content_type_for(key: str) -> str:
+    """The Content-Type to upload ``key`` with, from its suffix."""
+    suffix = PurePosixPath(key).suffix
+    return CONTENT_TYPES.get(suffix) or mimetypes.guess_type(key)[0] or "application/octet-stream"
+
+
+def file_md5(path: Path) -> str:
+    """The MD5 of ``path``, which equals a single-part upload's ETag."""
+    digest = hashlib.md5(usedforsecurity=False)
+    with path.open("rb") as fh:
+        for chunk in iter(lambda: fh.read(1 << 20), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+@dataclass(frozen=True)
+class R2Config:
+    """The R2 account and credentials to sync with."""
+
+    account_id: str
+    access_key_id: str
+    secret_access_key: str
+
+    @property
+    def endpoint_url(self) -> str:
+        """The account's S3 API endpoint."""
+        return f"https://{self.account_id}.r2.cloudflarestorage.com"
+
+    @classmethod
+    def from_env(cls, env: Mapping[str, str] = os.environ) -> R2Config:
+        """Read the config from the environment; raise SyncError when a variable is missing."""
+        names = ("CF_ACCOUNT_ID", "R2_ACCESS_KEY_ID", "R2_SECRET_ACCESS_KEY")
+        missing = [n for n in names if not env.get(n)]
+        if missing:
+            raise SyncError(f"Missing environment variables: {', '.join(missing)}")
+        return cls(*(env[n] for n in names))
+
+
+def make_client(config: R2Config) -> Any:
+    """An S3 client for R2, made with ``config``."""
+    import boto3
+
+    return boto3.client(
+        "s3",
+        endpoint_url=config.endpoint_url,
+        aws_access_key_id=config.access_key_id,
+        aws_secret_access_key=config.secret_access_key,
+        region_name="auto",
+    )
+
+
+@dataclass
+class SyncPlan:
+    """What a sync uploads, deletes, leaves alone and skips."""
+
+    uploads: list[str] = field(default_factory=list)
+    deletes: list[str] = field(default_factory=list)
+    unchanged: list[str] = field(default_factory=list)
+    skipped_derived: list[str] = field(default_factory=list)
+
+    @property
+    def changed(self) -> bool:
+        """True when the sync uploads or deletes anything."""
+        return bool(self.uploads or self.deletes)
+
+
+def local_files(dist_dir: Path, prefix: str = "") -> dict[str, Path]:
+    """Every file in ``dist_dir`` by its bucket key, except dotfiles."""
+    dist_dir = Path(dist_dir)
+    return {
+        prefix + path.relative_to(dist_dir).as_posix(): path
+        for path in sorted(dist_dir.rglob("*"))
+        if path.is_file() and not path.name.startswith(".")
+    }
+
+
+def remote_etags(client: Any, bucket: str, prefix: str = "") -> dict[str, str]:
+    """The ETag of every object in ``bucket`` under ``prefix``, by key."""
+    etags: dict[str, str] = {}
+    for page in client.get_paginator("list_objects_v2").paginate(Bucket=bucket, Prefix=prefix):
+        for obj in page.get("Contents", []):
+            etags[obj["Key"]] = obj["ETag"].strip('"')
+    return etags
+
+
+def _upload_order(key: str) -> tuple[int, str]:
+    if key.endswith(".html"):
+        return (2, key)
+    return (0 if cache_control_for(key) == IMMUTABLE_CACHE else 1, key)
+
+
+def plan_sync(
+    local: Mapping[str, Path],
+    remote: Mapping[str, str],
+    *,
+    delete: bool = False,
+    prefix: str = "",
+    derived: Sequence[str] = DERIVED_PREFIXES,
+    protected: Sequence[str] = PROTECTED_PREFIXES,
+) -> SyncPlan:
+    """Compare local files with the bucket, and decide what to upload and delete.
+
+    The module docstring gives the rules.
+    """
+    plan = SyncPlan()
+
+    def rel(key: str) -> str:
+        return key.removeprefix(prefix)
+
+    for key, path in local.items():
+        (plan.unchanged if remote.get(key) == file_md5(path) else plan.uploads).append(key)
+    if delete:
+        plan.deletes = sorted(
+            k for k in set(remote) - set(local) if not rel(k).startswith(tuple(protected))
+        )
+
+    # Derived files only follow a change to a source file; on their own they are churn.
+    is_derived = lambda k: rel(k).startswith(tuple(derived))  # noqa: E731
+    if not any(not is_derived(k) for k in plan.uploads + plan.deletes):
+        plan.skipped_derived = sorted(plan.uploads + plan.deletes)
+        plan.unchanged += plan.uploads
+        plan.uploads, plan.deletes = [], []
+
+    plan.uploads.sort(key=lambda k: _upload_order(rel(k)))
+    return plan
+
+
+def sync_directory(
+    dist_dir: Path,
+    bucket: str,
+    client: Any,
+    *,
+    prefix: str = "",
+    delete: bool = False,
+    dry_run: bool = False,
+    log: Callable[[str], None] = lambda _msg: None,
+) -> SyncPlan:
+    """Sync ``dist_dir`` to ``bucket``; with ``dry_run``, only log what would change."""
+    dist_dir = Path(dist_dir)
+    if not dist_dir.is_dir():
+        raise SyncError(f"{dist_dir} is not a directory; run `open-interpretability build` first")
+    if prefix and not prefix.endswith("/"):
+        prefix += "/"
+
+    local = local_files(dist_dir, prefix)
+    try:
+        plan = plan_sync(local, remote_etags(client, bucket, prefix), delete=delete, prefix=prefix)
+        for key in plan.skipped_derived:
+            log(f"skip    {key}  (derived; no source file changed)")
+        for key in plan.uploads:
+            log(f"upload  {key}  [{cache_control_for(key.removeprefix(prefix))}]")
+            if not dry_run:
+                with local[key].open("rb") as body:
+                    client.put_object(
+                        Bucket=bucket,
+                        Key=key,
+                        Body=body,
+                        ContentType=content_type_for(key),
+                        CacheControl=cache_control_for(key.removeprefix(prefix)),
+                    )
+        for start in range(0, len(plan.deletes), 1000):  # DeleteObjects takes <= 1000 keys
+            batch = plan.deletes[start : start + 1000]
+            for key in batch:
+                log(f"delete  {key}")
+            if not dry_run:
+                client.delete_objects(
+                    Bucket=bucket, Delete={"Objects": [{"Key": k} for k in batch], "Quiet": True}
+                )
+    except SyncError:
+        raise
+    except Exception as exc:  # botocore raises ClientError, EndpointConnectionError, ...
+        raise SyncError(f"Sync to bucket {bucket!r} failed: {exc}") from exc
+    return plan
+
+
+def purge_targets(plan: SyncPlan, site_url: str, prefix: str = "") -> list[str]:
+    """Public URLs whose cached copies are stale after ``plan`` ran.
+
+    Content-addressed (immutable) objects get new names when they change, so only
+    short-cache objects need purging.
+    """
+    base = site_url.rstrip("/") + "/"
+    urls = []
+    for key in plan.uploads + plan.deletes:
+        rel = key.removeprefix(prefix.rstrip("/") + "/" if prefix else "")
+        if cache_control_for(rel) == IMMUTABLE_CACHE:
+            continue
+        urls.append(base + rel)
+        if rel == "index.html" or rel.endswith("/index.html"):
+            urls.append(base + rel.removesuffix("index.html"))
+    return urls
+
+
+def purge_cloudflare_cache(
+    zone_id: str,
+    api_token: str,
+    urls: Sequence[str],
+    *,
+    opener: Callable[..., Any] = urllib.request.urlopen,
+) -> None:
+    """Purge ``urls`` from Cloudflare's edge cache (no-op for an empty list)."""
+    for start in range(0, len(urls), PURGE_BATCH):
+        request = urllib.request.Request(
+            f"https://api.cloudflare.com/client/v4/zones/{zone_id}/purge_cache",
+            data=json.dumps({"files": list(urls[start : start + PURGE_BATCH])}).encode(),
+            headers={"Authorization": f"Bearer {api_token}", "Content-Type": "application/json"},
+            method="POST",
+        )
+        try:
+            with opener(request, timeout=30) as response:
+                body = json.load(response)
+        except Exception as exc:
+            raise SyncError(f"Cloudflare cache purge failed: {exc}") from exc
+        if not body.get("success"):
+            raise SyncError(f"Cloudflare cache purge failed: {body.get('errors')}")