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

#!/usr/bin/env python3
"""Split the combined ARC forget-set parquet (recipe, doc_id, text) into the
per-recipe [doc_id, text] files the unlearning loader expects, matching the
existing forget_texts_faithful/ convention for the other 4 benchmarks.

Drop-in step for folding ARC-as-target into the faithful multi-seed grid.
Idempotent: overwrites only the 25 arc_challenge recipe files it owns.
"""

from pathlib import Path
import pandas as pd

SRC = Path.home() / "scratch/n16_selectivity/faithful_forget_text_arc.parquet"
OUT = Path.home() / "scratch/n16_selectivity/forget_texts_faithful"

d = pd.read_parquet(SRC)
assert set(d.columns) >= {"recipe", "doc_id", "text"}, d.columns
n = 0
for recipe, g in d.groupby("recipe"):
    dest = OUT / f"{recipe}.parquet"
    g[["doc_id", "text"]].reset_index(drop=True).to_parquet(dest, index=False)
    n += 1
print(f"wrote {n} per-recipe ARC parquets to {OUT}")
print("arc files now present:", len(list(OUT.glob("*arc_challenge.parquet"))))