hf_fetch_model/

lib.rs

// SPDX-License-Identifier: MIT OR Apache-2.0

//! # hf-fetch-model
//!
//! Fast `HuggingFace` model downloads for Rust.
//!
//! An embeddable library for downloading `HuggingFace` model repositories
//! with maximum throughput. Wraps [`hf_hub`] and adds repo-level orchestration.
//!
//! ## Quick Start
//!
//! ```rust,no_run
//! # async fn example() -> Result<(), hf_fetch_model::FetchError> {
//! let outcome = hf_fetch_model::download("julien-c/dummy-unknown".to_owned()).await?;
//! println!("Model at: {}", outcome.inner().display());
//! # Ok(())
//! # }
//! ```
//!
//! ## Configured Download
//!
//! ```rust,no_run
//! # async fn example() -> Result<(), hf_fetch_model::FetchError> {
//! use hf_fetch_model::FetchConfig;
//!
//! let config = FetchConfig::builder()
//!     .filter("*.safetensors")
//!     .filter("*.json")
//!     .on_progress(|e| {
//!         println!("{}: {:.1}%", e.filename, e.percent);
//!     })
//!     .build()?;
//!
//! let outcome = hf_fetch_model::download_with_config(
//!     "google/gemma-2-2b".to_owned(),
//!     &config,
//! ).await?;
//! // outcome.is_cached() tells you if it came from local cache
//! let path = outcome.into_inner();
//! # Ok(())
//! # }
//! ```
//!
//! ## Inspect Before Downloading
//!
//! Read tensor metadata from `.safetensors` headers via HTTP Range requests —
//! no weight data downloaded. Sharded repos (those with
//! `model.safetensors.index.json`) work transparently —
//! [`inspect::inspect_repo_safetensors`] reads every shard's header in parallel
//! and returns a flat per-file result list. See
//! [`examples/candle_inspect.rs`](https://github.com/PCfVW/hf-fetch-model/blob/main/examples/candle_inspect.rs)
//! for a runnable example, or the
//! [Inspect tutorial](https://github.com/PCfVW/hf-fetch-model/blob/main/docs/tutorials/inspect-before-downloading.md)
//! for a narrative walkthrough.
//!
//! ```rust,no_run
//! # async fn example() -> Result<(), hf_fetch_model::FetchError> {
//! let results = hf_fetch_model::inspect::inspect_repo_safetensors(
//!     "EleutherAI/pythia-1.4b", None, None,
//! ).await?;
//!
//! for (filename, header, _source) in &results {
//!     println!("{filename}: {} tensors", header.tensors.len());
//! }
//! # Ok(())
//! # }
//! ```
//!
//! The CLI also exposes `hf-fm inspect <repo> [FILE] --check-gpu [N]` (v0.10.1)
//! to print a one-line GPU-fit verdict against device `N` (default 0) using
//! the `hypomnesis` crate (NVML on Linux/Windows, DXGI on Windows). Adding
//! `--context N` (v0.10.4) folds in the KV cache at a context length and
//! reports a real fit against `weights + KV` instead of weights alone — the
//! difference between "fits" and "out-of-memory at token 8000" on a consumer
//! card. The architecture parameters come from the model's `config.json`,
//! parsed by the library API that v0.10.4 exposes for downstream reuse:
//! [`inspect::ModelConfig`] plus [`inspect::fetch_model_config`] /
//! [`inspect::fetch_model_config_cached`] (cache-first or cache-only) and the
//! [`inspect::torch_dtype_bytes`] helper. The KV math itself — `GQA`,
//! sliding-window, `MLA`-skip, and hybrid Mamba/attention layer counting —
//! and the verdict rendering stay binary-only; depend on `hypomnesis`
//! directly for the raw device-info numbers.
//!
//! ## Cached-file Inspection
//!
//! Beyond the remote-or-cached `.safetensors` path above,
//! [`inspect::inspect_gguf_cached`] (v0.10.2),
//! [`inspect::inspect_npz_cached`], and [`inspect::inspect_pth_cached`]
//! (both v0.10.3) extend inspect to `GGUF` / `NumPy` `.npz` / `PyTorch`
//! `.pth` files in the local cache via the `anamnesis` parser crate. All
//! four formats return the same format-agnostic
//! [`inspect::SafetensorsHeaderInfo`] shape, so downstream pipeline steps
//! (filter, tree, dtypes aggregation) work uniformly across formats.
//!
//! For cached `.safetensors` files, v0.10.3 also surfaces quantization
//! detection. When [`inspect::inspect_safetensors_local`] sees a quantized
//! header (`FP8` variants, `GPTQ`, `AWQ`, `BnB-NF4`, `BnB-INT8`), it
//! populates the new [`inspect::QuantInfo`] field with the scheme name and
//! both stored + dequantised byte sizes. Unquantized safetensors and
//! non-safetensors formats leave `quant_info` as `None`.
//!
//! ```rust,no_run
//! # fn example() -> Result<(), hf_fetch_model::FetchError> {
//! use hf_fetch_model::inspect;
//! use std::path::Path;
//!
//! let header = inspect::inspect_safetensors_local(
//!     Path::new("/path/to/cached/file.safetensors"),
//! )?;
//! if let Some(q) = &header.quant_info {
//!     println!(
//!         "Quantized as {}: {} stored -> {} dequantised",
//!         q.scheme, q.stored_bytes, q.dequantized_bytes,
//!     );
//! }
//! # Ok(())
//! # }
//! ```
//!
//! Remote inspect for `GGUF` / `NPZ` / `PTH` (via HTTP Range, without
//! going through the cache) is planned for v0.11; until then those formats
//! error early with a "pass --cached after downloading" recovery hint.
//!
//! ## `HuggingFace` Cache
//!
//! Downloaded files are stored in the standard `HuggingFace` cache directory
//! (`~/.cache/huggingface/hub/`), ensuring compatibility with Python tooling.
//!
//! ## Cache Management
//!
//! v0.10.0 adds library APIs for inspecting, verifying, and pruning the local
//! cache. [`cache::cache_summary`] enumerates every cached repo with size and
//! file counts; [`cache::repo_status`] gives a per-file `Complete` / `Partial` /
//! `Missing` breakdown for one repo; [`cache::verify_cache`] re-checks `SHA256`
//! digests of cached files against `HuggingFace` LFS metadata; and
//! [`cache::find_partial_files`] locates `.chunked.part` orphans from
//! interrupted downloads.
//!
//! For long verifications (multi-GiB safetensors files), drive
//! [`cache::verify_cache_with_progress`] with an [`Fn`] callback that receives
//! [`cache::VerifyEvent`]s so a CLI or GUI can render a spinner or progress
//! bar without polling.
//!
//! ```rust,no_run
//! # async fn example() -> Result<(), hf_fetch_model::FetchError> {
//! use hf_fetch_model::cache::{self, VerifyStatus};
//!
//! let results = cache::verify_cache("google/gemma-2-2b-it", None, None).await?;
//! let ok = results
//!     .iter()
//!     .filter(|r| matches!(r.status, VerifyStatus::Ok))
//!     .count();
//! let mismatch = results
//!     .iter()
//!     .filter(|r| matches!(r.status, VerifyStatus::Mismatch { .. }))
//!     .count();
//! println!("{}/{} files verified, {} mismatches", ok, results.len(), mismatch);
//! # Ok(())
//! # }
//! ```
//!
//! ## Download Durability
//!
//! Multi-connection downloads survive interruption. When a download is
//! aborted by [`FetchConfigBuilder::timeout_per_file`] (default 300 s),
//! Ctrl-C, panic, or a transient chunk error, the partial `.chunked.part`
//! file plus a small per-chunk progress sidecar are kept on disk. The next
//! call to [`download_with_config`] for the same file picks up where it
//! stopped — each parallel chunk sends a fresh `Range` request that skips
//! the bytes it already has — provided the upstream etag still matches.
//! On etag change, schema-version mismatch, or a different
//! [`FetchConfigBuilder::connections_per_file`] count, the partial is
//! discarded and a fresh download starts.
//!
//! For slow connections on multi-GiB files, raise the per-file budget to
//! match real throughput:
//!
//! ```rust,no_run
//! # async fn example() -> Result<(), hf_fetch_model::FetchError> {
//! use std::time::Duration;
//! use hf_fetch_model::FetchConfig;
//!
//! let config = FetchConfig::builder()
//!     .timeout_per_file(Duration::from_secs(1800))
//!     .build()?;
//! # let _ = hf_fetch_model::download_with_config(
//! #     "google/gemma-4-E2B-it".to_owned(),
//! #     &config,
//! # ).await?;
//! # Ok(())
//! # }
//! ```
//!
//! ## Authentication
//!
//! Set the `HF_TOKEN` environment variable to access private or gated models,
//! or use [`FetchConfig::builder().token()`](FetchConfigBuilder::token).

