brainpatch 1.2.0__py3-none-any.whl

brainpatch/schemas/sae.py

@@ -0,0 +1,146 @@
+"""Top-K sparse autoencoder configuration.
+
+This module holds only the *description* of an SAE. The torch implementation
+lives in :mod:`brainpatch.research.ml.sae` and is never imported locally.
+
+Design notes captured here because they matter for correctness of downstream
+interventions:
+
+Input normalization
+    Residual-stream activations of different models/layers have wildly
+    different scales. We rescale inputs so that ``E[||x||_2] == sqrt(d_in)``
+    and store the resulting ``input_scale`` in the checkpoint. Any intervention
+    that injects a decoder direction back into the *raw* residual stream must
+    multiply by ``input_scale`` to undo the normalization -- otherwise a
+    "strength of 1.0" means something different for every SAE.
+
+Decoder normalization
+    Decoder columns are constrained to unit L2 norm. Without this, the network
+    can trivially shrink the decoder and inflate feature activations (or vice
+    versa), which makes activation magnitudes -- and therefore any strength
+    parameter defined in terms of them -- meaningless. Unit-norm columns give
+    ``strength`` a stable interpretation: "add ``strength * input_scale`` units
+    of length along this direction".
+
+Dead features
+    Top-K SAEs reliably produce features that stop firing. We track a rolling
+    fire count and optionally apply the AuxK auxiliary loss (reconstruct the
+    residual error using only currently-dead features), which is the standard
+    mitigation from the OpenAI Top-K SAE work.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+SAE_FORMAT_VERSION = "0.1"
+
+
+@dataclass
+class SAEConfig:
+    """Architecture + training configuration of a Top-K SAE.
+
+    Attributes
+    ----------
+    d_in:
+        Residual-stream width of the base model (e.g. 1536 for Qwen2.5-1.5B).
+    d_sae:
+        Dictionary size (number of learned features).
+    k:
+        Upper bound on features kept active per token by the Top-K operator.
+        ``torch.topk`` always selects ``k`` indices, but the preceding ReLU can
+        make some of the selected values zero, so measured L0 is ``<= k`` rather
+        than identically ``k``.
+    normalize_decoder:
+        Constrain decoder columns to unit L2 norm after every optimizer step.
+    tied_init:
+        Initialise ``W_dec = W_enc.T``; a standard and stable starting point.
+    auxk_alpha:
+        Weight of the AuxK dead-feature revival loss. ``0.0`` disables it.
+    auxk_k:
+        How many dead features AuxK reconstructs the residual with.
+    dead_feature_window:
+        A feature is "dead" if it has not fired in this many training tokens.
+    """
+
+    # architecture
+    d_in: int
+    d_sae: int
+    k: int = 32
+    normalize_decoder: bool = True
+    tied_init: bool = True
+    auxk_alpha: float = 1.0 / 32.0
+    auxk_k: int = 256
+    dead_feature_window: int = 200_000
+
+    # provenance of the activations this SAE is defined over
+    model: str = ""
+    model_revision: str = ""
+    layer: int = -1
+    hook: str = ""
+
+    # optimisation
+    lr: float = 3e-4
+    beta1: float = 0.9
+    beta2: float = 0.999
+    batch_size: int = 512
+    epochs: int = 1
+    max_steps: int | None = None
+    grad_clip: float = 1.0
+    lr_warmup_steps: int = 50
+    seed: int = 0
+
+    # data handling
+    shuffle_buffer: int = 8192
+    val_fraction: float = 0.05
+
+    #: Filled in during training: multiply raw activations by this to normalize.
+    #: ``None`` until measured on a sample of the corpus.
+    input_scale: float | None = None
+
+    format_version: str = SAE_FORMAT_VERSION
+    notes: dict[str, Any] = field(default_factory=dict)
+
+    def validate(self) -> None:
+        """Raise :class:`ValueError` on an unusable configuration."""
+        if self.d_in <= 0:
+            raise ValueError(f"d_in must be positive, got {self.d_in}")
+        if self.d_sae <= 0:
+            raise ValueError(f"d_sae must be positive, got {self.d_sae}")
+        if not 0 < self.k <= self.d_sae:
+            raise ValueError(f"k must satisfy 0 < k <= d_sae, got k={self.k}, d_sae={self.d_sae}")
+        if self.auxk_alpha < 0:
+            raise ValueError(f"auxk_alpha must be non-negative, got {self.auxk_alpha}")
+        if self.auxk_k <= 0 or self.auxk_k > self.d_sae:
+            raise ValueError(f"auxk_k must satisfy 0 < auxk_k <= d_sae, got {self.auxk_k}")
+        if not 0.0 <= self.val_fraction < 1.0:
+            raise ValueError(f"val_fraction must be in [0, 1), got {self.val_fraction}")
+        if self.batch_size <= 0:
+            raise ValueError(f"batch_size must be positive, got {self.batch_size}")
+
+    @property
+    def expansion_factor(self) -> float:
+        """Dictionary size relative to the residual width."""
+        return self.d_sae / self.d_in
+
+    @property
+    def num_parameters(self) -> int:
+        """Parameter count: encoder + decoder weights plus both bias vectors."""
+        return 2 * self.d_in * self.d_sae + self.d_sae + self.d_in
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+    def to_json(self, *, indent: int = 2) -> str:
+        return json.dumps(self.to_dict(), indent=indent, sort_keys=True)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "SAEConfig":
+        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]
+
+    @classmethod
+    def from_json(cls, text: str) -> "SAEConfig":
+        return cls.from_dict(json.loads(text))

brainpatch/server/__init__.py

@@ -0,0 +1,11 @@
+"""OpenAI-compatible HTTP serving. Requires ``pip install 'brainpatch[server]'``."""
+
+__all__ = ["build_app"]
+
+
+def __getattr__(name: str):
+    if name == "build_app":
+        from brainpatch.server.app import build_app
+
+        return build_app
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

brainpatch/server/app.py

@@ -0,0 +1,269 @@
+"""OpenAI-compatible HTTP server.
+
+Existing OpenAI clients work unchanged: point ``base_url`` at this server and
+every request is served by the patched model. That is the whole point -- a
+BrainPatch should slot into the stack an application already has.
+
+Patch state is configured at **startup**, not per request. With a shared model
+object, a per-request strength change would be visible to every other in-flight
+request, so accepting one would be a correctness bug affecting other users'
+output. The ``brainpatch`` extra field is therefore accepted only when it
+matches the server's configuration, and rejected with a clear 400 otherwise,
+rather than silently ignored.
+
+Security posture: no patch loading over HTTP, no filesystem paths from request
+bodies, no code execution, strengths clamped to each patch's declared envelope.
+
+.. note::
+   This module deliberately has **no** ``from __future__ import annotations``
+   and defines its request models at module scope. FastAPI resolves parameter
+   annotations at runtime against the module's globals; with postponed
+   annotations and locally-defined models it cannot find them, silently demotes
+   the request body to a query parameter, and every POST fails with
+   ``422 Field required``. Keep both properties as they are.
+"""
+
+import time
+import uuid
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, Field
+
+from brainpatch.runtime.base import GenerationConfig
+
+#: Cap on tokens a single request may ask for, whatever it sends.
+MAX_REQUEST_TOKENS = 4096
+
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+
+class ChatRequest(BaseModel):
+    messages: List[ChatMessage]
+    model: Optional[str] = None
+    max_tokens: Optional[int] = Field(default=None, ge=1, le=MAX_REQUEST_TOKENS)
+    temperature: float = Field(default=0.0, ge=0.0, le=2.0)
+    top_p: float = Field(default=1.0, gt=0.0, le=1.0)
+    stream: bool = False
+    stop: Optional[List[str]] = None
+    #: Namespaced extension; see the module docstring on why it is validated
+    #: rather than honoured per request.
+    brainpatch: Optional[Dict[str, float]] = None
+
+
+class CompletionRequest(BaseModel):
+    prompt: str
+    model: Optional[str] = None
+    max_tokens: Optional[int] = Field(default=None, ge=1, le=MAX_REQUEST_TOKENS)
+    temperature: float = Field(default=0.0, ge=0.0, le=2.0)
+    top_p: float = Field(default=1.0, gt=0.0, le=1.0)
+    stream: bool = False
+    brainpatch: Optional[Dict[str, float]] = None
+
+
+def build_app(model: Any, served_model_name: Optional[str] = None) -> Any:
+    """Build the FastAPI application around an already-loaded model."""
+    try:
+        from fastapi import FastAPI, HTTPException
+        from fastapi.responses import StreamingResponse
+    except ModuleNotFoundError as exc:  # pragma: no cover
+        raise ModuleNotFoundError(
+            "the server needs FastAPI -- pip install 'brainpatch[server]'"
+        ) from exc
+
+    descriptor = model.backend.describe_model()
+    model_name = served_model_name or descriptor.model_id
+    capabilities = model.capabilities()
+
+    # Freeze patch state for the server's lifetime where the backend supports it.
+    begin = getattr(model.backend, "begin_serving", None)
+    if callable(begin):
+        begin()
+
+    api = FastAPI(title="BrainPatch", version="1.0")
+
+    def check_patch_override(requested: Optional[Dict[str, float]]) -> None:
+        """Accept a per-request patch spec only if it matches the server's."""
+        if not requested:
+            return
+        if capabilities.per_request_strength:
+            return
+        current = {name: model.backend.patches[name].strength for name in model.list_patches()}
+        mismatched = {
+            name: value
+            for name, value in requested.items()
+            if abs(current.get(name, 0.0) - float(value)) > 1e-9
+        }
+        if mismatched:
+            raise HTTPException(
+                status_code=400,
+                detail={
+                    "message": (
+                        f"the '{model.backend.name}' backend does not support "
+                        "per-request patch strength; requests share one model, so "
+                        "honouring this would change other users' output."
+                    ),
+                    "server_configuration": current,
+                    "requested": requested,
+                    "hint": "restart the server with the strengths you want",
+                },
+            )
+
+    def make_config(
+        max_tokens: Optional[int], temperature: float, top_p: float, stop: Optional[List[str]]
+    ) -> GenerationConfig:
+        return GenerationConfig(
+            max_new_tokens=min(max_tokens or 256, MAX_REQUEST_TOKENS),
+            temperature=temperature,
+            top_p=top_p,
+            stop=list(stop or []),
+        )
+
+    def render(messages: List[ChatMessage]) -> tuple:
+        system = next((m.content for m in messages if m.role == "system"), None)
+        user = next((m.content for m in reversed(messages) if m.role == "user"), None)
+        if user is None:
+            raise HTTPException(status_code=400, detail="no user message in request")
+        return user, system
+
+    @api.get("/health")
+    def health() -> Dict[str, Any]:
+        return {
+            "status": "ok",
+            "backend": model.backend.name,
+            "model": model_name,
+            "patches": {
+                name: {
+                    "strength": model.backend.patches[name].strength,
+                    "enabled": model.backend.patches[name].enabled,
+                    "evidence_level": model.backend.patches[name].manifest.evidence_level,
+                }
+                for name in model.list_patches()
+            },
+        }
+
+    @api.get("/v1/models")
+    def list_models() -> Dict[str, Any]:
+        return {
+            "object": "list",
+            "data": [
+                {
+                    "id": model_name,
+                    "object": "model",
+                    "created": int(time.time()),
+                    "owned_by": "brainpatch",
+                    "brainpatch": {"patches": model.list_patches()},
+                }
+            ],
+        }
+
+    @api.post("/v1/chat/completions")
+    def chat_completions(request: ChatRequest) -> Any:
+        check_patch_override(request.brainpatch)
+        prompt, system = render(request.messages)
+        cfg = make_config(request.max_tokens, request.temperature, request.top_p, request.stop)
+
+        if request.stream:
+            return StreamingResponse(
+                stream_chat(prompt, cfg, system, model_name), media_type="text/event-stream"
+            )
+        text = model.generate(prompt, cfg, system=system)
+        return chat_response(
+            text,
+            model_name,
+            prompt_tokens=count_tokens(model.backend, prompt),
+            completion_tokens=count_tokens(model.backend, text),
+        )
+
+    @api.post("/v1/completions")
+    def completions(request: CompletionRequest) -> Dict[str, Any]:
+        check_patch_override(request.brainpatch)
+        cfg = make_config(request.max_tokens, request.temperature, request.top_p, None)
+        text = model.generate(request.prompt, cfg, use_chat_template=False)
+        return {
+            "id": "cmpl-" + uuid.uuid4().hex[:24],
+            "object": "text_completion",
+            "created": int(time.time()),
+            "model": model_name,
+            "choices": [{"text": text, "index": 0, "finish_reason": "stop", "logprobs": None}],
+        }
+
+    def stream_chat(prompt: str, cfg: GenerationConfig, system: Optional[str], name: str):
+        import json as _json
+
+        request_id = "chatcmpl-" + uuid.uuid4().hex[:24]
+        created = int(time.time())
+
+        def chunk(delta: Dict[str, Any], finish: Optional[str]) -> str:
+            payload = {
+                "id": request_id,
+                "object": "chat.completion.chunk",
+                "created": created,
+                "model": name,
+                "choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
+            }
+            return "data: " + _json.dumps(payload) + "\n\n"
+
+        try:
+            for piece in model.stream(prompt, cfg, system=system):
+                yield chunk({"content": piece}, None)
+        except NotImplementedError:
+            # Backends without streaming still serve a valid SSE response.
+            yield chunk({"content": model.generate(prompt, cfg, system=system)}, None)
+        yield chunk({}, "stop")
+        yield "data: [DONE]\n\n"
+
+    return api
+
+
+def count_tokens(backend: Any, text: str) -> int:
+    """Token count for ``text``, or 0 if the backend cannot tokenise.
+
+    Reported zeros were a real problem, not a cosmetic one: any client that
+    meters usage, or any benchmark that normalises by generated tokens, silently
+    reads "nothing happened". Falling back to 0 is still possible for backends
+    with no tokenizer, but a backend that has one now reports the truth.
+    """
+    tokenizer = getattr(backend, "tokenizer", None)
+    if tokenizer is None:
+        # vLLM keeps its tokenizer behind the engine handle rather than exposing
+        # an attribute, so a plain getattr finds nothing and every count would
+        # come back 0 on the backend most likely to be metered.
+        engine = getattr(backend, "llm", None)
+        getter = getattr(engine, "get_tokenizer", None)
+        if callable(getter):
+            try:
+                tokenizer = getter()
+            except Exception:
+                tokenizer = None
+    if tokenizer is None:
+        return 0
+    try:
+        return len(tokenizer(text, add_special_tokens=False).input_ids)
+    except Exception:
+        return 0
+
+
+def chat_response(
+    text: str, model_name: str, *, prompt_tokens: int = 0, completion_tokens: int = 0
+) -> Dict[str, Any]:
+    return {
+        "id": "chatcmpl-" + uuid.uuid4().hex[:24],
+        "object": "chat.completion",
+        "created": int(time.time()),
+        "model": model_name,
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": text},
+                "finish_reason": "stop",
+            }
+        ],
+        "usage": {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": prompt_tokens + completion_tokens,
+        },
+    }

