Extend C-style API with full inference methods

This commit is contained in:
Georgi Gerganov
2022-10-04 22:43:37 +03:00
parent 6b77124e01
commit eba33adadd
4 changed files with 328 additions and 348 deletions

188
main.cpp
View File

@ -5,17 +5,11 @@
#define DR_WAV_IMPLEMENTATION
#include "dr_wav.h"
#include <cassert>
#include <cstdio>
#include <string>
#include <thread>
#include <vector>
int64_t get_time_us() {
return std::chrono::duration_cast<std::chrono::microseconds>(
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<whisper_token> prompt_past = { };
// run the inference
{
whisper_full_params wparams = whisper_full_default_params(WHISPER_DECODE_GREEDY);
// these tokens determine the task that will be performed
std::vector<whisper_token> 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());
}
}
wparams.print_special_tokens = params.print_special_tokens;
// the generated text including timestamps
//std::vector<whisper_result> result_all;
// main loop
int seek = 0;
while (true) {
if (seek >= whisper_n_len(ctx)) {
break;
if (whisper_full(ctx, wparams, pcmf32.data(), pcmf32.size()) != 0) {
fprintf(stderr, "%s: failed to process audio\n", argv[0]);
return 6;
}
// 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;
}
// print result;
{
printf("\n");
std::vector<whisper_token> prompt;
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);
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<whisper_result> 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());
// end of text token
if (id == whisper_token_eot(ctx)) {
break;
}
}
if (done) {
break;
}
}
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 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)) {
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);