milestone_interp.py

"""SAELens + frozen attention-pattern patching (milestone interpreters).

Frozen attention: graft **clean-run** ``hook_pattern`` onto a **corrupt** forward at chosen
layer(s). When sequence lengths differ, apply a **suffix-aligned** block: overwrite the
bottom-right ``m×m`` submatrix of patterns (``m = min(clean_len, corrupt_len)``) so Q/K
indices match the shared final positions.

SAELens: load a pretrained / disk sparse autoencoder and apply **top-k latent coordinate
transfer** (same recipe as frontier ``TinySparseAE``) at ``metadata.hook_name``.

Enable with truthy ``PG_FROZEN_ATTN``, ``PG_SAELENS``, or umbrella ``PG_MILESTONE``.
SAE load env: ``PG_SAE_RELEASE`` + ``PG_SAE_ID`` (HF hub via SAELens registry) or ``PG_SAE_DISK``
(path). Optional tuning: ``PG_SAELENS_TOPK``, ``PG_FROZEN_ATTN_LAYER`` (comma-separated ints;
default = path-best head layer from caller).
"""

from __future__ import annotations

import os
from typing import Callable, List, Sequence, Tuple

import torch
from transformer_lens import ActivationCache, HookedTransformer

from frontier_lab import _nt_nm_denoise, _token_logit


def _truthy(key: str) -> bool:
    return os.environ.get(key, "").strip().lower() in ("1", "true", "yes", "on", "all")


def milestone_any_env() -> bool:
    return _truthy("PG_MILESTONE") or _truthy("PG_FROZEN_ATTN") or _truthy("PG_SAELENS")


def _parse_layers(s: str) -> List[int]:
    out: List[int] = []
    for part in s.split(","):
        p = part.strip()
        if p:
            out.append(int(p))
    return out


def frozen_pattern_hooks_from_cache(
    donor_cache: ActivationCache,
    layers: Sequence[int],
) -> List[Tuple[str, Callable]]:
    """Build ``fwd_hooks`` that replace attn patterns from *donor* cache when shapes match."""

    hooks: List[Tuple[str, Callable]] = []
    for lyr in layers:
        hk = f"blocks.{lyr}.attn.hook_pattern"
        ref = donor_cache[hk]

        def factory(reference: torch.Tensor, layer_ix: int, hook_key: str) -> Callable:
            reference_det = reference.detach()

            def _fn(patt: torch.Tensor, *, hook) -> torch.Tensor:
                if hook.name != hook_key:
                    return patt

                ref = reference_det.to(device=patt.device, dtype=patt.dtype)
                rt = patt
                if ref.shape[:-2] != rt.shape[:-2]:
                    print(
                        f"[PG_FROZEN_ATTN] skip L{layer_ix}: batch/heads mismatch "
                        f"{tuple(ref.shape)} vs {tuple(rt.shape)}"
                    )
                    return patt

                seq_donor = ref.shape[-1]
                seq_run = rt.shape[-1]
                m = min(seq_donor, seq_run)
                if m <= 0:
                    return patt

                out = rt.clone()
                sd0 = seq_donor - m
                sr0 = seq_run - m
                out[:, :, sr0:seq_run, sr0:seq_run] = ref[:, :, sd0:seq_donor, sd0:seq_donor]

                if seq_donor != seq_run:
                    print(
                        f"[PG_FROZEN_ATTN] suffix-aligned pattern graft L{layer_ix}: "
                        f"m={m} (donor_seq={seq_donor}, runtime_seq={seq_run})"
                    )
                return out

            return _fn

        hooks.append((hk, factory(ref, lyr, hk)))
    return hooks


