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

from __future__ import annotations

import argparse
import logging
from pathlib import Path
from typing import Sequence

from data_attribution.config import ALLOWED_MODEL_IDS, DEFAULT_MODEL_ID
from dolma.constants import DOLMA_DATASET_ID
from data_attribution.hf_hub import load_env_secret

from .coordinator import precache_assets


def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Prefetch dataset and model assets into the cache."
    )
    parser.add_argument("--cache-root", type=Path, help="Override cache root")
    parser.add_argument(
        "--model-id",
        dest="model_ids",
        nargs="+",
        default=[DEFAULT_MODEL_ID],
        help=f"Model repo id(s) to download (allowed: {', '.join(ALLOWED_MODEL_IDS)})",
    )
    parser.add_argument(
        "--dataset-repo-id",
        default=DOLMA_DATASET_ID,
        help="Dataset repo id to download",
    )
    parser.add_argument("--skip-model", action="store_true", help="Skip model download")
    parser.add_argument(
        "--skip-dataset", action="store_true", help="Skip dataset download"
    )
    parser.add_argument("--subset", help="Optional dataset subset")
    parser.add_argument(
        "--force-download",
        action="store_true",
        help="Force redownload of dataset shards",
    )
    parser.add_argument(
        "--allow-full-download", action="store_true", help="Allow full dataset download"
    )
    parser.add_argument(
        "--snapshot",
        action="store_true",
        help="Use huggingface_hub snapshot_download to fetch dataset files",
    )
    parser.add_argument(
        "--include",
        action="append",
        help="Optional glob pattern(s) to limit snapshot download; repeatable",
    )
    parser.add_argument(
        "--max-workers",
        type=int,
        default=2,
        help="Limit parallel HF downloads to reduce 429 risk",
    )
    parser.add_argument(
        "--retry-attempts",
        type=int,
        default=6,
        help="Retry attempts for HF 429 responses",
    )
    return parser.parse_args(argv)


def main(argv: Sequence[str] | None = None) -> int:
    args = _parse_args(argv)
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s"
    )
    logger = logging.getLogger("precache_hf_assets")

    load_env_secret()
    precache_assets(
        logger=logger,
        cache_root=args.cache_root,
        model_ids=args.model_ids,
        dataset_repo_id=args.dataset_repo_id,
        subset=args.subset,
        skip_model=args.skip_model,
        skip_dataset=args.skip_dataset,
        force_download=args.force_download,
        allow_full_download=args.allow_full_download,
        snapshot=args.snapshot,
        include=args.include,
        max_workers=args.max_workers,
        retry_attempts=args.retry_attempts,
    )
    return 0