stream.cpp improved

[shadow-wxh]

Here is a bit more developed version of stream.cpp which is less crude and offer more accuracy. It is developed for windows command console you will need to alter the configuration to suit your environment.

Things that I have changed:
1 VAD has been remove because it doesn't seem to be very accurate.
2 Lines changes with alternating color.
3 Transcription between sections are spliced correctly.
4 Timestamp uses local time instead of transcription time.
5 Section length must be integer times of of step length.
6 Removed coding redundancy.

#define WIN32_LEAN_AND_MEAN 
#define NOMINMAX
#define WHISPER_SAMPLE_RATE 16000
#define WHISPER_SAMPLE_RATE_MS 16
#include "common-sdl.h"
#include "common.h"
#include "whisper.h"
#include <cassert>
#include <cstdio>
#include <string>
#include <thread>
#include <vector>
#include <fstream>
#include <Windows.h>


struct whisper_params {
int32_t n_threads  = std::min(4, (int32_t) std::thread::hardware_concurrency());
int32_t step_ms    = 3000;
int32_t step_no    = 3;
int32_t line_sz    = 100;
int32_t capture_id = -1;
int32_t max_tokens = 32;
int32_t audio_ctx  = 0;
bool translate     = false;
bool no_fallback   = false;
bool print_special = false;
bool keep_context  = false;
bool print_time    = false;
bool tinydiarize   = false;
bool save_audio    = false; // save audio to wav file
bool use_gpu       = true;
bool flash_attn    = false;
std::string language  = "en";
std::string model     = "openai-mutilan-large-v3-turbo.bin";
std::string logfile;
};

void whisper_print_usage(int argc, char ** argv, const whisper_params & params);

static bool whisper_params_parse(int argc, char ** argv, whisper_params & params) {
for (int i = 1; i < argc; i++) {
    std::string arg = argv[i];
    if (arg == "-h" || arg == "--help") {
        whisper_print_usage(argc, argv, params);
        exit(0);
    }
    else if (arg == "-t"    || arg == "--threads")       { params.n_threads     = std::stoi(argv[++i]); }
    else if (                  arg == "--step-length")   { params.step_ms       = std::stoi(argv[++i]); }
    else if (                  arg == "--step-number")   { params.step_no       = std::stoi(argv[++i]); }
    else if (                  arg == "--line-size")     { params.line_sz       = std::stoi(argv[++i]); }
    else if (arg == "-c"    || arg == "--capture")       { params.capture_id    = std::stoi(argv[++i]); }
    else if (arg == "-mt"   || arg == "--max-tokens")    { params.max_tokens    = std::stoi(argv[++i]); }
    else if (arg == "-ac"   || arg == "--audio-ctx")     { params.audio_ctx     = std::stoi(argv[++i]); }
    else if (arg == "-tr"   || arg == "--translate")     { params.translate     = true                ; }
    else if (arg == "-nf"   || arg == "--no-fallback")   { params.no_fallback   = true                ; }
    else if (arg == "-ps"   || arg == "--print-special") { params.print_special = true                ; }
    else if (arg == "-pt"   || arg == "--print-time")    { params.print_time    = true                ; }
    else if (arg == "-kc"   || arg == "--keep-context")  { params.keep_context  = true                ; }
    else if (arg == "-l"    || arg == "--language")      { params.language      = argv[++i]           ; }
    else if (arg == "-m"    || arg == "--model")         { params.model         = argv[++i]           ; }
    else if (arg == "-f"    || arg == "--file")          { params.logfile       = argv[++i]           ; }
    else if (arg == "-tdrz" || arg == "--tinydiarize")   { params.tinydiarize   = true                ; }
    else if (arg == "-sa"   || arg == "--save-audio")    { params.save_audio    = true                ; }
    else if (arg == "-ng"   || arg == "--no-gpu")        { params.use_gpu       = false               ; }
    else if (arg == "-fa"   || arg == "--flash-attn")    { params.flash_attn    = true                ; }
    else {
        fprintf(stderr, "error: unknown argument: %s\n", arg.c_str());
        whisper_print_usage(argc, argv, params);
        exit(0);
    }
}
return true;
}

