whisper.rn 0.4.0-rc.0 → 0.4.0-rc.10

package/android/src/main/java/com/rnwhisper/WhisperContext.java

@@ -42,7 +42,6 @@ public class WhisperContext {
   private AudioRecord recorder = null;
   private int bufferSize;
   private int nSamplesTranscribing = 0;
-  private ArrayList<short[]> shortBufferSlices;
   // Remember number of samples in each slice
   private ArrayList<Integer> sliceNSamples;
   // Current buffer slice index
@@ -54,6 +53,7 @@ public class WhisperContext {
   private boolean isCapturing = false;
   private boolean isStoppedByAction = false;
   private boolean isTranscribing = false;
+  private boolean isTdrzEnable = false;
   private Thread rootFullHandler = null;
   private Thread fullHandler = null;
 
@@ -66,7 +66,6 @@ public class WhisperContext {
   }
 
   private void rewind() {
-    shortBufferSlices = null;
     sliceNSamples = null;
     sliceIndex = 0;
     transcribeSliceIndex = 0;
@@ -75,29 +74,19 @@ public class WhisperContext {
     isCapturing = false;
     isStoppedByAction = false;
     isTranscribing = false;
+    isTdrzEnable = false;
     rootFullHandler = null;
     fullHandler = null;
   }
 
-  private boolean vad(ReadableMap options, short[] shortBuffer, int nSamples, int n) {
-    boolean isSpeech = true;
-    if (!isTranscribing && options.hasKey("useVad") && options.getBoolean("useVad")) {
-      int vadSec = options.hasKey("vadMs") ? options.getInt("vadMs") / 1000 : 2;
-      int sampleSize = vadSec * SAMPLE_RATE;
-      if (nSamples + n > sampleSize) {
-        int start = nSamples + n - sampleSize;
-        float[] audioData = new float[sampleSize];
-        for (int i = 0; i < sampleSize; i++) {
-          audioData[i] = shortBuffer[i + start] / 32768.0f;
-        }
-        float vadThold = options.hasKey("vadThold") ? (float) options.getDouble("vadThold") : 0.6f;
-        float vadFreqThold = options.hasKey("vadFreqThold") ? (float) options.getDouble("vadFreqThold") : 0.6f;
-        isSpeech = vadSimple(audioData, sampleSize, vadThold, vadFreqThold);
-      } else {
-        isSpeech = false;
-      }
-    }
-    return isSpeech;
+  private boolean vad(int sliceIndex, int nSamples, int n) {
+    if (isTranscribing) return true;
+    return vadSimple(jobId, sliceIndex, nSamples, n);
+  }
+
+  private void finishRealtimeTranscribe(WritableMap result) {
+    emitTranscribeEvent("@RNWhisper_onRealtimeTranscribeEnd", Arguments.createMap());
+    finishRealtimeTranscribeJob(jobId, context, sliceNSamples.stream().mapToInt(i -> i).toArray());
   }
 
   public int startRealtimeTranscribe(int jobId, ReadableMap options) {
@@ -119,19 +108,20 @@ public class WhisperContext {
 
     int realtimeAudioSec = options.hasKey("realtimeAudioSec") ? options.getInt("realtimeAudioSec") : 0;
     final int audioSec = realtimeAudioSec > 0 ? realtimeAudioSec : DEFAULT_MAX_AUDIO_SEC;
-
     int realtimeAudioSliceSec = options.hasKey("realtimeAudioSliceSec") ? options.getInt("realtimeAudioSliceSec") : 0;
     final int audioSliceSec = realtimeAudioSliceSec > 0 && realtimeAudioSliceSec < audioSec ? realtimeAudioSliceSec : audioSec;
-
     isUseSlices = audioSliceSec < audioSec;
 
-    String audioOutputPath = options.hasKey("audioOutputPath") ? options.getString("audioOutputPath") : null;
+    double realtimeAudioMinSec = options.hasKey("realtimeAudioMinSec") ? options.getDouble("realtimeAudioMinSec") : 0;
+    final double audioMinSec = realtimeAudioMinSec > 0.5 && realtimeAudioMinSec <= audioSliceSec ? realtimeAudioMinSec : 1;
+
+    this.isTdrzEnable = options.hasKey("tdrzEnable") && options.getBoolean("tdrzEnable");
+
+    createRealtimeTranscribeJob(jobId, context, options);
 
-    shortBufferSlices = new ArrayList<short[]>();
-    shortBufferSlices.add(new short[audioSliceSec * SAMPLE_RATE]);
     sliceNSamples = new ArrayList<Integer>();
     sliceNSamples.add(0);
-  
+
     isCapturing = true;
     recorder.startRecording();
 
@@ -159,49 +149,43 @@ public class WhisperContext {
                   nSamples == nSamplesTranscribing &&
                   sliceIndex == transcribeSliceIndex
                 ) {
-                  emitTranscribeEvent("@RNWhisper_onRealtimeTranscribeEnd", Arguments.createMap());
+                  finishRealtimeTranscribe(Arguments.createMap());
                 } else if (!isTranscribing) {
-                  short[] shortBuffer = shortBufferSlices.get(sliceIndex);
-                  boolean isSpeech = vad(options, shortBuffer, nSamples, 0);
-                  if (!isSpeech) {
-                    emitTranscribeEvent("@RNWhisper_onRealtimeTranscribeEnd", Arguments.createMap());
+                  boolean isSamplesEnough = nSamples / SAMPLE_RATE >= audioMinSec;
+                  if (!isSamplesEnough || !vad(sliceIndex, nSamples, 0)) {
+                    finishRealtimeTranscribe(Arguments.createMap());
                     break;
                   }
                   isTranscribing = true;
-                  fullTranscribeSamples(options, true);
+                  fullTranscribeSamples(true);
                 }
                 break;
               }
 
               // Append to buffer
-              short[] shortBuffer = shortBufferSlices.get(sliceIndex);
               if (nSamples + n > audioSliceSec * SAMPLE_RATE) {
                 Log.d(NAME, "next slice");
 
                 sliceIndex++;
                 nSamples = 0;
-                shortBuffer = new short[audioSliceSec * SAMPLE_RATE];
-                shortBufferSlices.add(shortBuffer);
                 sliceNSamples.add(0);
               }
+              putPcmData(jobId, buffer, sliceIndex, nSamples, n);
 
-              for (int i = 0; i < n; i++) {
-                shortBuffer[nSamples + i] = buffer[i];
-              }
-
-              boolean isSpeech = vad(options, shortBuffer, nSamples, n);
+              boolean isSpeech = vad(sliceIndex, nSamples, n);
 
               nSamples += n;
               sliceNSamples.set(sliceIndex, nSamples);
 
-              if (!isSpeech) continue;
+              boolean isSamplesEnough = nSamples / SAMPLE_RATE >= audioMinSec;
+              if (!isSamplesEnough || !isSpeech) continue;
 
               if (!isTranscribing && nSamples > SAMPLE_RATE / 2) {
                 isTranscribing = true;
                 fullHandler = new Thread(new Runnable() {
                   @Override
                   public void run() {
-                    fullTranscribeSamples(options, false);
+                    fullTranscribeSamples(false);
                   }
                 });
                 fullHandler.start();
@@ -210,11 +194,9 @@ public class WhisperContext {
               Log.e(NAME, "Error transcribing realtime: " + e.getMessage());
             }
           }
-          // TODO: Append in real time so we don't need to keep all slices & also reduce memory usage
-          Log.d(NAME, "Begin saving wav file to " + audioOutputPath);
-          AudioUtils.saveWavFile(AudioUtils.concatShortBuffers(shortBufferSlices), audioOutputPath);
+
           if (!isTranscribing) {
-            emitTranscribeEvent("@RNWhisper_onRealtimeTranscribeEnd", Arguments.createMap());
+            finishRealtimeTranscribe(Arguments.createMap());
           }
           if (fullHandler != null) {
             fullHandler.join(); // Wait for full transcribe to finish
@@ -232,26 +214,16 @@ public class WhisperContext {
     return state;
   }
 
-  private void fullTranscribeSamples(ReadableMap options, boolean skipCapturingCheck) {
+  private void fullTranscribeSamples(boolean skipCapturingCheck) {
     int nSamplesOfIndex = sliceNSamples.get(transcribeSliceIndex);
 
     if (!isCapturing && !skipCapturingCheck) return;
 
-    short[] shortBuffer = shortBufferSlices.get(transcribeSliceIndex);
-    int nSamples = sliceNSamples.get(transcribeSliceIndex);
-
     nSamplesTranscribing = nSamplesOfIndex;
-
-    // convert I16 to F32
-    float[] nSamplesBuffer32 = new float[nSamplesTranscribing];
-    for (int i = 0; i < nSamplesTranscribing; i++) {
-      nSamplesBuffer32[i] = shortBuffer[i] / 32768.0f;
-    }
-
     Log.d(NAME, "Start transcribing realtime: " + nSamplesTranscribing);
 
     int timeStart = (int) System.currentTimeMillis();
-    int code = full(jobId, options, nSamplesBuffer32, nSamplesTranscribing);
+    int code = fullWithJob(jobId, context, transcribeSliceIndex, nSamplesTranscribing);
     int timeEnd = (int) System.currentTimeMillis();
     int timeRecording = (int) (nSamplesTranscribing / SAMPLE_RATE * 1000);
 
@@ -264,7 +236,7 @@ public class WhisperContext {
 
     if (code == 0) {
       payload.putMap("data", getTextSegments(0, getTextSegmentCount(context)));
-    } else {
+    } else if (code != -999) { // Not aborted
       payload.putString("error", "Transcribe failed with code " + code);
     }
 
@@ -283,12 +255,12 @@ public class WhisperContext {
       nSamplesTranscribing = 0;
     }
 
-    boolean continueNeeded = !isCapturing && nSamplesTranscribing != nSamplesOfIndex;
+    boolean continueNeeded = !isCapturing && nSamplesTranscribing != nSamplesOfIndex && code != -999;
 
     if (isStopped && !continueNeeded) {
       payload.putBoolean("isCapturing", false);
       payload.putBoolean("isStoppedByAction", isStoppedByAction);
-      emitTranscribeEvent("@RNWhisper_onRealtimeTranscribeEnd", payload);
+      finishRealtimeTranscribe(payload);
     } else if (code == 0) {
       payload.putBoolean("isCapturing", true);
       emitTranscribeEvent("@RNWhisper_onRealtimeTranscribe", payload);
@@ -299,7 +271,7 @@ public class WhisperContext {
 
     if (continueNeeded) {
       // If no more capturing, continue transcribing until all slices are transcribed
-      fullTranscribeSamples(options, true);
+      fullTranscribeSamples(true);
     } else if (isStopped) {
       // No next, cleanup
       rewind();
@@ -360,71 +332,39 @@ public class WhisperContext {
     }
   }
 
-  public WritableMap transcribeInputStream(int jobId, InputStream inputStream, ReadableMap options) throws IOException, Exception {
+  public WritableMap transcribe(int jobId, float[] audioData, ReadableMap options) throws IOException, Exception {
     if (isCapturing || isTranscribing) {
       throw new Exception("Context is already in capturing or transcribing");
     }
     rewind();
-
     this.jobId = jobId;
+    this.isTdrzEnable = options.hasKey("tdrzEnable") && options.getBoolean("tdrzEnable");
+
     isTranscribing = true;
-    float[] audioData = AudioUtils.decodeWaveFile(inputStream);
-    int code = full(jobId, options, audioData, audioData.length);
-    isTranscribing = false;
-    this.jobId = -1;
-    if (code != 0) {
-      throw new Exception("Failed to transcribe the file. Code: " + code);
-    }
-    WritableMap result = getTextSegments(0, getTextSegmentCount(context));
-    result.putBoolean("isAborted", isStoppedByAction);
-    return result;
-  }
 
-  private int full(int jobId, ReadableMap options, float[] audioData, int audioDataLen) {
     boolean hasProgressCallback = options.hasKey("onProgress") && options.getBoolean("onProgress");
     boolean hasNewSegmentsCallback = options.hasKey("onNewSegments") && options.getBoolean("onNewSegments");
-    return fullTranscribe(
+    int code = fullWithNewJob(
       jobId,
       context,
       // float[] audio_data,
       audioData,
       // jint audio_data_len,
-      audioDataLen,
-      // jint n_threads,
-      options.hasKey("maxThreads") ? options.getInt("maxThreads") : -1,
-      // jint max_context,
-      options.hasKey("maxContext") ? options.getInt("maxContext") : -1,
-
-      // jint word_thold,
-      options.hasKey("wordThold") ? options.getInt("wordThold") : -1,
-      // jint max_len,
-      options.hasKey("maxLen") ? options.getInt("maxLen") : -1,
-      // jboolean token_timestamps,
-      options.hasKey("tokenTimestamps") ? options.getBoolean("tokenTimestamps") : false,
-  
-      // jint offset,
-      options.hasKey("offset") ? options.getInt("offset") : -1,
-      // jint duration,
-      options.hasKey("duration") ? options.getInt("duration") : -1,
-      // jfloat temperature,
-      options.hasKey("temperature") ? (float) options.getDouble("temperature") : -1.0f,
-      // jfloat temperature_inc,
-      options.hasKey("temperatureInc") ? (float) options.getDouble("temperatureInc") : -1.0f,
-      // jint beam_size,
-      options.hasKey("beamSize") ? options.getInt("beamSize") : -1,
-      // jint best_of,
-      options.hasKey("bestOf") ? options.getInt("bestOf") : -1,
-      // jboolean speed_up,
-      options.hasKey("speedUp") ? options.getBoolean("speedUp") : false,
-      // jboolean translate,
-      options.hasKey("translate") ? options.getBoolean("translate") : false,
-      // jstring language,
-      options.hasKey("language") ? options.getString("language") : "auto",
-      // jstring prompt
-      options.hasKey("prompt") ? options.getString("prompt") : null,
+      audioData.length,
+      // ReadableMap options,
+      options,
       // Callback callback
       hasProgressCallback || hasNewSegmentsCallback ? new Callback(this, hasProgressCallback, hasNewSegmentsCallback) : null
     );
+
+    isTranscribing = false;
+    this.jobId = -1;
+    if (code != 0 && code != 999) {
+      throw new Exception("Failed to transcribe the file. Code: " + code);
+    }
+    WritableMap result = getTextSegments(0, getTextSegmentCount(context));
+    result.putBoolean("isAborted", isStoppedByAction);
+    return result;
   }
 
   private WritableMap getTextSegments(int start, int count) {
@@ -432,8 +372,15 @@ public class WhisperContext {
 
     WritableMap data = Arguments.createMap();
     WritableArray segments = Arguments.createArray();
+
     for (int i = 0; i < count; i++) {
       String text = getTextSegment(context, i);
+
+      // If tdrzEnable is enabled and speaker turn is detected
+      if (this.isTdrzEnable && getTextSegmentSpeakerTurnNext(context, i)) {
+          text += " [SPEAKER_TURN]";
+      }
+
       builder.append(text);
 
       WritableMap segment = Arguments.createMap();
@@ -475,6 +422,10 @@ public class WhisperContext {
     stopTranscribe(this.jobId);
   }
 
+  public String bench(int n_threads) {
+    return bench(context, n_threads);
+  }
+
   public void release() {
     stopCurrentTranscribe();
     freeContext(context);
@@ -543,31 +494,18 @@ public class WhisperContext {
     }
   }
 
-
+  // JNI methods
   protected static native long initContext(String modelPath);
   protected static native long initContextWithAsset(AssetManager assetManager, String modelPath);
   protected static native long initContextWithInputStream(PushbackInputStream inputStream);
-  protected static native boolean vadSimple(float[] audio_data, int audio_data_len, float vad_thold, float vad_freq_thold);
-  protected static native int fullTranscribe(
+  protected static native void freeContext(long contextPtr);
+
+  protected static native int fullWithNewJob(
     int job_id,
     long context,
     float[] audio_data,
     int audio_data_len,
-    int n_threads,
-    int max_context,
-    int word_thold,
-    int max_len,
-    boolean token_timestamps,
-    int offset,
-    int duration,
-    float temperature,
-    float temperature_inc,
-    int beam_size,
-    int best_of,
-    boolean speed_up,
-    boolean translate,
-    String language,
-    String prompt,
+    ReadableMap options,
     Callback Callback
   );
   protected static native void abortTranscribe(int jobId);
@@ -576,5 +514,21 @@ public class WhisperContext {
   protected static native String getTextSegment(long context, int index);
   protected static native int getTextSegmentT0(long context, int index);
   protected static native int getTextSegmentT1(long context, int index);
-  protected static native void freeContext(long contextPtr);
-}
+  protected static native boolean getTextSegmentSpeakerTurnNext(long context, int index);
+
+  protected static native void createRealtimeTranscribeJob(
+    int job_id,
+    long context,
+    ReadableMap options
+  );
+  protected static native void finishRealtimeTranscribeJob(int job_id, long context, int[] sliceNSamples);
+  protected static native boolean vadSimple(int job_id, int slice_index, int n_samples, int n);
+  protected static native void putPcmData(int job_id, short[] buffer, int slice_index, int n_samples, int n);
+  protected static native int fullWithJob(
+    int job_id,
+    long context,
+    int slice_index,
+    int n_samples
+  );
+  protected static native String bench(long context, int n_threads);
+}

package/android/src/main/jni-utils.h

@@ -0,0 +1,76 @@
+#include <jni.h>
+
+// ReadableMap utils
+
+namespace readablemap {
+
+bool hasKey(JNIEnv *env, jobject readableMap, const char *key) {
+    jclass mapClass = env->GetObjectClass(readableMap);
+    jmethodID hasKeyMethod = env->GetMethodID(mapClass, "hasKey", "(Ljava/lang/String;)Z");
+    jstring jKey = env->NewStringUTF(key);
+    jboolean result = env->CallBooleanMethod(readableMap, hasKeyMethod, jKey);
+    env->DeleteLocalRef(jKey);
+    return result;
+}
+
+int getInt(JNIEnv *env, jobject readableMap, const char *key, jint defaultValue) {
+    if (!hasKey(env, readableMap, key)) {
+        return defaultValue;
+    }
+    jclass mapClass = env->GetObjectClass(readableMap);
+    jmethodID getIntMethod = env->GetMethodID(mapClass, "getInt", "(Ljava/lang/String;)I");
+    jstring jKey = env->NewStringUTF(key);
+    jint result = env->CallIntMethod(readableMap, getIntMethod, jKey);
+    env->DeleteLocalRef(jKey);
+    return result;
+}
+
+bool getBool(JNIEnv *env, jobject readableMap, const char *key, jboolean defaultValue) {
+    if (!hasKey(env, readableMap, key)) {
+        return defaultValue;
+    }
+    jclass mapClass = env->GetObjectClass(readableMap);
+    jmethodID getBoolMethod = env->GetMethodID(mapClass, "getBoolean", "(Ljava/lang/String;)Z");
+    jstring jKey = env->NewStringUTF(key);
+    jboolean result = env->CallBooleanMethod(readableMap, getBoolMethod, jKey);
+    env->DeleteLocalRef(jKey);
+    return result;
+}
+
+long getLong(JNIEnv *env, jobject readableMap, const char *key, jlong defaultValue) {
+    if (!hasKey(env, readableMap, key)) {
+        return defaultValue;
+    }
+    jclass mapClass = env->GetObjectClass(readableMap);
+    jmethodID getLongMethod = env->GetMethodID(mapClass, "getLong", "(Ljava/lang/String;)J");
+    jstring jKey = env->NewStringUTF(key);
+    jlong result = env->CallLongMethod(readableMap, getLongMethod, jKey);
+    env->DeleteLocalRef(jKey);
+    return result;
+}
+
+float getFloat(JNIEnv *env, jobject readableMap, const char *key, jfloat defaultValue) {
+    if (!hasKey(env, readableMap, key)) {
+        return defaultValue;
+    }
+    jclass mapClass = env->GetObjectClass(readableMap);
+    jmethodID getFloatMethod = env->GetMethodID(mapClass, "getDouble", "(Ljava/lang/String;)D");
+    jstring jKey = env->NewStringUTF(key);
+    jfloat result = env->CallDoubleMethod(readableMap, getFloatMethod, jKey);
+    env->DeleteLocalRef(jKey);
+    return result;
+}
+
+jstring getString(JNIEnv *env, jobject readableMap, const char *key, jstring defaultValue) {
+    if (!hasKey(env, readableMap, key)) {
+        return defaultValue;
+    }
+    jclass mapClass = env->GetObjectClass(readableMap);
+    jmethodID getStringMethod = env->GetMethodID(mapClass, "getString", "(Ljava/lang/String;)Ljava/lang/String;");
+    jstring jKey = env->NewStringUTF(key);
+    jstring result = (jstring) env->CallObjectMethod(readableMap, getStringMethod, jKey);
+    env->DeleteLocalRef(jKey);
+    return result;
+}
+
+}