mozu-audio 0.3.0

//! Live sliding-window ASR, not model-native cached streaming.
//! Feed 16 kHz mono blocks on the engine's worker thread and consume returned updates.
#[cfg(test)]
use crate::vad::VoiceDetector;
use crate::{
    Cancellation, Error, HaasEngine, Result,
    recording::SAMPLE_RATE,
    vad::{SileroVad, SpeechProcessor},
    vocabulary::{CompiledVocabulary, VocabularyRecognizer},
    window::{
        AsrConfig, AsrUpdate, RecognizerRef, Transcript, WindowConfig, WindowProcessor,
        WindowRecognizer,
    },
};
use serde::Serialize;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum SessionState {
    Active,
    Finished,
    Cancelled,
    Failed,
}
#[derive(Clone)]
enum SessionSource<'a> {
    Windowed(WindowConfig),
    Silero {
        model: &'a SileroVad,
        config: AsrConfig,
    },
    #[cfg(test)]
    Always(AsrConfig),
}
impl<'a> SessionSource<'a> {
    fn processor(self, engine: RecognizerRef<'a>) -> Result<Processor<'a>> {
        match self {
            Self::Windowed(windows) => Ok(Processor::Windowed(WindowProcessor::from_recognizer(
                engine, windows,
            )?)),
            Self::Silero { model, config } => Ok(Processor::Speech(
                SpeechProcessor::from_recognizer(engine, config, model.detector())?,
            )),
            #[cfg(test)]
            Self::Always(config) => Ok(Processor::Speech(SpeechProcessor::from_recognizer(
                engine,
                config,
                Box::new(AlwaysVoice),
            )?)),
        }
    }
}
enum Processor<'a> {
    Windowed(WindowProcessor<'a>),
    Speech(SpeechProcessor<'a>),
}
impl Processor<'_> {
    fn received(&self) -> u64 {
        match self {
            Self::Windowed(window) => window.received(),
            Self::Speech(speech) => speech.received(),
        }
    }
    fn peak_buffer(&self) -> usize {
        match self {
            Self::Windowed(window) => window.peak_buffer(),
            Self::Speech(speech) => speech.peak_buffer(),
        }
    }
    fn push(&mut self, audio: &[f32], cancellation: &Cancellation) -> Result<Vec<AsrUpdate>> {
        match self {
            Self::Windowed(window) => window.feed(audio, cancellation),
            Self::Speech(speech) => {
                let mut updates = speech.feed(audio, cancellation)?;
                if let Some(update) = speech.partial(cancellation)? {
                    updates.push(update);
                }
                Ok(updates)
            }
        }
    }
    fn finish(&mut self, cancellation: &Cancellation) -> Result<(Transcript, AsrUpdate)> {
        match self {
            Self::Windowed(window) => window.finish(cancellation),
            Self::Speech(speech) => speech.finish(cancellation),
        }
    }
    fn result(&self, transcript: &Transcript) -> Result<crate::result::AsrResult> {
        match self {
            Self::Windowed(window) => window.result(transcript),
            Self::Speech(speech) => speech.result(transcript),
        }
    }
    fn clear(&mut self) {
        match self {
            Self::Windowed(window) => window.clear(),
            Self::Speech(speech) => speech.clear(),
        }
    }
}
pub struct StreamingSession<'a> {
    engine: RecognizerRef<'a>,
    source: SessionSource<'a>,
    processor: Processor<'a>,
    state: SessionState,
    generation: u64,
}
impl<'a> StreamingSession<'a> {
    pub fn new(engine: &'a HaasEngine, vad: &'a SileroVad, config: AsrConfig) -> Result<Self> {
        Self::with_source(engine, SessionSource::Silero { model: vad, config })
    }
    /// Decode only when center plus right context is available, or on finish.
    /// This path has no VAD model, endpoint timer or periodic preview inference.
    pub fn windowed(engine: &'a HaasEngine, windows: WindowConfig) -> Result<Self> {
        Self::with_windowed_recognizer(engine, windows)
    }
    pub fn new_with_vocabulary(
        engine: &'a HaasEngine,
        vad: &'a SileroVad,
        config: AsrConfig,
        vocabulary: std::sync::Arc<CompiledVocabulary>,
    ) -> Result<Self> {
        let recognizer = RecognizerRef::Shared(std::rc::Rc::new(VocabularyRecognizer::new(
            engine, vocabulary,
        )?));
        Self::with_source_ref(recognizer, SessionSource::Silero { model: vad, config })
    }
    pub fn windowed_with_vocabulary(
        engine: &'a HaasEngine,
        windows: WindowConfig,
        vocabulary: std::sync::Arc<CompiledVocabulary>,
    ) -> Result<Self> {
        let recognizer = RecognizerRef::Shared(std::rc::Rc::new(VocabularyRecognizer::new(
            engine, vocabulary,
        )?));
        Self::with_source_ref(recognizer, SessionSource::Windowed(windows))
    }
    pub(crate) fn with_windowed_recognizer(
        engine: &'a dyn WindowRecognizer,
        windows: WindowConfig,
    ) -> Result<Self> {
        Self::with_source(engine, SessionSource::Windowed(windows))
    }
    fn with_source(engine: &'a dyn WindowRecognizer, source: SessionSource<'a>) -> Result<Self> {
        Self::with_source_ref(RecognizerRef::Borrowed(engine), source)
    }
    fn with_source_ref(engine: RecognizerRef<'a>, source: SessionSource<'a>) -> Result<Self> {
        Ok(Self {
            processor: source.clone().processor(engine.clone())?,
            engine,
            source,
            state: SessionState::Active,
            generation: 0,
        })
    }
    #[cfg(test)]
    pub(crate) fn with_recognizer(
        engine: &'a dyn WindowRecognizer,
        config: AsrConfig,
    ) -> Result<Self> {
        Self::with_source(engine, SessionSource::Always(config))
    }
    pub fn state(&self) -> SessionState {
        self.state
    }
    pub fn received_samples(&self) -> u64 {
        self.processor.received()
    }
    pub fn peak_buffered_samples(&self) -> usize {
        self.processor.peak_buffer()
    }
    fn require_active(&self) -> Result<()> {
        if self.state != SessionState::Active {
            return Err(Error::Invalid(format!(
                "ASR session is {:?}; reset before reuse",
                self.state
            )));
        }
        Ok(())
    }
    /// At most one second per call. Overload is rejected, never silently dropped.
    pub fn push_audio(
        &mut self,
        audio: &[f32],
        cancellation: &Cancellation,
    ) -> Result<Vec<AsrUpdate>> {
        self.require_active()?;
        if audio.len() > SAMPLE_RATE as usize || audio.iter().any(|v| !v.is_finite()) {
            return Err(Error::Invalid(format!(
                "Live input must be finite and <=16000 samples per call; actual={}",
                audio.len()
            )));
        }
        let result = (|| {
            cancellation.check()?;
            let mut updates = self.processor.push(audio, cancellation)?;
            for update in &mut updates {
                update.session_generation = self.generation;
            }
            Ok(updates)
        })();
        if let Err(error) = &result {
            self.state = if matches!(error, Error::Cancelled) {
                SessionState::Cancelled
            } else {
                SessionState::Failed
            };
            self.processor.clear();
        }
        result
    }
    pub fn finish(&mut self, cancellation: &Cancellation) -> Result<(Transcript, AsrUpdate)> {
        self.require_active()?;
        let mut result = self.processor.finish(cancellation);
        if let Ok((_, update)) = &mut result {
            update.session_generation = self.generation;
        }
        self.state = match &result {
            Ok(_) => SessionState::Finished,
            Err(Error::Cancelled) => SessionState::Cancelled,
            Err(_) => SessionState::Failed,
        };
        if result.is_err() {
            self.processor.clear();
        }
        result
    }
    /// Finish with shared ASR metadata. Computation includes previews and overlap;
    /// time awaiting subsequent audio is excluded. Reset starts fresh statistics.
    pub fn finish_result(
        &mut self,
        cancellation: &Cancellation,
    ) -> Result<(crate::result::AsrResult, AsrUpdate)> {
        let (transcript, update) = self.finish(cancellation)?;
        let result = self.processor.result(&transcript).and_then(|result| {
            cancellation.check()?;
            Ok(result)
        });
        if let Err(error) = &result {
            self.state = if matches!(error, Error::Cancelled) {
                SessionState::Cancelled
            } else {
                SessionState::Failed
            };
            self.processor.clear();
        }
        Ok((result?, update))
    }
    pub fn cancel(&mut self) -> Result<()> {
        self.require_active()?;
        self.processor.clear();
        self.state = SessionState::Cancelled;
        Ok(())
    }
    /// Begin a new session only after finish, failure or cancellation; retains loaded models.
    pub fn reset(&mut self) -> Result<()> {
        if self.state == SessionState::Active {
            return Err(Error::Invalid(
                "Finish or cancel the active ASR session before reset".into(),
            ));
        }
        let generation = self
            .generation
            .checked_add(1)
            .ok_or_else(|| Error::Invalid("Streaming generation overflow".into()))?;
        self.processor = self.source.clone().processor(self.engine.clone())?;
        self.generation = generation;
        self.state = SessionState::Active;
        Ok(())
    }
}

