zsfm_hub/
discover.rs

1//! Format-agnostic HuggingFace Hub download: list whatever files a repo
2//! actually has, pick a checkpoint format (auto or requested), and fetch it —
3//! without assuming `config.json` + `model.safetensors` like [`crate::download_model`].
4//!
5//! Feeds straight into `zsfm-checkpoint::load_checkpoint`, which reads every
6//! format this module can select.
7
8use std::collections::HashMap;
9use std::path::{Path, PathBuf};
10
11use anyhow::Context;
12use futures_util::{stream, StreamExt};
13use serde::Deserialize;
14
15use crate::download::{fetch_file, SHARD_DOWNLOAD_CONCURRENCY};
16use crate::http::{build_client, HF_BASE};
17
18/// Checkpoint file extensions this tool knows how to read, in the order
19/// preferred when a repo ships more than one format and the caller doesn't
20/// request a specific one.
21pub const FORMAT_PRIORITY: &[&str] = &[
22    "safetensors", "bin", "pt", "pth", "ckpt", "onnx", "h5", "hdf5", "keras", "npz", "npy", "gguf",
23];
24
25pub struct DownloadedRepo {
26    pub dir: PathBuf,
27    /// Extension of the format that was selected (e.g. "safetensors").
28    pub format: String,
29    /// Local paths of every checkpoint file downloaded, ready to pass to
30    /// `zsfm-checkpoint::load_checkpoint`.
31    pub checkpoint_files: Vec<PathBuf>,
32    /// `config.json`, if the repo has one (best-effort, not required).
33    pub config_json: Option<PathBuf>,
34}
35
36#[derive(Deserialize)]
37struct Sibling {
38    rfilename: String,
39}
40
41#[derive(Deserialize)]
42struct RepoInfo {
43    siblings: Vec<Sibling>,
44}
45
46/// List every file path in a repo at `revision` (default branch: `"main"`).
47pub async fn list_repo_files(
48    repo_id: &str,
49    revision: &str,
50    hf_token: Option<&str>,
51) -> anyhow::Result<Vec<String>> {
52    let client = build_client(hf_token)?;
53    let url = if revision == "main" {
54        format!("{HF_BASE}/api/models/{repo_id}")
55    } else {
56        format!("{HF_BASE}/api/models/{repo_id}/revision/{revision}")
57    };
58    let resp = client
59        .get(&url)
60        .send()
61        .await
62        .with_context(|| format!("GET {url}"))?;
63    let status = resp.status();
64    if !status.is_success() {
65        anyhow::bail!(
66            "HTTP {status} listing files for {repo_id}@{revision} — check the repo id, \
67             revision, and (for gated/private repos) that --token / HF_TOKEN is set"
68        );
69    }
70    let info: RepoInfo = resp.json().await.context("parse repo file listing")?;
71    Ok(info.siblings.into_iter().map(|s| s.rfilename).collect())
72}
73
74fn ext_of(path: &str) -> Option<String> {
75    Path::new(path)
76        .extension()
77        .and_then(|e| e.to_str())
78        .map(str::to_ascii_lowercase)
79}
80
81/// Download a repo's checkpoint in whatever format it's stored in (or a
82/// specific one, if `format` is given). Resolves HF's standard sharding
83/// indices for `safetensors` and `bin`; for every other format, if more than
84/// one file matches, `file_filter` (a case-insensitive substring) is required
85/// to pick which ones — this avoids silently pulling every alternative
86/// quantization/variant of a multi-GB checkpoint.
87pub async fn download_any_format(
88    repo_id: &str,
89    format: Option<&str>,
90    file_filter: Option<&str>,
91    revision: &str,
92    hf_token: Option<&str>,
93    dest_dir: &Path,
94) -> anyhow::Result<DownloadedRepo> {
95    let client = build_client(hf_token)?;
96    std::fs::create_dir_all(dest_dir).context("create dest dir")?;
97
98    let files = list_repo_files(repo_id, revision, hf_token).await?;
99    anyhow::ensure!(
100        !files.is_empty(),
101        "repo {repo_id}@{revision} has no files (wrong repo id/revision, or empty repo)"
102    );
103
104    let chosen_format = match format {
105        Some(f) => {
106            let f = f.trim_start_matches('.').to_ascii_lowercase();
107            anyhow::ensure!(
108                files.iter().any(|p| ext_of(p).as_deref() == Some(f.as_str())),
109                "repo {repo_id} has no .{f} files. Found extensions: {}",
110                describe_extensions(&files)
111            );
112            f
113        }
114        None => FORMAT_PRIORITY
115            .iter()
116            .find(|ext| files.iter().any(|p| ext_of(p).as_deref() == Some(**ext)))
117            .map(|s| s.to_string())
118            .with_context(|| {
119                format!(
120                    "repo {repo_id} has no recognized checkpoint format. Found extensions: {}",
121                    describe_extensions(&files)
122                )
123            })?,
124    };
125
126    let mut matches: Vec<&String> = files
127        .iter()
128        .filter(|p| ext_of(p).as_deref() == Some(chosen_format.as_str()))
129        .collect();
130    if let Some(filter) = file_filter {
131        let filter_lower = filter.to_ascii_lowercase();
132        matches.retain(|p| p.to_ascii_lowercase().contains(&filter_lower));
133        anyhow::ensure!(
134            !matches.is_empty(),
135            "no .{chosen_format} file in {repo_id} matches --file {filter}"
136        );
137    }
138
139    let checkpoint_files = if chosen_format == "safetensors"
140        && files.iter().any(|f| f == "model.safetensors.index.json")
141    {
142        download_sharded(&client, repo_id, "model.safetensors.index.json", dest_dir).await?
143    } else if (chosen_format == "bin" || chosen_format == "pt")
144        && files.iter().any(|f| f == "pytorch_model.bin.index.json")
145    {
146        download_sharded(&client, repo_id, "pytorch_model.bin.index.json", dest_dir).await?
147    } else if matches.len() == 1 {
148        vec![fetch_file(&client, repo_id, matches[0], dest_dir, None).await?]
149    } else {
150        anyhow::ensure!(
151            file_filter.is_some(),
152            "repo {repo_id} has {} .{chosen_format} files and none is HF-indexed sharding — \
153             pass --file <substring> to pick which one(s) to download. Candidates:\n  {}",
154            matches.len(),
155            matches
156                .iter()
157                .map(|s| s.as_str())
158                .collect::<Vec<_>>()
159                .join("\n  ")
160        );
161        let multi = indicatif::MultiProgress::new();
162        let mut results: Vec<(usize, PathBuf)> = stream::iter(matches.iter().enumerate())
163            .map(|(i, m)| {
164                let multi = &multi;
165                let client = client.clone();
166                async move {
167                    fetch_file(&client, repo_id, m, dest_dir, Some(multi))
168                        .await
169                        .map(|p| (i, p))
170                }
171            })
172            .buffer_unordered(SHARD_DOWNLOAD_CONCURRENCY)
173            .collect::<Vec<anyhow::Result<(usize, PathBuf)>>>()
174            .await
175            .into_iter()
176            .collect::<anyhow::Result<Vec<_>>>()?;
177        results.sort_by_key(|(i, _)| *i);
178        results.into_iter().map(|(_, p)| p).collect()
179    };
180
181    let config_json = if files.iter().any(|f| f == "config.json") {
182        fetch_file(&client, repo_id, "config.json", dest_dir, None).await.ok()
183    } else {
184        None
185    };
186
187    Ok(DownloadedRepo {
188        dir: dest_dir.to_path_buf(),
189        format: chosen_format,
190        checkpoint_files,
191        config_json,
192    })
193}
194
195/// Download an HF sharding index (`*.index.json`) and every shard it names.
196async fn download_sharded(
197    client: &reqwest::Client,
198    repo_id: &str,
199    index_relpath: &str,
200    dest_dir: &Path,
201) -> anyhow::Result<Vec<PathBuf>> {
202    let index_path = fetch_file(client, repo_id, index_relpath, dest_dir, None).await?;
203    #[derive(Deserialize)]
204    struct Index {
205        weight_map: HashMap<String, String>,
206    }
207    let raw = std::fs::read_to_string(&index_path).context("read sharding index")?;
208    let index: Index = serde_json::from_str(&raw).context("parse sharding index")?;
209
210    let mut shard_names: Vec<String> = index.weight_map.into_values().collect();
211    shard_names.sort();
212    shard_names.dedup();
213    anyhow::ensure!(!shard_names.is_empty(), "{index_relpath} lists no shards");
214
215    // See `download::resolve_shards` — same bounded-concurrency rationale.
216    let multi = indicatif::MultiProgress::new();
217    let mut results: Vec<(usize, PathBuf)> = stream::iter(shard_names.iter().enumerate())
218        .map(|(i, name)| {
219            let multi = &multi;
220            async move {
221                fetch_file(client, repo_id, name, dest_dir, Some(multi))
222                    .await
223                    .map(|p| (i, p))
224            }
225        })
226        .buffer_unordered(SHARD_DOWNLOAD_CONCURRENCY)
227        .collect::<Vec<anyhow::Result<(usize, PathBuf)>>>()
228        .await
229        .into_iter()
230        .collect::<anyhow::Result<Vec<_>>>()?;
231    results.sort_by_key(|(i, _)| *i);
232    Ok(results.into_iter().map(|(_, p)| p).collect())
233}
234
235fn describe_extensions(files: &[String]) -> String {
236    let mut exts: Vec<String> = files.iter().filter_map(|f| ext_of(f)).collect();
237    exts.sort();
238    exts.dedup();
239    if exts.is_empty() {
240        "(no file extensions found)".to_string()
241    } else {
242        exts.join(", ")
243    }
244}
zsfm_hub/discover.rs

zsfm_hub/
discover.rs
