coarse_localization.py

"""
Coarse activation-localization sweep for month modular arithmetic prompts.

This script estimates where a model represents the two causal variables in the
month-offset task: the starting month (`input`) and the number of months to move
(`offset`). For each row in `datasets/months.jsonl`, it treats
`prompt_fixed` as the receiver prompt and `prompt_cfact + answer_cfact` as the
donor computation. It caches donor activations from the counterfactual prompt
with its answer, then intervenes on the receiver prompt at one activation site at
a time.

For every layer and for each coarse position, it replaces one receiver activation
with the corresponding donor activation and greedily reads the next generated
token from the patched receiver prompt:

  - `input`: patch the receiver starting-month token position using the donor
    counterfactual starting-month token position.
  - `offset`: patch the receiver offset token position using the donor
    counterfactual offset token position.
  - `last_token`: cache the donor activations from `prompt_cfact + answer_cfact`,
    but use the donor final prompt token (`<|assistant|>`) activation, because in
    a decoder-only model that is the position whose logits generate the first
    answer token. This patches into the receiver final prompt token position,
    then evaluates the newly generated next token.

The metric is interchange intervention accuracy (IIA). The script checks whether
the patched next token equals one of three concept effects:

  - `input`: `answer_input_cfact`
  - `offset`: `answer_offset_cfact`
  - `output`: `answer_cfact`

It writes CSV rows to `results/coarse_localization.csv`, one row per
`(position, layer, concept_effect)`, with schema:

  concept_effect,position,layer,average_iia
"""

from __future__ import annotations

import argparse
import csv
import os
import string
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Literal

import matplotlib.pyplot as plt
import torch
from transformer_lens.model_bridge.bridge import TransformerBridge
from transformer_lens.utilities import get_act_name

from utils import (
    find_text_token_position,
    load_rows,
    n_layers_from_hook_dict,
    prompt_tokens,
    rendered_prompt_text,
    rotate_existing_output,
    tokenize_continuation,
)

MODEL = "microsoft/Phi-mini-MoE-instruct"

DTYPE_MAP = {
    "bfloat16": torch.bfloat16,
    "bf16": torch.bfloat16,
    "float16": torch.float16,
    "fp16": torch.float16,
    "float32": torch.float32,
    "fp32": torch.float32,
}

PatchPosition = Literal["input", "offset", "last_token"]
ConceptEffect = Literal["input", "offset", "output"]
PatchSiteKind = Literal["mlp_out", "resid_pre", "resid_mid", "resid_post"]

CONCEPT_ANSWER_FIELD: dict[ConceptEffect, str] = {
    "input": "answer_input_cfact",
    "offset": "answer_offset_cfact",
    "output": "answer_cfact",
}

POSITIONS: list[PatchPosition] = ["input", "offset", "last_token"]


@dataclass(frozen=True)
class PatchSpec:
    """Activation vector and receiver position for a single intervention."""

    site: str
    receiver_pos: int
    donor_vec: torch.Tensor


def boot_model(
    model_name_or_path: str,
    dtype_name: str,
    device_map: str | None,
    trust_remote_code: bool,
):
    """Load a TransformerBridge model using the same convention as probing_ff.py."""

    if dtype_name not in DTYPE_MAP:
        valid = ", ".join(sorted(DTYPE_MAP))
        raise ValueError(f"Unknown dtype {dtype_name!r}. Valid choices: {valid}")

    kwargs = {
        "dtype": DTYPE_MAP[dtype_name],
        "trust_remote_code": trust_remote_code,
    }
    if device_map is not None and device_map.lower() != "none":
        kwargs["device_map"] = device_map
    else:
        kwargs["device"] = "cuda" if torch.cuda.is_available() else "cpu"

    model = TransformerBridge.boot_transformers(model_name_or_path, **kwargs)
    model.eval()
    return model


def site_candidates(layer: int, kind: PatchSiteKind) -> list[str]:
    """Return plausible hook names for a layer and activation site kind."""

    if kind == "mlp_out":
        return [
            f"blocks.{layer}.mlp.hook_out",
            get_act_name("mlp_out", layer),
            f"blocks.{layer}.hook_mlp_out",
        ]
    if kind == "resid_pre":
        return [get_act_name("resid_pre", layer), f"blocks.{layer}.hook_resid_pre"]
    if kind == "resid_mid":
        return [get_act_name("resid_mid", layer), f"blocks.{layer}.hook_resid_mid"]
    if kind == "resid_post":
        return [get_act_name("resid_post", layer), f"blocks.{layer}.hook_resid_post"]

    raise ValueError(f"Unsupported site kind: {kind!r}")


