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

from __future__ import annotations

import logging
from typing import Sequence

from data_attribution.config import configure_hf_cache

from .download import download_dataset, download_model
from .policy import (
    normalize_model_ids,
    validate_dataset_id,
    validate_subset,
)


def precache_assets(
    *,
    logger: logging.Logger,
    cache_root,
    model_ids: Sequence[str] | None,
    dataset_repo_id: str,
    subset: str | None,
    skip_model: bool,
    skip_dataset: bool,
    force_download: bool,
    allow_full_download: bool,
    snapshot: bool,
    include: list[str] | None,
    max_workers: int,
    retry_attempts: int,
) -> None:
    model_ids = normalize_model_ids(model_ids)
    dataset_repo_id = validate_dataset_id(dataset_repo_id)
    subset = validate_subset(subset)
    paths = configure_hf_cache(cache_root, create_dirs=True)
    logger.info("Using cache root %s", paths.root)

    if not skip_model:
        for model_id in model_ids:
            logger.info("Downloading model snapshot: %s", model_id)
            download_model(
                logger=logger,
                model_id=model_id,
                max_workers=max_workers,
                retry_attempts=retry_attempts,
            )
        logger.info("Model snapshot complete")

    if not skip_dataset:
        logger.info("Downloading dataset shards: %s", dataset_repo_id)
        download_dataset(
            logger=logger,
            subset=subset,
            repo_id=dataset_repo_id,
            cache_dir=paths.datasets,
            force_download=force_download,
            allow_full_download=allow_full_download,
            snapshot=snapshot,
            include=include,
            max_workers=max_workers,
            retry_attempts=retry_attempts,
        )
        logger.info("Dataset download complete")


__all__ = ["precache_assets"]