brainpatch 1.2.0__py3-none-any.whl

brainpatch/runtime/model.py

@@ -0,0 +1,260 @@
+"""``BrainPatchedModel`` -- the user-facing Python API.
+
+::
+
+    from brainpatch import BrainPatchedModel
+
+    model = BrainPatchedModel.from_pretrained(
+        "Qwen/Qwen2.5-1.5B-Instruct",
+        backend="transformers",
+        device="auto",
+    )
+
+    patch = model.install("09Catho/example-patch")
+    patch.strength = 0.8
+
+    print(model.generate("Evaluate my idea."))
+
+This is a thin facade over a :class:`~brainpatch.runtime.base.BrainPatchBackend`.
+It owns no intervention logic of its own -- that lives in the backend contract,
+so every engine behaves identically where it can and reports honestly where it
+cannot.
+
+Nothing here knows or cares where a patch was trained.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any, Iterator
+
+from brainpatch.patch.loader import LoadedPatch, load_patch
+from brainpatch.patch.registry import PatchRegistry, default_registry
+from brainpatch.patch.validation import CompatibilityMode
+from brainpatch.runtime.auto import select_backend
+from brainpatch.runtime.base import ActivePatch, BrainPatchBackend, GenerationConfig
+from brainpatch.runtime.capabilities import Capabilities
+from brainpatch.runtime.scheduling import StrengthSchedule
+
+
+class PatchHandle:
+    """Live control surface for one installed patch.
+
+    ``patch.strength = 0.8`` and ``patch.schedule = {...}`` are the ergonomic
+    forms; both route to the backend so clamping and capability checks still
+    apply.
+    """
+
+    def __init__(self, model: "BrainPatchedModel", name: str) -> None:
+        self._model = model
+        self._name = name
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def _active(self) -> ActivePatch:
+        return self._model.backend.patches[self._name]
+
+    @property
+    def strength(self) -> float:
+        return self._active.strength
+
+    @strength.setter
+    def strength(self, value: float) -> None:
+        self._model.backend.set_strength(self._name, value)
+
+    @property
+    def enabled(self) -> bool:
+        return self._active.enabled
+
+    @enabled.setter
+    def enabled(self, value: bool) -> None:
+        self._model.backend.set_enabled(self._name, bool(value))
+
+    @property
+    def schedule(self) -> StrengthSchedule | None:
+        return self._active.schedule
+
+    @schedule.setter
+    def schedule(self, value: StrengthSchedule | dict[int, float] | None) -> None:
+        self._model.backend.set_schedule(self._name, value)
+
+    @property
+    def manifest(self) -> Any:
+        return self._active.manifest
+
+    @property
+    def evidence_level(self) -> str:
+        return str(self._active.manifest.evidence_level)
+
+    def describe(self) -> dict[str, Any]:
+        return {
+            "name": self._name,
+            "strength": self.strength,
+            "enabled": self.enabled,
+            "evidence_level": self.evidence_level,
+            "layers": self._active.manifest.layers,
+            "scheduled": self.schedule is not None,
+        }
+
+    def __repr__(self) -> str:
+        state = "on" if self.enabled else "off"
+        return f"<PatchHandle {self._name!r} strength={self.strength:+.3f} [{state}]>"
+
+
+class BrainPatchedModel:
+    """A frozen language model with installable activation patches."""
+
+    def __init__(self, backend: BrainPatchBackend, registry: PatchRegistry | None = None) -> None:
+        self.backend = backend
+        self.registry = registry or default_registry()
+        self.compatibility_mode: CompatibilityMode = "strict"
+
+    # -- construction ----------------------------------------------------------
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model: str,
+        *,
+        backend: str = "auto",
+        device: str = "auto",
+        dtype: str = "auto",
+        revision: str | None = None,
+        registry: PatchRegistry | None = None,
+        compatibility_mode: CompatibilityMode = "strict",
+        **kwargs: Any,
+    ) -> "BrainPatchedModel":
+        """Load a frozen base model on the chosen (or best available) backend.
+
+        ``backend="auto"`` picks the first available engine. Naming a backend
+        that is unavailable raises rather than substituting a different one.
+        """
+        backend_cls = select_backend(backend)
+        instance = backend_cls()
+        instance.load_model(model, revision=revision, device=device, dtype=dtype, **kwargs)
+        patched = cls(instance, registry=registry)
+        patched.compatibility_mode = compatibility_mode
+        return patched
+
+    # -- patch management ------------------------------------------------------
+
+    def install(
+        self,
+        ref: str | os.PathLike[str] | LoadedPatch,
+        *,
+        strength: float | None = None,
+        compatibility_mode: CompatibilityMode | None = None,
+    ) -> PatchHandle:
+        """Install a patch by installed name, file path, HF reference, or object.
+
+        Resolution order: an already-installed registry name, then a filesystem
+        path, then a Hugging Face ``owner/repo`` reference (which downloads only
+        the patch artifact, never the base model).
+        """
+        if isinstance(ref, LoadedPatch):
+            loaded = ref
+        else:
+            text = str(ref)
+            try:
+                loaded = load_patch(self.registry.resolve(text))
+            except Exception:
+                # Not installed and not a local file: try the Hub, then install.
+                installed = self.registry.install(text)
+                loaded = installed.load()
+
+        active = self.backend.install_patch(
+            loaded,
+            strength=strength,
+            mode=compatibility_mode or self.compatibility_mode,
+        )
+        return PatchHandle(self, active.name)
+
+    def remove_patch(self, name: str) -> None:
+        self.backend.remove_patch(name)
+
+    def enable_patch(self, name: str) -> None:
+        self.backend.set_enabled(name, True)
+
+    def disable_patch(self, name: str) -> None:
+        self.backend.set_enabled(name, False)
+
+    def set_patch_strength(self, name: str, strength: float) -> float:
+        return self.backend.set_strength(name, strength)
+
+    def set_patch_schedule(
+        self, name: str, schedule: StrengthSchedule | dict[int, float] | None
+    ) -> None:
+        self.backend.set_schedule(name, schedule)
+
+    def list_patches(self) -> list[str]:
+        return self.backend.list_patches()
+
+    def patch(self, name: str) -> PatchHandle:
+        if name not in self.backend.patches:
+            raise KeyError(f"no patch named {name!r} is installed")
+        return PatchHandle(self, name)
+
+    # -- generation ------------------------------------------------------------
+
+    def generate(self, prompt: str, config: GenerationConfig | None = None, **kwargs: Any) -> str:
+        return self.backend.generate(prompt, config, **kwargs)
+
+    def stream(
+        self, prompt: str, config: GenerationConfig | None = None, **kwargs: Any
+    ) -> Iterator[str]:
+        return self.backend.stream(prompt, config, **kwargs)
+
+    def compare(
+        self, prompt: str, config: GenerationConfig | None = None, **kwargs: Any
+    ) -> dict[str, str]:
+        """Generate the same prompt with patches off and on.
+
+        Disables patches rather than uninstalling them, so strengths and
+        schedules survive the comparison.
+        """
+        states = {name: p.enabled for name, p in self.backend.patches.items()}
+        try:
+            for name in states:
+                self.backend.set_enabled(name, False)
+            baseline = self.generate(prompt, config, **kwargs)
+            for name in states:
+                self.backend.set_enabled(name, True)
+            patched = self.generate(prompt, config, **kwargs)
+        finally:
+            for name, was_enabled in states.items():
+                self.backend.set_enabled(name, was_enabled)
+        return {"baseline": baseline, "patched": patched}
+
+    # -- introspection ---------------------------------------------------------
+
+    def capabilities(self) -> Capabilities:
+        return type(self.backend).capabilities()
+
+    def describe(self) -> dict[str, Any]:
+        descriptor = self.backend.describe_model()
+        return {
+            "backend": self.backend.name,
+            "model": descriptor.model_id,
+            "architecture": descriptor.architecture,
+            "hidden_size": descriptor.hidden_size,
+            "num_layers": descriptor.num_layers,
+            "revision": descriptor.revision,
+            "compatibility_mode": self.compatibility_mode,
+            "patches": [PatchHandle(self, n).describe() for n in self.backend.patches],
+        }
+
+    def unload(self) -> None:
+        self.backend.unload()
+
+    def __enter__(self) -> "BrainPatchedModel":
+        return self
+
+    def __exit__(self, *exc: Any) -> None:
+        self.unload()
+
+    def __repr__(self) -> str:
+        patches = ", ".join(self.backend.patches) or "none"
+        return f"<BrainPatchedModel backend={self.backend.name!r} patches=[{patches}]>"

brainpatch/runtime/scheduling.py

@@ -0,0 +1,13 @@
+"""Token-indexed strength schedules.
+
+Re-exported from :mod:`brainpatch.steering.schedule`, which already implements
+and tests this. It lives under ``runtime`` too because scheduling is a *runtime*
+concern, not a research one, and a user reading the runtime package should not
+have to know that the code happens to predate the split.
+"""
+
+from __future__ import annotations
+
+from brainpatch.steering.schedule import StrengthSchedule
+
+__all__ = ["StrengthSchedule"]

brainpatch/schemas/__init__.py

@@ -0,0 +1,35 @@
+"""Serializable schemas shared between the local control plane and Modal.
+
+Every schema here is a plain :mod:`dataclasses` structure with explicit
+``to_dict`` / ``from_dict`` methods. No pydantic, no torch: these types travel
+between a laptop with no ML stack and a GPU container, so they must be
+constructible from nothing but the standard library.
+"""
+
+from brainpatch.schemas.contrast import ContrastExample, ContrastSet
+from brainpatch.schemas.feature import FeatureContext, FeatureRecord, FeatureStats
+from brainpatch.schemas.manifest import ActivationManifest, ShardRecord
+from brainpatch.schemas.patch import (
+    BrainPatchSpec,
+    FeatureEdit,
+    PatchCompatibilityError,
+    PatchValidationError,
+    SAEReference,
+)
+from brainpatch.schemas.sae import SAEConfig
+
+__all__ = [
+    "ActivationManifest",
+    "BrainPatchSpec",
+    "ContrastExample",
+    "ContrastSet",
+    "FeatureContext",
+    "FeatureEdit",
+    "FeatureRecord",
+    "FeatureStats",
+    "PatchCompatibilityError",
+    "PatchValidationError",
+    "SAEConfig",
+    "SAEReference",
+    "ShardRecord",
+]

brainpatch/schemas/contrast.py

@@ -0,0 +1,161 @@
+"""Behavioural contrast datasets.
+
+A contrast example pairs one prompt with a *positive* response (the behaviour we
+want more of) and a *negative* response (the behaviour we want less of). The
+difference in internal activations between the two is the starting point for
+candidate-feature search.
+
+These are **development fixtures**, not benchmarks. The sets shipped in
+``examples/contrast/`` are small, hand-written, synthetic, and were never
+validated against human judgement or an external standard. They exist to
+exercise the pipeline; any number computed on them is a smoke-test number.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass, field
+from typing import Any, Iterable, Iterator
+
+
+@dataclass
+class ContrastExample:
+    """One prompt with a contrasting pair of responses."""
+
+    prompt: str
+    positive_response: str
+    negative_response: str
+    category: str = "uncategorized"
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def validate(self) -> None:
+        if not self.prompt.strip():
+            raise ValueError("contrast example has an empty prompt")
+        if not self.positive_response.strip():
+            raise ValueError(f"empty positive_response for prompt {self.prompt[:60]!r}")
+        if not self.negative_response.strip():
+            raise ValueError(f"empty negative_response for prompt {self.prompt[:60]!r}")
+        if self.positive_response.strip() == self.negative_response.strip():
+            raise ValueError(
+                f"positive and negative responses are identical for prompt {self.prompt[:60]!r}"
+            )
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "ContrastExample":
+        missing = {"prompt", "positive_response", "negative_response"} - set(data)
+        if missing:
+            raise ValueError(f"contrast example is missing keys: {sorted(missing)}")
+        return cls(
+            prompt=str(data["prompt"]),
+            positive_response=str(data["positive_response"]),
+            negative_response=str(data["negative_response"]),
+            category=str(data.get("category", "uncategorized")),
+            metadata=dict(data.get("metadata", {})),
+        )
+
+
+@dataclass
+class ContrastSet:
+    """A named collection of contrast examples.
+
+    Attributes
+    ----------
+    synthetic:
+        True for hand-written development fixtures. Always keep this honest --
+        it is what stops a fixture from being reported as a benchmark.
+    """
+
+    name: str
+    description: str
+    examples: list[ContrastExample] = field(default_factory=list)
+    synthetic: bool = True
+    version: str = "0.1"
+
+    def __len__(self) -> int:
+        return len(self.examples)
+
+    def __iter__(self) -> Iterator[ContrastExample]:
+        return iter(self.examples)
+
+    def validate(self) -> None:
+        if not self.name:
+            raise ValueError("contrast set must have a name")
+        if not self.examples:
+            raise ValueError(f"contrast set {self.name!r} is empty")
+        for example in self.examples:
+            example.validate()
+
+    def categories(self) -> list[str]:
+        """Distinct categories present, in sorted order."""
+        return sorted({e.category for e in self.examples})
+
+    def filter(self, category: str) -> "ContrastSet":
+        """A new set containing only examples from ``category``."""
+        return ContrastSet(
+            name=f"{self.name}:{category}",
+            description=self.description,
+            examples=[e for e in self.examples if e.category == category],
+            synthetic=self.synthetic,
+            version=self.version,
+        )
+
+    def split(self, holdout_fraction: float, seed: int = 0) -> tuple["ContrastSet", "ContrastSet"]:
+        """Deterministically split into (train, holdout).
+
+        Uses a seeded :class:`random.Random` so that patch search and held-out
+        evaluation never accidentally see the same examples across runs.
+        """
+        import random
+
+        if not 0.0 < holdout_fraction < 1.0:
+            raise ValueError(f"holdout_fraction must be in (0, 1), got {holdout_fraction}")
+        indices = list(range(len(self.examples)))
+        random.Random(seed).shuffle(indices)
+        n_holdout = max(1, int(round(len(indices) * holdout_fraction)))
+        holdout_idx = set(indices[:n_holdout])
+
+        train = [e for i, e in enumerate(self.examples) if i not in holdout_idx]
+        holdout = [e for i, e in enumerate(self.examples) if i in holdout_idx]
+        return (
+            ContrastSet(f"{self.name}:train", self.description, train, self.synthetic, self.version),
+            ContrastSet(
+                f"{self.name}:holdout", self.description, holdout, self.synthetic, self.version
+            ),
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "name": self.name,
+            "description": self.description,
+            "synthetic": self.synthetic,
+            "version": self.version,
+            "examples": [e.to_dict() for e in self.examples],
+        }
+
+    def to_json(self, *, indent: int = 2) -> str:
+        return json.dumps(self.to_dict(), indent=indent)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "ContrastSet":
+        if "name" not in data:
+            raise ValueError("contrast set is missing 'name'")
+        return cls(
+            name=str(data["name"]),
+            description=str(data.get("description", "")),
+            examples=[ContrastExample.from_dict(e) for e in data.get("examples", [])],
+            synthetic=bool(data.get("synthetic", True)),
+            version=str(data.get("version", "0.1")),
+        )
+
+    @classmethod
+    def from_json(cls, text: str) -> "ContrastSet":
+        return cls.from_dict(json.loads(text))
+
+    @classmethod
+    def from_examples(
+        cls, name: str, description: str, examples: Iterable[ContrastExample]
+    ) -> "ContrastSet":
+        return cls(name=name, description=description, examples=list(examples))

