runtime.py

"""Pure-transformers steering runtime for the America AI deployment.

Consumes ``steering_bundle.pt`` (schema_version 1) produced by
``america_export.py``. No TransformerLens dependency: forward hooks are
registered on ``model.model.layers[L]`` and reproduce the harness semantics

    resid += multiplier * base_strength * typical_norm * unit_vector

at each concept's layer. This module is intentionally self-contained (no
imports from the ``america_ai`` package) so it can be copied verbatim into
the Hugging Face Space repo as a flat ``runtime.py``.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path

import torch

SCHEMA_VERSION = 1
HOOK_MODES = {"all_positions", "generation_only"}
CONCEPT_KEYS = {"unit_vector", "layer", "typical_norm", "base_strength"}


@dataclass
class Bundle:
    schema_version: int
    source_model: str
    target_model: str
    d_model: int
    concepts: dict[str, dict]
    presets: dict[str, dict[str, float]]
    hook_mode: str
    provenance: dict = field(default_factory=dict)


def save_bundle(bundle: Bundle, path: Path | str) -> None:
    validate_bundle(bundle)
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(
        {
            "schema_version": bundle.schema_version,
            "source_model": bundle.source_model,
            "target_model": bundle.target_model,
            "d_model": bundle.d_model,
            "concepts": bundle.concepts,
            "presets": bundle.presets,
            "hook_mode": bundle.hook_mode,
            "provenance": bundle.provenance,
        },
        path,
    )


def load_bundle(path: Path | str) -> Bundle:
    raw = torch.load(Path(path), map_location="cpu", weights_only=False)
    bundle = Bundle(
        schema_version=int(raw["schema_version"]),
        source_model=raw["source_model"],
        target_model=raw["target_model"],
        d_model=int(raw["d_model"]),
        concepts=raw["concepts"],
        presets=raw["presets"],
        hook_mode=raw["hook_mode"],
        provenance=raw.get("provenance", {}),
    )
    validate_bundle(bundle)
    return bundle


def validate_bundle(bundle: Bundle) -> None:
    if bundle.schema_version != SCHEMA_VERSION:
        raise ValueError(f"unsupported schema_version: {bundle.schema_version}")
    if bundle.hook_mode not in HOOK_MODES:
        raise ValueError(f"unknown hook mode: {bundle.hook_mode}")
    if not bundle.concepts:
        raise ValueError("bundle has no concepts")
    for name, concept in bundle.concepts.items():
        missing = CONCEPT_KEYS - concept.keys()
        if missing:
            raise ValueError(f"concept {name} missing keys: {sorted(missing)}")
        vec = concept["unit_vector"]
        if not isinstance(vec, torch.Tensor) or vec.shape != (bundle.d_model,):
            raise ValueError(f"concept {name} unit_vector must be shape ({bundle.d_model},)")
        if abs(float(vec.float().norm()) - 1.0) > 1e-3:
            raise ValueError(f"concept {name} unit_vector is not unit norm")
    for preset, multipliers in bundle.presets.items():
        unknown = set(multipliers) - set(bundle.concepts)
        if unknown:
            raise ValueError(f"preset {preset} references unknown concepts: {sorted(unknown)}")


class SteeringState:
    """Holds slider multipliers and the per-layer injection tensors they imply.

    Injection tensors live on CPU in float32; ``apply`` casts to the hidden
    state's device/dtype so the same state works on CPU, CUDA, and ZeroGPU.
    """

    def __init__(self, bundle: Bundle):
        self.bundle = bundle
        self.hook_mode = bundle.hook_mode
        self.multipliers = {name: 0.0 for name in bundle.concepts}
        self._by_layer: dict[int, list[str]] = {}
        for name, concept in bundle.concepts.items():
            self._by_layer.setdefault(int(concept["layer"]), []).append(name)
        self._injections: dict[int, torch.Tensor | None] = {}
        self._rebuild()

    def layers(self) -> list[int]:
        return sorted(self._by_layer)

    def set_strengths(self, multipliers: dict[str, float]) -> None:
        unknown = set(multipliers) - set(self.multipliers)
        if unknown:
            raise KeyError(f"unknown concepts: {sorted(unknown)}")
        self.multipliers.update({name: float(value) for name, value in multipliers.items()})
        self._rebuild()

    def set_preset(self, name: str) -> None:
        self.set_strengths(self.bundle.presets[name])

    def injection(self, layer: int) -> torch.Tensor | None:
        return self._injections.get(layer)

    def _rebuild(self) -> None:
        for layer, names in self._by_layer.items():
            total = torch.zeros(self.bundle.d_model, dtype=torch.float32)
            active = False
            for name in names:
                multiplier = self.multipliers[name]
                if multiplier == 0.0:
                    continue
                concept = self.bundle.concepts[name]
                scale = multiplier * float(concept["base_strength"]) * float(concept["typical_norm"])
                total = total + scale * concept["unit_vector"].float()
                active = True
            self._injections[layer] = total if active else None

    def apply(self, hidden: torch.Tensor, layer: int) -> torch.Tensor:
        """Return steered hidden states; returns ``hidden`` unchanged if inactive."""
        vec = self._injections.get(layer)
        if vec is None:
            return hidden
        vec = vec.to(device=hidden.device, dtype=hidden.dtype)
        out = hidden.clone()
        if self.hook_mode == "all_positions":
            out = out + vec
        else:  # generation_only: prompt pass -> final position; cached decode -> that token
            out[:, -1, :] = out[:, -1, :] + vec
        return out


def attach_hooks(model, state: SteeringState) -> list:
    """Register one forward hook per steered layer on ``model.model.layers``."""
    layer_modules = model.model.layers
    handles = []
    for layer in state.layers():
        handles.append(layer_modules[layer].register_forward_hook(_make_hook(state, layer)))
    return handles


def _make_hook(state: SteeringState, layer: int):
    def hook(module, args, output):
        hidden = output[0] if isinstance(output, tuple) else output
        steered = state.apply(hidden, layer)
        if steered is hidden:
            return output
        if isinstance(output, tuple):
            return (steered,) + tuple(output[1:])
        return steered

    return hook


class SteeredGemma:
    """Bundle + plain-transformers model with steering hooks attached."""

    def __init__(
        self,
        bundle: Bundle,
        model_id: str | None = None,
        device: str | None = None,
        dtype: torch.dtype = torch.bfloat16,
        hf_token: str | None = None,
    ):
        from transformers import AutoModelForCausalLM, AutoTokenizer

        self.bundle = bundle
        self.model_id = model_id or bundle.target_model
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id, token=hf_token)
        self.model = AutoModelForCausalLM.from_pretrained(
            self.model_id,
            torch_dtype=dtype,
            attn_implementation="eager",  # recommended for gemma-2
            token=hf_token,
        ).to(self.device)
        self.model.eval()
        self.state = SteeringState(bundle)
        self._handles = attach_hooks(self.model, self.state)

    def set_strengths(self, multipliers: dict[str, float]) -> None:
        self.state.set_strengths(multipliers)

    def set_preset(self, name: str) -> None:
        self.state.set_preset(name)

    def build_input_ids(self, messages: list[dict[str, str]]) -> torch.Tensor:
        return self.tokenizer.apply_chat_template(
            messages, add_generation_prompt=True, return_tensors="pt"
        ).to(self.device)

    @torch.no_grad()
    def generate(self, messages: list[dict[str, str]], **generate_kwargs) -> str:
        input_ids = self.build_input_ids(messages)
        output = self.model.generate(input_ids=input_ids, **generate_kwargs)
        return self.tokenizer.decode(output[0, input_ids.shape[-1] :], skip_special_tokens=True)