social-data-attribution-w2/scripts/attribution/consolidate_preconditioner_sample.py

"""Consolidate .jsonl.zst files into a single plain JSONL for bergson.

Reads all .jsonl and .jsonl.zst files from an input directory, strips
metadata to keep only id and text fields, and writes one consolidated
JSONL file. Intended for preparing the SOC-151 100K preconditioner
sample for bergson preconditioner building.

Usage:
    python scripts/attribution/consolidate_preconditioner_sample.py \
        --input-dir ~/scratch/soc152/preconditioner_100k_raw \
        --output ~/scratch/soc152/data/preconditioner_100k.jsonl
"""

import argparse
import logging
from pathlib import Path

from data_attribution.attribution.trackstar.sharding import _open_lines, _strip_fields

log = logging.getLogger(__name__)

JSONL_GLOBS = ("*.jsonl", "*.jsonl.zst")


def find_jsonl_files(input_dir: Path) -> list[Path]:
    files: list[Path] = []
    for pattern in JSONL_GLOBS:
        files.extend(input_dir.glob(pattern))
    if not files:
        for pattern in JSONL_GLOBS:
            files.extend(input_dir.glob(f"**/{pattern}"))
    if not files:
        log.error("No .jsonl or .jsonl.zst files found in %s", input_dir)
        raise SystemExit(1)
    return sorted(set(files))


def consolidate(input_dir: Path, output: Path) -> int:
    sources = find_jsonl_files(input_dir)
    log.info("Found %d source files in %s", len(sources), input_dir)

    output.parent.mkdir(parents=True, exist_ok=True)
    count = 0

    with open(output, "w", encoding="utf-8") as out:
        for i, source in enumerate(sources):
            for line in _open_lines(source):
                if not line.strip():
                    continue
                out.write(_strip_fields(line))
                count += 1
            if (i + 1) % 5000 == 0:
                log.info(
                    "Processed %d / %d files (%d docs)", i + 1, len(sources), count
                )

    log.info("Wrote %d docs to %s", count, output)
    return count


def main() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
    )
    parser = argparse.ArgumentParser(
        description="Consolidate .jsonl.zst files into single JSONL"
    )
    parser.add_argument(
        "--input-dir",
        type=Path,
        required=True,
        help="Directory containing .jsonl.zst files",
    )
    parser.add_argument(
        "--output",
        type=Path,
        required=True,
        help="Output path for consolidated JSONL",
    )
    args = parser.parse_args()
    consolidate(args.input_dir, args.output)


if __name__ == "__main__":
    main()