From 8787a44ae4accb9b6039c69b76c6a24388c3f41f Mon Sep 17 00:00:00 2001
From: ckaznable <kami25565@gmail.com>
Date: Fri, 9 Jun 2023 15:37:26 +0800
Subject: [PATCH] feat: support zh-tw zh-hk zh-cn

---
 Cargo.lock  | 42 ++++++++++++++++++++++++++++++++++++++++++
 Cargo.toml  |  5 +++++
 src/main.rs | 15 ++++++++++++++-
 src/zh.rs   | 35 +++++++++++++++++++++++++++++++++++
 4 files changed, 96 insertions(+), 1 deletion(-)
 create mode 100644 src/zh.rs

diff --git a/Cargo.lock b/Cargo.lock
index bc8005f..75ab494 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -521,6 +521,17 @@ dependencies = [
  "static_assertions",
 ]
 
+[[package]]
+name = "lazy-static-include"
+version = "3.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b111244b70d4cf22aaaf8e0461ede19d623880f0f4779ee33dc35850a825bec3"
+dependencies = [
+ "lazy_static",
+ "manifest-dir-macros",
+ "syn 2.0.16",
+]
+
 [[package]]
 name = "lazy_static"
 version = "1.4.0"
@@ -637,6 +648,18 @@ dependencies = [
  "cfg-if",
 ]
 
+[[package]]
+name = "manifest-dir-macros"
+version = "0.1.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "450e5ef583bc05177c4975b9ea907047091a9f62e74e81fcafb99dbffac51e7e"
+dependencies = [
+ "once_cell",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.16",
+]
+
 [[package]]
 name = "maplit"
 version = "1.0.2"
@@ -763,6 +786,18 @@ version = "1.17.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b7e5500299e16ebb147ae15a00a942af264cf3688f47923b8fc2cd5858f23ad3"
 
+[[package]]
+name = "opencc-rust"
+version = "1.1.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ba09724567a3268572f2be89ac29678a9f9894231045061d65433625b1aab774"
+dependencies = [
+ "lazy-static-include",
+ "lazy_static",
+ "libc",
+ "pkg-config",
+]
+
 [[package]]
 name = "owo-colors"
 version = "3.5.0"
@@ -844,6 +879,12 @@ dependencies = [
  "sha2",
 ]
 
+[[package]]
+name = "pkg-config"
+version = "0.3.27"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "26072860ba924cbfa98ea39c8c19b4dd6a4a25423dbdf219c1eca91aa0cf6964"
+
 [[package]]
 name = "ppv-lite86"
 version = "0.2.17"
@@ -1817,6 +1858,7 @@ name = "yt-cli-live"
 version = "0.1.1"
 dependencies = [
  "clap",
+ "opencc-rust",
  "owo-colors",
  "ringbuf",
  "rubato",
diff --git a/Cargo.toml b/Cargo.toml
index 7db0bfd..ba93f23 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -17,6 +17,7 @@ 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
@@ -24,3 +25,7 @@ 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"]
diff --git a/src/main.rs b/src/main.rs
index c23d2e7..4360bdf 100644
--- a/src/main.rs
+++ b/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,7 +238,7 @@ 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();
 
@@ -240,6 +246,13 @@ fn evoke_whisper_thread(
                 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(),
diff --git a/src/zh.rs b/src/zh.rs
new file mode 100644
index 0000000..056441d
--- /dev/null
+++ b/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)
+    }
+}