brainpatch 1.2.0__py3-none-any.whl

brainpatch/patch/tensors.py

@@ -0,0 +1,236 @@
+"""Minimal, dependency-free safetensors reader and writer.
+
+Why this exists
+---------------
+A BrainPatch runtime artifact is a handful of vectors. Reading them should not
+force a user to install numpy, let alone torch: ``pip install brainpatch`` has
+to stay light enough that inspecting, validating and installing a patch works
+on any Python 3.10+ with nothing else present.
+
+The safetensors container is simple enough to implement directly:
+
+===============  =========================================================
+8 bytes          little-endian ``uint64`` header length ``N``
+``N`` bytes      UTF-8 JSON header
+remainder        raw little-endian tensor data
+===============  =========================================================
+
+Each header entry is ``{"dtype": ..., "shape": [...], "data_offsets": [a, b]}``
+with offsets relative to the end of the header. ``__metadata__`` is an optional
+reserved key holding a flat ``str -> str`` map.
+
+This module deliberately supports **only** the float dtypes a patch vector can
+use. It is not a general safetensors implementation, and it never executes
+anything from the file -- the whole point of the format choice is that a patch
+is inert data.
+"""
+
+from __future__ import annotations
+
+import json
+import struct
+from dataclasses import dataclass
+from typing import Iterable, Sequence
+
+#: dtype name -> (bytes per element, struct format or None for bf16)
+_DTYPES: dict[str, tuple[int, str | None]] = {
+    "F64": (8, "<d"),
+    "F32": (4, "<f"),
+    "F16": (2, "<e"),
+    "BF16": (2, None),  # handled specially: upper 16 bits of an f32
+}
+
+#: Refuse absurd headers before allocating anything.
+MAX_HEADER_BYTES = 16 * 1024 * 1024
+
+
+class SafetensorsError(ValueError):
+    """The byte stream is not a well-formed safetensors container."""
+
+
+@dataclass(frozen=True)
+class Tensor:
+    """A dense float tensor as plain Python floats.
+
+    Attributes
+    ----------
+    dtype:
+        The dtype it was *stored* as. Values are always decoded to Python
+        floats; this records the on-disk precision so a round trip can preserve
+        it and so a caller can report the real artifact size.
+    """
+
+    dtype: str
+    shape: tuple[int, ...]
+    data: list[float]
+
+    def __post_init__(self) -> None:
+        expected = 1
+        for dim in self.shape:
+            expected *= dim
+        if len(self.data) != expected:
+            raise SafetensorsError(
+                f"shape {self.shape} implies {expected} elements but got {len(self.data)}"
+            )
+
+    @property
+    def numel(self) -> int:
+        return len(self.data)
+
+    @property
+    def nbytes(self) -> int:
+        return self.numel * _DTYPES[self.dtype][0]
+
+
+def _decode_bf16(raw: bytes, count: int) -> list[float]:
+    """bfloat16 is the top 16 bits of an IEEE-754 float32."""
+    out: list[float] = []
+    for i in range(count):
+        bits = raw[2 * i] | (raw[2 * i + 1] << 8)
+        out.append(struct.unpack("<f", struct.pack("<I", bits << 16))[0])
+    return out
+
+
+def _encode_bf16(values: Sequence[float]) -> bytes:
+    """Round-to-nearest-even truncation of float32 to its top 16 bits."""
+    out = bytearray()
+    for value in values:
+        bits = struct.unpack("<I", struct.pack("<f", float(value)))[0]
+        # round-to-nearest-even on the discarded low 16 bits
+        rounding = 0x7FFF + ((bits >> 16) & 1)
+        bits = (bits + rounding) >> 16
+        out += struct.pack("<H", bits & 0xFFFF)
+    return bytes(out)
+
+
+def load(blob: bytes) -> tuple[dict[str, Tensor], dict[str, str]]:
+    """Parse a safetensors byte string.
+
+    Returns
+    -------
+    (tensors, metadata)
+
+    Raises
+    ------
+    SafetensorsError
+        On any structural problem. Offsets are validated against the actual
+        payload length, so a truncated or overlapping file is rejected rather
+        than silently yielding garbage.
+    """
+    if len(blob) < 8:
+        raise SafetensorsError("file is too short to contain a header length")
+
+    (header_len,) = struct.unpack("<Q", blob[:8])
+    if header_len > MAX_HEADER_BYTES:
+        raise SafetensorsError(f"header claims {header_len} bytes, refusing to read")
+    if 8 + header_len > len(blob):
+        raise SafetensorsError("header length exceeds file size")
+
+    try:
+        header = json.loads(blob[8 : 8 + header_len].decode("utf-8"))
+    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
+        raise SafetensorsError(f"header is not valid UTF-8 JSON: {exc}") from exc
+    if not isinstance(header, dict):
+        raise SafetensorsError("header must be a JSON object")
+
+    payload = blob[8 + header_len :]
+    metadata_raw = header.pop("__metadata__", {})
+    if not isinstance(metadata_raw, dict):
+        raise SafetensorsError("__metadata__ must be a JSON object")
+    metadata = {str(k): str(v) for k, v in metadata_raw.items()}
+
+    tensors: dict[str, Tensor] = {}
+    for name, spec in header.items():
+        if not isinstance(spec, dict):
+            raise SafetensorsError(f"entry {name!r} is not an object")
+        dtype = spec.get("dtype")
+        if dtype not in _DTYPES:
+            raise SafetensorsError(
+                f"tensor {name!r} has unsupported dtype {dtype!r}; "
+                f"supported: {sorted(_DTYPES)}"
+            )
+        shape = spec.get("shape")
+        if not isinstance(shape, list) or not all(
+            isinstance(d, int) and d >= 0 for d in shape
+        ):
+            raise SafetensorsError(f"tensor {name!r} has an invalid shape {shape!r}")
+        offsets = spec.get("data_offsets")
+        if (
+            not isinstance(offsets, list)
+            or len(offsets) != 2
+            or not all(isinstance(o, int) for o in offsets)
+        ):
+            raise SafetensorsError(f"tensor {name!r} has invalid data_offsets")
+
+        start, end = offsets
+        if not 0 <= start <= end <= len(payload):
+            raise SafetensorsError(
+                f"tensor {name!r} offsets [{start}, {end}] fall outside the "
+                f"{len(payload)}-byte payload"
+            )
+
+        item_size, fmt = _DTYPES[dtype]
+        raw = payload[start:end]
+        if len(raw) % item_size:
+            raise SafetensorsError(f"tensor {name!r} byte length is not a multiple of {item_size}")
+        count = len(raw) // item_size
+
+        expected = 1
+        for dim in shape:
+            expected *= dim
+        if count != expected:
+            raise SafetensorsError(
+                f"tensor {name!r} shape {shape} implies {expected} elements "
+                f"but the byte range holds {count}"
+            )
+
+        if fmt is None:
+            values = _decode_bf16(raw, count)
+        else:
+            values = [v[0] for v in struct.iter_unpack(fmt, raw)]
+
+        tensors[name] = Tensor(dtype=dtype, shape=tuple(shape), data=values)
+
+    return tensors, metadata
+
+
+def dump(tensors: dict[str, Tensor], metadata: dict[str, str] | None = None) -> bytes:
+    """Serialize tensors into a safetensors byte string.
+
+    Tensor names are written in sorted order so the output is byte-deterministic
+    for a given input -- which is what lets a patch carry a stable checksum.
+    """
+    header: dict[str, object] = {}
+    if metadata:
+        header["__metadata__"] = {str(k): str(v) for k, v in metadata.items()}
+
+    body = bytearray()
+    for name in sorted(tensors):
+        tensor = tensors[name]
+        item_size, fmt = _DTYPES[tensor.dtype]
+        if fmt is None:
+            raw = _encode_bf16(tensor.data)
+        else:
+            raw = b"".join(struct.pack(fmt, float(v)) for v in tensor.data)
+        start = len(body)
+        body += raw
+        header[name] = {
+            "dtype": tensor.dtype,
+            "shape": list(tensor.shape),
+            "data_offsets": [start, len(body)],
+        }
+
+    header_bytes = json.dumps(header, sort_keys=True, separators=(",", ":")).encode("utf-8")
+    return struct.pack("<Q", len(header_bytes)) + header_bytes + bytes(body)
+
+
+def vector(values: Iterable[float], dtype: str = "F32") -> Tensor:
+    """Build a 1-D tensor from an iterable of floats."""
+    if dtype not in _DTYPES:
+        raise SafetensorsError(f"unsupported dtype {dtype!r}")
+    data = [float(v) for v in values]
+    return Tensor(dtype=dtype, shape=(len(data),), data=data)
+
+
+def supported_dtypes() -> tuple[str, ...]:
+    return tuple(sorted(_DTYPES))

