brainpatch 1.2.0__py3-none-any.whl

brainpatch/config.py

@@ -0,0 +1,245 @@
+"""YAML experiment configuration with dotted CLI overrides.
+
+Every BrainPatch experiment is fully described by a YAML file plus a list of
+``key.path=value`` overrides. The resolved dictionary is written into the
+experiment directory on the Volume, so a run can always be reconstructed from
+its own artifacts.
+
+The loader is intentionally boring: no interpolation, no plugins, no imports.
+It has to run on a machine with nothing but ``pyyaml`` installed.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Iterable, Mapping
+
+try:  # pragma: no cover - exercised implicitly; guard keeps import failures readable
+    import yaml
+except ModuleNotFoundError as exc:  # pragma: no cover
+    raise ModuleNotFoundError(
+        "brainpatch.config requires PyYAML. Install the tiny control-plane extra: "
+        "pip install 'brainpatch[modal]'"
+    ) from exc
+
+
+class ConfigError(ValueError):
+    """Raised when a configuration file or override is unusable."""
+
+
+# ---------------------------------------------------------------------------
+# loading
+# ---------------------------------------------------------------------------
+
+
+def load_yaml(path: str | os.PathLike[str]) -> dict[str, Any]:
+    """Load a YAML file into a plain dict.
+
+    Raises
+    ------
+    ConfigError
+        If the file is missing, unparseable, or does not contain a mapping.
+    """
+    p = Path(path)
+    if not p.is_file():
+        raise ConfigError(f"config file not found: {p}")
+    try:
+        data = yaml.safe_load(p.read_text(encoding="utf-8"))
+    except yaml.YAMLError as exc:
+        raise ConfigError(f"could not parse {p}: {exc}") from exc
+    if data is None:
+        return {}
+    if not isinstance(data, dict):
+        raise ConfigError(f"{p} must contain a YAML mapping at the top level, got {type(data).__name__}")
+    return data
+
+
+def deep_merge(base: Mapping[str, Any], override: Mapping[str, Any]) -> dict[str, Any]:
+    """Recursively merge ``override`` into ``base``, returning a new dict.
+
+    Nested mappings merge key-by-key; every other type is replaced wholesale.
+
+    >>> deep_merge({"a": {"x": 1, "y": 2}}, {"a": {"y": 3}})
+    {'a': {'x': 1, 'y': 3}}
+    """
+    result: dict[str, Any] = dict(base)
+    for key, value in override.items():
+        existing = result.get(key)
+        if isinstance(existing, Mapping) and isinstance(value, Mapping):
+            result[key] = deep_merge(existing, value)
+        else:
+            result[key] = value
+    return result
+
+
+def _coerce_scalar(raw: str, value: Any) -> Any:
+    """Recover numeric types that YAML 1.1 leaves as strings.
+
+    ``yaml.safe_load`` follows the YAML 1.1 float grammar, which requires a
+    decimal point and a signed exponent: ``1.0e-4`` parses as a float but
+    ``1e-4`` parses as the *string* ``"1e-4"``. That is a nasty failure mode for
+    a config override -- ``--set sae.lr=1e-4`` would silently hand a string to
+    the optimizer instead of a learning rate. Python's own float grammar is more
+    permissive, so retry with it.
+    """
+    if not isinstance(value, str):
+        return value
+    text = value.strip()
+    if not text:
+        return value
+    try:
+        return int(text)
+    except ValueError:
+        pass
+    try:
+        return float(text)
+    except ValueError:
+        return value
+
+
+def parse_override(text: str) -> tuple[list[str], Any]:
+    """Parse a ``a.b.c=value`` override into a key path and a typed value.
+
+    Values are parsed as YAML scalars with a numeric-coercion fallback, so
+    ``k=32`` gives an int, ``lr=1e-4`` a float, ``flag=true`` a bool, and
+    ``name=smoke_v0`` a string.
+
+    >>> parse_override("sae.k=32")
+    (['sae', 'k'], 32)
+    >>> parse_override("run.force=true")
+    (['run', 'force'], True)
+    >>> parse_override("sae.lr=1e-4")
+    (['sae', 'lr'], 0.0001)
+    """
+    if "=" not in text:
+        raise ConfigError(f"override {text!r} must look like 'key.path=value'")
+    key, _, raw = text.partition("=")
+    key = key.strip()
+    if not key:
+        raise ConfigError(f"override {text!r} has an empty key")
+    try:
+        value = yaml.safe_load(raw)
+    except yaml.YAMLError as exc:
+        raise ConfigError(f"could not parse override value in {text!r}: {exc}") from exc
+    return key.split("."), _coerce_scalar(raw, value)
+
+
+def apply_overrides(config: Mapping[str, Any], overrides: Iterable[str]) -> dict[str, Any]:
+    """Apply dotted ``key=value`` overrides on top of ``config``."""
+    result = json.loads(json.dumps(config)) if config else {}
+    for override in overrides:
+        path, value = parse_override(override)
+        cursor = result
+        for part in path[:-1]:
+            nxt = cursor.get(part)
+            if not isinstance(nxt, dict):
+                nxt = {}
+                cursor[part] = nxt
+            cursor = nxt
+        cursor[path[-1]] = value
+    return result
+
+
+def load_config(
+    path: str | os.PathLike[str] | None,
+    overrides: Iterable[str] = (),
+    *,
+    defaults: Mapping[str, Any] | None = None,
+) -> dict[str, Any]:
+    """Resolve defaults -> YAML file -> CLI overrides into one dict."""
+    config: dict[str, Any] = dict(defaults or {})
+    if path is not None:
+        config = deep_merge(config, load_yaml(path))
+    return apply_overrides(config, overrides)
+
+
+def require(config: Mapping[str, Any], dotted_key: str) -> Any:
+    """Fetch ``a.b.c`` from a nested config, with a clear error if absent."""
+    cursor: Any = config
+    for part in dotted_key.split("."):
+        if not isinstance(cursor, Mapping) or part not in cursor:
+            raise ConfigError(f"missing required config key {dotted_key!r}")
+        cursor = cursor[part]
+    return cursor
+
+
+def get(config: Mapping[str, Any], dotted_key: str, default: Any = None) -> Any:
+    """Fetch ``a.b.c`` from a nested config, returning ``default`` if absent."""
+    try:
+        return require(config, dotted_key)
+    except ConfigError:
+        return default
+
+
+# ---------------------------------------------------------------------------
+# provenance
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class RunProvenance:
+    """Everything needed to say *what exactly* produced an artifact."""
+
+    git_commit: str | None = None
+    git_dirty: bool | None = None
+    package_versions: dict[str, str] = field(default_factory=dict)
+    gpu: str | None = None
+    hostname: str | None = None
+    started_at: str | None = None
+    finished_at: str | None = None
+    extra: dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "git_commit": self.git_commit,
+            "git_dirty": self.git_dirty,
+            "package_versions": dict(self.package_versions),
+            "gpu": self.gpu,
+            "hostname": self.hostname,
+            "started_at": self.started_at,
+            "finished_at": self.finished_at,
+            **self.extra,
+        }
+
+
+def git_commit(repo: str | os.PathLike[str] | None = None) -> str | None:
+    """Return the current git commit SHA, or ``None`` outside a repository.
+
+    Never raises: provenance capture must not be able to fail a run.
+    """
+    try:
+        out = subprocess.run(
+            ["git", "rev-parse", "HEAD"],
+            cwd=str(repo) if repo else None,
+            capture_output=True,
+            text=True,
+            timeout=10,
+            check=False,
+        )
+    except (OSError, subprocess.SubprocessError):
+        return None
+    if out.returncode != 0:
+        return None
+    return out.stdout.strip() or None
+
+
+def git_is_dirty(repo: str | os.PathLike[str] | None = None) -> bool | None:
+    """True if the working tree has uncommitted changes; ``None`` if unknown."""
+    try:
+        out = subprocess.run(
+            ["git", "status", "--porcelain"],
+            cwd=str(repo) if repo else None,
+            capture_output=True,
+            text=True,
+            timeout=10,
+            check=False,
+        )
+    except (OSError, subprocess.SubprocessError):
+        return None
+    if out.returncode != 0:
+        return None
+    return bool(out.stdout.strip())

brainpatch/datasets/__init__.py

@@ -0,0 +1,20 @@
+"""Dataset helpers.
+
+Contrast-set loading is pure Python and lives here. Text-corpus ingestion for
+activation extraction needs ``datasets`` and lives in
+:mod:`brainpatch.research.ml.corpus`, which is only imported inside Modal.
+"""
+
+from brainpatch.datasets.contrast_sets import (
+    CONTRAST_SET_NAMES,
+    default_contrast_dir,
+    list_contrast_sets,
+    load_contrast_set,
+)
+
+__all__ = [
+    "CONTRAST_SET_NAMES",
+    "default_contrast_dir",
+    "list_contrast_sets",
+    "load_contrast_set",
+]

brainpatch/datasets/contrast_sets.py

@@ -0,0 +1,64 @@
+"""Loading behavioural contrast fixtures.
+
+The sets shipped in ``examples/contrast/`` are small, synthetic, hand-written
+development fixtures. They are the input to candidate-feature search, and they
+are *not* benchmarks -- see the module docstring of
+:mod:`brainpatch.schemas.contrast`.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+from brainpatch.schemas.contrast import ContrastSet
+
+#: Fixtures shipped with the repository.
+CONTRAST_SET_NAMES: tuple[str, ...] = (
+    "sycophancy",
+    "verification",
+    "verbosity",
+    "contradiction",
+)
+
+
+def default_contrast_dir() -> Path:
+    """Directory holding the shipped contrast fixtures.
+
+    Resolved relative to the installed package so it works from a checkout and
+    from inside a Modal container where the repo lives at ``/root``.
+    """
+    here = Path(__file__).resolve()
+    # brainpatch/datasets/contrast_sets.py -> repo root
+    repo_root = here.parent.parent.parent
+    return repo_root / "examples" / "contrast"
+
+
+def list_contrast_sets(directory: str | os.PathLike[str] | None = None) -> list[str]:
+    """Names of every contrast set found in ``directory``."""
+    d = Path(directory) if directory is not None else default_contrast_dir()
+    if not d.is_dir():
+        return []
+    return sorted(p.stem for p in d.glob("*.json"))
+
+
+def load_contrast_set(
+    name: str, directory: str | os.PathLike[str] | None = None
+) -> ContrastSet:
+    """Load a contrast set by name, validating its contents.
+
+    Raises
+    ------
+    FileNotFoundError
+        If no such set exists, listing what is available.
+    """
+    d = Path(directory) if directory is not None else default_contrast_dir()
+    path = d / f"{name}.json"
+    if not path.is_file():
+        available = list_contrast_sets(d)
+        raise FileNotFoundError(
+            f"contrast set {name!r} not found at {path}. Available: {available or 'none'}"
+        )
+    contrast_set = ContrastSet.from_json(path.read_text(encoding="utf-8"))
+    contrast_set.validate()
+    return contrast_set

