diff options
Diffstat (limited to 'makima/src/server/state.rs')
| -rw-r--r-- | makima/src/server/state.rs | 12 |
1 files changed, 9 insertions, 3 deletions
diff --git a/makima/src/server/state.rs b/makima/src/server/state.rs index c38359d..31e1518 100644 --- a/makima/src/server/state.rs +++ b/makima/src/server/state.rs @@ -3,14 +3,16 @@ use std::sync::Arc; use tokio::sync::Mutex; -use crate::listen::{DiarizationConfig, ParakeetTDT, Sortformer}; +use crate::listen::{DiarizationConfig, ParakeetEOU, ParakeetTDT, Sortformer}; /// Shared application state containing ML models. /// /// Models are wrapped in `Mutex` for thread-safe mutable access during inference. pub struct AppState { - /// Speech-to-text model (Parakeet) + /// Speech-to-text model (Parakeet TDT) pub parakeet: Mutex<ParakeetTDT>, + /// End-of-Utterance detection model for streaming + pub parakeet_eou: Mutex<ParakeetEOU>, /// Speaker diarization model (Sortformer) pub sortformer: Mutex<Sortformer>, } @@ -19,13 +21,16 @@ impl AppState { /// Load all ML models from the specified directories. /// /// # Arguments - /// * `parakeet_model_dir` - Path to the Parakeet STT model directory + /// * `parakeet_model_dir` - Path to the Parakeet TDT model directory + /// * `parakeet_eou_dir` - Path to the Parakeet EOU model directory /// * `sortformer_model_path` - Path to the Sortformer diarization model file pub fn new( parakeet_model_dir: &str, + parakeet_eou_dir: &str, sortformer_model_path: &str, ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> { let parakeet = ParakeetTDT::from_pretrained(parakeet_model_dir, None)?; + let parakeet_eou = ParakeetEOU::from_pretrained(parakeet_eou_dir, None)?; let sortformer = Sortformer::with_config( sortformer_model_path, None, @@ -34,6 +39,7 @@ impl AppState { Ok(Self { parakeet: Mutex::new(parakeet), + parakeet_eou: Mutex::new(parakeet_eou), sortformer: Mutex::new(sortformer), }) } |
