command : wip in progress, improve guided decoding

examples/addon.node/addon.cpp

@@ -72,7 +72,7 @@ int timestamp_to_sample(int64_t t, int n_samples) {
     return std::max(0, std::min((int) n_samples - 1, (int) ((t*WHISPER_SAMPLE_RATE)/100)));
 }
 
-void whisper_print_segment_callback(struct whisper_context * ctx, int n_new, void * user_data) {
+void whisper_print_segment(struct whisper_context * ctx, int n_new, void * user_data) {
     const auto & params  = *((whisper_print_user_data *) user_data)->params;
     const auto & pcmf32s = *((whisper_print_user_data *) user_data)->pcmf32s;
 
@@ -250,7 +250,7 @@ int run(whisper_params &params, std::vector<std::vector<std::string>> &result) {
 
             // this callback is called on each new segment
             if (!wparams.print_realtime) {
-                wparams.new_segment_callback           = whisper_print_segment_callback;
+                wparams.new_segment_callback           = whisper_print_segment;
                 wparams.new_segment_callback_user_data = &user_data;
             }
 

examples/command/command.cpp

@@ -109,6 +109,73 @@ void whisper_print_usage(int /*argc*/, char ** argv, const whisper_params & para
     fprintf(stderr, "\n");
 }
 
+struct whisper_logits_filter_user_data {
+    std::vector<std::string>                * allowed_commands;
+    std::vector<std::vector<whisper_token>> * allowed_tokens;
+};
+
+void whisper_logits_filter(
+            struct whisper_context * ctx,
+          const whisper_token_data * tokens,
+                               int   n_tokens,
+                             float * logits,
+                              void * user_data){
+    const auto & allowed_tokens = *((whisper_logits_filter_user_data *) user_data)->allowed_tokens;
+
+    printf("n_tokens = %d\n", n_tokens);
+    for (int i = 0; i < n_tokens; i++) {
+        printf(" - '%s' (%.2f)\n", whisper_token_to_str(ctx, tokens[i].id), logits[i]);
+    }
+
+    if (n_tokens == 0) {
+        return;
+    }
+
+    std::vector<std::pair<whisper_token, float>> pool;
+    for (int i = 0; i < (int) allowed_tokens.size(); i++) {
+        const int n = (int) allowed_tokens[i].size();
+        if (n_tokens > n) {
+            continue;
+        }
+
+        const whisper_token id = allowed_tokens[i][n_tokens - 1];
+        pool.push_back({ id, logits[id] });
+    }
+
+    if (pool.empty()) {
+        return;
+    }
+
+    printf("applying logits filter, pool size = %d\n", (int) pool.size());
+
+    const int ibeg = whisper_token_beg(ctx);
+
+    double sum_all = 0.0;
+    for (int i = 0; i < ibeg; ++i) {
+        if (logits[i] == -INFINITY) {
+            continue;
+        }
+        sum_all += logits[i];
+    }
+
+    double sum_pool = 0.0;
+    for (int i = 0; i < (int) pool.size(); ++i) {
+        sum_pool += pool[i].second;
+    }
+
+    printf("sum_all = %.2f, sum_pool = %.2f\n", sum_all, sum_pool);
+
+    for (int i = 0; i < ibeg; ++i) {
+        logits[i] = -INFINITY;
+    }
+
+    for (int i = 0; i < (int) pool.size(); ++i) {
+        //logits[pool[i].first] = pool[i].second / sum_pool * sum_all;
+        logits[pool[i].first] = pool[i].second;
+        printf(" - '%s' (%.2f)\n", whisper_token_to_str(ctx, pool[i].first), logits[pool[i].first]);
+    }
+}
+
 std::string transcribe(whisper_context * ctx, const whisper_params & params, const std::vector<float> & pcmf32, float & prob, int64_t & t_ms) {
     const auto t_start = std::chrono::high_resolution_clock::now();
 
@@ -131,6 +198,8 @@ std::string transcribe(whisper_context * ctx, const whisper_params & params, con
     wparams.audio_ctx        = params.audio_ctx;
     wparams.speed_up         = params.speed_up;
 
+    wparams.temperature_inc  = -1.0f;
+
     if (whisper_full(ctx, wparams, pcmf32.data(), pcmf32.size()) != 0) {
         return "";
     }
@@ -334,22 +403,31 @@ int process_command_list(struct whisper_context * ctx, audio_async &audio, const
             wparams.translate        = params.translate;
             wparams.no_context       = true;
             wparams.single_segment   = true;
-            wparams.max_tokens       = 1;
+            //wparams.max_tokens       = 1;
             wparams.language         = params.language.c_str();
             wparams.n_threads        = params.n_threads;
 
             wparams.audio_ctx        = params.audio_ctx;
             wparams.speed_up         = params.speed_up;
 
+            wparams.temperature_inc  = -1.0f;
+
             wparams.prompt_tokens    = k_tokens.data();
             wparams.prompt_n_tokens  = k_tokens.size();
 
+            whisper_logits_filter_user_data user_data = { &allowed_commands, &allowed_tokens };
+
+            wparams.logits_filter_callback           = whisper_logits_filter;
+            wparams.logits_filter_callback_user_data = &user_data;
+
             // run the transformer and a single decoding pass
             if (whisper_full(ctx, wparams, pcmf32_cur.data(), pcmf32_cur.size()) != 0) {
                 fprintf(stderr, "%s: ERROR: whisper_full() failed\n", __func__);
                 break;
             }
 
+            fprintf(stdout, "%s: text - '%s'\n", __func__, whisper_full_get_segment_text(ctx, 0));
+
             // estimate command probability
             // NOTE: not optimal
             {
@@ -436,7 +514,7 @@ int process_command_list(struct whisper_context * ctx, audio_async &audio, const
 
 // always-prompt mode
 // transcribe the voice into text after valid prompt
-int always_prompt_transcription(struct whisper_context * ctx, audio_async & audio, const whisper_params & params) {
+int process_always_prompt(struct whisper_context * ctx, audio_async & audio, const whisper_params & params) {
     bool is_running = true;
     bool ask_prompt = true;
 
@@ -496,7 +574,7 @@ int always_prompt_transcription(struct whisper_context * ctx, audio_async & audi
                 const float sim = similarity(prompt, k_prompt);
 
                 //debug
-                //fprintf(stdout, "command size: %i\n", command_length);
+                //fprintf(stdout, "command size: %d, sim: %f\n", (int) command.size(), sim);
 
                 if ((sim > 0.7f) && (command.size() > 0)) {
                     fprintf(stdout, "%s: Command '%s%s%s', (t = %d ms)\n", __func__, "\033[1m", command.c_str(), "\033[0m", (int) t_ms);
@@ -676,7 +754,7 @@ int main(int argc, char ** argv) {
     if (!params.commands.empty()) {
         ret_val = process_command_list(ctx, audio, params);
     } else if (!params.prompt.empty()) {
-        ret_val = always_prompt_transcription(ctx, audio, params);
+        ret_val = process_always_prompt(ctx, audio, params);
     } else {
         ret_val = process_general_transcription(ctx, audio, params);
     }

examples/common-sdl.h

@@ -1,13 +1,13 @@
 #pragma once
 
-#include <SDL.h>
-#include <SDL_audio.h>
-
 #include <atomic>
 #include <cstdint>
 #include <vector>
 #include <mutex>
 
+#include <SDL.h>
+#include <SDL_audio.h>
+
 //
 // SDL Audio capture
 //

examples/main/main.cpp

@@ -193,7 +193,7 @@ struct whisper_print_user_data {
     const std::vector<std::vector<float>> * pcmf32s;
 };
 
-void whisper_print_segment_callback(struct whisper_context * ctx, int n_new, void * user_data) {
+void whisper_print_segment(struct whisper_context * ctx, int n_new, void * user_data) {
     const auto & params  = *((whisper_print_user_data *) user_data)->params;
     const auto & pcmf32s = *((whisper_print_user_data *) user_data)->pcmf32s;
 
@@ -597,7 +597,7 @@ int main(int argc, char ** argv) {
 
             // this callback is called on each new segment
             if (!wparams.print_realtime) {
-                wparams.new_segment_callback           = whisper_print_segment_callback;
+                wparams.new_segment_callback           = whisper_print_segment;
                 wparams.new_segment_callback_user_data = &user_data;
             }
 
@@ -618,8 +618,6 @@ int main(int argc, char ** argv) {
                 fprintf(stderr, "%s: failed to process audio\n", argv[0]);
                 return 10;
             }
-
-            whisper_full_cluster_segments(ctx);
         }
 
         // output stuff

ggml.c

@@ -8517,195 +8517,6 @@ enum ggml_opt_result ggml_opt(
 
 ////////////////////////////////////////////////////////////////////////////////
 
-void ggml_svd_reduce_dims(
-        int ne0,
-        int ne1,
-        float * a,
-        int nd) {
-    int n = ne1;
-    int m = ne0;
-
-    float * A = a;
-    float * A0 = (float *) malloc(n * m * sizeof(float));
-
-    // average vector
-    //float * M = (float *) malloc(m * sizeof(float));
-
-    //{
-    //    for (int j = 0; j < m; ++j) {
-    //        M[j] = 0.0f;
-    //    }
-    //    for (int i = 0; i < n; ++i) {
-    //        for (int j = 0; j < m; ++j) {
-    //            M[j] += A[i * m + j];
-    //        }
-    //    }
-    //    for (int j = 0; j < m; ++j) {
-    //        M[j] /= (float) n;
-    //    }
-    //}
-
-    //// subtract average vector
-    //for (int i = 0; i < n; ++i) {
-    //    for (int j = 0; j < m; ++j) {
-    //        A[i * m + j] -= M[j];
-    //    }
-    //}
-
-    //free(M);
-
-    memcpy(A0, A, n * m * sizeof(float));
-
-    // print A
-    //printf("A:\n");
-    //for (int i = 0; i < n; ++i) {
-    //    printf("col %d : ", i);
-    //    for (int j = 0; j < m; ++j) {
-    //        printf("%9.5f ", A[i * m + j]);
-    //    }
-    //    printf("\n");
-    //}
-    //printf("\n");
-
-    // SVD
-    // A = U * S * V^T
-
-    float * U = (float *) malloc(n * m * sizeof(float));
-    float * S = (float *) malloc(n * sizeof(float));
-    float * V = (float *) malloc(n * n * sizeof(float));
-
-    int lda = m;
-    int ldu = m;
-    int ldvt = n;
-
-    float work_size;
-    int lwork = -1;
-    int info = 0;
-
-    sgesvd_("S", "S", &m, &n, A, &lda, S, U, &ldu, V, &ldvt, &work_size, &lwork, &info);
-
-    lwork = (int) work_size;
-
-    //printf("work_size = %f, info = %d, lwork = %d\n", work_size, info, lwork);
-
-    float * work = (float *) malloc(lwork * sizeof(float));
-
-    sgesvd_("S", "S", &m, &n, A, &lda, S, U, &ldu, V, &ldvt, work, &lwork, &info);
-
-    free(work);
-
-    // print U
-    //printf("U:\n");
-    //for (int i = 0; i < n; ++i) {
-    //    printf("col %d : ", i);
-    //    for (int j = 0; j < m; ++j) {
-    //        printf("%9.5f ", U[i * m + j]);
-    //    }
-    //    printf("\n");
-    //}
-    //printf("\n");
-
-    // normalize S
-    {
-        double sum = 0.0;
-        for (int i = 0; i < n; ++i) {
-            sum += S[i];
-        }
-        sum *= sqrt((double) m);
-        for (int i = 0; i < n; ++i) {
-            S[i] /= sum;
-        }
-    }
-
-    // print S
-    printf("S:\n");
-    for (int i = 0; i < n; ++i) {
-        printf("- %d = %9.5f\n", i, S[i]);
-    }
-    printf("\n");
-
-    // print V
-    //printf("V:\n");
-    //for (int i = 0; i < n; ++i) {
-    //    printf("col %d : ", i);
-    //    for (int j = 0; j < n; ++j) {
-    //        printf("%9.5f ", V[i * n + j]);
-    //    }
-    //    printf("\n");
-    //}
-    //printf("\n");
-
-    // print A
-    //printf("A:\n");
-    //for (int i = 0; i < n; ++i) {
-    //    printf("col %d : ", i);
-    //    for (int j = 0; j < m; ++j) {
-    //        printf("%9.5f ", A[i * m + j]);
-    //    }
-    //    printf("\n");
-    //}
-    //printf("\n");
-
-    // compute singular vectors in U
-    for (int i = 0; i < n; ++i) {
-        for (int j = 0; j < m; ++j) {
-            U[i * m + j] *= S[i];
-        }
-    }
-
-    // normalize U
-    for (int i = 0; i < n; ++i) {
-        double sum = 0.0;
-        for (int j = 0; j < m; ++j) {
-            sum += U[i * m + j] * U[i * m + j];
-        }
-        sum = sqrt(sum);
-        for (int j = 0; j < m; ++j) {
-            U[i * m + j] /= sum*sqrt((double) m);
-        }
-    }
-
-    // print U
-    //printf("U:\n");
-    //for (int i = 0; i < n; ++i) {
-    //    printf("col %d : ", i);
-    //    for (int j = 0; j < m; ++j) {
-    //        printf("%9.5f ", U[i * m + j]);
-    //    }
-    //    printf("\n");
-    //}
-    //printf("\n");
-
-    // project A0 onto U
-    for (int i = 0; i < n; ++i) {
-        for (int j = 0; j < nd; ++j) {
-            A[i * nd + j] = 0.0f;
-            //if (j == 0) continue;
-            for (int k = 0; k < m; ++k) {
-                A[i * nd + j] += A0[i * m + k] * U[j * m + k];
-            }
-        }
-    }
-
-    // print A
-    //printf("A:\n");
-    //for (int i = 0; i < n; ++i) {
-    //    printf("col %d : ", i);
-    //    for (int j = 0; j < n; ++j) {
-    //        printf("%9.5f ", A[i * n + j]);
-    //    }
-    //    printf("\n");
-    //}
-    //printf("\n");
-
-    free(U);
-    free(S);
-    free(V);
-    free(A0);
-}
-
-////////////////////////////////////////////////////////////////////////////////
-
 int ggml_cpu_has_avx(void) {
 #if defined(__AVX__)
     return 1;

ggml.h

@@ -726,16 +726,6 @@ enum ggml_opt_result ggml_opt(
         struct ggml_opt_params params,
         struct ggml_tensor * f);
 
-//
-// Temp stuff
-//
-
-void ggml_svd_reduce_dims(
-        int ne0,
-        int ne1,
-        float * a,
-        int nd);
-
 //
 // system info
 //

whisper.cpp

@@ -268,14 +268,6 @@ static const std::map<e_model, size_t> MEM_REQ_KV_SELF = {
     { MODEL_LARGE,    71ull*MB },
 };
 
-static const std::map<e_model, size_t> MEM_REQ_KV_ENC_SELF = {
-    { MODEL_TINY,     23ull*MB },
-    { MODEL_BASE,     26ull*MB },
-    { MODEL_SMALL,   216ull*MB },
-    { MODEL_MEDIUM,  243ull*MB },
-    { MODEL_LARGE,   271ull*MB },
-};
-
 static const std::map<e_model, size_t> MEM_REQ_KV_CROSS = {
     { MODEL_TINY,      9ull*MB },
     { MODEL_BASE,     18ull*MB },
@@ -579,7 +571,6 @@ struct whisper_context {
     // cross-attention KV cache for the decoders
     // shared between all decoders
     whisper_kv_cache kv_cross;
-    whisper_kv_cache kv_enc_self;
 
     whisper_decoder decoders[WHISPER_MAX_DECODERS] = {};
 
@@ -612,8 +603,6 @@ struct whisper_context {
     // [EXPERIMENTAL] speed-up techniques
     int32_t exp_n_audio_ctx; // 0 - use default
 
-    std::vector<float> audio_embd;
-
     void use_buf(struct ggml_context * ctx, int i) {
 #if defined(WHISPER_USE_SCRATCH)
         size_t last_size = 0;
@@ -847,11 +836,6 @@ static bool whisper_model_load(struct whisper_model_loader * loader, whisper_con
             return false;
         }
 
-        if (!kv_cache_init(model.hparams, scale*MEM_REQ_KV_ENC_SELF.at(model.type), wctx.kv_enc_self, wctx.wtype, model.hparams.n_audio_ctx)) {
-            fprintf(stderr, "%s: kv_cache_init() failed for cross-attention cache\n", __func__);
-            return false;
-        }
-
         {
             const size_t memory_size = ggml_nbytes(wctx.kv_cross.k) + ggml_nbytes(wctx.kv_cross.v);
             fprintf(stderr, "%s: kv cross size = %7.2f MB\n", __func__, memory_size/1024.0/1024.0);
@@ -1374,8 +1358,7 @@ static bool whisper_model_load(struct whisper_model_loader * loader, whisper_con
 static bool whisper_encode(
         whisper_context & wctx,
               const int   mel_offset,
-              const int   n_threads,
-              bool repeat = false) {
+              const int   n_threads) {
     const int64_t t_start_us = ggml_time_us();
 
     const auto & model   = wctx.model;
@@ -1407,31 +1390,13 @@ static bool whisper_encode(
         const int i0 = std::min(mel_offset, mel_inp.n_len);
         const int i1 = std::min(mel_offset + 2*n_ctx, mel_inp.n_len);
 
-        if (repeat == false) {
-            for (int j = 0; j < mel_inp.n_mel; ++j) {
-                for (int i = i0; i < i1; ++i) {
-                    dst[j*2*n_ctx + (i - i0)] = mel_inp.data[j*mel_inp.n_len + i];
-                }
-            }
-        } else {
-            for (int j = 0; j < mel_inp.n_mel; ++j) {
-                int k = 0;
-                while (k < 2*n_ctx) {
-                    for (int i = i0; i < i1; ++i) {
-                        dst[j*2*n_ctx + k] = mel_inp.data[j*mel_inp.n_len + i];
-                        k++;
-                        if (k >= 2*n_ctx) {
-                            break;
-                        }
-                    }
-                }
+        for (int j = 0; j < mel_inp.n_mel; ++j) {
+            for (int i = i0; i < i1; ++i) {
+                dst[j*2*n_ctx + (i - i0)] = mel_inp.data[j*mel_inp.n_len + i];
             }
         }
     }
 
-    struct ggml_cgraph gf = {};
-    gf.n_threads = n_threads;
-
     struct ggml_tensor * cur;
 
     // convolution + gelu
@@ -1459,18 +1424,6 @@ static bool whisper_encode(
         cur = ggml_gelu(ctx0, cur);
     }
 
-    //{
-    //    //printf("cur: %d %d %d %d, size element = %d\n", cur->ne[0], cur->ne[1], cur->ne[2], cur->ne[3], ggml_element_size(cur));
-
-    //    wctx.use_buf(ctx0, -1);
-
-    //    struct ggml_tensor * k = ggml_view_1d(ctx0, wctx.kv_enc_self.k, n_state*n_ctx, (ggml_element_size(wctx.kv_enc_self.k)*n_state)*(0*n_ctx));
-    //    //struct ggml_tensor * v = ggml_view_1d(ctx0, wctx.kv_enc_self.v, n_state*n_ctx, (ggml_element_size(wctx.kv_enc_self.v)*n_state)*(il*n_ctx));
-
-    //    ggml_build_forward_expand(&gf, ggml_cpy(ctx0, cur, k));
-    //    //ggml_build_forward_expand(&gf, ggml_cpy(ctx0, Vcur, v));
-    //}
-
     wctx.use_buf(ctx0, 3);
 
     // ===================================================================
@@ -1551,18 +1504,6 @@ static bool whisper_encode(
                         Vcur),
                     Vcur);
 
-            //{
-            //    //printf("Kcur: %d %d %d %d, size element = %d\n", Kcur->ne[0], Kcur->ne[1], Kcur->ne[2], Kcur->ne[3], ggml_element_size(Kcur));
-
-            //    wctx.use_buf(ctx0, -1);
-
-            //    struct ggml_tensor * k = ggml_view_1d(ctx0, wctx.kv_enc_self.k, n_state*n_ctx, (ggml_element_size(wctx.kv_enc_self.k)*n_state)*(il*n_ctx));
-            //    struct ggml_tensor * v = ggml_view_1d(ctx0, wctx.kv_enc_self.v, n_state*n_ctx, (ggml_element_size(wctx.kv_enc_self.v)*n_state)*(il*n_ctx));
-
-            //    ggml_build_forward_expand(&gf, ggml_cpy(ctx0, Kcur, k));
-            //    ggml_build_forward_expand(&gf, ggml_cpy(ctx0, Vcur, v));
-            //}
-
             // ------
 
             wctx.use_buf(ctx0, 0);
@@ -1647,18 +1588,6 @@ static bool whisper_encode(
             cur = ggml_cpy(ctx0,
                     KQV_merged,
                     ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_state, n_ctx));
-
-            {
-                //printf("cur: %d %d %d %d, size element = %d\n", cur->ne[0], cur->ne[1], cur->ne[2], cur->ne[3], ggml_element_size(cur));
-
-                wctx.use_buf(ctx0, -1);
-
-                struct ggml_tensor * k = ggml_view_1d(ctx0, wctx.kv_enc_self.k, n_state*n_ctx, (ggml_element_size(wctx.kv_enc_self.k)*n_state)*(il*n_ctx));
-                //struct ggml_tensor * v = ggml_view_1d(ctx0, wctx.kv_enc_self.v, n_state*n_ctx, (ggml_element_size(wctx.kv_enc_self.v)*n_state)*(il*n_ctx));
-
-                ggml_build_forward_expand(&gf, ggml_cpy(ctx0, cur, k));
-                //ggml_build_forward_expand(&gf, ggml_cpy(ctx0, Vcur, v));
-            }
         }
 
         // projection
@@ -1768,6 +1697,8 @@ static bool whisper_encode(
 
     // run the computation
     {
+        struct ggml_cgraph gf = {};
+        gf.n_threads = n_threads;
 
         ggml_build_forward_expand(&gf, cur);
         ggml_graph_compute       (ctx0, &gf);
@@ -1789,24 +1720,6 @@ static bool whisper_encode(
     //    printf("\n");
     //}
 
-    {
-        //const int i0 = std::min(mel_offset, mel_inp.n_len);
-        //const int i1 = std::min(mel_offset + 2*n_ctx, mel_inp.n_len);
-        const int i0 = 0;
-        const int i1 = cur->ne[1];
-
-        //printf("i0 = %d, i1 = %d, (i1 - i0) = %d, embd size = %d\n", i0, i1, i1 - i0, cur->ne[0]);
-
-        wctx.audio_embd.clear();
-        wctx.audio_embd.resize(cur->ne[0], 0.0f);
-        for (int j = 0; j < cur->ne[0]; ++j) {
-            for (int i = i0; i < i1; ++i) {
-                wctx.audio_embd[j] += ((float *)(cur->data))[(i - i0)*cur->ne[0] + j];
-            }
-            wctx.audio_embd[j] /= (i1 - i0);
-        }
-    }
-
     // pre-compute cross-attention memory
     {
         struct ggml_cgraph gf = {};
@@ -3049,6 +2962,9 @@ struct whisper_full_params whisper_full_default_params(enum whisper_sampling_str
 
         /*.encoder_begin_callback           =*/ nullptr,
         /*.encoder_begin_callback_user_data =*/ nullptr,
+
+        /*.logits_filter_callback           =*/ nullptr,
+        /*.logits_filter_callback_user_data =*/ nullptr,
     };
 
     switch (strategy) {
@@ -3176,7 +3092,7 @@ static const std::vector<std::string> non_speech_tokens = {
 // - applies logit filters
 // - computes logprobs and probs
 static void whisper_process_logits(
-        const struct whisper_context & ctx,
+              struct whisper_context & ctx,
     const struct whisper_full_params   params,
               struct whisper_decoder & decoder,
                                float   temperature) {
@@ -3232,6 +3148,9 @@ static void whisper_process_logits(
         logits[vocab.token_translate]  = -INFINITY;
         logits[vocab.token_transcribe] = -INFINITY;
 
+        if (params.logits_filter_callback) {
+            params.logits_filter_callback(&ctx, tokens_cur.data(), tokens_cur.size(), logits.data(), params.logits_filter_callback_user_data);
+        }
 
         // suppress non-speech tokens
         // ref: https://github.com/openai/whisper/blob/7858aa9c08d98f75575035ecd6481f462d66ca27/whisper/tokenizer.py#L224-L253
@@ -3935,7 +3854,7 @@ int whisper_full(
                         return a.sequence.sum_logprobs_all > b.sequence.sum_logprobs_all;
                     });
 
-                    unsigned int cur_c = 0;
+                    uint32_t cur_c = 0;
 
                     for (int j = 0; j < n_decoders_cur; ++j) {
                         auto & decoder = ctx->decoders[j];
@@ -4893,258 +4812,3 @@ static void whisper_exp_compute_token_level_timestamps(
     //    }
     //}
 }
-
-//
-// diarization stuff
-//
-
-void whisper_full_cluster_segments(struct whisper_context * ctx) {
-    const int n_segments = ctx->result_all.size();
-    printf("%s: clustering %d segments\n", __func__, n_segments);
-
-    const auto mel_len_save = ctx->mel.n_len;
-    printf("%s: mel_len_save = %d\n", __func__, mel_len_save);
-
-    const int n_ctx   = ctx->model.hparams.n_audio_ctx;
-    const int n_state = ctx->model.hparams.n_audio_state;
-    const int n_layer = ctx->model.hparams.n_audio_layer;
-
-#if 0
-    // use the last layer of the encoder
-    {
-        std::vector<float> embd(n_segments*n_state);
-
-        for (int i = 0; i < n_segments; ++i) {
-            const auto & segment_i = ctx->result_all[i];
-            printf("%s: segment %3d: t0 = %7d, t1 = %7d, text = %s\n", __func__, i, (int) segment_i.t0, (int) segment_i.t1, segment_i.text.c_str());
-
-            ctx->mel.n_len = segment_i.t1;
-            whisper_encode(*ctx, segment_i.t0, 7, true);
-
-            for (int j = 0; j < n_state; ++j) {
-                embd[i*n_state + j] = ctx->audio_embd[j];
-            }
-        }
-
-        const int n_features = std::min(4, n_segments);
-
-        ggml_svd_reduce_dims(n_state, n_segments, embd.data(), n_features);
-#elif 0
-    // use cross kv cache of various layers
-    for (int il = 0; il < n_layer; ++il) {
-        std::vector<float> embd(n_segments*n_ctx*n_state);
-
-        for (int i = 0; i < n_segments; ++i) {
-            const auto & segment_i = ctx->result_all[i];
-            printf("%s: layer %2d, segment %3d: t0 = %7d, t1 = %7d, text = %s\n", __func__, il, i, (int) segment_i.t0, (int) segment_i.t1, segment_i.text.c_str());
-
-            ctx->mel.n_len = segment_i.t1;
-            whisper_encode(*ctx, segment_i.t0, 7, true);
-
-            const size_t offs = ggml_element_size(ctx->kv_cross.k)*(il*n_ctx*n_state);
-            const ggml_fp16_t * f = (const ggml_fp16_t * )((const char *) ctx->kv_cross.k->data + offs);
-
-            for (int j = 0; j < n_ctx*n_state; ++j) {
-                embd[i*n_ctx*n_state + j] = ggml_fp16_to_fp32(f[j]);
-            }
-        }
-
-        const int n_features = std::min(4, n_segments);
-
-        ggml_svd_reduce_dims(n_ctx*n_state, n_segments, embd.data(), n_features);
-#elif 0
-    // use conv embedding
-    for (int il = 0; il < 1; ++il) {
-        std::vector<float> embd(n_segments*n_ctx*n_state);
-
-        for (int i = 0; i < n_segments; ++i) {
-            const auto & segment_i = ctx->result_all[i];
-            printf("%s: layer %2d, segment %3d: t0 = %7d, t1 = %7d, text = %s\n", __func__, il, i, (int) segment_i.t0, (int) segment_i.t1, segment_i.text.c_str());
-
-            ctx->mel.n_len = segment_i.t1;
-            whisper_encode(*ctx, segment_i.t0, 7, true);
-
-            const size_t offs = ggml_element_size(ctx->kv_enc_self.k)*(il*n_ctx*n_state);
-            const ggml_fp16_t * f = (const ggml_fp16_t * )((const char *) ctx->kv_enc_self.k->data + offs);
-
-            for (int j = 0; j < n_ctx*n_state; ++j) {
-                embd[i*n_ctx*n_state + j] = ggml_fp16_to_fp32(f[j]);
-            }
-        }
-
-        const int n_features = std::min(3, n_segments);
-
-        ggml_svd_reduce_dims(n_ctx*n_state, n_segments, embd.data(), n_features);
-#else
-    // use enc self kv cache of various layers
-    for (int il = 0; il < n_layer; ++il) {
-        std::vector<float> embd(n_segments*n_ctx*n_state);
-
-        for (int i = 0; i < n_segments; ++i) {
-            const auto & segment_i = ctx->result_all[i];
-            printf("%s: layer %2d, segment %3d: t0 = %7d, t1 = %7d, text = %s\n", __func__, il, i, (int) segment_i.t0, (int) segment_i.t1, segment_i.text.c_str());
-
-            ctx->mel.n_len = segment_i.t1;
-            whisper_encode(*ctx, segment_i.t0, 7, true);
-
-            const size_t offs = ggml_element_size(ctx->kv_enc_self.k)*(il*n_ctx*n_state);
-            const ggml_fp16_t * f = (const ggml_fp16_t * )((const char *) ctx->kv_enc_self.k->data + offs);
-
-            for (int j = 0; j < n_ctx*n_state; ++j) {
-                embd[i*n_ctx*n_state + j] = ggml_fp16_to_fp32(f[j]);
-            }
-        }
-
-        const int n_features = std::min(4, n_segments);
-
-        ggml_svd_reduce_dims(n_ctx*n_state, n_segments, embd.data(), n_features);
-#endif
-
-        std::vector<std::vector<double>> features(n_segments);
-
-        for (int i = 0; i < n_segments; ++i) {
-            features[i].resize(n_features);
-            for (int j = 0; j < n_features; ++j) {
-                features[i][j] = embd[i*n_features + j];
-            }
-        }
-
-        // fuzzy c-means clustering
-        const int n_clusters = 2;
-
-        std::vector<std::vector<double>> centroids(n_clusters,  std::vector<double>(n_features, 0.0));
-        std::vector<std::vector<double>> membership(n_segments, std::vector<double>(n_clusters, 0.0));
-
-        // initialize the centroids
-        for (int i = 0; i < n_clusters; ++i) {
-            for (int j = 0; j < n_features; ++j) {
-                centroids[i][j] = features[i][j];
-            }
-        }
-
-        // initialize the membership
-        for (int i = 0; i < n_segments; ++i) {
-            //membership[i][i % n_clusters] = 1.0;
-            //for (int j = 0; j < n_clusters; ++j) {
-            //    membership[i][j] = rand() / (float) RAND_MAX;
-            //}
-            for (int j = 0; j < n_clusters; ++j) {
-                membership[i][j] = 1.0 / n_clusters;
-            }
-        }
-
-        const int niter = 10000;
-
-        // iterate
-        for (int i = 0; i < niter; ++i) {
-            // print the membership
-            if (i == niter - 1) {
-            //{
-                for (int i = 0; i < n_segments; ++i) {
-#if 1
-                    printf("%s: membership %3d: ", __func__, i);
-                    for (int j = 0; j < n_clusters; ++j) {
-                        printf("%.1f ", membership[i][j]);
-                    }
-                    printf(" '%s'\n", ctx->result_all[i].text.c_str());
-#else
-                    printf("%s: features      : ", __func__);
-                    for (int j = 0; j < n_features; ++j) {
-                        printf("%8.3f ", features[i][j]);
-                    }
-                    printf(" '%s'\n", ctx->result_all[i].text.c_str());
-#endif
-                }
-                printf("----------------\n");
-
-                // print the centroids
-                for (int i = 0; i < n_clusters; ++i) {
-                    printf("%s: centroid %d: ", __func__, i);
-                    for (int j = 0; j < n_features; ++j) {
-                        printf("%f ", centroids[i][j]);
-                    }
-                    printf("\n");
-                }
-            }
-
-            // update the membership
-            for (int j = 0; j < n_segments; ++j) {
-                for (int k = 0; k < n_clusters; ++k) {
-                    double sum = 0.0;
-                    for (int l = 0; l < n_clusters; ++l) {
-                        //sum += std::pow(whisper_distance(features[j], centroids[k])/whisper_distance(features[j], centroids[l]), 2.0/(2.0 - 1.0));
-
-                        double d0 = 0.0;
-                        double d1 = 0.0;
-
-#if 1
-                        // use the euclidean distance
-                        {
-                            for (int m = 0; m < n_features; ++m) {
-                                d0 += std::pow(features[j][m] - centroids[k][m], 2.0);
-                            }
-                            d0 = std::sqrt(d0);
-
-                            for (int m = 0; m < n_features; ++m) {
-                                d1 += std::pow(features[j][m] - centroids[l][m], 2.0);
-                            }
-                            d1 = std::sqrt(d1);
-                        }
-#else
-                        // use the cosine distance
-                        {
-                            double dot = 0.0;
-                            double norm0 = 0.0;
-                            double norm1 = 0.0;
-
-                            for (int m = 0; m < n_features; ++m) {
-                                dot += features[j][m]*centroids[k][m];
-                                norm0 += std::pow(features[j][m], 2.0);
-                                norm1 += std::pow(centroids[k][m], 2.0);
-                            }
-
-                            d0 = 1.0 - dot/(std::sqrt(norm0)*std::sqrt(norm1));
-
-                            dot = 0.0;
-                            norm0 = 0.0;
-                            norm1 = 0.0;
-
-                            for (int m = 0; m < n_features; ++m) {
-                                dot += features[j][m]*centroids[l][m];
-                                norm0 += std::pow(features[j][m], 2.0);
-                                norm1 += std::pow(centroids[l][m], 2.0);
-                            }
-
-                            d1 = 1.0 - dot/(std::sqrt(norm0)*std::sqrt(norm1));
-                        }
-#endif
-
-                        if (d1 > 0.0) {
-                            sum += std::pow(d0/d1, 2.0/(1.20 - 1.0));
-                        } else {
-                            sum += 1.0;
-                        }
-                    }
-
-                    membership[j][k] = sum == 0.0 ? 1.0 : 1.0/sum;
-                }
-            }
-
-            // update the centroids
-            for (int j = 0; j < n_clusters; ++j) {
-                for (int k = 0; k < n_features; ++k) {
-                    double sum  = 0.0;
-                    double sum2 = 0.0;
-                    for (int l = 0; l < n_segments; ++l) {
-                        sum  += membership[l][j]*features[l][k];
-                        sum2 += membership[l][j];
-                    }
-                    centroids[j][k] = sum2 == 0.0 ? 0.0 : sum/sum2;
-                }
-            }
-        }
-    }
-
-    // restore the mel length
-    ctx->mel.n_len = mel_len_save;
-}

whisper.h

@@ -243,6 +243,16 @@ extern "C" {
     // If it returns false, the computation is aborted
     typedef bool (*whisper_encoder_begin_callback)(struct whisper_context * ctx, void * user_data);
 
+    // Logits filter callback
+    // Can be used to modify the logits before sampling
+    // If not NULL, called after applying temperature to logits
+    typedef void (*whisper_logits_filter_callback)(
+            struct whisper_context * ctx,
+          const whisper_token_data * tokens,
+                               int   n_tokens,
+                             float * logits,
+                              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:
     // whisper_full_default_params()
@@ -315,6 +325,10 @@ extern "C" {
         // called each time before the encoder starts
         whisper_encoder_begin_callback encoder_begin_callback;
         void * encoder_begin_callback_user_data;
+
+        // called by each decoder to filter obtained logits
+        whisper_logits_filter_callback logits_filter_callback;
+        void * logits_filter_callback_user_data;
     };
 
     WHISPER_API struct whisper_full_params whisper_full_default_params(enum whisper_sampling_strategy strategy);
@@ -372,10 +386,6 @@ extern "C" {
     WHISPER_API int whisper_bench_memcpy(int n_threads);
     WHISPER_API int whisper_bench_ggml_mul_mat(int n_threads);
 
-    // Temporary experimental API
-
-    WHISPER_API void whisper_full_cluster_segments(struct whisper_context * ctx);
-
 #ifdef __cplusplus
 }
 #endif