From 45a473dc9fa8fe3ae67fc5314039eaf81f478643 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Tue, 9 Sep 2025 18:33:57 -0700 Subject: [PATCH] Parakeet V3 (#111) * wip: model manager support for parakeet * working parakeet??git status * Update LanguageSelector.tsx * Use the published version of the crate. * move word correction into audio toolkit --- src-tauri/Cargo.lock | 146 ++++++--- src-tauri/Cargo.toml | 12 +- src-tauri/src/audio_toolkit/mod.rs | 2 + src-tauri/src/audio_toolkit/text.rs | 176 ++++++++++ src-tauri/src/managers/model.rs | 202 ++++++++++-- src-tauri/src/managers/transcription.rs | 305 ++++++------------ .../model-selector/ModelSelector.tsx | 63 +++- src/components/settings/LanguageSelector.tsx | 36 ++- src/hooks/useModels.ts | 53 ++- src/lib/types.ts | 3 + 10 files changed, 705 insertions(+), 293 deletions(-) create mode 100644 src-tauri/src/audio_toolkit/text.rs diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index 064bf93..73cc2ab 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -597,7 +597,7 @@ dependencies = [ "semver", "serde", "serde_json", - "thiserror 2.0.11", + "thiserror 2.0.16", ] [[package]] @@ -1400,6 +1400,19 @@ dependencies = [ "regex", ] +[[package]] +name = "env_logger" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cd405aab171cb85d6735e5c8d9db038c17d3ca007a4d2c25f337935c3d90580" +dependencies = [ + "humantime", + "is-terminal", + "log", + "regex", + "termcolor", +] + [[package]] name = "env_logger" version = "0.11.6" @@ -2026,7 +2039,7 @@ dependencies = [ "objc2-app-kit 0.3.1", "once_cell", "serde", - "thiserror 2.0.11", + "thiserror 2.0.16", "windows-sys 0.59.0", "x11-dl", ] @@ -2121,7 +2134,8 @@ dependencies = [ "chrono", "cpal", "enigo", - "env_logger", + "env_logger 0.11.6", + "flate2", "futures-util", "hound", "log", @@ -2136,6 +2150,7 @@ dependencies = [ "serde", "serde_json", "strsim", + "tar", "tauri", "tauri-build", "tauri-plugin-autostart", @@ -2151,8 +2166,8 @@ dependencies = [ "tauri-plugin-store", "tauri-plugin-updater", "tokio", + "transcribe-rs", "vad-rs", - "whisper-rs", ] [[package]] @@ -2217,6 +2232,12 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fbf6a919d6cf397374f7dfeeea91d974c7c0a7221d0d0f4f20d859d329e53fcc" +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + [[package]] name = "hex" version = "0.4.3" @@ -2704,6 +2725,17 @@ dependencies = [ "once_cell", ] +[[package]] +name = "is-terminal" +version = "0.4.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e04d7f318608d35d4b61ddd75cbdaee86b023ebe2bd5a66ee0915f0bf93095a9" +dependencies = [ + "hermit-abi 0.5.2", + "libc", + "windows-sys 0.59.0", +] + [[package]] name = "is-wsl" version = "0.4.0" @@ -2973,9 +3005,9 @@ dependencies = [ [[package]] name = "log" -version = "0.4.25" +version = "0.4.28" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04cbf5b083de1c7e0222a7a51dbfdba1cbe1c6ab0b15e29fff3f6c077fd9cd9f" +checksum = "34080505efa8e45a4b816c349525ebe327ceaa8559756f0356cba97ef3bf7432" [[package]] name = "mac" @@ -3154,7 +3186,7 @@ dependencies = [ "once_cell", "png", "serde", - "thiserror 2.0.11", + "thiserror 2.0.16", "windows-sys 0.59.0", ] @@ -3711,9 +3743,9 @@ dependencies = [ [[package]] name = "once_cell" -version = "1.20.2" +version = "1.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1261fe7e33c73b354eab43b1273a57c8f967d0391e80353e51f764ac02cf6775" +checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" [[package]] name = "open" @@ -3834,7 +3866,7 @@ dependencies = [ "objc2-osa-kit", "serde", "serde_json", - "thiserror 2.0.11", + "thiserror 2.0.16", ] [[package]] @@ -4130,7 +4162,7 @@ checksum = "a604568c3202727d1507653cb121dbd627a58684eb09a820fd746bee38b4442f" dependencies = [ "cfg-if", "concurrent-queue", - "hermit-abi", + "hermit-abi 0.4.0", "pin-project-lite", "rustix 0.38.44", "tracing", @@ -4281,7 +4313,7 @@ dependencies = [ "rustc-hash 2.1.0", "rustls", "socket2", - "thiserror 2.0.11", + "thiserror 2.0.16", "tokio", "tracing", ] @@ -4300,7 +4332,7 @@ dependencies = [ "rustls", "rustls-pki-types", "slab", - "thiserror 2.0.11", + "thiserror 2.0.16", "tinyvec", "tracing", "web-time", @@ -4483,14 +4515,14 @@ checksum = "dd6f9d3d47bdd2ad6945c5015a226ec6155d0bcdfd8f7cd29f86b71f8de99d2b" dependencies = [ "getrandom 0.2.15", "libredox", - "thiserror 2.0.11", + "thiserror 2.0.16", ] [[package]] name = "regex" -version = "1.11.1" +version = "1.11.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b544ef1b4eac5dc2db33ea63606ae9ffcfac26c1416a2806ae0bf5f56b201191" +checksum = "23d7fd106d8c02486a8d64e778353d1cffe08ce79ac2e82f540c86d0facf6912" dependencies = [ "aho-corasick", "memchr", @@ -5296,7 +5328,7 @@ dependencies = [ "serde_json", "sha2", "smallvec 1.13.2", - "thiserror 2.0.11", + "thiserror 2.0.16", "time", "tokio", "tokio-stream", @@ -5379,7 +5411,7 @@ dependencies = [ "smallvec 1.13.2", "sqlx-core", "stringprep", - "thiserror 2.0.11", + "thiserror 2.0.16", "time", "tracing", "whoami", @@ -5417,7 +5449,7 @@ dependencies = [ "smallvec 1.13.2", "sqlx-core", "stringprep", - "thiserror 2.0.11", + "thiserror 2.0.16", "time", "tracing", "whoami", @@ -5442,7 +5474,7 @@ dependencies = [ "serde", "serde_urlencoded", "sqlx-core", - "thiserror 2.0.11", + "thiserror 2.0.16", "time", "tracing", "url", @@ -5891,7 +5923,7 @@ dependencies = [ "tauri-runtime", "tauri-runtime-wry", "tauri-utils", - "thiserror 2.0.11", + "thiserror 2.0.16", "tokio", "tray-icon", "url", @@ -5944,7 +5976,7 @@ dependencies = [ "sha2", "syn 2.0.96", "tauri-utils", - "thiserror 2.0.11", + "thiserror 2.0.16", "time", "url", "uuid", @@ -5993,7 +6025,7 @@ dependencies = [ "serde_json", "tauri", "tauri-plugin", - "thiserror 2.0.11", + "thiserror 2.0.16", ] [[package]] @@ -6008,7 +6040,7 @@ dependencies = [ "serde_json", "tauri", "tauri-plugin", - "thiserror 2.0.11", + "thiserror 2.0.16", ] [[package]] @@ -6028,7 +6060,7 @@ dependencies = [ "tauri", "tauri-plugin", "tauri-utils", - "thiserror 2.0.11", + "thiserror 2.0.16", "toml", "url", ] @@ -6045,7 +6077,7 @@ dependencies = [ "serde_json", "tauri", "tauri-plugin", - "thiserror 2.0.11", + "thiserror 2.0.16", ] [[package]] @@ -6058,7 +6090,7 @@ dependencies = [ "serde", "tauri", "tauri-plugin", - "thiserror 2.0.11", + "thiserror 2.0.16", ] [[package]] @@ -6077,7 +6109,7 @@ dependencies = [ "serde_json", "tauri", "tauri-plugin", - "thiserror 2.0.11", + "thiserror 2.0.16", "url", "windows 0.58.0", "zbus", @@ -6098,7 +6130,7 @@ dependencies = [ "sys-locale", "tauri", "tauri-plugin", - "thiserror 2.0.11", + "thiserror 2.0.16", ] [[package]] @@ -6120,7 +6152,7 @@ dependencies = [ "serde", "serde_json", "tauri", - "thiserror 2.0.11", + "thiserror 2.0.16", "tracing", "windows-sys 0.60.2", "zbus", @@ -6140,7 +6172,7 @@ dependencies = [ "sqlx", "tauri", "tauri-plugin", - "thiserror 2.0.11", + "thiserror 2.0.16", "time", "tokio", ] @@ -6156,7 +6188,7 @@ dependencies = [ "serde_json", "tauri", "tauri-plugin", - "thiserror 2.0.11", + "thiserror 2.0.16", "tokio", "tracing", ] @@ -6185,7 +6217,7 @@ dependencies = [ "tauri", "tauri-plugin", "tempfile", - "thiserror 2.0.11", + "thiserror 2.0.16", "time", "tokio", "url", @@ -6210,7 +6242,7 @@ dependencies = [ "serde", "serde_json", "tauri-utils", - "thiserror 2.0.11", + "thiserror 2.0.16", "url", "windows 0.61.3", ] @@ -6272,7 +6304,7 @@ dependencies = [ "serde_json", "serde_with", "swift-rs", - "thiserror 2.0.11", + "thiserror 2.0.16", "toml", "url", "urlpattern", @@ -6316,6 +6348,15 @@ dependencies = [ "utf-8", ] +[[package]] +name = "termcolor" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06794f8f6c5c898b3275aebefa6b8a1cb24cd2c6c79397ab15774837a0bc5755" +dependencies = [ + "winapi-util", +] + [[package]] name = "thiserror" version = "1.0.69" @@ -6327,11 +6368,11 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.11" +version = "2.0.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d452f284b73e6d76dd36758a0c8684b1d5be31f92b89d07fd5822175732206fc" +checksum = "3467d614147380f2e4e374161426ff399c91084acd2363eaf549172b3d5e60c0" dependencies = [ - "thiserror-impl 2.0.11", + "thiserror-impl 2.0.16", ] [[package]] @@ -6347,9 +6388,9 @@ dependencies = [ [[package]] name = "thiserror-impl" -version = "2.0.11" +version = "2.0.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26afc1baea8a989337eeb52b6e72a039780ce45c3edfcc9c5b9d112feeb173c2" +checksum = "6c5e1be1c48b9172ee610da68fd9cd2770e7a4056cb3fc98710ee6906f0c7960" dependencies = [ "proc-macro2", "quote", @@ -6609,6 +6650,23 @@ dependencies = [ "once_cell", ] +[[package]] +name = "transcribe-rs" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eaeba5648789d05eedce5478069195639a0770dcc0338146037d4b85163df4bd" +dependencies = [ + "env_logger 0.10.2", + "hound", + "log", + "ndarray", + "once_cell", + "ort", + "regex", + "thiserror 2.0.16", + "whisper-rs", +] + [[package]] name = "transpose" version = "0.2.3" @@ -6637,7 +6695,7 @@ dependencies = [ "once_cell", "png", "serde", - "thiserror 2.0.11", + "thiserror 2.0.16", "windows-sys 0.59.0", ] @@ -7133,7 +7191,7 @@ version = "0.38.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "36695906a1b53a3bf5c4289621efedac12b73eeb0b89e7e1a89b517302d5d75c" dependencies = [ - "thiserror 2.0.11", + "thiserror 2.0.16", "windows 0.61.3", "windows-core 0.61.2", ] @@ -7861,7 +7919,7 @@ dependencies = [ "sha2", "soup3", "tao-macros", - "thiserror 2.0.11", + "thiserror 2.0.16", "url", "webkit2gtk", "webkit2gtk-sys", @@ -8124,7 +8182,7 @@ dependencies = [ "displaydoc", "indexmap 2.7.1", "memchr", - "thiserror 2.0.11", + "thiserror 2.0.16", ] [[package]] diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index ec7d9d6..3830237 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -62,15 +62,9 @@ natural = "0.5.0" chrono = "0.4" rusqlite = { version = "0.32", features = ["bundled"] } tauri-plugin-sql = { version = "2", features = ["sqlite"] } - -[target.'cfg(target_os = "macos")'.dependencies] -whisper-rs = { version = "0.13.2", features = ["metal"] } - -[target.'cfg(target_os = "windows")'.dependencies] -whisper-rs = { version = "0.13.2", features = ["vulkan"] } - -[target.'cfg(target_os = "linux")'.dependencies] -whisper-rs = { version = "0.13.2", features = ["vulkan"] } +tar = "0.4.44" +flate2 = "1.0" +transcribe-rs = "0.1.0" [target.'cfg(not(any(target_os = "android", target_os = "ios")))'.dependencies] tauri-plugin-autostart = "2" diff --git a/src-tauri/src/audio_toolkit/mod.rs b/src-tauri/src/audio_toolkit/mod.rs index 0f88bdc..54c648c 100644 --- a/src-tauri/src/audio_toolkit/mod.rs +++ b/src-tauri/src/audio_toolkit/mod.rs @@ -1,8 +1,10 @@ pub mod audio; pub mod constants; +pub mod text; pub mod vad; pub use audio::{ list_input_devices, list_output_devices, save_wav_file, AudioRecorder, CpalDeviceInfo, }; +pub use text::apply_custom_words; pub use vad::{SileroVad, VoiceActivityDetector}; diff --git a/src-tauri/src/audio_toolkit/text.rs b/src-tauri/src/audio_toolkit/text.rs new file mode 100644 index 0000000..6f30ae9 --- /dev/null +++ b/src-tauri/src/audio_toolkit/text.rs @@ -0,0 +1,176 @@ +use natural::phonetics::soundex; +use strsim::levenshtein; + +/// Applies custom word corrections to transcribed text using fuzzy matching +/// +/// This function corrects words in the input text by finding the best matches +/// from a list of custom words using a combination of: +/// - Levenshtein distance for string similarity +/// - Soundex phonetic matching for pronunciation similarity +/// +/// # Arguments +/// * `text` - The input text to correct +/// * `custom_words` - List of custom words to match against +/// * `threshold` - Maximum similarity score to accept (0.0 = exact match, 1.0 = any match) +/// +/// # Returns +/// The corrected text with custom words applied +pub fn apply_custom_words(text: &str, custom_words: &[String], threshold: f64) -> String { + if custom_words.is_empty() { + return text.to_string(); + } + + // Pre-compute lowercase versions to avoid repeated allocations + let custom_words_lower: Vec = custom_words.iter().map(|w| w.to_lowercase()).collect(); + + let words: Vec<&str> = text.split_whitespace().collect(); + let mut corrected_words = Vec::new(); + + for word in words { + let cleaned_word = word + .trim_matches(|c: char| !c.is_alphabetic()) + .to_lowercase(); + + if cleaned_word.is_empty() { + corrected_words.push(word.to_string()); + continue; + } + + // Skip extremely long words to avoid performance issues + if cleaned_word.len() > 50 { + corrected_words.push(word.to_string()); + continue; + } + + let mut best_match: Option<&String> = None; + let mut best_score = f64::MAX; + + for (i, custom_word_lower) in custom_words_lower.iter().enumerate() { + // Skip if lengths are too different (optimization) + let len_diff = (cleaned_word.len() as i32 - custom_word_lower.len() as i32).abs(); + if len_diff > 5 { + continue; + } + + // Calculate Levenshtein distance (normalized by length) + let levenshtein_dist = levenshtein(&cleaned_word, custom_word_lower); + let max_len = cleaned_word.len().max(custom_word_lower.len()) as f64; + let levenshtein_score = if max_len > 0.0 { + levenshtein_dist as f64 / max_len + } else { + 1.0 + }; + + // Calculate phonetic similarity using Soundex + let phonetic_match = soundex(&cleaned_word, custom_word_lower); + + // Combine scores: favor phonetic matches, but also consider string similarity + let combined_score = if phonetic_match { + levenshtein_score * 0.3 // Give significant boost to phonetic matches + } else { + levenshtein_score + }; + + // Accept if the score is good enough (configurable threshold) + if combined_score < threshold && combined_score < best_score { + best_match = Some(&custom_words[i]); + best_score = combined_score; + } + } + + if let Some(replacement) = best_match { + // Preserve the original case pattern as much as possible + let corrected = preserve_case_pattern(word, replacement); + + // Preserve punctuation from original word + let (prefix, suffix) = extract_punctuation(word); + corrected_words.push(format!("{}{}{}", prefix, corrected, suffix)); + } else { + corrected_words.push(word.to_string()); + } + } + + corrected_words.join(" ") +} + +/// Preserves the case pattern of the original word when applying a replacement +fn preserve_case_pattern(original: &str, replacement: &str) -> String { + if original.chars().all(|c| c.is_uppercase()) { + replacement.to_uppercase() + } else if original.chars().next().map_or(false, |c| c.is_uppercase()) { + let mut chars: Vec = replacement.chars().collect(); + if let Some(first_char) = chars.get_mut(0) { + *first_char = first_char.to_uppercase().next().unwrap_or(*first_char); + } + chars.into_iter().collect() + } else { + replacement.to_string() + } +} + +/// Extracts punctuation prefix and suffix from a word +fn extract_punctuation(word: &str) -> (&str, &str) { + let prefix_end = word.chars().take_while(|c| !c.is_alphabetic()).count(); + let suffix_start = word + .char_indices() + .rev() + .take_while(|(_, c)| !c.is_alphabetic()) + .count(); + + let prefix = if prefix_end > 0 { + &word[..prefix_end] + } else { + "" + }; + + let suffix = if suffix_start > 0 { + &word[word.len() - suffix_start..] + } else { + "" + }; + + (prefix, suffix) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_apply_custom_words_exact_match() { + let text = "hello world"; + let custom_words = vec!["Hello".to_string(), "World".to_string()]; + let result = apply_custom_words(text, &custom_words, 0.5); + assert_eq!(result, "Hello World"); + } + + #[test] + fn test_apply_custom_words_fuzzy_match() { + let text = "helo wrold"; + let custom_words = vec!["hello".to_string(), "world".to_string()]; + let result = apply_custom_words(text, &custom_words, 0.5); + assert_eq!(result, "hello world"); + } + + #[test] + fn test_preserve_case_pattern() { + assert_eq!(preserve_case_pattern("HELLO", "world"), "WORLD"); + assert_eq!(preserve_case_pattern("Hello", "world"), "World"); + assert_eq!(preserve_case_pattern("hello", "WORLD"), "WORLD"); + } + + #[test] + fn test_extract_punctuation() { + assert_eq!(extract_punctuation("hello"), ("", "")); + assert_eq!(extract_punctuation("!hello?"), ("!", "?")); + assert_eq!(extract_punctuation("...hello..."), ("...", "...")); + } + + #[test] + fn test_empty_custom_words() { + let text = "hello world"; + let custom_words = vec![]; + let result = apply_custom_words(text, &custom_words, 0.5); + assert_eq!(result, "hello world"); + } +} diff --git a/src-tauri/src/managers/model.rs b/src-tauri/src/managers/model.rs index 96a988c..6607a81 100644 --- a/src-tauri/src/managers/model.rs +++ b/src-tauri/src/managers/model.rs @@ -8,6 +8,15 @@ use std::io::Write; use std::path::PathBuf; use std::sync::Mutex; use tauri::{App, AppHandle, Emitter, Manager}; +use tar::Archive; +use std::fs::File; +use flate2::read::GzDecoder; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum EngineType { + Whisper, + Parakeet, +} #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ModelInfo { @@ -20,6 +29,8 @@ pub struct ModelInfo { pub is_downloaded: bool, pub is_downloading: bool, pub partial_size: u64, + pub is_directory: bool, + pub engine_type: EngineType, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -65,6 +76,8 @@ impl ModelManager { is_downloaded: false, is_downloading: false, partial_size: 0, + is_directory: false, + engine_type: EngineType::Whisper, }, ); @@ -81,6 +94,8 @@ impl ModelManager { is_downloaded: false, is_downloading: false, partial_size: 0, + is_directory: false, + engine_type: EngineType::Whisper, }, ); @@ -96,6 +111,8 @@ impl ModelManager { is_downloaded: false, is_downloading: false, partial_size: 0, + is_directory: false, + engine_type: EngineType::Whisper, }, ); @@ -111,6 +128,26 @@ impl ModelManager { is_downloaded: false, is_downloading: false, partial_size: 0, + is_directory: false, + engine_type: EngineType::Whisper, + }, + ); + + // Add NVIDIA Parakeet model (directory-based) + available_models.insert( + "parakeet-tdt-0.6b-v3".to_string(), + ModelInfo { + id: "parakeet-tdt-0.6b-v3".to_string(), + name: "Parakeet V3".to_string(), + description: "The fastest and most accurate model".to_string(), + filename: "parakeet-tdt-0.6b-v3-int8".to_string(), // Directory name + url: Some("https://blob.handy.computer/parakeet-v3-int8.tar.gz".to_string()), + size_mb: 850, // Approximate size for int8 quantized model + is_downloaded: false, + is_downloading: false, + partial_size: 0, + is_directory: true, + engine_type: EngineType::Parakeet, }, ); @@ -173,17 +210,41 @@ impl ModelManager { let mut models = self.available_models.lock().unwrap(); for model in models.values_mut() { - let model_path = self.models_dir.join(&model.filename); - let partial_path = self.models_dir.join(format!("{}.partial", &model.filename)); + if model.is_directory { + // For directory-based models, check if the directory exists + let model_path = self.models_dir.join(&model.filename); + let partial_path = self.models_dir.join(format!("{}.partial", &model.filename)); + let extracting_path = self.models_dir.join(format!("{}.extracting", &model.filename)); - model.is_downloaded = model_path.exists(); - model.is_downloading = partial_path.exists(); + // Clean up any leftover .extracting directories from interrupted extractions + if extracting_path.exists() { + println!("Cleaning up interrupted extraction for model: {}", model.id); + let _ = fs::remove_dir_all(&extracting_path); + } - // Get partial file size if it exists - if partial_path.exists() { - model.partial_size = partial_path.metadata().map(|m| m.len()).unwrap_or(0); + model.is_downloaded = model_path.exists() && model_path.is_dir(); + model.is_downloading = partial_path.exists(); + + // Get partial file size if it exists (for the .tar.gz being downloaded) + if partial_path.exists() { + model.partial_size = partial_path.metadata().map(|m| m.len()).unwrap_or(0); + } else { + model.partial_size = 0; + } } else { - model.partial_size = 0; + // For file-based models (existing logic) + let model_path = self.models_dir.join(&model.filename); + let partial_path = self.models_dir.join(format!("{}.partial", &model.filename)); + + model.is_downloaded = model_path.exists(); + model.is_downloading = partial_path.exists(); + + // Get partial file size if it exists + if partial_path.exists() { + model.partial_size = partial_path.metadata().map(|m| m.len()).unwrap_or(0); + } else { + model.partial_size = 0; + } } } @@ -359,8 +420,74 @@ impl ModelManager { file.flush()?; drop(file); // Ensure file is closed before moving - // Move partial file to final location - fs::rename(&partial_path, &model_path)?; + // Handle directory-based models (extract tar.gz) vs file-based models + if model_info.is_directory { + // Emit extraction started event + let _ = self.app_handle.emit("model-extraction-started", model_id); + println!("Extracting archive for directory-based model: {}", model_id); + + // Use a temporary extraction directory to ensure atomic operations + let temp_extract_dir = self.models_dir.join(format!("{}.extracting", &model_info.filename)); + let final_model_dir = self.models_dir.join(&model_info.filename); + + // Clean up any previous incomplete extraction + if temp_extract_dir.exists() { + let _ = fs::remove_dir_all(&temp_extract_dir); + } + + // Create temporary extraction directory + fs::create_dir_all(&temp_extract_dir)?; + + // Open the downloaded tar.gz file + let tar_gz = File::open(&partial_path)?; + let tar = GzDecoder::new(tar_gz); + let mut archive = Archive::new(tar); + + // Extract to the temporary directory first + archive.unpack(&temp_extract_dir).map_err(|e| { + let error_msg = format!("Failed to extract archive: {}", e); + // Clean up failed extraction + let _ = fs::remove_dir_all(&temp_extract_dir); + let _ = self.app_handle.emit("model-extraction-failed", &serde_json::json!({ + "model_id": model_id, + "error": error_msg + })); + anyhow::anyhow!(error_msg) + })?; + + // Find the actual extracted directory (archive might have a nested structure) + let extracted_dirs: Vec<_> = fs::read_dir(&temp_extract_dir)? + .filter_map(|entry| entry.ok()) + .filter(|entry| entry.file_type().map(|ft| ft.is_dir()).unwrap_or(false)) + .collect(); + + if extracted_dirs.len() == 1 { + // Single directory extracted, move it to the final location + let source_dir = extracted_dirs[0].path(); + if final_model_dir.exists() { + fs::remove_dir_all(&final_model_dir)?; + } + fs::rename(&source_dir, &final_model_dir)?; + // Clean up temp directory + let _ = fs::remove_dir_all(&temp_extract_dir); + } else { + // Multiple items or no directories, rename the temp directory itself + if final_model_dir.exists() { + fs::remove_dir_all(&final_model_dir)?; + } + fs::rename(&temp_extract_dir, &final_model_dir)?; + } + + println!("Successfully extracted archive for model: {}", model_id); + // Emit extraction completed event + let _ = self.app_handle.emit("model-extraction-completed", model_id); + + // Remove the downloaded tar.gz file + let _ = fs::remove_file(&partial_path); + } else { + // Move partial file to final location for file-based models + fs::rename(&partial_path, &model_path)?; + } // Update download status { @@ -405,15 +532,28 @@ impl ModelManager { let mut deleted_something = false; - // Delete complete model file if it exists - if model_path.exists() { - println!("ModelManager: Deleting model file at: {:?}", model_path); - fs::remove_file(&model_path)?; - println!("ModelManager: Model file deleted successfully"); - deleted_something = true; + if model_info.is_directory { + // Delete complete model directory if it exists + if model_path.exists() && model_path.is_dir() { + println!( + "ModelManager: Deleting model directory at: {:?}", + model_path + ); + fs::remove_dir_all(&model_path)?; + println!("ModelManager: Model directory deleted successfully"); + deleted_something = true; + } + } else { + // Delete complete model file if it exists + if model_path.exists() { + println!("ModelManager: Deleting model file at: {:?}", model_path); + fs::remove_file(&model_path)?; + println!("ModelManager: Model file deleted successfully"); + deleted_something = true; + } } - // Delete partial file if it exists + // Delete partial file if it exists (same for both types) if partial_path.exists() { println!("ModelManager: Deleting partial file at: {:?}", partial_path); fs::remove_file(&partial_path)?; @@ -441,7 +581,7 @@ impl ModelManager { return Err(anyhow::anyhow!("Model not available: {}", model_id)); } - // Ensure we don't return partial files + // Ensure we don't return partial files/directories if model_info.is_downloading { return Err(anyhow::anyhow!( "Model is currently downloading: {}", @@ -454,14 +594,26 @@ impl ModelManager { .models_dir .join(format!("{}.partial", &model_info.filename)); - // Ensure we only return complete model files, not partial ones - if model_path.exists() && !partial_path.exists() { - Ok(model_path) + if model_info.is_directory { + // For directory-based models, ensure the directory exists and is complete + if model_path.exists() && model_path.is_dir() && !partial_path.exists() { + Ok(model_path) + } else { + Err(anyhow::anyhow!( + "Complete model directory not found: {}", + model_id + )) + } } else { - Err(anyhow::anyhow!( - "Complete model file not found: {}", - model_id - )) + // For file-based models (existing logic) + if model_path.exists() && !partial_path.exists() { + Ok(model_path) + } else { + Err(anyhow::anyhow!( + "Complete model file not found: {}", + model_id + )) + } } } diff --git a/src-tauri/src/managers/transcription.rs b/src-tauri/src/managers/transcription.rs index 93ebf2a..2d753b8 100644 --- a/src-tauri/src/managers/transcription.rs +++ b/src-tauri/src/managers/transcription.rs @@ -1,17 +1,22 @@ -use crate::managers::model::ModelManager; +use crate::audio_toolkit::apply_custom_words; +use crate::managers::model::{EngineType, ModelManager}; use crate::settings::{get_settings, ModelUnloadTimeout}; use anyhow::Result; use log::debug; -use natural::phonetics::soundex; use serde::Serialize; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; use std::thread; use std::time::{Duration, SystemTime}; -use strsim::levenshtein; use tauri::{App, AppHandle, Emitter, Manager}; -use whisper_rs::{ - FullParams, SamplingStrategy, WhisperContext, WhisperContextParameters, WhisperState, +use transcribe_rs::{ + engines::{ + parakeet::{ + ParakeetEngine, ParakeetInferenceParams, ParakeetModelParams, TimestampGranularity, + }, + whisper::{WhisperEngine, WhisperInferenceParams}, + }, + TranscriptionEngine, }; #[derive(Clone, Debug, Serialize)] @@ -22,10 +27,14 @@ pub struct ModelStateEvent { pub error: Option, } +enum LoadedEngine { + Whisper(WhisperEngine), + Parakeet(ParakeetEngine), +} + #[derive(Clone)] pub struct TranscriptionManager { - state: Arc>>, - context: Arc>>, + engine: Arc>>, model_manager: Arc, app_handle: AppHandle, current_model_id: Arc>>, @@ -34,121 +43,12 @@ pub struct TranscriptionManager { watcher_handle: Arc>>>, } -fn apply_custom_words(text: &str, custom_words: &[String], threshold: f64) -> String { - if custom_words.is_empty() { - return text.to_string(); - } - - // Pre-compute lowercase versions to avoid repeated allocations - let custom_words_lower: Vec = custom_words.iter().map(|w| w.to_lowercase()).collect(); - - let words: Vec<&str> = text.split_whitespace().collect(); - let mut corrected_words = Vec::new(); - - for word in words { - let cleaned_word = word - .trim_matches(|c: char| !c.is_alphabetic()) - .to_lowercase(); - - if cleaned_word.is_empty() { - corrected_words.push(word.to_string()); - continue; - } - - // Skip extremely long words to avoid performance issues - if cleaned_word.len() > 50 { - corrected_words.push(word.to_string()); - continue; - } - - let mut best_match: Option<&String> = None; - let mut best_score = f64::MAX; - - for (i, custom_word_lower) in custom_words_lower.iter().enumerate() { - // Skip if lengths are too different (optimization) - let len_diff = (cleaned_word.len() as i32 - custom_word_lower.len() as i32).abs(); - if len_diff > 5 { - continue; - } - - // Calculate Levenshtein distance (normalized by length) - let levenshtein_dist = levenshtein(&cleaned_word, custom_word_lower); - let max_len = cleaned_word.len().max(custom_word_lower.len()) as f64; - let levenshtein_score = if max_len > 0.0 { - levenshtein_dist as f64 / max_len - } else { - 1.0 - }; - - // Calculate phonetic similarity using Soundex - let phonetic_match = soundex(&cleaned_word, custom_word_lower); - - // Combine scores: favor phonetic matches, but also consider string similarity - let combined_score = if phonetic_match { - levenshtein_score * 0.3 // Give significant boost to phonetic matches - } else { - levenshtein_score - }; - - // Accept if the score is good enough (configurable threshold) - if combined_score < threshold && combined_score < best_score { - best_match = Some(&custom_words[i]); - best_score = combined_score; - } - } - - if let Some(replacement) = best_match { - // Preserve the original case pattern as much as possible - let corrected = if word.chars().all(|c| c.is_uppercase()) { - replacement.to_uppercase() - } else if word.chars().next().map_or(false, |c| c.is_uppercase()) { - let mut chars: Vec = replacement.chars().collect(); - if let Some(first_char) = chars.get_mut(0) { - *first_char = first_char.to_uppercase().next().unwrap_or(*first_char); - } - chars.into_iter().collect() - } else { - replacement.clone() - }; - - // Preserve punctuation from original word - optimized version - let prefix_end = word.chars().take_while(|c| !c.is_alphabetic()).count(); - let suffix_start = word - .char_indices() - .rev() - .take_while(|(_, c)| !c.is_alphabetic()) - .count(); - - let original_prefix = if prefix_end > 0 { - &word[..prefix_end] - } else { - "" - }; - let original_suffix = if suffix_start > 0 { - &word[word.len() - suffix_start..] - } else { - "" - }; - - corrected_words.push(format!( - "{}{}{}", - original_prefix, corrected, original_suffix - )); - } else { - corrected_words.push(word.to_string()); - } - } - - corrected_words.join(" ") -} - impl TranscriptionManager { pub fn new(app: &App, model_manager: Arc) -> Result { let app_handle = app.app_handle().clone(); let manager = Self { - state: Arc::new(Mutex::new(None)), - context: Arc::new(Mutex::new(None)), + engine: Arc::new(Mutex::new(None)), model_manager, app_handle: app_handle.clone(), current_model_id: Arc::new(Mutex::new(None)), @@ -230,8 +130,8 @@ impl TranscriptionManager { } pub fn is_model_loaded(&self) -> bool { - let state = self.state.lock().unwrap(); - state.is_some() + let engine = self.engine.lock().unwrap(); + engine.is_some() } pub fn unload_model(&self) -> Result<()> { @@ -239,12 +139,14 @@ impl TranscriptionManager { debug!("Starting to unload model"); { - let mut state = self.state.lock().unwrap(); - *state = None; // Dropping state frees GPU/CPU memory - } - { - let mut context = self.context.lock().unwrap(); - *context = None; // Dropping context frees additional memory + let mut engine = self.engine.lock().unwrap(); + if let Some(ref mut loaded_engine) = *engine { + match loaded_engine { + LoadedEngine::Whisper(ref mut whisper) => whisper.unload_model(), + LoadedEngine::Parakeet(ref mut parakeet) => parakeet.unload_model(), + } + } + *engine = None; // Drop the engine to free memory } { let mut current_model = self.current_model_id.lock().unwrap(); @@ -306,19 +208,16 @@ impl TranscriptionManager { let model_path = self.model_manager.get_model_path(model_id)?; - let path_str = model_path - .to_str() - .ok_or_else(|| anyhow::anyhow!("Invalid path for model: {}", model_id))?; - println!( - "Loading transcription model {} from: {}", - model_id, path_str + "Loading transcription model {} from: {:?}", + model_id, model_path ); - // Create new context - let context = - WhisperContext::new_with_params(path_str, WhisperContextParameters::default()) - .map_err(|e| { + // Create appropriate engine based on model type + let loaded_engine = match model_info.engine_type { + EngineType::Whisper => { + let mut engine = WhisperEngine::new(); + engine.load_model(&model_path).map_err(|e| { let error_msg = format!("Failed to load whisper model {}: {}", model_id, e); let _ = self.app_handle.emit( "model-state-changed", @@ -331,30 +230,34 @@ impl TranscriptionManager { ); anyhow::anyhow!(error_msg) })?; + LoadedEngine::Whisper(engine) + } + EngineType::Parakeet => { + let mut engine = ParakeetEngine::new(); + engine + .load_model_with_params(&model_path, ParakeetModelParams::int8()) + .map_err(|e| { + let error_msg = + format!("Failed to load parakeet model {}: {}", model_id, e); + let _ = self.app_handle.emit( + "model-state-changed", + ModelStateEvent { + event_type: "loading_failed".to_string(), + model_id: Some(model_id.to_string()), + model_name: Some(model_info.name.clone()), + error: Some(error_msg.clone()), + }, + ); + anyhow::anyhow!(error_msg) + })?; + LoadedEngine::Parakeet(engine) + } + }; - // Create new state - let state = context.create_state().map_err(|e| { - let error_msg = format!("Failed to create state for model {}: {}", model_id, e); - let _ = self.app_handle.emit( - "model-state-changed", - ModelStateEvent { - event_type: "loading_failed".to_string(), - model_id: Some(model_id.to_string()), - model_name: Some(model_info.name.clone()), - error: Some(error_msg.clone()), - }, - ); - anyhow::anyhow!(error_msg) - })?; - - // Update the current context and state + // Update the current engine and model ID { - let mut current_context = self.context.lock().unwrap(); - *current_context = Some(context); - } - { - let mut current_state = self.state.lock().unwrap(); - *current_state = Some(state); + let mut engine = self.engine.lock().unwrap(); + *engine = Some(loaded_engine); } { let mut current_model = self.current_model_id.lock().unwrap(); @@ -398,18 +301,17 @@ impl TranscriptionManager { let st = std::time::Instant::now(); - let mut result = String::new(); println!("Audio vector length: {}", audio.len()); if audio.len() == 0 { println!("Empty audio vector"); - return Ok(result); + return Ok(String::new()); } // Check if model is loaded, if not try to load it { - let state_guard = self.state.lock().unwrap(); - if state_guard.is_none() { + let engine_guard = self.engine.lock().unwrap(); + if engine_guard.is_none() { // Model not loaded, try to load the selected model from settings let settings = get_settings(&self.app_handle); println!( @@ -418,7 +320,7 @@ impl TranscriptionManager { ); // Drop the guard before calling load_model to avoid deadlock - drop(state_guard); + drop(engine_guard); // Try to load the model if let Err(e) = self.load_model(&settings.selected_model) { @@ -430,60 +332,55 @@ impl TranscriptionManager { } } - let mut state_guard = self.state.lock().unwrap(); - let state = state_guard.as_mut().ok_or_else(|| { - anyhow::anyhow!( - "Model failed to load after auto-load attempt. Please check your model settings." - ) - })?; - - // Get current settings to check translation preference + // Get current settings for configuration let settings = get_settings(&self.app_handle); - // Initialize parameters - let mut params = FullParams::new(SamplingStrategy::BeamSearch { - beam_size: 3, - patience: -1.0, - }); - let language = Some(settings.selected_language.as_str()); - params.set_language(language); - params.set_print_special(false); - params.set_print_progress(false); - params.set_print_realtime(false); - params.set_print_timestamps(false); - params.set_suppress_blank(true); - params.set_suppress_non_speech_tokens(true); - params.set_no_speech_thold(0.2); + // Perform transcription with the appropriate engine + let result = { + let mut engine_guard = self.engine.lock().unwrap(); + let engine = engine_guard.as_mut().ok_or_else(|| { + anyhow::anyhow!( + "Model failed to load after auto-load attempt. Please check your model settings." + ) + })?; - // Enable translation to English if requested - if settings.translate_to_english { - params.set_translate(true); - } + match engine { + LoadedEngine::Whisper(whisper_engine) => { + let params = WhisperInferenceParams { + language: if settings.selected_language == "auto" { + None + } else { + Some(settings.selected_language.clone()) + }, + ..Default::default() + }; - state - .full(params, &audio) - .expect("failed to convert samples"); + whisper_engine + .transcribe_samples(audio, Some(params)) + .map_err(|e| anyhow::anyhow!("Whisper transcription failed: {}", e))? + } + LoadedEngine::Parakeet(parakeet_engine) => { + let params = ParakeetInferenceParams { + timestamp_granularity: TimestampGranularity::Segment, + ..Default::default() + }; - let num_segments = state - .full_n_segments() - .expect("failed to get number of segments"); - - for i in 0..num_segments { - let segment = state - .full_get_segment_text(i) - .expect("failed to get segment"); - result.push_str(&segment); - } + parakeet_engine + .transcribe_samples(audio, Some(params)) + .map_err(|e| anyhow::anyhow!("Parakeet transcription failed: {}", e))? + } + } + }; // Apply word correction if custom words are configured let corrected_result = if !settings.custom_words.is_empty() { apply_custom_words( - &result, + &result.text, &settings.custom_words, settings.word_correction_threshold, ) } else { - result + result.text }; let et = std::time::Instant::now(); @@ -497,8 +394,6 @@ impl TranscriptionManager { // Check if we should immediately unload the model after transcription if settings.model_unload_timeout == ModelUnloadTimeout::Immediately { println!("⚡ Immediately unloading model after transcription"); - // Drop the state guard first to avoid deadlock - drop(state_guard); if let Err(e) = self.unload_model() { eprintln!("Failed to immediately unload model: {}", e); } diff --git a/src/components/model-selector/ModelSelector.tsx b/src/components/model-selector/ModelSelector.tsx index 0880923..88a7b71 100644 --- a/src/components/model-selector/ModelSelector.tsx +++ b/src/components/model-selector/ModelSelector.tsx @@ -20,7 +20,7 @@ interface DownloadProgress { percentage: number; } -type ModelStatus = "ready" | "loading" | "downloading" | "error" | "none"; +type ModelStatus = "ready" | "loading" | "downloading" | "extracting" | "error" | "none"; interface DownloadStats { startTime: number; @@ -45,6 +45,9 @@ const ModelSelector: React.FC = ({ onError }) => { const [downloadStats, setDownloadStats] = useState< Map >(new Map()); + const [extractingModels, setExtractingModels] = useState>( + new Set(), + ); const dropdownRef = useRef(null); @@ -156,6 +159,49 @@ const ModelSelector: React.FC = ({ onError }) => { }, ); + // Listen for extraction events + const extractionStartedUnlisten = listen( + "model-extraction-started", + (event) => { + const modelId = event.payload; + setExtractingModels((prev) => new Set(prev.add(modelId))); + setModelStatus("extracting"); + }, + ); + + const extractionCompletedUnlisten = listen( + "model-extraction-completed", + (event) => { + const modelId = event.payload; + setExtractingModels((prev) => { + const next = new Set(prev); + next.delete(modelId); + return next; + }); + loadModels(); // Refresh models list + + // Auto-select the newly extracted model + setTimeout(() => { + loadCurrentModel(); + handleModelSelect(modelId); + }, 500); + }, + ); + + const extractionFailedUnlisten = listen<{model_id: string, error: string}>( + "model-extraction-failed", + (event) => { + const modelId = event.payload.model_id; + setExtractingModels((prev) => { + const next = new Set(prev); + next.delete(modelId); + return next; + }); + setModelError(`Failed to extract model: ${event.payload.error}`); + setModelStatus("error"); + }, + ); + // Click outside to close dropdown const handleClickOutside = (event: MouseEvent) => { if ( @@ -173,6 +219,9 @@ const ModelSelector: React.FC = ({ onError }) => { modelStateUnlisten.then((fn) => fn()); downloadProgressUnlisten.then((fn) => fn()); downloadCompleteUnlisten.then((fn) => fn()); + extractionStartedUnlisten.then((fn) => fn()); + extractionCompletedUnlisten.then((fn) => fn()); + extractionFailedUnlisten.then((fn) => fn()); }; }, []); @@ -241,6 +290,16 @@ const ModelSelector: React.FC = ({ onError }) => { }; const getModelDisplayText = (): string => { + if (extractingModels.size > 0) { + if (extractingModels.size === 1) { + const [modelId] = Array.from(extractingModels); + const model = models.find(m => m.id === modelId); + return `Extracting ${model?.name || 'Model'}...`; + } else { + return `Extracting ${extractingModels.size} models...`; + } + } + if (modelDownloadProgress.size > 0) { if (modelDownloadProgress.size === 1) { const [progress] = Array.from(modelDownloadProgress.values()); @@ -261,6 +320,8 @@ const ModelSelector: React.FC = ({ onError }) => { return currentModel?.name || "Model Ready"; case "loading": return currentModel ? `Loading ${currentModel.name}...` : "Loading..."; + case "extracting": + return currentModel ? `Extracting ${currentModel.name}...` : "Extracting..."; case "error": return modelError || "Model Error"; case "none": diff --git a/src/components/settings/LanguageSelector.tsx b/src/components/settings/LanguageSelector.tsx index c0ec705..0cd6c01 100644 --- a/src/components/settings/LanguageSelector.tsx +++ b/src/components/settings/LanguageSelector.tsx @@ -1,7 +1,9 @@ import React, { useState, useRef, useEffect, useMemo } from "react"; +import { listen } from "@tauri-apps/api/event"; import { SettingContainer } from "../ui/SettingContainer"; import { ResetButton } from "../ui/ResetButton"; import { useSettings } from "../../hooks/useSettings"; +import { useModels } from "../../hooks/useModels"; import { LANGUAGES } from "../../lib/constants/languages"; interface LanguageSelectorProps { @@ -14,12 +16,15 @@ export const LanguageSelector: React.FC = ({ grouped = false, }) => { const { getSetting, updateSetting, resetSetting, isUpdating } = useSettings(); + const { currentModel, loadCurrentModel } = useModels(); const [isOpen, setIsOpen] = useState(false); const [searchQuery, setSearchQuery] = useState(""); const dropdownRef = useRef(null); const searchInputRef = useRef(null); const selectedLanguage = getSetting("selected_language") || "auto"; + const isParakeetModel = currentModel === "parakeet-tdt-0.6b-v3"; + const isLanguageSelectionDisabled = isParakeetModel; useEffect(() => { const handleClickOutside = (event: MouseEvent) => { @@ -38,6 +43,17 @@ export const LanguageSelector: React.FC = ({ }; }, []); + // Listen for model state changes to update UI reactively + useEffect(() => { + const modelStateUnlisten = listen("model-state-changed", () => { + loadCurrentModel(); + }); + + return () => { + modelStateUnlisten.then((fn) => fn()); + }; + }, [loadCurrentModel]); + useEffect(() => { if (isOpen && searchInputRef.current) { searchInputRef.current.focus(); @@ -51,8 +67,9 @@ export const LanguageSelector: React.FC = ({ [searchQuery] ); - const selectedLanguageName = - LANGUAGES.find((lang) => lang.value === selectedLanguage)?.label || "Auto"; + const selectedLanguageName = isParakeetModel + ? "Auto" + : LANGUAGES.find((lang) => lang.value === selectedLanguage)?.label || "Auto"; const handleLanguageSelect = async (languageCode: string) => { await updateSetting("selected_language", languageCode); @@ -65,7 +82,7 @@ export const LanguageSelector: React.FC = ({ }; const handleToggle = () => { - if (isUpdating("selected_language")) return; + if (isUpdating("selected_language") || isLanguageSelectionDisabled) return; setIsOpen(!isOpen); }; @@ -86,7 +103,10 @@ export const LanguageSelector: React.FC = ({ return ( @@ -95,12 +115,12 @@ export const LanguageSelector: React.FC = ({ - {isOpen && !isUpdating("selected_language") && ( + {isOpen && !isUpdating("selected_language") && !isLanguageSelectionDisabled && (
{/* Search input */}
@@ -164,7 +184,7 @@ export const LanguageSelector: React.FC = ({
{isUpdating("selected_language") && ( diff --git a/src/hooks/useModels.ts b/src/hooks/useModels.ts index a799c2a..7559b23 100644 --- a/src/hooks/useModels.ts +++ b/src/hooks/useModels.ts @@ -5,11 +5,14 @@ import { listen } from "@tauri-apps/api/event"; interface ModelInfo { id: string; name: string; + description: string; filename: string; url?: string; size_mb: number; is_downloaded: boolean; - is_bundled: boolean; + is_downloading: boolean; + partial_size: number; + is_directory: boolean; } interface DownloadProgress { @@ -25,6 +28,9 @@ export const useModels = () => { const [downloadingModels, setDownloadingModels] = useState>( new Set(), ); + const [extractingModels, setExtractingModels] = useState>( + new Set(), + ); const [downloadProgress, setDownloadProgress] = useState< Map >(new Map()); @@ -117,6 +123,10 @@ export const useModels = () => { return downloadingModels.has(modelId); }; + const isModelExtracting = (modelId: string): boolean => { + return extractingModels.has(modelId); + }; + const getDownloadProgress = ( modelId: string, ): DownloadProgress | undefined => { @@ -158,9 +168,48 @@ export const useModels = () => { }, ); + // Listen for extraction events + const extractionStartedUnlisten = listen( + "model-extraction-started", + (event) => { + const modelId = event.payload; + setExtractingModels((prev) => new Set(prev.add(modelId))); + }, + ); + + const extractionCompletedUnlisten = listen( + "model-extraction-completed", + (event) => { + const modelId = event.payload; + setExtractingModels((prev) => { + const next = new Set(prev); + next.delete(modelId); + return next; + }); + // Refresh models list to update download status + loadModels(); + }, + ); + + const extractionFailedUnlisten = listen<{model_id: string, error: string}>( + "model-extraction-failed", + (event) => { + const modelId = event.payload.model_id; + setExtractingModels((prev) => { + const next = new Set(prev); + next.delete(modelId); + return next; + }); + setError(`Failed to extract model: ${event.payload.error}`); + }, + ); + return () => { progressUnlisten.then((fn) => fn()); completeUnlisten.then((fn) => fn()); + extractionStartedUnlisten.then((fn) => fn()); + extractionCompletedUnlisten.then((fn) => fn()); + extractionFailedUnlisten.then((fn) => fn()); }; }, []); @@ -170,6 +219,7 @@ export const useModels = () => { loading, error, downloadingModels, + extractingModels, downloadProgress, hasAnyModels, isFirstRun, @@ -181,6 +231,7 @@ export const useModels = () => { deleteModel, getModelInfo, isModelDownloading, + isModelExtracting, getDownloadProgress, }; }; diff --git a/src/lib/types.ts b/src/lib/types.ts index 3f99dcf..68000dd 100644 --- a/src/lib/types.ts +++ b/src/lib/types.ts @@ -68,10 +68,13 @@ export const ModelInfoSchema = z.object({ id: z.string(), name: z.string(), description: z.string(), + filename: z.string(), + url: z.string().optional(), size_mb: z.number(), is_downloaded: z.boolean(), is_downloading: z.boolean(), partial_size: z.number(), + is_directory: z.boolean(), }); export type ModelInfo = z.infer;