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

#!/usr/bin/env python3
# Verify §I (Extended Benchmark Analysis) caption claims for GSM8K + ARC-Easy.
# Specifically:
#  - Documentation as dominant positive format for STEM-oriented benchmarks
#  - Literature distinctive to SocialIQA (i.e., not positive for GSM8K/ARC-Easy)
#  - Industrial / Sci.&Tech. positive across STEM tasks

import csv
import os
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 Sci.": "zscored_mmlu_social_science.csv",
    "ARC-Challenge": "zscored_arc_challenge.csv",
    "MMLU STEM": "zscored_mmlu_stem.csv",
    "GSM8K": "zscored_gsm8k.csv",
    "ARC-Easy": "zscored_arc_easy.csv",
}

TOPICS_OF_INTEREST = [
    "literature",
    "industrial",
    "science_math_and_technology",
    "social_life",
    "software_development",
]

FORMATS_OF_INTEREST = [
    "documentation",
    "academic_writing",
    "knowledge_article",
    "customer_support",
    "q_a_forum",
]


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 topic_marg(z):
    s, c = defaultdict(float), defaultdict(int)
    for (t, _), v in z.items():
        s[t] += v
        c[t] += 1
    return {t: s[t] / c[t] for t in s}


def format_marg(z):
    s, c = defaultdict(float), defaultdict(int)
    for (_, f), v in z.items():
        s[f] += v
        c[f] += 1
    return {f: s[f] / c[f] for f in s}


def main() -> None:
    data = {n: read_zscores(DATA_ROOT / fn) for n, fn in BENCHMARKS.items()}
    tmargs = {n: topic_marg(z) for n, z in data.items()}
    fmargs = {n: format_marg(z) for n, z in data.items()}

    # Top-3 positive and bottom-3 negative formats per benchmark
    print("# Format-level top-3 positive per benchmark")
    for name in BENCHMARKS:
        ranked = sorted(fmargs[name].items(), key=lambda kv: kv[1], reverse=True)
        top3 = ranked[:3]
        print(f"  {name:18s}  " + ", ".join(f"{f} ({v:+.2f})" for f, v in top3))
    print()

    # Topic-level top-3 positive per benchmark
    print("# Topic-level top-3 positive per benchmark")
    for name in BENCHMARKS:
        ranked = sorted(tmargs[name].items(), key=lambda kv: kv[1], reverse=True)
        top3 = ranked[:3]
        print(f"  {name:18s}  " + ", ".join(f"{t} ({v:+.2f})" for t, v in top3))
    print()

    # Specific topics/formats of interest
    print("# Topics of interest (marginal z-scores per benchmark)")
    print("  " + " " * 30 + "  ".join(f"{n:>15s}" for n in BENCHMARKS))
    for topic in TOPICS_OF_INTEREST:
        row = [f"{tmargs[n].get(topic, float('nan')):+.2f}" for n in BENCHMARKS]
        print(f"  {topic:30s}" + "  ".join(f"{v:>15s}" for v in row))
    print()
    print("# Formats of interest (marginal z-scores per benchmark)")
    print("  " + " " * 30 + "  ".join(f"{n:>15s}" for n in BENCHMARKS))
    for fmt in FORMATS_OF_INTEREST:
        row = [f"{fmargs[n].get(fmt, float('nan')):+.2f}" for n in BENCHMARKS]
        print(f"  {fmt:30s}" + "  ".join(f"{v:>15s}" for v in row))


if __name__ == "__main__":
    main()