transcribe-rs

By default, engines use the CPU. You can enable GPU acceleration by enabling the appropriate feature flag at compile time and setting the accelerator preference at runtime before loading any models.

ORT (ONNX Runtime) Engines

Use set_ort_accelerator with OrtAccelerator values. Note that OrtAccelerator::Auto does not include DirectML due to specific session requirements; you must select OrtAccelerator::DirectMl explicitly if using DirectML.

• OrtAccelerator::Cuda: NVIDIA CUDA
• OrtAccelerator::Auto: Auto-detect best available GPU
• OrtAccelerator::DirectMl: Microsoft DirectML (Windows)
• OrtAccelerator::available(): Query which accelerators are compiled in.

whisper.cpp Engines

GPU backends (Metal, Vulkan, CUDA) are selected at compile time via feature flags. Use set_whisper_accelerator to control usage at runtime.

• WhisperAccelerator::CpuOnly: Force CPU usage.

use transcribe_rs::{set_ort_accelerator, OrtAccelerator};

// Use CUDA for all ORT engines (SenseVoice, GigaAM, Parakeet, Moonshine)
set_ort_accelerator(OrtAccelerator::Cuda);

// Or auto-detect the best available GPU
set_ort_accelerator(OrtAccelerator::Auto);

use transcribe_rs::{set_whisper_accelerator, WhisperAccelerator};

set_whisper_accelerator(WhisperAccelerator::CpuOnly); // force CPU

Use one of two patterns depending on model loading cost:

Pattern A: Lightweight models

For models that load quickly (e.g., most ONNX models), load a new instance per test. Always use common::require_paths() to skip tests gracefully if model files are missing.

// tests/your_model.rs
mod common;

use std::path::PathBuf;
use transcribe_rs::your_engine::your_model::YourModel;
use transcribe_rs::SpeechModel;

#[test]
fn test_your_model_transcribe() {
    env_logger::init();

    let model_dir = PathBuf::from("models/your-model");
    let wav_path = PathBuf::from("samples/jfk.wav");

    if !common::require_paths(&[&model_dir, &wav_path]) {
        return;
    }

    let mut model = YourModel::load(&model_dir).expect("Failed to load model");

    let result = model
        .transcribe_file(&wav_path, None)
        .expect("Failed to transcribe");

    assert!(!result.text.is_empty(), "Transcription should not be empty");
}

Pattern B: Expensive models

For slow-loading models (e.g., whisper.cpp), use a shared instance via once_cell::sync::Lazy and std::sync::Mutex to avoid reloading the model for every test.

// tests/your_model.rs
mod common;

use once_cell::sync::Lazy;
use std::path::PathBuf;
use std::sync::Mutex;
use transcribe_rs::your_engine::YourEngine;
use transcribe_rs::SpeechModel;

static ENGINE: Lazy<Mutex<Option<YourEngine>>> = Lazy::new(|| {
    let model = PathBuf::from("models/your-model.bin");
    if !common::require_paths(&[&model]) {
        return Mutex::new(None);
    }

    match YourEngine::load(&model) {
        Ok(engine) => Mutex::new(Some(engine)),
        Err(_) => Mutex::new(None),
    }
});

#[test]
fn test_transcription() {
    let mut guard = match ENGINE.lock().unwrap().as_mut() {
        Some(e) => e,
        None => return,
    };

    let audio_path = PathBuf::from("samples/jfk.wav");
    let result = guard.transcribe_file(&audio_path, None).expect("Failed");
    assert!(!result.text.is_empty());
}

Registration

Register tests in Cargo.toml with the required feature:

[[test]]
name = "your_model"
required-features = ["your-engine"]

A new engine family is required when integrating a new inference runtime (e.g., Candle, Burn, MLX, TensorRT). Follow these steps:

1. Create the source directory

Create src/your_engine/ with a mod.rs for engine-level types and re-exports, and subdirectories for individual models (e.g., src/your_engine/your_model/mod.rs). For simple engines, a single file src/your_engine.rs is sufficient.

2. Configure Cargo.toml

Add a feature flag for your engine and include it in the all feature. If your engine requires shared audio features (like mel spectrograms), include the audio-features feature.

[features]
your-engine = ["dep:your-runtime-crate"]

# Update the "all" feature
all = ["onnx", "whisper-cpp", "whisperfile", "openai", "your-engine"]

[dependencies]
your-runtime-crate = { version = "...", optional = true }

Note: To include audio features: your-engine = ["audio-features", "dep:your-runtime-crate"]

3. Register the module

In src/lib.rs, use a conditional compilation guard:

#[cfg(feature = "your-engine")]
pub mod your_engine;

4. Implement the SpeechModel trait

Every local model must implement the SpeechModel trait. Follow these conventions:

• CAPABILITIES: Define const CAPABILITIES: ModelCapabilities.
• Constructor: Implement Model::load(...) as a single-step constructor.
• Params: Create a {Model}Params struct with #[derive(Debug, Clone, Default)].
• Methods: Implement transcribe_with(&mut self, samples, &Params) for engine-specific parameters, and the standard transcribe() for the generic interface.
• Errors: Map runtime errors to TranscribeError variants like ModelNotFound, Inference, Audio, or Config.

5. Error Conversions

If your runtime has a custom error type, implement From in src/error.rs:

#[cfg(feature = "your-engine")]
impl From<your_runtime::Error> for TranscribeError {
    fn from(e: your_runtime::Error) -> Self {
        TranscribeError::Inference(e.to_string())
    }
}

