brainpatch 1.2.0__py3-none-any.whl

brainpatch/runtime/__init__.py

@@ -0,0 +1,50 @@
+"""The BrainPatch runtime: backends, capabilities, scheduling, and the model API.
+
+This subpackage is **infrastructure-independent**. It knows how to apply a
+vector to a layer of a frozen model; it knows nothing about SAEs, activation
+corpora, Modal, or how the patch it is applying came to exist.
+
+That separation is the product: research happens wherever the author has GPUs,
+and the resulting artifact runs on the user's own machine with whatever
+inference engine they already use.
+
+``brainpatch.runtime.model`` and ``brainpatch.runtime.base`` are importable with
+no ML stack present; the engine-specific code under
+:mod:`brainpatch.backends` imports torch/vllm/llama.cpp only when instantiated.
+"""
+
+from brainpatch.runtime.auto import (
+    BackendNotAvailable,
+    BackendStatus,
+    available_backends,
+    backend_class,
+    environment_report,
+    select_backend,
+)
+from brainpatch.runtime.base import (
+    ActivePatch,
+    BrainPatchBackend,
+    GenerationConfig,
+    ResolvedEdit,
+)
+from brainpatch.runtime.capabilities import CAPABILITY_FLAGS, Capabilities
+from brainpatch.runtime.model import BrainPatchedModel, PatchHandle
+from brainpatch.runtime.scheduling import StrengthSchedule
+
+__all__ = [
+    "ActivePatch",
+    "BackendNotAvailable",
+    "BackendStatus",
+    "BrainPatchBackend",
+    "BrainPatchedModel",
+    "CAPABILITY_FLAGS",
+    "Capabilities",
+    "GenerationConfig",
+    "PatchHandle",
+    "ResolvedEdit",
+    "StrengthSchedule",
+    "available_backends",
+    "backend_class",
+    "environment_report",
+    "select_backend",
+]

brainpatch/runtime/auto.py