def resolve_site(model, layer: int, kind: PatchSiteKind) -> str:
    """Resolve the concrete hook name for a layer and activation site kind."""

    for name in site_candidates(layer, kind):
        if name in model.hook_dict:
            return name

    nearby = [key for key in model.hook_dict.keys() if f"blocks.{layer}." in key]
    nearby_text = "\n".join(nearby[:80])
    raise KeyError(
        f"No hook site found for layer={layer}, kind={kind!r}. "
        f"Nearby hooks:\n{nearby_text}"
    )


def prompt_plus_answer_tokens(model, prompt: str, answer: str) -> torch.Tensor:
    """Tokenize a month prompt and append its answer continuation tokens."""

    tokens = prompt_tokens(model, prompt)
    continuation = tokenize_continuation(model, answer)
    continuation = continuation.to(device=tokens.device).unsqueeze(0)
    return torch.cat([tokens, continuation], dim=1)


def run_with_cache_one_site(model, tokens: torch.Tensor, site: str):
    """Run the model and cache one requested activation site."""

    with torch.no_grad():
        try:
            return model.run_with_cache(tokens, names_filter=lambda name: name == site)
        except TypeError:
            try:
                return model.run_with_cache(tokens, names_filter=[site])
            except TypeError:
                return model.run_with_cache(tokens)


def make_patch_hook(patch_spec: PatchSpec) -> Callable:
    """Build a hook that replaces one receiver activation vector."""

    def patch_hook(value: torch.Tensor, hook):
        if value.ndim != 3:
            raise ValueError(
                f"Expected rank-3 activation at {patch_spec.site}, "
                f"got shape {tuple(value.shape)}."
            )
        if patch_spec.receiver_pos >= value.shape[1]:
            raise ValueError(
                f"Patch position {patch_spec.receiver_pos} exceeds sequence length "
                f"{value.shape[1]} at {patch_spec.site}."
            )

        patched = value.clone()
        donor_vec = patch_spec.donor_vec.to(device=value.device, dtype=value.dtype)
        patched[:, patch_spec.receiver_pos, :] = donor_vec
        return patched

    return patch_hook


def patched_logits(
    model, receiver_tokens: torch.Tensor, patch_spec: PatchSpec
) -> torch.Tensor:
    """Run a patched receiver forward pass and return logits."""

    with torch.no_grad():
        return model.run_with_hooks(
            receiver_tokens,
            fwd_hooks=[(patch_spec.site, make_patch_hook(patch_spec))],
        )


def patched_next_token_from_logits(model, logits: torch.Tensor) -> str:
    """Decode the greedy next token from patched final-position logits."""

    token_id = int(logits[0, -1].argmax().item())
    return decode_answer_token(model, token_id)


def decode_answer_token(model, token_id: int) -> str:
    """Decode and normalize one generated answer token for exact matching."""

    try:
        text = model.to_string(token_id)
    except Exception:
        text = model.tokenizer.decode([token_id])

    text = text.strip()
    if not text:
        return text
    return text.split()[0].strip(string.punctuation + string.whitespace)


def token_text(model, token_id: int) -> str:
    """Decode one token without normalization for debugging token boundaries."""

    try:
        return model.to_string(token_id)
    except Exception:
        return model.tokenizer.decode([token_id], skip_special_tokens=False)


def token_table(model, tokens: torch.Tensor) -> list[tuple[int, int, str]]:
    """Return `(position, token_id, token_text)` rows for a token tensor."""

    token_ids = tokens[0].detach().cpu().tolist()
    return [
        (pos, token_id, token_text(model, token_id))
        for pos, token_id in enumerate(token_ids)
    ]


def print_token_table(model, label: str, tokens: torch.Tensor) -> None:
    """Print a readable token table for debugging prompt formatting."""

    print(f"\n--- {label}: seq_len={tokens.shape[1]} ---")
    for pos, token_id, text in token_table(model, tokens):
        print(f"{pos:>4} | {token_id:>8} | {text!r}")


