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

#!/usr/bin/env python3
# pyright: reportArgumentType=false, reportAttributeAccessIssue=false
"""Build bin-level influence-targeted forget-sets on the LOCAL 6T cache (self-service;
no per-doc influence parquet / no doc_id->text mapping needed).

For each (benchmark, topic) we rank the topic's 6T-cache docs by their (topic, format)
bin z-score for that benchmark and take the top-K doc_ids. This is the "bin-targeted"
operationalization of expA (influence-ranked within one topic). expC = same ranking
corpus-wide (no topic filter). Writes one doc_id-per-line file per recipe.

Inputs (all local):
  - 6T cache: doc_id, weborganizer_topic, weborganizer_format  (DOLMA_CACHE)
  - z-scored bins: ~/dev/data-attribution/.../zscored_<bench>.csv (topic_label, format_label, zscore)
Outputs: ~/scratch/n16_selectivity/forgetsets_binlevel/{expA__<topic>__<bench>.txt, expC__<bench>.txt}
"""

import os
import sys
from pathlib import Path
import pandas as pd
from datasets import load_from_disk

HOME = Path.home()
ZS = HOME / "dev/data-attribution/artifacts/zscored_bin_scores/aggregated"
OUT = HOME / "scratch/n16_selectivity/forgetsets_binlevel"
DOLMA = os.environ.get(
    "DOLMA_CACHE",
    "/storage/ice-shared/cs7634/seom35/hf_cache/datasets/dolma3_6t_filtered",
)
ZSCORED_FILE = {
    "socialiqa": "socialiqa",
    "mmlu_social_science": "mmlu_social_science",
    "mmlu_stem": "mmlu_stem",
    "arc_challenge": "arc_challenge",
}
K = int(os.environ.get("FORGET_K", "200"))


def bin_z(bench):
    z = pd.read_csv(ZS / f"zscored_{ZSCORED_FILE[bench]}.csv")
    return {(r.topic_label, r.format_label): r.zscore for r in z.itertuples()}


def main():
    benches = sys.argv[1:] if len(sys.argv) > 1 else list(ZSCORED_FILE)
    OUT.mkdir(parents=True, exist_ok=True)
    ds = load_from_disk(DOLMA)
    cache = pd.DataFrame(
        {
            "doc_id": ds["doc_id"],
            "topic": ds["weborganizer_topic"],
            "format": ds["weborganizer_format"],
        }
    )
    print(f"6T cache: {len(cache)} docs, {cache.topic.nunique()} topics, K={K}")

    for bench in benches:
        zdf = pd.read_csv(ZS / f"zscored_{ZSCORED_FILE[bench]}.csv")
        zmap = {(r.topic_label, r.format_label): r.zscore for r in zdf.itertuples()}
        cache["binz"] = [
            zmap.get((t, f), float("-inf")) for t, f in zip(cache.topic, cache.format)
        ]

        def stratified_topk(group, zt, topic):
            """Spread K docs across the topic's highest-z format bins, capped per format so
            the forget-set spans >=~4 formats (mirrors the diversity of per-doc influence
            selection within a topic) instead of collapsing to a single dominant bin."""
            pos = zt[zt.zscore > 0].sort_values("zscore", ascending=False)
            if pos.empty:
                pos = zt.nlargest(min(6, len(zt)), "zscore")
            cap = max(
                1, K // 4
            )  # no single format contributes more than K/4 -> >=4 formats
            picked, remaining = [], K
            for fmt in pos.format_label:
                if remaining <= 0:
                    break
                fg = group[group.format == fmt]
                take = min(cap, remaining, len(fg))
                picked.extend(fg.doc_id.head(take).tolist())
                remaining -= take
            return picked[:K]

        # expC: corpus-wide top-K by bin z (naive top-K is inherently cross-topic-diverse)
        (OUT / f"expC__{bench}.txt").write_text(
            "\n".join(cache.nlargest(K, "binz").doc_id) + "\n"
        )
        # expA: per-topic, stratified across the topic's positive-z formats
        for topic, g in cache.groupby("topic"):
            zt = zdf[zdf.topic_label == topic]
            ids = stratified_topk(g, zt, topic)
            (OUT / f"expA__{topic}__{bench}.txt").write_text("\n".join(ids) + "\n")
        # sanity: which formats does the top-1 topic's forget-set draw from?
        z = pd.read_csv(ZS / f"zscored_{ZSCORED_FILE[bench]}.csv")
        k1 = z.groupby("topic_label").zscore.mean().idxmax()
        sub = cache[cache.topic == k1].nlargest(K, "binz")
        print(
            f"  {bench:22s} top1_topic={k1:<28s} forget-set formats: "
            f"{sub.format.value_counts().head(3).to_dict()}"
        )
    print(f"\nWrote forget-sets to {OUT}")


if __name__ == "__main__":
    main()