patch.py

"""
patch.py — Steered generation and control directions.
Spec: §5.2–5.4, §8.2 (core skeleton provided).
"""
from __future__ import annotations
import torch
from transformer_lens import HookedTransformer


def make_steer_hook(direction: torch.Tensor, alpha: float):
    """Hook that adds alpha * direction to the residual stream."""
    def hook(resid, hook):          # resid: [batch, seq, d_model]
        return resid + alpha * direction.to(resid.dtype)
    return hook


def generate_steered(
    model: HookedTransformer,
    prompt: str,
    direction: torch.Tensor,
    alpha: float,
    layer: int,
    max_new: int = 80,
) -> str:
    """Generate with activation patch at blocks.{layer}.hook_resid_post."""
    name = f"blocks.{layer}.hook_resid_post"
    with model.hooks(fwd_hooks=[(name, make_steer_hook(direction, alpha))]):
        return model.generate(prompt, max_new_tokens=max_new, do_sample=False)


def norm_matched_random(direction: torch.Tensor, seed: int) -> torch.Tensor:
    """Random control direction r, norm-matched to direction. (Spec §5.3)"""
    g = torch.Generator().manual_seed(seed)
    r = torch.randn(direction.shape, generator=g)
    return r / r.norm()            # same unit norm; alpha supplies magnitude