void whisper_print_usage(int /*argc*/, char ** argv, const whisper_params & params) {
fprintf(stderr, "\n"                                                                     );
fprintf(stderr, "usage: %s [options]\n", argv[0]                                         );
fprintf(stderr, "\n"                                                                     );
fprintf(stderr, "options:\n"                                                             );
fprintf(stderr, "  -h,       --help          [default] show this help message and exit\n");
fprintf(stderr, "  -t N,     --threads N     [%-7d] number of threads to use during computation\n",    params.n_threads );
fprintf(stderr, "            --step-length N [%-7d] audio step size in milliseconds\n",                params.step_ms   );
fprintf(stderr, "            --step-number N [%-7d] total number of steps before splicing\n",          params.step_no   );
fprintf(stderr, "            --line-size   N [%-7d] maxium number of letters per line\n",              params.line_sz   );
fprintf(stderr, "  -c ID,    --capture ID    [%-7d] capture device ID\n",                              params.capture_id);
fprintf(stderr, "  -mt N,    --max-tokens N  [%-7d] maximum number of tokens per audio chunk\n",       params.max_tokens);
fprintf(stderr, "  -ac N,    --audio-ctx N   [%-7d] audio context size (0 - all)\n",                   params.audio_ctx );
fprintf(stderr, "  -tr,      --translate     [%-7s] translate from source language to english\n",      params.translate ?     "true" : "false");
fprintf(stderr, "  -nf,      --no-fallback   [%-7s] do not use temperature fallback while decoding\n", params.no_fallback ?   "true" : "false");
fprintf(stderr, "  -ps,      --print-special [%-7s] print special tokens\n",                           params.print_special ? "true" : "false");
fprintf(stderr, "  -kc,      --keep-context  [%-7s] keep context between audio chunks\n",              params.keep_context ?  "true" : "false");
fprintf(stderr, "  -l LANG,  --language LANG [%-7s] spoken language\n",                                params.language.c_str()                );
fprintf(stderr, "  -m FNAME, --model FNAME   [%-7s] model path\n",                                     params.model.c_str()                   );
fprintf(stderr, "  -f FNAME, --file FNAME    [%-7s] text output file name\n",                          params.logfile.c_str()                 );
fprintf(stderr, "  -tdrz,    --tinydiarize   [%-7s] enable tinydiarize (requires a tdrz model)\n",     params.tinydiarize ?   "true" : "false");
fprintf(stderr, "  -sa,      --save-audio    [%-7s] save the recorded audio to a file\n",              params.save_audio ?    "true" : "false");
fprintf(stderr, "  -ng,      --no-gpu        [%-7s] disable GPU inference\n",                          params.use_gpu ?       "false": "true" );
fprintf(stderr, "  -fa,      --flash-attn    [%-7s] flash attention during inference\n",               params.flash_attn ?    "true" : "false");
fprintf(stderr, "\n");
}

bool cmp(std::pair<int, float>& a, std::pair<int, float>& b)
{    return a.second < b.second;}