pub mod cache;
pub mod cache_layout;
pub mod checksum;
mod chunked;
mod chunked_state;
pub mod config;
pub mod discover;
pub mod download;
pub mod error;
pub mod inspect;
pub mod plan;
pub mod progress;
pub mod repo;
mod retry;

pub use chunked::build_client;
pub use config::{
    compile_glob_patterns, file_matches, has_glob_chars, FetchConfig, FetchConfigBuilder, Filter,
};
pub use discover::{DiscoveredFamily, GateStatus, ModelCardMetadata, SearchResult};
pub use download::DownloadOutcome;
pub use error::{FetchError, FileFailure};
pub use inspect::{AdapterConfig, ModelConfig};
pub use plan::{download_plan, DownloadPlan, FilePlan};
pub use progress::{ProgressEvent, ProgressReceiver};

use std::collections::HashMap;
use std::path::PathBuf;

use hf_hub::{Repo, RepoType};

/// Pre-flight check for gated model access.
///
/// Two cases:
/// - **No token**: checks the model metadata (unauthenticated) for gating
///   status and rejects with a clear message if gated.
/// - **Token present**: if the model is gated, makes one authenticated
///   metadata request to verify the token actually grants access. Catches
///   invalid tokens and unaccepted licenses before the download starts.
///
/// If the metadata request itself fails (network error, private repo),
/// the check is silently skipped so that normal download error handling
/// can take over.
async fn preflight_gated_check(repo_id: &str, config: &FetchConfig) -> Result<(), FetchError> {
    // Best-effort: if the metadata call fails, let the download proceed.
    let Ok(metadata) = discover::fetch_model_card(repo_id).await else {
        return Ok(());
    };

    if !metadata.gated.is_gated() {
        return Ok(());
    }

    // Model is gated — check auth.
    if config.token.is_none() {
        return Err(FetchError::Auth {
            reason: format!(
                "{repo_id} is a gated model — accept the license at \
                 https://huggingface.co/{repo_id} and set HF_TOKEN or pass --token"
            ),
        });
    }

    // Token is present — verify it grants access with a lightweight probe.
    let probe_client = chunked::build_client(config.token.as_deref())?;
    let probe = repo::list_repo_files_with_metadata(
        repo_id,
        config.token.as_deref(),
        config.revision.as_deref(),
        &probe_client,
    )
    .await;

    if let Err(ref e) = probe {
        // BORROW: explicit .to_string() for error Display formatting
        let msg = e.to_string();
        if msg.contains("401") || msg.contains("403") {
            return Err(FetchError::Auth {
                reason: format!(
                    "{repo_id} is a gated model and your token was rejected — \
                     accept the license at https://huggingface.co/{repo_id} \
                     and check that your token is valid"
                ),
            });
        }
    }

    Ok(())
}

