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

#!/usr/bin/env python3
# Compute top-3 / bottom-3 format-level z-score marginals per benchmark.
#
# Reads the 4 primary-benchmark z-scored bin CSVs from the social-data-attribution
# repo, averages z-scores across topics for each format (24 topics -> 1 format-level
# mean), ranks the 24 formats, and emits the top-3 positive and bottom-3 negative
# per benchmark as a LaTeX-ready summary.

import csv
import os
import sys
from collections import defaultdict
from pathlib import Path

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

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

# Pretty-print labels matching the WebOrganizer taxonomy used in the paper.
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 format_marginals(csv_path: Path) -> dict[str, float]:
    sums: dict[str, float] = defaultdict(float)
    counts: dict[str, int] = defaultdict(int)
    with csv_path.open() as fh:
        reader = csv.DictReader(fh)
        for row in reader:
            fmt = row["format_label"]
            sums[fmt] += float(row["zscore"])
            counts[fmt] += 1
    return {fmt: sums[fmt] / counts[fmt] for fmt in sums}


def main() -> None:
    print("# Format-level z-score marginals (mean across 24 topics per format)")
    print()
    rows_for_table = []
    for bench_name, fname in BENCHMARKS.items():
        marginals = format_marginals(DATA_ROOT / fname)
        ranked = sorted(marginals.items(), key=lambda kv: kv[1], reverse=True)
        top3 = ranked[:3]
        bot3 = ranked[-3:]
        print(f"== {bench_name} ==")
        for fmt, z in top3:
            print(f"  +{z:+.2f}  {FORMAT_DISPLAY.get(fmt, fmt)}")
        print("  ...")
        for fmt, z in bot3:
            print(f"  {z:+.2f}  {FORMAT_DISPLAY.get(fmt, fmt)}")
        print()
        rows_for_table.append((bench_name, top3, bot3))

    # Emit a LaTeX tabular block for direct paste.
    print()
    print("# LaTeX tabular (top-3 positive, top-3 negative per benchmark)")
    print()
    print(r"\begin{tabular}{lll}")
    print(r"\toprule")
    print(
        r"\textbf{Benchmark} & \textbf{Top-3 positive (z)} & \textbf{Top-3 negative (z)} \\"
    )
    print(r"\midrule")
    for bench_name, top3, bot3 in rows_for_table:
        top_str = ", ".join(f"{FORMAT_DISPLAY.get(f, f)}~({z:+.2f})" for f, z in top3)
        bot_str = ", ".join(f"{FORMAT_DISPLAY.get(f, f)}~({z:+.2f})" for f, z in bot3)
        print(f"{bench_name} & {top_str} & {bot_str} \\\\")
    print(r"\bottomrule")
    print(r"\end{tabular}")


if __name__ == "__main__":
    sys.exit(main())