social-data-attribution-w2/src/data_attribution/cli/ingest.py

from __future__ import annotations

import argparse
import datetime as _dt
import logging
from pathlib import Path
from typing import Any, Dict

from dolma.constants import DEFAULT_SPLIT


def _default_run_id(prefix: str = "run") -> str:
    ts = _dt.datetime.now(tz=_dt.timezone.utc)
    return f"{prefix}-{ts.strftime('%Y%m%dT%H%M%SZ')}"


def register_ingest_subcommand(
    subparsers: argparse._SubParsersAction[argparse.ArgumentParser],
) -> None:
    parser = subparsers.add_parser("ingest", help="Ingest Dolma data")
    parser.add_argument("--output-path", type=Path, required=True)
    parser.add_argument("--manifest-path", type=Path, required=True)
    parser.add_argument("--subset")
    parser.add_argument("--split", default=DEFAULT_SPLIT)
    parser.add_argument("--num-documents", type=int)
    parser.add_argument("--token-budget", type=int)
    parser.add_argument("--stratify-by", nargs="+", action="extend")
    parser.add_argument(
        "--output-format", choices=("jsonl", "parquet"), default="jsonl"
    )
    parser.add_argument("--cache-dir", type=Path)
    parser.add_argument("--download-dir", type=Path)
    parser.add_argument("--shuffle-buffer", type=int, default=10_000)
    parser.add_argument("--seed", type=int, default=0)
    parser.set_defaults(func=_ingest_handler)


def _ingest_handler(args: argparse.Namespace, config: Dict[str, Any]) -> int:
    logger = logging.getLogger(args.command)
    run_id = args.run_id or _default_run_id()
    logger.info("Starting %s (run_id=%s)", args.command, run_id)
    if config:
        logger.debug("Loaded config: %s", config)
    if args.num_documents is None and args.token_budget is None:
        raise ValueError(
            "Provide --num-documents or --token-budget to control sampling"
        )
    if args.cache_dir is None and args.download_dir is None:
        raise ValueError(
            "Ingest requires --cache-dir or --download-dir pointing to local shards."
        )
    from dolma import materialize_sample

    materialize_sample(
        args.output_path,
        args.manifest_path,
        subset=args.subset,
        split=args.split,
        cache_dir=args.cache_dir,
        download_dir=args.download_dir,
        use_local_shards=True,
        ensure_download=False,
        allow_full_download=False,
        shuffle_buffer=args.shuffle_buffer,
        seed=args.seed,
        num_documents=args.num_documents,
        token_budget=args.token_budget,
        stratify_by=args.stratify_by,
        output_format=args.output_format,
    )
    logger.info("Wrote sample to %s", args.output_path)
    logger.info("Wrote manifest to %s", args.manifest_path)
    return 0