diff --git a/apps/stage-tamagotchi/src-tauri/src/app/models/silero_vad/mod.rs b/apps/stage-tamagotchi/src-tauri/src/app/models/silero_vad/mod.rs index a411e86b0..e2fa85598 100644 --- a/apps/stage-tamagotchi/src-tauri/src/app/models/silero_vad/mod.rs +++ b/apps/stage-tamagotchi/src-tauri/src/app/models/silero_vad/mod.rs @@ -1,7 +1,8 @@ use std::path::PathBuf; use anyhow::Result; -use log::{debug, info}; +use hf_hub::Repo; +use log::info; use ndarray::{Array2, Array3}; use ort::{ execution_providers::{ @@ -33,17 +34,30 @@ impl Processor { _device: candle_core::Device, window: tauri::WebviewWindow, ) -> Result { - let api = hf_hub::api::sync::ApiBuilder::new().build()?; - let repo = api.repo(hf_hub::Repo::with_revision( - "onnx-community/silero-vad".into(), + let model_id = "onnx-community/silero-vad"; + let revision = "main"; + + let cache_api = hf_hub::Cache::from_env(); + let cache_repo = cache_api.repo(Repo::with_revision( + model_id.into(), hf_hub::RepoType::Model, - "main".into(), + revision.into(), )); - let model_path = repo.download_with_progress( - "onnx/model.onnx", - create_progress_emitter(window.clone(), "onnx/model.onnx".to_string()), - )?; + let api = hf_hub::api::sync::ApiBuilder::new().build()?; + let repo = api.repo(hf_hub::Repo::with_revision( + model_id.into(), + hf_hub::RepoType::Model, + revision.into(), + )); + + let model_path = match cache_repo.get("model.onnx") { + Some(path) => path, + None => repo.download_with_progress( + "onnx/model.onnx", + create_progress_emitter(window.clone(), "onnx/model.onnx".to_string()), + )?, + }; let session = Self::create_optimized_session(model_path.clone())?; let (frame_size, context_size) = (512, 64); diff --git a/apps/stage-tamagotchi/src-tauri/src/app/models/whisper/languages.rs b/apps/stage-tamagotchi/src-tauri/src/app/models/whisper/languages.rs new file mode 100644 index 000000000..42f41370d --- /dev/null +++ b/apps/stage-tamagotchi/src-tauri/src/app/models/whisper/languages.rs @@ -0,0 +1,102 @@ +// Language constants +pub const LANGUAGES: [(&str, &str); 99] = [ + ("en", "english"), + ("zh", "chinese"), + ("de", "german"), + ("es", "spanish"), + ("ru", "russian"), + ("ko", "korean"), + ("fr", "french"), + ("ja", "japanese"), + ("pt", "portuguese"), + ("tr", "turkish"), + ("pl", "polish"), + ("ca", "catalan"), + ("nl", "dutch"), + ("ar", "arabic"), + ("sv", "swedish"), + ("it", "italian"), + ("id", "indonesian"), + ("hi", "hindi"), + ("fi", "finnish"), + ("vi", "vietnamese"), + ("he", "hebrew"), + ("uk", "ukrainian"), + ("el", "greek"), + ("ms", "malay"), + ("cs", "czech"), + ("ro", "romanian"), + ("da", "danish"), + ("hu", "hungarian"), + ("ta", "tamil"), + ("no", "norwegian"), + ("th", "thai"), + ("ur", "urdu"), + ("hr", "croatian"), + ("bg", "bulgarian"), + ("lt", "lithuanian"), + ("la", "latin"), + ("mi", "maori"), + ("ml", "malayalam"), + ("cy", "welsh"), + ("sk", "slovak"), + ("te", "telugu"), + ("fa", "persian"), + ("lv", "latvian"), + ("bn", "bengali"), + ("sr", "serbian"), + ("az", "azerbaijani"), + ("sl", "slovenian"), + ("kn", "kannada"), + ("et", "estonian"), + ("mk", "macedonian"), + ("br", "breton"), + ("eu", "basque"), + ("is", "icelandic"), + ("hy", "armenian"), + ("ne", "nepali"), + ("mn", "mongolian"), + ("bs", "bosnian"), + ("kk", "kazakh"), + ("sq", "albanian"), + ("sw", "swahili"), + ("gl", "galician"), + ("mr", "marathi"), + ("pa", "punjabi"), + ("si", "sinhala"), + ("km", "khmer"), + ("sn", "shona"), + ("yo", "yoruba"), + ("so", "somali"), + ("af", "afrikaans"), + ("oc", "occitan"), + ("ka", "georgian"), + ("be", "belarusian"), + ("tg", "tajik"), + ("sd", "sindhi"), + ("gu", "gujarati"), + ("am", "amharic"), + ("yi", "yiddish"), + ("lo", "lao"), + ("uz", "uzbek"), + ("fo", "faroese"), + ("ht", "haitian creole"), + ("ps", "pashto"), + ("tk", "turkmen"), + ("nn", "nynorsk"), + ("mt", "maltese"), + ("sa", "sanskrit"), + ("lb", "luxembourgish"), + ("my", "myanmar"), + ("bo", "tibetan"), + ("tl", "tagalog"), + ("mg", "malagasy"), + ("as", "assamese"), + ("tt", "tatar"), + ("haw", "hawaiian"), + ("ln", "lingala"), + ("ha", "hausa"), + ("ba", "bashkir"), + ("jw", "javanese"), + ("su", "sundanese"), +]; diff --git a/apps/stage-tamagotchi/src-tauri/src/app/models/whisper/mod.rs b/apps/stage-tamagotchi/src-tauri/src/app/models/whisper/mod.rs index d771bab42..2d787a8c9 100644 --- a/apps/stage-tamagotchi/src-tauri/src/app/models/whisper/mod.rs +++ b/apps/stage-tamagotchi/src-tauri/src/app/models/whisper/mod.rs @@ -1,8 +1,7 @@ -// https://github.com/proj-airi/candle-examples/blob/3a6783a4788333e7a117f4f8548f02b2fbfec2ed/apps/silero-vad-whisper-realtime/src/whisper.rs use anyhow::Result; use byteorder::{ByteOrder, LittleEndian}; -use candle_core::{Device, IndexOp, Tensor}; -use candle_nn::VarBuilder; +use candle_core::{D, Device, IndexOp, Tensor}; +use candle_nn::{VarBuilder, ops::softmax}; use candle_transformers::models::whisper::{self as whisper_model, Config, audio}; use clap::ValueEnum; use hf_hub::{Repo, RepoType, api::sync::ApiBuilder}; @@ -12,6 +11,8 @@ use tokenizers::Tokenizer; use crate::helpers::huggingface::create_progress_emitter; +mod languages; + pub enum WhisperModel { Normal(whisper_model::model::Whisper), } @@ -46,6 +47,12 @@ impl WhisperModel { Self::Normal(model) => model.decoder.final_linear(x), } } + + pub fn reset_kv_cache(&mut self) { + match self { + Self::Normal(model) => model.reset_kv_cache(), + } + } } #[derive(Clone, Copy, Debug, PartialEq, Eq, ValueEnum)] @@ -91,6 +98,21 @@ impl WhichWhisperModel { Self::DistilLargeV2 => ("distil-whisper/distil-large-v2", "main"), } } + + pub const fn is_multilingual(self) -> bool { + match self { + Self::Tiny + | Self::Base + | Self::Small + | Self::Medium + | Self::Large + | Self::LargeV2 + | Self::LargeV3 + | Self::LargeV3Turbo + | Self::DistilLargeV2 => true, + Self::TinyEn | Self::BaseEn | Self::SmallEn | Self::MediumEn | Self::DistilMediumEn => false, + } + } } pub struct Processor { @@ -99,37 +121,64 @@ pub struct Processor { pub config: Config, pub mel_filters: Vec, pub device: Device, + pub model_type: WhichWhisperModel, + + // Special tokens + pub sot_token: u32, + pub transcribe_token: u32, + pub translate_token: u32, + pub eot_token: u32, + pub no_timestamps_token: u32, + pub suppress_tokens: Tensor, } impl Processor { pub fn new( - model: WhichWhisperModel, + model_type: WhichWhisperModel, device: Device, window: tauri::WebviewWindow, ) -> Result { - // Load the Whisper model based on the provided model type + let (model_id, revision) = model_type.model_and_revision(); + + let cache_api = hf_hub::Cache::from_env(); + let cache_repo = cache_api.repo(Repo::with_revision( + model_id.to_string(), + RepoType::Model, + revision.to_string(), + )); + let api = ApiBuilder::new().with_progress(false).build()?; - let (model_id, revision) = model.model_and_revision(); let repo = api.repo(Repo::with_revision( model_id.to_string(), RepoType::Model, revision.to_string(), )); - let config_filename = repo.download_with_progress( - "config.json", - create_progress_emitter(window.clone(), "config.json".to_string()), - )?; + let config_filename = match cache_repo.get("config.json") { + None => repo.download_with_progress( + "config.json", + create_progress_emitter(window.clone(), "config.json".to_string()), + )?, + Some(p) => p, + }; info!("config_filename: {:?}", config_filename.display()); - let tokenizer_filename = repo.download_with_progress( - "tokenizer.json", - create_progress_emitter(window.clone(), "tokenizer.json".to_string()), - )?; + + let tokenizer_filename = match cache_repo.get("tokenizer.json") { + None => repo.download_with_progress( + "tokenizer.json", + create_progress_emitter(window.clone(), "tokenizer.json".to_string()), + )?, + Some(p) => p, + }; info!("tokenizer_filename: {:?}", tokenizer_filename.display()); - let model_filename = repo.download_with_progress( - "model.safetensors", - create_progress_emitter(window.clone(), "model.safetensors".to_string()), - )?; + + let model_filename = match cache_repo.get("model.safetensors") { + None => repo.download_with_progress( + "model.safetensors", + create_progress_emitter(window.clone(), "model.safetensors".to_string()), + )?, + Some(p) => p, + }; info!("model_filename: {:?}", model_filename.display()); let config: Config = serde_json::from_str(&std::fs::read_to_string(config_filename)?)?; @@ -156,19 +205,48 @@ impl Processor { let mut mel_filters = vec![0f32; mel_bytes.len() / 4]; ::read_f32_into(mel_bytes, &mut mel_filters); + // Initialize special tokens + let sot_token = Self::token_id_static(&tokenizer, whisper_model::SOT_TOKEN)?; + let transcribe_token = Self::token_id_static(&tokenizer, whisper_model::TRANSCRIBE_TOKEN)?; + let translate_token = Self::token_id_static(&tokenizer, whisper_model::TRANSLATE_TOKEN)?; + let eot_token = Self::token_id_static(&tokenizer, whisper_model::EOT_TOKEN)?; + let no_timestamps_token = + Self::token_id_static(&tokenizer, whisper_model::NO_TIMESTAMPS_TOKEN)?; + + // Create suppress tokens + let suppress_values: Vec = (0..config.vocab_size as u32) + .map(|i| { + if config.suppress_tokens.contains(&i) { + f32::NEG_INFINITY + } else { + 0.0f32 + } + }) + .collect(); + + let suppress_tokens = Tensor::new(suppress_values.as_slice(), &device)?; + Ok(Self { model, tokenizer, config, mel_filters, device, + model_type, + sot_token, + transcribe_token, + translate_token, + eot_token, + no_timestamps_token, + suppress_tokens, }) } pub fn transcribe( &mut self, audio: &[f32], - ) -> Result { + language: Option<&str>, + ) -> Result<(String, String)> { // Convert PCM to mel spectrogram let mel = audio::pcm_to_mel(&self.config, audio, &self.mel_filters); let mel_len = mel.len(); @@ -182,30 +260,134 @@ impl Processor { &self.device, )?; - // Run inference + // Get language token - either specified or detected + let (language_token, detected_language) = match language { + Some(lang) => { + let token = self.get_language_token(lang)?; + (Some(token), lang.to_string()) + }, + None => { + if self.model_type.is_multilingual() { + let (token, lang) = self.detect_language(&mel)?; + (Some(token), lang) + } else { + // English-only models don't use language tokens + (None, "english".to_string()) + } + }, + }; + + // Run encoder let audio_features = self.model.encoder_forward(&mel, true)?; - // Simple greedy decoding - let tokens = self.decode_greedy(&audio_features)?; + + // Decode with the determined language + let tokens = self.decode_whisper_with_language_token(&audio_features, language_token)?; + + // Filter out special tokens + let filtered_tokens: Vec = tokens + .into_iter() + .filter(|&token| !self.is_special_token(token)) + .collect(); let text = self .tokenizer - .decode(&tokens, true) + .decode(&filtered_tokens, true) .map_err(anyhow::Error::msg)?; - Ok(text) + // Reset KV cache for next inference + self.model.reset_kv_cache(); + + Ok((text.trim().to_string(), detected_language)) } - fn decode_greedy( + fn detect_language( + &mut self, + mel: &Tensor, + ) -> Result<(u32, String)> { + let (_bsize, _, seq_len) = mel.dims3()?; + let mel = mel.narrow(2, 0, usize::min(seq_len, self.config.max_source_positions))?; + let device = mel.device(); + + // Get all language token IDs + let language_token_ids = languages::LANGUAGES + .iter() + .map(|(code, _)| Self::token_id_static(&self.tokenizer, &format!("<|{}|>", code))) + .collect::, _>>()?; + + let audio_features = self.model.encoder_forward(&mel, true)?; + let tokens = Tensor::new(&[[self.sot_token]], device)?; + let language_token_ids_tensor = Tensor::new(language_token_ids.as_slice(), device)?; + + let ys = self + .model + .decoder_forward(&tokens, &audio_features, true)?; + let logits = self + .model + .decoder_final_linear(&ys.i(..1)?)? + .i(0)? + .i(0)?; + + // Get logits for language tokens only + let logits = logits.index_select(&language_token_ids_tensor, 0)?; + let probs = softmax(&logits, D::Minus1)?; + let probs = probs.to_vec1::()?; + + // Combine with language info and sort by probability + let mut language_probs = languages::LANGUAGES + .iter() + .zip(probs.iter()) + .collect::>(); + + language_probs.sort_by(|(_, p1), (_, p2)| p2.total_cmp(p1)); + + // Log top 5 detected languages + info!("Language detection results:"); + for ((_, name), prob) in language_probs.iter().take(5) { + info!(" {}: {:.3}", name, prob); + } + + // Return the most likely language + let best_language = language_probs[0].0; + let language_token = + Self::token_id_static(&self.tokenizer, &format!("<|{}|>", best_language.0))?; + + Ok((language_token, best_language.1.to_string())) + } + + fn get_language_token( + &self, + language: &str, + ) -> Result { + let language_lower = language.to_lowercase(); + + // Check if it's a language code or name + for (code, name) in languages::LANGUAGES.iter() { + if code == &language_lower || name == &language_lower { + let token = format!("<|{}|>", code); + return Self::token_id_static(&self.tokenizer, &token); + } + } + + anyhow::bail!("Unsupported language: {}", language) + } + + fn decode_whisper_with_language_token( &mut self, audio_features: &Tensor, + language_token: Option, ) -> Result> { - let mut tokens = vec![ - self.token_id(whisper_model::SOT_TOKEN)?, - self.token_id(whisper_model::TRANSCRIBE_TOKEN)?, - self.token_id(whisper_model::NO_TIMESTAMPS_TOKEN)?, - ]; + // Initialize with proper Whisper token sequence + let mut tokens = vec![self.sot_token]; - let max_len = 50; // Short sequence for real-time processing + // Add language token if provided (for multilingual models) + if let Some(lang_token) = language_token { + tokens.push(lang_token); + } + + tokens.push(self.transcribe_token); + tokens.push(self.no_timestamps_token); + + let max_len = 200; for i in 0..max_len { let tokens_t = Tensor::new(tokens.as_slice(), &self.device)?.unsqueeze(0)?; @@ -220,31 +402,61 @@ impl Processor { .i(0)? .i(0)?; - // Get most likely token - let logits_v: Vec = logits.to_vec1()?; - let next_token = logits_v + // Apply suppression tokens + let logits = logits.broadcast_add(&self.suppress_tokens)?; + + // Apply softmax and get most likely token + let probs = softmax(&logits, 0)?; + let probs_vec: Vec = probs.to_vec1()?; + + let next_token = probs_vec .iter() .enumerate() .max_by(|(_, a), (_, b)| a.total_cmp(b)) - .map(|(i, _)| u32::try_from(i).unwrap()) + .map(|(i, _)| i as u32) .unwrap(); - if next_token == self.token_id(whisper_model::EOT_TOKEN)? { + tokens.push(next_token); + + if next_token == self.eot_token { break; } - - tokens.push(next_token); } Ok(tokens) } - fn token_id( + fn is_special_token( &self, + token: u32, + ) -> bool { + // Check basic special tokens + if token == self.sot_token + || token == self.eot_token + || token == self.transcribe_token + || token == self.translate_token + || token == self.no_timestamps_token + { + return true; + } + + // Check if it's a language token + for (code, _) in languages::LANGUAGES.iter() { + if let Ok(lang_token) = Self::token_id_static(&self.tokenizer, &format!("<|{}|>", code)) { + if token == lang_token { + return true; + } + } + } + + false + } + + fn token_id_static( + tokenizer: &Tokenizer, token: &str, ) -> Result { - self - .tokenizer + tokenizer .token_to_id(token) .ok_or_else(|| anyhow::anyhow!("Token not found: {}", token)) } diff --git a/apps/stage-tamagotchi/src-tauri/src/plugins/audio_transcription/mod.rs b/apps/stage-tamagotchi/src-tauri/src/plugins/audio_transcription/mod.rs index cffe8c5fa..172b7f26a 100644 --- a/apps/stage-tamagotchi/src-tauri/src/plugins/audio_transcription/mod.rs +++ b/apps/stage-tamagotchi/src-tauri/src/plugins/audio_transcription/mod.rs @@ -21,6 +21,15 @@ pub async fn load_model_whisper( ) -> Result<(), String> { info!("Loading models..."); + { + let data = app.state::>(); + let data = data.lock().unwrap(); + if data.whisper_processor.is_some() { + info!("Whisper model already loaded, skipping..."); + return Ok(()); + } + } + // Load the traditional whisper models first match new_whisper_processor(window) { Ok(p) => { @@ -44,7 +53,8 @@ pub async fn load_model_whisper( pub async fn audio_transcription( app: tauri::AppHandle, chunk: Vec, -) -> Result { + language: Option, +) -> Result<(String, String), String> { info!("Processing audio transcription..."); let data = app.state::>(); @@ -61,11 +71,13 @@ pub async fn audio_transcription( let mut data = data.lock().unwrap(); let processor = data.whisper_processor.as_mut().unwrap(); - let transcription = processor - .transcribe(chunk.as_slice()) + let (transcription, language) = processor + .transcribe(chunk.as_slice(), language.as_deref()) .map_err(|e| e.to_string())?; - Ok(transcription) + info!("Transcription completed: {}", transcription); + + Ok((transcription, language)) } pub fn init() -> TauriPlugin { diff --git a/apps/stage-tamagotchi/src-tauri/src/plugins/audio_vad/mod.rs b/apps/stage-tamagotchi/src-tauri/src/plugins/audio_vad/mod.rs index 31a4b1e1a..fbcf39eed 100644 --- a/apps/stage-tamagotchi/src-tauri/src/plugins/audio_vad/mod.rs +++ b/apps/stage-tamagotchi/src-tauri/src/plugins/audio_vad/mod.rs @@ -21,6 +21,15 @@ pub async fn load_model_silero_vad( ) -> Result<(), String> { info!("Loading models..."); + { + let data = app.state::>(); + let data = data.lock().unwrap(); + if data.silero_vad_processor.is_some() { + info!("Silero VAD model already loaded, skipping..."); + return Ok(()); + } + } + match new_silero_vad_processor(window) { Ok(p) => { let data = app.state::>();