From 31b5976b175c09495216078ce5c8f9e6dc4bebcb Mon Sep 17 00:00:00 2001
From: Niko <zero@imaskeleton.me>
Date: Tue, 9 Jan 2024 17:38:28 -0700
Subject: [PATCH] Add some logging utilities to capture whisper.cpp output

---
 Cargo.toml                 |  4 ++++
 src/lib.rs                 |  7 ++++++
 src/whisper_sys_log.rs     | 46 ++++++++++++++++++++++++++++++++++++++
 src/whisper_sys_tracing.rs | 43 +++++++++++++++++++++++++++++++++++
 4 files changed, 100 insertions(+)
 create mode 100644 src/whisper_sys_log.rs
 create mode 100644 src/whisper_sys_tracing.rs

diff --git a/Cargo.toml b/Cargo.toml
index e92fefd..b70b5df 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -15,6 +15,8 @@ repository = "https://github.com/tazz4843/whisper-rs"
 
 [dependencies]
 whisper-rs-sys = { path = "sys", version = "0.8" }
+log = { version = "0.4", optional = true }
+tracing = { version = "0.1", optional = true }
 
 [dev-dependencies]
 hound = "3.5.0"
@@ -30,6 +32,8 @@ openblas = ["whisper-rs-sys/openblas"]
 metal = ["whisper-rs-sys/metal", "_gpu"]
 _gpu = []
 test-with-tiny-model = []
+whisper-cpp-log = ["dep:log"]
+whisper-cpp-tracing = ["dep:tracing"]
 
 [package.metadata.docs.rs]
 features = ["simd"]
diff --git a/src/lib.rs b/src/lib.rs
index d93e588..45a91d2 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -7,9 +7,16 @@ mod whisper_ctx;
 mod whisper_grammar;
 mod whisper_params;
 mod whisper_state;
+#[cfg(feature = "whisper-cpp-log")]
+mod whisper_sys_log;
+#[cfg(feature = "whisper-cpp-tracing")]
+mod whisper_sys_tracing;
+
+static LOG_TRAMPOLINE_INSTALL: Once = Once::new();
 
 pub use error::WhisperError;
 pub use standalone::*;
+use std::sync::Once;
 pub use utilities::*;
 pub use whisper_ctx::WhisperContext;
 pub use whisper_ctx::WhisperContextParameters;
diff --git a/src/whisper_sys_log.rs b/src/whisper_sys_log.rs
new file mode 100644
index 0000000..838506f
--- /dev/null
+++ b/src/whisper_sys_log.rs
@@ -0,0 +1,46 @@
+use log::{debug, error, info, warn};
+use std::sync::Once;
+use whisper_rs_sys::ggml_log_level;
+
+unsafe extern "C" fn whisper_cpp_log_trampoline(
+    level: ggml_log_level,
+    text: *const std::os::raw::c_char,
+    _: *mut std::os::raw::c_void, // user_data
+) {
+    if text.is_null() {
+        error!("whisper_cpp_log_trampoline: text is nullptr");
+    }
+
+    // SAFETY: we must trust whisper.cpp that it will not pass us a string that does not satisfy
+    // from_ptr's requirements.
+    let log_str = unsafe { std::ffi::CStr::from_ptr(text) }
+        .to_string_lossy()
+        // whisper.cpp gives newlines at the end of its log messages, so we trim them
+        .trim();
+
+    match level {
+        whisper_rs_sys::ggml_log_level_GGML_LOG_LEVEL_DEBUG => debug!("{}", log_str),
+        whisper_rs_sys::ggml_log_level_GGML_LOG_LEVEL_INFO => info!("{}", log_str),
+        whisper_rs_sys::ggml_log_level_GGML_LOG_LEVEL_WARN => warn!("{}", log_str),
+        whisper_rs_sys::ggml_log_level_GGML_LOG_LEVEL_ERROR => error!("{}", log_str),
+        _ => {
+            warn!(
+                "whisper_cpp_log_trampoline: unknown log level {}: message: {}",
+                level, log_str
+            )
+        }
+    }
+}
+
+static LOG_TRAMPOLINE_INSTALL: Once = Once::new();
+
+/// Shortcut utility to redirect all whisper.cpp logging to the `log` crate.
+///
+/// Filter for logs from the `whisper-rs` crate to see all log output from whisper.cpp.
+///
+/// You should only call this once (subsequent calls have no ill effect).
+pub fn install_whisper_log_trampoline() {
+    LOG_TRAMPOLINE_INSTALL.call_once(|| unsafe {
+        whisper_rs_sys::whisper_log_set(Some(whisper_cpp_log_trampoline), std::ptr::null_mut())
+    });
+}
diff --git a/src/whisper_sys_tracing.rs b/src/whisper_sys_tracing.rs
new file mode 100644
index 0000000..d77769d
--- /dev/null
+++ b/src/whisper_sys_tracing.rs
@@ -0,0 +1,43 @@
+use std::sync::Once;
+use tracing::{debug, error, info, warn};
+use whisper_rs_sys::ggml_log_level;
+
+unsafe extern "C" fn whisper_cpp_tracing_trampoline(
+    level: ggml_log_level,
+    text: *const std::os::raw::c_char,
+    _: *mut std::os::raw::c_void, // user_data
+) {
+    if text.is_null() {
+        error!("whisper_cpp_tracing_trampoline: text is nullptr");
+    }
+
+    // SAFETY: we must trust whisper.cpp that it will not pass us a string that does not satisfy
+    // from_ptr's requirements.
+    let log_str = unsafe { std::ffi::CStr::from_ptr(text) }
+        .to_string_lossy()
+        .trim();
+
+    match level {
+        whisper_rs_sys::ggml_log_level_GGML_LOG_LEVEL_DEBUG => debug!("{}", log_str),
+        whisper_rs_sys::ggml_log_level_GGML_LOG_LEVEL_INFO => info!("{}", log_str),
+        whisper_rs_sys::ggml_log_level_GGML_LOG_LEVEL_WARN => warn!("{}", log_str),
+        whisper_rs_sys::ggml_log_level_GGML_LOG_LEVEL_ERROR => error!("{}", log_str),
+        _ => {
+            warn!(
+                "whisper_cpp_tracing_trampoline: unknown log level {}: message: {}",
+                level, log_str
+            )
+        }
+    }
+}
+
+/// Shortcut utility to redirect all whisper.cpp logging to the `tracing` crate.
+///
+/// Filter for logs from the `whisper-rs` crate to see all log output from whisper.cpp.
+///
+/// You should only call this once (subsequent calls have no effect).
+pub fn install_whisper_tracing_trampoline() {
+    crate::LOG_TRAMPOLINE_INSTALL.call_once(|| unsafe {
+        whisper_rs_sys::whisper_log_set(Some(whisper_cpp_tracing_trampoline), std::ptr::null_mut())
+    });
+}