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

#!/usr/bin/env python3
# Compute top-5 bins by correctness differential (correct minus incorrect)
# for each of the four primary benchmarks.

import csv
import os
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[2]
DATA_ROOT = Path(
    os.environ.get(
        "SDA_ZSCORED_SPLITS_ROOT",
        str(REPO_ROOT / "artifacts/zscored_bin_scores/splits"),
    )
)

BENCHMARKS = {
    "SocialIQA": "zscored_socialiqa",
    "MMLU Social Sci.": "zscored_mmlu_social_science",
    "ARC-Challenge": "zscored_arc_challenge",
    "MMLU STEM": "zscored_mmlu_stem",
}

TOPIC_DISPLAY = {
    "adult_content": "Adult Content",
    "art_and_design": "Art \\& Design",
    "crime_and_law": "Crime \\& Law",
    "education_and_jobs": "Education \\& Jobs",
    "electronics_and_hardware": "Hardware",
    "entertainment": "Entertainment",
    "fashion_and_beauty": "Fashion \\& Beauty",
    "finance_and_business": "Finance \\& Business",
    "food_and_dining": "Food \\& Dining",
    "games": "Games",
    "health": "Health",
    "history_and_geography": "History",
    "home_and_hobbies": "Home \\& Hobbies",
    "industrial": "Industrial",
    "literature": "Literature",
    "politics": "Politics",
    "religion": "Religion",
    "science_math_and_technology": "Sci.~\\& Tech.",
    "social_life": "Social Life",
    "software": "Software",
    "software_development": "Software Dev.",
    "sports_and_fitness": "Sports \\& Fitness",
    "transportation": "Transportation",
    "travel_and_tourism": "Travel",
}

FORMAT_DISPLAY = {
    "about_org": "About (Org.)",
    "about_pers": "About (Pers.)",
    "academic_writing": "Academic Writing",
    "audio_transcript": "Audio Transcript",
    "comment_section": "Comment Section",
    "content_listing": "Content Listing",
    "creative_writing": "Creative Writing",
    "customer_support": "Customer Support",
    "documentation": "Documentation",
    "faq": "FAQ",
    "knowledge_article": "Knowledge Article",
    "legal_notices": "Legal Notices",
    "listicle": "Listicle",
    "news_article": "News Article",
    "news_org": "News (Org.)",
    "nonfiction_writing": "Nonfiction Writing",
    "personal_blog": "Personal Blog",
    "product_page": "Product Page",
    "q_a_forum": "Q\\&A Forum",
    "spam_ads": "Spam / Ads",
    "structured_data": "Structured Data",
    "truncated": "Truncated",
    "tutorial": "Tutorial",
    "user_review": "User Review",
}


def read_zscores(path: Path) -> dict[tuple[str, str], float]:
    out = {}
    with path.open() as fh:
        for row in csv.DictReader(fh):
            out[(row["topic_label"], row["format_label"])] = float(row["zscore"])
    return out


def main() -> None:
    rows_for_table = []
    for bench_name, prefix in BENCHMARKS.items():
        correct = read_zscores(DATA_ROOT / f"{prefix}_correct.csv")
        incorrect = read_zscores(DATA_ROOT / f"{prefix}_incorrect.csv")
        bins = set(correct) & set(incorrect)
        diffs = {b: correct[b] - incorrect[b] for b in bins}
        ranked = sorted(diffs.items(), key=lambda kv: kv[1], reverse=True)
        top5 = ranked[:5]
        bot5 = ranked[-5:]
        print(f"== {bench_name} ==  ({len(bins)} bins compared)")
        print("  Top-5 correct > incorrect (positive differential):")
        for (t, f), d in top5:
            print(
                f"    {d:+.2f}   {TOPIC_DISPLAY.get(t, t)}~$\\times$~"
                f"{FORMAT_DISPLAY.get(f, f)}"
            )
        print("  Top-5 incorrect > correct (negative differential):")
        for (t, f), d in bot5:
            print(
                f"    {d:+.2f}   {TOPIC_DISPLAY.get(t, t)}~$\\times$~"
                f"{FORMAT_DISPLAY.get(f, f)}"
            )
        print()
        rows_for_table.append((bench_name, top5, bot5))


if __name__ == "__main__":
    main()