feat: enable tdrz

Enables tinydiarize models ggerganov/whisper.cpp#1058
JEF1056 · Sep 29, 2023 · e2999d5 · e2999d5
1 parent d07a114
commit e2999d5
Show file tree

Hide file tree

Showing 12 changed files with 12,445 additions and 8,888 deletions.
diff --git a/android/src/main/java/com/rnwhisper/WhisperContext.java b/android/src/main/java/com/rnwhisper/WhisperContext.java
@@ -474,6 +474,8 @@ private int full(int jobId, ReadableMap options, float[] audioData, int audioDat
       options.hasKey("speedUp") ? options.getBoolean("speedUp") : false,
       // jboolean translate,
       options.hasKey("translate") ? options.getBoolean("translate") : false,
+      // jboolean tdrz_enable,
+      options.hasKey("tdrzEnable") ? options.getBoolean("tdrzEnable") : false,
       // jstring language,
       options.hasKey("language") ? options.getString("language") : "auto",
       // jstring prompt
@@ -645,6 +647,7 @@ protected static native int fullTranscribe(
     int best_of,
     boolean speed_up,
     boolean translate,
+    boolean tdrz_enable,
     String language,
     String prompt,
     ProgressCallback progressCallback

diff --git a/android/src/main/jni.cpp b/android/src/main/jni.cpp
@@ -232,6 +232,7 @@ Java_com_rnwhisper_WhisperContext_fullTranscribe(
     jint best_of,
     jboolean speed_up,
     jboolean translate,
+    jboolean tdrz_enable,
     jstring language,
     jstring prompt,
     jobject progress_callback_instance
@@ -256,7 +257,7 @@ Java_com_rnwhisper_WhisperContext_fullTranscribe(
     params.print_realtime = false;
     params.print_progress = false;
     params.print_timestamps = false;
-    params.print_special = false;
+    params.print_special = true;
     params.translate = translate;
     const char *language_chars = env->GetStringUTFChars(language, nullptr);
     params.language = language_chars;
@@ -265,6 +266,7 @@ Java_com_rnwhisper_WhisperContext_fullTranscribe(
     params.offset_ms = 0;
     params.no_context = true;
     params.single_segment = false;
+    params.tdrz_enable = tdrz_enable;
 
     if (max_len > -1) {
         params.max_len = max_len;

diff --git a/cpp/coreml/whisper-encoder.mm b/cpp/coreml/whisper-encoder.mm
@@ -22,7 +22,13 @@
 
     NSURL * url_model = [NSURL fileURLWithPath: path_model_str];
 
-    const void * data = CFBridgingRetain([[whisper_encoder_impl alloc] initWithContentsOfURL:url_model error:nil]);
+    // select which device to run the Core ML model on
+    MLModelConfiguration *config = [[MLModelConfiguration alloc] init];
+    config.computeUnits = MLComputeUnitsCPUAndGPU;
+    //config.computeUnits = MLComputeUnitsCPUAndNeuralEngine;
+    //config.computeUnits = MLComputeUnitsAll;
+
+    const void * data = CFBridgingRetain([[whisper_encoder_impl alloc] initWithContentsOfURL:url_model configuration:config error:nil]);
 
     if (data == NULL) {
         return NULL;