Skip to main content

voice_bird_cli/transcription/
models.rs

1use std::path::{Path, PathBuf};
2
3use anyhow::{anyhow, Context};
4use sha2::{Digest, Sha256};
5
6#[derive(Debug, Clone)]
7pub struct ModelEntry {
8    pub id: &'static str,
9    pub size_mb: u32,
10    pub language: &'static str,
11    pub gguf_url: &'static str,
12    pub gguf_sha256: &'static str,
13    pub coreml_url: Option<&'static str>,
14    pub coreml_sha256: Option<&'static str>,
15    pub is_default: bool,
16}
17
18pub struct Catalog(Vec<ModelEntry>);
19
20impl Catalog {
21    pub fn builtin() -> Self {
22        Catalog(vec![
23            ModelEntry {
24                id: "distil-small.en", size_mb: 250, language: "en",
25                gguf_url: "https://huggingface.co/distil-whisper/distil-small.en/resolve/main/ggml-distil-small.en.bin",
26                gguf_sha256: "<FILL>",
27                coreml_url: Some("https://huggingface.co/distil-whisper/distil-small.en/resolve/main/coreml-distil-small.en.zip"),
28                coreml_sha256: Some("<FILL>"),
29                is_default: true,
30            },
31            ModelEntry {
32                id: "distil-large-v3", size_mb: 1_500, language: "multi",
33                gguf_url: "https://huggingface.co/distil-whisper/distil-large-v3/resolve/main/ggml-distil-large-v3.bin",
34                gguf_sha256: "<FILL>",
35                coreml_url: Some("https://huggingface.co/distil-whisper/distil-large-v3/resolve/main/coreml-distil-large-v3.zip"),
36                coreml_sha256: Some("<FILL>"),
37                is_default: false,
38            },
39            ModelEntry {
40                id: "large-v3-turbo", size_mb: 1_600, language: "multi",
41                gguf_url: "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-large-v3-turbo.bin",
42                gguf_sha256: "<FILL>",
43                coreml_url: Some("https://huggingface.co/argmaxinc/whisperkit-coreml/resolve/main/openai_whisper-large-v3-turbo.zip"),
44                coreml_sha256: Some("<FILL>"),
45                is_default: false,
46            },
47            ModelEntry {
48                id: "base.en", size_mb: 150, language: "en",
49                gguf_url: "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-base.en.bin",
50                gguf_sha256: "<FILL>",
51                coreml_url: None, coreml_sha256: None,
52                is_default: false,
53            },
54            ModelEntry {
55                id: "tiny.en", size_mb: 75, language: "en",
56                gguf_url: "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-tiny.en.bin",
57                gguf_sha256: "<FILL>",
58                coreml_url: None, coreml_sha256: None,
59                is_default: false,
60            },
61        ])
62    }
63
64    pub fn get(&self, id: &str) -> Option<&ModelEntry> {
65        self.0.iter().find(|m| m.id == id)
66    }
67
68    pub fn default_id(&self) -> &'static str {
69        self.0
70            .iter()
71            .find(|m| m.is_default)
72            .map(|m| m.id)
73            .unwrap_or("distil-small.en")
74    }
75
76    pub fn all(&self) -> &[ModelEntry] {
77        &self.0
78    }
79}
80
81/// Reject combinations where a non-English language is requested but the
82/// configured local Whisper model is English-only. Without this, users
83/// selecting Russian/Polish/etc with a `.en` ggml file silently get
84/// gibberish English output. Returns `Ok(())` for known good combinations
85/// and for anything we can't verify (missing catalog entry, "auto" /
86/// English / empty language).
87pub fn validate_local_language(model_id: &str, language: &str) -> Result<(), String> {
88    let lang = language.trim();
89    if lang.is_empty() || lang == "en" || lang == "auto" {
90        return Ok(());
91    }
92    let catalog = Catalog::builtin();
93    let Some(entry) = catalog.get(model_id) else {
94        // Unknown model id: don't block — let the engine surface its own error.
95        return Ok(());
96    };
97    if entry.language == "en" {
98        return Err(format!(
99            "Model '{}' is English-only; pick distil-large-v3 or large-v3-turbo for {}.",
100            entry.id, lang
101        ));
102    }
103    Ok(())
104}
105
106pub fn cache_dir() -> anyhow::Result<PathBuf> {
107    let base = dirs::cache_dir().ok_or_else(|| anyhow!("no cache dir"))?;
108    Ok(base.join("voice-bird").join("models"))
109}
110
111pub fn gguf_path(id: &str) -> anyhow::Result<PathBuf> {
112    Ok(cache_dir()?.join(format!("{id}.gguf")))
113}
114
115pub fn verify_sha256(path: &Path, expected_hex: &str) -> anyhow::Result<()> {
116    let data = std::fs::read(path).with_context(|| format!("read {}", path.display()))?;
117    let mut h = Sha256::new();
118    h.update(&data);
119    let got = hex::encode(h.finalize());
120    if got != expected_hex {
121        return Err(anyhow!(
122            "sha256 mismatch for {}: got {} expected {}",
123            path.display(),
124            got,
125            expected_hex
126        ));
127    }
128    Ok(())
129}
130
131pub fn download_with_verify(
132    url: &str,
133    dest: &Path,
134    expected_sha: &str,
135    progress: &mut dyn FnMut(u64, Option<u64>),
136) -> anyhow::Result<()> {
137    if let Some(parent) = dest.parent() {
138        std::fs::create_dir_all(parent)?;
139    }
140    let resp = reqwest::blocking::get(url)?.error_for_status()?;
141    let total = resp.content_length();
142    let mut downloaded = 0u64;
143    let mut out = std::fs::File::create(dest)?;
144    let mut src = resp;
145    let mut buf = [0u8; 1 << 16];
146    loop {
147        let n = std::io::Read::read(&mut src, &mut buf)?;
148        if n == 0 {
149            break;
150        }
151        std::io::Write::write_all(&mut out, &buf[..n])?;
152        downloaded += n as u64;
153        progress(downloaded, total);
154    }
155    drop(out);
156    if !expected_sha.starts_with("<FILL") {
157        verify_sha256(dest, expected_sha)?;
158    }
159    Ok(())
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165
166    #[test]
167    fn catalog_has_default_and_required_ids() {
168        let catalog = Catalog::builtin();
169        let default = catalog.default_id();
170        assert_eq!(default, "distil-small.en");
171        for id in [
172            "distil-small.en",
173            "distil-large-v3",
174            "large-v3-turbo",
175            "base.en",
176            "tiny.en",
177        ] {
178            assert!(catalog.get(id).is_some(), "missing {id}");
179        }
180    }
181
182    #[test]
183    fn sha256_verify_detects_mismatch() {
184        let tmp = tempfile::NamedTempFile::new().unwrap();
185        std::fs::write(tmp.path(), b"hello").unwrap();
186        let wrong = "0".repeat(64);
187        assert!(verify_sha256(tmp.path(), &wrong).is_err());
188    }
189
190    #[test]
191    fn validate_local_language_rejects_english_model_for_russian() {
192        let err = validate_local_language("tiny.en", "ru").unwrap_err();
193        assert!(err.contains("tiny.en"));
194        assert!(err.contains("ru"));
195    }
196
197    #[test]
198    fn validate_local_language_accepts_multilingual_for_russian() {
199        assert!(validate_local_language("distil-large-v3", "ru").is_ok());
200        assert!(validate_local_language("large-v3-turbo", "pl").is_ok());
201    }
202
203    #[test]
204    fn validate_local_language_passes_through_english_and_auto() {
205        assert!(validate_local_language("tiny.en", "en").is_ok());
206        assert!(validate_local_language("tiny.en", "auto").is_ok());
207        assert!(validate_local_language("tiny.en", "").is_ok());
208    }
209
210    #[test]
211    fn validate_local_language_ignores_unknown_model_id() {
212        // Unknown ids slip through — engine layer surfaces its own error.
213        assert!(validate_local_language("custom-user-model", "ru").is_ok());
214    }
215}