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

package/ios/RNWhisperContext.mm

@@ -1,30 +1,97 @@
 #import "RNWhisperContext.h"
-#import "RNWhisperAudioUtils.h"
+#import <Metal/Metal.h>
 #include <vector>
 
 #define NUM_BYTES_PER_BUFFER 16 * 1024
 
 @implementation RNWhisperContext
 
-+ (instancetype)initWithModelPath:(NSString *)modelPath contextId:(int)contextId noCoreML:(BOOL)noCoreML {
++ (instancetype)initWithModelPath:(NSString *)modelPath
+    contextId:(int)contextId
+    noCoreML:(BOOL)noCoreML
+    noMetal:(BOOL)noMetal
+    useFlashAttn:(BOOL)useFlashAttn
+{
     RNWhisperContext *context = [[RNWhisperContext alloc] init];
     context->contextId = contextId;
-#ifdef WHISPER_USE_COREML
-    if (noCoreML) {
-       context->ctx = whisper_init_from_file_no_coreml([modelPath UTF8String]);
-    } else {
-       context->ctx = whisper_init_from_file([modelPath UTF8String]);
+    struct whisper_context_params cparams;
+    NSString *reasonNoMetal = @"";
+    cparams.use_gpu = !noMetal;
+    cparams.flash_attn = useFlashAttn;
+
+    // TODO: Expose dtw_token_timestamps and dtw_aheads_preset
+    cparams.dtw_token_timestamps = false;
+    // cparams.dtw_aheads_preset = WHISPER_AHEADS_BASE;
+
+    cparams.use_coreml = !noCoreML;
+#ifndef WHISPER_USE_COREML
+    if (cparams.use_coreml) {
+        NSLog(@"[RNWhisper] CoreML is not enabled in this build, ignoring use_coreml option");
+        cparams.use_coreml = false;
+    }
+#endif
+
+#ifndef WSP_GGML_USE_METAL
+    if (cparams.use_gpu) {
+        NSLog(@"[RNWhisper] ggml-metal is not enabled in this build, ignoring use_gpu option");
+        cparams.use_gpu = false;
     }
-#else
-    context->ctx = whisper_init_from_file([modelPath UTF8String]);
 #endif
+
+#ifdef WSP_GGML_USE_METAL
+    if (cparams.use_gpu) {
+#if TARGET_OS_SIMULATOR
+        NSLog(@"[RNWhisper] ggml-metal is not available in simulator, ignoring use_gpu option: %@", reasonNoMetal);
+        cparams.use_gpu = false;
+#else // TARGET_OS_SIMULATOR
+        // Check ggml-metal availability
+        NSError * error = nil;
+        id<MTLDevice> device = MTLCreateSystemDefaultDevice();
+        id<MTLLibrary> library = [device
+            newLibraryWithSource:@"#include <metal_stdlib>\n"
+                                    "using namespace metal;"
+                                    "kernel void test() { simd_sum(0); }"
+            options:nil
+            error:&error
+        ];
+        if (error) {
+            reasonNoMetal = [error localizedDescription];
+        } else {
+            id<MTLFunction> kernel = [library newFunctionWithName:@"test"];
+            id<MTLComputePipelineState> pipeline = [device newComputePipelineStateWithFunction:kernel error:&error];
+            if (pipeline == nil) {
+                reasonNoMetal = [error localizedDescription];
+                NSLog(@"[RNWhisper] ggml-metal is not available, ignoring use_gpu option: %@", reasonNoMetal);
+                cparams.use_gpu = false;
+            }
+        }
+#endif // TARGET_OS_SIMULATOR
+    }
+#endif // WSP_GGML_USE_METAL
+
+    if (cparams.use_gpu && cparams.use_coreml) {
+        NSLog(@"[RNWhisper] Both use_gpu and use_coreml are enabled, ignoring use_coreml option");
+        cparams.use_coreml = false; // Skip CoreML if Metal is enabled
+    }
+
+    context->ctx = whisper_init_from_file_with_params([modelPath UTF8String], cparams);
     context->dQueue = dispatch_queue_create(
         [[NSString stringWithFormat:@"RNWhisperContext-%d", contextId] UTF8String],
         DISPATCH_QUEUE_SERIAL
     );
+    context->isMetalEnabled = cparams.use_gpu;
+    context->reasonNoMetal = reasonNoMetal;
     return context;
 }
 
+- (bool)isMetalEnabled {
+    return isMetalEnabled;
+}
+
+- (NSString *)reasonNoMetal {
+    return reasonNoMetal;
+}
+
 - (struct whisper_context *)getContext {
     return self->ctx;
 }
@@ -33,7 +100,7 @@
     return self->dQueue;
 }
 
-- (void)prepareRealtime:(NSDictionary *)options {
+- (void)prepareRealtime:(int)jobId options:(NSDictionary *)options {
     self->recordState.options = options;
 
     self->recordState.dataFormat.mSampleRate = WHISPER_SAMPLE_RATE; // 16000
@@ -46,68 +113,40 @@
     self->recordState.dataFormat.mReserved = 0;
     self->recordState.dataFormat.mFormatFlags = kLinearPCMFormatFlagIsSignedInteger;
 
-    int maxAudioSecOpt = options[@"realtimeAudioSec"] != nil ? [options[@"realtimeAudioSec"] intValue] : 0;
-    int maxAudioSec = maxAudioSecOpt > 0 ? maxAudioSecOpt : DEFAULT_MAX_AUDIO_SEC;
-    self->recordState.maxAudioSec = maxAudioSec;
-
-    int realtimeAudioSliceSec = options[@"realtimeAudioSliceSec"] != nil ? [options[@"realtimeAudioSliceSec"] intValue] : 0;
-    int audioSliceSec = realtimeAudioSliceSec > 0 && realtimeAudioSliceSec < maxAudioSec ? realtimeAudioSliceSec : maxAudioSec;
-
-    self->recordState.audioSliceSec = audioSliceSec;
-    self->recordState.isUseSlices = audioSliceSec < maxAudioSec;
+    self->recordState.isRealtime = true;
+    self->recordState.isTranscribing = false;
+    self->recordState.isCapturing = false;
+    self->recordState.isStoppedByAction = false;
 
     self->recordState.sliceIndex = 0;
     self->recordState.transcribeSliceIndex = 0;
     self->recordState.nSamplesTranscribing = 0;
 
-    [self freeBufferIfNeeded];
-    self->recordState.shortBufferSlices = [NSMutableArray new];
-
-    int16_t *audioBufferI16 = (int16_t *) malloc(audioSliceSec * WHISPER_SAMPLE_RATE * sizeof(int16_t));
-    [self->recordState.shortBufferSlices addObject:[NSValue valueWithPointer:audioBufferI16]];
-
-    self->recordState.sliceNSamples = [NSMutableArray new];
-    [self->recordState.sliceNSamples addObject:[NSNumber numberWithInt:0]];
-
-    self->recordState.isRealtime = true;
-    self->recordState.isTranscribing = false;
-    self->recordState.isCapturing = false;
-    self->recordState.isStoppedByAction = false;
+    self->recordState.sliceNSamples.clear();
+    self->recordState.sliceNSamples.push_back(0);
+
+    self->recordState.job = rnwhisper::job_new(jobId, [self createParams:options jobId:jobId]);
+    self->recordState.job->set_realtime_params(
+        {
+            .use_vad = options[@"useVad"] != nil ? [options[@"useVad"] boolValue] : false,
+            .vad_ms = options[@"vadMs"] != nil ? [options[@"vadMs"] intValue] : 2000,
+            .vad_thold = options[@"vadThold"] != nil ? [options[@"vadThold"] floatValue] : 0.6f,
+            .freq_thold = options[@"vadFreqThold"] != nil ? [options[@"vadFreqThold"] floatValue] : 100.0f
+        },
+        options[@"realtimeAudioSec"] != nil ? [options[@"realtimeAudioSec"] intValue] : 0,
+        options[@"realtimeAudioSliceSec"] != nil ? [options[@"realtimeAudioSliceSec"] intValue] : 0,
+        options[@"realtimeAudioMinSec"] != nil ? [options[@"realtimeAudioMinSec"] floatValue] : 0,
+        options[@"audioOutputPath"] != nil ? [options[@"audioOutputPath"] UTF8String] : nullptr
+    );
+    self->recordState.isUseSlices = self->recordState.job->audio_slice_sec < self->recordState.job->audio_sec;
 
     self->recordState.mSelf = self;
 }
 
-- (void)freeBufferIfNeeded {
-    if (self->recordState.shortBufferSlices != nil) {
-        for (int i = 0; i < [self->recordState.shortBufferSlices count]; i++) {
-            int16_t *audioBufferI16 = (int16_t *) [self->recordState.shortBufferSlices[i] pointerValue];
-            free(audioBufferI16);
-        }
-        self->recordState.shortBufferSlices = nil;
-    }
-}
-
-bool vad(RNWhisperContextRecordState *state, int16_t* audioBufferI16, int nSamples, int n)
+bool vad(RNWhisperContextRecordState *state, int sliceIndex, int nSamples, int n)
 {
-    bool isSpeech = true;
-    if (!state->isTranscribing && state->options[@"useVad"]) {
-        int vadSec = state->options[@"vadMs"] != nil ? [state->options[@"vadMs"] intValue] / 1000 : 2;
-        int sampleSize = vadSec * WHISPER_SAMPLE_RATE;
-        if (nSamples + n > sampleSize) {
-            int start = nSamples + n - sampleSize;
-            std::vector<float> audioBufferF32Vec(sampleSize);
-            for (int i = 0; i < sampleSize; i++) {
-                audioBufferF32Vec[i] = (float)audioBufferI16[i + start] / 32768.0f;
-            }
-            float vadThold = state->options[@"vadThold"] != nil ? [state->options[@"vadThold"] floatValue] : 0.6f;
-            float vadFreqThold = state->options[@"vadFreqThold"] != nil ? [state->options[@"vadFreqThold"] floatValue] : 100.0f;
-            isSpeech = rn_whisper_vad_simple(audioBufferF32Vec, WHISPER_SAMPLE_RATE, 1000, vadThold, vadFreqThold, false);
-            NSLog(@"[RNWhisper] VAD result: %d", isSpeech);
-        } else {
-            isSpeech = false;
-        }
-    }
-    return isSpeech;
+    if (state->isTranscribing) return true;
+    return state->job->vad_simple(sliceIndex, nSamples, n);
 }
 
 void AudioInputCallback(void * inUserData,
@@ -122,21 +161,21 @@ void AudioInputCallback(void * inUserData,
     if (!state->isCapturing) {
         NSLog(@"[RNWhisper] Not capturing, ignoring audio");
         if (!state->isTranscribing) {
-            state->transcribeHandler(state->jobId, @"end", @{});
+            [state->mSelf finishRealtimeTranscribe:state result:@{}];
         }
         return;
     }
 
     int totalNSamples = 0;
-    for (int i = 0; i < [state->sliceNSamples count]; i++) {
-        totalNSamples += [[state->sliceNSamples objectAtIndex:i] intValue];
+    for (int i = 0; i < state->sliceNSamples.size(); i++) {
+        totalNSamples += state->sliceNSamples[i];
     }
 
     const int n = inBuffer->mAudioDataByteSize / 2;
 
-    int nSamples = [state->sliceNSamples[state->sliceIndex] intValue];
+    int nSamples = state->sliceNSamples[state->sliceIndex];
 
-    if (totalNSamples + n > state->maxAudioSec * WHISPER_SAMPLE_RATE) {
+    if (totalNSamples + n > state->job->audio_sec * WHISPER_SAMPLE_RATE) {
         NSLog(@"[RNWhisper] Audio buffer is full, stop capturing");
         state->isCapturing = false;
         [state->mSelf stopAudio];
@@ -145,14 +184,14 @@ void AudioInputCallback(void * inUserData,
             nSamples == state->nSamplesTranscribing &&
             state->sliceIndex == state->transcribeSliceIndex
         ) {
-            state->transcribeHandler(state->jobId, @"end", @{});
+            [state->mSelf finishRealtimeTranscribe:state result:@{}];
         } else if (
             !state->isTranscribing &&
             nSamples != state->nSamplesTranscribing
         ) {
-            int16_t* audioBufferI16 = (int16_t*) [state->shortBufferSlices[state->sliceIndex] pointerValue];
-            if (!vad(state, audioBufferI16, nSamples, 0)) {
-                state->transcribeHandler(state->jobId, @"end", @{});
+            bool isSamplesEnough = nSamples / WHISPER_SAMPLE_RATE >= state->job->audio_min_sec;
+            if (!isSamplesEnough || !vad(state, state->sliceIndex, nSamples, 0)) {
+                [state->mSelf finishRealtimeTranscribe:state result:@{}];
                 return;
             }
             state->isTranscribing = true;
@@ -163,31 +202,25 @@ void AudioInputCallback(void * inUserData,
         return;
     }
 
-    int audioSliceSec = state->audioSliceSec;
-    if (nSamples + n > audioSliceSec * WHISPER_SAMPLE_RATE) {
+    if (nSamples + n > state->job->audio_slice_sec * WHISPER_SAMPLE_RATE) {
         // next slice
         state->sliceIndex++;
         nSamples = 0;
-        int16_t* audioBufferI16 = (int16_t*) malloc(audioSliceSec * WHISPER_SAMPLE_RATE * sizeof(int16_t));
-        [state->shortBufferSlices addObject:[NSValue valueWithPointer:audioBufferI16]];
-        [state->sliceNSamples addObject:[NSNumber numberWithInt:0]];
+        state->sliceNSamples.push_back(0);
     }
 
-    // Append to buffer
-    NSLog(@"[RNWhisper] Slice %d has %d samples", state->sliceIndex, nSamples);
+    NSLog(@"[RNWhisper] Slice %d has %d samples, put %d samples", state->sliceIndex, nSamples, n);
 
-    int16_t* audioBufferI16 = (int16_t*) [state->shortBufferSlices[state->sliceIndex] pointerValue];
-    for (int i = 0; i < n; i++) {
-        audioBufferI16[nSamples + i] = ((short*)inBuffer->mAudioData)[i];
-    }
+    state->job->put_pcm_data((short*) inBuffer->mAudioData, state->sliceIndex, nSamples, n);
 
-    bool isSpeech = vad(state, audioBufferI16, nSamples, n);
+    bool isSpeech = vad(state, state->sliceIndex, nSamples, n);
     nSamples += n;
-    state->sliceNSamples[state->sliceIndex] = [NSNumber numberWithInt:nSamples];
+    state->sliceNSamples[state->sliceIndex] = nSamples;
 
     AudioQueueEnqueueBuffer(state->queue, inBuffer, 0, NULL);
 
-    if (!isSpeech) return;
+    bool isSamplesEnough = nSamples / WHISPER_SAMPLE_RATE >= state->job->audio_min_sec;
+    if (!isSamplesEnough || !isSpeech) return;
 
     if (!state->isTranscribing) {
         state->isTranscribing = true;
@@ -197,21 +230,29 @@ void AudioInputCallback(void * inUserData,
     }
 }
 
+- (void)finishRealtimeTranscribe:(RNWhisperContextRecordState*) state result:(NSDictionary*)result {
+    // Save wav if needed
+    if (state->job->audio_output_path != nullptr) {
+        // 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(state->job->pcm_slices, state->sliceNSamples),
+            state->job->audio_output_path
+        );
+    }
+    state->transcribeHandler(state->job->job_id, @"end", result);
+    rnwhisper::job_remove(state->job->job_id);
+}
+
 - (void)fullTranscribeSamples:(RNWhisperContextRecordState*) state {
-    int nSamplesOfIndex = [[state->sliceNSamples objectAtIndex:state->transcribeSliceIndex] intValue];
+    int nSamplesOfIndex = state->sliceNSamples[state->transcribeSliceIndex];
     state->nSamplesTranscribing = nSamplesOfIndex;
     NSLog(@"[RNWhisper] Transcribing %d samples", state->nSamplesTranscribing);
 
-    int16_t* audioBufferI16 = (int16_t*) [state->shortBufferSlices[state->transcribeSliceIndex] pointerValue];
-    float* audioBufferF32 = (float*) malloc(state->nSamplesTranscribing * sizeof(float));
-    // convert I16 to F32
-    for (int i = 0; i < state->nSamplesTranscribing; i++) {
-        audioBufferF32[i] = (float)audioBufferI16[i] / 32768.0f;
-    }
+    float* pcmf32 = state->job->pcm_slice_to_f32(state->transcribeSliceIndex, state->nSamplesTranscribing);
+
     CFTimeInterval timeStart = CACurrentMediaTime();
-    struct whisper_full_params params = [state->mSelf getParams:state->options jobId:state->jobId];
-    int code = [state->mSelf fullTranscribe:state->jobId params:params audioData:audioBufferF32 audioDataCount:state->nSamplesTranscribing];
-    free(audioBufferF32);
+    int code = [state->mSelf fullTranscribe:state->job audioData:pcmf32 audioDataCount:state->nSamplesTranscribing];
+    free(pcmf32);
     CFTimeInterval timeEnd = CACurrentMediaTime();
     const float timeRecording = (float) state->nSamplesTranscribing / (float) state->dataFormat.mSampleRate;
 
@@ -231,7 +272,7 @@ void AudioInputCallback(void * inUserData,
         result[@"error"] = [NSString stringWithFormat:@"Transcribe failed with code %d", code];
     }
 
-    nSamplesOfIndex = [[state->sliceNSamples objectAtIndex:state->transcribeSliceIndex] intValue];
+    nSamplesOfIndex = state->sliceNSamples[state->transcribeSliceIndex];
 
     bool isStopped = state->isStoppedByAction || (
         !state->isCapturing &&
@@ -256,23 +297,13 @@ void AudioInputCallback(void * inUserData,
         result[@"isStoppedByAction"] = @(state->isStoppedByAction);
         result[@"isCapturing"] = @(false);
 
-        // Save wav if needed
-        if (state->options[@"audioOutputPath"] != nil) {
-            // TODO: Append in real time so we don't need to keep all slices & also reduce memory usage
-            [RNWhisperAudioUtils
-                saveWavFile:[RNWhisperAudioUtils concatShortBuffers:state->shortBufferSlices
-                                sliceNSamples:state->sliceNSamples]
-                audioOutputFile:state->options[@"audioOutputPath"]
-            ];
-        }
-
-        state->transcribeHandler(state->jobId, @"end", result);
+        [state->mSelf finishRealtimeTranscribe:state result:result];
     } else if (code == 0) {
         result[@"isCapturing"] = @(true);
-        state->transcribeHandler(state->jobId, @"transcribe", result);
+        state->transcribeHandler(state->job->job_id, @"transcribe", result);
     } else {
         result[@"isCapturing"] = @(true);
-        state->transcribeHandler(state->jobId, @"transcribe", result);
+        state->transcribeHandler(state->job->job_id, @"transcribe", result);
     }
 
     if (continueNeeded) {
@@ -300,8 +331,7 @@ void AudioInputCallback(void * inUserData,
     onTranscribe:(void (^)(int, NSString *, NSDictionary *))onTranscribe
 {
     self->recordState.transcribeHandler = onTranscribe;
-    self->recordState.jobId = jobId;
-    [self prepareRealtime:options];
+    [self prepareRealtime:jobId options:options];
 
     OSStatus status = AudioQueueNewInput(
         &self->recordState.dataFormat,
@@ -329,9 +359,10 @@ void AudioInputCallback(void * inUserData,
 struct rnwhisper_segments_callback_data {
     void (^onNewSegments)(NSDictionary *);
     int total_n_new;
+    bool tdrzEnable;
 };
 
-- (void)transcribeFile:(int)jobId
+- (void)transcribeData:(int)jobId
     audioData:(float *)audioData
     audioDataCount:(int)audioDataCount
     options:(NSDictionary *)options
@@ -342,9 +373,9 @@ struct rnwhisper_segments_callback_data {
     dispatch_async(dQueue, ^{
         self->recordState.isStoppedByAction = false;
         self->recordState.isTranscribing = true;
-        self->recordState.jobId = jobId;
 
-        whisper_full_params params = [self getParams:options jobId:jobId];
+        whisper_full_params params = [self createParams:options jobId:jobId];
+
         if (options[@"onProgress"] && [options[@"onProgress"] boolValue]) {
             params.progress_callback = [](struct whisper_context * /*ctx*/, struct whisper_state * /*state*/, int progress, void * user_data) {
                 void (^onProgress)(int) = (__bridge void (^)(int))user_data;
@@ -362,12 +393,18 @@ struct rnwhisper_segments_callback_data {
                 NSMutableArray *segments = [[NSMutableArray alloc] init];
                 for (int i = data->total_n_new - n_new; i < data->total_n_new; i++) {
                     const char * text_cur = whisper_full_get_segment_text(ctx, i);
-                    text = [text stringByAppendingString:[NSString stringWithUTF8String:text_cur]];
+                    NSMutableString *mutable_ns_text = [NSMutableString stringWithUTF8String:text_cur];
+
+                    if (data->tdrzEnable && whisper_full_get_segment_speaker_turn_next(ctx, i)) {
+                        [mutable_ns_text appendString:@" [SPEAKER_TURN]"];
+                    }
+
+                    text = [text stringByAppendingString:mutable_ns_text];
 
                     const int64_t t0 = whisper_full_get_segment_t0(ctx, i);
                     const int64_t t1 = whisper_full_get_segment_t1(ctx, i);
                     NSDictionary *segment = @{
-                        @"text": [NSString stringWithUTF8String:text_cur],
+                        @"text": [NSString stringWithString:mutable_ns_text],
                         @"t0": [NSNumber numberWithLongLong:t0],
                         @"t1": [NSNumber numberWithLongLong:t1]
                     };
@@ -385,12 +422,16 @@ struct rnwhisper_segments_callback_data {
             };
             struct rnwhisper_segments_callback_data user_data = {
                 .onNewSegments = onNewSegments,
-                .total_n_new = 0
+                .tdrzEnable = options[@"tdrzEnable"] && [options[@"tdrzEnable"] boolValue],
+                .total_n_new = 0,
             };
             params.new_segment_callback_user_data = &user_data;
         }
-        int code = [self fullTranscribe:jobId params:params audioData:audioData audioDataCount:audioDataCount];
-        self->recordState.jobId = -1;
+
+        rnwhisper::job* job = rnwhisper::job_new(jobId, params);
+        self->recordState.job = job;
+        int code = [self fullTranscribe:job audioData:audioData audioDataCount:audioDataCount];
+        rnwhisper::job_remove(jobId);
         self->recordState.isTranscribing = false;
         onEnd(code);
     });
@@ -405,9 +446,13 @@ struct rnwhisper_segments_callback_data {
 }
 
 - (void)stopTranscribe:(int)jobId {
-    rn_whisper_abort_transcribe(jobId);
+    if (self->recordState.job) self->recordState.job->abort();
     if (self->recordState.isRealtime && self->recordState.isCapturing) {
         [self stopAudio];
+        if (!self->recordState.isTranscribing) {
+            // Handle for VAD case
+            self->recordState.transcribeHandler(jobId, @"end", @{});
+        }
     }
     self->recordState.isCapturing = false;
     self->recordState.isStoppedByAction = true;
@@ -415,13 +460,11 @@ struct rnwhisper_segments_callback_data {
 }
 
 - (void)stopCurrentTranscribe {
-    if (!self->recordState.jobId) {
-        return;
-    }
-    [self stopTranscribe:self->recordState.jobId];
+    if (self->recordState.job == nullptr) return;
+    [self stopTranscribe:self->recordState.job->job_id];
 }
 
-- (struct whisper_full_params)getParams:(NSDictionary *)options jobId:(int)jobId {
+- (struct whisper_full_params)createParams:(NSDictionary *)options jobId:(int)jobId {
     struct whisper_full_params params = whisper_full_default_params(WHISPER_SAMPLING_GREEDY);
 
     const int n_threads = options[@"maxThreads"] != nil ?
@@ -440,9 +483,8 @@ struct rnwhisper_segments_callback_data {
     params.print_progress   = false;
     params.print_timestamps = false;
     params.print_special    = false;
-    params.speed_up         = options[@"speedUp"] != nil ? [options[@"speedUp"] boolValue] : false;
     params.translate        = options[@"translate"] != nil ? [options[@"translate"] boolValue] : false;
-    params.language         = options[@"language"] != nil ? [options[@"language"] UTF8String] : "auto";
+    params.language         = options[@"language"] != nil ? strdup([options[@"language"] UTF8String]) : "auto";
     params.n_threads        = n_threads > 0 ? n_threads : default_n_threads;
     params.offset_ms        = 0;
     params.no_context       = true;
@@ -452,6 +494,7 @@ struct rnwhisper_segments_callback_data {
         params.max_len = [options[@"maxLen"] intValue];
     }
     params.token_timestamps = options[@"tokenTimestamps"] != nil ? [options[@"tokenTimestamps"] boolValue] : false;
+    params.tdrz_enable = options[@"tdrzEnable"] != nil ? [options[@"tdrzEnable"] boolValue] : false;
 
     if (options[@"bestOf"] != nil) {
         params.greedy.best_of = [options[@"bestOf"] intValue];
@@ -459,7 +502,6 @@ struct rnwhisper_segments_callback_data {
     if (options[@"maxContext"] != nil) {
         params.n_max_text_ctx = [options[@"maxContext"] intValue];
     }
-
     if (options[@"offset"] != nil) {
         params.offset_ms = [options[@"offset"] intValue];
     }
@@ -475,35 +517,20 @@ struct rnwhisper_segments_callback_data {
     if (options[@"temperatureInc"] != nil) {
         params.temperature_inc = [options[@"temperature_inc"] floatValue];
     }
-
     if (options[@"prompt"] != nil) {
-        params.initial_prompt = [options[@"prompt"] UTF8String];
+        params.initial_prompt = strdup([options[@"prompt"] UTF8String]);
     }
 
-    // abort handler
-    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(jobId);
-    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(jobId);
-
     return params;
 }
 
-- (int)fullTranscribe:(int)jobId
-  params:(struct whisper_full_params)params
+- (int)fullTranscribe:(rnwhisper::job *)job
   audioData:(float *)audioData
   audioDataCount:(int)audioDataCount
 {
     whisper_reset_timings(self->ctx);
-
-    int code = whisper_full(self->ctx, params, audioData, audioDataCount);
-    rn_whisper_remove_abort_map(jobId);
+    int code = whisper_full(self->ctx, job->params, audioData, audioDataCount);
+    if (job && job->is_aborted()) code = -999;
     // if (code == 0) {
     //     whisper_print_timings(self->ctx);
     // }
@@ -517,12 +544,21 @@ struct rnwhisper_segments_callback_data {
     NSMutableArray *segments = [[NSMutableArray alloc] init];
     for (int i = 0; i < n_segments; i++) {
         const char * text_cur = whisper_full_get_segment_text(self->ctx, i);
-        text = [text stringByAppendingString:[NSString stringWithUTF8String:text_cur]];
+        NSMutableString *mutable_ns_text = [NSMutableString stringWithUTF8String:text_cur];
+
+        // Simplified condition
+        if (self->recordState.options[@"tdrzEnable"] &&
+            [self->recordState.options[@"tdrzEnable"] boolValue] &&
+            whisper_full_get_segment_speaker_turn_next(self->ctx, i)) {
+            [mutable_ns_text appendString:@" [SPEAKER_TURN]"];
+        }
+
+        text = [text stringByAppendingString:mutable_ns_text];
 
         const int64_t t0 = whisper_full_get_segment_t0(self->ctx, i);
         const int64_t t1 = whisper_full_get_segment_t1(self->ctx, i);
         NSDictionary *segment = @{
-            @"text": [NSString stringWithUTF8String:text_cur],
+            @"text": [NSString stringWithString:mutable_ns_text],
             @"t0": [NSNumber numberWithLongLong:t0],
             @"t1": [NSNumber numberWithLongLong:t1]
         };
@@ -534,10 +570,19 @@ struct rnwhisper_segments_callback_data {
     return result;
 }
 
+- (NSString *)bench:(int)maxThreads {
+    const int n_threads = maxThreads > 0 ? maxThreads : 0;
+
+    const int max_threads = (int) [[NSProcessInfo processInfo] processorCount];
+    // Use 2 threads by default on 4-core devices, 4 threads on more cores
+    const int default_n_threads = max_threads == 4 ? 2 : MIN(4, max_threads);
+    NSString *result = [NSString stringWithUTF8String:rnwhisper::bench(self->ctx, n_threads).c_str()];
+    return result;
+}
+
 - (void)invalidate {
     [self stopCurrentTranscribe];
     whisper_free(self->ctx);
-    [self freeBufferIfNeeded];
 }
 
 @end

package/jest/mock.js

@@ -2,12 +2,17 @@ const { NativeModules, DeviceEventEmitter } = require('react-native')
 
 if (!NativeModules.RNWhisper) {
   NativeModules.RNWhisper = {
-    initContext: jest.fn(() => Promise.resolve(1)),
+    initContext: jest.fn(() => Promise.resolve({ contextId: 1 })),
     transcribeFile: jest.fn(() => Promise.resolve({
       result: ' Test',
       segments: [{ text: ' Test', t0: 0, t1: 33 }],
       isAborted: false,
     })),
+    transcribeData: jest.fn(() => Promise.resolve({
+      result: ' Test',
+      segments: [{ text: ' Test', t0: 0, t1: 33 }],
+      isAborted: false,
+    })),
     startRealtimeTranscribe: jest.fn((contextId, jobId) => {
       setTimeout(() => {
         // Start
@@ -45,11 +50,19 @@ if (!NativeModules.RNWhisper) {
         })
       })
     }),
+    bench: jest.fn(() => Promise.resolve({
+      config: 'NEON',
+      nThreads: 1,
+      encodeMs: 1,
+      decodeMs: 1,
+      batchMs: 1,
+      promptMs: 1,
+    })),
     releaseContext: jest.fn(() => Promise.resolve()),
     releaseAllContexts: jest.fn(() => Promise.resolve()),
 
     // iOS AudioSession utils
-    getAudioSessionCurrentCategory: jest.fn(() => Promise.resolve({ 
+    getAudioSessionCurrentCategory: jest.fn(() => Promise.resolve({
       category: 'AVAudioSessionCategoryPlayAndRecord',
       options: [],
     })),

package/lib/commonjs/NativeRNWhisper.js.map

@@ -1 +1 @@
-{"version":3,"names":["_reactNative","require","_default","TurboModuleRegistry","get","exports","default"],"sourceRoot":"../../src","sources":["NativeRNWhisper.ts"],"mappings":";;;;;;AACA,IAAAA,YAAA,GAAAC,OAAA;AAAkD,IAAAC,QAAA,GA0FnCC,gCAAmB,CAACC,GAAG,CAAO,WAAW,CAAC;AAAAC,OAAA,CAAAC,OAAA,GAAAJ,QAAA"}
+{"version":3,"names":["_reactNative","require","_default","TurboModuleRegistry","get","exports","default"],"sourceRoot":"../../src","sources":["NativeRNWhisper.ts"],"mappings":";;;;;;AACA,IAAAA,YAAA,GAAAC,OAAA;AAAkD,IAAAC,QAAA,GA0GnCC,gCAAmB,CAACC,GAAG,CAAO,WAAW,CAAC;AAAAC,OAAA,CAAAC,OAAA,GAAAJ,QAAA"}