| 1 | //! HuggingFace repo inspection and base-model cloning for onde-cli. |
| 2 | //! |
| 3 | //! Provides helpers to: |
| 4 | //! - Check whether a HuggingFace repo exists and if it contains model files |
| 5 | //! - Fetch the list of recommended base models for fine-tuning |
| 6 | //! - Clone (download) a base model's safetensors files into the local HF cache |
| 7 | |
| 8 | use anyhow::{Context, Result}; |
| 9 | use serde::Deserialize; |
| 10 | |
| 11 | // --------------------------------------------------------------------------- |
| 12 | // Public types |
| 13 | // --------------------------------------------------------------------------- |
| 14 | |
| 15 | /// Result of inspecting a HuggingFace repository. |
| 16 | #[derive(Debug, Clone)] |
| 17 | #[allow(dead_code)] |
| 18 | pub enum RepoStatus { |
| 19 | /// Repo does not exist on HuggingFace (404). |
| 20 | NotFound, |
| 21 | /// Repo exists but contains no model files (safetensors/GGUF). |
| 22 | Empty { repo_id: String, files: Vec<String> }, |
| 23 | /// Repo exists and has model files ready for use. |
| 24 | HasModel { |
| 25 | repo_id: String, |
| 26 | files: Vec<String>, |
| 27 | /// Approximate total size in bytes of all model files. |
| 28 | model_size_bytes: u64, |
| 29 | }, |
| 30 | } |
| 31 | |
| 32 | /// A base model that can be used as a starting point for fine-tuning. |
| 33 | #[derive(Debug, Clone)] |
| 34 | pub struct BaseModel { |
| 35 | /// HuggingFace repo ID, e.g. `Qwen/Qwen3-0.6B` |
| 36 | pub repo_id: &'static str, |
| 37 | /// Human-readable display name |
| 38 | pub display_name: &'static str, |
| 39 | /// Approximate size of the safetensors model files |
| 40 | pub size_display: &'static str, |
| 41 | /// Parameter count description |
| 42 | pub params: &'static str, |
| 43 | /// Short description of recommended use case |
| 44 | pub description: &'static str, |
| 45 | } |
| 46 | |
| 47 | /// Progress updates emitted while checking or cloning a repo. |
| 48 | #[derive(Debug, Clone)] |
| 49 | pub enum CloneProgress { |
| 50 | /// Checking if the target repo exists on HuggingFace. |
| 51 | CheckingRepo, |
| 52 | /// Repo status has been determined. |
| 53 | RepoChecked(RepoStatus), |
| 54 | /// Creating the target repo on HuggingFace. |
| 55 | CreatingRepo, |
| 56 | /// Repo was created (or already existed). |
| 57 | RepoReady, |
| 58 | /// Failed with an error message. |
| 59 | Failed(String), |
| 60 | } |
| 61 | |
| 62 | // --------------------------------------------------------------------------- |
| 63 | // Base model catalog |
| 64 | // --------------------------------------------------------------------------- |
| 65 | |
| 66 | /// Recommended base models for LoRA fine-tuning. |
| 67 | /// |
| 68 | /// These are full-precision safetensors models (not GGUF) since fine-tuning |
| 69 | /// requires autograd-capable weights. Ordered by size (smallest first). |
| 70 | pub const BASE_MODELS: &[BaseModel] = &[ |
| 71 | BaseModel { |
| 72 | repo_id: "Qwen/Qwen3-0.6B", |
| 73 | display_name: "Qwen3 0.6B", |
| 74 | size_display: "~1.2 GB", |
| 75 | params: "751M", |
| 76 | description: "Smallest Qwen3 — ideal for mobile & quick experiments", |
| 77 | }, |
| 78 | BaseModel { |
| 79 | repo_id: "Qwen/Qwen2.5-1.5B-Instruct", |
| 80 | display_name: "Qwen2.5 1.5B Instruct", |
| 81 | size_display: "~3.0 GB", |
| 82 | params: "1.5B", |
| 83 | description: "Balanced size — good for instruction-following tasks", |
| 84 | }, |
| 85 | BaseModel { |
| 86 | repo_id: "Qwen/Qwen3-1.7B", |
| 87 | display_name: "Qwen3 1.7B", |
| 88 | size_display: "~3.4 GB", |
| 89 | params: "1.7B", |
| 90 | description: "Latest Qwen3 small model — strong reasoning", |
| 91 | }, |
| 92 | BaseModel { |
| 93 | repo_id: "Qwen/Qwen3-4B", |
| 94 | display_name: "Qwen3 4B", |
| 95 | size_display: "~8.0 GB", |
| 96 | params: "4B", |
| 97 | description: "Larger Qwen3 — best quality, macOS recommended", |
| 98 | }, |
| 99 | ]; |
| 100 | |
| 101 | // --------------------------------------------------------------------------- |
| 102 | // HF API response types |
| 103 | // --------------------------------------------------------------------------- |
| 104 | |
| 105 | #[derive(Debug, Deserialize)] |
| 106 | struct HfRepoInfo { |
| 107 | #[serde(default)] |
| 108 | siblings: Vec<HfFileSibling>, |
| 109 | } |
| 110 | |
| 111 | #[derive(Debug, Deserialize)] |
| 112 | struct HfFileSibling { |
| 113 | rfilename: String, |
| 114 | #[serde(default)] |
| 115 | size: Option<u64>, |
| 116 | } |
| 117 | |
| 118 | // --------------------------------------------------------------------------- |
| 119 | // Public API |
| 120 | // --------------------------------------------------------------------------- |
| 121 | |
| 122 | /// Check whether a HuggingFace repo exists and inspect its contents. |
| 123 | /// |
| 124 | /// Returns [`RepoStatus::NotFound`] if the repo doesn't exist (HTTP 404), |
| 125 | /// [`RepoStatus::Empty`] if it exists but has no model files, or |
| 126 | /// [`RepoStatus::HasModel`] if it contains safetensors or GGUF files. |
| 127 | pub async fn check_repo(repo_id: &str, hf_token: &str) -> Result<RepoStatus> { |
| 128 | let client = build_client()?; |
| 129 | |
| 130 | let mut request = client.get(format!("https://huggingface.co/api/models/{repo_id}")); |
| 131 | if !hf_token.is_empty() { |
| 132 | request = request.bearer_auth(hf_token); |
| 133 | } |
| 134 | |
| 135 | let response = request |
| 136 | .send() |
| 137 | .await |
| 138 | .context("Failed to reach HuggingFace API")?; |
| 139 | |
| 140 | if response.status() == reqwest::StatusCode::NOT_FOUND { |
| 141 | return Ok(RepoStatus::NotFound); |
| 142 | } |
| 143 | |
| 144 | let response = response |
| 145 | .error_for_status() |
| 146 | .context("HuggingFace API error")?; |
| 147 | |
| 148 | let info: HfRepoInfo = response |
| 149 | .json() |
| 150 | .await |
| 151 | .context("Failed to parse repo metadata")?; |
| 152 | |
| 153 | let all_files: Vec<String> = info.siblings.iter().map(|s| s.rfilename.clone()).collect(); |
| 154 | |
| 155 | // Check for model files: safetensors or GGUF |
| 156 | let model_files: Vec<&HfFileSibling> = info |
| 157 | .siblings |
| 158 | .iter() |
| 159 | .filter(|s| is_model_file(&s.rfilename)) |
| 160 | .collect(); |
| 161 | |
| 162 | if model_files.is_empty() { |
| 163 | Ok(RepoStatus::Empty { |
| 164 | repo_id: repo_id.to_string(), |
| 165 | files: all_files, |
| 166 | }) |
| 167 | } else { |
| 168 | let total_size: u64 = model_files.iter().filter_map(|s| s.size).sum(); |
| 169 | Ok(RepoStatus::HasModel { |
| 170 | repo_id: repo_id.to_string(), |
| 171 | files: all_files, |
| 172 | model_size_bytes: total_size, |
| 173 | }) |
| 174 | } |
| 175 | } |
| 176 | |
| 177 | /// Create a HuggingFace repo if it doesn't already exist. |
| 178 | /// |
| 179 | /// Returns `Ok(())` on success or if the repo already exists (HTTP 409). |
| 180 | pub async fn create_repo(repo_id: &str, hf_token: &str) -> Result<()> { |
| 181 | let client = build_client()?; |
| 182 | |
| 183 | let body = if let Some((org, name)) = repo_id.split_once('/') { |
| 184 | serde_json::json!({ |
| 185 | "type": "model", |
| 186 | "name": name, |
| 187 | "organization": org, |
| 188 | "private": false, |
| 189 | "sdk": "onde-cli" |
| 190 | }) |
| 191 | } else { |
| 192 | serde_json::json!({ |
| 193 | "type": "model", |
| 194 | "name": repo_id, |
| 195 | "private": false, |
| 196 | "sdk": "onde-cli" |
| 197 | }) |
| 198 | }; |
| 199 | |
| 200 | let response = client |
| 201 | .post("https://huggingface.co/api/repos/create") |
| 202 | .bearer_auth(hf_token) |
| 203 | .json(&body) |
| 204 | .send() |
| 205 | .await |
| 206 | .context("Failed to reach HuggingFace API")?; |
| 207 | |
| 208 | let status = response.status(); |
| 209 | |
| 210 | // 200 = created, 409 = already exists — both are fine. |
| 211 | if !status.is_success() && status.as_u16() != 409 { |
| 212 | let body_text = response.text().await.unwrap_or_default(); |
| 213 | anyhow::bail!("Failed to create repo {repo_id}: {status} {body_text}"); |
| 214 | } |
| 215 | |
| 216 | Ok(()) |
| 217 | } |
| 218 | |
| 219 | /// Spawn-friendly entry point that checks a repo and sends progress via channel. |
| 220 | pub async fn start_check_repo( |
| 221 | repo_id: String, |
| 222 | hf_token: String, |
| 223 | tx: tokio::sync::mpsc::UnboundedSender<CloneProgress>, |
| 224 | ) { |
| 225 | let _ = tx.send(CloneProgress::CheckingRepo); |
| 226 | |
| 227 | match check_repo(&repo_id, &hf_token).await { |
| 228 | Ok(status) => { |
| 229 | let _ = tx.send(CloneProgress::RepoChecked(status)); |
| 230 | } |
| 231 | Err(e) => { |
| 232 | let _ = tx.send(CloneProgress::Failed(format!("{e:#}"))); |
| 233 | } |
| 234 | } |
| 235 | } |
| 236 | |
| 237 | /// Create the repo (if needed) and report progress. |
| 238 | pub async fn start_create_repo( |
| 239 | repo_id: String, |
| 240 | hf_token: String, |
| 241 | tx: tokio::sync::mpsc::UnboundedSender<CloneProgress>, |
| 242 | ) { |
| 243 | let _ = tx.send(CloneProgress::CreatingRepo); |
| 244 | |
| 245 | match create_repo(&repo_id, &hf_token).await { |
| 246 | Ok(()) => { |
| 247 | let _ = tx.send(CloneProgress::RepoReady); |
| 248 | } |
| 249 | Err(e) => { |
| 250 | let _ = tx.send(CloneProgress::Failed(format!("{e:#}"))); |
| 251 | } |
| 252 | } |
| 253 | } |
| 254 | |
| 255 | // --------------------------------------------------------------------------- |
| 256 | // Private helpers |
| 257 | // --------------------------------------------------------------------------- |
| 258 | |
| 259 | fn build_client() -> Result<reqwest::Client> { |
| 260 | reqwest::Client::builder() |
| 261 | .user_agent("onde-cli/0.2") |
| 262 | .redirect(reqwest::redirect::Policy::limited(10)) |
| 263 | .build() |
| 264 | .context("Failed to build HTTP client") |
| 265 | } |
| 266 | |
| 267 | /// Check if a filename looks like a model weight file. |
| 268 | fn is_model_file(filename: &str) -> bool { |
| 269 | let lower = filename.to_lowercase(); |
| 270 | lower.ends_with(".safetensors") |
| 271 | || lower.ends_with(".gguf") |
| 272 | || lower.ends_with(".bin") |
| 273 | || lower.ends_with(".pt") |
| 274 | || lower.ends_with(".pth") |
| 275 | } |