brainpatch/schemas/feature.py

@@ -0,0 +1,193 @@
+"""Feature-database records.
+
+A :class:`FeatureRecord` describes one SAE feature: how often it fires, how
+strongly, and which token contexts drive it hardest.
+
+Scientific note
+---------------
+``FeatureRecord.hypothesis`` is exactly that -- a *hypothesis*. Top activating
+examples are correlational evidence and nothing more. The field
+:attr:`FeatureRecord.evidence_level` records how much support a semantic
+description actually has, and it never advances past ``"correlational"``
+automatically. Only an intervention experiment with scale-matched controls can
+move a feature to ``"controlled_interventional"``, and only an independent
+repetition can move it to ``"replicated"``. Both transitions are performed by
+the validation pipeline writing a measured result, never by a labelling
+heuristic.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass, field
+from typing import Any, Literal
+
+#: Ordered strength of evidence behind a semantic claim about a feature.
+#:
+#: Deliberately conservative. An earlier version of this ladder topped out at
+#: ``causal``, which claimed more than a single controlled experiment can
+#: deliver: passing scale-matched controls once, on one model, at one layer,
+#: with one prompt set, is evidence *consistent with* a causal effect -- not a
+#: demonstration of causation. The top two rungs now name what was actually
+#: done (controls were run; the result was reproduced) rather than what one
+#: might wish to infer from it.
+EvidenceLevel = Literal[
+    # no description offered
+    "none",
+    # top-activating contexts look suggestive; nothing more
+    "correlational",
+    # feature activation predicts a behaviour on held-out data
+    "predictive",
+    # steering it changes behaviour; controls absent, incomplete, or not yet run
+    "interventional",
+    # steering changes behaviour AND scale-matched controls did not, in one
+    # experiment with adequate statistical power
+    "controlled_interventional",
+    # the controlled result held up on independent repetition -- a different
+    # prompt set, seed, or run
+    "replicated",
+]
+
+EVIDENCE_ORDER: tuple[str, ...] = (
+    "none",
+    "correlational",
+    "predictive",
+    "interventional",
+    "controlled_interventional",
+    "replicated",
+)
+
+#: Levels at which scale-matched controls have actually been run and passed.
+CONTROLLED_LEVELS: frozenset[str] = frozenset({"controlled_interventional", "replicated"})
+
+
+@dataclass
+class FeatureContext:
+    """One high-activation occurrence of a feature, with surrounding text."""
+
+    example_index: int
+    token_position: int
+    token_id: int
+    token_text: str
+    activation: float
+    #: Decoded text a few tokens either side, for human inspection.
+    context_before: str = ""
+    context_after: str = ""
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "FeatureContext":
+        known = {f for f in cls.__dataclass_fields__}  # noqa: SLF001
+        return cls(**{k: v for k, v in data.items() if k in known})  # type: ignore[arg-type]
+
+
+@dataclass
+class FeatureStats:
+    """Activation statistics for one feature over an activation corpus."""
+
+    fire_count: int = 0
+    total_tokens: int = 0
+    mean_activation: float = 0.0
+    """Mean over *firing* tokens only (zeros excluded)."""
+    max_activation: float = 0.0
+    std_activation: float = 0.0
+    decoder_norm: float = 0.0
+
+    @property
+    def firing_rate(self) -> float:
+        """Fraction of tokens on which this feature is active."""
+        if self.total_tokens == 0:
+            return 0.0
+        return self.fire_count / self.total_tokens
+
+    @property
+    def is_dead(self) -> bool:
+        """A feature that never fired over the analysed corpus."""
+        return self.fire_count == 0
+
+    def to_dict(self) -> dict[str, Any]:
+        data = asdict(self)
+        data["firing_rate"] = self.firing_rate
+        data["is_dead"] = self.is_dead
+        return data
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "FeatureStats":
+        known = {f for f in cls.__dataclass_fields__}  # noqa: SLF001
+        return cls(**{k: v for k, v in data.items() if k in known})  # type: ignore[arg-type]
+
+
+@dataclass
+class FeatureRecord:
+    """A single entry in the feature database."""
+
+    feature_id: int
+    stats: FeatureStats = field(default_factory=FeatureStats)
+    top_contexts: list[FeatureContext] = field(default_factory=list)
+
+    #: Tentative, human- or machine-suggested description. NOT a validated label.
+    hypothesis: str | None = None
+    evidence_level: EvidenceLevel = "none"
+    #: Free-form pointers to experiments that produced the evidence.
+    evidence_refs: list[str] = field(default_factory=list)
+
+    def __post_init__(self) -> None:
+        if self.evidence_level not in EVIDENCE_ORDER:
+            raise ValueError(
+                f"unknown evidence_level {self.evidence_level!r}; "
+                f"expected one of {EVIDENCE_ORDER}"
+            )
+        if self.hypothesis is not None and self.evidence_level == "none":
+            # A description always carries at least correlational weight; being
+            # explicit here prevents an unlabelled-looking record from silently
+            # shipping a semantic claim.
+            self.evidence_level = "correlational"
+
+    @property
+    def has_controlled_evidence(self) -> bool:
+        """True once scale-matched controls have been run and passed."""
+        return self.evidence_level in CONTROLLED_LEVELS
+
+    @property
+    def is_validated(self) -> bool:
+        """True only for a controlled result that survived independent repetition.
+
+        Intentionally strict: one passing experiment is
+        ``controlled_interventional``, not validation.
+        """
+        return self.evidence_level == "replicated"
+
+    def label_for_display(self) -> str:
+        """Human-facing label that never overstates the evidence."""
+        if self.hypothesis is None:
+            return f"feature {self.feature_id} (no description)"
+        if self.is_validated:
+            return f"feature {self.feature_id}: {self.hypothesis} [replicated]"
+        return f"feature {self.feature_id}: {self.hypothesis} [{self.evidence_level}, unvalidated]"
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "feature_id": self.feature_id,
+            "stats": self.stats.to_dict(),
+            "top_contexts": [c.to_dict() for c in self.top_contexts],
+            "hypothesis": self.hypothesis,
+            "evidence_level": self.evidence_level,
+            "evidence_refs": list(self.evidence_refs),
+        }
+
+    def to_json(self) -> str:
+        """Single-line JSON, suitable for a ``features.jsonl`` row."""
+        return json.dumps(self.to_dict(), sort_keys=True)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "FeatureRecord":
+        return cls(
+            feature_id=int(data["feature_id"]),
+            stats=FeatureStats.from_dict(data.get("stats", {})),
+            top_contexts=[FeatureContext.from_dict(c) for c in data.get("top_contexts", [])],
+            hypothesis=data.get("hypothesis"),
+            evidence_level=data.get("evidence_level", "none"),
+            evidence_refs=list(data.get("evidence_refs", [])),
+        )