social-data-attribution-w2/scripts/visualization/experiment1_single_bin.py

#!/usr/bin/env python3
"""Experiment 1: Single-Bin Unlearning — Figures 1-4."""
from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
import numpy as np
import pandas as pd
import seaborn as sns

from scripts.visualization._shared import (
    ACCURACY_METRICS, BASELINE, BENCH_COLORS,
    METRIC_LABELS, METRIC_LABELS_WRAP,
    paper_rc, save_fig,
)

RESULTS = [
    ("Entertainment", 1200, 0.757392, 0.745788, 0.588831, 0.7996, 10.2434, 11.7224),
    ("Politics", 1000, 0.753601, 0.754052, 0.592684, 0.8029, 10.7691, 14.0575),
    ("Finance & Business", 1200, 0.761941, 0.750664, 0.592612, 0.8032, 9.9664, 8.5489),
    ("Sports & Fitness", 1200, 0.750569, 0.743890, 0.588798, 0.7962, 11.0423, 12.4779),
    ("Health", 800, 0.742987, 0.753453, 0.594509, 0.8028, 10.1688, 9.6236),
    ("Home & Hobbies", 1400, 0.767248, 0.748103, 0.596145, 0.7986, 10.0434, 9.4269),
    ("Education & Jobs", 1400, 0.733889, 0.726606, 0.586996, 0.7763, 10.7598, 17.4806),
    ("Literature", 1800, 0.757392, 0.742136, 0.589454, 0.7925, 10.3608, 13.8343),
    ("Social Life", 1400, 0.752843, 0.750519, 0.598368, 0.7850, 10.0129, 13.1067),
    ("Religion", 1200, 0.755876, 0.745517, 0.584907, 0.7979, 10.2700, 13.2857),
    ("Science, Math & Technology", 1200, 0.755876, 0.744075, 0.591968, 0.8005, 11.7832, 15.1513),
    ("Food & Dining", 800, 0.750569, 0.750041, 0.595229, 0.8021, 9.9274, 9.1037),
    ("Travel & Tourism", 800, 0.761941, 0.750773, 0.593320, 0.7997, 9.9584, 10.1610),
    ("Crime & Law", 1200, 0.752843, 0.747759, 0.595534, 0.8025, 10.1236, 8.0123),
    ("Games", 1200, 0.750569, 0.754392, 0.592882, 0.7986, 10.0591, 10.8993),
    ("Transportation", 1200, 0.755118, 0.751964, 0.596603, 0.8025, 10.0331, 9.3671),
    ("Software", 1000, 0.763457, 0.751781, 0.594540, 0.8006, 9.9473, 8.8681),
    ("Art & Design", 1200, 0.754359, 0.743585, 0.587855, 0.7938, 10.8357, 13.0337),
    ("Fashion & Beauty", 1200, 0.758908, 0.754093, 0.597638, 0.8013, 10.0202, 10.8949),
    ("History & Geography", 1200, 0.752085, 0.739648, 0.589411, 0.7973, 12.0295, 13.4971),
    ("Software Development", 1000, 0.752843, 0.756395, 0.596281, 0.8021, 9.9568, 10.0005),
    ("Electronics & Hardware", 1200, 0.757392, 0.750078, 0.599743, 0.8009, 9.9506, 8.5438),
    ("Industrial", 1000, 0.747536, 0.751317, 0.594910, 0.7999, 10.5536, 10.8478),
    ("Adult Content", 800, 0.765732, 0.756617, 0.594293, 0.7977, 10.0062, 16.3168),
]

COLS = [
    "topic", "checkpoint", "gsm8k", "mmlu_socsci",
    "mmlu_stem", "socialiqa", "wikitext_ppl", "forget_ppl",
]


def build_dataframe() -> pd.DataFrame:
    df = pd.DataFrame(RESULTS, columns=COLS)
    for m in ACCURACY_METRICS:
        df[f"{m}_gamma"] = (df[m] - BASELINE[m]) / abs(BASELINE[m])
    base_ppl = BASELINE["wikitext_ppl"]
    df["wikitext_ppl_gamma"] = (df["wikitext_ppl"] - base_ppl) / abs(base_ppl)
    df["forget_ppl_gamma"] = (df["forget_ppl"] - base_ppl) / abs(base_ppl)
    df["mean_acc_gamma"] = df[[f"{m}_gamma" for m in ACCURACY_METRICS]].mean(axis=1)
    return df