def print_continuation_variants(model, answer: str) -> None:
    """Print tokenization of answer continuations with and without a space."""

    print(f"\n--- answer continuation variants for {answer!r} ---")
    for prefix in ["", " "]:
        text = prefix + answer
        ids = tokenize_continuation(model, text).detach().cpu().tolist()
        pieces = [token_text(model, token_id) for token_id in ids]
        print(f"prefix={prefix!r} text={text!r} ids={ids} pieces={pieces!r}")


def topk_next_tokens_from_logits(
    model, logits: torch.Tensor, k: int = 10
) -> list[dict]:
    """Return top-k decoded next tokens from final-position logits."""

    probs = torch.softmax(logits[0, -1].detach().float(), dim=-1)
    values, indices = torch.topk(probs, k=k)
    rows = []
    for prob, token_id in zip(values.tolist(), indices.tolist()):
        rows.append(
            {
                "token_id": int(token_id),
                "token_text": token_text(model, int(token_id)),
                "normalized": decode_answer_token(model, int(token_id)),
                "prob": float(prob),
            }
        )
    return rows


def print_debug_row(
    model,
    row: dict,
    receiver_tokens: torch.Tensor,
    donor_prompt_only_tokens: torch.Tensor,
    donor_tokens: torch.Tensor,
    donor_pos: int,
    receiver_pos: int,
    logits: torch.Tensor,
    position: PatchPosition,
    layer: int,
    site: str,
) -> None:
    """Print tokenization and patch-position details for one intervention."""

    print("\n================ TOKEN DEBUG ================")
    print(f"layer={layer} site={site} position={position}")
    print(f"receiver prompt_fixed: {row['prompt_fixed']!r}")
    print(f"donor prompt_cfact:    {row['prompt_cfact']!r}")
    print(f"donor answer_cfact:    {row['answer_cfact']!r}")
    print(
        f"rendered receiver:     {rendered_prompt_text(model.tokenizer, row['prompt_fixed'])!r}"
    )
    print(
        f"rendered donor:        {rendered_prompt_text(model.tokenizer, row['prompt_cfact'])!r}"
    )
    print_continuation_variants(model, row["answer_cfact"])
    print_token_table(model, "receiver prompt tokens", receiver_tokens)
    print_token_table(model, "donor prompt-only tokens", donor_prompt_only_tokens)
    print_token_table(model, "donor prompt + answer tokens", donor_tokens)
    print(
        "selected positions:",
        f"donor_pos={donor_pos}",
        f"donor_token={token_text(model, int(donor_tokens[0, donor_pos]))!r}",
        f"receiver_pos={receiver_pos}",
        f"receiver_token={token_text(model, int(receiver_tokens[0, receiver_pos]))!r}",
    )
    print("patched top next tokens:")
    for item in topk_next_tokens_from_logits(model, logits):
        print(item)
    print("=============================================\n")


def donor_and_receiver_positions(
    model,
    row: dict,
    position: PatchPosition,
    receiver_tokens: torch.Tensor,
    donor_prompt_tokens: torch.Tensor,
) -> tuple[int, int]:
    """Return `(donor_pos, receiver_pos)` for a coarse patch position."""

    if position == "input":
        return (
            find_text_token_position(model, row["prompt_cfact"], row["input_cfact"]),
            find_text_token_position(model, row["prompt_fixed"], row["input"]),
        )
    if position == "offset":
        return (
            find_text_token_position(model, row["prompt_cfact"], row["offset_cfact"]),
            find_text_token_position(model, row["prompt_fixed"], row["offset"]),
        )
    if position == "last_token":
        return donor_prompt_tokens.shape[1] - 1, receiver_tokens.shape[1] - 1

    raise ValueError(f"Unsupported position: {position!r}")


