social-data-attribution-w2/scripts/validation/split_for_verification.py

"""Split socialiqa_verify_input.jsonl into N equal chunks for parallel verification."""

from __future__ import annotations

import json
from pathlib import Path

INPUT = Path("runs/manifests/socialiqa_verify_input.jsonl")
CHUNK_DIR = Path("runs/manifests/socialiqa_verify_chunks")
N_CHUNKS = 100


def main() -> None:
    records = [
        json.loads(line) for line in INPUT.read_text().splitlines() if line.strip()
    ]
    CHUNK_DIR.mkdir(parents=True, exist_ok=True)

    chunk_size = (len(records) + N_CHUNKS - 1) // N_CHUNKS
    for i in range(N_CHUNKS):
        chunk = records[i * chunk_size : (i + 1) * chunk_size]
        chunk_path = CHUNK_DIR / f"chunk_{i:02d}.jsonl"
        with chunk_path.open("w") as fh:
            for r in chunk:
                fh.write(json.dumps(r) + "\n")
        print(f"chunk_{i:02d}.jsonl: {len(chunk)} records")

    print(f"\nTotal: {len(records)} records across {N_CHUNKS} chunks in {CHUNK_DIR}")


if __name__ == "__main__":
    main()