HEATMAP_COL_ORDER = ["socialiqa", "mmlu_socsci", "gsm8k", "mmlu_stem"]


def fig1_heatmap(df: pd.DataFrame, output_dir: Path) -> None:
    gamma_cols = [f"{m}_gamma" for m in HEATMAP_COL_ORDER]
    df_s = df.sort_values("mean_acc_gamma", ascending=True)
    topics = df_s["topic"].values

    fig, (ax_a, ax_p) = plt.subplots(
        1, 2, figsize=(9, 8),
        gridspec_kw={"width_ratios": [4, 1.3], "wspace": 0.08},
    )

    acc = df_s[gamma_cols].values
    col_labels = [METRIC_LABELS_WRAP[m] for m in HEATMAP_COL_ORDER]
    vmax_a = max(abs(acc.min()), abs(acc.max()), 0.01)
    sns.heatmap(
        acc, ax=ax_a,
        xticklabels=col_labels,
        yticklabels=topics,
        cmap="RdBu", center=0, vmin=-vmax_a, vmax=vmax_a,
        annot=True, fmt=".1%", annot_kws={"fontsize": 6.5},
        linewidths=0.4, linecolor="white",
        cbar_kws={"label": r"Accuracy $\gamma$", "shrink": 0.35, "aspect": 18, "pad": 0.02},
    )
    ax_a.axvline(x=2, color="black", linewidth=1.0)
    ax_a.set_yticklabels(ax_a.get_yticklabels(), fontsize=8, rotation=0)
    ax_a.tick_params(axis="both", length=0)
    ax_a.set_xlabel("")
    ax_a.set_ylabel("")

    ppl = df_s[["wikitext_ppl_gamma"]].values
    sns.heatmap(
        ppl, ax=ax_p,
        xticklabels=["Wikitext\nPPL"],
        yticklabels=False,
        cmap="Purples", vmin=0, vmax=ppl.max() * 1.05,
        annot=True, fmt=".1%", annot_kws={"fontsize": 6.5},
        linewidths=0.4, linecolor="white",
        cbar_kws={"label": r"PPL $\gamma$", "shrink": 0.35, "aspect": 18, "pad": 0.15},
    )
    ax_p.set_yticks([])
    ax_p.tick_params(axis="x", length=0)
    fig.subplots_adjust(left=0.17)
    save_fig(fig, output_dir, "fig1_cross_benchmark_heatmap")


def fig2_bars(df: pd.DataFrame, output_dir: Path) -> None:
    fig, axes = plt.subplots(2, 2, figsize=(10, 9))
    for idx, m in enumerate(ACCURACY_METRICS):
        ax = axes.flat[idx]
        col = f"{m}_gamma"
        ds = df.sort_values(col)
        vals = ds[col].values
        ax.barh(range(len(ds)), vals, color=BENCH_COLORS[m], edgecolor="none", height=0.7)
        ax.set_yticks(range(len(ds)))
        ax.set_yticklabels(ds["topic"], fontsize=7)
        ax.axvline(0, color="black", linewidth=0.5)
        ax.set_xlabel(r"$\gamma$")
        ax.set_title(METRIC_LABELS[m], fontweight="bold")
        ax.xaxis.set_major_formatter(mticker.PercentFormatter(xmax=1, decimals=1))
        ax.grid(True, axis="x", alpha=0.2, linewidth=0.3)
        lim = max(abs(vals.min()), abs(vals.max()), 0.005) * 1.3
        ax.set_xlim(-lim, lim)
    fig.tight_layout()
    save_fig(fig, output_dir, "fig2_per_metric_gamma_bars")