6. Documentation

Create src/your_engine/PORTING.md to document how to add models to this specific engine family.

To add a new speech recognition model to the onnx engine family, follow these steps:

1. Create the directory structure: Create a new directory under src/onnx/ (e.g., src/onnx/your_model/) and add a mod.rs file.
2. Register the module: Add pub mod your_model; to src/onnx/mod.rs.
3. Implement the model: Implement the SpeechModel trait and provide a load method. Use src/onnx/gigaam/mod.rs as a simple reference.
4. Add a test: Create a test file in tests/your_model.rs and register it in Cargo.toml with required-features = ["onnx"].
5. Add an example: Create an example file in examples/your_model.rs and register it in Cargo.toml with required-features = ["onnx"].
6. Prepare model files: Place your .onnx files and vocabulary files in a directory (e.g., models/your-model/).

mkdir src/onnx/your_model
touch src/onnx/your_model/mod.rs

When implementing a new model, follow this structure in src/onnx/your_model/mod.rs. The model must define its ModelCapabilities, a Params struct for model-specific settings, and implement the SpeechModel trait.

Key requirements:

• Constructor: Use pub fn load(model_dir: &Path, quantization: &Quantization) -> Result<Self, TranscribeError>.
• Capabilities: Define a const CAPABILITIES: ModelCapabilities containing the model's name, engine ID, sample rate, and supported languages.
• Inference: Implement the infer method to handle feature extraction, tensor preparation, the ONNX forward pass, and decoding.
• Trait Implementation: Implement SpeechModel to allow generic usage via the transcribe method.

use ort::inputs;
use ort::session::Session;
use ort::value::TensorRef;
use std::path::Path;

use crate::decode::{ctc_greedy_decode, sentencepiece_to_text};
use crate::decode::tokens::load_vocab;
use crate::features::{compute_mel, MelConfig, WindowType};
use crate::TranscribeError;
use super::session;
use super::Quantization;
use crate::{ModelCapabilities, SpeechModel, TranscriptionResult};

const CAPABILITIES: ModelCapabilities = ModelCapabilities {
    name: "Your Model",           // Human-readable name
    engine_id: "your_model",      // Machine identifier, matches directory name
    sample_rate: 16000,           // Expected input sample rate
    languages: &["en"],
    supports_timestamps: false,
    supports_translation: false,
    supports_streaming: false,
};

#[derive(Debug, Clone, Default)]
pub struct YourModelParams {
    pub language: Option<String>,
}

pub struct YourModel {
    session: Session,
}

impl YourModel {
    pub fn load(model_dir: &Path, quantization: &Quantization) -> Result<Self, TranscribeError> {
        let model_path = session::resolve_model_path(model_dir, "model", quantization);

        if !model_path.exists() {
            return Err(TranscribeError::ModelNotFound(model_path));
        }

        let session = session::create_session(&model_path)?;

        Ok(Self { session })
    }

    pub fn transcribe_with(
        &mut self, 
        samples: &[f32], 
        params: &YourModelParams
    ) -> Result<TranscriptionResult, TranscribeError> {
        self.infer(samples)
    }

    fn infer(&mut self, samples: &[f32]) -> Result<TranscriptionResult, TranscribeError> {
        // 1. Feature extraction
        // 2. Prepare input tensors
        // 3. Run ONNX forward pass: self.session.run(inputs)?
        // 4. Decode output
        // 5. Return TranscriptionResult
        todo!()
    }
}

impl SpeechModel for YourModel {
    fn capabilities(&self) -> ModelCapabilities {
        CAPABILITIES
    }

    fn transcribe(
        &mut self, 
        samples: &[f32], 
        _options: &TranscribeOptions
    ) -> Result<TranscriptionResult, TranscribeError> {
        self.infer(samples)
    }
}

Models must be downloaded and placed in your local directory before use. The required directory structure depends on the engine.

Model Download Links

Add transcribe-rs to your Cargo.toml. Note that no features are enabled by default; you must explicitly enable the engines you intend to use via feature flags.

Engine Features

GPU Accelerator Features for whisper.cpp

GPU Accelerator Features for ORT (ONNX Runtime) engines

[dependencies]
transcribe-rs = { version = "0.3", features = ["onnx"] }

Every model should include an example demonstrating loading and transcription. Use transcribe_with() to show how to pass model-specific parameters.

Conventions:

• Show load and transcription timing.
• Calculate and print the real-time speedup factor.
• Print segments if the model supports timestamps.
• Register the example in Cargo.toml with required-features.

// examples/your_model.rs
use std::path::PathBuf;
use std::time::Instant;
use transcribe_rs::your_engine::your_model::{YourModel, YourModelParams};
use transcribe_rs::SpeechModel;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    env_logger::init();
    let model_path = PathBuf::from("models/your-model");
    let wav_path = PathBuf::from("samples/jfk.wav");

    // Load
    let mut model = YourModel::load(&model_path)?;

    // Transcribe
    let samples = transcribe_rs::audio::read_wav_samples(&wav_path)?;
    let result = model.transcribe_with(
        &samples,
        &YourModelParams {
            language: Some("en".to_string()),
            ..Default::default()
        },
    )?;

    println!("Transcription result: {}", result.text);
    Ok(())
}

Registration in Cargo.toml:

[[example]]
name = "your_model"
required-features = ["your-engine"]