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

zsfm_hub/
discover.rs
