mri_live.py

"""
MRI Live — streaming rank-capture for mid-inference monitoring.

Two projection methods, side by side:
  - lm   : residual / per-head vectors -> final_norm -> lm_head  (the standard
           logit-lens readout; expensive, so by default I only project the
           'Layer' residual column).
  - interp : the same vectors linearly interpolated up to vocab length and
           softmaxed — a projection-free readout that needs NO lm_head. Cheap
           enough that by default I project EVERY head + Layer.

For each generated token it yields a TraceStep carrying, per active method, the
rank of each tracked token at each (layer, col) cell, plus the projection
latency for that method (so the cost difference is visible).

Interventions: pass an InterventionConfig. The trigger is screened against each
active method; if any fires I suppress/stop the (single, real) generation and
record which method(s) tripped.
"""

from __future__ import annotations

import os
import sys
import time
from dataclasses import dataclass, field
from typing import Callable, Iterator, Optional

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import numpy as np
import torch
import torch.nn.functional as F

from projection_map import load_model, build_col_specs


# ---------------------------------------------------------------------------
# Configs
# ---------------------------------------------------------------------------

@dataclass
class TrackedWord:
    """A word to track. Tokenizes to 1..N tokens; anchor picks which one."""
    text: str
    token_ids: list[int]
    anchor_idx: int = 0

    @property
    def anchor_id(self) -> int:
        return self.token_ids[self.anchor_idx]

    @property
    def label(self) -> str:
        if len(self.token_ids) == 1:
            return self.text
        return f"{self.text} [tok{self.anchor_idx}]"


@dataclass
class InterventionConfig:
    """When to intervene, and how.

    trigger_fn: given a MethodTrace, returns True if intervention fires.
    action: 'log' | 'suppress' | 'stop'
    """
    trigger_fn: Callable[["MethodTrace"], bool]
    action: str = "log"


@dataclass
class CellTrace:
    """Per-cell trace data for a single (layer, col)."""
    layer: int
    col: str
    ranks: list[int]            # rank per tracked word, 1-indexed
    probs: list[float]          # prob per tracked word
    top30_ids: list[int] = field(default_factory=list)
    top30_probs: list[float] = field(default_factory=list)


@dataclass
class MethodTrace:
    """One projection method's view of a single generated token."""
    method: str                 # 'lm' | 'interp'
    cells: list[CellTrace]
    layers_used: list[int]
    cols_used: list[str]
    proj_ms: float = 0.0

    def rank_matrix(self, tracked_idx: int) -> np.ndarray:
        n_layers, n_cols = len(self.layers_used), len(self.cols_used)
        m = np.full((n_layers, n_cols), 1, dtype=np.int32)
        lr = {l: i for i, l in enumerate(self.layers_used)}
        cc = {c: j for j, c in enumerate(self.cols_used)}
        for cell in self.cells:
            if cell.layer in lr and cell.col in cc:
                m[lr[cell.layer], cc[cell.col]] = cell.ranks[tracked_idx]
        return m

    def prob_matrix(self, tracked_idx: int) -> np.ndarray:
        n_layers, n_cols = len(self.layers_used), len(self.cols_used)
        m = np.zeros((n_layers, n_cols), dtype=np.float32)
        lr = {l: i for i, l in enumerate(self.layers_used)}
        cc = {c: j for j, c in enumerate(self.cols_used)}
        for cell in self.cells:
            if cell.layer in lr and cell.col in cc:
                m[lr[cell.layer], cc[cell.col]] = cell.probs[tracked_idx]
        return m


@dataclass
class TraceStep:
    """One generated token's full trace across the active method(s)."""
    step: int
    token_id: int
    token_str: str
    token_prob: float
    fwd_ms: float
    methods: dict                       # 'lm'/'interp' -> MethodTrace
    intervened: bool = False
    intervention_action: str = ""
    intervention_methods: list = field(default_factory=list)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def tokenize_word(tokenizer, word: str, with_leading_space: bool = True) -> list[int]:
    text = (" " + word) if with_leading_space else word
    return tokenizer.encode(text, add_special_tokens=False)


def word_preview(tokenizer, word: str) -> dict:
    with_space = tokenize_word(tokenizer, word, with_leading_space=True)
    no_space = tokenize_word(tokenizer, word, with_leading_space=False)
    return {
        "with_leading_space": [(tid, tokenizer.decode([tid])) for tid in with_space],
        "no_leading_space": [(tid, tokenizer.decode([tid])) for tid in no_space],
    }


def _is_cuda(device) -> bool:
    if isinstance(device, str):
        return device.startswith("cuda")
    return getattr(device, "type", "") == "cuda"


# ---------------------------------------------------------------------------
# Streaming trace
# ---------------------------------------------------------------------------

