social-data-attribution-w2/src/dolma/pool_sample/cli.py

"""CLI entry point for pool sampling."""

from __future__ import annotations

import argparse
import logging
import os
from pathlib import Path

from dolma.pool_sample.recipe import (
    DEFAULT_SEED,
    DEFAULT_SHARD_FRACTION,
    DEFAULT_TOKEN_BUDGET,
)


def main(argv: list[str] | None = None) -> None:
    parser = argparse.ArgumentParser(
        description="Draw a uniform random sample from dolma3_pool."
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        required=True,
        help="Directory for output shards and manifest",
    )
    parser.add_argument(
        "--token-budget",
        type=int,
        default=DEFAULT_TOKEN_BUDGET,
        help="Target token count (default: 150B)",
    )
    parser.add_argument("--seed", type=int, default=DEFAULT_SEED, help="Random seed")
    parser.add_argument(
        "--shard-fraction",
        type=float,
        default=DEFAULT_SHARD_FRACTION,
        help="Fraction of pool shards to download (default: 0.05)",
    )
    parser.add_argument(
        "--num-workers",
        type=int,
        default=None,
        help="Number of processing workers (default: cpu_count)",
    )
    parser.add_argument(
        "--download-dir",
        type=Path,
        default=None,
        help="Directory for downloaded shards (default: TMPDIR or output-dir/download)",
    )
    parser.add_argument(
        "--download-workers",
        type=int,
        default=4,
        help="Number of parallel download workers",
    )
    parser.add_argument(
        "--lines-per-shard",
        type=int,
        default=10_000_000,
        help="Lines per output shard (default: 10M)",
    )
    parser.add_argument(
        "--log-every",
        type=int,
        default=100_000,
        help="Log progress every N documents per worker",
    )
    parser.add_argument(
        "--shard-paths-file",
        type=Path,
        default=None,
        help="File with one shard path per line (skips HF API listing)",
    )
    parser.add_argument("--verbose", action="store_true", default=False)
    args = parser.parse_args(argv)

    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s %(levelname)s [%(name)s] %(message)s",
    )

    download_dir = args.download_dir
    if download_dir is None:
        tmpdir = os.environ.get("TMPDIR")
        if tmpdir:
            download_dir = Path(tmpdir) / "pool_download"

    shard_paths = None
    if args.shard_paths_file is not None:
        shard_paths = [
            line.strip()
            for line in args.shard_paths_file.read_text().splitlines()
            if line.strip()
        ]

    from dolma.pool_sample.recipe import sample_pool

    manifest_path = sample_pool(
        output_dir=args.output_dir,
        token_budget=args.token_budget,
        seed=args.seed,
        num_workers=args.num_workers,
        shard_fraction=args.shard_fraction,
        download_dir=download_dir,
        download_workers=args.download_workers,
        lines_per_shard=args.lines_per_shard,
        log_every=args.log_every,
        shard_paths=shard_paths,
    )

    logging.getLogger(__name__).info("Manifest written to %s", manifest_path)


if __name__ == "__main__":
    main()