Separated context and state for easier parallelization

whisper.h

@@ -32,7 +32,7 @@ extern "C" {
     //
     // C interface
     //
-    // The following interface is thread-safe as long as the sample whisper_context is not used by multiple threads
+    // The following interface is thread-safe as long as the sample whisper_state is not used by multiple threads
     // concurrently.
     //
     // Basic usage:
@@ -66,6 +66,7 @@ extern "C" {
     //
 
     struct whisper_context;
+    struct whisper_state;
 
     typedef int whisper_token;
 
@@ -101,14 +102,21 @@ extern "C" {
     WHISPER_API struct whisper_context * whisper_init_from_buffer(void * buffer, size_t buffer_size);
     WHISPER_API struct whisper_context * whisper_init(struct whisper_model_loader * loader);
 
+    // Function to create a new state used to run a transformation.
+    WHISPER_API struct whisper_state* whisper_init_state(whisper_context* ctx);
+
     // Frees all memory allocated by the model.
     WHISPER_API void whisper_free(struct whisper_context * ctx);
 
+    // Frees all memory allocated by the current state.
+    WHISPER_API void whisper_free_state(struct whisper_state * state);
+
     // Convert RAW PCM audio to log mel spectrogram.
     // The resulting spectrogram is stored inside the provided whisper context.
     // Returns 0 on success
     WHISPER_API int whisper_pcm_to_mel(
             struct whisper_context * ctx,
+              struct whisper_state * state,
                        const float * samples,
                                int   n_samples,
                                int   n_threads);
@@ -118,27 +126,29 @@ extern "C" {
     // Returns 0 on success
     WHISPER_API int whisper_pcm_to_mel_phase_vocoder(
             struct whisper_context * ctx,
+              struct whisper_state * state,
                        const float * samples,
                                int   n_samples,
                                int   n_threads);
 
 
-    // This can be used to set a custom log mel spectrogram inside the provided whisper context.
+    // This can be used to set a custom log mel spectrogram inside the provided whisper state.
     // Use this instead of whisper_pcm_to_mel() if you want to provide your own log mel spectrogram.
     // n_mel must be 80
     // Returns 0 on success
     WHISPER_API int whisper_set_mel(
-            struct whisper_context * ctx,
+              struct whisper_state * ctx,
                        const float * data,
                                int   n_len,
                                int   n_mel);
 
-    // Run the Whisper encoder on the log mel spectrogram stored inside the provided whisper context.
+    // Run the Whisper encoder on the log mel spectrogram stored inside the provided whisper state.
     // Make sure to call whisper_pcm_to_mel() or whisper_set_mel() first.
     // offset can be used to specify the offset of the first frame in the spectrogram.
     // Returns 0 on success
     WHISPER_API int whisper_encode(
             struct whisper_context * ctx,
+              struct whisper_state * state,
                                int   offset,
                                int   n_threads);
 
@@ -150,6 +160,7 @@ extern "C" {
     // TODO: add support for multiple decoders
     WHISPER_API int whisper_decode(
             struct whisper_context * ctx,
+              struct whisper_state * state,
                const whisper_token * tokens,
                                int   n_tokens,
                                int   n_past,
@@ -186,11 +197,12 @@ extern "C" {
     // ref: https://github.com/openai/whisper/blob/main/whisper/decoding.py#L18-L69
     WHISPER_API int whisper_lang_auto_detect(
             struct whisper_context * ctx,
+              struct whisper_state * state,
                                int   offset_ms,
                                int   n_threads,
                              float * lang_probs);
 
-    WHISPER_API int whisper_n_len          (struct whisper_context * ctx); // mel length
+    WHISPER_API int whisper_n_len          (struct whisper_state * ctx); // mel length
     WHISPER_API int whisper_n_vocab        (struct whisper_context * ctx);
     WHISPER_API int whisper_n_text_ctx     (struct whisper_context * ctx);
     WHISPER_API int whisper_n_audio_ctx    (struct whisper_context * ctx);
@@ -200,7 +212,7 @@ extern "C" {
     // The logits for the last token are stored in the last row
     // Rows: n_tokens
     // Cols: n_vocab
-    WHISPER_API float * whisper_get_logits(struct whisper_context * ctx);
+    WHISPER_API float * whisper_get_logits(struct whisper_state * ctx);
 
     // Token Id -> String. Uses the vocabulary in the provided context
     WHISPER_API const char * whisper_token_to_str(struct whisper_context * ctx, whisper_token token);
@@ -219,8 +231,8 @@ extern "C" {
     WHISPER_API whisper_token whisper_token_transcribe(void);
 
     // Performance information
-    WHISPER_API void whisper_print_timings(struct whisper_context * ctx);
+    WHISPER_API void whisper_print_timings(struct whisper_context * ctx, struct whisper_state * state);
-    WHISPER_API void whisper_reset_timings(struct whisper_context * ctx);
+    WHISPER_API void whisper_reset_timings(struct whisper_state * ctx);
 
     // Print system information
     WHISPER_API const char * whisper_print_system_info(void);
@@ -236,12 +248,12 @@ extern "C" {
     // Text segment callback
     // Called on every newly generated text segment
     // Use the whisper_full_...() functions to obtain the text segments
-    typedef void (*whisper_new_segment_callback)(struct whisper_context * ctx, int n_new, void * user_data);
+    typedef void (*whisper_new_segment_callback)(struct whisper_context * ctx, struct whisper_state * state, int n_new, void * user_data);
 
     // Encoder begin callback
     // If not NULL, called before the encoder starts
     // If it returns false, the computation is aborted
-    typedef bool (*whisper_encoder_begin_callback)(struct whisper_context * ctx, void * user_data);
+    typedef bool (*whisper_encoder_begin_callback)(struct whisper_context * ctx, struct whisper_state * state, void * user_data);
 
     // Parameters for the whisper_full() function
     // If you chnage the order or add new parameters, make sure to update the default values in whisper.cpp:
@@ -339,31 +351,31 @@ extern "C" {
 
     // Number of generated text segments.
     // A segment can be a few words, a sentence, or even a paragraph.
-    WHISPER_API int whisper_full_n_segments(struct whisper_context * ctx);
+    WHISPER_API int whisper_full_n_segments(struct whisper_state * state);
 
     // Language id associated with the current context
-    WHISPER_API int whisper_full_lang_id(struct whisper_context * ctx);
+    WHISPER_API int whisper_full_lang_id(struct whisper_state * state);
 
     // Get the start and end time of the specified segment.
-    WHISPER_API int64_t whisper_full_get_segment_t0(struct whisper_context * ctx, int i_segment);
+    WHISPER_API int64_t whisper_full_get_segment_t0(struct whisper_state * state, int i_segment);
-    WHISPER_API int64_t whisper_full_get_segment_t1(struct whisper_context * ctx, int i_segment);
+    WHISPER_API int64_t whisper_full_get_segment_t1(struct whisper_state * state, int i_segment);
 
     // Get the text of the specified segment.
-    WHISPER_API const char * whisper_full_get_segment_text(struct whisper_context * ctx, int i_segment);
+    WHISPER_API const char * whisper_full_get_segment_text(struct whisper_state * state, int i_segment);
 
     // Get number of tokens in the specified segment.
-    WHISPER_API int whisper_full_n_tokens(struct whisper_context * ctx, int i_segment);
+    WHISPER_API int whisper_full_n_tokens(struct whisper_state* state, int i_segment);
 
     // Get the token text of the specified token in the specified segment.
-    WHISPER_API const char * whisper_full_get_token_text(struct whisper_context * ctx, int i_segment, int i_token);
+    WHISPER_API const char * whisper_full_get_token_text(struct whisper_context * ctx, struct whisper_state * state, int i_segment, int i_token);
-    WHISPER_API whisper_token whisper_full_get_token_id (struct whisper_context * ctx, int i_segment, int i_token);
+    WHISPER_API whisper_token whisper_full_get_token_id (struct whisper_state * ctx, int i_segment, int i_token);
 
     // Get token data for the specified token in the specified segment.
     // This contains probabilities, timestamps, etc.
-    WHISPER_API whisper_token_data whisper_full_get_token_data(struct whisper_context * ctx, int i_segment, int i_token);
+    WHISPER_API whisper_token_data whisper_full_get_token_data(struct whisper_state * ctx, int i_segment, int i_token);
 
     // Get the probability of the specified token in the specified segment.
-    WHISPER_API float whisper_full_get_token_p(struct whisper_context * ctx, int i_segment, int i_token);
+    WHISPER_API float whisper_full_get_token_p(struct whisper_state * state, int i_segment, int i_token);
 
     ////////////////////////////////////////////////////////////////////////////