Files
searchi/tools/speech-runtime/worker.cpp
T

184 lines
12 KiB
C++
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
// Isolated pinned whisper.cpp C API. Binary mono PCM input; strict UTF-8 JSON.
#include "whisper.h"
#include "ggml-backend.h"
#include <algorithm>
#include <chrono>
#include <cmath>
#include <cfloat>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <iostream>
#include <sstream>
#include <stdexcept>
#include <string>
#include <vector>
#ifdef _WIN32
#ifndef NOMINMAX
#define NOMINMAX
#endif
#include <windows.h>
#include <io.h>
#include <fcntl.h>
#endif
using Clock=std::chrono::steady_clock;
constexpr size_t MAX_SAMPLES=480000, MAX_OUTPUT=4*1024*1024;
constexpr const char* POLICY="live-window-vad-dtw-v1;silero-v6.2.0;base;30s;vad-reset;threshold0.5;minspeech250ms;minsilence100ms;maxspeechFLT_MAX;pad30ms;vad-overlap0.1s;source-sample-intersections;beam5;bestof5;temp0;tempinc0.2;entropy2.4;logprob-1;nospeech0.6;suppressblank;no-suppress-nst;auto-language;no-context;dtw-base;no-flash;contiguous-intervals;first-last-token-points;exact-utf8";
static double elapsed(Clock::time_point start){return std::chrono::duration<double,std::milli>(Clock::now()-start).count();}
static std::string backend="unknown";
static void logger(ggml_log_level level,const char * text,void*) {
// Native logs can contain recognized bytes. Never forward them in production.
const std::string value(text);
const auto using_at=value.find(": using "), end=value.find(" backend");
if(using_at!=std::string::npos && end!=std::string::npos && end>using_at+8) backend=value.substr(using_at+8,end-using_at-8);
if(value.find("no GPU found")!=std::string::npos)backend="cpu";
(void)level;
}
// Reject malformed sequences, overlong encodings, surrogates and > U+10FFFF.
static bool utf8(const std::string & text) {
size_t i=0;
while(i<text.size()) {
const uint8_t c=static_cast<uint8_t>(text[i++]);
if(c<128)continue;
unsigned n;uint32_t point,minimum;
if(c>=0xc2 && c<=0xdf){n=1;point=c&31;minimum=128;}
else if(c>=0xe0 && c<=0xef){n=2;point=c&15;minimum=2048;}
else if(c>=0xf0 && c<=0xf4){n=3;point=c&7;minimum=65536;}
else return false;
if(i+n>text.size())return false;
for(unsigned k=0;k<n;++k){const uint8_t next=static_cast<uint8_t>(text[i++]);if((next&0xc0)!=0x80)return false;point=(point<<6)|(next&63);}
if(point<minimum || point>0x10ffff || (point>=0xd800 && point<=0xdfff))return false;
}
return true;
}
static std::string quoted(const std::string & text) {
if(!utf8(text))throw std::runtime_error("malformed_recognition_utf8");
std::string result="\"";const char * hex="0123456789abcdef";
for(unsigned char c:text){if(c=='"'||c=='\\'){result+='\\';result+=static_cast<char>(c);}else if(c<32){result+="\\u00";result+=hex[c>>4];result+=hex[c&15];}else result+=static_cast<char>(c);}
return result+'"';
}
static void send(const std::string & json) {
if(json.empty() || json.size()>MAX_OUTPUT)throw std::runtime_error("output_limit");
uint32_t n=static_cast<uint32_t>(json.size());char prefix[4];for(int i=0;i<4;i++)prefix[i]=static_cast<char>((n>>(i*8))&255);
std::cout.write(prefix,4);std::cout.write(json.data(),json.size());std::cout.flush();
if(!std::cout)throw std::runtime_error("output_closed");
}
static bool boundary(std::string text) {
while(!text.empty() && std::string("\"')]").find(text.back())!=std::string::npos)text.pop_back();
for(const std::string suffix:{".","!","?","。","!","?"})if(text.size()>=suffix.size() && text.compare(text.size()-suffix.size(),suffix.size(),suffix)==0)return true;
return false;
}
struct Runtime {
std::string base_path,vad_path;bool auto_device;
whisper_context * base=nullptr;whisper_vad_context * vad=nullptr;
~Runtime(){if(base)whisper_free(base);if(vad)whisper_vad_free(vad);}
void load_base(bool gpu) {
if(!gpu)backend="cpu";
auto params=whisper_context_default_params();params.use_gpu=gpu;params.flash_attn=false;
params.dtw_token_timestamps=true;params.dtw_aheads_preset=WHISPER_AHEADS_BASE;
base=whisper_init_from_file_with_params(base_path.c_str(),params);
if(!base)throw std::runtime_error("base_load_failed");
}
std::string request(unsigned op,const std::vector<float>&pcm) {
double load=0,vad_time=0,asr_time=0;std::ostringstream body;
if(op==1) {
if(!vad){auto start=Clock::now();auto params=whisper_vad_default_context_params();params.n_threads=1;params.use_gpu=auto_device;
if(!auto_device)backend="cpu";
vad=whisper_vad_init_from_file_with_params(vad_path.c_str(),params);
if(!vad)throw std::runtime_error("vad_load_failed");load=elapsed(start);}
auto start=Clock::now();
if(!whisper_vad_detect_speech(vad,pcm.data(),static_cast<int>(pcm.size())))throw std::runtime_error("vad_failed");
const whisper_vad_params policy{0.5f,250,100,FLT_MAX,30,0.1f};
auto segments=whisper_vad_segments_from_probs(vad,policy);if(!segments)throw std::runtime_error("vad_intervals_failed");
body<<",\"intervals_samples\":[";size_t previous=0;bool comma=false;
const int n=whisper_vad_segments_n_segments(segments);
if(n<0||n>1000){whisper_vad_free_segments(segments);throw std::runtime_error("vad_interval_limit");}
for(int i=0;i<n;i++){
const double a=whisper_vad_segments_get_segment_t0(segments,i),b=whisper_vad_segments_get_segment_t1(segments,i);
if(!std::isfinite(a)||!std::isfinite(b)||a<0||a>=b||std::floor(a)!=a||std::floor(b)!=b){whisper_vad_free_segments(segments);throw std::runtime_error("vad_invalid_interval");}
// Native VAD tail padding intersects source samples; never invent PCM.
const size_t first=static_cast<size_t>(std::min(a*160.0,static_cast<double>(pcm.size())));
const size_t last=static_cast<size_t>(std::min(b*160.0,static_cast<double>(pcm.size())));
if(first<previous){whisper_vad_free_segments(segments);throw std::runtime_error("vad_overlapping_interval");}
previous=last;if(first==last)continue;if(comma)body<<',';comma=true;body<<'['<<first<<','<<last<<']';
}
whisper_vad_free_segments(segments);body<<']';vad_time=elapsed(start);
}else if(op==2) {
if(!base){auto start=Clock::now();load_base(auto_device);load=elapsed(start);}
auto params=whisper_full_default_params(WHISPER_SAMPLING_BEAM_SEARCH);
params.n_threads=1;params.language="auto";params.detect_language=false;
params.greedy.best_of=5;params.beam_search.beam_size=5;
params.temperature=0.0f;params.temperature_inc=0.2f;
params.entropy_thold=2.4f;params.logprob_thold=-1.0f;params.no_speech_thold=0.6f;
params.suppress_blank=true;params.suppress_nst=false;params.translate=false;
params.print_progress=false;params.print_realtime=false;params.print_timestamps=false;params.print_special=false;
params.token_timestamps=true;params.no_context=true;
auto start=Clock::now();int status=whisper_full(base,params,pcm.data(),static_cast<int>(pcm.size()));
if(status!=0)throw std::runtime_error("recognition_failed");
body<<",\"sentences\":[";bool comma=false;int64_t previous=0;
const int count=whisper_full_n_segments(base);if(count<0||count>1000)throw std::runtime_error("segment_limit");
for(int i=0;i<count;i++){
const std::string original=whisper_full_get_segment_text(base,i);
if(original.size()>4096||!utf8(original))throw std::runtime_error("malformed_recognition_utf8");
std::string exact,phrase,merged;int64_t first=-1,last=-1,merge_first=-1;
const int tokens=whisper_full_n_tokens(base,i);if(tokens<0||tokens>4096)throw std::runtime_error("token_limit");
auto append=[&](){if(phrase.find_first_not_of(" \r\n\t")!=std::string::npos){if(comma)body<<',';comma=true;body<<"{\"start_cs\":"<<first<<",\"end_cs\":"<<last<<",\"text\":"<<quoted(phrase)<<'}';}phrase.clear();first=-1;};
for(int j=0;j<tokens;j++){
auto token=whisper_full_get_token_data(base,i,j);if(token.id>=whisper_token_eot(base))continue;
if(token.t_dtw<previous || token.t_dtw<0 || static_cast<uint64_t>(token.t_dtw)*160>pcm.size())throw std::runtime_error("invalid_dtw_point");
previous=token.t_dtw;if(merge_first<0)merge_first=token.t_dtw;
merged+=whisper_token_to_str(base,token.id);
if(merged.size()>4096)throw std::runtime_error("token_text_limit");
if(!utf8(merged))continue;
exact+=merged;if(first<0)first=merge_first;last=token.t_dtw;phrase+=merged;
if(boundary(merged))append();merged.clear();merge_first=-1;
}
if(!merged.empty() || exact!=original)throw std::runtime_error("malformed_recognition_utf8");append();
}
body<<']';asr_time=elapsed(start);
}else throw std::runtime_error("unknown_operation");
std::ostringstream output;output<<"{\"ok\":true,\"operation\":"<<op<<",\"samples\":"<<pcm.size()<<",\"vad_ms\":"<<vad_time<<",\"asr_ms\":"<<asr_time<<",\"load_ms\":"<<load<<",\"backend_observed\":"<<quoted(backend)<<body.str()<<'}';return output.str();
}
};
static int run(const std::vector<std::string>&args){
if(args.size()==2 && args[1]=="--self-test"){
if(!utf8("\xe2\x80\x8d") || utf8("\xe2\x80") || utf8("\xed\xa0\x80") || utf8("\xc0\x80") || utf8("\xf4\x90\x80\x80"))return 5;
if(quoted(std::string("a\n\"\\"))!="\"a\\u000a\\\"\\\\\"")return 5;
return 0;
}
if(args.size()!=4)return 2;
const uint16_t endian=1;if(*reinterpret_cast<const uint8_t*>(&endian)!=1||sizeof(float)!=4)return 2;
#ifdef _WIN32
_setmode(_fileno(stdin),_O_BINARY);_setmode(_fileno(stdout),_O_BINARY);
SetDefaultDllDirectories(LOAD_LIBRARY_SEARCH_DEFAULT_DIRS);
#endif
if(args[3]!="auto" && args[3]!="cpu")return 2;
const char* force_cpu=std::getenv("SEARCHI_FORCE_CPU");
const bool auto_device=args[3]=="auto" && !(force_cpu && std::string(force_cpu)=="1");
whisper_log_set(logger,nullptr);
const auto directory=std::filesystem::u8path(args[0]).parent_path().u8string();
if(auto_device)ggml_backend_load_all_from_path(directory.c_str());
else ggml_backend_load_cpu_only_from_path(directory.c_str());
Runtime runtime{args[1],args[2],auto_device};
#ifdef SEARCHI_DISTRIBUTION_AUTO
const bool distribution_auto=true;
#else
const bool distribution_auto=false;
#endif
send(std::string("{\"ready\":true,\"protocol\":1,\"runtime\":\"d1be6fde11ac6e0407606b4e42fe72d34add8037\",\"distribution_auto\":")+(distribution_auto?"true":"false")+",\"policy\":"+quoted(POLICY)+'}');
while(true){char header[5];std::cin.read(header,5);if(std::cin.gcount()==0&&std::cin.eof())return 0;if(std::cin.gcount()!=5)return 3;
uint32_t count=0;for(int i=0;i<4;i++)count|=static_cast<uint32_t>(static_cast<uint8_t>(header[i+1]))<<(i*8);
if(count==0||count>MAX_SAMPLES)return 3;std::vector<float> pcm(count);
std::cin.read(reinterpret_cast<char*>(pcm.data()),count*4);if(!std::cin)return 3;
if(!std::all_of(pcm.begin(),pcm.end(),[](float v){return std::isfinite(v);} ))return 3;
try{send(runtime.request(static_cast<uint8_t>(header[0]),pcm));}
catch(const std::exception&error){send(std::string("{\"ok\":false,\"error\":")+quoted(error.what())+'}');}
}
}
#ifdef _WIN32
int wmain(int argc,wchar_t**argv){std::vector<std::string> args;for(int i=0;i<argc;i++){int n=WideCharToMultiByte(CP_UTF8,WC_ERR_INVALID_CHARS,argv[i],-1,nullptr,0,nullptr,nullptr);if(!n)return 2;std::string value(n,'\0');WideCharToMultiByte(CP_UTF8,WC_ERR_INVALID_CHARS,argv[i],-1,value.data(),n,nullptr,nullptr);value.pop_back();args.push_back(value);}return run(args);}
#else
int main(int argc,char**argv){return run(std::vector<std::string>(argv,argv+argc));}
#endif