Add EOU detection and streaming diarization

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),
         })
     }