/// Downloads all files from a `HuggingFace` model repository.
///
/// Uses high-throughput mode for maximum download speed, including
/// auto-tuned concurrency, chunked multi-connection downloads for large
/// files, and plan-optimized settings based on file size distribution.
/// Files are stored in the standard `HuggingFace` cache layout
/// (`~/.cache/huggingface/hub/`).
///
/// Authentication is handled via the `HF_TOKEN` environment variable when set.
///
/// For filtering, progress, and other options, use [`download_with_config()`].
///
/// # Arguments
///
/// * `repo_id` — The repository identifier (e.g., `"google/gemma-2-2b-it"`).
///
/// # Returns
///
/// The path to the snapshot directory containing all downloaded files.
///
/// # Errors
///
/// * [`FetchError::Auth`] — if the repository is gated and access is denied (no token, invalid token, or license not accepted).
/// * [`FetchError::Api`] — if the `HuggingFace` API or download fails (includes auth failures).
/// * [`FetchError::RepoNotFound`] — if the repository does not exist.
/// * [`FetchError::InvalidPattern`] — if the default config fails to build (should not happen).
pub async fn download(repo_id: String) -> Result<DownloadOutcome<PathBuf>, FetchError> {
    let config = FetchConfig::builder().build()?;
    download_with_config(repo_id, &config).await
}

