Skip to content

Commit dfc6ca6

Browse files
authored
stream : add beam size parameter(#2836)
* feat: Add beam size parameter to stream.cpp for beam search configuration * feat: Add beam size parameter to whisper full params in stream example * fix: Remove duplicate beam search size assignment in server.cpp
1 parent 47e14c0 commit dfc6ca6

File tree

1 file changed

+5
-1
lines changed

1 file changed

+5
-1
lines changed

examples/stream/stream.cpp

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ struct whisper_params {
2323
int32_t capture_id = -1;
2424
int32_t max_tokens = 32;
2525
int32_t audio_ctx = 0;
26+
int32_t beam_size = -1;
2627

2728
float vad_thold = 0.6f;
2829
float freq_thold = 100.0f;
@@ -59,6 +60,7 @@ static bool whisper_params_parse(int argc, char ** argv, whisper_params & params
5960
else if (arg == "-c" || arg == "--capture") { params.capture_id = std::stoi(argv[++i]); }
6061
else if (arg == "-mt" || arg == "--max-tokens") { params.max_tokens = std::stoi(argv[++i]); }
6162
else if (arg == "-ac" || arg == "--audio-ctx") { params.audio_ctx = std::stoi(argv[++i]); }
63+
else if (arg == "-bs" || arg == "--beam-size") { params.beam_size = std::stoi(argv[++i]); }
6264
else if (arg == "-vth" || arg == "--vad-thold") { params.vad_thold = std::stof(argv[++i]); }
6365
else if (arg == "-fth" || arg == "--freq-thold") { params.freq_thold = std::stof(argv[++i]); }
6466
else if (arg == "-tr" || arg == "--translate") { params.translate = true; }
@@ -96,6 +98,7 @@ void whisper_print_usage(int /*argc*/, char ** argv, const whisper_params & para
9698
fprintf(stderr, " -c ID, --capture ID [%-7d] capture device ID\n", params.capture_id);
9799
fprintf(stderr, " -mt N, --max-tokens N [%-7d] maximum number of tokens per audio chunk\n", params.max_tokens);
98100
fprintf(stderr, " -ac N, --audio-ctx N [%-7d] audio context size (0 - all)\n", params.audio_ctx);
101+
fprintf(stderr, " -bs N, --beam-size N [%-7d] beam size for beam search\n", params.beam_size);
99102
fprintf(stderr, " -vth N, --vad-thold N [%-7.2f] voice activity detection threshold\n", params.vad_thold);
100103
fprintf(stderr, " -fth N, --freq-thold N [%-7.2f] high-pass frequency cutoff\n", params.freq_thold);
101104
fprintf(stderr, " -tr, --translate [%-7s] translate from source language to english\n", params.translate ? "true" : "false");
@@ -298,7 +301,7 @@ int main(int argc, char ** argv) {
298301

299302
// run the inference
300303
{
301-
whisper_full_params wparams = whisper_full_default_params(WHISPER_SAMPLING_GREEDY);
304+
whisper_full_params wparams = whisper_full_default_params(params.beam_size > 1 ? WHISPER_SAMPLING_BEAM_SEARCH : WHISPER_SAMPLING_GREEDY);
302305

303306
wparams.print_progress = false;
304307
wparams.print_special = params.print_special;
@@ -309,6 +312,7 @@ int main(int argc, char ** argv) {
309312
wparams.max_tokens = params.max_tokens;
310313
wparams.language = params.language.c_str();
311314
wparams.n_threads = params.n_threads;
315+
wparams.beam_search.beam_size = params.beam_size;
312316

313317
wparams.audio_ctx = params.audio_ctx;
314318

0 commit comments

Comments
 (0)