fix(stage-tamagotchi): cache models & language detection

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<R>,
   ) -> Result<Self> {
-    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);

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"),
+];

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<f32>,
   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<R: Runtime>(
-    model: WhichWhisperModel,
+    model_type: WhichWhisperModel,
     device: Device,
     window: tauri::WebviewWindow<R>,
   ) -> Result<Self> {
-    // 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];
     <LittleEndian as ByteOrder>::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<f32> = (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<String> {
+    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<u32> = 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::<Result<Vec<_>, _>>()?;
+
+    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::<f32>()?;
+
+    // Combine with language info and sort by probability
+    let mut language_probs = languages::LANGUAGES
+      .iter()
+      .zip(probs.iter())
+      .collect::<Vec<_>>();
+
+    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<u32> {
+    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<u32>,
   ) -> Result<Vec<u32>> {
-    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<f32> = 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<f32> = 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<u32> {
-    self
-      .tokenizer
+    tokenizer
       .token_to_id(token)
       .ok_or_else(|| anyhow::anyhow!("Token not found: {}", token))
   }

apps/stage-tamagotchi/src-tauri/src/plugins/audio_transcription/mod.rs

@@ -21,6 +21,15 @@ pub async fn load_model_whisper<R: Runtime>(
 ) -> Result<(), String> {
   info!("Loading models...");
 
+  {
+    let data = app.state::<Mutex<AppDataWhisperProcessor>>();
+    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<R: Runtime>(
 pub async fn audio_transcription<R: Runtime>(
   app: tauri::AppHandle<R>,
   chunk: Vec<f32>,
-) -> Result<String, String> {
+  language: Option<String>,
+) -> Result<(String, String), String> {
   info!("Processing audio transcription...");
 
   let data = app.state::<Mutex<AppDataWhisperProcessor>>();
@@ -61,11 +71,13 @@ pub async fn audio_transcription<R: Runtime>(
   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<R: Runtime>() -> TauriPlugin<R> {

apps/stage-tamagotchi/src-tauri/src/plugins/audio_vad/mod.rs

@@ -21,6 +21,15 @@ pub async fn load_model_silero_vad<R: Runtime>(
 ) -> Result<(), String> {
   info!("Loading models...");
 
+  {
+    let data = app.state::<Mutex<AppDataSileroVadProcessor>>();
+    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::<Mutex<AppDataSileroVadProcessor>>();