aria-kit 0.2.0

//! One transcription run: a recording the session owns, the transcription
//! over it, and the lifecycle between them. A [`Streaming`] session yields
//! words as they are committed; a `Whole` session has only its end.

use std::marker::PhantomData;

use super::{
    Language, Layer, Mode, Operation, Streaming, Transcriber, Transcript, TranscriptSegment,
    Transcription, TranscriptionError,
};
use crate::audio::RecordingReader;

#[cfg(not(target_family = "wasm"))]
use crate::audio::Stopper;
#[cfg(all(feature = "microphone", not(target_family = "wasm")))]
use crate::audio::{Microphone, RecordingMode};
#[cfg(all(feature = "silero", not(target_family = "wasm")))]
use crate::speaker::{Diarizer, Embedder, Heard, SegmentLabeler, SegmentReader, Verdict};
#[cfg(all(feature = "silero", not(target_family = "wasm")))]
use crate::voice_activity::speech_gate;

/// One transcription run over audio that arrives while the run is live.
///
/// Constructed over a recording the caller feeds, or straight from the
/// microphone where one exists. `stop` waits for input to end, drains the
/// tail, and returns the completed transcript.
pub struct Session<'a, T, M: Mode, L = ()> {
    operation: Operation<'a, T, L>,
    recording: RecordingReader,
    /// Held so the capture worker outlives the run.
    #[cfg(all(feature = "microphone", not(target_family = "wasm")))]
    _microphone: Option<Microphone>,
    #[cfg(not(target_family = "wasm"))]
    stopper: Option<Stopper>,
    mode: PhantomData<M>,
}

impl<'a, T: Transcriber, M: Mode, L: Layer> Session<'a, T, M, L> {
    /// Runs against a recording the caller already has. Whoever feeds it also
    /// ends it; `stop` only waits for that end and finishes the transcript.
    pub fn over(
        transcription: &'a Transcription<T, M, L>,
        language: &'a Language,
        recording: RecordingReader,
    ) -> Result<Self, TranscriptionError> {
        Ok(Self {
            operation: transcription.start(language)?,
            recording,
            #[cfg(all(feature = "microphone", not(target_family = "wasm")))]
            _microphone: None,
            #[cfg(not(target_family = "wasm"))]
            stopper: None,
            mode: PhantomData,
        })
    }

    /// Waits for input to end, drains whatever remains, and returns the
    /// completed transcript — or [`TranscriptionError::NoSpeech`] when
    /// nothing was said.
    pub fn stop(mut self) -> Result<Transcript, TranscriptionError> {
        #[cfg(not(target_family = "wasm"))]
        if let Some(stopper) = &self.stopper {
            stopper.stop();
        }
        self.operation.finish(&mut self.recording)
    }
}

#[cfg(all(feature = "microphone", not(target_family = "wasm")))]
impl<'a, T: Transcriber, M: Mode, L: Layer> Session<'a, T, M, L> {
    /// Opens the microphone and records one utterance until stopped. `stop`
    /// ends the recording; until then every [`Streaming`] step blocks on
    /// captured audio.
    pub fn microphone(
        transcription: &'a Transcription<T, M, L>,
        language: &'a Language,
    ) -> Result<Self, TranscriptionError> {
        let mut microphone = Microphone::open().map_err(TranscriptionError::Capture)?;
        let (recording, stopper) = microphone
            .record(RecordingMode::UntilStopped)
            .map_err(TranscriptionError::Capture)?;
        let operation = match transcription.start(language) {
            Ok(operation) => operation,
            Err(error) => {
                stopper.stop();
                return Err(error);
            }
        };
        Ok(Self {
            operation,
            recording,
            _microphone: Some(microphone),
            stopper: Some(stopper),
            mode: PhantomData,
        })
    }
}

impl<'a, T: Transcriber, L: Layer> Session<'a, T, Streaming, L> {
    /// The next committed stretch of words, or `None` once the input has
    /// ended and the tail is delivered. Committed words are never revised.
    pub fn next_finalized(&mut self) -> Result<Option<TranscriptSegment>, TranscriptionError> {
        self.operation.next_finalized(&mut self.recording)
    }

