social-data-attribution-w2/scripts/manifests/pre_split_stratified_ids.py

"""Pre-split stratified manifest into per-shard files.

Scans stratified_combined.jsonl once and writes one file per shard.
Default mode writes full JSONL rows so downstream extraction can
merge HF text with the original metadata. Legacy mode (--ids-only)
writes plain doc_id lists.

Usage:
    python scripts/manifests/pre_split_stratified_ids.py \
        --manifest stratified_data/stratified_combined.jsonl \
        --output-dir stratified_data/shard_rows
"""

import argparse
import json
import logging
from pathlib import Path

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
)
log = logging.getLogger(__name__)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Pre-split manifest into per-shard files"
    )
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument(
        "--ids-only",
        action="store_true",
        help="Write plain doc_id lists instead of full JSONL rows",
    )
    args = parser.parse_args()

    args.output_dir.mkdir(parents=True, exist_ok=True)

    handles: dict[str, object] = {}
    total = 0
    ext = ".ids.txt" if args.ids_only else ".rows.jsonl"

    log.info(
        "Scanning manifest: %s (mode: %s)",
        args.manifest,
        "ids-only" if args.ids_only else "jsonl",
    )
    try:
        with open(args.manifest) as f:
            for line in f:
                rec = json.loads(line)
                shard = rec["shard_path"]

                if shard not in handles:
                    out_file = args.output_dir / shard.replace(".jsonl.zst", ext)
                    handles[shard] = open(out_file, "w")

                if args.ids_only:
                    handles[shard].write(rec["doc_id"] + "\n")
                else:
                    handles[shard].write(line)
                total += 1

                if total % 5_000_000 == 0:
                    log.info("  %d lines scanned...", total)
    finally:
        for fh in handles.values():
            fh.close()

    log.info("Total lines: %d, unique shards: %d", total, len(handles))


if __name__ == "__main__":
    main()