social-data-attribution-w2/scripts/unlearning/reaggregate_mmlu_macro.py

"""Re-aggregate stitched unlearning prediction files.

MMLU suites are macro-averaged across subtasks. Other retained unlearning
benchmark suites are micro-averaged. The output CSV is intended for local
review of existing prediction JSONL files, not for running evaluations.
"""

from __future__ import annotations

import argparse
import csv
import json
from collections import defaultdict
from pathlib import Path

SUITE_TO_KEY = {
    "mmlu_social_sciences": "mmlu_social_science",
    "mmlu_social_science": "mmlu_social_science",
    "mmlu_stem": "mmlu_stem",
    "social_iqa": "socialiqa",
    "socialiqa": "socialiqa",
}

BASELINES = {
    "mmlu_social_science": 0.750846,
    "mmlu_stem": 0.597871,
    "socialiqa": 0.802900,
}

BENCHES = tuple(BASELINES)


def reaggregate(pred_path: Path) -> dict[str, float]:
    """Read a stitched predictions JSONL and return benchmark accuracies."""
    task_correct: dict[str, dict[str, int]] = defaultdict(lambda: defaultdict(int))
    task_total: dict[str, dict[str, int]] = defaultdict(lambda: defaultdict(int))
    correct: dict[str, int] = defaultdict(int)
    total: dict[str, int] = defaultdict(int)

    for line in pred_path.read_text(encoding="utf-8").splitlines():
        if not line.strip():
            continue
        row = json.loads(line)
        suite = SUITE_TO_KEY.get(str(row.get("task_suite", "")))
        if suite is None:
            continue
        task = str(row.get("task_name", suite))
        hit = int(bool(row.get("is_correct")))
        if "mmlu" in suite:
            task_correct[suite][task] += hit
            task_total[suite][task] += 1
        else:
            correct[suite] += hit
            total[suite] += 1

    results: dict[str, float] = {}
    for suite, counts in task_total.items():
        per_task = [
            task_correct[suite][task] / counts[task]
            for task in counts
            if counts[task] > 0
        ]
        results[suite] = sum(per_task) / len(per_task) if per_task else float("nan")

    for suite, count in total.items():
        results[suite] = correct[suite] / count if count else float("nan")

    return results


def gamma(acc: float, bench: str) -> float | None:
    """Return relative accuracy change against the retained baseline."""
    baseline = BASELINES.get(bench)
    if baseline is None or acc != acc:
        return None
    return (acc - baseline) / baseline


def infer_label(pred_path: Path) -> tuple[str, str]:
    """Infer a topic and context label from common run path shapes."""
    parts = pred_path.parts
    for index, part in enumerate(parts):
        if part == "expA" and index + 2 < len(parts):
            return parts[index + 1], parts[index + 2]
    for index, part in enumerate(parts):
        if part in {"unlearn", "runs"} and index + 1 < len(parts):
            return parts[index + 1], "all"
    return pred_path.parent.name, "unknown"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--search-root",
        default="runs/unlearn",
        help="Root directory to search for prediction JSONL files.",
    )
    parser.add_argument(
        "--output",
        default="results/unlearning_mmlu_macro.csv",
        help="Output CSV path.",
    )
    parser.add_argument(
        "--pattern",
        default="*_predictions.jsonl",
        help="Glob pattern for prediction files.",
    )
    parser.add_argument(
        "--include-olmes-runs",
        action="store_true",
        help="Include per-task files inside olmes_runs/ subdirectories.",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    root = Path(args.search_root).expanduser()
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)

    pred_files = sorted(root.rglob(args.pattern))
    if not args.include_olmes_runs:
        pred_files = [path for path in pred_files if "olmes_runs" not in path.parts]

    rows = []
    for pred_path in pred_files:
        topic, bench_context = infer_label(pred_path)
        scores = reaggregate(pred_path)
        if not scores:
            continue

        row: dict[str, str] = {
            "topic": topic,
            "bench_context": bench_context,
            "predictions_file": str(pred_path),
        }
        for bench in BENCHES:
            acc = scores.get(bench, float("nan"))
            change = gamma(acc, bench)
            row[bench] = f"{acc:.6f}" if acc == acc else ""
            row[f"gamma_{bench}"] = f"{change * 100:+.3f}" if change is not None else ""
        rows.append(row)

    fieldnames = (
        ["topic", "bench_context"]
        + list(BENCHES)
        + [f"gamma_{bench}" for bench in BENCHES]
        + ["predictions_file"]
    )
    with output.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    print(f"Wrote {len(rows)} rows to {output}")
    return 0


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