@@ -0,0 +1,157 @@
+"""Backend discovery, selection, and the ``doctor`` report.
+
+Every backend is imported lazily and probed defensively: ``brainpatch doctor``
+has to work on a machine with none of the engines installed, and its whole job
+is to say *which* are missing and what to do about it.
+"""
+
+from __future__ import annotations
+
+import importlib
+import platform
+import sys
+from dataclasses import dataclass
+from typing import Any
+
+from brainpatch.runtime.base import BrainPatchBackend
+from brainpatch.runtime.capabilities import Capabilities
+
+#: backend name -> module path. Order is preference order for ``backend="auto"``.
+BACKEND_MODULES: dict[str, str] = {
+    "transformers": "brainpatch.backends.transformers_backend",
+    "llamacpp": "brainpatch.backends.llamacpp",
+    "vllm": "brainpatch.backends.vllm_backend",
+    "mlx": "brainpatch.backends.mlx_backend",
+}
+
+#: Friendly aliases users actually type.
+BACKEND_ALIASES: dict[str, str] = {
+    "hf": "transformers",
+    "torch": "transformers",
+    "pytorch": "transformers",
+    "llama.cpp": "llamacpp",
+    "llama_cpp": "llamacpp",
+    "gguf": "llamacpp",
+    "mlx-lm": "mlx",
+    "mlx_lm": "mlx",
+}
+
+
+class BackendNotAvailable(RuntimeError):
+    """The requested backend cannot run in this environment."""
+
+
+def normalize_backend_name(name: str) -> str:
+    key = name.strip().lower()
+    return BACKEND_ALIASES.get(key, key)
+
+
+def backend_class(name: str) -> type[BrainPatchBackend]:
+    """Import and return a backend class by name."""
+    key = normalize_backend_name(name)
+    if key not in BACKEND_MODULES:
+        raise BackendNotAvailable(
+            f"unknown backend {name!r}. Known: {', '.join(sorted(BACKEND_MODULES))}"
+        )
+    module = importlib.import_module(BACKEND_MODULES[key])
+    cls = getattr(module, "BACKEND", None)
+    if cls is None:  # pragma: no cover - guards a malformed backend module
+        raise BackendNotAvailable(f"backend module for {key!r} exposes no BACKEND class")
+    return cls
+
+
+@dataclass
+class BackendStatus:
+    """One row of the doctor report."""
+
+    name: str
+    available: bool
+    detail: str
+    capabilities: Capabilities | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "backend": self.name,
+            "available": self.available,
+            "detail": self.detail,
+            "capabilities": self.capabilities.to_dict() if self.capabilities else None,
+        }
+
+
+def probe_backend(name: str) -> BackendStatus:
+    """Check one backend without raising, whatever is or is not installed."""
+    key = normalize_backend_name(name)
+    try:
+        cls = backend_class(key)
+    except BackendNotAvailable as exc:
+        return BackendStatus(name=key, available=False, detail=str(exc))
+    except Exception as exc:  # noqa: BLE001 - a broken backend must not break doctor
+        return BackendStatus(name=key, available=False, detail=f"import failed: {exc}")
+
+    try:
+        available, detail = cls.is_available()
+    except Exception as exc:  # noqa: BLE001
+        return BackendStatus(name=key, available=False, detail=f"probe failed: {exc}")
+
+    try:
+        caps = cls.capabilities()
+    except Exception:  # noqa: BLE001
+        caps = None
+    return BackendStatus(name=key, available=available, detail=detail, capabilities=caps)
+
+
+def available_backends() -> list[BackendStatus]:
+    """Probe every known backend, in preference order."""
+    return [probe_backend(name) for name in BACKEND_MODULES]
+
+
+def select_backend(preferred: str = "auto") -> type[BrainPatchBackend]:
+    """Resolve ``preferred`` to a usable backend class.
+
+    ``"auto"`` picks the first available backend in preference order. An
+    explicitly named backend that is unavailable raises with the reason, rather
+    than silently falling back to a different engine -- a silent substitution
+    would make "I tested it on vLLM" untrue.
+    """
+    if preferred != "auto":
+        key = normalize_backend_name(preferred)
+        cls = backend_class(key)
+        ok, detail = cls.is_available()
+        if not ok:
+            raise BackendNotAvailable(f"backend {key!r} is not available: {detail}")
+        return cls
+
+    problems: list[str] = []
+    for name in BACKEND_MODULES:
+        status = probe_backend(name)
+        if status.available:
+            return backend_class(name)
+        problems.append(f"  {name}: {status.detail}")
+    raise BackendNotAvailable(
+        "no inference backend is available. Install one:\n"
+        "  pip install 'brainpatch[transformers]'\n\n"
+        "Probed:\n" + "\n".join(problems)
+    )
+
+
+def environment_report() -> dict[str, Any]:
+    """Everything ``brainpatch doctor`` prints."""
+    from brainpatch import __version__
+    from brainpatch.patch.registry import default_registry
+
+    registry = default_registry()
+    try:
+        installed = [p.name for p in registry.list_patches()]
+    except OSError:
+        installed = []
+
+    return {
+        "brainpatch_version": __version__,
+        "python": sys.version.split()[0],
+        "platform": platform.platform(),
+        "machine": platform.machine(),
+        "is_apple_silicon": platform.system() == "Darwin" and platform.machine() == "arm64",
+        "registry_home": str(registry.home),
+        "installed_patches": installed,
+        "backends": [s.to_dict() for s in available_backends()],
+    }

brainpatch/runtime/base.py

