fn_registry/python/functions/ml/hf_snapshot_download.py

"""Wrapper de huggingface_hub.snapshot_download con manejo de ImportError descriptivo."""

from __future__ import annotations


def hf_snapshot_download(
    repo_id: str,
    allow_patterns: list[str] | None = None,
    ignore_patterns: list[str] | None = None,
    local_dir: str | None = None,
    token: str | None = None,
) -> str:
    """Descarga un snapshot completo (o filtrado) de un repo de HuggingFace Hub.

    Wrapper sobre `huggingface_hub.snapshot_download` con manejo de ImportError
    descriptivo. Si `local_dir` se especifica, el snapshot se descarga alli en
    lugar del cache global de HuggingFace (~/.cache/huggingface/).

    Args:
        repo_id: identificador del repositorio en HuggingFace Hub
                 (ej: "runwayml/stable-diffusion-v1-5", "meta-llama/Llama-2-7b-hf").
        allow_patterns: lista opcional de patrones glob para incluir solo ciertos
                        archivos (ej: ["*.safetensors", "*.json"]).
                        Si None, se descargan todos los archivos.
        ignore_patterns: lista opcional de patrones glob para excluir archivos
                         (ej: ["*.bin", "flax_*"]). Util para evitar descargar
                         pesos en formato pytorch si ya se tienen en safetensors.
        local_dir: directorio local donde guardar el snapshot. Si None, usa
                   el cache global de HuggingFace Hub.
        token: token de acceso a HuggingFace (para repos privados o con gated
               access como Llama). Si None, usa HF_TOKEN del entorno.

    Returns:
        Path local (str) donde quedo almacenado el snapshot.

    Raises:
        ImportError: si huggingface_hub no esta instalado, con sugerencia de instalacion.
        Exception: cualquier error de red o autenticacion propagado desde snapshot_download.
    """
    try:
        from huggingface_hub import snapshot_download
    except ImportError as exc:
        raise ImportError(
            "huggingface_hub no esta instalado. "
            "Instalar con: pip install huggingface_hub"
        ) from exc

    kwargs: dict = {"repo_id": repo_id}
    if allow_patterns is not None:
        kwargs["allow_patterns"] = allow_patterns
    if ignore_patterns is not None:
        kwargs["ignore_patterns"] = ignore_patterns
    if local_dir is not None:
        kwargs["local_dir"] = local_dir
    if token is not None:
        kwargs["token"] = token

    result = snapshot_download(**kwargs)
    return str(result)