whisper.rn 0.1.3 → 0.1.5

package/LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2023 Jhen <developer@jhen.me>
+Copyright (c) 2023 Jhen-Jie Hong
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
 in the Software without restriction, including without limitation the rights

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

@@ -12,6 +12,7 @@ import com.facebook.react.bridge.ReactContextBaseJavaModule;
 import com.facebook.react.bridge.ReactMethod;
 import com.facebook.react.bridge.LifecycleEventListener;
 import com.facebook.react.bridge.ReadableMap;
+import com.facebook.react.bridge.WritableMap;
 import com.facebook.react.module.annotations.ReactModule;
 
 import java.util.HashMap;
@@ -72,11 +73,11 @@ public class RNWhisperModule extends ReactContextBaseJavaModule implements Lifec
 
   @ReactMethod
   public void transcribe(int id, String filePath, ReadableMap options, Promise promise) {
-    new AsyncTask<Void, Void, String>() {
+    new AsyncTask<Void, Void, WritableMap>() {
       private Exception exception;
 
       @Override
-      protected String doInBackground(Void... voids) {
+      protected WritableMap doInBackground(Void... voids) {
         try {
           WhisperContext context = contexts.get(id);
           if (context == null) {
@@ -90,12 +91,12 @@ public class RNWhisperModule extends ReactContextBaseJavaModule implements Lifec
       }
 
       @Override
-      protected void onPostExecute(String result) {
+      protected void onPostExecute(WritableMap data) {
         if (exception != null) {
           promise.reject(exception);
           return;
         }
-        promise.resolve(result);
+        promise.resolve(data);
       }
     }.execute();
   }

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

@@ -1,5 +1,8 @@
 package com.rnwhisper;
 
+import com.facebook.react.bridge.Arguments;
+import com.facebook.react.bridge.WritableArray;
+import com.facebook.react.bridge.WritableMap;
 import com.facebook.react.bridge.ReadableMap;
 
 import android.util.Log;
@@ -29,7 +32,7 @@ public class WhisperContext {
     this.context = context;
   }
 
-  public String transcribe(final String filePath, final ReadableMap options) throws IOException, Exception {
+  public WritableMap transcribe(final String filePath, final ReadableMap options) throws IOException, Exception {
     int code = fullTranscribe(
       context,
       decodeWaveFile(new File(filePath)),
@@ -37,14 +40,18 @@ public class WhisperContext {
       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,
-      // jint word_thold,
-      options.hasKey("wordThold") ? options.getInt("wordThold") : -1,
       // jfloat temperature,
       options.hasKey("temperature") ? (float) options.getDouble("temperature") : -1.0f,
       // jfloat temperature_inc,
@@ -58,17 +65,31 @@ public class WhisperContext {
       // jboolean translate,
       options.hasKey("translate") ? options.getBoolean("translate") : false,
       // jstring language,
-      options.hasKey("language") ? options.getString("language") : "auto"
+      options.hasKey("language") ? options.getString("language") : "auto",
+      // jstring prompt
+      options.hasKey("prompt") ? options.getString("prompt") : null
     );
     if (code != 0) {
       throw new Exception("Transcription failed with code " + code);
     }
     Integer count = getTextSegmentCount(context);
     StringBuilder builder = new StringBuilder();
+
+    WritableMap data = Arguments.createMap();
+    WritableArray segments = Arguments.createArray();
     for (int i = 0; i < count; i++) {
-      builder.append(getTextSegment(context, i));
+      String text = getTextSegment(context, i);
+      builder.append(text);
+
+      WritableMap segment = Arguments.createMap();
+      segment.putString("text", text);
+      segment.putInt("t0", getTextSegmentT0(context, i));
+      segment.putInt("t1", getTextSegmentT1(context, i));
+      segments.pushMap(segment);
     }
-    return builder.toString();
+    data.putString("result", builder.toString());
+    data.putArray("segments", segments);
+    return data;
   }
 
   public void release() {
@@ -168,19 +189,23 @@ public class WhisperContext {
     float[] audio_data,
     int n_threads,
     int max_context,
+    int word_thold,
     int max_len,
+    boolean token_timestamps,
     int offset,
     int duration,
-    int word_thold,
     float temperature,
     float temperature_inc,
     int beam_size,
     int best_of,
     boolean speed_up,
     boolean translate,
-    String language
+    String language,
+    String prompt
   );
   protected static native int getTextSegmentCount(long context);
   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);
 }

package/android/src/main/jni/whisper/Whisper.mk

@@ -15,4 +15,5 @@ LOCAL_CFLAGS    += -DSTDC_HEADERS -std=c11 -I $(WHISPER_LIB_DIR)
 LOCAL_CPPFLAGS  += -std=c++11
 LOCAL_SRC_FILES := $(WHISPER_LIB_DIR)/ggml.c \
                    $(WHISPER_LIB_DIR)/whisper.cpp \
-                   $(LOCAL_PATH)/jni.c
+                   $(WHISPER_LIB_DIR)/rn-whisper.cpp \
+                   $(LOCAL_PATH)/jni.cpp

package/android/src/main/jni/whisper/{jni.c → jni.cpp}

@@ -2,10 +2,11 @@
 #include <android/asset_manager.h>
 #include <android/asset_manager_jni.h>
 #include <android/log.h>
-#include <stdlib.h>
+#include <cstdlib>
 #include <sys/sysinfo.h>
-#include <string.h>
+#include <string>
 #include "whisper.h"
+#include "rn-whisper.h"
 #include "ggml.h"
 
 #define UNUSED(x) (void)(x)
@@ -18,31 +19,17 @@ static inline int min(int a, int b) {
     return (a < b) ? a : b;
 }
 
-static inline int max(int a, int b) {
-    return (a > b) ? a : b;
-}
-
-static size_t asset_read(void *ctx, void *output, size_t read_size) {
-    return AAsset_read((AAsset *) ctx, output, read_size);
-}
-
-static bool asset_is_eof(void *ctx) {
-    return AAsset_getRemainingLength64((AAsset *) ctx) <= 0;
-}
-
-static void asset_close(void *ctx) {
-    AAsset_close((AAsset *) ctx);
-}
+extern "C" {
 
 JNIEXPORT jlong JNICALL
 Java_com_rnwhisper_WhisperContext_initContext(
         JNIEnv *env, jobject thiz, jstring model_path_str) {
     UNUSED(thiz);
-    struct whisper_context *context = NULL;
-    const char *model_path_chars = (*env)->GetStringUTFChars(env, model_path_str, NULL);
+    struct whisper_context *context = nullptr;
+    const char *model_path_chars = env->GetStringUTFChars(model_path_str, nullptr);
     context = whisper_init_from_file(model_path_chars);
-    (*env)->ReleaseStringUTFChars(env, model_path_str, model_path_chars);
-    return (jlong) context;
+    env->ReleaseStringUTFChars(model_path_str, model_path_chars);
+    return reinterpret_cast<jlong>(context);
 }
 
 JNIEXPORT jint JNICALL
@@ -53,29 +40,31 @@ Java_com_rnwhisper_WhisperContext_fullTranscribe(
     jfloatArray audio_data,
     jint n_threads,
     jint max_context,
-    jint max_len,
+    int word_thold,
+    int max_len,
+    jboolean token_timestamps,
     jint offset,
     jint duration,
-    jint word_thold,
     jfloat temperature,
     jfloat temperature_inc,
     jint beam_size,
     jint best_of,
     jboolean speed_up,
     jboolean translate,
-    jstring language
+    jstring language,
+    jstring prompt
 ) {
     UNUSED(thiz);
-    struct whisper_context *context = (struct whisper_context *) context_ptr;
-    jfloat *audio_data_arr = (*env)->GetFloatArrayElements(env, audio_data, NULL);
-    const jsize audio_data_length = (*env)->GetArrayLength(env, audio_data);
+    struct whisper_context *context = reinterpret_cast<struct whisper_context *>(context_ptr);
+    jfloat *audio_data_arr = env->GetFloatArrayElements(audio_data, nullptr);
+    const jsize audio_data_length = env->GetArrayLength(audio_data);
 
-    int max_threads = max(1, min(8, get_nprocs() - 2));
+    int max_threads = min(4, get_nprocs());
 
     LOGI("About to create params");
 
     struct whisper_full_params params = whisper_full_default_params(WHISPER_SAMPLING_GREEDY);
-    
+
     if (beam_size > -1) {
         params.strategy = WHISPER_SAMPLING_BEAM_SEARCH;
         params.beam_search.beam_size = beam_size;
@@ -86,22 +75,25 @@ Java_com_rnwhisper_WhisperContext_fullTranscribe(
     params.print_timestamps = false;
     params.print_special = false;
     params.translate = translate;
-    params.language = language;
+    const char *language_chars = env->GetStringUTFChars(language, nullptr);
+    params.language = language_chars;
     params.n_threads = n_threads > 0 ? n_threads : max_threads;
     params.speed_up = speed_up;
     params.offset_ms = 0;
     params.no_context = true;
     params.single_segment = false;
 
+    if (max_len > -1) {
+        params.max_len = max_len;
+    }
+    params.token_timestamps = token_timestamps;
+
     if (best_of > -1) {
         params.greedy.best_of = best_of;
     }
     if (max_context > -1) {
         params.n_max_text_ctx = max_context;
     }
-    if (max_len > -1) {
-        params.max_len = max_len;
-    }
     if (offset > -1) {
         params.offset_ms = offset;
     }
@@ -117,6 +109,13 @@ Java_com_rnwhisper_WhisperContext_fullTranscribe(
     if (temperature_inc > -1) {
         params.temperature_inc = temperature_inc;
     }
+    if (prompt != nullptr) {
+        rn_whisper_convert_prompt(
+            context,
+            params,
+            new std::string(env->GetStringUTFChars(prompt, nullptr))
+        );
+    }
 
     LOGI("About to reset timings");
     whisper_reset_timings(context);
@@ -126,7 +125,8 @@ Java_com_rnwhisper_WhisperContext_fullTranscribe(
     if (code == 0) {
         // whisper_print_timings(context);
     }
-    (*env)->ReleaseFloatArrayElements(env, audio_data, audio_data_arr, JNI_ABORT);
+    env->ReleaseFloatArrayElements(audio_data, audio_data_arr, JNI_ABORT);
+    env->ReleaseStringUTFChars(language, language_chars);
     return code;
 }
 
@@ -135,7 +135,7 @@ Java_com_rnwhisper_WhisperContext_getTextSegmentCount(
         JNIEnv *env, jobject thiz, jlong context_ptr) {
     UNUSED(env);
     UNUSED(thiz);
-    struct whisper_context *context = (struct whisper_context *) context_ptr;
+    struct whisper_context *context = reinterpret_cast<struct whisper_context *>(context_ptr);
     return whisper_full_n_segments(context);
 }
 
@@ -143,17 +143,37 @@ JNIEXPORT jstring JNICALL
 Java_com_rnwhisper_WhisperContext_getTextSegment(
         JNIEnv *env, jobject thiz, jlong context_ptr, jint index) {
     UNUSED(thiz);
-    struct whisper_context *context = (struct whisper_context *) context_ptr;
+    struct whisper_context *context = reinterpret_cast<struct whisper_context *>(context_ptr);
     const char *text = whisper_full_get_segment_text(context, index);
-    jstring string = (*env)->NewStringUTF(env, text);
+    jstring string = env->NewStringUTF(text);
     return string;
 }
 
+JNIEXPORT jint JNICALL
+Java_com_rnwhisper_WhisperContext_getTextSegmentT0(
+        JNIEnv *env, jobject thiz, jlong context_ptr, jint index) {
+    UNUSED(env);
+    UNUSED(thiz);
+    struct whisper_context *context = reinterpret_cast<struct whisper_context *>(context_ptr);
+    return whisper_full_get_segment_t0(context, index);
+}
+
+JNIEXPORT jint JNICALL
+Java_com_rnwhisper_WhisperContext_getTextSegmentT1(
+        JNIEnv *env, jobject thiz, jlong context_ptr, jint index) {
+    UNUSED(env);
+    UNUSED(thiz);
+    struct whisper_context *context = reinterpret_cast<struct whisper_context *>(context_ptr);
+    return whisper_full_get_segment_t1(context, index);
+}
+
 JNIEXPORT void JNICALL
 Java_com_rnwhisper_WhisperContext_freeContext(
         JNIEnv *env, jobject thiz, jlong context_ptr) {
     UNUSED(env);
     UNUSED(thiz);
-    struct whisper_context *context = (struct whisper_context *) context_ptr;
+    struct whisper_context *context = reinterpret_cast<struct whisper_context *>(context_ptr);
     whisper_free(context);
 }
+
+} // extern "C"

package/cpp/rn-whisper.cpp

@@ -0,0 +1,31 @@
+#include <cstdio>
+#include <string>
+#include <vector>
+#include "whisper.h"
+
+extern "C" {
+
+void rn_whisper_convert_prompt(
+  struct whisper_context * ctx,
+  struct whisper_full_params params,
+  std::string * prompt
+) {
+  std::vector<whisper_token> prompt_tokens;
+  if (!prompt->empty()) {
+    prompt_tokens.resize(1024);
+    prompt_tokens.resize(whisper_tokenize(ctx, prompt->c_str(), prompt_tokens.data(), prompt_tokens.size()));
+
+    // fprintf(stderr, "\n");
+    // fprintf(stderr, "initial prompt: '%s'\n", prompt->c_str());
+    // fprintf(stderr, "initial tokens: [ ");
+    // for (int i = 0; i < (int) prompt_tokens.size(); ++i) {
+    //     fprintf(stderr, "%d ", prompt_tokens[i]);
+    // }
+    // fprintf(stderr, "]\n");
+
+    params.prompt_tokens = prompt_tokens.data();
+    params.prompt_n_tokens = prompt_tokens.size();
+  }
+}
+
+}

package/cpp/rn-whisper.h

@@ -0,0 +1,16 @@
+
+#ifdef __cplusplus
+#include <string>
+#include <whisper.h>
+extern "C" {
+#endif
+
+void rn_whisper_convert_prompt(
+  struct whisper_context * ctx,
+  struct whisper_full_params params,
+  std::string * prompt
+);
+
+#ifdef __cplusplus
+}
+#endif