#[cfg(test)]
struct AlwaysVoice;
#[cfg(test)]
impl VoiceDetector for AlwaysVoice {
    fn probability(&mut self, _audio: &[f32]) -> Result<f32> {
        Ok(1.0)
    }
}
#[cfg(test)]
mod tests {
    use super::*;
    use crate::window::test_support::{FULL_WINDOW, Markers, SHORT_WINDOW};
    fn feed(
        session: &mut StreamingSession<'_>,
        audio: &[f32],
        cancellation: &Cancellation,
    ) -> Vec<AsrUpdate> {
        let mut updates = Vec::new();
        for block in audio.chunks(16000) {
            updates.extend(session.push_audio(block, cancellation).unwrap());
        }
        updates
    }
    #[test]
    fn windowed_waits_for_full_context_without_periodic_or_silence_flushes() {
        for limits in [FULL_WINDOW, SHORT_WINDOW] {
            let recognizer = Markers::with_limits(limits);
            let windows = WindowConfig::interactive_for(limits).unwrap();
            let threshold = windows.center_samples + windows.right_samples;
            let mut session =
                StreamingSession::with_windowed_recognizer(&recognizer, windows.clone()).unwrap();
            let cancellation = Cancellation::new();
            let mut audio = vec![0.0; threshold - 1];
            audio[8000] = 7.0;
            assert!(feed(&mut session, &audio, &cancellation).is_empty());
            assert_eq!(recognizer.calls.get(), 0);
            assert_eq!(session.received_samples(), audio.len() as u64);
            let updates = session.push_audio(&[0.0], &cancellation).unwrap();
            assert_eq!(recognizer.calls.get(), 1);
            assert_eq!(updates.len(), 1);
            assert_eq!(updates[0].sequence, 1);
            let span = updates[0].recognition_span.as_ref().unwrap();
            assert_eq!(span.window, 0..threshold as u64);
            assert_eq!(span.center, 0..windows.center_samples as u64);
            let (transcript, terminal) = session.finish(&cancellation).unwrap();
            assert_eq!(transcript.text, "7");
            assert_eq!(terminal.sequence, 2);
            assert_eq!(terminal.kind, crate::window::UpdateKind::Final);
        }
    }
    #[test]
    fn windowed_short_tail_and_reset_keep_only_the_new_job_audio() {
        let recognizer = Markers::new();
        let windows = WindowConfig::interactive_for(FULL_WINDOW).unwrap();
        let mut session = StreamingSession::with_windowed_recognizer(&recognizer, windows).unwrap();
        let cancellation = Cancellation::new();
        for generation in 0..3 {
            let mut audio = vec![0.0001; 16001];
            audio[15000] = (generation + 7) as f32;
            assert!(feed(&mut session, &audio, &cancellation).is_empty());
            let (transcript, terminal) = session.finish(&cancellation).unwrap();
            assert_eq!(transcript.text, (generation + 7).to_string());
            assert_eq!(transcript.audio_samples, audio.len() as u64);
            assert_eq!(terminal.session_generation, generation);
            assert_eq!(terminal.sequence, 1);
            assert_eq!(terminal.revision.previous_token_count, 0);
            assert_eq!(
                terminal.revision.stable_token_count,
                transcript.tokens.len()
            );
            assert!(session.push_audio(&[0.0], &cancellation).is_err());
            session.reset().unwrap();
        }
        session.cancel().unwrap();
        assert_eq!(session.state(), SessionState::Cancelled);
        session.reset().unwrap();
        assert!(session.finish(&cancellation).unwrap().0.tokens.is_empty());
    }
    #[test]
    fn finish_cancel_reset_and_short_tail_do_not_leak_session_state() {
        for limits in [FULL_WINDOW, SHORT_WINDOW] {
            let marker = Markers::with_limits(limits);
            let mut s = StreamingSession::with_recognizer(
                &marker,
                AsrConfig::interactive_for(limits).unwrap(),
            )
            .unwrap();
            let c = Cancellation::new();
            let mut audio = vec![0.0; 16000];
            audio[4000] = 3.0;
            s.push_audio(&audio, &c).unwrap();
            assert_eq!(s.finish(&c).unwrap().0.text, "3");
            assert_eq!(s.state(), SessionState::Finished);
            assert!(s.push_audio(&audio, &c).is_err());
            s.reset().unwrap();
            s.push_audio(&audio, &c).unwrap();
            s.cancel().unwrap();
            assert_eq!(s.state(), SessionState::Cancelled);
            s.reset().unwrap();
            audio[4000] = 9.0;
            s.push_audio(&audio, &c).unwrap();
            assert_eq!(s.finish(&c).unwrap().0.text, "9");
            s.reset().unwrap();
            let cancelled = Cancellation::new();
            cancelled.cancel();
            assert!(matches!(
                s.push_audio(&audio, &cancelled),
                Err(Error::Cancelled)
            ));
            assert_eq!(s.state(), SessionState::Cancelled);
            s.reset().unwrap();
            s.push_audio(&[0.0; 10], &c).unwrap();
            assert_eq!(s.finish(&c).unwrap().0.audio_samples, 10);
            assert_eq!(s.state(), SessionState::Finished);
            s.reset().unwrap();
            assert_eq!(s.finish(&c).unwrap().0.audio_samples, 0);
            assert!(marker.largest_window.get() <= limits.maximum_samples());
        }
    }
    #[test]
    fn reset_begins_a_new_generation_with_an_empty_revision_cursor() {
        let recognizer = Markers::new();
        let mut session =
            StreamingSession::with_recognizer(&recognizer, AsrConfig::interactive()).unwrap();
        let cancellation = Cancellation::new();
        let mut audio = vec![0.0001; 16000];
        audio[8000] = 7.0;
        for generation in 0..3 {
            let updates = session.push_audio(&audio, &cancellation).unwrap();
            for update in &updates {
                assert_eq!(update.session_generation, generation);
            }
            let (_, final_update) = session.finish(&cancellation).unwrap();
            assert_eq!(final_update.session_generation, generation);
            assert_eq!(final_update.revision.stable_token_count, 1);
            assert_eq!(final_update.revision.previous_token_count, 0);
            assert_eq!(final_update.sequence, 1);
            session.reset().unwrap();
        }
    }

