social-data-attribution-w2/scripts/analysis/multiseed/check_cache_doc_ids.py

"""Compare doc_id formats of the two 6T-filtered Arrow caches and test
membership of sample bin-level rescore forget ids in each."""
from pathlib import Path

from datasets import load_from_disk

CACHES = {
    "seom35": "/storage/ice-shared/cs7634/seom35/hf_cache/datasets/dolma3_6t_filtered",
    "arc_rescore": "/storage/ice-shared/cs7634/staff/TDA/arc_rescore/dolma_cache/dolma3_6t_filtered",
}
SAMPLE_TXT = Path.home() / (
    "scratch/n16_selectivity/forgetsets_binlevel_arc_rescore_20260608/expC__arc_challenge.txt"
)

sample = [ln.strip() for ln in SAMPLE_TXT.read_text().splitlines() if ln.strip()]
print(f"sample forget ids: {len(sample)} e.g. {sample[:2]}")

for name, path in CACHES.items():
    ds = load_from_disk(path)
    head = ds.select(range(3))["doc_id"]
    ids = set(ds["doc_id"])
    hit = sum(1 for s in sample if s in ids)
    print(f"{name}: rows={len(ds)} cols={ds.column_names}")
    print(f"  head doc_ids: {head}")
    print(f"  forget-id membership: {hit}/{len(sample)}")