Skip to main content

uni_xervo/traits/
asr.rs

1//! Automatic speech recognition types and trait.
2
3use crate::error::{Result, RuntimeError};
4use crate::traits::ModelInfo;
5use crate::traits::multimodal::AudioInput;
6use async_trait::async_trait;
7
8/// Options for a [`TranscriptionModel::transcribe`] call.
9#[derive(Debug, Clone, Default)]
10pub struct TranscribeOptions {
11    /// ISO 639-1 language code (e.g. `"en"`). `None` requests auto-detection.
12    pub language: Option<String>,
13    /// Whether to populate per-word timestamps on each segment.
14    pub word_timestamps: bool,
15    /// Whether to populate speaker labels via diarization.
16    pub diarize: bool,
17    /// Optional initial prompt for biasing decoding (e.g. domain-specific
18    /// terminology, named-entity priming). Supported by whisper.cpp et al.
19    pub initial_prompt: Option<String>,
20}
21
22/// Result of a transcription call.
23#[derive(Debug, Clone)]
24pub struct TranscribeResult {
25    /// Detected (or supplied) ISO 639-1 language code.
26    pub language: String,
27    /// Ordered speech segments.
28    pub segments: Vec<TranscribeSegment>,
29}
30
31/// One transcribed segment.
32#[derive(Debug, Clone)]
33pub struct TranscribeSegment {
34    /// Start timestamp in milliseconds from the audio start.
35    pub start_ms: u64,
36    /// End timestamp in milliseconds from the audio start (exclusive).
37    pub end_ms: u64,
38    /// Recognized text.
39    pub text: String,
40    /// Speaker label, populated when [`TranscribeOptions::diarize`] is set
41    /// and supported by the provider.
42    pub speaker: Option<String>,
43    /// Per-word timestamps, populated when
44    /// [`TranscribeOptions::word_timestamps`] is set and supported.
45    pub words: Vec<TranscribeWord>,
46}
47
48/// One word-level timestamp entry.
49#[derive(Debug, Clone)]
50pub struct TranscribeWord {
51    /// Start timestamp in milliseconds.
52    pub start_ms: u64,
53    /// End timestamp in milliseconds (exclusive).
54    pub end_ms: u64,
55    /// Word surface form.
56    pub text: String,
57    /// Provider-reported confidence in `[0.0, 1.0]`, if available.
58    pub confidence: Option<f32>,
59}
60
61/// A model that transcribes speech audio into text with timing information.
62///
63/// Targets whisper.cpp / whisper-rs, OpenAI Whisper API, AssemblyAI, and
64/// similar ASR engines. The primary method is batch (matching the
65/// batch-in/batch-out convention of the other tasks); use
66/// [`transcribe_one`](TranscriptionModel::transcribe_one) for the single-stream
67/// convenience.
68#[async_trait]
69pub trait TranscriptionModel: ModelInfo {
70    /// Transcribe a batch of audio inputs.
71    ///
72    /// This is the canonical primitive every implementation provides. Results
73    /// are returned in the same order as `audios`; the shared `options` apply to
74    /// every input. Single-stream engines (whisper.cpp, remote APIs) loop or fan
75    /// out internally; GPU-batched backends batch natively.
76    ///
77    /// # Errors
78    /// Returns an error if the provider cannot decode an input or fails internally.
79    async fn transcribe(
80        &self,
81        audios: Vec<AudioInput>,
82        options: TranscribeOptions,
83    ) -> Result<Vec<TranscribeResult>>;
84
85    /// Transcribe a single audio stream — convenience over
86    /// [`transcribe`](TranscriptionModel::transcribe).
87    ///
88    /// # Errors
89    /// Returns an error if transcription fails, or (a provider contract
90    /// violation) if no result is returned for the single input.
91    async fn transcribe_one(
92        &self,
93        audio: AudioInput,
94        options: TranscribeOptions,
95    ) -> Result<TranscribeResult> {
96        self.transcribe(vec![audio], options)
97            .await?
98            .into_iter()
99            .next()
100            .ok_or_else(|| {
101                RuntimeError::InferenceError(
102                    "transcribe returned no result for a single input".to_string(),
103                )
104            })
105    }
106
107    /// ISO 639-1 language codes the model can handle. May be empty for
108    /// providers that report language only at runtime via auto-detection.
109    fn supported_languages(&self) -> &[String];
110
111    /// Optional warmup hook. The default is a no-op.
112    ///
113    /// # Errors
114    /// Returns an error if the underlying model fails to initialize.
115    async fn warmup(&self) -> Result<()> {
116        Ok(())
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123    use std::sync::Arc;
124
125    #[test]
126    fn transcription_model_is_dyn_safe() {
127        fn _accept(_: Arc<dyn TranscriptionModel>) {}
128    }
129}