def stream_trace(
    model,
    tokenizer,
    device,
    num_layers: int,
    num_heads: int,
    prompt: str,
    tracked_words: list[TrackedWord],
    max_new_tokens: int = 30,
    layer_mask: Optional[list[int]] = None,
    lm_cols: Optional[list[str]] = None,
    interp_cols: Optional[list[str]] = None,
    methods: tuple = ("lm",),
    scale_mode: str = "full",
    intervention: Optional[InterventionConfig] = None,
    apply_chat_template: bool = True,
    precise_timing: bool = False,
) -> Iterator[TraceStep]:
    """Generate tokens one-by-one, yielding a TraceStep after each.

    methods: any of ('lm',), ('interp',), ('lm','interp').
    lm_cols / interp_cols: column labels each method projects. Defaults:
        lm     -> ['Layer']     (residual only; lm_head is expensive)
        interp -> all cols      (every head + Layer; interp is cheap)
    """
    col_specs, head_indices, _, _ = build_col_specs(num_heads)
    default_cols = [c for c in col_specs if c[0] in ("single", "layer")]
    all_labels = [c[2] for c in default_cols]

    if lm_cols is None:
        lm_cols = ["Layer"]
    if interp_cols is None:
        interp_cols = all_labels  # everything

    def _resolve(wanted):
        return [c for c in default_cols if c[2] in set(wanted)]

    method_cols = {}
    if "lm" in methods:
        method_cols["lm"] = _resolve(lm_cols)
    if "interp" in methods:
        method_cols["interp"] = _resolve(interp_cols)

    if layer_mask is None:
        layer_mask = list(range(num_layers))

    # Do I need per-head decomposition at all? (only if some active col is a head)
    need_heads = any(
        ct != "layer"
        for cols in method_cols.values()
        for (ct, _, _) in cols
    )

    final_norm = model.model.norm
    lm_head = model.lm_head
    vocab_size = lm_head.weight.shape[0]
    tracked_ids = [tw.anchor_id for tw in tracked_words]
    tids_t = torch.tensor(tracked_ids, device=device)

    # Hook o_proj inputs (per-head pre-projection) — only used if need_heads.
    captured: dict[int, torch.Tensor] = {}

    def make_hook(layer_idx):
        def hook_fn(module, args, output):
            captured[layer_idx] = args[0].detach()
        return hook_fn

    handles = []
    if need_heads:
        for l in range(num_layers):
            handles.append(
                model.model.layers[l].self_attn.o_proj.register_forward_hook(make_hook(l))
            )

    # ---- project a list of (layer, col_label, vec[hidden]) through a method ----
    def project_method(entries, method):
        if not entries:
            return [], -1.0
        if precise_timing and _is_cuda(device):
            torch.cuda.synchronize()
        t0 = time.time()
        stack = torch.stack([v for (_, _, v) in entries], dim=0)  # [N, hidden], model dtype
        with torch.no_grad():
            if method == "lm":
                dist = lm_head(final_norm(stack)).float()                 # [N, vocab] logits
            else:  # interp — linear upsample to vocab; align_corners=True == np.interp(linspace)
                dist = F.interpolate(
                    stack.float().unsqueeze(1), size=vocab_size, mode="linear", align_corners=True
                ).squeeze(1)                                              # [N, vocab]
            sel = dist.index_select(1, tids_t)                           # [N, T]
            ranks = torch.empty((dist.shape[0], len(tracked_ids)), dtype=torch.long, device=dist.device)
            for ti in range(len(tracked_ids)):
                ranks[:, ti] = (dist > sel[:, ti:ti + 1]).sum(dim=1) + 1
            # prob of the tracked tokens without materializing the full [N,vocab] softmax
            lse = torch.logsumexp(dist, dim=1, keepdim=True)             # [N, 1]
            probs = torch.exp(sel - lse)                                 # [N, T]
        if precise_timing and _is_cuda(device):
            torch.cuda.synchronize()
        proj_ms = (time.time() - t0) * 1000.0 if precise_timing else -1.0
        ranks = ranks.cpu().numpy()
        probs = probs.float().cpu().numpy()
        cells = []
        for i, (layer, col_label, _) in enumerate(entries):
            cells.append(CellTrace(
                layer=layer, col=col_label,
                ranks=[int(x) for x in ranks[i]],
                probs=[float(x) for x in probs[i]],
            ))
        return cells, proj_ms

    # ---- build input ----
    if prompt and apply_chat_template:
        try:
            full_prompt = tokenizer.apply_chat_template(
                [{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True
            )
        except Exception:
            full_prompt = prompt
    else:
        full_prompt = prompt or ""
    if full_prompt:
        step_input = tokenizer(full_prompt, return_tensors="pt").input_ids.to(device)
    else:
        step_input = torch.tensor([[tokenizer.bos_token_id]], device=device)

    # KV cache: first pass processes the whole prompt; every later pass feeds
    # ONLY the new token with past_key_values, so generation is O(n) not O(n^2).
    # The projection code only ever reads the last position, so it is unchanged.
    past = None

    try:
        for step_idx in range(max_new_tokens):
            captured.clear()
            if precise_timing and _is_cuda(device):
                torch.cuda.synchronize()
            t_fwd = time.time()
            with torch.no_grad():
                outputs = model(step_input, past_key_values=past, use_cache=True,
                                output_hidden_states=True)
            past = outputs.past_key_values
            if precise_timing and _is_cuda(device):
                torch.cuda.synchronize()
            fwd_ms = (time.time() - t_fwd) * 1000.0 if precise_timing else -1.0

            logits = outputs.logits[0, -1, :]
            probs_final = F.softmax(logits.float(), dim=-1)

            # Gather per-layer residual + (if needed) per-head vectors at last pos.
            residual_by_layer = {}
            head_by_layer = {}
            for l_idx in layer_mask:
                residual_by_layer[l_idx] = outputs.hidden_states[l_idx + 1][:, -1, :][0]
                if need_heads and l_idx in captured:
                    inp = captured[l_idx]
                    ow = model.model.layers[l_idx].self_attn.o_proj.weight
                    hidden_size, attn_out_dim = ow.shape
                    per_head_dim = attn_out_dim // num_heads
                    # all heads at the last position in one shot, then slice:
                    #   proj[h, o] = sum_d o_proj_weight[o, h, d] * head_out[h, d]
                    mh_last = inp[0, -1].view(num_heads, per_head_dim)        # [H, d]
                    wv = ow.view(hidden_size, num_heads, per_head_dim)        # [O, H, d]
                    proj = torch.einsum("ohd,hd->ho", wv, mh_last)           # [H, O]
                    head_by_layer[l_idx] = {h: proj[h] for h in head_indices}

            # Build entries + project, per method.
            step_methods = {}
            for m, cols in method_cols.items():
                entries = []
                for l_idx in layer_mask:
                    hv = head_by_layer.get(l_idx, {})
                    for (col_type, col_heads, col_label) in cols:
                        if col_type == "layer":
                            entries.append((l_idx, col_label, residual_by_layer[l_idx]))
                        else:
                            if not hv:
                                continue
                            combined = sum(hv[h] for h in col_heads)
                            if scale_mode == "full":
                                combined = combined * (num_heads / len(col_heads))
                            elif scale_mode == "mean":
                                combined = combined / len(col_heads)
                            entries.append((l_idx, col_label, combined))
                cells, proj_ms = project_method(entries, m)
                cols_used = [lbl for (_, _, lbl) in
                             ([(0, 0, c[2]) for c in cols])]
                step_methods[m] = MethodTrace(
                    method=m, cells=cells,
                    layers_used=list(layer_mask),
                    cols_used=[c[2] for c in cols],
                    proj_ms=proj_ms,
                )

            # Pick next token (greedy, from the real logits).
            top_prob, top_id = torch.max(probs_final, dim=-1)
            chosen_id = int(top_id.item())
            chosen_str = tokenizer.decode(chosen_id)
            chosen_prob = float(top_prob.item())

            step = TraceStep(
                step=step_idx, token_id=chosen_id, token_str=chosen_str,
                token_prob=chosen_prob, fwd_ms=fwd_ms, methods=step_methods,
            )

            # Intervention — screen each active method.
            if intervention is not None:
                fired = [m for m, mt in step_methods.items() if intervention.trigger_fn(mt)]
                if fired:
                    step.intervened = True
                    step.intervention_action = intervention.action
                    step.intervention_methods = fired
                    if intervention.action == "suppress":
                        masked = logits.clone()
                        masked[chosen_id] = float("-inf")
                        np2 = F.softmax(masked.float(), dim=-1)
                        p2, i2 = torch.max(np2, dim=-1)
                        chosen_id = int(i2.item())
                        chosen_str = tokenizer.decode(chosen_id)
                        chosen_prob = float(p2.item())
                        step.token_id, step.token_str, step.token_prob = chosen_id, chosen_str, chosen_prob

            yield step

            if step.intervened and intervention.action == "stop":
                break

            # Next pass: feed ONLY the new token (the cache holds the rest).
            step_input = torch.tensor([[chosen_id]], device=device)
            if chosen_id == tokenizer.eos_token_id:
                break
    finally:
        for h in handles:
            h.remove()


# ---------------------------------------------------------------------------
# Threshold-based trigger (screened per method)
# ---------------------------------------------------------------------------

def make_rank_threshold_trigger(
    tracked_idx: int,
    layers: list[int],
    cols: list[str],
    rank_below: Optional[int] = None,
    rank_above: Optional[int] = None,
) -> Callable[["MethodTrace"], bool]:
    """Fire if the tracked token's rank at any matching cell crosses a
    threshold. `rank_below=10` -> fires when rank <= 10 (token surfaced)."""
    layers_s = set(layers)
    cols_s = set(cols)

    def trigger(mt: "MethodTrace") -> bool:
        for cell in mt.cells:
            if cell.layer not in layers_s or cell.col not in cols_s:
                continue
            r = cell.ranks[tracked_idx]
            if rank_below is not None and r <= rank_below:
                return True
            if rank_above is not None and r >= rank_above:
                return True
        return False
    return trigger