brainpatch/evaluation/__init__.py

@@ -0,0 +1,28 @@
+"""Evaluation utilities.
+
+Pure-Python text metrics live in :mod:`brainpatch.evaluation.metrics` and can be
+computed anywhere. Model-dependent measurements (log-probabilities, capability
+probes) require torch and live in :mod:`brainpatch.research.ml.evaluation`.
+"""
+
+from brainpatch.evaluation.metrics import (
+    GenerationMetrics,
+    compare_generations,
+    distinct_n,
+    jaccard_similarity,
+    longest_repeated_ngram,
+    most_common_ngram_fraction,
+    repetition_rate,
+    score_generation,
+)
+
+__all__ = [
+    "GenerationMetrics",
+    "compare_generations",
+    "distinct_n",
+    "jaccard_similarity",
+    "longest_repeated_ngram",
+    "most_common_ngram_fraction",
+    "repetition_rate",
+    "score_generation",
+]

brainpatch/evaluation/metrics.py

@@ -0,0 +1,223 @@
+"""Model-free text metrics for detecting intervention side effects.
+
+A large enough residual-stream perturbation degrades a model's fluency, and it
+can do so before it changes the behaviour you were aiming at -- observed here at
+steering strength 32, where the model looped. These metrics are the cheap
+tripwires: they catch degeneration (loops, single-token spam, truncation)
+without needing a judge model or a paid API.
+
+They are *not* measures of quality. A high `distinct_2` does not mean the answer
+is good; a low one strongly suggests the answer is broken. Use them to reject
+interventions, not to award them.
+"""
+
+from __future__ import annotations
+
+import math
+import re
+from dataclasses import asdict, dataclass
+from typing import Any, Sequence
+
+_WORD_RE = re.compile(r"\w+(?:'\w+)?|[^\w\s]")
+
+
+def tokenize_words(text: str) -> list[str]:
+    """Cheap whitespace/punctuation tokenizer, lowercased.
+
+    Deliberately model-independent so a metric never depends on which tokenizer
+    happened to be loaded.
+
+    >>> tokenize_words("Hello, world! Hello.")
+    ['hello', ',', 'world', '!', 'hello', '.']
+    """
+    return _WORD_RE.findall(text.lower())
+
+
+def distinct_n(tokens: Sequence[str], n: int = 2) -> float:
+    """Ratio of unique n-grams to total n-grams; 1.0 means no repetition.
+
+    Returns 0.0 when the text is too short to contain an n-gram.
+    """
+    if n <= 0:
+        raise ValueError(f"n must be positive, got {n}")
+    if len(tokens) < n:
+        return 0.0
+    grams = [tuple(tokens[i : i + n]) for i in range(len(tokens) - n + 1)]
+    return len(set(grams)) / len(grams)
+
+
+def repetition_rate(tokens: Sequence[str], n: int = 3) -> float:
+    """Fraction of n-grams that occur more than once.
+
+    The complement of a "novel n-gram" rate. High values indicate looping.
+    """
+    if len(tokens) < n:
+        return 0.0
+    grams = [tuple(tokens[i : i + n]) for i in range(len(tokens) - n + 1)]
+    counts: dict[tuple[str, ...], int] = {}
+    for gram in grams:
+        counts[gram] = counts.get(gram, 0) + 1
+    repeated = sum(count for count in counts.values() if count > 1)
+    return repeated / len(grams)
+
+
+def longest_repeated_ngram(tokens: Sequence[str], max_n: int = 20) -> int:
+    """Length of the longest n-gram that appears at least twice.
+
+    A blunt but reliable degeneration detector: healthy prose rarely repeats a
+    10-gram, a looping model repeats one constantly.
+    """
+    best = 0
+    upper = min(max_n, len(tokens) // 2)
+    for n in range(1, upper + 1):
+        seen: set[tuple[str, ...]] = set()
+        found = False
+        for i in range(len(tokens) - n + 1):
+            gram = tuple(tokens[i : i + n])
+            if gram in seen:
+                found = True
+                break
+            seen.add(gram)
+        if found:
+            best = n
+        else:
+            # No repeat of length n means no repeat of any longer length.
+            break
+    return best
+
+
+def most_common_ngram_fraction(tokens: Sequence[str], n: int = 2) -> float:
+    """Share of all n-grams taken by the single most frequent one.
+
+    Added after an observed miss: at high steering strength the model produced
+    ``"as they encounter each other, as they interact with each other, as they
+    collide, as they merge, ..."`` -- obviously degenerate, yet it passed the
+    distinct-n and longest-repeat checks because each clause ends differently.
+    A single bigram occupying ~18% of all bigrams catches that pattern, where
+    set-based diversity measures do not.
+    """
+    if len(tokens) < n:
+        return 0.0
+    grams = [tuple(tokens[i : i + n]) for i in range(len(tokens) - n + 1)]
+    counts: dict[tuple[str, ...], int] = {}
+    for gram in grams:
+        counts[gram] = counts.get(gram, 0) + 1
+    return max(counts.values()) / len(grams)
+
+
+def type_token_ratio(tokens: Sequence[str]) -> float:
+    """Unique tokens / total tokens. 0.0 for empty input."""
+    if not tokens:
+        return 0.0
+    return len(set(tokens)) / len(tokens)
+
+
+def shannon_entropy(tokens: Sequence[str]) -> float:
+    """Unigram entropy in bits. Collapsed output has near-zero entropy."""
+    if not tokens:
+        return 0.0
+    counts: dict[str, int] = {}
+    for token in tokens:
+        counts[token] = counts.get(token, 0) + 1
+    total = len(tokens)
+    return -sum((c / total) * math.log2(c / total) for c in counts.values())
+
+
+@dataclass
+class GenerationMetrics:
+    """Model-free summary of one generated string."""
+
+    num_chars: int
+    num_words: int
+    distinct_1: float
+    distinct_2: float
+    distinct_3: float
+    repetition_3: float
+    longest_repeat: int
+    type_token_ratio: float
+    entropy: float
+    top_bigram_fraction: float
+    is_empty: bool
+
+    @property
+    def degeneration_flag(self) -> bool:
+        """Heuristic tripwire for obviously broken output.
+
+        Thresholds are conservative: they fire on text that is plainly looping
+        or empty, not on text that is merely repetitive. A True here means
+        "inspect this generation", not "this intervention failed".
+
+        These are heuristics, and they have been observed to miss things -- the
+        ``top_bigram_fraction`` clause exists because the earlier rules scored a
+        clearly-looping generation as clean. Treat a False as "no obvious
+        breakage detected", not as "output is fine".
+        """
+        if self.is_empty:
+            return True
+        if self.num_words >= 30 and self.distinct_2 < 0.35:
+            return True
+        if self.longest_repeat >= 10:
+            return True
+        if self.num_words >= 20 and self.entropy < 2.0:
+            return True
+        if self.num_words >= 30 and self.top_bigram_fraction > 0.10:
+            return True
+        return False
+
+    def to_dict(self) -> dict[str, Any]:
+        data = asdict(self)
+        data["degeneration_flag"] = self.degeneration_flag
+        return data
+
+
+def score_generation(text: str) -> GenerationMetrics:
+    """Compute all model-free metrics for one generated string."""
+    tokens = tokenize_words(text)
+    return GenerationMetrics(
+        num_chars=len(text),
+        num_words=len(tokens),
+        distinct_1=distinct_n(tokens, 1),
+        distinct_2=distinct_n(tokens, 2),
+        distinct_3=distinct_n(tokens, 3),
+        repetition_3=repetition_rate(tokens, 3),
+        longest_repeat=longest_repeated_ngram(tokens),
+        type_token_ratio=type_token_ratio(tokens),
+        entropy=shannon_entropy(tokens),
+        top_bigram_fraction=most_common_ngram_fraction(tokens, 2),
+        is_empty=not text.strip(),
+    )
+
+
+def jaccard_similarity(a: str, b: str, n: int = 3) -> float:
+    """n-gram Jaccard overlap between two strings, in ``[0, 1]``.
+
+    Used to quantify *semantic drift*: how far an intervened generation moved
+    from its baseline. 1.0 means identical n-gram sets.
+    """
+    ta, tb = tokenize_words(a), tokenize_words(b)
+    if len(ta) < n or len(tb) < n:
+        return 1.0 if a.strip() == b.strip() else 0.0
+    ga = {tuple(ta[i : i + n]) for i in range(len(ta) - n + 1)}
+    gb = {tuple(tb[i : i + n]) for i in range(len(tb) - n + 1)}
+    union = ga | gb
+    if not union:
+        return 1.0
+    return len(ga & gb) / len(union)
+
+
+def compare_generations(baseline: str, intervened: str) -> dict[str, Any]:
+    """Side-by-side comparison of a baseline and an intervened generation."""
+    base_metrics = score_generation(baseline)
+    int_metrics = score_generation(intervened)
+    return {
+        "baseline": base_metrics.to_dict(),
+        "intervened": int_metrics.to_dict(),
+        "identical": baseline == intervened,
+        "jaccard_3": jaccard_similarity(baseline, intervened, n=3),
+        "length_ratio": (
+            int_metrics.num_words / base_metrics.num_words if base_metrics.num_words else None
+        ),
+        "distinct_2_delta": int_metrics.distinct_2 - base_metrics.distinct_2,
+        "entropy_delta": int_metrics.entropy - base_metrics.entropy,
+        "degeneration_introduced": int_metrics.degeneration_flag and not base_metrics.degeneration_flag,
+    }

brainpatch/patch/__init__.py

@@ -0,0 +1,64 @@
+"""The portable BrainPatch artifact: format, loading, validation, registry.
+
+Everything here is importable with **no ML stack installed**. Parsing a patch,
+checking its compatibility metadata, installing it, and reporting its size all
+work on a bare Python 3.10+.
+
+The one exception is :mod:`brainpatch.patch.compiler`, which reads SAE
+checkpoints and therefore needs torch. It is not imported here; import it
+explicitly when you mean to compile.
+"""
+
+from brainpatch.patch.format import (
+    ABSOLUTE_MAX_STRENGTH,
+    FORMAT_VERSION,
+    SUFFIX,
+    BaseModelSpec,
+    Intervention,
+    Manifest,
+    PatchFormatError,
+)
+from brainpatch.patch.loader import (
+    LoadedPatch,
+    PatchLoadError,
+    load_patch,
+    patch_size_report,
+    save_patch,
+)
+from brainpatch.patch.registry import (
+    InstalledPatch,
+    PatchRegistry,
+    RegistryError,
+    default_registry,
+    registry_home,
+)
+from brainpatch.patch.validation import (
+    CompatibilityReport,
+    ModelDescriptor,
+    PatchCompatibilityError,
+    check_compatibility,
+)
+
+__all__ = [
+    "ABSOLUTE_MAX_STRENGTH",
+    "BaseModelSpec",
+    "CompatibilityReport",
+    "FORMAT_VERSION",
+    "InstalledPatch",
+    "Intervention",
+    "LoadedPatch",
+    "Manifest",
+    "ModelDescriptor",
+    "PatchCompatibilityError",
+    "PatchFormatError",
+    "PatchLoadError",
+    "PatchRegistry",
+    "RegistryError",
+    "SUFFIX",
+    "check_compatibility",
+    "default_registry",
+    "load_patch",
+    "patch_size_report",
+    "registry_home",
+    "save_patch",
+]