def fig3_lollipop(df: pd.DataFrame, output_dir: Path) -> None:
    ds = df.sort_values("forget_ppl", ascending=True)
    topics = ds["topic"].values
    forget_ppl = ds["forget_ppl"].values
    mean_gamma = ds["mean_acc_gamma"].values
    y = np.arange(len(ds))

    fig, (ax_l, ax_r) = plt.subplots(
        1, 2, figsize=(7, 7), sharey=True,
        gridspec_kw={"width_ratios": [1, 1], "wspace": 0.12},
    )

    ax_l.hlines(y, 0, forget_ppl - BASELINE["wikitext_ppl"], color="#636363", linewidth=0.6)
    ax_l.scatter(
        forget_ppl - BASELINE["wikitext_ppl"], y,
        c=forget_ppl, cmap="YlOrRd", s=35,
        edgecolors="black", linewidth=0.3, zorder=3,
        vmin=forget_ppl.min(), vmax=forget_ppl.max(),
    )
    ax_l.axvline(0, color="black", linewidth=0.4)
    ax_l.set_xlabel(r"Forget PPL increase over baseline")
    ax_l.set_yticks(y)
    ax_l.set_yticklabels(topics, fontsize=7)
    ax_l.grid(True, axis="x", alpha=0.2, linewidth=0.3)

    colors = [BENCH_COLORS["gsm8k"] if v < -0.005 else "#bdbdbd" for v in mean_gamma]
    ax_r.hlines(y, 0, mean_gamma, color="#636363", linewidth=0.6)
    ax_r.scatter(mean_gamma, y, color=colors, s=35, edgecolors="black", linewidth=0.3, zorder=3)
    ax_r.axvline(0, color="black", linewidth=0.4)
    ax_r.set_xlabel(r"Mean accuracy $\gamma$")
    ax_r.xaxis.set_major_formatter(mticker.PercentFormatter(xmax=1, decimals=1))
    ax_r.grid(True, axis="x", alpha=0.2, linewidth=0.3)

    save_fig(fig, output_dir, "fig3_forget_ppl_vs_retention")


def fig4_ppl_bars(df: pd.DataFrame, output_dir: Path) -> None:
    ds = df.sort_values("wikitext_ppl", ascending=True)
    fig, ax = plt.subplots(figsize=(5, 7))
    increase = ds["wikitext_ppl"].values - BASELINE["wikitext_ppl"]
    colors = [
        "#d73027" if v > 1.5 else "#fc8d59" if v > 0.8 else "#fee08b"
        for v in increase
    ]
    ax.barh(range(len(ds)), ds["wikitext_ppl"], color=colors, height=0.7, edgecolor="none")
    ax.axvline(
        BASELINE["wikitext_ppl"], color="black", linestyle="--", linewidth=0.7,
        label=f"Baseline ({BASELINE['wikitext_ppl']:.2f})",
    )
    ax.set_yticks(range(len(ds)))
    ax.set_yticklabels(ds["topic"], fontsize=7)
    ax.set_xlabel(r"Wikitext-2 perplexity $\longrightarrow$ (higher = worse language modeling)")
    ax.legend(loc="lower right", fontsize=7, framealpha=0.9)
    ax.grid(True, axis="x", alpha=0.2, linewidth=0.3)
    for i, (ppl, inc) in enumerate(zip(ds["wikitext_ppl"], increase)):
        ax.text(ppl + 0.05, i, f"+{inc:.2f}", va="center", ha="left", fontsize=6, color="#333")
    ax.set_xlim(BASELINE["wikitext_ppl"] - 0.3, ds["wikitext_ppl"].max() + 0.5)
    save_fig(fig, output_dir, "fig4_wikitext_ppl_by_topic")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", type=Path, default=Path("artifacts/experiment1_figures"))
    args = parser.parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)
    paper_rc()
    df = build_dataframe()
    print("Generating Experiment 1 figures...")
    fig1_heatmap(df, args.output_dir)
    fig2_bars(df, args.output_dir)
    fig3_lollipop(df, args.output_dir)
    fig4_ppl_bars(df, args.output_dir)
    print(f"All figures saved to {args.output_dir}/")


if __name__ == "__main__":
    main()