/// Downloads files from a `HuggingFace` model repository using the given configuration.
///
/// Supports filtering, progress reporting, custom revision, authentication,
/// and concurrency settings via [`FetchConfig`].
///
/// # Arguments
///
/// * `repo_id` — The repository identifier (e.g., `"google/gemma-2-2b-it"`).
/// * `config` — Download configuration (see [`FetchConfig::builder()`]).
///
/// # Returns
///
/// The path to the snapshot directory containing all downloaded files.
///
/// # Errors
///
/// * [`FetchError::Auth`] — if the repository is gated and access is denied (no token, invalid token, or license not accepted).
/// * [`FetchError::Api`] — if the `HuggingFace` API or download fails (includes auth failures).
/// * [`FetchError::RepoNotFound`] — if the repository does not exist.
pub async fn download_with_config(
    repo_id: String,
    config: &FetchConfig,
) -> Result<DownloadOutcome<PathBuf>, FetchError> {
    // BORROW: explicit .as_str() instead of Deref coercion
    preflight_gated_check(repo_id.as_str(), config).await?;

    let mut builder = hf_hub::api::tokio::ApiBuilder::new().high();

    if let Some(ref token) = config.token {
        // BORROW: explicit .clone() to pass owned String
        builder = builder.with_token(Some(token.clone()));
    }

    if let Some(ref dir) = config.output_dir {
        // BORROW: explicit .clone() for owned PathBuf
        builder = builder.with_cache_dir(dir.clone());
    }

    let api = builder.build().map_err(FetchError::Api)?;

    let hf_repo = match config.revision {
        Some(ref rev) => {
            // BORROW: explicit .clone() for owned String arguments
            Repo::with_revision(repo_id.clone(), RepoType::Model, rev.clone())
        }
        None => Repo::new(repo_id.clone(), RepoType::Model),
    };

    let repo = api.repo(hf_repo);
    download::download_all_files(repo, repo_id, Some(config)).await
}

/// Blocking version of [`download()`] for non-async callers.
///
/// Creates a Tokio runtime internally. Do not call from within
/// an existing async context (use [`download()`] instead).
///
/// # Errors
///
/// Same as [`download()`].
pub fn download_blocking(repo_id: String) -> Result<DownloadOutcome<PathBuf>, FetchError> {
    let rt = tokio::runtime::Runtime::new().map_err(|e| FetchError::Io {
        path: PathBuf::from("<runtime>"),
        source: e,
    })?;
    rt.block_on(download(repo_id))
}

/// Blocking version of [`download_with_config()`] for non-async callers.
///
/// Creates a Tokio runtime internally. Do not call from within
/// an existing async context (use [`download_with_config()`] instead).
///
/// # Errors
///
/// Same as [`download_with_config()`].
pub fn download_with_config_blocking(
    repo_id: String,
    config: &FetchConfig,
) -> Result<DownloadOutcome<PathBuf>, FetchError> {
    let rt = tokio::runtime::Runtime::new().map_err(|e| FetchError::Io {
        path: PathBuf::from("<runtime>"),
        source: e,
    })?;
    rt.block_on(download_with_config(repo_id, config))
}

/// Downloads all files from a `HuggingFace` model repository and returns
/// a filename → path map.
///
/// Each key is the relative filename within the repository (e.g.,
/// `"config.json"`, `"model.safetensors"`), and each value is the
/// absolute local path to the downloaded file.
///
/// Uses the same high-throughput defaults as [`download()`]: auto-tuned
/// concurrency and chunked multi-connection downloads for large files.
///
/// For filtering, progress, and other options, use
/// [`download_files_with_config()`].
///
/// # Arguments
///
/// * `repo_id` — The repository identifier (e.g., `"google/gemma-2-2b-it"`).
///
/// # Errors
///
/// * [`FetchError::Api`] — if the `HuggingFace` API or download fails (includes auth failures).
/// * [`FetchError::RepoNotFound`] — if the repository does not exist.
/// * [`FetchError::InvalidPattern`] — if the default config fails to build (should not happen).
pub async fn download_files(
    repo_id: String,
) -> Result<DownloadOutcome<HashMap<String, PathBuf>>, FetchError> {
    let config = FetchConfig::builder().build()?;
    download_files_with_config(repo_id, &config).await
}