def evaluate_layer_position(
    model,
    rows: list[dict],
    layer: int,
    site: str,
    position: PatchPosition,
    debug_tokens: bool,
) -> dict[ConceptEffect, float]:
    """Evaluate IIA for all concept effects at one layer and coarse position."""

    correct = {concept_effect: 0 for concept_effect in CONCEPT_ANSWER_FIELD}

    for row in rows:
        receiver_tokens = prompt_tokens(model, row["prompt_fixed"])
        donor_prompt_only_tokens = prompt_tokens(model, row["prompt_cfact"])
        donor_tokens = prompt_plus_answer_tokens(
            model, row["prompt_cfact"], row["answer_cfact"]
        )
        donor_pos, receiver_pos = donor_and_receiver_positions(
            model=model,
            row=row,
            position=position,
            receiver_tokens=receiver_tokens,
            donor_prompt_tokens=donor_prompt_only_tokens,
        )

        _, donor_cache = run_with_cache_one_site(model, donor_tokens, site)
        if site not in donor_cache:
            available = "\n".join(str(key) for key in donor_cache.keys())
            raise KeyError(
                f"Requested site {site!r} was not cached. Available keys:\n{available}"
            )

        patch_spec = PatchSpec(
            site=site,
            receiver_pos=receiver_pos,
            donor_vec=donor_cache[site][:, donor_pos, :].detach(),
        )
        logits = patched_logits(model, receiver_tokens, patch_spec)
        predicted = patched_next_token_from_logits(model, logits)

        if debug_tokens:
            print_debug_row(
                model=model,
                row=row,
                receiver_tokens=receiver_tokens,
                donor_prompt_only_tokens=donor_prompt_only_tokens,
                donor_tokens=donor_tokens,
                donor_pos=donor_pos,
                receiver_pos=receiver_pos,
                logits=logits,
                position=position,
                layer=layer,
                site=site,
            )

        for concept_effect, answer_field in CONCEPT_ANSWER_FIELD.items():
            if debug_tokens:
                print(
                    "predicted:",
                    predicted,
                    "| expected:",
                    row[answer_field],
                    "| original:",
                    row["answer"],
                )
            if predicted == row[answer_field]:
                correct[concept_effect] += 1

    return {
        concept_effect: count / len(rows) if rows else 0.0
        for concept_effect, count in correct.items()
    }


def write_results(path: Path, rows: list[dict]) -> None:
    """Rotate any existing output file and write aggregate results as CSV."""

    path.parent.mkdir(parents=True, exist_ok=True)
    rotate_existing_output(path)
    with path.open("w", encoding="utf-8") as handle:
        writer = csv.DictWriter(
            handle,
            fieldnames=["concept_effect", "position", "layer", "average_iia"],
            lineterminator="\n",
        )
        writer.writeheader()
        writer.writerows(sort_result_rows(rows))


def load_result_rows(path: Path) -> list[dict]:
    """Load aggregate localization rows from a CSV results file."""

    with path.open("r", encoding="utf-8", newline="") as handle:
        reader = csv.DictReader(handle)
        rows = []
        for row in reader:
            rows.append(
                {
                    "concept_effect": row["concept_effect"],
                    "position": row["position"],
                    "layer": int(row["layer"]),
                    "average_iia": float(row["average_iia"]),
                }
            )
    return rows


def plot_heat_tables(rows: list[dict], output_dir: Path) -> None:
    """Save all concept-effect average-IIA heat tables in one figure."""

    output_dir.mkdir(parents=True, exist_ok=True)
    concept_effects = sorted({row["concept_effect"] for row in rows})
    layers = sorted({row["layer"] for row in rows})
    positions = [
        position
        for position in POSITIONS
        if any(row["position"] == position for row in rows)
    ]

    fig_width = max(
        5.0 * len(concept_effects), 1.5 * len(positions) * len(concept_effects)
    )
    fig_height = max(6.0, 0.28 * len(layers))
    fig, axes = plt.subplots(
        1,
        len(concept_effects),
        figsize=(fig_width, fig_height),
        sharey=True,
        squeeze=False,
        constrained_layout=True,
    )

    image = None
    for ax, concept_effect in zip(axes[0], concept_effects):
        values_by_key = {
            (row["layer"], row["position"]): row["average_iia"]
            for row in rows
            if row["concept_effect"] == concept_effect
        }
        heat = [
            [
                values_by_key.get((layer, position), float("nan"))
                for position in positions
            ]
            for layer in layers
        ]

        image = ax.imshow(
            heat, aspect="auto", origin="lower", vmin=0.0, vmax=1.0, cmap="viridis"
        )

        ax.set_title(f"IIA target concept: {concept_effect}")
        ax.set_xlabel("Patched position")
        ax.set_xticks(range(len(positions)), labels=positions)
        ax.set_yticks(range(len(layers)), labels=layers)

        for row_idx, layer in enumerate(layers):
            for col_idx, position in enumerate(positions):
                value = values_by_key.get((layer, position))
                if value is not None:
                    ax.text(
                        col_idx,
                        row_idx,
                        f"{value:.2f}",
                        ha="center",
                        va="center",
                        color="white" if value < 0.65 else "black",
                        fontsize=7,
                    )

    axes[0][0].set_ylabel("layer")
    if image is not None:
        fig.colorbar(image, ax=axes[0].tolist(), label="average_iia")
    output_path = output_dir / "coarse_localization_heat_tables.png"
    fig.savefig(output_path, dpi=200)
    plt.close(fig)
    print(f"Wrote {output_path}")


