one

Author: yujonglee

crates/whisper-local/src/model.rs

@@ -189,11 +189,26 @@ impl Whisper {
 
         let lang_str = {
             self.state.pcm_to_mel(audio, 1)?;
-            let (lang_id, _lang_probs) = self.state.lang_detect(0, 1)?;
-            whisper_rs::get_lang_str(lang_id)
+            let (_lang_id, lang_probs) = self.state.lang_detect(0, 1)?;
+
+            let mut best_lang = None;
+            let mut best_prob = f32::NEG_INFINITY;
+
+            for lang in &self.languages {
+                let lang_id = lang.whisper_index();
+                if lang_id < lang_probs.len() {
+                    let prob = lang_probs[lang_id];
+                    if prob > best_prob {
+                        best_prob = prob;
+                        best_lang = Some(lang.as_ref().to_string());
+                    }
+                }
+            }
+
+            best_lang
         };
 
-        Ok(lang_str.map(|s| s.to_owned()))
+        Ok(lang_str)
     }
 
     fn filter_segments(segments: Vec<Segment>) -> Vec<Segment> {

crates/whisper/src/lib.rs

@@ -1,5 +1,6 @@
 // https://github.com/openai/whisper/blob/ba3f3cd/whisper/tokenizer.py#L10-L128
-#[derive(Debug, strum::EnumString, strum::Display, strum::AsRefStr)]
+#[repr(u8)]
+#[derive(Debug, Copy, Clone, strum::EnumString, strum::Display, strum::AsRefStr)]
 pub enum Language {
     #[strum(serialize = "en")]
     En,
@@ -202,3 +203,9 @@ pub enum Language {
     #[strum(serialize = "yue")]
     Yue,
 }
+
+impl Language {
+    pub fn whisper_index(self) -> usize {
+        self as usize
+    }
+}