voice_bird_cli/transcription/
models.rs1use 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
81pub 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 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 assert!(validate_local_language("custom-user-model", "ru").is_ok());
214 }
215}