brainpatch/patch/validation.py

@@ -0,0 +1,157 @@
+"""Patch-to-model compatibility checking.
+
+A patch vector is a direction in one model's residual basis at one layer. The
+same architecture with different weights has a different basis, so "it is also a
+Qwen2ForCausalLM" is not evidence that a direction transfers. Applying a patch
+to the wrong model does not degrade gracefully -- it adds an arbitrary vector and
+produces confident nonsense, which is worse than an error.
+
+Hence three explicit modes, defaulting to the strictest:
+
+``strict`` (default)
+    Model id must match. Revision must match when both are known. Hidden size
+    and layer count must match.
+
+``architecture``
+    Model id may differ but the architecture string, hidden size and layer count
+    must match. For fine-tunes and merges of the same base. Warns.
+
+``unsafe``
+    Only the geometry is checked -- hidden size must match and the layer must
+    exist, because anything else cannot be executed at all. Warns loudly.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+from brainpatch.patch.format import Manifest
+
+CompatibilityMode = Literal["strict", "architecture", "unsafe"]
+COMPATIBILITY_MODES: tuple[str, ...] = ("strict", "architecture", "unsafe")
+
+
+class PatchCompatibilityError(ValueError):
+    """The patch does not match the model it is being applied to."""
+
+
+@dataclass
+class ModelDescriptor:
+    """What the backend knows about the loaded model."""
+
+    model_id: str
+    hidden_size: int
+    num_layers: int
+    architecture: str = ""
+    revision: str | None = None
+
+
+@dataclass
+class CompatibilityReport:
+    """Outcome of a compatibility check."""
+
+    ok: bool
+    mode: str
+    errors: list[str] = field(default_factory=list)
+    warnings: list[str] = field(default_factory=list)
+
+    def raise_if_failed(self) -> None:
+        if not self.ok:
+            raise PatchCompatibilityError(
+                "patch is not compatible with this model:\n  - "
+                + "\n  - ".join(self.errors)
+                + f"\n\nChecked in '{self.mode}' mode. If you understand the risk, "
+                "re-run with compatibility_mode='architecture' or 'unsafe'."
+            )
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "ok": self.ok,
+            "mode": self.mode,
+            "errors": list(self.errors),
+            "warnings": list(self.warnings),
+        }
+
+
+def check_compatibility(
+    manifest: Manifest,
+    model: ModelDescriptor,
+    *,
+    mode: CompatibilityMode = "strict",
+) -> CompatibilityReport:
+    """Check a patch against a loaded model. Never raises; inspect ``.ok``."""
+    if mode not in COMPATIBILITY_MODES:
+        raise ValueError(f"unknown compatibility mode {mode!r}; expected {COMPATIBILITY_MODES}")
+
+    errors: list[str] = []
+    warnings: list[str] = []
+    spec = manifest.base_model
+
+    # -- geometry: required in every mode, because it is what makes the
+    #    arithmetic executable at all.
+    if spec.hidden_size != model.hidden_size:
+        errors.append(
+            f"hidden size mismatch: patch was built for {spec.hidden_size}, "
+            f"model has {model.hidden_size}"
+        )
+    for layer in manifest.layers:
+        if layer >= model.num_layers:
+            errors.append(
+                f"patch targets layer {layer} but the model has {model.num_layers} layers"
+            )
+
+    if mode == "unsafe":
+        if spec.model_id != model.model_id:
+            warnings.append(
+                f"UNSAFE MODE: applying a patch built for {spec.model_id!r} to "
+                f"{model.model_id!r}. Feature directions are not transferable between "
+                "models; output is not meaningful."
+            )
+        return CompatibilityReport(ok=not errors, mode=mode, errors=errors, warnings=warnings)
+
+    # -- architecture-level checks
+    if spec.num_layers and spec.num_layers != model.num_layers:
+        errors.append(
+            f"layer count mismatch: patch declares {spec.num_layers}, "
+            f"model has {model.num_layers}"
+        )
+    if spec.architecture and model.architecture and spec.architecture != model.architecture:
+        errors.append(
+            f"architecture mismatch: patch is for {spec.architecture!r}, "
+            f"model is {model.architecture!r}"
+        )
+
+    if mode == "architecture":
+        if spec.model_id != model.model_id:
+            warnings.append(
+                f"model id differs (patch: {spec.model_id!r}, loaded: {model.model_id!r}). "
+                "Allowed in 'architecture' mode, but the direction was fitted on "
+                "different weights and may not transfer."
+            )
+        return CompatibilityReport(ok=not errors, mode=mode, errors=errors, warnings=warnings)
+
+    # -- strict
+    if spec.model_id != model.model_id:
+        errors.append(
+            f"model mismatch: patch targets {spec.model_id!r} but {model.model_id!r} "
+            "is loaded"
+        )
+    if spec.revision and model.revision and spec.revision != model.revision:
+        errors.append(
+            f"revision mismatch: patch was derived from {spec.revision[:12]}… but "
+            f"{model.revision[:12]}… is loaded"
+        )
+    elif spec.revision and not model.revision:
+        warnings.append(
+            f"patch pins revision {spec.revision[:12]}… but the loaded model's "
+            "revision is unknown, so it could not be checked"
+        )
+
+    return CompatibilityReport(ok=not errors, mode=mode, errors=errors, warnings=warnings)
+
+
+def validate_strength(manifest: Manifest, strength: float) -> float:
+    """Clamp a live strength to the patch's declared envelope, warning if clipped."""
+    clamped = manifest.clamp_strength(strength)
+    return clamped