brainpatch/steering/__init__.py

@@ -0,0 +1,13 @@
+"""Steering control logic.
+
+Pure-Python scheduling and intervention *planning* live here. The torch hooks
+that actually mutate activations live in :mod:`brainpatch.research.ml.intervention`.
+
+Keeping the two apart means the trickiest part to get right -- when a strength
+changes and by how much -- is unit-testable on a laptop with no GPU.
+"""
+
+from brainpatch.steering.plan import InterventionPlan, PlannedEdit
+from brainpatch.steering.schedule import StrengthSchedule
+
+__all__ = ["InterventionPlan", "PlannedEdit", "StrengthSchedule"]

brainpatch/steering/plan.py

@@ -0,0 +1,177 @@
+"""Resolving installed patches into a concrete per-token intervention.
+
+The runtime keeps a set of installed :class:`~brainpatch.schemas.patch.BrainPatchSpec`
+objects, each with a user-controllable strength multiplier and enabled flag. On
+every forward pass the hook needs one question answered: *at generated-token
+index n, which feature directions do I add, and with what coefficients?*
+
+:class:`InterventionPlan` answers that with no torch involved, which is what
+makes the semantics testable locally. In particular the guarantee that
+
+    total strength == 0  =>  no tensor is modified at all
+
+is enforced here by returning an empty edit list, so "strength 0" is baseline by
+construction rather than by floating-point luck.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from brainpatch.schemas.patch import BrainPatchSpec
+from brainpatch.steering.schedule import StrengthSchedule
+
+#: Coefficients whose absolute value is below this are treated as exactly zero.
+#: Chosen well below any strength a user would set, but above float noise.
+STRENGTH_EPSILON = 1e-12
+
+
+@dataclass
+class PlannedEdit:
+    """One feature direction to apply, with its resolved coefficient."""
+
+    feature_id: int
+    coefficient: float
+    mode: str = "add"
+    source_patch: str = ""
+
+
+@dataclass
+class InstalledPatch:
+    """A patch registered in the runtime, plus its live control state."""
+
+    spec: BrainPatchSpec
+    strength: float = 1.0
+    enabled: bool = True
+    schedule: StrengthSchedule | None = None
+
+    def __post_init__(self) -> None:
+        if self.schedule is None and self.spec.schedule is not None:
+            self.schedule = StrengthSchedule.from_dict(self.spec.schedule)
+
+    @property
+    def name(self) -> str:
+        return self.spec.name
+
+    def multiplier_at(self, token_index: int) -> float:
+        """Global multiplier for this patch at a generated-token index."""
+        if not self.enabled:
+            return 0.0
+        if self.schedule is None:
+            return self.strength
+        return self.strength * self.schedule.strength_at(token_index)
+
+
+@dataclass
+class InterventionPlan:
+    """The set of installed patches and the layers they touch.
+
+    A plan is intentionally cheap to build and query: the hot path
+    (:meth:`edits_at`) runs once per forward pass.
+    """
+
+    patches: dict[str, InstalledPatch] = field(default_factory=dict)
+    #: Restrict interventions to ``[start, end)`` in generated-token index.
+    #: ``None`` means "every generated token".
+    token_range: tuple[int, int] | None = None
+
+    # -- registry --------------------------------------------------------------
+
+    def install(
+        self,
+        spec: BrainPatchSpec,
+        *,
+        strength: float = 1.0,
+        enabled: bool = True,
+    ) -> None:
+        """Register a patch. Re-installing the same name replaces it."""
+        self.patches[spec.name] = InstalledPatch(spec=spec, strength=strength, enabled=enabled)
+
+    def uninstall(self, name: str) -> None:
+        if name not in self.patches:
+            raise KeyError(f"no patch named {name!r} is installed")
+        del self.patches[name]
+
+    def set_strength(self, name: str, strength: float) -> None:
+        if name not in self.patches:
+            raise KeyError(f"no patch named {name!r} is installed")
+        self.patches[name].strength = float(strength)
+
+    def set_enabled(self, name: str, enabled: bool) -> None:
+        if name not in self.patches:
+            raise KeyError(f"no patch named {name!r} is installed")
+        self.patches[name].enabled = bool(enabled)
+
+    def set_schedule(self, name: str, schedule: StrengthSchedule | None) -> None:
+        if name not in self.patches:
+            raise KeyError(f"no patch named {name!r} is installed")
+        self.patches[name].schedule = schedule
+
+    # -- queries ---------------------------------------------------------------
+
+    def layers(self) -> list[int]:
+        """Every layer index that at least one installed patch hooks."""
+        return sorted({p.spec.sae.layer for p in self.patches.values()})
+
+    def is_active(self, token_index: int) -> bool:
+        """True if any edit would be applied at this generated-token index."""
+        return bool(self.edits_at(token_index))
+
+    def in_token_range(self, token_index: int) -> bool:
+        if self.token_range is None:
+            return True
+        start, end = self.token_range
+        return start <= token_index < end
+
+    def edits_at(self, token_index: int, *, layer: int | None = None) -> list[PlannedEdit]:
+        """Resolve the edits to apply at ``token_index``.
+
+        Parameters
+        ----------
+        token_index:
+            Index of the token being generated, 0-based, prompt excluded.
+        layer:
+            When given, only edits for patches hooking this layer are returned.
+
+        Returns
+        -------
+        list[PlannedEdit]
+            Possibly empty. An empty list is the signal to leave the residual
+            stream untouched, which is what makes ``strength=0`` bit-identical
+            to baseline rather than merely close to it.
+        """
+        if not self.in_token_range(token_index):
+            return []
+
+        edits: list[PlannedEdit] = []
+        for patch in self.patches.values():
+            if layer is not None and patch.spec.sae.layer != layer:
+                continue
+            multiplier = patch.multiplier_at(token_index)
+            if abs(multiplier) < STRENGTH_EPSILON:
+                continue
+            for edit in patch.spec.features:
+                coefficient = multiplier * edit.strength
+                if abs(coefficient) < STRENGTH_EPSILON:
+                    continue
+                edits.append(
+                    PlannedEdit(
+                        feature_id=edit.feature_id,
+                        coefficient=coefficient,
+                        mode=edit.mode,
+                        source_patch=patch.name,
+                    )
+                )
+        return edits
+
+    def describe(self) -> list[str]:
+        """One human-readable line per installed patch."""
+        lines = []
+        for patch in self.patches.values():
+            state = "on" if patch.enabled else "off"
+            sched = "scheduled" if patch.schedule is not None else "constant"
+            lines.append(
+                f"{patch.name}: strength={patch.strength:+.3f} [{state}, {sched}] "
+                f"-> {patch.spec.summary()}"
+            )
+        return lines