    /// The same session, with every committed segment's voice named by
    /// `voices`, behind the bundled speech gate.
    #[cfg(all(feature = "silero", not(target_family = "wasm")))]
    pub fn hearing<E: Embedder>(self, voices: Diarizer<E>) -> HeardSession<'a, T, E, L> {
        self.with_reader(SegmentReader::new(SegmentLabeler::new(
            speech_gate(),
            Some(voices),
        )))
    }

    /// The same session, with every committed segment judged by a reader
    /// the caller built.
    #[cfg(all(feature = "silero", not(target_family = "wasm")))]
    pub fn with_reader<E: Embedder>(self, reader: SegmentReader<E>) -> HeardSession<'a, T, E, L> {
        HeardSession {
            session: self,
            reader,
            kept: Vec::new(),
        }
    }
}

/// A streaming session whose segments come with a verdict on whose voice
/// they are in.
#[cfg(all(feature = "silero", not(target_family = "wasm")))]
pub struct HeardSession<'a, T, E, L = ()> {
    session: Session<'a, T, Streaming, L>,
    reader: SegmentReader<E>,
    /// What passed the reader, for the transcript of a session that keeps
    /// strangers out.
    kept: Vec<String>,
}

#[cfg(all(feature = "silero", not(target_family = "wasm")))]
impl<T: Transcriber, E: Embedder, L: Layer> HeardSession<'_, T, E, L> {
    /// Keeps out every segment in a voice that is not enrolled.
    pub fn only_own_voice(mut self) -> Self {
        self.reader.labeler_mut().set_only_own_voice(true);
        self
    }

    /// The next committed segment and whose voice it was, or `None` once
    /// the input has ended and the tail is delivered.
    pub fn next_heard(&mut self) -> Result<Option<Heard>, TranscriptionError> {
        let heard = self
            .reader
            .next(&mut self.session.operation, &mut self.session.recording)?;
        if let Some(heard) = &heard
            && matches!(heard.verdict, Verdict::Keep { .. })
        {
            self.kept.push(heard.segment.as_str().to_owned());
        }
        Ok(heard)
    }

    pub fn reader(&self) -> &SegmentReader<E> {
        &self.reader
    }

    /// The completed transcript. A session that keeps strangers out answers
    /// with what it kept, and [`TranscriptionError::NoSpeech`] when that is
    /// nothing.
    pub fn stop(self) -> Result<Transcript, TranscriptionError> {
        let transcript = self.session.stop();
        if !self.reader.labeler().only_own_voice() {
            return transcript;
        }
        // Wait for the input to end first, and only then hand back what was
        // kept.
        transcript?;
        Transcript::new(self.kept.join(" "))
    }
}

#[cfg(not(target_family = "wasm"))]
impl<T, M: Mode, L> Drop for Session<'_, T, M, L> {
    fn drop(&mut self) {
        // A session dropped without stop() must not leave the microphone
        // recording to its ceiling. Stopping an already-stopped input is a
        // no-op, so this composes with stop().
        if let Some(stopper) = &self.stopper {
            stopper.stop();
        }
    }
}

#[cfg(test)]
mod tests {
    use std::collections::VecDeque;
    use std::sync::Mutex;

    use super::*;
    use crate::audio::CANONICAL_SAMPLE_RATE;
    use crate::audio::testing::{Pace, paced, recorded};
    use crate::cancel::CancelToken;
    use crate::transcription::{Cadence, TranscriptionAudio};

    const STRIDE_MS: usize = 100;
    const STRIDE: usize = STRIDE_MS * CANONICAL_SAMPLE_RATE as usize / 1000;

    /// Answers each decode with the next entry of its script, releasing the
    /// next paced batch as it goes.
    struct Scripted {
        answers: Mutex<VecDeque<&'static str>>,
        pace: Option<Pace>,
    }

