social-data-attribution-w2/scripts/analysis/make_enriched_manifest.py

"""Generate manifests for enriched Dolma shards."""

from __future__ import annotations

import argparse
from pathlib import Path
from typing import Iterable


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Generate enriched shard manifest")
    parser.add_argument("--input-dir", type=Path, default=Path("runs/dolma_enriched"))
    parser.add_argument("--pattern", default="*.enriched.jsonl.zst")
    parser.add_argument("--group-size", type=int, default=1)
    parser.add_argument("--output", type=Path, required=True)
    return parser


def _list_paths(input_dir: Path, pattern: str) -> list[Path]:
    return sorted(input_dir.rglob(pattern))


def _chunk(items: list[Path], size: int) -> Iterable[list[Path]]:
    for start in range(0, len(items), size):
        yield items[start : start + size]


def main() -> int:
    args = build_parser().parse_args()
    if args.group_size <= 0:
        raise ValueError("--group-size must be >= 1")

    paths = _list_paths(args.input_dir, args.pattern)
    if not paths:
        raise ValueError("No enriched shards found for manifest")

    lines = [
        " ".join(str(path) for path in chunk)
        for chunk in _chunk(paths, args.group_size)
    ]
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text("\n".join(lines) + "\n", encoding="utf-8")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())