AI models shared registry + Code cleanup + Better async handling + Some fixes, etc

This commit is contained in:
Priler 2026-02-18 21:08:48 +05:00
parent a8ff3442ff
commit 520b98143f
62 changed files with 1683 additions and 1239 deletions

View file

@ -1,47 +1,50 @@
use once_cell::sync::OnceCell;
use vosk::{DecodingState, Model, Recognizer};
use vosk::{DecodingState, Recognizer};
use std::sync::Arc;
use parking_lot::Mutex;
use std::sync::Mutex;
// use crate::config::VOSK_MODEL_PATH;
use crate::{stt::vosk_models, i18n, config};
use crate::{vosk_models, i18n, config, models};
use crate::models::vosk::VoskModel;
use crate::DB;
static MODEL: OnceCell<Model> = OnceCell::new();
// the model Arc keeps the vosk::Model alive for the recognizers
static VOSK_MODEL: OnceCell<Arc<VoskModel>> = OnceCell::new();
static WAKE_RECOGNIZER: OnceCell<Mutex<Recognizer>> = OnceCell::new();
static SPEECH_RECOGNIZER: OnceCell<Mutex<Recognizer>> = OnceCell::new();
pub fn init_vosk() -> Result<(), String> {
if MODEL.get().is_some() {
if VOSK_MODEL.get().is_some() {
return Ok(());
} // already initialized
}
let model_path = get_configured_model_path()?;
info!("Loading Vosk model from: {}", model_path.display());
let model_id = format!("vosk:{}", model_path.display());
let model = Model::new(model_path.to_str().unwrap())
.ok_or_else(|| format!("Failed to load Vosk model from: {}", model_path.display()))?;
// load through registry (shared if anything else needs the same model)
let vosk = models::vosk::load(
models::registry(),
&model_id,
model_path.to_str().unwrap(),
)?;
// language-specific wake grammar
let lang = i18n::get_language();
let wake_grammar = config::get_wake_grammar(&lang);
info!("Wake grammar for '{}': {:?}", lang, wake_grammar);
//let mut recognizer = Recognizer::new(&model, 16000.0)
// .ok_or("Failed to create Vosk recognizer")?;
let mut wake_recognizer = Recognizer::new_with_grammar(&model, 16000.0, wake_grammar)
let mut wake_recognizer = Recognizer::new_with_grammar(&vosk.model, 16000.0, wake_grammar)
.ok_or("Failed to create wake word recognizer")?;
wake_recognizer.set_max_alternatives(1); // required for confidence check later on
wake_recognizer.set_max_alternatives(1);
let mut speech_recognizer = Recognizer::new(&model, 16000.0)
let mut speech_recognizer = Recognizer::new(&vosk.model, 16000.0)
.ok_or("Failed to create speech recognizer")?;
speech_recognizer.set_max_alternatives(config::VOSK_SPEECH_RECOGNIZER_MAX_ALTERNATIVES);
speech_recognizer.set_words(config::VOSK_SPEECH_RECOGNIZER_WORDS);
speech_recognizer.set_partial_words(config::VOSK_SPEECH_PARTIAL_WORDS);
MODEL.set(model).map_err(|_| "Model already set")?;
VOSK_MODEL.set(vosk).map_err(|_| "Model already set")?;
WAKE_RECOGNIZER.set(Mutex::new(wake_recognizer)).map_err(|_| "Wake recognizer already set")?;
SPEECH_RECOGNIZER.set(Mutex::new(speech_recognizer)).map_err(|_| "Speech recognizer already set")?;
@ -50,17 +53,15 @@ pub fn init_vosk() -> Result<(), String> {
pub fn recognize_wake_word(data: &[i16]) -> Option<(String, f32)> {
let mut recognizer = WAKE_RECOGNIZER.get()?.lock().unwrap();
let mut recognizer = WAKE_RECOGNIZER.get()?.lock();
match recognizer.accept_waveform(data) {
Ok(DecodingState::Running) => {
// partials don't have confidence, skip them
None
}
Ok(DecodingState::Finalized) => {
let result = recognizer.result();
// compensate confidence issues
if let Some(alternatives) = result.multiple() {
if let Some(best) = alternatives.alternatives.first() {
if !best.text.is_empty() {
@ -77,7 +78,7 @@ pub fn recognize_wake_word(data: &[i16]) -> Option<(String, f32)> {
pub fn recognize_speech(data: &[i16]) -> Option<String> {
let mut recognizer = SPEECH_RECOGNIZER.get()?.lock().unwrap();
let mut recognizer = SPEECH_RECOGNIZER.get()?.lock();
match recognizer.accept_waveform(data) {
Ok(DecodingState::Finalized) => {
@ -92,65 +93,16 @@ pub fn recognize_speech(data: &[i16]) -> Option<String> {
pub fn reset_speech_recognizer() {
if let Some(recognizer) = SPEECH_RECOGNIZER.get() {
recognizer.lock().unwrap().reset();
recognizer.lock().reset();
}
}
pub fn reset_wake_recognizer() {
if let Some(recognizer) = WAKE_RECOGNIZER.get() {
recognizer.lock().unwrap().reset();
recognizer.lock().reset();
}
}
// pub fn recognize(data: &[i16], include_partial: bool) -> Option<String> {
// let state = RECOGNIZER
// .get()
// .unwrap()
// .lock()
// .unwrap()
// .accept_waveform(data);
// match state {
// Ok(ds) => {
// match ds {
// DecodingState::Running => {
// if include_partial {
// Some(
// RECOGNIZER
// .get()
// .unwrap()
// .lock()
// .unwrap()
// .partial_result()
// .partial
// .into(),
// )
// } else {
// None
// }
// }
// DecodingState::Finalized => {
// // Result will always be multiple because we called set_max_alternatives
// RECOGNIZER
// .get()
// .unwrap()
// .lock()
// .unwrap()
// .result()
// .multiple()
// .and_then(|m| m.alternatives.first().map(|a| a.text.to_string()))
// }
// DecodingState::Failed => None,
// }
// },
// Err(err) => {
// error!("Vosk accept waveform error.\nError details: {}", err);
// None
// }
// }
// }
fn get_configured_model_path() -> Result<std::path::PathBuf, String> {
// try to get from settings
if let Some(db) = DB.get() {
@ -167,11 +119,10 @@ fn get_configured_model_path() -> Result<std::path::PathBuf, String> {
let available = vosk_models::scan_vosk_models();
let language = i18n::get_language();
// try language match first
let lang_code = match language.as_str() {
"ru" => "ru",
"en" => "us", // vosk uses "us" not "en"
"ua" => "uk", // vosk uses "uk" not "ua"
"en" => "us",
"ua" => "uk",
other => other,
};
@ -180,7 +131,6 @@ fn get_configured_model_path() -> Result<std::path::PathBuf, String> {
return Ok(matched.path.clone());
}
// fallback to first available
if let Some(first) = available.first() {
info!("Auto-detected Vosk model (no language match): {}", first.name);
return Ok(first.path.clone());
@ -194,14 +144,3 @@ fn get_configured_model_path() -> Result<std::path::PathBuf, String> {
Err("No Vosk models found".into())
}
// pub fn stereo_to_mono(input_data: &[i16]) -> Vec<i16> {
// let mut result = Vec::with_capacity(input_data.len() / 2);
// result.extend(
// input_data
// .chunks_exact(2)
// .map(|chunk| chunk[0] / 2 + chunk[1] / 2),
// );
// result
// }