parakeet-rs 0.4.2

use crate::audio::{self, FeatureCache};
use crate::config::PreprocessorConfig;
use crate::decoder::TranscriptionResult;
use crate::decoder_tdt::ParakeetTDTDecoder;
use crate::error::{Error, Result};
use crate::execution::ModelConfig as ExecutionConfig;
use crate::model_tdt::ParakeetTDTModel;
use crate::timestamps::{process_timestamps, rebuild_text, TimestampMode};
use crate::transcriber::Transcriber;
use crate::vocab::Vocabulary;
use std::path::{Path, PathBuf};

/// Parakeet TDT model for multilingual ASR
pub struct ParakeetTDT {
    model: ParakeetTDTModel,
    decoder: ParakeetTDTDecoder,
    preprocessor_config: PreprocessorConfig,
    feature_cache: FeatureCache,
    model_dir: PathBuf,
}

impl ParakeetTDT {
    /// Load Parakeet TDT model from path with optional configuration.
    ///
    /// # Arguments
    /// * `path` - Directory containing encoder-model.onnx, decoder_joint-model.onnx, and vocab.txt
    /// * `config` - Optional execution configuration (defaults to CPU if None)
    ///
    /// With a burn GPU provider (e.g. [`ExecutionProvider::BurnWgpu`](crate::ExecutionProvider))
    /// the encoder runs on the GPU and the decoder/joint, which runs once per token, on burn's CPU
    /// backend: a GPU round trip per token costs more than the step itself. Use
    /// [`from_pretrained_with_joint_config`](Self::from_pretrained_with_joint_config) to choose.
    pub fn from_pretrained<P: AsRef<Path>>(
        path: P,
        config: Option<ExecutionConfig>,
    ) -> Result<Self> {
        let mut joint_config = config.clone();
        if let Some(c) = joint_config.as_mut() {
            c.execution_provider = c.execution_provider.per_token_provider();
        }
        Self::from_pretrained_with_joint_config(path, config, joint_config)
    }

    /// Load Parakeet TDT model from path, running the decoder/joint graph under its own execution
    /// configuration.
    ///
    /// The encoder and the decoder/joint are separate ONNX graphs and do not have to run on the
    /// same execution provider. Splitting them matters when a provider can take one graph and not
    /// the other, and the joint runs once per token, so where it runs is a decision of its own.
    ///
    /// # Arguments
    /// * `path` - Directory containing encoder-model.onnx, decoder_joint-model.onnx, and vocab.txt
    /// * `config` - Optional execution configuration for the encoder (defaults to CPU if None)
    /// * `joint_config` - Optional execution configuration for the decoder/joint (defaults to CPU
    ///   if None)
    pub fn from_pretrained_with_joint_config<P: AsRef<Path>>(
        path: P,
        config: Option<ExecutionConfig>,
        joint_config: Option<ExecutionConfig>,
    ) -> Result<Self> {
        let path = path.as_ref();

        if !path.is_dir() {
            return Err(Error::Config(format!(
                "TDT model path must be a directory: {}",
                path.display()
            )));
        }

        let vocab_path = path.join("vocab.txt");
        if !vocab_path.exists() {
            return Err(Error::Config(format!(
                "vocab.txt not found in {}",
                path.display()
            )));
        }

        // TDT-specific preprocessor config (128 features instead of 80)
        let preprocessor_config = PreprocessorConfig {
            feature_extractor_type: "ParakeetFeatureExtractor".to_string(),
            feature_size: 128,
            hop_length: 160,
            n_fft: 512,
            padding_side: "right".to_string(),
            padding_value: 0.0,
            preemphasis: 0.97,
            processor_class: "ParakeetProcessor".to_string(),
            return_attention_mask: true,
            sampling_rate: 16000,
            win_length: 400,
        };

        let exec_config = config.unwrap_or_default();
        let joint_exec_config = joint_config.unwrap_or_else(|| {
            #[allow(unused_mut)]
            let mut joint = ExecutionConfig::default();
            #[cfg(feature = "ort")]
            {
                joint.environment = exec_config.environment.clone();
            }
            joint
        });

        // Load vocab first to get the actual vocabulary size
        let vocab = Vocabulary::from_file(&vocab_path)?;
        let vocab_size = vocab.size();

        let model = ParakeetTDTModel::from_pretrained_with_configs(
            path,
            exec_config,
            joint_exec_config,
            vocab_size,
        )?;
        let decoder = ParakeetTDTDecoder::from_vocab(vocab);
        let feature_cache = FeatureCache::from_config(&preprocessor_config);

        Ok(Self {
            model,
            decoder,
            preprocessor_config,
            feature_cache,
            model_dir: path.to_path_buf(),
        })
    }

    pub fn model_dir(&self) -> &Path {
        &self.model_dir
    }

    pub fn preprocessor_config(&self) -> &PreprocessorConfig {
        &self.preprocessor_config
    }
}

impl Transcriber for ParakeetTDT {
    fn transcribe_samples(
        &mut self,
        audio: Vec<f32>,
        sample_rate: u32,
        channels: u16,
        mode: Option<TimestampMode>,
    ) -> Result<TranscriptionResult> {
        let features = audio::extract_features_with_cache(
            audio,
            sample_rate,
            channels,
            &self.preprocessor_config,
            &self.feature_cache,
        )?;
        let (tokens, frame_indices, durations) = self.model.forward(features)?;

        let mut result = self.decoder.decode_with_timestamps(
            &tokens,
            &frame_indices,
            &durations,
            self.preprocessor_config.hop_length,
            self.preprocessor_config.sampling_rate,
        )?;

        // Apply timestamp mode conversion
        let mode = mode.unwrap_or(TimestampMode::Tokens);
        result.tokens = process_timestamps(&result.tokens, mode);

        // Rebuild full text from processed tokens
        result.text = rebuild_text(&result.tokens, mode);

        Ok(result)
    }
}

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

    // Apps share models across threads, on every backend.
    #[test]
    fn parakeet_tdt_is_send_and_sync() {
        fn assert_send_sync<T: Send + Sync>() {}
        assert_send_sync::<ParakeetTDT>();
    }
}