fim_engine/
download.rs

1//! Model-file acquisition — downloads the quantized GGUF weights + the
2//! tokenizer from the HuggingFace CDN into a cache directory, skipping
3//! files already present. Plain blocking HTTP; no `hf-hub` dependency.
4
5use std::fs;
6use std::io::{Read, Write};
7use std::path::{Path, PathBuf};
8
9const TOKENIZER_FILE: &str = "tokenizer.json";
10
11/// Which qwen2.5-coder size to run. 1.5B is the fast default; 3B is
12/// noticeably smarter at multi-line completion but ~2x slower + a
13/// bigger download. (Instruct GGUFs — the base GGUF repos are gated;
14/// instruct retains FIM capability.)
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum ModelChoice {
17    Qwen1_5B,
18    Qwen3B,
19}
20
21impl ModelChoice {
22    /// Parse a config string (`"1.5b"` / `"3b"`); unknown ⇒ 1.5B.
23    pub fn parse(s: &str) -> Self {
24        match s.trim().to_ascii_lowercase().as_str() {
25            "3b" | "3" | "qwen3b" => ModelChoice::Qwen3B,
26            _ => ModelChoice::Qwen1_5B,
27        }
28    }
29    /// `(gguf_repo, gguf_file, tokenizer_repo)`.
30    fn sources(self) -> (&'static str, &'static str, &'static str) {
31        match self {
32            ModelChoice::Qwen1_5B => (
33                "Qwen/Qwen2.5-Coder-1.5B-Instruct-GGUF",
34                "qwen2.5-coder-1.5b-instruct-q4_k_m.gguf",
35                "Qwen/Qwen2.5-Coder-1.5B",
36            ),
37            ModelChoice::Qwen3B => (
38                "Qwen/Qwen2.5-Coder-3B-Instruct-GGUF",
39                "qwen2.5-coder-3b-instruct-q4_k_m.gguf",
40                "Qwen/Qwen2.5-Coder-3B",
41            ),
42        }
43    }
44    /// Per-size tokenizer cache name so 1.5B + 3B can coexist.
45    fn tokenizer_cache_name(self) -> &'static str {
46        match self {
47            ModelChoice::Qwen1_5B => "tokenizer-1.5b.json",
48            ModelChoice::Qwen3B => "tokenizer-3b.json",
49        }
50    }
51}
52
53/// Progress callback payload — emitted periodically during a download so
54/// the host can paint a progress bar.
55#[derive(Debug, Clone)]
56pub struct DownloadProgress {
57    /// Human label for the file in flight (`weights` / `tokenizer`).
58    pub label: &'static str,
59    /// Bytes received so far.
60    pub received: u64,
61    /// Total bytes, when the server reported a Content-Length.
62    pub total: Option<u64>,
63}
64
65/// Resolved on-disk paths to the two model files.
66#[derive(Debug, Clone)]
67pub struct ModelPaths {
68    pub gguf: PathBuf,
69    pub tokenizer: PathBuf,
70}
71
72/// Ensure both model files for `choice` exist in `cache_dir`,
73/// downloading whichever are missing. `progress` is called periodically
74/// during a download. Blocking — run on a worker thread.
75pub fn ensure_model(
76    cache_dir: &Path,
77    choice: ModelChoice,
78    progress: &(dyn Fn(DownloadProgress) + Sync),
79) -> Result<ModelPaths, String> {
80    fs::create_dir_all(cache_dir).map_err(|e| format!("create {}: {e}", cache_dir.display()))?;
81    let (gguf_repo, gguf_file, tok_repo) = choice.sources();
82    let gguf = cache_dir.join(gguf_file);
83    let tokenizer = cache_dir.join(choice.tokenizer_cache_name());
84
85    if !tokenizer.exists() {
86        let url = hf_url(tok_repo, TOKENIZER_FILE);
87        download(&url, &tokenizer, "tokenizer", progress)?;
88    }
89    if !gguf.exists() {
90        let url = hf_url(gguf_repo, gguf_file);
91        download(&url, &gguf, "weights", progress)?;
92    }
93    Ok(ModelPaths { gguf, tokenizer })
94}
95
96/// True when both model files for `choice` are already cached.
97pub fn is_model_cached(cache_dir: &Path, choice: ModelChoice) -> bool {
98    let (_, gguf_file, _) = choice.sources();
99    cache_dir.join(gguf_file).exists() && cache_dir.join(choice.tokenizer_cache_name()).exists()
100}
101
102fn hf_url(repo: &str, file: &str) -> String {
103    format!("https://huggingface.co/{repo}/resolve/main/{file}")
104}
105
106/// Stream a URL to `dest`, writing to a `.part` temp file first and
107/// renaming on success so an interrupted download never leaves a
108/// half-file that looks complete.
109fn download(
110    url: &str,
111    dest: &Path,
112    label: &'static str,
113    progress: &(dyn Fn(DownloadProgress) + Sync),
114) -> Result<(), String> {
115    let client = reqwest::blocking::Client::builder()
116        .build()
117        .map_err(|e| format!("http client: {e}"))?;
118    let mut resp = client
119        .get(url)
120        .send()
121        .map_err(|e| format!("GET {url}: {e}"))?;
122    if !resp.status().is_success() {
123        return Err(format!("GET {url}: HTTP {}", resp.status()));
124    }
125    let total = resp.content_length();
126    let part = dest.with_extension("part");
127    let mut file =
128        fs::File::create(&part).map_err(|e| format!("create {}: {e}", part.display()))?;
129    let mut buf = [0u8; 64 * 1024];
130    let mut received: u64 = 0;
131    let mut last_report: u64 = 0;
132    loop {
133        let n = resp
134            .read(&mut buf)
135            .map_err(|e| format!("read {label}: {e}"))?;
136        if n == 0 {
137            break;
138        }
139        file.write_all(&buf[..n])
140            .map_err(|e| format!("write {label}: {e}"))?;
141        received += n as u64;
142        // Report every ~4 MB so the callback isn't hammered.
143        if received - last_report >= 4 * 1024 * 1024 {
144            last_report = received;
145            progress(DownloadProgress {
146                label,
147                received,
148                total,
149            });
150        }
151    }
152    file.flush().map_err(|e| format!("flush {label}: {e}"))?;
153    drop(file);
154    fs::rename(&part, dest).map_err(|e| format!("finalize {}: {e}", dest.display()))?;
155    progress(DownloadProgress {
156        label,
157        received,
158        total,
159    });
160    Ok(())
161}
fim_engine/download.rs

fim_engine/
download.rs
