DOCSEARCHbyAlgolia
Submit Docs

parakeet-rs

repository·master·Indexed 18 days ago

https://github.com/altunenes/parakeet-rs

A high-performance Rust library for Automatic Speech Recognition (ASR) and speaker diarization using NVIDIA Parakeet models via ONNX Runtime. It supports CTC and TDT models for English and multilingual transcription, real-time streaming ASR via ParakeetEOU and Nemotron, offline multilingual ASR via Cohere, and speaker-attributed transcription with Multitalker and Sortformer. The library provides GPU acceleration support for CUDA, TensorRT, WebGPU, DirectML, and MIGraphX.

Tokens
13.9K
Snippets
52
Records
60
Agent score
60%

What's inside parakeet-rs

  1. Configure timestamp output modes

    master

    When calling transcription methods, you can specify the granularity of the returned timestamps using the TimestampMode enum. This determines whether timestamps are provided for individual tokens, words, or entire sentences.

    Supported modes (via TimestampMode):

    • Tokens
    • Words
    • Sentences

  2. Understand Sortformer streaming latency

    master

    The streaming latency of the Sortformer engine is determined by the chunk length and the lookahead (right context).

    Formula: (chunk_len + right_context) * 80ms

    For example, with a chunk_len of 124 and right_context of 1, the latency is approximately 10.0 seconds. You can check the current latency using the latency() method.

  3. Configure GPU execution providers

    master

    By default, parakeet-rs attempts to use GPU acceleration and falls back to CPU if it fails. You can explicitly request an execution provider via ExecutionConfig.

    Setup: Enable the appropriate feature in Cargo.toml (e.g., cuda, tensorrt, webgpu, directml, migraphx).

    Note for Apple Users: CoreML is unstable. Use the webgpu feature (which uses Metal under the hood) or standard cpu.

    Advanced Configuration: You can use with_custom_configure to access the underlying ort SessionBuilder for fine-grained control, such as disabling the memory pattern.

    use parakeet_rs::{Parakeet, ExecutionConfig, ExecutionProvider};

// Explicitly request CUDA
let config = ExecutionConfig::new().with_execution_provider(ExecutionProvider::Cuda);
let mut parakeet = Parakeet::from_pretrained(".", Some(config))?;

// Advanced configuration via ort SessionBuilder
let config = ExecutionConfig::new()
    .with_custom_configure(|builder| builder.with_memory_pattern(false));

  4. Manage KV cache with CoherePastKv

    master

    The CoherePastKv struct manages the Key/Value caches for both the decoder (self-attention) and the encoder (cross-attention) across all 8 layers.

    Lifecycle:

    1. Initialization: Start with an empty cache using CoherePastKv::empty(). At this stage, all cache tensors have a sequence length of 0.
    2. First Decoder Call: On the first call to run_decoder_step, the model populates the cross-attention caches (encoder_k, encoder_v) using the encoder_hidden_states.
    3. Subsequent Calls: The model writes new self-attention K/V into the growing decoder_k and decoder_v caches and reuses the existing encoder caches.

    Key Methods:

    • past_decoder_len(): Returns the number of tokens currently stored in the self-attention cache (derived from the shape of decoder_k[0]).
    let mut past_kv = CoherePastKv::empty();
// ... after first decoder step ...
let current_len = past_kv.past_decoder_len();

  5. Understand Nemotron model variants

    master

    The Nemotron ASR model comes in two variants, which are automatically detected from the ONNX encoder graph during loading:

    • EnglishOnly: The English-only 0.6B model. It uses a smaller vocabulary (1024) and does not support language conditioning.
    • Multilingual: The Nemotron 3.5 0.6B multilingual model. It has a larger vocabulary (~13k) and supports selecting a target language via a prompt_index.

    You can check the detected mode using the .mode() method on either a NemotronHandle or a Nemotron instance.

    let mode = nemotron_instance.mode(); // Returns NemotronMode::EnglishOnly or NemotronMode::Multilingual

  6. Choose a LatencyMode for the ASR pipeline

    master

    The LatencyMode enum controls the trade-off between transcription accuracy and processing latency by adjusting the encoder chunk size.

    ModeLatency (approx)Accuracy
    Normal (default)1.12sHighest
    Low0.56sMedium
    VeryLow0.16sLow
    Ultra0.08sLowest

    Use asr.chunk_audio_samples() to determine the exact number of audio samples required per chunk for your current mode.

    // Get the required number of samples for the current mode
