social-data-attribution-w2/scripts/analysis/tombench_attribution_alignment.py

"""Compare ToMBench attribution rankings with unlearning damage rankings."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import cast

import pandas as pd
from scipy.stats import spearmanr

PRAGMATIC = ["hinting_task_test", "faux_pas_recognition_test", "strange_story_task"]
ALL_TOMBENCH = [
    "ambiguous_story_task",
    "completion_of_failed_actions",
    "discrepant_desires",
    "discrepant_emotions",
    "discrepant_intentions",
    "emotion_regulation",
    "false_belief_task",
    "faux_pas_recognition_test",
    "hidden_emotions",
    "hinting_task_test",
    "knowledge_attention_links",
    "knowledge_pretend_play_links",
    "moral_emotions",
    "multiple_desires",
    "percepts_knowledge_links",
    "persuasion_story_task",
    "prediction_of_actions",
    "scalar_implicature_test",
    "strange_story_task",
    "unexpected_outcome_test",
]


def topic_attr(bin_scores_dir: Path, probe: str) -> pd.Series:
    df = pd.read_csv(bin_scores_dir / f"queries_tombench_{probe}_bin_scores.csv")
    df["mass"] = df["mean_score"] * df["doc_count"]
    return cast(pd.Series, df.groupby("topic_label")["mass"].sum())


def gamma_by_topic(gamma_cells: dict, probe: str) -> dict[str, float]:
    cells = gamma_cells[f"tombench_{probe}"]
    return {
        topic: cell["mean"] for topic, cell in cells.items() if cell["mean"] is not None
    }


def analyze(
    gamma_cells: dict, bin_scores_dir: Path, probes: list[str], label: str
) -> None:
    rhos = []
    social_attr_ranks: list[int] = []
    social_gamma_ranks: list[int] = []

    for probe in probes:
        attr = topic_attr(bin_scores_dir, probe)
        gamma = gamma_by_topic(gamma_cells, probe)
        attr_topics = {str(topic) for topic in attr.index}
        common = [topic for topic in gamma if topic in attr_topics]
        rho, _ = spearmanr(
            [float(attr.loc[topic]) for topic in common],
            [gamma[topic] for topic in common],
        )
        rhos.append(rho)

        attr_sorted = [str(topic) for topic in attr.sort_values(ascending=False).index]
        gamma_sorted = sorted(common, key=lambda topic: gamma[topic])
        if "social_life" in attr_sorted:
            social_attr_ranks.append(attr_sorted.index("social_life") + 1)
        if "social_life" in gamma_sorted:
            social_gamma_ranks.append(gamma_sorted.index("social_life") + 1)

    mean_rho = sum(rhos) / len(rhos)
    social_attr_rank = sorted(social_attr_ranks)[len(social_attr_ranks) // 2]
    social_gamma_rank = sorted(social_gamma_ranks)[len(social_gamma_ranks) // 2]

    print(f"\n== {label} ({len(probes)} probes) ==")
    print(
        f"   mean Spearman(attr, gamma) = {mean_rho:+.2f} "
        "(negative means attribution support predicts unlearning damage)"
    )
    print(f"   social_life median attribution-support rank: {social_attr_rank}/24")
    print(f"   social_life median unlearning-damage rank:    {social_gamma_rank}/24")
    print(f"   per-probe social_life attr-support ranks: {social_attr_ranks}")
    print(f"   per-probe social_life damage ranks:       {social_gamma_ranks}")


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--gamma-ci", type=Path, required=True)
    parser.add_argument("--bin-scores-dir", type=Path, required=True)
    args = parser.parse_args(argv)

    gamma_ci = json.loads(args.gamma_ci.read_text())
    gamma_cells = gamma_ci["results"]["acc_uncond"]["absolute"]["net"]
    analyze(gamma_cells, args.bin_scores_dir, PRAGMATIC, "social-pragmatic subtasks")
    analyze(gamma_cells, args.bin_scores_dir, ALL_TOMBENCH, "all 20 ToMBench")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())