feat: support zh-tw zh-hk zh-cn

Cargo.toml

@@ -17,10 +17,15 @@ tract-onnx = "0.17.9"
 rubato = "0.13.0"
 owo-colors = "3.5.0"
 yt_tsu = "0.1.0"
+opencc-rust = { version = "1.1.15", optional = true }
 
 [profile.release]
 opt-level = 'z'     # Optimize for size
 lto = true          # Enable link-time optimization
 codegen-units = 1   # Reduce number of codegen units to increase optimizations
 panic = 'abort'     # Abort on panic
 strip = true        # Strip symbols from binary*
+
+[features]
+default = []
+zh = ["opencc-rust/static-dictionaries"]

src/main.rs

@@ -25,6 +25,8 @@ mod audio;
 mod speech;
 mod util;
 mod vad;
+#[cfg(feature = "zh")]
+mod zh;
 
 type F32Consumer = Consumer<f32, Arc<SharedRb<f32, Vec<MaybeUninit<f32>>>>>;
 type SegmentProducer =
@@ -224,6 +226,10 @@ fn evoke_whisper_thread(
     thread::spawn(move || {
         let mut state = ctx.create_state().expect("failed to create state");
         let mut streaming_time = 0.0f64;
+        let lang = if args.lang.starts_with("zh") { "zh" } else { &args.lang };
+
+        #[cfg(feature = "zh")]
+        let zh_transformer = zh::ZHTransformer::from(&args.lang);
 
         while let Ok(ThreadState::Sync) = rx.recv() {
             if cons.is_empty() {
@@ -232,14 +238,21 @@ fn evoke_whisper_thread(
             }
 
             cons.pop_iter().for_each(|segment| {
-                let config = SpeechConfig::new(args.threads as c_int, Some(&args.lang));
+                let config = SpeechConfig::new(args.threads as c_int, Some(lang));
                 let mut payload: WhisperPayload = WhisperPayload::new(&segment.data, config);
                 let running_calc = Instant::now();
 
                 let segment_time = (streaming_time * 1000.0) as i64;
                 streaming_time += segment.duration as f64;
 
                 speech::process(&mut state, &mut payload, &mut |segment, start| {
+                    #[cfg(feature = "zh")]
+                    let segment = if let Ok(zh) = &zh_transformer {
+                        zh.convert(segment)
+                    } else {
+                        segment.to_string()
+                    };
+
                     println!(
                         "[{}] {}",
                         util::format_timestamp_to_time(segment_time + start).bright_yellow(),

src/zh.rs

@@ -0,0 +1,35 @@
+use opencc_rust::{DefaultConfig, OpenCC};
+
+pub struct ZHTransformer {
+    convert: OpenCC
+}
+
+impl ZHTransformer {
+    pub fn from(s: &str) -> Result<Self, &'static str> {
+        let s = s.to_string().to_lowercase();
+
+        if s == "zh" {
+            return Err("auto detection chinese language code");
+        }
+
+        if !s.starts_with("zh") {
+            return Err("is not chinese language code");
+        }
+
+        let config = if matches!(s.as_str(), "zh_tw" | "zh-tw") {
+            DefaultConfig::S2TW
+        } else if matches!(s.as_str(), "zh_hk" | "zh-hk") {
+            DefaultConfig::S2HK
+        } else {
+            DefaultConfig::T2S
+        };
+
+        Ok(Self {
+            convert: OpenCC::new(config)?,
+        })
+    }
+
+    pub fn convert<S: AsRef<str>>(&self, s: S) -> String {
+        self.convert.convert(s)
+    }
+}