@@ -0,0 +1,311 @@
+"""The backend contract every inference engine adapter implements.
+
+Design rule: this module imports nothing heavy. It defines the interface and the
+bookkeeping that is identical across engines -- which patches are installed, what
+their live strength is, how a schedule resolves at token *n* -- so each backend
+only has to implement the part that is genuinely engine-specific: loading a
+model, injecting a vector, and generating.
+
+That split is what keeps the tricky logic (strength resolution, the guarantee
+that strength 0 is exactly baseline) in one tested place rather than
+reimplemented four times with four different bugs.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Iterator
+
+from brainpatch.patch.format import Manifest
+from brainpatch.patch.loader import LoadedPatch
+from brainpatch.patch.validation import (
+    CompatibilityMode,
+    CompatibilityReport,
+    ModelDescriptor,
+    check_compatibility,
+)
+from brainpatch.runtime.capabilities import Capabilities
+from brainpatch.runtime.scheduling import StrengthSchedule
+
+#: Coefficients below this are treated as exactly zero, so a zeroed patch is
+#: bit-identical to baseline rather than merely close to it.
+STRENGTH_EPSILON = 1e-12
+
+
+@dataclass
+class GenerationConfig:
+    """Sampling settings, shared verbatim across compared conditions."""
+
+    max_new_tokens: int = 128
+    temperature: float = 0.0
+    top_p: float = 1.0
+    top_k: int = 0
+    repetition_penalty: float = 1.0
+    seed: int = 0
+    stop: list[str] = field(default_factory=list)
+
+    @property
+    def do_sample(self) -> bool:
+        return self.temperature > 0.0
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "max_new_tokens": self.max_new_tokens,
+            "temperature": self.temperature,
+            "top_p": self.top_p,
+            "top_k": self.top_k,
+            "repetition_penalty": self.repetition_penalty,
+            "seed": self.seed,
+            "stop": list(self.stop),
+        }
+
+
+@dataclass
+class ActivePatch:
+    """An installed patch plus its live, user-controllable state."""
+
+    patch: LoadedPatch
+    strength: float = 1.0
+    enabled: bool = True
+    schedule: StrengthSchedule | None = None
+
+    def __post_init__(self) -> None:
+        if self.schedule is None and self.patch.manifest.schedule is not None:
+            self.schedule = StrengthSchedule.from_dict(self.patch.manifest.schedule)
+
+    @property
+    def name(self) -> str:
+        return self.patch.manifest.name
+
+    @property
+    def manifest(self) -> Manifest:
+        return self.patch.manifest
+
+    def multiplier_at(self, token_index: int) -> float:
+        """Live multiplier at a generated-token index, clamped to the envelope."""
+        if not self.enabled:
+            return 0.0
+        value = self.manifest.clamp_strength(self.strength)
+        if self.schedule is not None:
+            value *= self.schedule.strength_at(token_index)
+        return value
+
+
+@dataclass
+class ResolvedEdit:
+    """One vector to add, with its final coefficient, at one layer."""
+
+    layer: int
+    hook: str
+    vector_key: str
+    coefficient: float
+    patch_name: str
+
+
+class BrainPatchBackend(ABC):
+    """Common interface across Transformers, llama.cpp, vLLM and MLX."""
+
+    #: Short identifier used by ``--backend`` and in capability tables.
+    name: str = "abstract"
+
+    def __init__(self) -> None:
+        self._patches: dict[str, ActivePatch] = {}
+        self._compatibility_mode: CompatibilityMode = "strict"
+
+    # -- engine-specific -------------------------------------------------------
+
+    @classmethod
+    @abstractmethod
+    def is_available(cls) -> tuple[bool, str]:
+        """``(available, reason)``. Must not raise, even with nothing installed."""
+
+    @classmethod
+    @abstractmethod
+    def capabilities(cls) -> Capabilities:
+        """What this backend can do. Callable without the engine installed."""
+
+    @abstractmethod
+    def load_model(self, model: str, **kwargs: Any) -> None:
+        """Load the frozen base model. Never modifies weights on disk."""
+
+    @abstractmethod
+    def describe_model(self) -> ModelDescriptor:
+        """Architecture facts discovered from the loaded model."""
+
+    @abstractmethod
+    def generate(self, prompt: str, config: GenerationConfig | None = None, **kwargs: Any) -> str:
+        """Generate a completion with all enabled patches active."""
+
+    def stream(
+        self, prompt: str, config: GenerationConfig | None = None, **kwargs: Any
+    ) -> Iterator[str]:
+        """Yield incremental text. Backends without streaming yield once."""
+        self.capabilities().require("streaming")
+        yield self.generate(prompt, config, **kwargs)
+
+    # -- shared patch bookkeeping ---------------------------------------------
+
+    @property
+    def patches(self) -> dict[str, ActivePatch]:
+        return self._patches
+
+    def validate_patch(
+        self, patch: LoadedPatch, *, mode: CompatibilityMode | None = None
+    ) -> CompatibilityReport:
+        return check_compatibility(
+            patch.manifest,
+            self.describe_model(),
+            mode=mode or self._compatibility_mode,
+        )
+
+    def install_patch(
+        self,
+        patch: LoadedPatch,
+        *,
+        strength: float | None = None,
+        mode: CompatibilityMode | None = None,
+    ) -> ActivePatch:
+        """Validate then register a patch. Nothing changes if validation fails."""
+        report = self.validate_patch(patch, mode=mode)
+        report.raise_if_failed()
+        for warning in report.warnings:
+            self._warn(warning)
+
+        backend_status = patch.manifest.backend_status(self.name)
+        if backend_status in {"unsupported", "implemented"}:
+            self._warn(
+                f"patch {patch.manifest.name!r} declares backend '{self.name}' as "
+                f"'{backend_status}' -- it has not been verified on this engine."
+            )
+
+        active = ActivePatch(
+            patch=patch,
+            strength=(
+                patch.manifest.default_strength if strength is None else float(strength)
+            ),
+        )
+        if len(self._patches) >= 1 and not self.capabilities().multiple_patches:
+            existing = next(iter(self._patches))
+            if existing != active.name:
+                raise NotImplementedError(
+                    f"the {self.name!r} backend supports one patch at a time; "
+                    f"{existing!r} is already installed"
+                )
+        self._patches[active.name] = active
+        self._on_patches_changed()
+        return active
+
+    def remove_patch(self, name: str) -> None:
+        if name not in self._patches:
+            raise KeyError(f"no patch named {name!r} is installed on this backend")
+        del self._patches[name]
+        self._on_patches_changed()
+
+    def set_strength(self, name: str, strength: float) -> float:
+        """Set live strength; returns the value after clamping."""
+        active = self._require(name)
+        clamped = active.manifest.clamp_strength(strength)
+        if clamped != float(strength):
+            self._warn(
+                f"strength {strength} clamped to {clamped} by patch "
+                f"{name!r} (max_abs_strength={active.manifest.max_abs_strength})"
+            )
+        active.strength = clamped
+        self._on_patches_changed()
+        return clamped
+
+    def set_enabled(self, name: str, enabled: bool) -> None:
+        self._require(name).enabled = bool(enabled)
+        self._on_patches_changed()
+
+    def set_schedule(self, name: str, schedule: StrengthSchedule | dict[int, float] | None) -> None:
+        self.capabilities().require("dynamic_schedule")
+        active = self._require(name)
+        if isinstance(schedule, dict):
+            schedule = StrengthSchedule(schedule)
+        active.schedule = schedule
+        self._on_patches_changed()
+
+    def list_patches(self) -> list[str]:
+        return list(self._patches)
+
+    def _require(self, name: str) -> ActivePatch:
+        if name not in self._patches:
+            installed = ", ".join(self._patches) or "none"
+            raise KeyError(f"no patch named {name!r} is installed (installed: {installed})")
+        return self._patches[name]
+
+    # -- edit resolution -------------------------------------------------------
+
+    def resolve_edits(
+        self,
+        token_index: int = 0,
+        layer: int | None = None,
+        *,
+        is_prompt_pass: bool | None = None,
+    ) -> list[ResolvedEdit]:
+        """Vectors to add at a generated-token index.
+
+        An empty list is the signal to leave activations completely untouched.
+        That is what makes strength 0 identical to baseline rather than a no-op
+        addition that still round-trips through floating point.
+
+        ``is_prompt_pass`` lets an intervention restrict itself to the prompt or
+        to generated tokens. Backends that cannot distinguish the two pass
+        ``None``, in which case every intervention applies and the caller is
+        responsible for knowing that a site-restricted patch is being applied
+        more broadly than it was measured.
+        """
+        edits: list[ResolvedEdit] = []
+        for active in self._patches.values():
+            multiplier = active.multiplier_at(token_index)
+            if abs(multiplier) < STRENGTH_EPSILON:
+                continue
+            for intervention in active.manifest.interventions:
+                if layer is not None and intervention.layer != layer:
+                    continue
+                site = getattr(intervention, "site", "all")
+                if is_prompt_pass is not None and site != "all":
+                    if site == "prompt" and not is_prompt_pass:
+                        continue
+                    if site == "continuation" and is_prompt_pass:
+                        continue
+                coefficient = multiplier * intervention.coefficient
+                if abs(coefficient) < STRENGTH_EPSILON:
+                    continue
+                edits.append(
+                    ResolvedEdit(
+                        layer=intervention.layer,
+                        hook=intervention.hook,
+                        vector_key=intervention.vector,
+                        coefficient=coefficient,
+                        patch_name=active.name,
+                    )
+                )
+        return edits
+
+    def active_layers(self) -> list[int]:
+        return sorted({e.layer for e in self.resolve_edits(0)})
+
+    def vector_values(self, patch_name: str, key: str) -> list[float]:
+        return list(self._require(patch_name).patch.vector_for(key).data)
+
+    # -- hooks for subclasses --------------------------------------------------
+
+    def _on_patches_changed(self) -> None:
+        """Called after any patch-state mutation. Override to rebuild caches."""
+
+    def _warn(self, message: str) -> None:
+        import warnings
+
+        warnings.warn(f"[brainpatch:{self.name}] {message}", stacklevel=3)
+
+    def unload(self) -> None:
+        """Release engine resources. Safe to call more than once."""
+
+    def __enter__(self) -> "BrainPatchBackend":
+        return self
+
+    def __exit__(self, *exc: Any) -> None:
+        self.unload()

brainpatch/runtime/capabilities.py

@@ -0,0 +1,96 @@
+"""Backend capability reporting.
+
+Backends genuinely differ in what they can do, and the honest thing is to say so
+rather than emulate a feature badly. llama.cpp applies a control vector for a
+whole run; vLLM batches concurrent requests with shared model state; MLX has no
+CI hardware here. Each of those is a real constraint, and a user deserves to
+know before they design around a capability that is not there.
+
+:class:`Capabilities` is what ``brainpatch backends`` and ``brainpatch doctor``
+print, and what the runtime consults before accepting an operation.
+"""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+#: Every capability flag, in display order.
+CAPABILITY_FLAGS: tuple[str, ...] = (
+    "static_intervention",
+    "dynamic_schedule",
+    "multiple_patches",
+    "streaming",
+    "cpu",
+    "cuda",
+    "mps",
+    "apple_silicon",
+    "server",
+    "concurrent_requests",
+    "quantization",
+    "per_request_strength",
+)
+
+
+@dataclass
+class Capabilities:
+    """What a backend can actually do.
+
+    Defaults are all False on purpose: a backend must opt in to each claim, so a
+    forgotten flag understates rather than overstates support.
+    """
+
+    name: str
+
+    static_intervention: bool = False
+    """Apply a fixed-strength patch for a whole generation."""
+
+    dynamic_schedule: bool = False
+    """Change strength between generated tokens."""
+
+    multiple_patches: bool = False
+    streaming: bool = False
+
+    cpu: bool = False
+    cuda: bool = False
+    mps: bool = False
+    apple_silicon: bool = False
+
+    server: bool = False
+    concurrent_requests: bool = False
+    per_request_strength: bool = False
+
+    quantization: tuple[str, ...] = ()
+    """Quantization formats actually exercised, e.g. ``("Q4_K_M",)``."""
+
+    notes: dict[str, str] = field(default_factory=dict)
+    """Per-capability explanation, especially for the False ones."""
+
+    def supports(self, flag: str) -> bool:
+        if flag not in CAPABILITY_FLAGS:
+            raise ValueError(f"unknown capability {flag!r}; expected one of {CAPABILITY_FLAGS}")
+        value = getattr(self, flag)
+        return bool(value)
+
+    def require(self, flag: str) -> None:
+        """Raise a useful error if the backend cannot do ``flag``."""
+        if not self.supports(flag):
+            note = self.notes.get(flag)
+            detail = f" {note}" if note else ""
+            raise NotImplementedError(
+                f"the {self.name!r} backend does not support {flag!r}.{detail}"
+            )
+
+    def to_dict(self) -> dict[str, Any]:
+        data = asdict(self)
+        data["quantization"] = list(self.quantization)
+        return data
+
+    def matrix_row(self) -> dict[str, Any]:
+        """Compact row for the capability matrix table."""
+        row: dict[str, Any] = {"backend": self.name}
+        for flag in CAPABILITY_FLAGS:
+            value = getattr(self, flag)
+            row[flag] = bool(value) if not isinstance(value, tuple) else bool(value)
+        row["quantization"] = ", ".join(self.quantization) or "-"
+        return row