@ondeinference / cli
Code Commits Issues Pulls Cloud Agent CI/CD
development
development main
rs 227 lines 7.75 KB
Raw
1 //! HuggingFace Hub model search and download for onde-cli.
2 //!
3 //! `search_hf` queries the public HF API (no auth required for public models).
4 //! `download_model` streams all files into the standard HF hub cache layout:
5 //! `{hub}/models--{org}--{name}/snapshots/{sha}/{filename}`
6
7 use anyhow::{Context, Result};
8 use futures::StreamExt;
9 use serde::{Deserialize, Serialize};
10 use std::path::PathBuf;
11 use tokio::io::AsyncWriteExt;
12
13 // ---------------------------------------------------------------------------
14 // Public types
15 // ---------------------------------------------------------------------------
16
17 /// Compact model info returned by the HF search API.
18 #[derive(Debug, Clone, Deserialize, Serialize)]
19 pub struct HfModelInfo {
20 #[serde(rename = "id")]
21 pub model_id: String,
22 #[serde(default)]
23 pub downloads: u64,
24 #[serde(default)]
25 #[allow(dead_code)]
26 pub tags: Vec<String>,
27 /// Commit SHA of the latest revision.
28 pub sha: Option<String>,
29 /// Files present in the repository.
30 #[serde(default)]
31 pub siblings: Vec<HfSibling>,
32 }
33
34 /// A file entry in a HuggingFace model repository.
35 #[derive(Debug, Clone, Deserialize, Serialize)]
36 pub struct HfSibling {
37 pub rfilename: String,
38 }
39
40 /// Progress update for an in-progress model download.
41 #[derive(Debug, Clone)]
42 pub struct DownloadProgress {
43 pub model_id: String,
44 pub filename: String,
45 pub file_index: usize,
46 pub total_files: usize,
47 pub file_bytes_done: u64,
48 pub file_bytes_total: u64,
49 }
50
51 /// Events emitted by the download background task.
52 pub enum DownloadEvent {
53 Progress(DownloadProgress),
54 Complete,
55 Failed(String),
56 }
57
58 // ---------------------------------------------------------------------------
59 // Public API
60 // ---------------------------------------------------------------------------
61
62 /// Search the HuggingFace Hub for models matching `query` (up to 20 results).
63 pub async fn search_hf(query: &str) -> Result<Vec<HfModelInfo>> {
64 let client = build_client()?;
65 let results: Vec<HfModelInfo> = client
66 .get("https://huggingface.co/api/models")
67 .query(&[("search", query), ("limit", "20"), ("full", "false")])
68 .send()
69 .await
70 .context("HuggingFace API unreachable")?
71 .error_for_status()
72 .context("HuggingFace API returned an error")?
73 .json()
74 .await
75 .context("Failed to parse HuggingFace API response")?;
76 Ok(results)
77 }
78
79 /// Download all files of a HuggingFace model into `hub_dir` using the standard
80 /// HF hub cache layout. Progress and completion are reported via `tx`.
81 ///
82 /// This function is infallible — all errors are sent as `DownloadEvent::Failed`.
83 pub async fn download_model(
84 model_id: String,
85 hub_dir: PathBuf,
86 tx: tokio::sync::mpsc::UnboundedSender<DownloadEvent>,
87 ) {
88 match run_download(model_id, hub_dir, &tx).await {
89 Ok(()) => {
90 let _ = tx.send(DownloadEvent::Complete);
91 }
92 Err(e) => {
93 let _ = tx.send(DownloadEvent::Failed(format!("{e:#}")));
94 }
95 }
96 }
97
98 // ---------------------------------------------------------------------------
99 // Private helpers
100 // ---------------------------------------------------------------------------
101
102 fn build_client() -> Result<reqwest::Client> {
103 reqwest::Client::builder()
104 .user_agent("onde-cli/0.1")
105 .redirect(reqwest::redirect::Policy::limited(10))
106 .build()
107 .context("Failed to build HTTP client")
108 }
109
110 /// Fetch full model metadata (commit SHA + file list).
111 async fn fetch_model_info(client: &reqwest::Client, model_id: &str) -> Result<HfModelInfo> {
112 let url = format!("https://huggingface.co/api/models/{model_id}");
113 let info: HfModelInfo = client
114 .get(&url)
115 .send()
116 .await
117 .context("Failed to fetch model metadata")?
118 .error_for_status()
119 .context("Model not found or inaccessible")?
120 .json()
121 .await
122 .context("Failed to parse model metadata")?;
123 Ok(info)
124 }
125
126 async fn run_download(
127 model_id: String,
128 hub_dir: PathBuf,
129 tx: &tokio::sync::mpsc::UnboundedSender<DownloadEvent>,
130 ) -> Result<()> {
131 let client = build_client()?;
132
133 // Fetch metadata to get the commit SHA and file list.
134 let info = fetch_model_info(&client, &model_id).await?;
135 let sha = info.sha.as_deref().unwrap_or("main");
136
137 // Build snapshot directory: {hub}/models--{org}--{name}/snapshots/{sha}/
138 let dir_name = format!("models--{}", model_id.replace('/', "--"));
139 let snapshot_dir = hub_dir.join(&dir_name).join("snapshots").join(sha);
140 tokio::fs::create_dir_all(&snapshot_dir)
141 .await
142 .context("Failed to create snapshot directory")?;
143
144 // Write refs/main so the standard HF cache tools can find the revision.
145 let refs_dir = hub_dir.join(&dir_name).join("refs");
146 tokio::fs::create_dir_all(&refs_dir).await?;
147 tokio::fs::write(refs_dir.join("main"), sha).await?;
148
149 // Download each file.
150 let files: Vec<String> = info.siblings.iter().map(|s| s.rfilename.clone()).collect();
151 let total_files = files.len();
152
153 for (file_index, filename) in files.iter().enumerate() {
154 let dest = snapshot_dir.join(filename);
155
156 // Create parent dirs for nested paths (e.g. "subfolder/file.json").
157 if let Some(parent) = dest.parent() {
158 tokio::fs::create_dir_all(parent).await?;
159 }
160
161 // Skip files that are already fully present.
162 if dest.exists() {
163 // Still send a progress event so the UI reflects the file count.
164 let _ = tx.send(DownloadEvent::Progress(DownloadProgress {
165 model_id: model_id.clone(),
166 filename: filename.clone(),
167 file_index,
168 total_files,
169 file_bytes_done: 0,
170 file_bytes_total: 0,
171 }));
172 continue;
173 }
174
175 let url = format!("https://huggingface.co/{model_id}/resolve/{sha}/{filename}");
176 let response = client
177 .get(&url)
178 .send()
179 .await
180 .with_context(|| format!("Failed to download {filename}"))?
181 .error_for_status()
182 .with_context(|| format!("Server error downloading {filename}"))?;
183
184 let total_bytes = response.content_length().unwrap_or(0);
185 let mut downloaded: u64 = 0;
186 let mut last_reported: u64 = 0;
187
188 // Write to a temp file first, then rename on success.
189 let tmp_dest = snapshot_dir.join(format!("{filename}.tmp"));
190 let mut file = tokio::fs::File::create(&tmp_dest)
191 .await
192 .with_context(|| format!("Failed to create {}", tmp_dest.display()))?;
193
194 let mut stream = response.bytes_stream();
195
196 while let Some(chunk) = stream.next().await {
197 let chunk = chunk.with_context(|| format!("Download error for {filename}"))?;
198 file.write_all(&chunk)
199 .await
200 .with_context(|| format!("Write error for {filename}"))?;
201 downloaded += chunk.len() as u64;
202
203 // Report progress roughly every 512 KB.
204 if downloaded.saturating_sub(last_reported) >= 512 * 1024
205 || (total_bytes > 0 && downloaded >= total_bytes)
206 {
207 last_reported = downloaded;
208 let _ = tx.send(DownloadEvent::Progress(DownloadProgress {
209 model_id: model_id.clone(),
210 filename: filename.clone(),
211 file_index,
212 total_files,
213 file_bytes_done: downloaded,
214 file_bytes_total: total_bytes,
215 }));
216 }
217 }
218
219 file.flush().await?;
220 drop(file);
221 tokio::fs::rename(&tmp_dest, &dest)
222 .await
223 .with_context(|| format!("Failed to finalize {filename}"))?;
224 }
225
226 Ok(())
227 }