def load_sae_optional(device: torch.device):
    release = os.environ.get("PG_SAE_RELEASE", "").strip()
    sae_id = os.environ.get("PG_SAE_ID", "").strip()
    disk = os.environ.get("PG_SAE_DISK", "").strip()

    if not disk and not (release and sae_id):
        return None

    try:
        from sae_lens.saes.sae import SAE  # deferred import for optional dependency
    except ImportError:
        print(
            "[PG_SAELENS] `sae-lens` import failed — install deps (see requirements.txt) "
            "or disable PG_SAELENS / PG_MILESTONE."
        )
        return None

    if disk:
        print(f"[PG_SAELENS] loading disk SAE from {disk!r}")
        return SAE.load_from_disk(disk, device=str(device))

    print(f"[PG_SAELENS] loading HF SAE release={release!r} id={sae_id!r}")
    dtype = os.environ.get("PG_SAE_DTYPE", "float32").strip() or "float32"
    return SAE.from_pretrained(release, sae_id, device=str(device), dtype=dtype)


def _make_saelens_transfer_hook(
    sae,
    clean_ca: ActivationCache,
    corr_ca: ActivationCache,
    *,
    hk: str,
    k_latents: int,
):

    target_device = torch.device(next(sae.parameters()).device)

    def _fn(act: torch.Tensor, *, hook) -> torch.Tensor:
        if hook.name != hk:
            return act

        corr_ref = corr_ca[hk][:, -1, :].detach()
        cl = clean_ca[hk][:, -1, :].detach()
        rn = act[:, -1, :].detach()

        rn2 = rn.to(device=target_device, dtype=torch.float32)
        cl2 = cl.to(device=target_device, dtype=torch.float32)
        cr2 = corr_ref.to(device=target_device, dtype=torch.float32)

        with torch.no_grad():
            z_rn = sae.encode(rn2.unsqueeze(0)).squeeze(0)
            z_cl = sae.encode(cl2.unsqueeze(0)).squeeze(0)
            z_cr = sae.encode(cr2.unsqueeze(0)).squeeze(0)
            focus = torch.abs(z_cl - z_cr)
            kk = max(1, min(k_latents, int(focus.numel())))
            _, ix = torch.topk(focus.reshape(-1), kk)
            z_patch = z_rn.clone().reshape(-1)
            flat_cl = z_cl.reshape(-1)
            z_patch[ix] = flat_cl[ix]
            z_patch = z_patch.reshape_as(z_rn)
            recon = sae.decode(z_patch.unsqueeze(0)).squeeze(0)

        act[:, -1, :] = recon.to(device=act.device, dtype=act.dtype)
        return act

    return _fn


def run_frozen_attn_block(
    model: HookedTransformer,
    *,
    corrupt_prompt: str,
    donor_cache_clean: ActivationCache,
    target_token: str,
    distractor_token: str,
    corrupt_tgt: float,
    clean_tgt: float,
    corrupt_margin: float,
    clean_margin: float,
    default_layer: int,
) -> None:
    layers_env = os.environ.get("PG_FROZEN_ATTN_LAYER", "").strip()
    layers = _parse_layers(layers_env) if layers_env else [default_layer]

    print("\n=== FROZEN ATTENTION PATTERN (donor=clean attn.hook_pattern on corrupt fwd) ===")
    fwd_hooks = frozen_pattern_hooks_from_cache(donor_cache_clean, layers)

    patched = model.run_with_hooks(corrupt_prompt, fwd_hooks=fwd_hooks)
    nt, nm = _nt_nm_denoise(
        model,
        patched,
        target_token=target_token,
        distractor_token=distractor_token,
        corrupt_tgt=corrupt_tgt,
        clean_tgt=clean_tgt,
        corrupt_margin=corrupt_margin,
        clean_margin=clean_margin,
    )
    tgt_log = _token_logit(model, patched, target_token)
    mr = tgt_log - _token_logit(model, patched, distractor_token)

    lyr_str = ",".join(str(x) for x in layers)
    print(
        f"Layers [{lyr_str}] (env PG_FROZEN_ATTN_LAYER overrides; default path head layer when unset) "
        f"→ tgt={tgt_log:.4f} margin={mr:.4f} nt={nt:.3f} nm={nm:.3f}"
    )


