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

#!/usr/bin/env python3
"""Rebuild per-document ARC-Challenge influence over the full 5.68M working set.

The old faithful per-doc influence parquet was built from only ~216/316 staged
shards (3.89M docs). The score matrices were always complete; this re-aggregates
per-document influence over all 316 shards. Per-doc influence = mean over the
benchmark's query columns (matching bin_aggregate). Doc UUIDs come from the
training shard JSONLs (positionally aligned to npy rows); the scores-dir
doc_ids.json holds only positional "shard:idx" ids, which do not join the manifest.

Emits doc_id, arc_challenge_score, weborganizer_topic for every scored doc.
"""
import argparse
import json
from pathlib import Path

import numpy as np
import pandas as pd


def shard_uuids(shard_dir: Path, name: str) -> list[str]:
    uuids: list[str] = []
    with open(shard_dir / f"{name}.jsonl", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                uuids.append(json.loads(line)["id"])
    return uuids


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--scores-dir", type=Path, required=True)
    ap.add_argument("--shard-dir", type=Path, required=True)
    ap.add_argument("--manifest", type=Path, required=True)
    ap.add_argument("--output", type=Path, required=True)
    args = ap.parse_args()

    m = pd.read_parquet(args.manifest, columns=["doc_id", "bin_topic"]).dropna(
        subset=["bin_topic"]
    )
    topic = dict(zip(m.doc_id, m.bin_topic))

    uuids: list[str] = []
    scores: list[float] = []
    shard_files = sorted(args.scores_dir.glob("shard_*.npy"))
    print(f"shards: {len(shard_files)}", flush=True)
    for i, npy in enumerate(shard_files):
        name = npy.stem
        ids = shard_uuids(args.shard_dir, name)
        s = np.asarray(np.load(npy, mmap_mode="r"))
        if s.shape[0] != len(ids):
            print(f"SKIP {name}: {s.shape[0]} scores vs {len(ids)} ids", flush=True)
            continue
        dm = s.mean(axis=1).astype(np.float64)
        uuids.extend(ids)
        scores.extend(dm.tolist())
        del s
        if (i + 1) % 50 == 0:
            print(f"  {i + 1}/{len(shard_files)} shards", flush=True)

    df = pd.DataFrame({"doc_id": uuids, "arc_challenge_score": scores})
    df["weborganizer_topic"] = df.doc_id.map(topic)
    mapped = int(df.weborganizer_topic.notna().sum())
    args.output.parent.mkdir(parents=True, exist_ok=True)
    df.to_parquet(args.output, index=False)
    print(f"docs={len(df)} topic_mapped={mapped} ({100 * mapped / max(len(df), 1):.1f}%)", flush=True)
    print(f"wrote {args.output}", flush=True)


if __name__ == "__main__":
    main()