    #[test]
    fn partials_preserve_confirmed_prefix_and_clock_for_long_input() {
        for limits in [FULL_WINDOW, SHORT_WINDOW] {
            let marker = Markers::with_limits(limits);
            let config = AsrConfig::interactive_for(limits).unwrap();
            let mut s = StreamingSession::with_recognizer(&marker, config.clone()).unwrap();
            let c = Cancellation::new();
            let mut confirmed = String::new();
            let mut last_sequence = 0;
            let mut partials = 0;
            for _ in 0..65 {
                let mut audio = vec![0.0001; 16000];
                audio[2000] = 7.0;
                for update in s.push_audio(&audio, &c).unwrap() {
                    assert!(update.confirmed_text.starts_with(&confirmed));
                    assert!(update.sequence > last_sequence);
                    last_sequence = update.sequence;
                    confirmed = update.confirmed_text;
                    if update.kind == crate::window::UpdateKind::Partial {
                        partials += 1;
                    }
                }
            }
            let (result, last) = s.finish(&c).unwrap();
            assert_eq!(result.tokens.len(), 65);
            assert_eq!(result.audio_samples, 65 * 16000);
            assert!(last.confirmed_text.starts_with(&confirmed));
            assert!(last.provisional_text.is_empty());
            assert!(partials > 1);
            assert!(s.peak_buffered_samples() <= config.maximum_buffered_samples() + 4096);
            assert!(marker.largest_window.get() <= limits.maximum_samples());
        }
    }
}

