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

"""Shared constants and utilities for unlearning experiment figures."""
from __future__ import annotations

from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt


BASELINE = {
    "gsm8k": 0.758908,
    "mmlu_socsci": 0.750846,
    "mmlu_stem": 0.597871,
    "socialiqa": 0.802900,
    "wikitext_ppl": 9.1559,
}

ACCURACY_METRICS = ["gsm8k", "mmlu_socsci", "mmlu_stem", "socialiqa"]
MAIN_METRICS = ["socialiqa", "mmlu_socsci", "mmlu_stem"]

METRIC_LABELS = {
    "gsm8k": "GSM8K",
    "mmlu_socsci": "MMLU Social Science",
    "mmlu_stem": "MMLU STEM",
    "socialiqa": "SocialIQA",
    "wikitext_ppl": "Wikitext PPL",
    "forget_ppl": "Forget PPL",
}

METRIC_LABELS_WRAP = {
    "gsm8k": "GSM8K",
    "mmlu_socsci": "MMLU\nSocial Sci.",
    "mmlu_stem": "MMLU\nSTEM",
    "socialiqa": "SocialIQA",
    "wikitext_ppl": "Wikitext\nPPL",
    "forget_ppl": "Forget\nPPL",
}

BENCH_COLORS = {
    "socialiqa": "#984EA3",
    "mmlu_socsci": "#984EA3",
    "gsm8k": "#1b7837",
    "mmlu_stem": "#1b7837",
    "wikitext_ppl": "#ff7f00",
    "forget_ppl": "#ff7f00",
}

BENCH_HATCHES = {
    "socialiqa": "",
    "mmlu_socsci": "//",
    "gsm8k": "",
    "mmlu_stem": "//",
}

BENCH_STYLE_LABELS = {
    "socialiqa": "SocialIQA (reasoning)",
    "mmlu_socsci": "MMLU Social Sci. (knowledge)",
    "gsm8k": "GSM8K (reasoning)",
    "mmlu_stem": "MMLU STEM (knowledge)",
}

EXP_COLORS = {
    "exp1": "#a6cee3",
    "exp2": "#fb9a99",
    "exp3": "#fdbf6f",
}

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

FULL_TOPIC_NAMES = {
    "Science & Tech.": "Science, Math & Technology",
    "Software Dev.": "Software Development",
    "History": "History & Geography",
    "Travel": "Travel & Tourism",
}


def paper_rc():
    plt.rcParams.update({
        "font.family": "serif",
        "font.serif": ["Liberation Serif", "DejaVu Serif", "Times New Roman"],
        "mathtext.fontset": "dejavuserif",
        "font.size": 9,
        "axes.titlesize": 10,
        "axes.labelsize": 9,
        "xtick.labelsize": 8,
        "ytick.labelsize": 8,
        "legend.fontsize": 8,
        "figure.dpi": 300,
        "savefig.dpi": 300,
        "savefig.bbox": "tight",
        "axes.spines.top": False,
        "axes.spines.right": False,
        "axes.linewidth": 0.5,
        "xtick.major.width": 0.5,
        "ytick.major.width": 0.5,
        "grid.linewidth": 0.3,
        "grid.alpha": 0.2,
    })


def save_fig(fig, output_dir: Path, name: str) -> None:
    for ext in ("pdf", "svg", "png"):
        fig.savefig(output_dir / f"{name}.{ext}", dpi=300, bbox_inches="tight")
    plt.close(fig)
    print(f"  {name} (pdf/svg/png)")


def gamma(score: float, metric: str) -> float:
    return (score - BASELINE[metric]) / abs(BASELINE[metric])