std:: string getTimeString()
{
SYSTEMTIME  now;
GetLocalTime(&now);
char timestring[32]{};
const char* format = { "[%04d-%02d-%02d %02d:%02d:%02d]" };
sprintf_s(timestring, 32, format, now.wYear, now.wMonth, now.wDay, now.wHour, now.wMinute, now.wSecond);
return std::string(timestring);
}
int main(int argc, char ** argv) {
HANDLE hConsole = GetStdHandle(STD_OUTPUT_HANDLE);
//SetConsoleMode(hConsole, ENABLE_VIRTUAL_TERMINAL_PROCESSING);
whisper_params params;
if (whisper_params_parse(argc, argv, params) == false) {
    return 1;
}
const int n_samples_step = params.step_ms *WHISPER_SAMPLE_RATE_MS;
const int n_samples_loop = n_samples_step * params.step_no;
params.max_tokens        = 0;
const bool log           = params.logfile.length() > 0;
// init audio
audio_async audio(300000);
if (!audio.init(params.capture_id, WHISPER_SAMPLE_RATE)) {
    fprintf(stderr, "%s: audio.init() failed!\n", __func__);
    return 1;
}
audio.resume();
// whisper init
if (params.language != "auto" && whisper_lang_id(params.language.c_str()) == -1){
    fprintf(stderr, "error: unknown language '%s'\n", params.language.c_str());
    whisper_print_usage(argc, argv, params);
    exit(0);
}
struct whisper_context_params cparams = whisper_context_default_params();
cparams.use_gpu              = params.use_gpu;
cparams.flash_attn           = params.flash_attn;
struct whisper_context * ctx = whisper_init_from_file_with_params(params.model.c_str(), cparams);
std::vector<float> audio_this    (n_samples_loop, 0.0f);
std::vector<float> audio_prev;
std::vector<float> audio_next    (n_samples_loop, 0.0f);
std::vector<whisper_token> prompt_tokens;

// print some info about the processing
{
    fprintf(stderr, "\n");
    if (!whisper_is_multilingual(ctx)) {
        if (params.language != "en" || params.translate) {
            params.language = "en";
            params.translate = false;
            fprintf(stderr, "%s: WARNING: model is not multilingual, ignoring language and translation options\n", __func__);
        }
    }
    fprintf(stderr, "%s: processing %d samples (step = %.1f sec / len = %.1f sec ), %d threads, lang = %s, task = %s, timestamps = %s ...\n",
            __func__,
            n_samples_step,
            float(n_samples_step)/WHISPER_SAMPLE_RATE,
            float(n_samples_loop)/WHISPER_SAMPLE_RATE,        
            params.n_threads,
            params.language.c_str(),
            params.translate ? "translate" : "transcribe",
            params.print_time ? "Y" : "N");

   fprintf(stderr, "%s: line_size = %d, keep_context = %d\n", __func__, params.line_sz, params.keep_context);
   fprintf(stderr, "\n");
}

std::ofstream fout;
if (log) {
    fout.open(params.logfile);
    if (!fout.is_open()) {
        fprintf(stderr, "%s: failed to open output file '%s'!\n", __func__, params.logfile.c_str());
        return 1;
    }
}
wav_writer wavWriter;

if (params.save_audio) {
    time_t now = time(0);
    char buffer[80];
    strftime(buffer, sizeof(buffer), "%Y%m%d%H%M%S", localtime(&now));
    std::string filename = std::string(buffer) + ".wav";
    wavWriter.open(filename, WHISPER_SAMPLE_RATE, 16, 1);
}
printf("[Start Speaking]\n");
fflush(stdout);

whisper_full_params wparams = whisper_full_default_params(WHISPER_SAMPLING_GREEDY);
wparams.print_progress   = false;
wparams.print_special    = params.print_special;
wparams.print_realtime   = false;
wparams.print_timestamps = params.print_time;
wparams.translate        = params.translate;
//wparams.single_segment   = !use_vad;
wparams.single_segment   = true;
wparams.max_tokens       = params.max_tokens;
wparams.language         = params.language.c_str();
wparams.n_threads        = params.n_threads;
wparams.audio_ctx        = params.audio_ctx;
wparams.tdrz_enable      = params.tinydiarize; // [TDRZ]
wparams.temperature_inc  = params.no_fallback ?   0.0f : wparams.temperature_inc;
wparams.prompt_tokens    = params.keep_context ?  prompt_tokens.data():nullptr ;
wparams.prompt_n_tokens  = params.keep_context ?  prompt_tokens.size():      0 ;

std::chrono::steady_clock::time_point      t_now   = std::chrono::high_resolution_clock::now();
std::chrono::steady_clock::time_point      t_start = t_now;
std::chrono::steady_clock::time_point      t_last  = t_now;
std::chrono::steady_clock::rep             t_diff;
std::string str_this = getTimeString();
std::string str_prev;
std::string str_next;
std::string str_mark;
std::string str_time_this;
std::string str_time_prev;
int int_trim  = 0;
int int_shif  = 0;
int  n_iter   = 1;
int  n_step   = params.step_no;
int  n_line   = 0;
int  n_nvad   = 3;
bool running  = true;
bool pausing  = false;
bool splicing = false;
int   color[] = { 7,8,15 };

// main audio loop
while (running) {
    std::this_thread::sleep_for(std::chrono::milliseconds(100));
    if (params.save_audio) {
        wavWriter.write(audio_next.data(), audio_next.size());
    }
    running = sdl_poll_events();
    short SpaceKeyPressed = GetAsyncKeyState(VK_SPACE) & 0x8001;
    short EscapKeyPressed = GetAsyncKeyState(VK_ESCAPE) & 0x8001;
    if       (SpaceKeyPressed &&  pausing) {
        printf("\n");
        n_line++;
        SetConsoleTextAttribute(hConsole, color[n_line % 3]);
        n_iter = 1;
        splicing = false;
        pausing  = false;
    } else if(SpaceKeyPressed && !pausing) {
        __noop();
    } else if(EscapKeyPressed &&  pausing) {
        __noop();
    } else if(EscapKeyPressed && !pausing) {
        printf("\n");
        n_line++;
        SetConsoleTextAttribute(hConsole, color[n_line % 3]);
        pausing = true;
    }    
   
    if (!running) {
        printf("[Terminating]\n");
        break;
    }
    if (pausing) {
        printf("\r        \r");
        printf("[Paused]");
        continue;
    }


    while (true) {  
        t_now = std::chrono::high_resolution_clock::now();
        t_diff = std::chrono::duration_cast<std::chrono::milliseconds>(t_now - t_last).count();
        if (t_diff >= params.step_ms) {
            audio.get(t_diff, audio_next);
            t_now = std::chrono::high_resolution_clock::now();
            t_last = t_now;
            break;
        } else {
            continue;
        }
    }
    const int n_samples_step = audio_next.size();
    const int n_samples_prev = audio_prev.size(); 
    audio_this.resize( n_samples_prev + n_samples_step);
    memcpy(audio_this.data() + 0             , audio_prev.data(), n_samples_prev * sizeof(float)); //Write prev audio into this 
    memcpy(audio_this.data() + n_samples_prev, audio_next.data(), n_samples_step * sizeof(float)); //Write next audio into this
    audio_prev = audio_this;
   
    if (whisper_full(ctx, wparams, audio_this.data(), audio_this.size()) != 0) {
        fprintf(stderr, "%s: failed to process audio\n", argv[0]);
        return 6;
    }

    str_next = std::string(whisper_full_get_segment_text(ctx, 0));

    if (str_next == " Thank you." || str_next == " .") 
        str_next = "[Silence]";


    if (splicing && str_next != "[Silence]") {
        float temp;
        float score;
        int prev_end = str_prev.length();
        int next_end = str_next.length();
        int limit = std::min(prev_end,next_end);
        std::vector< std::pair<int, float>> shift_score;
        for (int i = 1; i <= 40 && i <= limit; i++) {
            int_shif = i;
            temp = 0;
            score = 0;
            for (int j = 0; j < i; j++) {
                if (str_prev[prev_end - i + j] == str_next[j]) {
                    temp++;
                    score = std::max(temp, score);
                }
                else {
                    temp = 0;
                }
            }
            shift_score.emplace_back(int_shif, score / int_shif + score / 40);
        }
        std::sort(shift_score.begin(), shift_score.end(), cmp);
        int_shif = shift_score.back().first;
        score    = shift_score.back().second;
        if (score <= 0.6) { int_shif = 0; }
        int pos1 = str_next.find_last_of (" ", int_shif / 2);
        int pos2 = str_next.find_first_of(" ", int_shif / 2);
        if      (pos1 == 0 && pos2 == -1) {
            int_trim = pos1;
        }
        else if (pos1 == 0 && pos2 != -1) {
            int_trim = pos2;
        }
        else if (pos1 != 0 && pos2 == -1) {
            int_trim = pos1;
        }
        else if (pos1 != 0 && pos2 != -1) {
            if (pos2 + pos1 - int_shif >= 0) { int_trim = pos1; }
            else                             { int_trim = pos2; }
        }

        for (int i = 0; i < int_shif - int_trim; i++) {
            printf("\b \b");
        }
        printf("\n");
        n_line++;
        SetConsoleTextAttribute(hConsole, color[n_line % 3]);
        str_this = str_next.substr(int_trim);
        str_mark = str_this.substr(0, 5);
        std::transform(str_mark.begin(), str_mark.end(), str_mark.begin(), ::tolower);
    }
    else if (splicing && str_next == "[Silence]") {
        printf("%\n");
        n_line++;
        SetConsoleTextAttribute(hConsole, color[n_line % 3]);
        int_trim = 0;
        str_mark.clear();
        str_this = str_next;
    }
    else {
        printf("\r");
        printf("%s", std::string(100, ' ').c_str());
        printf("\r");
        str_this = str_next;
        std::transform(str_next.begin(), str_next.end(), str_this.begin(), ::tolower);
        int_trim = str_this.find(str_mark);
        if (int_trim != -1) {
            str_this = str_next.substr(int_trim);
        }
        else {
            str_this = str_next;
        }
    }

    if (splicing && log) {
        fout <<  std::endl;
        fout << "str_prev             " << str_prev << std::endl;
        fout << "str_next" << str_next << std::endl;
        fout << "str_this" << str_this << std::endl;
        fout << "int_shif " << int_shif << std::endl;
        fout << "int_trim " << int_trim << std::endl;
        fout << "str_mark" << str_mark << std::endl;
        fout << str_time_prev.c_str();
        fout << str_prev.c_str();
        fout.seekp(-int_shif + int_trim, std::ofstream::cur);
        fout << std::endl;
    }

    if (params.print_time && str_this != "[Silence]") {
        printf("%s", str_time_this.c_str());
    }
    printf("%s", str_this.c_str());
    fflush(stdout);
    str_prev = str_this;
    

    ++n_iter;

    if (n_iter % n_step == 0) {
        splicing   = true;
        audio_prev = std::vector<float>(audio_this.end() - n_samples_step / 2, audio_this.end());
        str_time_prev = str_time_this;
        str_time_this = getTimeString();
    } else {            
        splicing = false;
    }      

    if (params.keep_context) { // if has context then push token back in to the prompt
        prompt_tokens.clear();
        const int n_segments = whisper_full_n_segments(ctx);
        for (int i = 0; i < n_segments; ++i) {
            const int token_count = whisper_full_n_tokens(ctx, i);
            for (int j = 0; j < token_count; ++j) {
                prompt_tokens.push_back(whisper_full_get_token_id(ctx, i, j));
            }
        }
    }        
    fflush(stdout);       
}
audio.pause();
whisper_print_timings(ctx);
whisper_free(ctx);
system("pause");
return 0;
}