let samples_per_chunk = asr.chunk_audio_samples();

  7. How to scale to multiple concurrent streams with ParakeetEOUHandle

    master

    To handle multiple concurrent audio streams efficiently, do not create multiple ParakeetEOU instances from disk. Instead, use a shared handle to avoid reloading the heavy ONNX models into memory multiple times.

    1. Load a single ParakeetEOUHandle using from_pretrained or load.
    2. For each new stream, create a lightweight ParakeetEOU instance using ParakeetEOU::from_shared(&handle).

    Each ParakeetEOU instance maintains its own independent encoder cache, decoder state, and audio buffer, while the expensive ONNX session is shared via the ParakeetEOUHandle.

    use parakeet_rs::parakeet_eou::{ParakeetEOUHandle, ParakeetEOU};

// 1. Load the shared handle once
let handle = ParakeetEOUHandle::from_pretrained("path/to/model_dir", None)?;

// 2. Spawn multiple independent streams from the same handle
let mut stream1 = ParakeetEOU::from_shared(&handle);
let mut stream2 = ParakeetEOU::from_shared(&handle);

// Stream 1 processing
let text1 = stream1.transcribe(&chunk1, true)?;
// Stream 2 processing
let text2 = stream2.transcribe(&chunk2, true)?;

  8. Configure TimestampMode for transcription results

    master

    The TimestampMode enum determines how token-level timestamps are grouped and presented in your transcription output. Choosing the correct mode depends on the model you are using:

    • Tokens: Provides raw, most detailed token-level output. Use this for maximum granularity.
    • Words: Groups subword tokens into individual words. Recommended for Parakeet CTC (English) models, as they do not predict punctuation.
    • Sentences: Groups tokens by sentence boundaries (., ?, !). Recommended for Parakeet TDT (Multilingual) models, which predict punctuation.

    Note: Using Sentences mode with a CTC model will not work effectively because CTC models do not output punctuation.

    #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum TimestampMode {
    #[default]
    Tokens,
    Words,
    Sentences,
}

  9. Understand MultitalkerEncoderCache structure

    master

    The MultitalkerEncoderCache manages the state for streaming encoder inference. It is specifically designed for the batch-first requirements of the Multitalker ONNX export.

    Fields:

    • cache_last_channel: Array4<f32> with shape [1, n_layers, left_context, d_model].
    • cache_last_time: Array4<f32> with shape [1, n_layers, d_model, conv_context].
    • cache_last_channel_len: Array1<i64> representing the current cache length [1].

    Initialization: Use MultitalkerEncoderCache::new(num_layers, left_context, hidden_dim, conv_context) to create an empty cache initialized with zeros.

    let cache = MultitalkerEncoderCache::new(num_layers, left_context, hidden_dim, conv_context);

  10. How to use ParakeetEOU for single-stream ASR

    master

    For a single audio stream, use the ParakeetEOU::from_pretrained method. This is a convenience wrapper that loads the model, tokenizer, and mel filterbank from a specified directory and returns a ready-to-use instance.

    Required files in the directory:

    • encoder.onnx
    • decoder_joint.onnx
    • tokenizer.json
    use parakeet_rs::parakeet_eou::ParakeetEOU;

let mut model = ParakeetEOU::from_pretrained("path/to/model_dir", None)?;
// Transcribe chunks of audio (16kHz)
let text = model.transcribe(&audio_chunk, true)?;

  11. Quick Start with parakeet-rs

    master

    To perform speech-to-text transcription, load a Parakeet model from a directory containing the required ONNX and configuration files, then use transcribe_samples to process audio data.

    Model Directory Requirements: Your model directory must contain:

    • model.onnx: The ONNX model file
    • model.onnx_data: External model weights
    • config.json: Model configuration
    • preprocessor_config.json: Audio preprocessing configuration
    • tokenizer.json: Tokenizer vocabulary
    • tokenizer_config.json: Tokenizer configuration

    Audio Requirements:

    • Format: WAV
    • Sample Rate: 16kHz
    • Channels: Mono (stereo is automatically converted)
    • Bit Depth: 16-bit PCM or 32-bit float
    use parakeet_rs::{Parakeet, Transcriber, TimestampMode};

// Load the model
let mut parakeet = Parakeet::from_pretrained(".")?;

// Transcribe audio samples
let result = parakeet.transcribe_samples(audio, sample_rate, channels, Some(TimestampMode::Words))?;
println!("Transcription: {}", result.text);

  12. Use Parakeet TDT for multilingual transcription

    master

    The ParakeetTDT model supports 25 languages with automatic language detection. It also supports token-level timestamps.

    Setup: Download encoder-model.onnx, encoder-model.onnx.data, decoder_joint-model.onnx, and vocab.txt from HuggingFace.

    use parakeet_rs::{ParakeetTDT, Transcriber, TimestampMode};

let mut parakeet = ParakeetTDT::from_pretrained("./tdt", None)?;
let result = parakeet.transcribe_samples(audio, 16000, 1, Some(TimestampMode::Sentences))?;
println!("{}", result.text);

// Token-level timestamps
for token in result.tokens {
    println!("[{:.3}s - {:.3}s] {}", token.start, token.end, token.text);
}