    impl Scripted {
        fn new(answers: &[&'static str]) -> Self {
            Self {
                answers: Mutex::new(answers.iter().copied().collect()),
                pace: None,
            }
        }

        fn paced(answers: &[&'static str], pace: Pace) -> Self {
            Self {
                answers: Mutex::new(answers.iter().copied().collect()),
                pace: Some(pace),
            }
        }
    }

    impl Transcriber for Scripted {
        fn validate_language(
            &self,
            _language: &Language,
            _cancel: &CancelToken,
        ) -> Result<(), TranscriptionError> {
            Ok(())
        }

        fn transcribe_with(
            &self,
            _audio: TranscriptionAudio<'_>,
            _language: &Language,
            _cancel: &CancelToken,
        ) -> Result<super::super::Hypothesis, TranscriptionError> {
            if let Some(pace) = &self.pace {
                pace.advance();
            }
            let answer = self
                .answers
                .lock()
                .unwrap_or_else(|poisoned| poisoned.into_inner())
                .pop_front()
                .expect("the script covers every decode");
            Ok(super::super::Hypothesis::new(answer))
        }
    }

    fn stride_batch() -> Vec<i16> {
        vec![8_192; STRIDE]
    }

    #[test]
    fn a_streaming_session_commits_words_as_their_agreement_holds() {
        let (_directory, recording, pace) = paced(vec![stride_batch(); 4]);
        let cadence = Cadence::from_millis(STRIDE_MS, 20_000);
        let transcription = Transcription::agreement(
            cadence,
            Scripted::paced(&["one", "one two", "one two three", "one two three"], pace),
        );
        let language = Language::Automatic;
        let mut session =
            Session::over(&transcription, &language, recording).expect("the session starts");

        let mut segments = Vec::new();
        while let Some(segment) = session.next_finalized().expect("every step succeeds") {
            segments.push(segment.into_text());
        }
        let transcript = session.stop().expect("the transcript completes");

        assert_eq!(segments, ["one", "two", "three"]);
        assert_eq!(transcript.as_str(), "one two three");
    }

    #[cfg(feature = "silero")]
    #[test]
    fn a_session_with_a_reader_says_whose_voice_each_segment_is_in() {
        use crate::speaker::{Diarizer, Embedder, SegmentLabeler, SpeakerError, Verdict};

        const OWN: [f32; 2] = [1.0, 0.0];
        const OTHER: [f32; 2] = [0.0, 1.0];

        struct Voice([f32; 2]);

        impl Embedder for Voice {
            fn embed(&mut self, _window: &[f32]) -> Result<Option<Vec<f32>>, SpeakerError> {
                Ok(Some(self.0.to_vec()))
            }

            fn dimension(&self) -> usize {
                2
            }
        }

        /// Hears "one" spanning the whole window it is given.
        struct Timed;

        impl Transcriber for Timed {
            fn validate_language(
                &self,
                _language: &Language,
                _cancel: &CancelToken,
            ) -> Result<(), TranscriptionError> {
                Ok(())
            }

            fn transcribe_with(
                &self,
                audio: TranscriptionAudio<'_>,
                _language: &Language,
                _cancel: &CancelToken,
            ) -> Result<super::super::Hypothesis, TranscriptionError> {
                let TranscriptionAudio::Samples(samples) = audio else {
                    unreachable!("the session hands over samples");
                };
                let end = samples.len() as f32 / CANONICAL_SAMPLE_RATE as f32;
                Ok(super::super::Hypothesis::with_words(
                    "one",
                    vec![super::super::TimedWord { start: 0.0, end }],
                ))
            }
        }

        // A finished recording is heard in one decode, so the one segment
        // spans its whole 3.4 s: enough to judge.
        let judged = |voice: [f32; 2]| {
            let (_directory, recording) = recorded(&vec![8_192; 34 * STRIDE]);
            let cadence = Cadence::from_millis(STRIDE_MS, 20_000);
            let transcription = Transcription::agreement(cadence, Timed);
            let language = Language::Automatic;
            let enrolled = vec![crate::speaker::enrolled_voice(
                "alex",
                crate::speaker::Role::Owner,
                OWN.to_vec(),
            )];
            let labeler = SegmentLabeler::new(None, Some(Diarizer::new(Voice(voice), enrolled)));
            let mut session = Session::over(&transcription, &language, recording)
                .expect("the session starts")
                .with_reader(SegmentReader::new(labeler))
                .only_own_voice();

            let mut heard = Vec::new();
            while let Some(segment) = session.next_heard().expect("every step succeeds") {
                let stranger = matches!(segment.verdict, Verdict::Stranger);
                heard.push((segment.segment.into_text(), stranger));
            }
            (heard, session.stop().map(Transcript::into_text))
        };

        let (heard, transcript) = judged(OTHER);
        assert_eq!(heard, [("one".to_owned(), true)]);
        assert!(matches!(transcript, Err(TranscriptionError::NoSpeech)));

        let (heard, transcript) = judged(OWN);
        assert_eq!(heard, [("one".to_owned(), false)]);
        assert_eq!(transcript.expect("the kept words complete"), "one");
    }

    #[test]
    fn a_whole_session_has_only_its_end() {
        let (_directory, recording) = recorded(&[0; 64]);
        let transcription = Transcription::complete_samples(Scripted::new(&["hello world"]));
        let language = Language::Automatic;

        let transcript = Session::over(&transcription, &language, recording)
            .expect("the session starts")
            .stop()
            .expect("the transcript completes");

        assert_eq!(transcript.as_str(), "hello world");
    }
}