social-data-attribution-w2/scripts/figure1/build_panel4_unlearning.py

"""Panel-4 thumbnail for Fig 1 (unlearning validation).

Reproduces the 6-topic × 2-benchmark mini-table currently hand-coded in
``figure_1_hero.drawio``. Each cell shows the absolute change in
accuracy (post − baseline, percentage points) for influence-targeted
unlearning of that topic, on either SocialIQA or ARC-Challenge. The
curated 6 topics emphasize the *contrastive* benchmark sensitivity
pattern: topics where social reasoning suffers but factual recall is
spared (and vice versa).

Usage:
    uv run python scripts/figure1/build_panel4_unlearning.py
"""

from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.colors import TwoSlopeNorm

from dolma.distribution_report.style import style_size

# Curated topics (lowercase data label, display name) mirroring drawio panel 4.
_TOPICS: list[tuple[str, str]] = [
    ("social_life",            "Social Life"),
    ("politics",               "Politics"),
    ("literature",             "Literature"),
    ("health",                 "Health"),
    ("entertainment",          "Entertainment"),
    ("history_and_geography",  "History & Geography"),
]

# Benchmarks displayed in the panel. Each tuple is (criterion key in xlsx,
# accuracy column in xlsx, display heading).
_BENCHMARKS: list[tuple[str, str, str]] = [
    ("socialiqa",      "socialiqa",      "SocialIQA\nΔ acc (pp)"),
    ("arc_challenge",  "arc_challenge",  "ARC-Chal.\nΔ acc (pp)"),
]

# Cap the diverging color scale so a single very large value (e.g.
# Social Life on SocialIQA at ~−15 pp) does not wash out the other
# cells. Out-of-range cells still show the exact number but are clipped
# to the saturation color.
_COLOR_CAP_PP = 5.0


def _load_baseline_and_targeted(xlsx_path: Path) -> tuple[dict[str, float], pd.DataFrame]:
    """Return baseline accuracies and a long-format targeted-unlearning dataframe.

    ``targeted`` has columns: topic, criterion, accuracy_col, post_acc.
    """
    df = pd.read_excel(xlsx_path, sheet_name="Unlearning With Influence Score")
    base_row = df.iloc[0]
    baseline = {bench: float(base_row[bench]) for _, bench, _ in _BENCHMARKS}

    body = df.iloc[1:].copy()
    body["experiment"] = body["Unnamed: 0"].ffill()
    body = body[body["experiment"].astype(str).str.startswith("Experiment A")]

    records = []
    for _, row in body.iterrows():
        topic = str(row["Bin(s)"]).strip().lower()
        criterion = row.get("top-2000 criteria")
        if pd.isna(criterion):
            continue
        criterion = str(criterion).strip()
        for _, bench_col, _ in _BENCHMARKS:
            if criterion != bench_col:
                continue
            val = row.get(bench_col)
            if pd.notna(val):
                records.append({
                    "topic": topic,
                    "criterion": criterion,
                    "post_acc": float(val),
                })
    return baseline, pd.DataFrame(records)


def build_panel(
    xlsx_path: Path,
    out_dir: Path,
    figsize: tuple[float, float] = (3.6, 2.6),
    dpi: int = 300,
) -> None:
    baseline, targeted = _load_baseline_and_targeted(xlsx_path)

    # Build the delta matrix: rows = topics, cols = benchmarks, value = pp change.
    matrix = np.full((len(_TOPICS), len(_BENCHMARKS)), np.nan)
    for i, (topic_key, _) in enumerate(_TOPICS):
        for j, (crit, bench_col, _) in enumerate(_BENCHMARKS):
            base = baseline[bench_col]
            row = targeted[(targeted["topic"] == topic_key) & (targeted["criterion"] == crit)]
            if not row.empty:
                # Δ in percentage points: (post − baseline) × 100
                matrix[i, j] = (row["post_acc"].iloc[0] - base) * 100.0

    fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
    # Diverging scale: blue = capability preserved/positive, red = degraded.
    # Centered at zero. Capped at ±_COLOR_CAP_PP so a single dominant value
    # does not visually dilute the others.
    norm = TwoSlopeNorm(vmin=-_COLOR_CAP_PP, vcenter=0.0, vmax=_COLOR_CAP_PP)
    ax.imshow(matrix, cmap="RdBu", norm=norm, aspect="auto")

    ax.set_xticks(range(len(_BENCHMARKS)))
    ax.set_xticklabels([disp for _, _, disp in _BENCHMARKS],
                        fontsize=style_size("TICK"), family="serif")
    ax.set_yticks(range(len(_TOPICS)))
    ax.set_yticklabels([disp for _, disp in _TOPICS],
                        fontsize=style_size("TICK"), family="serif")
    ax.tick_params(axis="both", which="both", length=0)
    for spine in ax.spines.values():
        spine.set_visible(False)

    for i in range(len(_TOPICS)):
        for j in range(len(_BENCHMARKS)):
            val = matrix[i, j]
            if np.isnan(val):
                continue
            text_color = "white" if abs(val) > _COLOR_CAP_PP * 0.55 else "black"
            # Format with sign so direction is obvious; 2-decimal precision.
            ax.text(j, i, f"{val:+.2f}", ha="center", va="center",
                    fontsize=style_size("ANNOTATION"), color=text_color, family="serif")

    ax.set_xticks(np.arange(len(_BENCHMARKS) + 1) - 0.5, minor=True)
    ax.set_yticks(np.arange(len(_TOPICS) + 1) - 0.5, minor=True)
    ax.grid(which="minor", color="white", linewidth=1.2)
    ax.tick_params(which="minor", length=0)

    fig.tight_layout(pad=0.4)
    out_dir.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_dir / "panel4_unlearning.pdf", bbox_inches="tight")
    fig.savefig(out_dir / "panel4_unlearning.png", bbox_inches="tight", dpi=dpi)
    plt.close(fig)
    print(f"Wrote {out_dir / 'panel4_unlearning.pdf'} (+ .png)")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--xlsx",
        type=Path,
        default=Path("outputs/cache/saehee_unlearning_results.xlsx"),
    )
    parser.add_argument("--out-dir", type=Path, default=Path("artifacts/figure_1"))
    args = parser.parse_args()
    build_panel(args.xlsx, args.out_dir)


if __name__ == "__main__":
    main()