def sort_result_rows(rows: list[dict]) -> list[dict]:
    """Sort result rows by concept effect, position, and layer."""

    return sorted(
        rows,
        key=lambda row: (row["concept_effect"], row["position"], row["layer"]),
    )


def parse_layers(layers_text: str, n_layers: int) -> list[int]:
    """Parse comma-separated layer IDs and validate them against the model depth."""

    layers = []
    for part in layers_text.split(","):
        part = part.strip()
        if not part:
            continue
        layer = int(part)
        if layer < 0 or layer >= n_layers:
            raise ValueError(
                f"Layer {layer} is outside valid range [0, {n_layers - 1}]."
            )
        layers.append(layer)

    if not layers:
        raise ValueError("--layers did not contain any valid layer IDs.")

    return layers


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments for the localization sweep."""

    parser = argparse.ArgumentParser(description="Coarse month IIA localization sweep.")
    parser.add_argument("--model", default=os.environ.get("MODEL", MODEL))
    parser.add_argument("--dataset", default="datasets/months.jsonl")
    parser.add_argument("--output", default="results/coarse_localization.csv")
    parser.add_argument("--dtype", default="bfloat16", choices=sorted(DTYPE_MAP))
    parser.add_argument("--device-map", default="auto")
    parser.add_argument("--trust-remote-code", action="store_true")
    parser.add_argument(
        "--site-kind",
        default="resid_post",
        choices=["mlp_out", "resid_pre", "resid_mid", "resid_post"],
    )
    parser.add_argument("--max-rows", type=int, default=None)
    parser.add_argument("--max-layer", type=int, default=None)
    parser.add_argument(
        "--layers",
        default=None,
        help="Comma-separated layer IDs to sweep. Overrides --max-layer when set.",
    )
    parser.add_argument(
        "--position",
        default=None,
        choices=["input", "offset", "last_token"],
        help="Run only one coarse patch position. Defaults to all positions.",
    )
    parser.add_argument(
        "--debug-tokens",
        action="store_true",
        help="Print rendered prompts, token tables, selected positions, and top tokens.",
    )
    parser.add_argument(
        "--plot-only",
        action="store_true",
        help="Only plot heat tables from --output; skip model loading and evaluation.",
    )
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()

    if args.plot_only:
        result_rows = load_result_rows(Path(args.output))
        plot_heat_tables(result_rows, Path(args.output).parent)
        raise SystemExit(0)

    rows = load_rows(Path(args.dataset), limit=args.max_rows)
    model = boot_model(args.model, args.dtype, args.device_map, args.trust_remote_code)

    n_layers = n_layers_from_hook_dict(model)
    if args.layers is not None:
        layers = parse_layers(args.layers, n_layers)
    else:
        stop_layer = (
            n_layers if args.max_layer is None else min(args.max_layer, n_layers)
        )
        layers = list(range(stop_layer))

    results = []
    positions = [args.position] if args.position is not None else POSITIONS
    for layer in layers:
        site = resolve_site(model, layer, args.site_kind)
        for position in positions:
            scores = evaluate_layer_position(
                model=model,
                rows=rows,
                layer=layer,
                site=site,
                position=position,
                debug_tokens=args.debug_tokens,
            )
            for concept_effect, average_iia in scores.items():
                result = {
                    "average_iia": average_iia,
                    "position": position,
                    "layer": layer,
                    "concept_effect": concept_effect,
                }
                results.append(result)
                print(result)

    write_results(Path(args.output), results)
    print(f"Wrote {args.output}")
    plot_heat_tables(results, Path(args.output).parent)