brainpatch/paths.py

@@ -0,0 +1,184 @@
+"""Canonical layout of the BrainPatch Modal Volume.
+
+The Volume ``brainpatch-data`` is mounted at ``/vol`` inside Modal containers.
+All expensive or large artifacts live there; none of it is ever copied to a
+local machine.
+
+::
+
+    /vol/
+    |-- hf-cache/                     Hugging Face model + dataset cache
+    |-- datasets/                     preprocessed text corpora
+    |-- activations/<experiment>/     immutable activation shards + manifest
+    |-- sae/<experiment>/             SAE checkpoints + training metrics
+    |-- feature-db/<experiment>/      per-feature statistics and contexts
+    |-- patches/                      BrainPatch json files
+    |-- experiments/<experiment>/     causal-validation artifacts
+    `-- reports/                      generated markdown/html reports
+
+This module is deliberately dependency-free so it can be used both locally
+(to build CLI messages) and remotely (to actually touch the filesystem).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import PurePosixPath
+
+#: Mount point of the ``brainpatch-data`` Volume inside Modal containers.
+DEFAULT_VOLUME_ROOT = "/vol"
+
+#: Name of the Modal Volume holding all persistent artifacts.
+VOLUME_NAME = "brainpatch-data"
+
+#: Name of the Modal Secret exposing ``HF_TOKEN``.
+HF_SECRET_NAME = "huggingface-secret"
+
+#: Modal environment this project is developed in.
+MODAL_ENVIRONMENT = "brainpatch-dev"
+
+
+@dataclass(frozen=True)
+class VolumePaths:
+    """Resolver for every path on the BrainPatch Volume.
+
+    Parameters
+    ----------
+    root:
+        Volume mount point. Defaults to ``/vol``.
+
+    Notes
+    -----
+    Paths are :class:`~pathlib.PurePosixPath` because the Volume is always
+    mounted in a Linux container, even when this code is *constructed* on
+    Windows. Converting to :class:`pathlib.Path` only happens on the remote
+    side, where the platform is known to be POSIX.
+    """
+
+    root: str = DEFAULT_VOLUME_ROOT
+
+    # -- top-level directories ------------------------------------------------
+
+    @property
+    def base(self) -> PurePosixPath:
+        return PurePosixPath(self.root)
+
+    @property
+    def hf_cache(self) -> PurePosixPath:
+        return self.base / "hf-cache"
+
+    @property
+    def datasets(self) -> PurePosixPath:
+        return self.base / "datasets"
+
+    @property
+    def activations_root(self) -> PurePosixPath:
+        return self.base / "activations"
+
+    @property
+    def sae_root(self) -> PurePosixPath:
+        return self.base / "sae"
+
+    @property
+    def feature_db_root(self) -> PurePosixPath:
+        return self.base / "feature-db"
+
+    @property
+    def patches(self) -> PurePosixPath:
+        return self.base / "patches"
+
+    @property
+    def experiments_root(self) -> PurePosixPath:
+        return self.base / "experiments"
+
+    @property
+    def reports(self) -> PurePosixPath:
+        return self.base / "reports"
+
+    def all_top_level(self) -> tuple[PurePosixPath, ...]:
+        """Every directory that should exist on a freshly initialised Volume."""
+        return (
+            self.hf_cache,
+            self.datasets,
+            self.activations_root,
+            self.sae_root,
+            self.feature_db_root,
+            self.patches,
+            self.experiments_root,
+            self.reports,
+        )
+
+    # -- per-experiment directories -------------------------------------------
+
+    def activations(self, experiment: str) -> PurePosixPath:
+        """Directory holding activation shards for ``experiment``."""
+        return self.activations_root / experiment
+
+    def activation_manifest(self, experiment: str) -> PurePosixPath:
+        return self.activations(experiment) / "manifest.json"
+
+    def activation_examples(self, experiment: str) -> PurePosixPath:
+        """JSONL of source examples; token metadata references these by index."""
+        return self.activations(experiment) / "examples.jsonl"
+
+    def activation_shard(self, experiment: str, index: int) -> PurePosixPath:
+        """Immutable shard path. Shard names are stable and never rewritten."""
+        return self.activations(experiment) / shard_filename(index)
+
+    def sae(self, experiment: str) -> PurePosixPath:
+        return self.sae_root / experiment
+
+    def sae_checkpoint(self, experiment: str, name: str = "sae_latest.pt") -> PurePosixPath:
+        return self.sae(experiment) / name
+
+    def sae_config(self, experiment: str) -> PurePosixPath:
+        return self.sae(experiment) / "config.json"
+
+    def sae_metrics(self, experiment: str) -> PurePosixPath:
+        return self.sae(experiment) / "metrics.jsonl"
+
+    def feature_db(self, experiment: str) -> PurePosixPath:
+        return self.feature_db_root / experiment
+
+    def features_jsonl(self, experiment: str) -> PurePosixPath:
+        return self.feature_db(experiment) / "features.jsonl"
+
+    def feature_summary(self, experiment: str) -> PurePosixPath:
+        return self.feature_db(experiment) / "summary.json"
+
+    def experiment(self, experiment: str) -> PurePosixPath:
+        return self.experiments_root / experiment
+
+    def experiment_file(self, experiment: str, filename: str) -> PurePosixPath:
+        return self.experiment(experiment) / filename
+
+    def patch(self, name: str) -> PurePosixPath:
+        return self.patches / f"{name}.json"
+
+
+def shard_filename(index: int) -> str:
+    """Return the immutable filename for activation shard ``index``.
+
+    Shard names are zero-padded to six digits so lexical order matches numeric
+    order, which lets a streaming reader glob-and-sort without parsing.
+
+    >>> shard_filename(0)
+    'shard_000000.safetensors'
+    >>> shard_filename(42)
+    'shard_000042.safetensors'
+    """
+    if index < 0:
+        raise ValueError(f"shard index must be non-negative, got {index}")
+    return f"shard_{index:06d}.safetensors"
+
+
+def parse_shard_index(filename: str) -> int:
+    """Inverse of :func:`shard_filename`.
+
+    >>> parse_shard_index("shard_000042.safetensors")
+    42
+    """
+    stem = filename.rsplit("/", 1)[-1]
+    if not stem.startswith("shard_") or not stem.endswith(".safetensors"):
+        raise ValueError(f"not a shard filename: {filename!r}")
+    return int(stem[len("shard_") : -len(".safetensors")])

brainpatch/research/__init__.py

@@ -0,0 +1,16 @@
+"""Research toolkit: how BrainPatches are *created*.
+
+Everything here is for patch **authors** -- activation extraction, SAE training,
+feature discovery, causal validation, patch search. It needs torch and is
+installed by the ``research`` extra::
+
+    pip install "brainpatch[research]"
+
+**The runtime never imports this.** A user applying a patch needs none of it: no
+SAE, no activation corpus, no training code. That separation is the product.
+
+Where the compute runs is an author's choice. This repository's own experiments
+used Modal (see ``modal_app/``) because the development machine deliberately
+carries no ML stack, but nothing here requires it -- a local GPU, a cluster, or
+a notebook works equally well.
+"""