whisper.rn 0.4.0-rc.1 → 0.4.0-rc.11

package/android/src/main/jni.cpp

@@ -10,6 +10,7 @@
 #include "whisper.h"
 #include "rn-whisper.h"
 #include "ggml.h"
+#include "jni-utils.h"
 
 #define UNUSED(x) (void)(x)
 #define TAG "JNI"
@@ -96,7 +97,8 @@ static void input_stream_close(void *ctx) {
 
 static struct whisper_context *whisper_init_from_input_stream(
     JNIEnv *env,
-    jobject input_stream // PushbackInputStream
+    jobject input_stream, // PushbackInputStream
+    struct whisper_context_params cparams
 ) {
     input_stream_context *context = new input_stream_context;
     context->env = env;
@@ -108,7 +110,7 @@ static struct whisper_context *whisper_init_from_input_stream(
         .eof = &input_stream_is_eof,
         .close = &input_stream_close
     };
-    return whisper_init(&loader);
+    return whisper_init_with_params(&loader, cparams);
 }
 
 // Load model from asset
@@ -127,7 +129,8 @@ static void asset_close(void *ctx) {
 static struct whisper_context *whisper_init_from_asset(
     JNIEnv *env,
     jobject assetManager,
-    const char *asset_path
+    const char *asset_path,
+    struct whisper_context_params cparams
 ) {
     LOGI("Loading model from asset '%s'\n", asset_path);
     AAssetManager *asset_manager = AAssetManager_fromJava(env, assetManager);
@@ -142,7 +145,7 @@ static struct whisper_context *whisper_init_from_asset(
         .eof = &asset_is_eof,
         .close = &asset_close
     };
-    return whisper_init(&loader);
+    return whisper_init_with_params(&loader, cparams);
 }
 
 extern "C" {
@@ -151,9 +154,15 @@ JNIEXPORT jlong JNICALL
 Java_com_rnwhisper_WhisperContext_initContext(
         JNIEnv *env, jobject thiz, jstring model_path_str) {
     UNUSED(thiz);
+    struct whisper_context_params cparams;
+
+    // TODO: Expose dtw_token_timestamps and dtw_aheads_preset
+    cparams.dtw_token_timestamps = false;
+    // cparams.dtw_aheads_preset = WHISPER_AHEADS_BASE;
+
     struct whisper_context *context = nullptr;
     const char *model_path_chars = env->GetStringUTFChars(model_path_str, nullptr);
-    context = whisper_init_from_file(model_path_chars);
+    context = whisper_init_from_file_with_params(model_path_chars, cparams);
     env->ReleaseStringUTFChars(model_path_str, model_path_chars);
     return reinterpret_cast<jlong>(context);
 }
@@ -166,9 +175,15 @@ Java_com_rnwhisper_WhisperContext_initContextWithAsset(
     jstring model_path_str
 ) {
     UNUSED(thiz);
+    struct whisper_context_params cparams;
+
+    // TODO: Expose dtw_token_timestamps and dtw_aheads_preset
+    cparams.dtw_token_timestamps = false;
+    // cparams.dtw_aheads_preset = WHISPER_AHEADS_BASE;
+
     struct whisper_context *context = nullptr;
     const char *model_path_chars = env->GetStringUTFChars(model_path_str, nullptr);
-    context = whisper_init_from_asset(env, asset_manager, model_path_chars);
+    context = whisper_init_from_asset(env, asset_manager, model_path_chars, cparams);
     env->ReleaseStringUTFChars(model_path_str, model_path_chars);
     return reinterpret_cast<jlong>(context);
 }
@@ -180,30 +195,70 @@ Java_com_rnwhisper_WhisperContext_initContextWithInputStream(
     jobject input_stream
 ) {
     UNUSED(thiz);
+    struct whisper_context_params cparams;
+
+    // TODO: Expose dtw_token_timestamps and dtw_aheads_preset
+    cparams.dtw_token_timestamps = false;
+    // cparams.dtw_aheads_preset = WHISPER_AHEADS_BASE;
+
     struct whisper_context *context = nullptr;
-    context = whisper_init_from_input_stream(env, input_stream);
+    context = whisper_init_from_input_stream(env, input_stream, cparams);
     return reinterpret_cast<jlong>(context);
 }
 
-JNIEXPORT jboolean JNICALL
-Java_com_rnwhisper_WhisperContext_vadSimple(
-    JNIEnv *env,
-    jobject thiz,
-    jfloatArray audio_data,
-    jint audio_data_len,
-    jfloat vad_thold,
-    jfloat vad_freq_thold
-) {
-    UNUSED(thiz);
 
-    std::vector<float> samples(audio_data_len);
-    jfloat *audio_data_arr = env->GetFloatArrayElements(audio_data, nullptr);
-    for (int i = 0; i < audio_data_len; i++) {
-        samples[i] = audio_data_arr[i];
+struct whisper_full_params createFullParams(JNIEnv *env, jobject options) {
+    struct whisper_full_params params = whisper_full_default_params(WHISPER_SAMPLING_GREEDY);
+
+    params.print_realtime = false;
+    params.print_progress = false;
+    params.print_timestamps = false;
+    params.print_special = false;
+
+    int max_threads = std::thread::hardware_concurrency();
+    // Use 2 threads by default on 4-core devices, 4 threads on more cores
+    int default_n_threads = max_threads == 4 ? 2 : min(4, max_threads);
+    int n_threads = readablemap::getInt(env, options, "maxThreads", default_n_threads);
+    params.n_threads = n_threads > 0 ? n_threads : default_n_threads;
+    params.translate = readablemap::getBool(env, options, "translate", false);
+    params.token_timestamps = readablemap::getBool(env, options, "tokenTimestamps", false);
+    params.tdrz_enable = readablemap::getBool(env, options, "tdrzEnable", false);
+    params.offset_ms = 0;
+    params.no_context = true;
+    params.single_segment = false;
+
+    int beam_size = readablemap::getInt(env, options, "beamSize", -1);
+    if (beam_size > -1) {
+        params.strategy = WHISPER_SAMPLING_BEAM_SEARCH;
+        params.beam_search.beam_size = beam_size;
     }
-    bool is_speech = rn_whisper_vad_simple(samples, WHISPER_SAMPLE_RATE, 1000, vad_thold, vad_freq_thold, false);
-    env->ReleaseFloatArrayElements(audio_data, audio_data_arr, JNI_ABORT);
-    return is_speech;
+    int best_of = readablemap::getInt(env, options, "bestOf", -1);
+    if (best_of > -1) params.greedy.best_of = best_of;
+    int max_len = readablemap::getInt(env, options, "maxLen", -1);
+    if (max_len > -1) params.max_len = max_len;
+    int max_context = readablemap::getInt(env, options, "maxContext", -1);
+    if (max_context > -1) params.n_max_text_ctx = max_context;
+    int offset = readablemap::getInt(env, options, "offset", -1);
+    if (offset > -1) params.offset_ms = offset;
+    int duration = readablemap::getInt(env, options, "duration", -1);
+    if (duration > -1) params.duration_ms = duration;
+    int word_thold = readablemap::getInt(env, options, "wordThold", -1);
+    if (word_thold > -1) params.thold_pt = word_thold;
+    float temperature = readablemap::getFloat(env, options, "temperature", -1);
+    if (temperature > -1) params.temperature = temperature;
+    float temperature_inc = readablemap::getFloat(env, options, "temperatureInc", -1);
+    if (temperature_inc > -1) params.temperature_inc = temperature_inc;
+    jstring prompt = readablemap::getString(env, options, "prompt", nullptr);
+    if (prompt != nullptr) {
+        params.initial_prompt = env->GetStringUTFChars(prompt, nullptr);
+        env->DeleteLocalRef(prompt);
+    }
+    jstring language = readablemap::getString(env, options, "language", nullptr);
+    if (language != nullptr) {
+        params.language = env->GetStringUTFChars(language, nullptr);
+        env->DeleteLocalRef(language);
+    }
+    return params;
 }
 
 struct callback_context {
@@ -212,101 +267,23 @@ struct callback_context {
 };
 
 JNIEXPORT jint JNICALL
-Java_com_rnwhisper_WhisperContext_fullTranscribe(
+Java_com_rnwhisper_WhisperContext_fullWithNewJob(
     JNIEnv *env,
     jobject thiz,
     jint job_id,
     jlong context_ptr,
     jfloatArray audio_data,
     jint audio_data_len,
-    jint n_threads,
-    jint max_context,
-    int word_thold,
-    int max_len,
-    jboolean token_timestamps,
-    jint offset,
-    jint duration,
-    jfloat temperature,
-    jfloat temperature_inc,
-    jint beam_size,
-    jint best_of,
-    jboolean speed_up,
-    jboolean translate,
-    jstring language,
-    jstring prompt,
+    jobject options,
     jobject callback_instance
 ) {
     UNUSED(thiz);
     struct whisper_context *context = reinterpret_cast<struct whisper_context *>(context_ptr);
     jfloat *audio_data_arr = env->GetFloatArrayElements(audio_data, nullptr);
 
-    int max_threads = std::thread::hardware_concurrency();
-    // Use 2 threads by default on 4-core devices, 4 threads on more cores
-    int default_n_threads = max_threads == 4 ? 2 : min(4, max_threads);
-
     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;
-    }
-
-    params.print_realtime = false;
-    params.print_progress = false;
-    params.print_timestamps = false;
-    params.print_special = false;
-    params.translate = translate;
-    const char *language_chars = env->GetStringUTFChars(language, nullptr);
-    params.language = language_chars;
-    params.n_threads = n_threads > 0 ? n_threads : default_n_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 (offset > -1) {
-        params.offset_ms = offset;
-    }
-    if (duration > -1) {
-        params.duration_ms = duration;
-    }
-    if (word_thold > -1) {
-        params.thold_pt = word_thold;
-    }
-    if (temperature > -1) {
-        params.temperature = temperature;
-    }
-    if (temperature_inc > -1) {
-        params.temperature_inc = temperature_inc;
-    }
-    if (prompt != nullptr) {
-        params.initial_prompt = env->GetStringUTFChars(prompt, nullptr);
-    }
-
-    // abort handlers
-    params.encoder_begin_callback = [](struct whisper_context * /*ctx*/, struct whisper_state * /*state*/, void * user_data) {
-        bool is_aborted = *(bool*)user_data;
-        return !is_aborted;
-    };
-    params.encoder_begin_callback_user_data = rn_whisper_assign_abort_map(job_id);
-    params.abort_callback = [](void * user_data) {
-        bool is_aborted = *(bool*)user_data;
-        return is_aborted;
-    };
-    params.abort_callback_user_data = rn_whisper_assign_abort_map(job_id);
+    whisper_full_params params = createFullParams(env, options);
 
     if (callback_instance != nullptr) {
         callback_context *cb_ctx = new callback_context;
@@ -334,6 +311,8 @@ Java_com_rnwhisper_WhisperContext_fullTranscribe(
         params.new_segment_callback_user_data = cb_ctx;
     }
 
+    rnwhisper::job* job = rnwhisper::job_new(job_id, params);
+
     LOGI("About to reset timings");
     whisper_reset_timings(context);
 
@@ -343,8 +322,123 @@ Java_com_rnwhisper_WhisperContext_fullTranscribe(
         // whisper_print_timings(context);
     }
     env->ReleaseFloatArrayElements(audio_data, audio_data_arr, JNI_ABORT);
-    env->ReleaseStringUTFChars(language, language_chars);
-    rn_whisper_remove_abort_map(job_id);
+
+    if (job->is_aborted()) code = -999;
+    rnwhisper::job_remove(job_id);
+    return code;
+}
+
+JNIEXPORT void JNICALL
+Java_com_rnwhisper_WhisperContext_createRealtimeTranscribeJob(
+    JNIEnv *env,
+    jobject thiz,
+    jint job_id,
+    jlong context_ptr,
+    jobject options
+) {
+    whisper_full_params params = createFullParams(env, options);
+    rnwhisper::job* job = rnwhisper::job_new(job_id, params);
+    rnwhisper::vad_params vad;
+    vad.use_vad = readablemap::getBool(env, options, "useVad", false);
+    vad.vad_ms = readablemap::getInt(env, options, "vadMs", 2000);
+    vad.vad_thold = readablemap::getFloat(env, options, "vadThold", 0.6f);
+    vad.freq_thold = readablemap::getFloat(env, options, "vadFreqThold", 100.0f);
+
+    jstring audio_output_path = readablemap::getString(env, options, "audioOutputPath", nullptr);
+    const char* audio_output_path_str = nullptr;
+    if (audio_output_path != nullptr) {
+        audio_output_path_str = env->GetStringUTFChars(audio_output_path, nullptr);
+        env->DeleteLocalRef(audio_output_path);
+    }
+    job->set_realtime_params(
+        vad,
+        readablemap::getInt(env, options, "realtimeAudioSec", 0),
+        readablemap::getInt(env, options, "realtimeAudioSliceSec", 0),
+        readablemap::getFloat(env, options, "realtimeAudioMinSec", 0),
+        audio_output_path_str
+    );
+}
+
+JNIEXPORT void JNICALL
+Java_com_rnwhisper_WhisperContext_finishRealtimeTranscribeJob(
+    JNIEnv *env,
+    jobject thiz,
+    jint job_id,
+    jlong context_ptr,
+    jintArray slice_n_samples
+) {
+    UNUSED(env);
+    UNUSED(thiz);
+    UNUSED(context_ptr);
+
+    rnwhisper::job *job = rnwhisper::job_get(job_id);
+    if (job->audio_output_path != nullptr) {
+        RNWHISPER_LOG_INFO("job->params.language: %s\n", job->params.language);
+        std::vector<int> slice_n_samples_vec;
+        jint *slice_n_samples_arr = env->GetIntArrayElements(slice_n_samples, nullptr);
+        slice_n_samples_vec = std::vector<int>(slice_n_samples_arr, slice_n_samples_arr + env->GetArrayLength(slice_n_samples));
+        env->ReleaseIntArrayElements(slice_n_samples, slice_n_samples_arr, JNI_ABORT);
+
+        // TODO: Append in real time so we don't need to keep all slices & also reduce memory usage
+        rnaudioutils::save_wav_file(
+            rnaudioutils::concat_short_buffers(job->pcm_slices, slice_n_samples_vec),
+            job->audio_output_path
+        );
+    }
+    rnwhisper::job_remove(job_id);
+}
+
+JNIEXPORT jboolean JNICALL
+Java_com_rnwhisper_WhisperContext_vadSimple(
+    JNIEnv *env,
+    jobject thiz,
+    jint job_id,
+    jint slice_index,
+    jint n_samples,
+    jint n
+) {
+    UNUSED(thiz);
+    rnwhisper::job* job = rnwhisper::job_get(job_id);
+    return job->vad_simple(slice_index, n_samples, n);
+}
+
+JNIEXPORT void JNICALL
+Java_com_rnwhisper_WhisperContext_putPcmData(
+    JNIEnv *env,
+    jobject thiz,
+    jint job_id,
+    jshortArray pcm,
+    jint slice_index,
+    jint n_samples,
+    jint n
+) {
+    UNUSED(thiz);
+    rnwhisper::job* job = rnwhisper::job_get(job_id);
+    jshort *pcm_arr = env->GetShortArrayElements(pcm, nullptr);
+    job->put_pcm_data(pcm_arr, slice_index, n_samples, n);
+    env->ReleaseShortArrayElements(pcm, pcm_arr, JNI_ABORT);
+}
+
+JNIEXPORT jint JNICALL
+Java_com_rnwhisper_WhisperContext_fullWithJob(
+    JNIEnv *env,
+    jobject thiz,
+    jint job_id,
+    jlong context_ptr,
+    jint slice_index,
+    jint n_samples
+) {
+    UNUSED(thiz);
+    struct whisper_context *context = reinterpret_cast<struct whisper_context *>(context_ptr);
+
+    rnwhisper::job* job = rnwhisper::job_get(job_id);
+    float* pcmf32 = job->pcm_slice_to_f32(slice_index, n_samples);
+    int code = whisper_full(context, job->params, pcmf32, n_samples);
+    free(pcmf32);
+    if (code == 0) {
+        // whisper_print_timings(context);
+    }
+    if (job->is_aborted()) code = -999;
     return code;
 }
 
@@ -355,7 +449,8 @@ Java_com_rnwhisper_WhisperContext_abortTranscribe(
     jint job_id
 ) {
     UNUSED(thiz);
-    rn_whisper_abort_transcribe(job_id);
+    rnwhisper::job *job = rnwhisper::job_get(job_id);
+    if (job) job->abort();
 }
 
 JNIEXPORT void JNICALL
@@ -364,7 +459,7 @@ Java_com_rnwhisper_WhisperContext_abortAllTranscribe(
     jobject thiz
 ) {
     UNUSED(thiz);
-    rn_whisper_abort_all_transcribe();
+    rnwhisper::job_abort_all();
 }
 
 JNIEXPORT jint JNICALL
@@ -413,4 +508,26 @@ Java_com_rnwhisper_WhisperContext_freeContext(
     whisper_free(context);
 }
 
+JNIEXPORT jboolean JNICALL
+Java_com_rnwhisper_WhisperContext_getTextSegmentSpeakerTurnNext(
+        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_speaker_turn_next(context, index);
+}
+
+JNIEXPORT jstring JNICALL
+Java_com_rnwhisper_WhisperContext_bench(
+    JNIEnv *env,
+    jobject thiz,
+    jlong context_ptr,
+    jint n_threads
+) {
+    UNUSED(thiz);
+    struct whisper_context *context = reinterpret_cast<struct whisper_context *>(context_ptr);
+    std::string result = rnwhisper::bench(context, n_threads);
+    return env->NewStringUTF(result.c_str());
+}
+
 } // extern "C"

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

@@ -47,6 +47,11 @@ public class RNWhisperModule extends NativeRNWhisperSpec {
     rnwhisper.transcribeFile(id, jobId, filePath, options, promise);
   }
 
+  @ReactMethod
+  public void transcribeData(double id, double jobId, String dataBase64, ReadableMap options, Promise promise) {
+    rnwhisper.transcribeData(id, jobId, dataBase64, options, promise);
+  }
+
   @ReactMethod
   public void startRealtimeTranscribe(double id, double jobId, ReadableMap options, Promise promise) {
     rnwhisper.startRealtimeTranscribe(id, jobId, options, promise);
@@ -57,6 +62,11 @@ public class RNWhisperModule extends NativeRNWhisperSpec {
     rnwhisper.abortTranscribe(contextId, jobId, promise);
   }
 
+  @ReactMethod
+  public void bench(double id, double nThreads, Promise promise) {
+    rnwhisper.bench(id, nThreads, promise);
+  }
+
   @ReactMethod
   public void releaseContext(double id, Promise promise) {
     rnwhisper.releaseContext(id, promise);

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

@@ -47,6 +47,11 @@ public class RNWhisperModule extends ReactContextBaseJavaModule {
     rnwhisper.transcribeFile(id, jobId, filePath, options, promise);
   }
 
+  @ReactMethod
+  public void transcribeData(double id, double jobId, String dataBase64, ReadableMap options, Promise promise) {
+    rnwhisper.transcribeData(id, jobId, dataBase64, options, promise);
+  }
+
   @ReactMethod
   public void startRealtimeTranscribe(double id, double jobId, ReadableMap options, Promise promise) {
     rnwhisper.startRealtimeTranscribe(id, jobId, options, promise);
@@ -57,6 +62,11 @@ public class RNWhisperModule extends ReactContextBaseJavaModule {
     rnwhisper.abortTranscribe(contextId, jobId, promise);
   }
 
+  @ReactMethod
+  public void bench(double id, double nThreads, Promise promise) {
+    rnwhisper.bench(id, nThreads, promise);
+  }
+
   @ReactMethod
   public void releaseContext(double id, Promise promise) {
     rnwhisper.releaseContext(id, promise);

package/cpp/coreml/whisper-encoder-impl.h

@@ -123,7 +123,7 @@ API_AVAILABLE(macos(12.0), ios(15.0), watchos(8.0), tvos(15.0)) __attribute__((v
 
 /**
     Make a prediction using the convenience interface
-    @param logmel_data as 1 × 80 × 3000 3-dimensional array of floats:
+    @param logmel_data as 1 × n_mel × 3000 3-dimensional array of floats:
     @param error If an error occurs, upon return contains an NSError object that describes the problem. If you are not interested in possible errors, pass in NULL.
     @return the prediction as whisper_encoder_implOutput
 */

package/cpp/coreml/whisper-encoder.h

@@ -3,6 +3,8 @@
 // Code is derived from the work of Github user @wangchou
 // ref: https://github.com/wangchou/callCoreMLFromCpp
 
+#include <stdint.h>
+
 #if __cplusplus
 extern "C" {
 #endif
@@ -14,6 +16,8 @@ void whisper_coreml_free(struct whisper_coreml_context * ctx);
 
 void whisper_coreml_encode(
         const whisper_coreml_context * ctx,
+                             int64_t   n_ctx,
+                             int64_t   n_mel,
                                float * mel,
                                float * out);
 

package/cpp/coreml/whisper-encoder.mm

@@ -24,7 +24,7 @@ struct whisper_coreml_context * whisper_coreml_init(const char * path_model) {
 
     // select which device to run the Core ML model on
     MLModelConfiguration *config = [[MLModelConfiguration alloc] init];
-    //config.computeUnits = MLComputeUnitsCPUAndGPU;
+    // config.computeUnits = MLComputeUnitsCPUAndGPU;
     //config.computeUnits = MLComputeUnitsCPUAndNeuralEngine;
     config.computeUnits = MLComputeUnitsAll;
 
@@ -48,13 +48,15 @@ void whisper_coreml_free(struct whisper_coreml_context * ctx) {
 
 void whisper_coreml_encode(
         const whisper_coreml_context * ctx,
+                             int64_t   n_ctx,
+                             int64_t   n_mel,
                                float * mel,
                                float * out) {
     MLMultiArray * inMultiArray = [
         [MLMultiArray alloc] initWithDataPointer: mel
-                                           shape: @[@1, @80, @3000]
+                                           shape: @[@1, @(n_mel), @(n_ctx)]
                                         dataType: MLMultiArrayDataTypeFloat32
-                                         strides: @[@(240000), @(3000), @1]
+                                         strides: @[@(n_ctx*n_mel), @(n_ctx), @1]
                                      deallocator: nil
                                            error: nil
     ];