euhadra 0.2.0

//! `ParaformerAdapter` — `AsrAdapter` implementation backed by an
//! ONNX-exported FunASR Paraformer-large model. Defaults match the
//! published `iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-onnx`
//! bundle.

use async_trait::async_trait;
use ndarray::{Array1, Array3};
use ort::session::Session;
use ort::value::Value;
use std::path::{Path, PathBuf};
use std::sync::Mutex;

use crate::traits::{AsrAdapter, AsrError};
use crate::types::{AudioChunk, Transcript};

use super::fbank::{Fbank, FbankOpts};
use super::frontend::{apply_cmvn, apply_lfr, load_cmvn, Cmvn};
use super::vocab::{ids_to_tokens, load_tokens_json, sentence_postprocess};

/// Configuration that mirrors the standard offline Paraformer-large
/// export. All knobs are exposed so non-default vocab sizes (e.g. the
/// 8358 hot-word variant) can be loaded without code changes.
#[derive(Debug, Clone)]
pub struct ParaformerConfig {
    pub fbank: FbankOpts,
    pub lfr_m: usize,
    pub lfr_n: usize,
    /// Offset subtracted from `valid_token_num`, matching FunASR's
    /// `predictor_bias` config. The standard checkpoint uses 1.
    pub predictor_bias: usize,
    /// Filename overrides. Defaults match the ModelScope bundle layout.
    pub model_filename: String,
    pub cmvn_filename: String,
    pub tokens_filename: String,
}

impl Default for ParaformerConfig {
    fn default() -> Self {
        Self {
            fbank: FbankOpts::paraformer_default(),
            lfr_m: 7,
            lfr_n: 6,
            predictor_bias: 1,
            model_filename: "model.onnx".to_string(),
            cmvn_filename: "am.mvn".to_string(),
            tokens_filename: "tokens.json".to_string(),
        }
    }
}

/// `AsrAdapter` for Paraformer-large via ONNX Runtime.
///
/// The adapter accumulates the full utterance, runs FBANK → LFR →
/// CMVN → encoder → argmax → token postprocess, and emits a single
/// final `AsrResult`.
pub struct ParaformerAdapter {
    session: Mutex<Session>,
    fbank: Fbank,
    cmvn: Cmvn,
    vocab: Vec<String>,
    cfg: ParaformerConfig,
    /// Cached at load time so we don't have to introspect the session
    /// on every utterance.
    input_names: (String, String),
}

impl ParaformerAdapter {
    /// Load a model bundle laid out as
    /// `<dir>/{model.onnx, am.mvn, tokens.json}`.
    pub fn load(model_dir: impl AsRef<Path>) -> Result<Self, AsrError> {
        Self::load_with_config(model_dir, ParaformerConfig::default())
    }

    pub fn load_with_config(
        model_dir: impl AsRef<Path>,
        cfg: ParaformerConfig,
    ) -> Result<Self, AsrError> {
        let dir: PathBuf = model_dir.as_ref().to_path_buf();
        let model_path = dir.join(&cfg.model_filename);
        let mvn_path = dir.join(&cfg.cmvn_filename);
        let tokens_path = dir.join(&cfg.tokens_filename);

        let session = Session::builder()
            .and_then(|mut b| b.commit_from_file(&model_path))
            .map_err(|e| AsrError::ModelLoad(format!(
                    "failed to load Paraformer ONNX {}: {e}",
                    model_path.display()
                )))?;

        let cmvn = load_cmvn(&mvn_path)?;
        let vocab = load_tokens_json(&tokens_path)?;

        // Sanity: post-LFR feature dim must equal CMVN dim.
        let expected_dim = cfg.fbank.n_mels * cfg.lfr_m;
        if cmvn.dim() != expected_dim {
            return Err(AsrError::ModelLoad(format!(
                    "CMVN dim {} does not match n_mels({}) * lfr_m({}) = {}",
                    cmvn.dim(),
                    cfg.fbank.n_mels,
                    cfg.lfr_m,
                    expected_dim
                )));
        }

        let mut input_iter = session.inputs().iter().map(|i| i.name().to_string());
        let speech_name = input_iter.next().ok_or_else(|| AsrError::Inference("Paraformer ONNX has no inputs".into()))?;
        let lengths_name = input_iter.next().ok_or_else(|| AsrError::Config("Paraformer ONNX has only one input (expected two)".into()))?;

        let fbank = Fbank::new(cfg.fbank.clone());

        Ok(Self {
            session: Mutex::new(session),
            fbank,
            cmvn,
            vocab,
            cfg,
            input_names: (speech_name, lengths_name),
        })
    }

