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}