/// Downloads files from a `HuggingFace` model repository using the given
/// configuration and returns a filename → path map.
///
/// Each key is the relative filename within the repository (e.g.,
/// `"config.json"`, `"model.safetensors"`), and each value is the
/// absolute local path to the downloaded file.
///
/// # Arguments
///
/// * `repo_id` — The repository identifier (e.g., `"google/gemma-2-2b-it"`).
/// * `config` — Download configuration (see [`FetchConfig::builder()`]).
///
/// # Errors
///
/// * [`FetchError::Auth`] — if the repository is gated and access is denied (no token, invalid token, or license not accepted).
/// * [`FetchError::Api`] — if the `HuggingFace` API or download fails (includes auth failures).
/// * [`FetchError::RepoNotFound`] — if the repository does not exist.
pub async fn download_files_with_config(
    repo_id: String,
    config: &FetchConfig,
) -> Result<DownloadOutcome<HashMap<String, PathBuf>>, FetchError> {
    // BORROW: explicit .as_str() instead of Deref coercion
    preflight_gated_check(repo_id.as_str(), config).await?;

    let mut builder = hf_hub::api::tokio::ApiBuilder::new().high();

    if let Some(ref token) = config.token {
        // BORROW: explicit .clone() to pass owned String
        builder = builder.with_token(Some(token.clone()));
    }

    if let Some(ref dir) = config.output_dir {
        // BORROW: explicit .clone() for owned PathBuf
        builder = builder.with_cache_dir(dir.clone());
    }

    let api = builder.build().map_err(FetchError::Api)?;

    let hf_repo = match config.revision {
        Some(ref rev) => {
            // BORROW: explicit .clone() for owned String arguments
            Repo::with_revision(repo_id.clone(), RepoType::Model, rev.clone())
        }
        None => Repo::new(repo_id.clone(), RepoType::Model),
    };

    let repo = api.repo(hf_repo);
    download::download_all_files_map(repo, repo_id, Some(config)).await
}

/// Blocking version of [`download_files()`] for non-async callers.
///
/// Creates a Tokio runtime internally. Do not call from within
/// an existing async context (use [`download_files()`] instead).
///
/// # Errors
///
/// Same as [`download_files()`].
pub fn download_files_blocking(
    repo_id: String,
) -> Result<DownloadOutcome<HashMap<String, PathBuf>>, FetchError> {
    let rt = tokio::runtime::Runtime::new().map_err(|e| FetchError::Io {
        path: PathBuf::from("<runtime>"),
        source: e,
    })?;
    rt.block_on(download_files(repo_id))
}

/// Downloads a single file from a `HuggingFace` model repository.
///
/// Returns the local cache path. If the file is already cached (and
/// checksums match when `verify_checksums` is enabled), the download
/// is skipped and the cached path is returned immediately.
///
/// Files at or above [`FetchConfig`]'s `chunk_threshold` (auto-tuned by
/// the download plan optimizer, or 100 MiB fallback) are downloaded using
/// multiple parallel HTTP Range connections (`connections_per_file`,
/// auto-tuned or 8 fallback). Smaller files use a single connection.
///
/// # Arguments
///
/// * `repo_id` — Repository identifier (e.g., `"mntss/clt-gemma-2-2b-426k"`).
/// * `filename` — Exact filename within the repository (e.g., `"W_enc_5.safetensors"`).
/// * `config` — Shared configuration for auth, progress, checksums, retries, and chunking.
///
/// # Errors
///
/// * [`FetchError::Auth`] — if the repository is gated and access is denied (no token, invalid token, or license not accepted).
/// * [`FetchError::Http`] — if the file does not exist in the repository.
/// * [`FetchError::Api`] — on download failure (after retries).
/// * [`FetchError::Checksum`] — if verification is enabled and fails.
pub async fn download_file(
    repo_id: String,
    filename: &str,
    config: &FetchConfig,
) -> Result<DownloadOutcome<PathBuf>, FetchError> {
    // BORROW: explicit .as_str() instead of Deref coercion
    preflight_gated_check(repo_id.as_str(), config).await?;

    let mut builder = hf_hub::api::tokio::ApiBuilder::new().high();

    if let Some(ref token) = config.token {
        // BORROW: explicit .clone() to pass owned String
        builder = builder.with_token(Some(token.clone()));
    }

    if let Some(ref dir) = config.output_dir {
        // BORROW: explicit .clone() for owned PathBuf
        builder = builder.with_cache_dir(dir.clone());
    }

    let api = builder.build().map_err(FetchError::Api)?;

    let hf_repo = match config.revision {
        Some(ref rev) => {
            // BORROW: explicit .clone() for owned String arguments
            Repo::with_revision(repo_id.clone(), RepoType::Model, rev.clone())
        }
        None => Repo::new(repo_id.clone(), RepoType::Model),
    };

    let repo = api.repo(hf_repo);
    download::download_file_by_name(repo, repo_id, filename, config).await
}

