zsfm_hub/
download.rs

1use std::path::{Path, PathBuf};
2
3use anyhow::Context;
4use futures_util::{stream, StreamExt};
5use serde::Deserialize;
6use tokio::io::AsyncWriteExt;
7
8use crate::http::{build_client, HF_BASE};
9
10/// How many shard files to fetch concurrently. HF's CDN comfortably serves this many
11/// parallel range/GET requests per client; higher offers diminishing returns and risks
12/// the server throttling or the local connection pool thrashing.
13pub(crate) const SHARD_DOWNLOAD_CONCURRENCY: usize = 4;
14
15/// Paths to all local files needed for conversion.
16pub struct ModelFiles {
17    pub config_json: PathBuf,
18    /// Ordered list of safetensors shard paths, already downloaded locally.
19    pub safetensors_shards: Vec<PathBuf>,
20}
21
22impl ModelFiles {
23    /// Delete the (potentially large) weight shards once they've been converted into a
24    /// GGUF, leaving `config_json` in place — several `infer` commands read it directly
25    /// (architecture parameters not embedded in the GGUF), and it's tiny, so there's no
26    /// disk-space reason to remove it. Best-effort: a failed delete is not fatal, since
27    /// the GGUF conversion has already succeeded by the time this is called.
28    pub fn cleanup_weights(&self) {
29        for f in &self.safetensors_shards {
30            if let Err(e) = std::fs::remove_file(f) {
31                eprintln!("warning: could not remove {}: {e}", f.display());
32            }
33        }
34    }
35}
36
37/// Where the canonical, always-F32 GGUF for `repo_id` lives once converted once. Every
38/// per-model `convert` command checks this path first: if present, it recasts straight
39/// from this cached GGUF to whatever dtype was requested instead of re-downloading and
40/// re-converting from HuggingFace.
41pub fn canonical_gguf_path(model_dir: &Path, repo_id: &str) -> PathBuf {
42    model_dir.join(repo_id.replace('/', "__")).join("model-f32.gguf")
43}
44
45/// Download (or locate from cache) `config.json` + `model.safetensors[.index.json]`
46/// for `repo_id` into `model_dir`.
47pub async fn download_model(
48    repo_id: &str,
49    hf_token: Option<&str>,
50    model_dir: &Path,
51) -> anyhow::Result<ModelFiles> {
52    download_model_prefixed(repo_id, "", hf_token, model_dir).await
53}
54
55/// Download (or locate from cache) a single arbitrary file from `repo_id`, resuming a partial
56/// download if one exists. For models that don't fit the `config.json` + `model.safetensors`
57/// shape — e.g. Lag-Llama's raw PyTorch Lightning `.ckpt` checkpoint.
58pub async fn download_file(
59    repo_id: &str,
60    relpath: &str,
61    hf_token: Option<&str>,
62    dest_dir: &Path,
63) -> anyhow::Result<PathBuf> {
64    let client = build_client(hf_token)?;
65    std::fs::create_dir_all(dest_dir).context("create dest dir")?;
66    fetch_file(&client, repo_id, relpath, dest_dir, None).await
67}
68
69/// Same as [`download_model`], but every remote path is joined under `prefix` first.
70///
71/// Pass `""` for a flat repo layout (`config.json`, `model.safetensors`).
72/// Pass e.g. `"classification"` for a repo that keeps multiple task variants
73/// in subfolders (`classification/config.json`, `classification/model.safetensors`),
74/// as TabFM does.
75pub async fn download_model_prefixed(
76    repo_id: &str,
77    prefix: &str,
78    hf_token: Option<&str>,
79    model_dir: &Path,
80) -> anyhow::Result<ModelFiles> {
81    let client = build_client(hf_token)?;
82    // Namespace by repo (matching the `owner__name` convention the generic `zsfm convert
83    // --repo` path already uses) so two different repos sharing the generic `config.json` +
84    // `model.safetensors` naming — which is most of them — never collide when a caller passes
85    // the same (often default) `model_dir` for both, whether run sequentially or concurrently.
86    // Without this, the second download's `config.json`/`model.safetensors` would either
87    // silently overwrite the first's files, or worse, get "(cached)" hit on the *wrong* model's
88    // bytes.
89    let model_dir = model_dir.join(repo_id.replace('/', "__"));
90    let model_dir = model_dir.as_path();
91    std::fs::create_dir_all(model_dir).context("create model dir")?;
92
93    let config_rel = joined(prefix, "config.json");
94    println!("Fetching {config_rel} …");
95    let config_json = fetch_file(&client, repo_id, &config_rel, model_dir, None).await?;
96
97    // Detect sharded model by fetching the index file.
98    let index_rel = joined(prefix, "model.safetensors.index.json");
99    let single_rel = joined(prefix, "model.safetensors");
100
101    let shards = match fetch_file(&client, repo_id, &index_rel, model_dir, None).await {
102        Ok(index_path) => {
103            println!("Found sharded model — reading index …");
104            resolve_shards(&client, repo_id, &index_path, model_dir, prefix).await?
105        }
106        Err(_) => {
107            println!("Fetching {single_rel} …");
108            let shard = fetch_file(&client, repo_id, &single_rel, model_dir, None).await?;
109            vec![shard]
110        }
111    };
112
113    Ok(ModelFiles {
114        config_json,
115        safetensors_shards: shards,
116    })
117}
118
119fn joined(prefix: &str, filename: &str) -> String {
120    if prefix.is_empty() {
121        filename.to_string()
122    } else {
123        format!("{prefix}/{filename}")
124    }
125}
126
127/// Download `relpath` (may include subfolders, e.g. `"classification/config.json"`)
128/// from `repo_id` into `dest_dir`, resuming if a partial `.tmp` file already exists.
129/// Returns the local path of the completed file.
130///
131/// `mp`: when several `fetch_file` calls run concurrently (see `resolve_shards`), pass a
132/// shared [`indicatif::MultiProgress`] so their bars stack cleanly instead of each fighting
133/// to redraw the same terminal line; `None` draws a lone standalone bar.
134pub(crate) async fn fetch_file(
135    client: &reqwest::Client,
136    repo_id: &str,
137    relpath: &str,
138    dest_dir: &Path,
139    mp: Option<&indicatif::MultiProgress>,
140) -> anyhow::Result<PathBuf> {
141    let dest = dest_dir.join(relpath.replace('/', "_"));
142    if dest.exists() {
143        println!("  (cached) {relpath}");
144        return Ok(dest);
145    }
146
147    let dest_tmp = dest.with_extension("tmp");
148    let already = if dest_tmp.exists() {
149        dest_tmp.metadata()?.len()
150    } else {
151        0
152    };
153
154    let url = format!("{HF_BASE}/{repo_id}/resolve/main/{relpath}");
155
156    let mut req = client.get(&url);
157    if already > 0 {
158        req = req.header(reqwest::header::RANGE, format!("bytes={already}-"));
159        println!("  Resuming {relpath} from {} MB …", already / 1_000_000);
160    }
161
162    let response = req
163        .send()
164        .await
165        .with_context(|| format!("GET {url}"))?;
166
167    let status = response.status();
168    // 206 = partial content (resume accepted), 200 = full content
169    if !status.is_success() {
170        anyhow::bail!("HTTP {status} fetching {relpath} from {repo_id}");
171    }
172
173    // If server ignored the Range header and sent 200, truncate the tmp file.
174    let (file, resume_offset) = if status == reqwest::StatusCode::PARTIAL_CONTENT {
175        let f = tokio::fs::OpenOptions::new()
176            .append(true)
177            .open(&dest_tmp)
178            .await
179            .with_context(|| format!("open tmp {}", dest_tmp.display()))?;
180        (f, already)
181    } else {
182        let f = tokio::fs::File::create(&dest_tmp)
183            .await
184            .with_context(|| format!("create tmp {}", dest_tmp.display()))?;
185        (f, 0)
186    };
187
188    let total = response
189        .content_length()
190        .map(|n| n + resume_offset)
191        .unwrap_or(0);
192
193    let pb = indicatif::ProgressBar::new(total);
194    let pb = match mp {
195        Some(mp) => mp.add(pb),
196        None => pb,
197    };
198    pb.set_style(
199        indicatif::ProgressStyle::with_template(
200            "  {msg} [{bar:40}] {bytes}/{total_bytes} ({bytes_per_sec}, eta {eta})",
201        )
202        .unwrap()
203        .progress_chars("=>-"),
204    );
205    pb.set_message(relpath.to_string());
206    pb.set_position(resume_offset);
207
208    {
209        let mut file = file;
210        let mut stream = response.bytes_stream();
211        while let Some(chunk) = stream.next().await {
212            let chunk = chunk.with_context(|| format!("stream chunk of {relpath}"))?;
213            pb.inc(chunk.len() as u64);
214            file.write_all(&chunk)
215                .await
216                .with_context(|| format!("write chunk to {}", dest_tmp.display()))?;
217        }
218    }
219
220    pb.finish_and_clear();
221    std::fs::rename(&dest_tmp, &dest)
222        .with_context(|| format!("rename tmp → {}", dest.display()))?;
223
224    Ok(dest)
225}
226
227/// Parse the shard index JSON and download every unique shard (joined under `prefix`).
228async fn resolve_shards(
229    client: &reqwest::Client,
230    repo_id: &str,
231    index_path: &Path,
232    cache_dir: &Path,
233    prefix: &str,
234) -> anyhow::Result<Vec<PathBuf>> {
235    #[derive(Deserialize)]
236    struct Index {
237        weight_map: std::collections::HashMap<String, String>,
238    }
239
240    let raw = std::fs::read_to_string(index_path).context("read index json")?;
241    let index: Index = serde_json::from_str(&raw).context("parse index json")?;
242
243    let mut shard_names: Vec<String> = index.weight_map.into_values().collect();
244    shard_names.sort();
245    shard_names.dedup();
246
247    // Fetch up to SHARD_DOWNLOAD_CONCURRENCY shards at once — each is an independent file,
248    // so there's no reason to serialize what's fundamentally a bandwidth-bound operation.
249    // `buffer_unordered` completes them in whatever order the network delivers them; each
250    // result is tagged with its original index so the returned `Vec<PathBuf>` still matches
251    // `shard_names`' sorted order (some downstream converters iterate shards positionally).
252    let multi = indicatif::MultiProgress::new();
253    let mut results: Vec<(usize, PathBuf)> = stream::iter(shard_names.iter().enumerate())
254        .map(|(i, name)| {
255            let rel = joined(prefix, name);
256            let multi = &multi;
257            async move {
258                println!("  Fetching {rel} …");
259                fetch_file(client, repo_id, &rel, cache_dir, Some(multi))
260                    .await
261                    .map(|p| (i, p))
262            }
263        })
264        .buffer_unordered(SHARD_DOWNLOAD_CONCURRENCY)
265        .collect::<Vec<anyhow::Result<(usize, PathBuf)>>>()
266        .await
267        .into_iter()
268        .collect::<anyhow::Result<Vec<_>>>()?;
269    results.sort_by_key(|(i, _)| *i);
270    Ok(results.into_iter().map(|(_, p)| p).collect())
271}
zsfm_hub/download.rs

zsfm_hub/
download.rs
