voice_bird_cli/transcription/
models.rs

1use std::path::{Path, PathBuf};
2
3use anyhow::{anyhow, Context};
4use flate2::read::GzDecoder;
5use sha2::{Digest, Sha256};
6use tar::Archive;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum ModelFormat {
10    WhisperGguf,
11    NemotronPackage,
12}
13
14#[derive(Debug, Clone)]
15pub struct ModelEntry {
16    pub id: &'static str,
17    pub size_mb: u32,
18    pub language: &'static str,
19    pub format: ModelFormat,
20    /// Source URL for the primary artifact: a Whisper GGUF `.bin` for
21    /// `WhisperGguf`, or the Nemotron `.tar.gz` package for `NemotronPackage`.
22    pub download_url: &'static str,
23    pub download_sha256: &'static str,
24    pub coreml_url: Option<&'static str>,
25    pub coreml_sha256: Option<&'static str>,
26    pub is_default: bool,
27}
28
29pub struct Catalog(Vec<ModelEntry>);
30
31impl Catalog {
32    pub fn builtin() -> Self {
33        Catalog(vec![
34            ModelEntry {
35                id: "distil-small.en",
36                size_mb: 250,
37                language: "en",
38                format: ModelFormat::WhisperGguf,
39                download_url: "https://huggingface.co/distil-whisper/distil-small.en/resolve/main/ggml-distil-small.en.bin",
40                download_sha256: "7691eb11167ab7aaf6b3e05d8266f2fd9ad89c550e433f86ac266ebdee6c970a",
41                // No standalone CoreML zip is published for this model;
42                // argmaxinc/whisperkit-coreml serves .mlmodelc directory trees, not zips.
43                coreml_url: None,
44                coreml_sha256: None,
45                is_default: true,
46            },
47            ModelEntry {
48                id: "distil-large-v3",
49                size_mb: 1_500,
50                language: "multi",
51                format: ModelFormat::WhisperGguf,
52                download_url: "https://huggingface.co/distil-whisper/distil-large-v3-ggml/resolve/main/ggml-distil-large-v3.bin",
53                download_sha256: "2883a11b90fb10ed592d826edeaee7d2929bf1ab985109fe9e1e7b4d2b69a298",
54                coreml_url: None,
55                coreml_sha256: None,
56                is_default: false,
57            },
58            ModelEntry {
59                id: "large-v3-turbo",
60                size_mb: 1_600,
61                language: "multi",
62                format: ModelFormat::WhisperGguf,
63                download_url: "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-large-v3-turbo.bin",
64                download_sha256: "1fc70f774d38eb169993ac391eea357ef47c88757ef72ee5943879b7e8e2bc69",
65                coreml_url: Some("https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-large-v3-turbo-encoder.mlmodelc.zip"),
66                coreml_sha256: Some("84bedfe895bd7b5de6e8e89a0803dfc5addf8c0c5bc4c937451716bf7cf7988a"),
67                is_default: false,
68            },
69            ModelEntry {
70                id: "nemotron-3.5-asr-streaming-0.6b",
71                size_mb: 740,
72                language: "multi",
73                format: ModelFormat::NemotronPackage,
74                download_url: "https://huggingface.co/smcleod/nemotron-3.5-asr-streaming-0.6b-int8/resolve/main/nemotron-3.5-asr-streaming-0.6b-int8.tar.gz",
75                download_sha256: "d1d57d86212528fa03dfdbb88979f1dd637814dec6db31257a603739c73bd9d2",
76                coreml_url: None,
77                coreml_sha256: None,
78                is_default: false,
79            },
80            ModelEntry {
81                id: "base.en",
82                size_mb: 150,
83                language: "en",
84                format: ModelFormat::WhisperGguf,
85                download_url: "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-base.en.bin",
86                download_sha256: "a03779c86df3323075f5e796cb2ce5029f00ec8869eee3fdfb897afe36c6d002",
87                coreml_url: Some("https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-base.en-encoder.mlmodelc.zip"),
88                coreml_sha256: Some("8cf860309e2449e2bdc8be834cf838ab2565747ecc8c0ef914ef5975115e192b"),
89                is_default: false,
90            },
91            ModelEntry {
92                id: "tiny.en",
93                size_mb: 75,
94                language: "en",
95                format: ModelFormat::WhisperGguf,
96                download_url: "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-tiny.en.bin",
97                download_sha256: "921e4cf8686fdd993dcd081a5da5b6c365bfde1162e72b08d75ac75289920b1f",
98                coreml_url: Some("https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-tiny.en-encoder.mlmodelc.zip"),
99                coreml_sha256: Some("82b32eef73c94bb0c432a776a047b757d9525c26d84038a15d8798d7c8d1ee58"),
100                is_default: false,
101            },
102        ])
103    }
104
105    pub fn all(&self) -> &[ModelEntry] {
106        &self.0
107    }
108
109    pub fn get(&self, id: &str) -> Option<&ModelEntry> {
110        self.0.iter().find(|m| m.id == id)
111    }
112}
113
114pub fn validate_local_language(model_id: &str, language: &str) -> Result<(), String> {
115    let lang = language.trim();
116    if lang.is_empty() || lang == "en" || lang == "auto" {
117        return Ok(());
118    }
119
120    let catalog = Catalog::builtin();
121    let Some(entry) = catalog.get(model_id) else {
122        return Err(format!(
123            "Model '{model_id}' is not supported by this release; pick one from the model picker."
124        ));
125    };
126
127    if entry.language == "en" {
128        return Err(format!(
129            "Model '{}' is English-only; pick distil-large-v3, large-v3-turbo, or nemotron-3.5-asr-streaming-0.6b for {}.",
130            entry.id, lang
131        ));
132    }
133
134    Ok(())
135}
136
137pub fn cache_dir() -> anyhow::Result<PathBuf> {
138    let base = dirs::cache_dir().ok_or_else(|| anyhow!("no cache dir"))?;
139    Ok(base.join("voice-bird").join("models"))
140}
141
142pub fn gguf_path(id: &str) -> anyhow::Result<PathBuf> {
143    Ok(cache_dir()?.join(format!("{id}.gguf")))
144}
145
146/// Directory where the CoreML encoder for a GGUF model must live for
147/// whisper.cpp to auto-load it. whisper.cpp derives this from the GGUF path by
148/// stripping the extension and appending `-encoder.mlmodelc`, so for
149/// `<cache>/{id}.gguf` it looks for `<cache>/{id}-encoder.mlmodelc`.
150pub fn coreml_path(id: &str) -> anyhow::Result<PathBuf> {
151    Ok(cache_dir()?.join(format!("{id}-encoder.mlmodelc")))
152}
153
154pub fn nemotron_model_dir(id: &str) -> anyhow::Result<PathBuf> {
155    Ok(cache_dir()?.join(id))
156}
157
158pub fn model_path(id: &str) -> anyhow::Result<PathBuf> {
159    let catalog = Catalog::builtin();
160    match catalog.get(id).map(|m| m.format) {
161        Some(ModelFormat::NemotronPackage) => nemotron_model_dir(id),
162        _ => gguf_path(id),
163    }
164}
165
166pub fn is_nemotron_model(id: &str) -> bool {
167    let catalog = Catalog::builtin();
168    matches!(
169        catalog.get(id).map(|m| m.format),
170        Some(ModelFormat::NemotronPackage)
171    )
172}
173
174/// Whether a model is present on disk and usable. For Whisper GGUF models this
175/// is a plain file-exists check; for the Nemotron package it verifies the
176/// unpacked directory actually contains the ONNX artifacts the engine loads,
177/// so a half-unpacked or empty directory is not mistaken for a ready model.
178pub fn is_model_available(id: &str) -> bool {
179    let Ok(path) = model_path(id) else {
180        return false;
181    };
182    if is_nemotron_model(id) {
183        path.join("encoder.onnx").exists() && path.join("decoder_joint.onnx").exists()
184    } else {
185        path.exists()
186    }
187}
188
189pub fn verify_sha256(path: &Path, expected_hex: &str) -> anyhow::Result<()> {
190    let data = std::fs::read(path).with_context(|| format!("read {}", path.display()))?;
191    let mut h = Sha256::new();
192    h.update(&data);
193    let got = hex::encode(h.finalize());
194    if got != expected_hex {
195        return Err(anyhow!(
196            "sha256 mismatch for {}: got {} expected {}",
197            path.display(),
198            got,
199            expected_hex
200        ));
201    }
202    Ok(())
203}
204
205pub fn download_with_verify(
206    url: &str,
207    dest: &Path,
208    expected_sha: &str,
209    progress: &mut dyn FnMut(u64, Option<u64>),
210) -> anyhow::Result<()> {
211    if let Some(parent) = dest.parent() {
212        std::fs::create_dir_all(parent)?;
213    }
214    let resp = reqwest::blocking::get(url)?.error_for_status()?;
215    let total = resp.content_length();
216    let mut downloaded = 0u64;
217    let mut out = std::fs::File::create(dest)?;
218    let mut src = resp;
219    let mut buf = [0u8; 1 << 16];
220    loop {
221        let n = std::io::Read::read(&mut src, &mut buf)?;
222        if n == 0 {
223            break;
224        }
225        std::io::Write::write_all(&mut out, &buf[..n])?;
226        downloaded += n as u64;
227        progress(downloaded, total);
228    }
229    drop(out);
230    if !expected_sha.starts_with("<FILL") {
231        verify_sha256(dest, expected_sha)?;
232    }
233    Ok(())
234}
235
236pub fn download_model_with_verify(
237    entry: &ModelEntry,
238    progress: &mut dyn FnMut(u64, Option<u64>),
239) -> anyhow::Result<()> {
240    match entry.format {
241        ModelFormat::WhisperGguf => {
242            download_with_verify(
243                entry.download_url,
244                &gguf_path(entry.id)?,
245                entry.download_sha256,
246                progress,
247            )?;
248            // Optional CoreML (ANE) encoder. Only published for some models;
249            // whisper.cpp auto-loads `{id}-encoder.mlmodelc` next to the GGUF
250            // and falls back to Metal/CPU when it is absent, so a missing or
251            // failed CoreML package never blocks the GGUF from working.
252            if let (Some(url), Some(sha)) = (entry.coreml_url, entry.coreml_sha256) {
253                let zip_path = cache_dir()?.join(format!("{}-encoder.mlmodelc.zip", entry.id));
254                download_with_verify(url, &zip_path, sha, progress)?;
255                progress(0, None);
256                unpack_coreml_zip(&zip_path, &coreml_path(entry.id)?)?;
257                progress(1, Some(1));
258            }
259            Ok(())
260        }
261        ModelFormat::NemotronPackage => {
262            let archive_path = cache_dir()?.join(format!("{}.tar.gz", entry.id));
263            download_with_verify(
264                entry.download_url,
265                &archive_path,
266                entry.download_sha256,
267                progress,
268            )?;
269            progress(0, None);
270            unpack_nemotron_archive(&archive_path, &nemotron_model_dir(entry.id)?)?;
271            progress(1, Some(1));
272            Ok(())
273        }
274    }
275}
276
277fn unpack_nemotron_archive(archive_path: &Path, dest_dir: &Path) -> anyhow::Result<()> {
278    let tmp_dir = dest_dir.with_extension("tmp");
279    if tmp_dir.exists() {
280        std::fs::remove_dir_all(&tmp_dir)?;
281    }
282    if dest_dir.exists() {
283        std::fs::remove_dir_all(dest_dir)?;
284    }
285    std::fs::create_dir_all(&tmp_dir)?;
286
287    let archive = std::fs::File::open(archive_path)?;
288    let decoder = GzDecoder::new(archive);
289    Archive::new(decoder).unpack(&tmp_dir)?;
290
291    let model_dir = locate_nemotron_dir(&tmp_dir).ok_or_else(|| {
292        anyhow!("Nemotron package did not contain encoder.onnx and decoder_joint.onnx")
293    })?;
294    if let Some(parent) = dest_dir.parent() {
295        std::fs::create_dir_all(parent)?;
296    }
297    std::fs::rename(model_dir, dest_dir)?;
298    let _ = std::fs::remove_dir_all(tmp_dir);
299    Ok(())
300}
301
302fn locate_nemotron_dir(root: &Path) -> Option<PathBuf> {
303    let mut stack = vec![root.to_path_buf()];
304    while let Some(dir) = stack.pop() {
305        if dir.join("encoder.onnx").exists() && dir.join("decoder_joint.onnx").exists() {
306            return Some(dir);
307        }
308        let entries = std::fs::read_dir(&dir).ok()?;
309        for entry in entries.flatten() {
310            if entry.file_type().ok()?.is_dir() {
311                stack.push(entry.path());
312            }
313        }
314    }
315    None
316}
317
318/// Extract a `*-encoder.mlmodelc.zip` and install the encoder directory at
319/// `dest_dir` (e.g. `<cache>/base.en-encoder.mlmodelc`). The zip's top-level
320/// directory is named `ggml-<id>-encoder.mlmodelc`, so we locate it by suffix
321/// and rename, dropping the `ggml-` prefix to match whisper.cpp's lookup path.
322fn unpack_coreml_zip(zip_path: &Path, dest_dir: &Path) -> anyhow::Result<()> {
323    let tmp_dir = dest_dir.with_extension("ziptmp");
324    if tmp_dir.exists() {
325        std::fs::remove_dir_all(&tmp_dir)?;
326    }
327    if dest_dir.exists() {
328        std::fs::remove_dir_all(dest_dir)?;
329    }
330    std::fs::create_dir_all(&tmp_dir)?;
331
332    let file = std::fs::File::open(zip_path)?;
333    let mut archive = zip::ZipArchive::new(file)
334        .with_context(|| format!("open coreml zip {}", zip_path.display()))?;
335    archive.extract(&tmp_dir)?;
336
337    let model_dir = locate_coreml_dir(&tmp_dir).ok_or_else(|| {
338        anyhow!("CoreML package did not contain a *-encoder.mlmodelc directory")
339    })?;
340    if let Some(parent) = dest_dir.parent() {
341        std::fs::create_dir_all(parent)?;
342    }
343    std::fs::rename(model_dir, dest_dir)?;
344    let _ = std::fs::remove_dir_all(tmp_dir);
345    Ok(())
346}
347
348fn locate_coreml_dir(root: &Path) -> Option<PathBuf> {
349    let mut stack = vec![root.to_path_buf()];
350    while let Some(dir) = stack.pop() {
351        let entries = std::fs::read_dir(&dir).ok()?;
352        for entry in entries.flatten() {
353            if !entry.file_type().ok()?.is_dir() {
354                continue;
355            }
356            let path = entry.path();
357            let name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
358            // Skip the AppleDouble metadata folder some macOS zips carry.
359            if name == "__MACOSX" {
360                continue;
361            }
362            if name.ends_with("-encoder.mlmodelc") {
363                return Some(path);
364            }
365            stack.push(path);
366        }
367    }
368    None
369}
370
371#[cfg(test)]
372mod tests {
373    use super::*;
374
375    #[test]
376    fn catalog_has_default_and_required_ids() {
377        let catalog = Catalog::builtin();
378        assert!(catalog.all().iter().any(|m| m.is_default));
379        assert!(catalog.get("distil-small.en").is_some());
380        assert!(catalog.get("large-v3-turbo").is_some());
381        assert!(catalog.get("nemotron-3.5-asr-streaming-0.6b").is_some());
382    }
383
384    #[test]
385    fn nemotron_uses_directory_model_path() {
386        let path = model_path("nemotron-3.5-asr-streaming-0.6b").unwrap();
387        assert!(path.ends_with("nemotron-3.5-asr-streaming-0.6b"));
388        assert!(is_nemotron_model("nemotron-3.5-asr-streaming-0.6b"));
389    }
390
391    #[test]
392    fn locate_nemotron_dir_uses_parakeet_rs_layout() {
393        let tmp = tempfile::tempdir().unwrap();
394        let nested = tmp.path().join("nested").join("model");
395        std::fs::create_dir_all(&nested).unwrap();
396        std::fs::write(nested.join("encoder.onnx"), b"encoder").unwrap();
397        std::fs::write(nested.join("decoder_joint.onnx"), b"decoder").unwrap();
398
399        assert_eq!(locate_nemotron_dir(tmp.path()).unwrap(), nested);
400    }
401
402    #[test]
403    fn coreml_path_matches_whisper_cpp_lookup() {
404        // whisper.cpp derives the CoreML dir from `{id}.gguf` by stripping the
405        // extension and appending `-encoder.mlmodelc`. The gguf and coreml
406        // paths must share a parent and the exact `{id}-encoder.mlmodelc` name.
407        let gguf = gguf_path("base.en").unwrap();
408        let coreml = coreml_path("base.en").unwrap();
409        assert_eq!(gguf.parent(), coreml.parent());
410        assert!(coreml.ends_with("base.en-encoder.mlmodelc"));
411    }
412
413    #[test]
414    fn coreml_models_have_paired_url_and_sha() {
415        let catalog = Catalog::builtin();
416        for m in catalog.all() {
417            assert_eq!(
418                m.coreml_url.is_some(),
419                m.coreml_sha256.is_some(),
420                "{} has a half-populated CoreML pair",
421                m.id
422            );
423        }
424        // The three Whisper models with published whisper.cpp CoreML encoders.
425        for id in ["base.en", "tiny.en", "large-v3-turbo"] {
426            assert!(catalog.get(id).unwrap().coreml_url.is_some(), "{id} missing CoreML url");
427        }
428    }
429
430    #[test]
431    fn locate_coreml_dir_finds_encoder_and_skips_macosx() {
432        let tmp = tempfile::tempdir().unwrap();
433        // Simulate the zip layout: an AppleDouble folder plus the encoder dir.
434        std::fs::create_dir_all(tmp.path().join("__MACOSX")).unwrap();
435        let enc = tmp.path().join("ggml-base.en-encoder.mlmodelc");
436        std::fs::create_dir_all(enc.join("weights")).unwrap();
437        std::fs::write(enc.join("weights").join("weight.bin"), b"w").unwrap();
438
439        assert_eq!(locate_coreml_dir(tmp.path()).unwrap(), enc);
440    }
441
442    #[test]
443    fn sha256_verify_detects_mismatch() {
444        let tmp = tempfile::NamedTempFile::new().unwrap();
445        std::fs::write(tmp.path(), b"hello").unwrap();
446        let wrong = "0".repeat(64);
447        assert!(verify_sha256(tmp.path(), &wrong).is_err());
448    }
449
450    #[test]
451    fn validate_local_language_rejects_english_model_for_russian() {
452        let err = validate_local_language("tiny.en", "ru").unwrap_err();
453        assert!(err.contains("tiny.en"));
454        assert!(err.contains("ru"));
455    }
456
457    #[test]
458    fn validate_local_language_accepts_multilingual_for_russian() {
459        assert!(validate_local_language("distil-large-v3", "ru").is_ok());
460        assert!(validate_local_language("large-v3-turbo", "pl").is_ok());
461        assert!(validate_local_language("nemotron-3.5-asr-streaming-0.6b", "pl").is_ok());
462    }
463
464    #[test]
465    fn validate_local_language_passes_through_english_and_auto() {
466        assert!(validate_local_language("tiny.en", "en").is_ok());
467        assert!(validate_local_language("tiny.en", "auto").is_ok());
468    }
469
470    #[test]
471    fn validate_local_language_rejects_unknown_model_id() {
472        let err = validate_local_language("custom-user-model", "ru").unwrap_err();
473        assert!(err.contains("custom-user-model"));
474        assert!(err.contains("not supported"));
475    }
476}
voice_bird_cli/transcription/models.rs

voice_bird_cli/transcription/
models.rs