def run_saelens_block(
    model: HookedTransformer,
    *,
    corrupt_prompt: str,
    donor_cache_clean: ActivationCache,
    corrupt_cache_corrupt: ActivationCache,
    target_token: str,
    distractor_token: str,
    corrupt_tgt: float,
    clean_tgt: float,
    corrupt_margin: float,
    clean_margin: float,
) -> None:
    print("\n=== SAELens PRETRAINED — TOP-K LATENT TRANSFER AT SAE METADATA HOOK ===")
    device = torch.device(next(model.parameters()).device)
    sae = load_sae_optional(device)
    if sae is None:
        print(
            "[PG_SAELENS] skipped — set PG_SAE_RELEASE + PG_SAE_ID or PG_SAE_DISK "
            "(and PG_SAELENS=1)."
        )
        return

    meta = sae.cfg.metadata
    hk_obj = getattr(meta, "hook_name", None)
    hk = hk_obj.strip() if isinstance(hk_obj, str) else None
    if not hk:
        print("[PG_SAELENS] SAE metadata has no hook_name; cannot graft.")
        return

    if hk not in model.hook_dict:
        print(
            f"[PG_SAELENS] hook {hk!r} missing on this TransformerLens model "
            f"({model.cfg.model_name}); skipping."
        )
        return

    _, probe = model.run_with_cache(corrupt_prompt, names_filter=[hk])

    probe_tensor = probe[hk]
    last_dim = int(probe_tensor.shape[-1])
    if getattr(sae.cfg, "d_in", None) is not None and int(sae.cfg.d_in) != last_dim:
        print(
            f"[PG_SAELENS] incompatible d_in={sae.cfg.d_in} vs activation dim {last_dim} "
            f"at {hk!r} for hooked model '{model.cfg.model_name}' — skipping."
        )
        return

    k_lat = int(os.environ.get("PG_SAELENS_TOPK", os.environ.get("PG_SAE_TOPK", "32")))
    fwd_fn = _make_saelens_transfer_hook(
        sae,
        donor_cache_clean,
        corrupt_cache_corrupt,
        hk=hk,
        k_latents=k_lat,
    )

    logits = model.run_with_hooks(corrupt_prompt, fwd_hooks=[(hk, fwd_fn)])
    nt, nm = _nt_nm_denoise(
        model,
        logits,
        target_token=target_token,
        distractor_token=distractor_token,
        corrupt_tgt=corrupt_tgt,
        clean_tgt=clean_tgt,
        corrupt_margin=corrupt_margin,
        clean_margin=clean_margin,
    )
    tgt_log = _token_logit(model, logits, target_token)
    mr = tgt_log - _token_logit(model, logits, distractor_token)

    mn = getattr(meta, "model_name", None)
    mn_s = repr(mn) if mn else "?"
    print(
        f"hook={hk!r} pretrained model tag {mn_s} topkXfer={k_lat} "
        f"→ tgt={tgt_log:.4f} margin={mr:.4f} nt={nt:.3f} nm={nm:.3f}"
    )


def run_milestone_bundle(
    model: HookedTransformer,
    *,
    corrupt_prompt: str,
    target_token: str,
    distractor_token: str,
    clean_act: ActivationCache,
    corr_act: ActivationCache,
    clean_tgt: float,
    corrupt_tgt: float,
    clean_margin: float,
    corrupt_margin: float,
    default_attn_layer_for_frozen_pattern: int,
) -> None:

    if _truthy("PG_FROZEN_ATTN") or _truthy("PG_MILESTONE"):
        run_frozen_attn_block(
            model,
            corrupt_prompt=corrupt_prompt,
            donor_cache_clean=clean_act,
            target_token=target_token,
            distractor_token=distractor_token,
            corrupt_tgt=corrupt_tgt,
            clean_tgt=clean_tgt,
            corrupt_margin=corrupt_margin,
            clean_margin=clean_margin,
            default_layer=default_attn_layer_for_frozen_pattern,
        )

    if _truthy("PG_SAELENS") or _truthy("PG_MILESTONE"):
        run_saelens_block(
            model,
            corrupt_prompt=corrupt_prompt,
            donor_cache_clean=clean_act,
            corrupt_cache_corrupt=corr_act,
            target_token=target_token,
            distractor_token=distractor_token,
            corrupt_tgt=corrupt_tgt,
            clean_tgt=clean_tgt,
            corrupt_margin=corrupt_margin,
            clean_margin=clean_margin,
        )