path: root/crates/asist-asr/src/lib.rs

//! Speech recognition on top of Canary (ONNX).
//!
//! The fact that shapes this crate: **decoding a window costs more than
//! recording it**. On this machine a 6 s window takes about 800 ms to decode
//! and only advances 400 ms of audio. Any design that processes every window
//! in order keeps falling behind the speaker without limit.
//!
//! The way out is dropping work: the decoding thread drains its whole queue,
//! keeps only the most recent window and discards the rest. Latency stays
//! bounded by one decode instead of growing unchecked, and what is lost are
//! partial transcriptions that were going to be overwritten anyway.

use std::time::{Duration, Instant};

use crossbeam_channel::{Receiver, Sender};

use asist_core::config::AsrConfig;
use asist_core::error::{Error, Result};
use asist_core::event::TurnId;

pub use canary_rs::{Canary, CanarySession, ExecutionConfig, ExecutionProvider, StreamConfig};

/// Work arriving at the decoder.
#[derive(Debug)]
pub enum AsrJob {
    /// New audio for the sliding window.
    Window {
        turn: TurnId,
        samples: Vec<f32>,
        at: Instant,
    },
    /// Closed utterance, to be transcribed in full.
    Utterance {
        turn: TurnId,
        samples: Vec<f32>,
        at: Instant,
    },
    /// The utterance was empty: resets the window state.
    Reset { turn: TurnId },
}

/// What the decoder returns.
#[derive(Debug, Clone)]
pub enum AsrResult {
    Partial {
        turn: TurnId,
        committed: String,
        volatile: String,
        /// Windows discarded as stale before this one.
        dropped: usize,
        decode: Duration,
    },
    Final {
        turn: TurnId,
        text: String,
        audio_secs: f32,
        decode: Duration,
        /// Instant the user stopped talking. It is the origin perceived latency
        /// is measured from, and it cannot be taken here: by the time the
        /// transcription is ready almost a second has passed.
        spoken_at: Instant,
    },
    Empty {
        turn: TurnId,
    },
    Error {
        turn: TurnId,
        message: String,
    },
}

/// Engine loaded and ready to decode.
pub struct Recognizer {
    model: Canary,
    config: AsrConfig,
}

impl Recognizer {
    /// Loads the model from `config.model_dir`.
    pub fn load(config: &AsrConfig) -> Result<Self> {
        if !config.model_dir.is_dir() {
            return Err(Error::Asr(format!(
                "the model directory does not exist: {}. Run scripts/bootstrap.sh",
                config.model_dir.display()
            )));
        }
        let started = Instant::now();
        let model = Canary::from_pretrained(
            config.model_dir.to_string_lossy().as_ref(),
            Some(execution_config(config)),
        )
        .map_err(|e| Error::Asr(format!("could not load Canary: {e}")))?;

        tracing::info!(
            target: "asr",
            dir = %config.model_dir.display(),
            provider = %config.execution_provider,
            ms = started.elapsed().as_millis(),
            "modelo cargado"
        );
        Ok(Self {
            model,
            config: config.clone(),
        })
    }

    /// Standalone session to transcribe in one go (tests and checks).
    pub fn transcribe(&self, samples: &[f32], sample_rate: u32) -> Result<String> {
        let mut session = self.model.session();
        session
            .transcribe_samples(
                samples,
                sample_rate as usize,
                1,
                &self.config.source_lang,
                &self.config.target_lang,
            )
            .map(|r| normalize(&r.text))
            .map_err(|e| Error::Asr(e.to_string()))
    }

    /// Decoder loop. Runs on its own thread until `jobs` is closed.
    pub fn run(self, jobs: Receiver<AsrJob>, out: Sender<AsrResult>) {
        let mut stream = match self.model.stream(
            self.config.source_lang.clone(),
            self.config.target_lang.clone(),
            stream_config(&self.config),
        ) {
            Ok(stream) => stream,
            Err(err) => {
                let _ = out.send(AsrResult::Error {
                    turn: TurnId::default(),
                    message: format!("could not open the stream: {err}"),
                });
                return;
            }
        };
        let mut session = self.model.session();
        let mut committed = String::new();
        let step_samples = (self.config.step * 16_000.0).max(1.0) as usize;

        while let Some(batch) = drain(&jobs) {
            // All the audio left behind an utterance close belongs to a turn that
            // is about to be transcribed in full: spending a window on it would
            // be work doomed to be overwritten.
            let mut pending: Vec<f32> = Vec::new();
            let mut pending_turn = TurnId::default();
            let mut newest = Instant::now();
            let mut dropped = 0usize;

            for job in batch {
                match job {
                    AsrJob::Window { turn, samples, at } => {
                        pending_turn = turn;
                        pending.extend_from_slice(&samples);
                        newest = at;
                    }
                    AsrJob::Reset { turn } => {
                        dropped += pending.len() / step_samples;
                        pending.clear();
                        stream.reset();
                        committed.clear();
                        if out.send(AsrResult::Empty { turn }).is_err() {
                            return;
                        }
                    }
                    AsrJob::Utterance { turn, samples, at } => {
                        dropped += pending.len() / step_samples;
                        pending.clear();
                        stream.reset();
                        committed.clear();
                        let result = decode_final(&mut session, &self.config, turn, &samples, at);
                        if out.send(result).is_err() {
                            return;
                        }
                    }
                }
            }

            if pending.is_empty() || !self.config.partials {
                continue;
            }
            // Only the newest window survives; the older ones no longer describe
            // what is being said now.
            dropped += (pending.len() / step_samples).saturating_sub(1);

            let started = Instant::now();
            let chunks = match stream.push_samples(&pending, 16_000, 1) {
                Ok(chunks) => chunks,
                Err(err) => {
                    let _ = out.send(AsrResult::Error {
                        turn: pending_turn,
                        message: err.to_string(),
                    });
                    continue;
                }
            };
            let Some(chunk) = chunks.last() else { continue };
            if is_degenerate(&chunk.result.text) {
                continue;
            }

            append_delta(&mut committed, chunk.delta_text.trim());
            let volatile = volatile_tail(&committed, chunk.result.text.trim());
            let decode = started.elapsed();
            tracing::debug!(
                target: "asr",
                turn = pending_turn.0,
                ms = decode.as_millis(),
                delay_ms = newest.elapsed().as_millis(),
                dropped = dropped,
                "ventana"
            );
            if out
                .send(AsrResult::Partial {
                    turn: pending_turn,
                    committed: committed.clone(),
                    volatile,
                    dropped,
                    decode,
                })
                .is_err()
            {
                return;
            }
        }
    }
}

fn decode_final(
    session: &mut CanarySession,
    config: &AsrConfig,
    turn: TurnId,
    samples: &[f32],
    at: Instant,
) -> AsrResult {
    let started = Instant::now();
    let audio_secs = samples.len() as f32 / 16_000.0;
    match session.transcribe_samples(samples, 16_000, 1, &config.source_lang, &config.target_lang) {
        Ok(result) => {
            let text = normalize(&result.text);
            let decode = started.elapsed();
            tracing::debug!(
                target: "asr",
                turn = turn.0,
                ms = decode.as_millis(),
                audio_s = audio_secs,
                rtf = decode.as_secs_f32() / audio_secs.max(0.001),
                delay_ms = at.elapsed().as_millis(),
                "final transcription"
            );
            if text.is_empty() {
                AsrResult::Empty { turn }
            } else {
                AsrResult::Final <