/// Blocking version of [`download_file()`] for non-async callers.
///
/// Creates a Tokio runtime internally. Do not call from within
/// an existing async context (use [`download_file()`] instead).
///
/// # Errors
///
/// Same as [`download_file()`].
pub fn download_file_blocking(
    repo_id: String,
    filename: &str,
    config: &FetchConfig,
) -> Result<DownloadOutcome<PathBuf>, FetchError> {
    let rt = tokio::runtime::Runtime::new().map_err(|e| FetchError::Io {
        path: PathBuf::from("<runtime>"),
        source: e,
    })?;
    rt.block_on(download_file(repo_id, filename, config))
}

/// Blocking version of [`download_files_with_config()`] for non-async callers.
///
/// Creates a Tokio runtime internally. Do not call from within
/// an existing async context (use [`download_files_with_config()`] instead).
///
/// # Errors
///
/// Same as [`download_files_with_config()`].
pub fn download_files_with_config_blocking(
    repo_id: String,
    config: &FetchConfig,
) -> Result<DownloadOutcome<HashMap<String, PathBuf>>, FetchError> {
    let rt = tokio::runtime::Runtime::new().map_err(|e| FetchError::Io {
        path: PathBuf::from("<runtime>"),
        source: e,
    })?;
    rt.block_on(download_files_with_config(repo_id, config))
}

/// Downloads files according to an existing [`DownloadPlan`].
///
/// Only uncached files in the plan are downloaded. The `config` controls
/// authentication, progress, timeouts, and performance settings.
/// Use [`DownloadPlan::recommended_config()`] to compute an optimized config,
/// or override specific fields via [`DownloadPlan::recommended_config_builder()`].
///
/// # Errors
///
/// Returns [`FetchError::Io`] if the cache directory cannot be resolved.
/// Same error conditions as [`download_with_config()`] for the download itself.
pub async fn download_with_plan(
    plan: &DownloadPlan,
    config: &FetchConfig,
) -> Result<DownloadOutcome<PathBuf>, FetchError> {
    if plan.fully_cached() {
        // Resolve snapshot path from cache and return immediately.
        let cache_dir = config
            .output_dir
            .clone()
            .map_or_else(cache::hf_cache_dir, Ok)?;
        let repo_dir = cache_layout::repo_dir(&cache_dir, plan.repo_id.as_str());
        let snapshot_dir = cache_layout::snapshot_dir(&repo_dir, plan.revision.as_str());
        return Ok(DownloadOutcome::Cached(snapshot_dir));
    }

    // Delegate to the standard download path which will re-check cache
    // internally. The plan's value is the dry-run preview and the
    // recommended config computed by the caller.
    // BORROW: explicit .clone() for owned String argument
    download_with_config(plan.repo_id.clone(), config).await
}

/// Blocking version of [`download_with_plan()`] for non-async callers.
///
/// Creates a Tokio runtime internally. Do not call from within
/// an existing async context (use [`download_with_plan()`] instead).
///
/// # Errors
///
/// Same as [`download_with_plan()`].
pub fn download_with_plan_blocking(
    plan: &DownloadPlan,
    config: &FetchConfig,
) -> Result<DownloadOutcome<PathBuf>, FetchError> {
    let rt = tokio::runtime::Runtime::new().map_err(|e| FetchError::Io {
        path: PathBuf::from("<runtime>"),
        source: e,
    })?;
    rt.block_on(download_with_plan(plan, config))
}