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

//! Reconocimiento de voz sobre Canary (ONNX).
//!
//! El hecho que da forma a este crate: **decodificar una ventana cuesta más
//! que grabarla**. En esta máquina, 6 s de ventana tardan cerca de 800 ms en
//! decodificarse y sólo avanzan 400 ms de audio. Cualquier diseño que procese
//! todas las ventanas en orden se va quedando atrás del hablante sin límite.
//!
//! La salida es tirar trabajo: el hilo de decodificación vacía su cola entera,
//! se queda sólo con la ventana más reciente y descarta el resto. La latencia
//! queda acotada por una decodificación en lugar de crecer sin freno, y lo que
//! se pierde son transcripciones provisionales que iban a ser sobreescritas.

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};

/// Trabajo que llega al decodificador.
#[derive(Debug)]
pub enum AsrJob {
    /// Audio nuevo para la ventana deslizante.
    Window {
        turn: TurnId,
        samples: Vec<f32>,
        at: Instant,
    },
    /// Intervención cerrada, para transcribir entera.
    Utterance {
        turn: TurnId,
        samples: Vec<f32>,
        at: Instant,
    },
    /// La intervención no tenía nada: reinicia el estado de la ventana.
    Reset { turn: TurnId },
}

/// Lo que el decodificador devuelve.
#[derive(Debug, Clone)]
pub enum AsrResult {
    Partial {
        turn: TurnId,
        committed: String,
        volatile: String,
        /// Ventanas descartadas por obsoletas antes de esta.
        dropped: usize,
        decode: Duration,
    },
    Final {
        turn: TurnId,
        text: String,
        audio_secs: f32,
        decode: Duration,
        /// Instante en que el usuario dejó de hablar. Es el origen del que se
        /// mide la latencia percibida, y no puede tomarse aquí: para cuando la
        /// transcripción está lista ya ha pasado casi un segundo.
        spoken_at: Instant,
    },
    Empty {
        turn: TurnId,
    },
    Error {
        turn: TurnId,
        message: String,
    },
}

/// Motor cargado y listo para decodificar.
pub struct Recognizer {
    model: Canary,
    config: AsrConfig,
}

impl Recognizer {
    /// Carga el modelo desde `config.model_dir`.
    pub fn load(config: &AsrConfig) -> Result<Self> {
        if !config.model_dir.is_dir() {
            return Err(Error::Asr(format!(
                "no existe la carpeta del modelo: {}. Ejecuta 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!("no se pudo cargar Canary: {e}")))?;

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

    /// Sesión suelta para transcribir de una vez (pruebas y comprobaciones).
    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()))
    }

    /// Bucle del decodificador. Se ejecuta en su propio hilo hasta que se
    /// cierre `jobs`.
    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!("no se pudo abrir el flujo: {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) {
            // Todo el audio que se quedó detrás de un cierre de intervención
            // pertenece a un turno que ya se va a transcribir entero: gastar
            // una ventana en él sería trabajo condenado a sobreescribirse.
            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;
            }
            // Sólo sobrevive la ventana más nueva; lo anterior ya no describe
            // lo que se está diciendo ahora.
            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",
                turno = pending_turn.0,
                ms = decode.as_millis(),
                retraso_ms = newest.elapsed().as_millis(),
                descartadas = 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",
                turno = turn.0,
                ms = decode.as_millis(),
                audio_s = audio_secs,
                rtf = decode.as_secs_f32() / audio_secs.max(0.001),
                retraso_ms = at.elapsed().as_millis(),
                "transcripción final"
            );
            if text.is_empty() {
                AsrResult::Empty { turn }
            } else <