    /// Run the full extract → encode → decode pipeline on a single
    /// concatenated waveform. Exposed primarily so tests in cousin
    /// modules can drive the model end-to-end without spinning up a
    /// pipeline session.
    pub fn transcribe_samples(&self, samples: &[f32]) -> Result<String, AsrError> {
        if samples.is_empty() {
            return Err(AsrError::NoAudio);
        }

        let (mel, n_frames) = self.fbank.compute(samples);
        if n_frames == 0 {
            return Err(AsrError::Inference(format!(
                    "audio too short for one FBANK frame ({} samples)",
                    samples.len()
                )));
        }

        let n_mels = self.fbank.n_mels();
        let (mut feats, t_lfr) = apply_lfr(&mel, n_frames, n_mels, self.cfg.lfr_m, self.cfg.lfr_n);
        let feat_dim = n_mels * self.cfg.lfr_m;
        apply_cmvn(&mut feats, feat_dim, &self.cmvn);

        let speech: Array3<f32> =
            Array3::from_shape_vec((1, t_lfr, feat_dim), feats).map_err(|e| AsrError::Inference(format!("speech tensor shape: {e}")))?;
        let lengths: Array1<i32> = Array1::from(vec![t_lfr as i32]);

        let speech_val = Value::from_array(speech).map_err(|e| AsrError::Inference(format!("speech Value: {e}")))?;
        let lengths_val = Value::from_array(lengths).map_err(|e| AsrError::Inference(format!("lengths Value: {e}")))?;

        let (speech_name, lengths_name) =
            (self.input_names.0.as_str(), self.input_names.1.as_str());

        // Hold the session lock for the duration of decoding so the
        // borrowed output tensors stay live; the encoder pass is the
        // hot path and we want to avoid copying its full [T, V] tensor
        // before argmax.
        let mut session = self.session.lock().map_err(|e| AsrError::Inference(format!("session lock poisoned: {e}")))?;
        let outputs = session
            .run(vec![
                (speech_name, speech_val.into_dyn()),
                (lengths_name, lengths_val.into_dyn()),
            ])
            .map_err(|e| AsrError::Inference(format!("Paraformer ONNX run: {e}")))?;

        // Output 0: am_scores [1, T, V] f32
        let am = outputs[0]
            .try_extract_array::<f32>()
            .map_err(|e| AsrError::Inference(format!("extract am_scores: {e}")))?;
        let view = am.view();
        let shape = view.shape().to_vec();
        if shape.len() != 3 {
            return Err(AsrError::Inference(format!("unexpected am_scores rank {}", shape.len())));
        }
        let t = shape[1];
        let v = shape[2];

        // Output 1: valid_token_num [1] (i32 or i64)
        let valid_n = extract_first_int(&outputs[1]).ok_or_else(|| AsrError::ModelLoad("could not read valid_token_num output".into()))?;

        let usable = (valid_n as usize)
            .saturating_sub(self.cfg.predictor_bias)
            .min(t);

        let mut ids = Vec::with_capacity(usable);
        for ti in 0..usable {
            let mut best = 0u32;
            let mut best_v = f32::NEG_INFINITY;
            for vi in 0..v {
                let val = view[[0, ti, vi]];
                if val > best_v {
                    best_v = val;
                    best = vi as u32;
                }
            }
            ids.push(best);
        }

        let tokens = ids_to_tokens(&ids, &self.vocab);
        Ok(sentence_postprocess(&tokens))
    }
}

fn extract_first_int(value: &ort::value::DynValue) -> Option<i64> {
    if let Ok(arr) = value.try_extract_array::<i32>() {
        return arr.iter().next().copied().map(|x| x as i64);
    }
    if let Ok(arr) = value.try_extract_array::<i64>() {
        return arr.iter().next().copied();
    }
    None
}

#[async_trait]
impl AsrAdapter for ParaformerAdapter {
    async fn transcribe(&self, audio: &[AudioChunk]) -> Result<Transcript, AsrError> {
        let all_samples = AudioChunk::concat(audio);
        if all_samples.is_empty() {
            return Err(AsrError::NoAudio);
        }

        tracing::info!(
            audio_samples = all_samples.len(),
            "transcribing with paraformer-large"
        );

        let text = self.transcribe_samples(&all_samples)?;
        Ok(Transcript::new(text))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn load_nonexistent_dir_returns_error() {
        let res = ParaformerAdapter::load("/nonexistent/path/to/paraformer");
        assert!(res.is_err());
        let err = res.err().unwrap();
        assert!(
            err.to_string().contains("failed to load") || err.to_string().contains("read"),
            "expected load-failure message, got: {}",
            err
        );
    }

    #[test]
    fn adapter_is_send_and_sync() {
        fn assert_send_sync<T: Send + Sync>() {}
        assert_send_sync::<ParaformerAdapter>();
    }

    #[tokio::test]
    async fn transcribe_empty_audio_errors_via_mock_contract() {
        // We can't construct a real ParaformerAdapter without a model
        // bundle, but we can verify the contract — empty audio yields
        // "no audio received" — through MockAsr the same way the
        // parakeet test does. This guards against changes that drop
        // empty-channel handling.
        use crate::mock::MockAsr;
        use crate::traits::AsrAdapter;
        let mock = MockAsr::new("");
        let transcript = mock.transcribe(&[]).await.expect("mock always answers");
        assert!(transcript.text.is_empty());
    }
}