Add EOU detection and streaming diarization

1 files changed, 263 insertions, 0 deletions

diff --git a/parakeet-rs/src/model_tdt.rs b/parakeet-rs/src/model_tdt.rs
new file mode 100644
index 0000000..e00ebdc
--- /dev/null
+++ b/parakeet-rs/src/model_tdt.rs
@@ -0,0 +1,263 @@
+use crate::error::{Error, Result};
+use crate::execution::ModelConfig as ExecutionConfig;
+use ndarray::{Array1, Array2, Array3};
+use ort::session::Session;
+use std::path::{Path, PathBuf};
+
+/// TDT model configs
+#[derive(Debug, Clone)]
+pub struct TDTModelConfig {
+    pub vocab_size: usize,
+}
+
+impl Default for TDTModelConfig {
+    fn default() -> Self {
+        Self {
+            vocab_size: 8193,
+        }
+    }
+}
+
+pub struct ParakeetTDTModel {
+    encoder: Session,
+    decoder_joint: Session,
+    config: TDTModelConfig,
+}
+
+impl ParakeetTDTModel {
+    /// Load TDT model from directory containing encoder and decoder_joint ONNX files
+    pub fn from_pretrained<P: AsRef<Path>>(
+        model_dir: P,
+        exec_config: ExecutionConfig,
+    ) -> Result<Self> {
+        let model_dir = model_dir.as_ref();
+
+        // Find encoder and decoder_joint files
+        let encoder_path = Self::find_encoder(model_dir)?;
+        let decoder_joint_path = Self::find_decoder_joint(model_dir)?;
+
+        let config = TDTModelConfig::default();
+
+        // Load encoder
+        let builder = Session::builder()?;
+        let builder = exec_config.apply_to_session_builder(builder)?;
+        let encoder = builder.commit_from_file(&encoder_path)?;
+
+        // Load decoder_joint
+        let builder = Session::builder()?;
+        let builder = exec_config.apply_to_session_builder(builder)?;
+        let decoder_joint = builder.commit_from_file(&decoder_joint_path)?;
+
+
+        Ok(Self {
+            encoder,
+            decoder_joint,
+            config,
+        })
+    }
+
+    fn find_encoder(dir: &Path) -> Result<PathBuf> {
+        let candidates = ["encoder-model.onnx", "encoder.onnx"];
+        for candidate in &candidates {
+            let path = dir.join(candidate);
+            if path.exists() {
+                return Ok(path);
+            }
+        }
+        Err(Error::Config(format!(
+            "No encoder model found in {}",
+            dir.display()
+        )))
+    }
+
+    fn find_decoder_joint(dir: &Path) -> Result<PathBuf> {
+        let candidates = [
+            "decoder_joint-model.onnx",
+            "decoder_joint.onnx",
+            "decoder-model.onnx",
+        ];
+        for candidate in &candidates {
+            let path = dir.join(candidate);
+            if path.exists() {
+                return Ok(path);
+            }
+        }
+        Err(Error::Config(format!(
+            "No decoder_joint model found in {}",
+            dir.display()
+        )))
+    }
+
+    /// Run greedy decoding - returns (token_ids, frame_indices, durations)
+    pub fn forward(&mut self, features: Array2<f32>) -> Result<(Vec<usize>, Vec<usize>, Vec<usize>)> {
+        // Run encoder
+        let (encoder_out, encoder_len) = self.run_encoder(&features)?;
+
+        // Run greedy decoding with decoder_joint
+        let (tokens, frame_indices, durations) = self.greedy_decode(&encoder_out, encoder_len)?;
+
+        Ok((tokens, frame_indices, durations))
+    }
+
+    fn run_encoder(&mut self, features: &Array2<f32>) -> Result<(Array3<f32>, i64)> {
+        let batch_size = 1;
+        let time_steps = features.shape()[0];
+        let feature_size = features.shape()[1];
+
+        // TDT encoder expects (batch, features, time) not (batch, time, features)
+        let input = features
+            .t()
+            .to_shape((batch_size, feature_size, time_steps))
+            .map_err(|e| Error::Model(format!("Failed to reshape encoder input: {e}")))?
+            .to_owned();
+
+        let input_length = Array1::from_vec(vec![time_steps as i64]);
+
+        let input_value = ort::value::Value::from_array(input)?;
+        let length_value = ort::value::Value::from_array(input_length)?;
+
+        let outputs = self.encoder.run(ort::inputs!(
+            "audio_signal" => input_value,
+            "length" => length_value
+        ))?;
+
+        let encoder_out = &outputs["outputs"];
+        let encoder_lens = &outputs["encoded_lengths"];
+
+        let (shape, data) = encoder_out
+            .try_extract_tensor::<f32>()
+            .map_err(|e| Error::Model(format!("Failed to extract encoder output: {e}")))?;
+
+        let (_, lens_data) = encoder_lens
+            .try_extract_tensor::<i64>()
+            .map_err(|e| Error::Model(format!("Failed to extract encoder lengths: {e}")))?;
+
+        let shape_dims = shape.as_ref();
+        if shape_dims.len() != 3 {
+            return Err(Error::Model(format!(
+                "Expected 3D encoder output, got shape: {shape_dims:?}"
+            )));
+        }
+
+        let b = shape_dims[0] as usize;
+        let t = shape_dims[1] as usize;
+        let d = shape_dims[2] as usize;
+
+        let encoder_array = Array3::from_shape_vec((b, t, d), data.to_vec())
+            .map_err(|e| Error::Model(format!("Failed to create encoder array: {e}")))?;
+
+        // TDT encoder outputs [batch, encoder_dim, time] directly
+        Ok((encoder_array, lens_data[0]))
+    }
+
+    fn greedy_decode(&mut self, encoder_out: &Array3<f32>, _encoder_len: i64) -> Result<(Vec<usize>, Vec<usize>, Vec<usize>)> {
+        // encoder_out shape: [batch, encoder_dim, time]
+        let encoder_dim = encoder_out.shape()[1];
+        let time_steps = encoder_out.shape()[2];
+        let vocab_size = self.config.vocab_size;
+        let max_tokens_per_step = 10;
+        let blank_id = vocab_size - 1;
+
+        // States: (num_layers=2, batch=1, hidden_dim=640)
+        let mut state_h = Array3::<f32>::zeros((2, 1, 640));
+        let mut state_c = Array3::<f32>::zeros((2, 1, 640));
+
+        let mut tokens = Vec::new();
+        let mut frame_indices = Vec::new();
+        let mut durations = Vec::new();
+
+        let mut t = 0;
+        let mut emitted_tokens = 0;
+        let mut last_emitted_token = blank_id as i32;
+
+        // Frame-by-frame RNN-T/TDT greedy decoding
+        while t < time_steps {
+            // Get single encoder frame: slice [0, :, t] and reshape to [1, encoder_dim, 1]
+            let frame = encoder_out.slice(ndarray::s![0, .., t]).to_owned();
+            let frame_reshaped = frame
+                .to_shape((1, encoder_dim, 1))
+                .map_err(|e| Error::Model(format!("Failed to reshape frame: {e}")))?
+                .to_owned();
+
+            // Current token for prediction network
+            let targets = Array2::from_shape_vec((1, 1), vec![last_emitted_token])
+                .map_err(|e| Error::Model(format!("Failed to create targets: {e}")))?;
+
+            // Run decoder_joint
+            let outputs = self.decoder_joint.run(ort::inputs!(
+                "encoder_outputs" => ort::value::Value::from_array(frame_reshaped)?,
+                "targets" => ort::value::Value::from_array(targets)?,
+                "target_length" => ort::value::Value::from_array(Array1::from_vec(vec![1i32]))?,
+                "input_states_1" => ort::value::Value::from_array(state_h.clone())?,
+                "input_states_2" => ort::value::Value::from_array(state_c.clone())?
+            ))?;
+
+            // Extract logits
+            let (_, logits_data) = outputs["outputs"]
+                .try_extract_tensor::<f32>()
+                .map_err(|e| Error::Model(format!("Failed to extract logits: {e}")))?;
+
+            // TDT outputs vocab_size + 5 durations (8193 + 5 = 8198)
+            let vocab_logits: Vec<f32> = logits_data.iter().take(vocab_size).copied().collect();
+            let duration_logits: Vec<f32> = logits_data.iter().skip(vocab_size).copied().collect();
+
+            let token_id = vocab_logits
+                .iter()
+                .enumerate()
+                .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
+                .map(|(idx, _)| idx)
+                .unwrap_or(blank_id);
+
+            let duration_step = if !duration_logits.is_empty() {
+                duration_logits
+                    .iter()
+                    .enumerate()
+                    .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
+                    .map(|(idx, _)| idx)
+                    .unwrap_or(0)
+            } else {
+                0
+            };
+
+            // Check if blank token
+            if token_id != blank_id {
+                // Update states when we emit a token
+                if let Ok((h_shape, h_data)) = outputs["output_states_1"].try_extract_tensor::<f32>() {
+                    let dims = h_shape.as_ref();
+                    state_h = Array3::from_shape_vec((dims[0] as usize, dims[1] as usize, dims[2] as usize), h_data.to_vec())
+                        .map_err(|e| Error::Model(format!("Failed to update state_h: {e}")))?;
+                }
+                if let Ok((c_shape, c_data)) = outputs["output_states_2"].try_extract_tensor::<f32>() {
+                    let dims = c_shape.as_ref();
+                    state_c = Array3::from_shape_vec((dims[0] as usize, dims[1] as usize, dims[2] as usize), c_data.to_vec())
+                        .map_err(|e| Error::Model(format!("Failed to update state_c: {e}")))?;
+                }
+
+                tokens.push(token_id);
+                frame_indices.push(t);
+                durations.push(duration_step);
+                last_emitted_token = token_id as i32;
+                emitted_tokens += 1;
+
+                // Don't advance yet - try to emit more tokens from the same frame
+            } else {
+                // Blank token - advance frame pointer
+                // Duration prediction applies when we finally move to next frame after emitting tokens
+                if duration_step > 0 && emitted_tokens > 0 {
+                    t += duration_step;
+                } else {
+                    t += 1;
+                }
+                emitted_tokens = 0;
+            }
+
+            // Safety check: if we've emitted too many tokens from the same frame, advance
+            if emitted_tokens >= max_tokens_per_step {
+                t += 1;
+                emitted_tokens = 0;
+            }
+        }
+
+        Ok((tokens, frame_indices, durations))
+    }
+}