#[cfg(test)]
mod failure_tests {
    use super::*;
    struct Broken;
    impl WindowRecognizer for Broken {
        fn input_limits(&self) -> crate::InputLimits {
            crate::window::test_support::FULL_WINDOW
        }
        fn infer(
            &self,
            _audio: &[f32],
            _: &Cancellation,
        ) -> Result<Vec<crate::window::TimedToken>> {
            Err(Error::Invalid("injected inference error".into()))
        }
        fn render(&self, _tokens: &[crate::window::TimedToken]) -> Result<String> {
            Ok(String::new())
        }
    }
    #[test]
    fn inference_failure_closes_session_and_reset_is_explicit() {
        let mut session =
            StreamingSession::with_recognizer(&Broken, AsrConfig::interactive()).unwrap();
        let c = Cancellation::new();
        let mut audio = vec![0.0001; 16000];
        audio[3000] = 7.0;
        assert!(session.reset().is_err());
        session.push_audio(&audio, &c).unwrap();
        assert!(session.finish(&c).is_err());
        assert_eq!(session.state(), SessionState::Failed);
        assert!(session.push_audio(&audio, &c).is_err());
        session.reset().unwrap();
        assert_eq!(session.state(), SessionState::Active);
    }
    #[test]
    fn rich_finish_profiles_each_generation_and_preserves_legacy_transcripts() {
        let recognizer = crate::window::test_support::Markers::new();
        let mut session = StreamingSession::with_windowed_recognizer(
            &recognizer,
            WindowConfig::interactive_for(crate::window::test_support::FULL_WINDOW).unwrap(),
        )
        .unwrap();
        let token = Cancellation::new();
        let mut input = vec![0.0001; 16000];
        input[8000] = 7.;
        session.push_audio(&input, &token).unwrap();
        let (result, update) = session.finish_result(&token).unwrap();
        assert_eq!(result.text, "7");
        assert_eq!(result.tokens[0].piece.text, "7");
        assert_eq!(result.metadata.audio_samples, 16000);
        assert_eq!(result.metadata.inference.calls, 1);
        assert_eq!(result.metadata.inference.audio_samples, 16000);
        assert!(result.metadata.inference.stages.is_none());
        assert_eq!(update.kind, crate::window::UpdateKind::Final);
        session.reset().unwrap();
        let (empty, update) = session.finish_result(&token).unwrap();
        assert_eq!(empty.metadata.inference.calls, 0);
        assert_eq!(empty.metadata.audio_samples, 0);
        assert_eq!(empty.metadata.mean_token_confidence, None);
        assert_eq!(empty.metadata.rtfx(), None);
        assert_eq!(update.session_generation, 1);
    }
}