From eba33adaddf90ada65396f4dee1715e033ff8d4a Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 4 Oct 2022 22:43:37 +0300 Subject: [PATCH] Extend C-style API with full inference methods --- main.cpp | 188 +++++-------------------------------- stream.cpp | 192 +++++-------------------------------- whisper.cpp | 266 ++++++++++++++++++++++++++++++++++++++++++++++++++-- whisper.h | 30 ++++-- 4 files changed, 328 insertions(+), 348 deletions(-) diff --git a/main.cpp b/main.cpp index cc36322..ca29b82 100644 --- a/main.cpp +++ b/main.cpp @@ -5,17 +5,11 @@ #define DR_WAV_IMPLEMENTATION #include "dr_wav.h" -#include #include #include #include #include -int64_t get_time_us() { - return std::chrono::duration_cast( - std::chrono::high_resolution_clock::now().time_since_epoch()).count(); -} - // 500 -> 00:05.000 // 6000 -> 01:00.000 std::string to_timestamp(int64_t t) { @@ -30,11 +24,6 @@ std::string to_timestamp(int64_t t) { return std::string(buf); } -struct whisper_result { - whisper_token id; - int64_t t; -}; - // command-line parameters struct whisper_params { int32_t seed = -1; // RNG seed, not used currently @@ -111,8 +100,6 @@ void whisper_print_usage(int argc, char ** argv, const whisper_params & params) } int main(int argc, char ** argv) { - const int64_t t_main_start_us = get_time_us(); - whisper_params params; if (whisper_params_parse(argc, argv, params) == false) { @@ -142,7 +129,7 @@ int main(int argc, char ** argv) { return 3; } - if (wav.sampleRate != SAMPLE_RATE) { + if (wav.sampleRate != WHISPER_SAMPLE_RATE) { fprintf(stderr, "%s: WAV file '%s' must be 16 kHz\n", argv[0], params.fname_inp.c_str()); return 4; } @@ -172,12 +159,6 @@ int main(int argc, char ** argv) { } } - // compute log mel spectrogram - if (whisper_pcm_to_mel(ctx, pcmf32.data(), pcmf32.size(), params.n_threads) != 0) { - fprintf(stderr, "%s: failed to compute log mel spectrogram\n", argv[0]); - return 6; - } - // print some info about the processing { printf("\n"); @@ -189,168 +170,43 @@ int main(int argc, char ** argv) { } } printf("%s: processing %d samples (%.1f sec), %d threads, lang = %s, task = %s, timestamps = %d ...\n", - __func__, int(pcmf32.size()), float(pcmf32.size())/SAMPLE_RATE, params.n_threads, + __func__, int(pcmf32.size()), float(pcmf32.size())/WHISPER_SAMPLE_RATE, params.n_threads, params.language.c_str(), params.translate ? "translate" : "transcribe", params.no_timestamps ? 0 : 1); printf("\n"); } - // the accumulated text context so far - std::vector prompt_past = { }; - - // these tokens determine the task that will be performed - std::vector prompt_init = { whisper_token_sot(ctx) }; - if (whisper_is_multilingual(ctx)) { - prompt_init.push_back(whisper_token_sot(ctx) + 1 + whisper_lang_id(params.language.c_str())); - if (params.translate) { - prompt_init.push_back(whisper_token_translate()); - } else { - prompt_init.push_back(whisper_token_transcribe()); - } - } - - // the generated text including timestamps - //std::vector result_all; - - // main loop - int seek = 0; - while (true) { - if (seek >= whisper_n_len(ctx)) { - break; - } - - // encode audio features starting at offset seek - if (whisper_encode(ctx, seek, params.n_threads) != 0) { - fprintf(stderr, "%s: failed to encode\n", __func__); - return 7; - } - - std::vector prompt; - - int n_past = 0; - - // if we have already generated some text, use it as a prompt to condition the next generation - if (prompt_past.size() > 0) { - int n_take = std::min(whisper_n_text_ctx(ctx)/2, int(prompt_past.size())); - - prompt = { whisper_token_prev(ctx) }; - prompt.insert(prompt.begin() + 1, prompt_past.end() - n_take, prompt_past.end()); - - prompt_past.clear(); - prompt_past.insert(prompt_past.end(), prompt.begin() + 1, prompt.end()); - } - - prompt.insert(prompt.end(), prompt_init.begin(), prompt_init.end()); - - bool done = false; - int seek_delta = 100*CHUNK_SIZE; - whisper_token last_id = 0; - - // print the prompt - //printf("\n\n"); - //for (int i = 0; i < prompt.size(); i++) { - // printf("%s: prompt[%d] = %s\n", __func__, i, vocab.id_to_token[prompt[i]].c_str()); - //} - //printf("\n\n"); - - // the accumulated transcription in the current interation - int result_len = 0; - std::vector result_cur; - - for (int i = 0; i < whisper_n_text_ctx(ctx)/2 - 4; ++i) { - if (whisper_decode(ctx, prompt.data(), prompt.size(), n_past, params.n_threads) != 0) { - fprintf(stderr, "%s: failed to decode\n", __func__); - return 8; - } - - n_past += prompt.size(); - prompt.clear(); - - // very basic greedy sampling strategy: - // - // - always take the most probable token - // - // more sophisticated sampling strategies could be implemented here, but we keep it simple - // feel free to experiment! - // - { - const int n_vocab = whisper_n_vocab(ctx); - - whisper_token id = 0; - whisper_token tid = whisper_token_beg(ctx); - - id = whisper_sample_best(ctx, result_len == 0); - if (i > 0) { - tid = whisper_sample_timestamp(ctx); - } - - // update sliding window - if (id > whisper_token_beg(ctx)) { - seek_delta = 2*(id - whisper_token_beg(ctx)); - result_len = i + 1; - } - last_id = id; - - // add it to the context - prompt.push_back(id); - result_cur.push_back({ id, seek + 2*(tid - whisper_token_beg(ctx)) }); - - //printf("%s: %s\n", __func__, vocab.id_to_token[id].c_str()); + // run the inference + { + whisper_full_params wparams = whisper_full_default_params(WHISPER_DECODE_GREEDY); - // end of text token - if (id == whisper_token_eot(ctx)) { - break; - } - } + wparams.print_special_tokens = params.print_special_tokens; - if (done) { - break; - } + if (whisper_full(ctx, wparams, pcmf32.data(), pcmf32.size()) != 0) { + fprintf(stderr, "%s: failed to process audio\n", argv[0]); + return 6; } - result_cur.resize(result_len); - //result_all.insert(result_all.end(), result_cur.begin(), result_cur.end()); - - for (const auto & r : result_cur) { - prompt_past.push_back(r.id); - } + // print result; + { + printf("\n"); - // print the text from this iteration - if (result_cur.size() > 0) { - auto t0 = result_cur.front().t; + const int n_segments = whisper_full_n_segments(ctx); + for (int i = 0; i < n_segments; ++i) { + const char * text = whisper_full_get_segment_text(ctx, i); - std::string text = ""; - for (int i = 0; i < result_cur.size(); i++) { - if (params.print_special_tokens == false && result_cur[i].id >= whisper_token_eot(ctx)) { + if (params.no_timestamps) { + printf ("%s", text); + fflush(stdout); } else { - text += whisper_token_to_str(ctx, result_cur[i].id); - } - if (result_cur[i].id > whisper_token_beg(ctx)) { - const auto t1 = result_cur[i].t; - if (!text.empty()) { - if (params.no_timestamps) { - printf ("%s", text.c_str()); - fflush(stdout); - } else { - printf ("[%s --> %s] %s\n", to_timestamp(t0).c_str(), to_timestamp(t1).c_str(), text.c_str()); - } - } - text = ""; - while (result_cur[i].id > whisper_token_beg(ctx) && i < result_cur.size()) { - i++; - } - i--; - t0 = result_cur[i].t; - } - } + const int64_t t0 = whisper_full_get_segment_t0(ctx, i); + const int64_t t1 = whisper_full_get_segment_t1(ctx, i); - if (!text.empty()) { - printf ("[%s --> %s] %s\n", to_timestamp(t0).c_str(), to_timestamp(seek + seek_delta).c_str(), text.c_str()); + printf ("[%s --> %s] %s\n", to_timestamp(t0).c_str(), to_timestamp(t1).c_str(), text); + } } } - - seek += seek_delta; } whisper_print_timings(ctx); diff --git a/stream.cpp b/stream.cpp index 0f5a09e..d0e40c2 100644 --- a/stream.cpp +++ b/stream.cpp @@ -18,11 +18,6 @@ #include #include -int64_t get_time_us() { - return std::chrono::duration_cast( - std::chrono::high_resolution_clock::now().time_since_epoch()).count(); -} - // 500 -> 00:05.000 // 6000 -> 01:00.000 std::string to_timestamp(int64_t t) { @@ -37,11 +32,6 @@ std::string to_timestamp(int64_t t) { return std::string(buf); } -struct whisper_result { - whisper_token id; - int64_t t; -}; - // command-line parameters struct whisper_params { int32_t seed = -1; // RNG seed, not used currently @@ -155,7 +145,7 @@ bool audio_sdl_init(const int capture_id) { SDL_zero(capture_spec_requested); SDL_zero(capture_spec_obtained); - capture_spec_requested.freq = SAMPLE_RATE; + capture_spec_requested.freq = WHISPER_SAMPLE_RATE; capture_spec_requested.format = AUDIO_F32; capture_spec_requested.channels = 1; capture_spec_requested.samples = 1024; @@ -186,8 +176,6 @@ bool audio_sdl_init(const int capture_id) { /////////////////////////// int main(int argc, char ** argv) { - const int64_t t_main_start_us = get_time_us(); - whisper_params params; if (whisper_params_parse(argc, argv, params) == false) { @@ -209,7 +197,7 @@ int main(int argc, char ** argv) { struct whisper_context * ctx = whisper_init(params.model.c_str()); - const int n_samples_30s = 30*SAMPLE_RATE; + const int n_samples_30s = 30*WHISPER_SAMPLE_RATE; std::vector pcmf32(n_samples_30s, 0.0f); std::vector pcmf32_old; @@ -224,7 +212,7 @@ int main(int argc, char ** argv) { } } printf("%s: processing %d samples (%.1f sec), %d threads, lang = %s, task = %s, timestamps = %d ...\n", - __func__, int(pcmf32.size()), float(pcmf32.size())/SAMPLE_RATE, params.n_threads, + __func__, int(pcmf32.size()), float(pcmf32.size())/WHISPER_SAMPLE_RATE, params.n_threads, params.language.c_str(), params.translate ? "translate" : "transcribe", params.no_timestamps ? 0 : 1); @@ -250,7 +238,7 @@ int main(int argc, char ** argv) { } // process 3 seconds of new audio - while ((int) SDL_GetQueuedAudioSize(g_dev_id_in) < 3*SAMPLE_RATE*sizeof(float)) { + while ((int) SDL_GetQueuedAudioSize(g_dev_id_in) < 3*WHISPER_SAMPLE_RATE*sizeof(float)) { SDL_Delay(1); } const int n_samples_new = SDL_GetQueuedAudioSize(g_dev_id_in)/sizeof(float); @@ -271,167 +259,37 @@ int main(int argc, char ** argv) { pcmf32_old = pcmf32; - // compute log mel spectrogram - if (whisper_pcm_to_mel(ctx, pcmf32.data(), pcmf32.size(), params.n_threads) != 0) { - fprintf(stderr, "%s: failed to compute log mel spectrogram\n", argv[0]); - return 6; - } - - // the accumulated text context so far - std::vector prompt_past = { }; - - // these tokens determine the task that will be performed - std::vector prompt_init = { whisper_token_sot(ctx) }; - if (whisper_is_multilingual(ctx)) { - prompt_init.push_back(whisper_token_sot(ctx) + 1 + whisper_lang_id(params.language.c_str())); - if (params.translate) { - prompt_init.push_back(whisper_token_translate()); - } else { - prompt_init.push_back(whisper_token_transcribe()); - } - } - - // the generated text including timestamps - //std::vector result_all; - - // main loop - int seek = 0; - while (true) { - if (seek >= whisper_n_len(ctx)) { - break; - } - - // encode audio features starting at offset seek - if (whisper_encode(ctx, seek, params.n_threads) != 0) { - fprintf(stderr, "%s: failed to encode\n", __func__); - return 7; - } - - std::vector prompt; - - int n_past = 0; - - // if we have already generated some text, use it as a prompt to condition the next generation - if (prompt_past.size() > 0) { - int n_take = std::min(whisper_n_text_ctx(ctx)/2, int(prompt_past.size())); - - prompt = { whisper_token_prev(ctx) }; - prompt.insert(prompt.begin() + 1, prompt_past.end() - n_take, prompt_past.end()); - - prompt_past.clear(); - prompt_past.insert(prompt_past.end(), prompt.begin() + 1, prompt.end()); - } - - prompt.insert(prompt.end(), prompt_init.begin(), prompt_init.end()); - - bool done = false; - int seek_delta = 100*CHUNK_SIZE; - whisper_token last_id = 0; - - // print the prompt - //printf("\n\n"); - //for (int i = 0; i < prompt.size(); i++) { - // printf("%s: prompt[%d] = %s\n", __func__, i, vocab.id_to_token[prompt[i]].c_str()); - //} - //printf("\n\n"); - - // the accumulated transcription in the current interation - int result_len = 0; - std::vector result_cur; - - for (int i = 0; i < whisper_n_text_ctx(ctx)/2 - 4; ++i) { - if (whisper_decode(ctx, prompt.data(), prompt.size(), n_past, params.n_threads) != 0) { - fprintf(stderr, "%s: failed to decode\n", __func__); - return 8; - } - - n_past += prompt.size(); - prompt.clear(); - - // very basic greedy sampling strategy: - // - // - always take the most probable token - // - // more sophisticated sampling strategies could be implemented here, but we keep it simple - // feel free to experiment! - // - { - const int n_vocab = whisper_n_vocab(ctx); - - whisper_token id = 0; - whisper_token tid = whisper_token_beg(ctx); - - id = whisper_sample_best(ctx, result_len == 0); - if (i > 0) { - tid = whisper_sample_timestamp(ctx); - } - - // update sliding window - if (id > whisper_token_beg(ctx)) { - seek_delta = 2*(id - whisper_token_beg(ctx)); - result_len = i + 1; - } - last_id = id; - - // add it to the context - prompt.push_back(id); - result_cur.push_back({ id, seek + 2*(tid - whisper_token_beg(ctx)) }); - - //printf("%s: %s\n", __func__, vocab.id_to_token[id].c_str()); + // run the inference + { + whisper_full_params wparams = whisper_full_default_params(WHISPER_DECODE_GREEDY); - // end of text token - if (id == whisper_token_eot(ctx)) { - break; - } - } + wparams.print_progress = false; + wparams.print_special_tokens = params.print_special_tokens; - if (done) { - break; - } + if (whisper_full(ctx, wparams, pcmf32.data(), pcmf32.size()) != 0) { + fprintf(stderr, "%s: failed to process audio\n", argv[0]); + return 6; } - result_cur.resize(result_len); - //result_all.insert(result_all.end(), result_cur.begin(), result_cur.end()); - - for (const auto & r : result_cur) { - prompt_past.push_back(r.id); - } + // print result; + { + printf("\n"); - // print the text from this iteration - if (result_cur.size() > 0) { - auto t0 = result_cur.front().t; + const int n_segments = whisper_full_n_segments(ctx); + for (int i = 0; i < n_segments; ++i) { + const char * text = whisper_full_get_segment_text(ctx, i); - std::string text = ""; - for (int i = 0; i < result_cur.size(); i++) { - if (params.print_special_tokens == false && result_cur[i].id >= whisper_token_eot(ctx)) { + if (params.no_timestamps) { + printf ("%s", text); + fflush(stdout); } else { - text += whisper_token_to_str(ctx, result_cur[i].id); - } - if (result_cur[i].id > whisper_token_beg(ctx)) { - const auto t1 = result_cur[i].t; - if (!text.empty()) { - if (params.no_timestamps) { - printf ("%s", text.c_str()); - fflush(stdout); - } else { - printf ("[%s --> %s] %s\n", to_timestamp(t0).c_str(), to_timestamp(t1).c_str(), text.c_str()); - } - } - text = ""; - while (result_cur[i].id > whisper_token_beg(ctx) && i < result_cur.size()) { - i++; - } - i--; - t0 = result_cur[i].t; - } - } + const int64_t t0 = whisper_full_get_segment_t0(ctx, i); + const int64_t t1 = whisper_full_get_segment_t1(ctx, i); - if (!text.empty()) { - printf ("[%s --> %s] %s\n", to_timestamp(t0).c_str(), to_timestamp(seek + seek_delta).c_str(), text.c_str()); + printf ("[%s --> %s] %s\n", to_timestamp(t0).c_str(), to_timestamp(t1).c_str(), text); + } } } - - seek += seek_delta; } } diff --git a/whisper.cpp b/whisper.cpp index c57ca60..583eb72 100644 --- a/whisper.cpp +++ b/whisper.cpp @@ -210,8 +210,15 @@ struct whisper_vocab { }; struct whisper_result { - whisper_vocab::id id; int64_t t; + whisper_token id; +}; + +struct whisper_segment { + int64_t t0; + int64_t t1; + + std::string text; }; // medium @@ -395,6 +402,9 @@ struct whisper_context { std::vector probs; std::vector logits; + + std::vector result_cur; + std::vector result_all; }; // load the model from a ggml file @@ -1946,8 +1956,8 @@ bool log_mel_spectrogram( const int n_fft = 1 + fft_size/2; - printf("%s: n_samples = %d, n_len = %d\n", __func__, n_samples, mel.n_len); - printf("%s: recording length: %f s\n", __func__, (float) n_samples/sample_rate); + //printf("%s: n_samples = %d, n_len = %d\n", __func__, n_samples, mel.n_len); + //printf("%s: recording length: %f s\n", __func__, (float) n_samples/sample_rate); std::vector workers(n_threads); for (int iw = 0; iw < n_threads; ++iw) { @@ -2066,7 +2076,7 @@ void whisper_free(struct whisper_context * ctx) { int whisper_pcm_to_mel(struct whisper_context * ctx, const float * samples, int n_samples, int n_threads) { const int64_t t_start_us = ggml_time_us(); - if (!log_mel_spectrogram(samples, n_samples, SAMPLE_RATE, N_FFT, HOP_LENGTH, N_MEL, n_threads, ctx->model.filters, ctx->mel)) { + if (!log_mel_spectrogram(samples, n_samples, WHISPER_SAMPLE_RATE, WHISPER_N_FFT, WHISPER_HOP_LENGTH, WHISPER_N_MEL, n_threads, ctx->model.filters, ctx->mel)) { fprintf(stderr, "%s: failed to compute mel spectrogram\n", __func__); return -1; } @@ -2081,8 +2091,8 @@ int whisper_set_mel( const float * data, int n_len, int n_mel) { - if (n_mel != N_MEL) { - fprintf(stderr, "%s: invalid number of mel bands: %d (expected %d)\n", __func__, n_mel, N_MEL); + if (n_mel != WHISPER_N_MEL) { + fprintf(stderr, "%s: invalid number of mel bands: %d (expected %d)\n", __func__, n_mel, WHISPER_N_MEL); return -1; } @@ -2219,3 +2229,247 @@ void whisper_print_timings(struct whisper_context * ctx) { printf("%s: decode time = %8.2f ms / %.2f ms per layer\n", __func__, ctx->t_decode_us/1000.0f, ctx->t_decode_us/1000.0f/ctx->model.hparams.n_text_layer); printf("%s: total time = %8.2f ms\n", __func__, (t_end_us - ctx->t_start_us)/1000.0f); } + +//////////////////////////////////////////////////////////////////////////// + +struct whisper_full_params whisper_full_default_params(enum whisper_decode_strategy strategy) { + struct whisper_full_params result; + + switch (strategy) { + case WHISPER_DECODE_GREEDY: + { + result = (struct whisper_full_params) { + .strategy = WHISPER_DECODE_GREEDY, + .n_threads = std::min(4, (int32_t) std::thread::hardware_concurrency()), + + .translate = false, + .print_special_tokens = false, + .print_progress = true, + + .language = "en", + + .greedy = { + .n_past = 0, + }, + }; + } break; + case WHISPER_DECODE_BEAM_SEARCH: + { + result = (struct whisper_full_params) { + .strategy = WHISPER_DECODE_GREEDY, + .n_threads = std::min(4, (int32_t) std::thread::hardware_concurrency()), + + .translate = false, + .print_special_tokens = false, + .print_progress = true, + + .language = "en", + + .beam_search = { + .n_past = 0, + .beam_width = 10, + .n_best = 5, + }, + }; + } break; + } + + return result; +} +int whisper_full( + struct whisper_context * ctx, + struct whisper_full_params params, + const float * samples, + int n_samples) { + // compute log mel spectrogram + if (whisper_pcm_to_mel(ctx, samples, n_samples, params.n_threads) != 0) { + fprintf(stderr, "%s: failed to compute log mel spectrogram\n", __func__); + return -1; + } + + // the accumulated text context so far + std::vector prompt_past = { }; + + // these tokens determine the task that will be performed + std::vector prompt_init = { whisper_token_sot(ctx) }; + if (whisper_is_multilingual(ctx)) { + prompt_init.push_back(whisper_token_sot(ctx) + 1 + whisper_lang_id(params.language)); + if (params.translate) { + prompt_init.push_back(whisper_token_translate()); + } else { + prompt_init.push_back(whisper_token_transcribe()); + } + } + + auto & result_all = ctx->result_all; + auto & result_cur = ctx->result_cur; + + result_all.clear(); + + int progress_prev = 0; + int progress_step = 5; + + // main loop + int seek = 0; + while (true) { + int progress_cur = (100*seek)/whisper_n_len(ctx); + while (progress_cur >= progress_prev + progress_step) { + progress_prev += progress_step; + if (params.print_progress) { + printf("%s: progress = %3d%%\n", __func__, progress_prev); + } + } + + if (seek >= whisper_n_len(ctx)) { + break; + } + + // encode audio features starting at offset seek + if (whisper_encode(ctx, seek, params.n_threads) != 0) { + fprintf(stderr, "%s: failed to encode\n", __func__); + return 7; + } + + std::vector prompt; + + int n_past = 0; + + // if we have already generated some text, use it as a prompt to condition the next generation + if (prompt_past.size() > 0) { + int n_take = std::min(whisper_n_text_ctx(ctx)/2, int(prompt_past.size())); + + prompt = { whisper_token_prev(ctx) }; + prompt.insert(prompt.begin() + 1, prompt_past.end() - n_take, prompt_past.end()); + + prompt_past.clear(); + prompt_past.insert(prompt_past.end(), prompt.begin() + 1, prompt.end()); + } + + prompt.insert(prompt.end(), prompt_init.begin(), prompt_init.end()); + + bool done = false; + int seek_delta = 100*WHISPER_CHUNK_SIZE; + whisper_token last_id = 0; + + // print the prompt + //printf("\n\n"); + //for (int i = 0; i < prompt.size(); i++) { + // printf("%s: prompt[%d] = %s\n", __func__, i, vocab.id_to_token[prompt[i]].c_str()); + //} + //printf("\n\n"); + + // the accumulated transcription in the current interation + int result_len = 0; + result_cur.clear(); + + for (int i = 0; i < whisper_n_text_ctx(ctx)/2 - 4; ++i) { + if (whisper_decode(ctx, prompt.data(), prompt.size(), n_past, params.n_threads) != 0) { + fprintf(stderr, "%s: failed to decode\n", __func__); + return 8; + } + + n_past += prompt.size(); + prompt.clear(); + + // very basic greedy sampling strategy: + // + // - always take the most probable token + // + // more sophisticated sampling strategies could be implemented here, but we keep it simple + // feel free to experiment! + // + { + const int n_vocab = whisper_n_vocab(ctx); + + whisper_token id = 0; + whisper_token tid = whisper_token_beg(ctx); + + id = whisper_sample_best(ctx, result_len == 0); + if (i > 0) { + tid = whisper_sample_timestamp(ctx); + } + + // update sliding window + if (id > whisper_token_beg(ctx)) { + seek_delta = 2*(id - whisper_token_beg(ctx)); + result_len = i + 1; + } + last_id = id; + + // add it to the context + prompt.push_back(id); + result_cur.push_back({ seek + 2*(tid - whisper_token_beg(ctx)), id }); + + //printf("%s: %s\n", __func__, ctx->vocab.id_to_token[id].c_str()); + + // end of text token + if (id == whisper_token_eot(ctx)) { + if (result_len == 0) { + result_len = i + 1; + } + break; + } + } + + if (done) { + break; + } + } + + result_cur.resize(result_len); + + for (const auto & r : result_cur) { + prompt_past.push_back(r.id); + } + + // store the text from this iteration + if (result_cur.size() > 0) { + auto t0 = result_cur.front().t; + + std::string text = ""; + + for (int i = 0; i < result_cur.size(); i++) { + if (params.print_special_tokens == false && result_cur[i].id >= whisper_token_eot(ctx)) { + } else { + text += whisper_token_to_str(ctx, result_cur[i].id); + } + if (result_cur[i].id > whisper_token_beg(ctx)) { + const auto t1 = result_cur[i].t; + if (!text.empty()) { + result_all.push_back({ t0, t1, text }); + } + text = ""; + while (result_cur[i].id > whisper_token_beg(ctx) && i < result_cur.size()) { + i++; + } + i--; + t0 = result_cur[i].t; + } + } + + if (!text.empty()) { + result_all.push_back({ t0, seek + seek_delta, text }); + } + } + + seek += seek_delta; + } + + return 0; +} + +int whisper_full_n_segments(struct whisper_context * ctx) { + return ctx->result_all.size(); +} + +int64_t whisper_full_get_segment_t0(struct whisper_context * ctx, int i_segment) { + return ctx->result_all[i_segment].t0; +} + +int64_t whisper_full_get_segment_t1(struct whisper_context * ctx, int i_segment) { + return ctx->result_all[i_segment].t1; +} + +const char * whisper_full_get_segment_text(struct whisper_context * ctx, int i_segment) { + return ctx->result_all[i_segment].text.c_str(); +} diff --git a/whisper.h b/whisper.h index c282aa2..cc24d71 100644 --- a/whisper.h +++ b/whisper.h @@ -1,6 +1,8 @@ #ifndef WHISPER_H #define WHISPER_H +#include + #ifdef WHISPER_SHARED # ifdef _WIN32 # ifdef WHISPER_BUILD @@ -15,6 +17,12 @@ # define WHISPER_API #endif +#define WHISPER_SAMPLE_RATE 16000 +#define WHISPER_N_FFT 400 +#define WHISPER_N_MEL 80 +#define WHISPER_HOP_LENGTH 160 +#define WHISPER_CHUNK_SIZE 30 + #ifdef __cplusplus extern "C" { #endif @@ -23,12 +31,6 @@ extern "C" { // C interface // -#define SAMPLE_RATE 16000 -#define N_FFT 400 -#define N_MEL 80 -#define HOP_LENGTH 160 -#define CHUNK_SIZE 30 - // TODO: documentation will come soon struct whisper_context; @@ -101,7 +103,9 @@ extern "C" { int n_threads; - bool transcribe; + bool translate; + bool print_special_tokens; + bool print_progress; const char * language; @@ -118,14 +122,22 @@ extern "C" { }; }; + WHISPER_API struct whisper_full_params whisper_full_default_params(enum whisper_decode_strategy strategy); + // full whisper run - encode + decode - // TODO: implement WHISPER_API int whisper_full( struct whisper_context * ctx, - struct whisper_full_params * params, + struct whisper_full_params params, const float * samples, int n_samples); + WHISPER_API int whisper_full_n_segments(struct whisper_context * ctx); + + WHISPER_API int64_t whisper_full_get_segment_t0(struct whisper_context * ctx, int i_segment); + WHISPER_API int64_t whisper_full_get_segment_t1(struct whisper_context * ctx, int i_segment); + + WHISPER_API const char * whisper_full_get_segment_text(struct whisper_context * ctx, int i_segment); + #ifdef __cplusplus } #endif