ravenrs_ai/
hearing.rs

1//! Speech to text with Whisper, on the CPU in plain Rust, the way OpenAI's reference
2//! reads it: each 30 seconds padded with silence, the language held, the likeliest word
3//! taken each step.
4
5use std::fmt;
6
7use candle_core::{D, Device, IndexOp, Tensor};
8use candle_nn::VarBuilder;
9use candle_transformers::models::whisper::{self, Config, audio, model::Whisper};
10use tokenizers::Tokenizer;
11
12use crate::engine::EngineError;
13use crate::{LocalModel, ModelFolder};
14
15/// The rate a transcriber hears at: give it mono samples at 16 kHz.
16pub const HEARING_RATE: usize = whisper::SAMPLE_RATE;
17
18/// What a transcription found.
19#[derive(Clone, Debug, PartialEq)]
20pub struct Transcribed {
21    /// The words, trimmed.
22    pub text: String,
23    /// How many tokens it wrote.
24    pub tokens: usize,
25    /// How long it took, in seconds.
26    pub seconds: f64,
27}
28
29/// Whisper, loaded and ready to hear.
30pub struct Transcriber {
31    model: Whisper,
32    tokenizer: Tokenizer,
33    filters: Vec<f32>,
34    suppressed: Tensor,
35    device: Device,
36    start: u32,
37    transcribe: u32,
38    no_timestamps: u32,
39    end: u32,
40}
41
42impl fmt::Debug for Transcriber {
43    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
44        f.debug_struct("Transcriber").field("device", &self.device).finish_non_exhaustive()
45    }
46}
47
48impl Transcriber {
49    /// Loads `model` from `folder`.
50    ///
51    /// # Errors
52    ///
53    /// When its files are missing or are not a Whisper model.
54    pub fn load(folder: &ModelFolder, model: &LocalModel) -> Result<Self, EngineError> {
55        let load = |why: String| EngineError::Load(why);
56        let named = |name: &str| {
57            let file = model.files.iter().find(|file| file.name == name);
58            file.map(|file| folder.path_of(model, file)).ok_or_else(|| load(format!("no {name}")))
59        };
60        let device = Device::Cpu;
61        let config: Config =
62            serde_json::from_slice(&std::fs::read(named("config.json")?).map_err(|e| load(e.to_string()))?)
63                .map_err(|e| load(e.to_string()))?;
64        let tensors = candle_core::safetensors::load(named("model.safetensors")?, &device)
65            .map_err(|e| load(e.to_string()))?;
66        let weights = VarBuilder::from_tensors(tensors, whisper::DTYPE, &device);
67        let tokenizer = Tokenizer::from_file(named("tokenizer.json")?).map_err(|e| load(e.to_string()))?;
68        let token = |text: &str| tokenizer.token_to_id(text).ok_or_else(|| load(format!("no token {text}")));
69        let (start, transcribe) = (token(whisper::SOT_TOKEN)?, token(whisper::TRANSCRIBE_TOKEN)?);
70        let (no_timestamps, end) = (token(whisper::NO_TIMESTAMPS_TOKEN)?, token(whisper::EOT_TOKEN)?);
71        let mask: Vec<f32> = (0..config.vocab_size as u32)
72            .map(|id| if config.suppress_tokens.contains(&id) { f32::NEG_INFINITY } else { 0.0 })
73            .collect();
74        let suppressed = Tensor::new(mask.as_slice(), &device).map_err(|e| load(e.to_string()))?;
75        let filters = mel_filters(config.num_mel_bins);
76        let model = Whisper::load(&weights, config).map_err(|e| load(e.to_string()))?;
77        Ok(Self { model, tokenizer, filters, suppressed, device, start, transcribe, no_timestamps, end })
78    }
79
80    /// The words in `samples`, mono at [`HEARING_RATE`], spoken in `language` (its
81    /// two-letter code, such as `pt`).
82    ///
83    /// # Errors
84    ///
85    /// When the language is not one Whisper knows, or it stops in the middle.
86    pub fn transcribe(&mut self, samples: &[f32], language: &str) -> Result<Transcribed, EngineError> {
87        let started = std::time::Instant::now();
88        let run = |why: candle_core::Error| EngineError::Run(why.to_string());
89        let spoken = self
90            .tokenizer
91            .token_to_id(&format!("<|{language}|>"))
92            .ok_or_else(|| EngineError::Run(format!("Whisper does not know the language {language}")))?;
93        let prompt = [self.start, spoken, self.transcribe, self.no_timestamps];
94        let bins = self.model.config.num_mel_bins;
95        let most = self.model.config.max_target_positions / 2;
96        let (mut text, mut written) = (String::new(), 0);
97        for chunk in samples.chunks(whisper::N_SAMPLES) {
98            let mut padded = chunk.to_vec();
99            padded.resize(whisper::N_SAMPLES, 0.0);
100            let mel = audio::pcm_to_mel(&self.model.config, &padded, &self.filters);
101            let frames = mel.len() / bins;
102            let mel = Tensor::from_vec(mel, (1, bins, frames), &self.device)
103                .and_then(|mel| mel.narrow(2, 0, frames.min(whisper::N_FRAMES)))
104                .map_err(run)?;
105            let heard = self.model.encoder.forward(&mel, true).map_err(run)?;
106            let mut tokens = prompt.to_vec();
107            for step in 0..most {
108                let input =
109                    Tensor::new(tokens.as_slice(), &self.device).and_then(|t| t.unsqueeze(0)).map_err(run)?;
110                let output = self.model.decoder.forward(&input, &heard, step == 0).map_err(run)?;
111                let last = output.dim(1).map_err(run)? - 1;
112                let next = output
113                    .i((..1, last..))
114                    .and_then(|last| self.model.decoder.final_linear(&last))
115                    .and_then(|logits| logits.i(0)?.i(0)?.broadcast_add(&self.suppressed))
116                    .and_then(|logits| logits.argmax(D::Minus1)?.to_scalar::<u32>())
117                    .map_err(run)?;
118                if next == self.end {
119                    break;
120                }
121                tokens.push(next);
122            }
123            written += tokens.len() - prompt.len();
124            let words = self
125                .tokenizer
126                .decode(&tokens[prompt.len()..], true)
127                .map_err(|e| EngineError::Run(e.to_string()))?;
128            text.push_str(words.trim());
129            text.push(' ');
130        }
131        Ok(Transcribed {
132            text: text.trim().to_owned(),
133            tokens: written,
134            seconds: started.elapsed().as_secs_f64(),
135        })
136    }
137}
138
139/// The mel filters Whisper was trained with: librosa's, with Slaney's scale and area
140/// normalization, for its 400-point FFT at 16 kHz. Worked out here rather than shipped.
141fn mel_filters(bins: usize) -> Vec<f32> {
142    const MIN_LOG_HZ: f64 = 1000.0;
143    const HZ_PER_MEL: f64 = 200.0 / 3.0;
144    let min_log_mel = MIN_LOG_HZ / HZ_PER_MEL;
145    let log_step = 6.4_f64.ln() / 27.0;
146    let to_mel = |hz: f64| {
147        if hz >= MIN_LOG_HZ { min_log_mel + (hz / MIN_LOG_HZ).ln() / log_step } else { hz / HZ_PER_MEL }
148    };
149    let to_hz = |mel: f64| {
150        if mel >= min_log_mel {
151            MIN_LOG_HZ * (log_step * (mel - min_log_mel)).exp()
152        } else {
153            HZ_PER_MEL * mel
154        }
155    };
156    let rate = HEARING_RATE as f64;
157    let columns = whisper::N_FFT / 2 + 1;
158    let top = to_mel(rate / 2.0);
159    let edges: Vec<f64> = (0..bins + 2).map(|i| to_hz(top * i as f64 / (bins + 1) as f64)).collect();
160    let mut filters = Vec::with_capacity(bins * columns);
161    for band in 0..bins {
162        let (low, middle, high) = (edges[band], edges[band + 1], edges[band + 2]);
163        let area = 2.0 / (high - low);
164        for column in 0..columns {
165            let hz = column as f64 * rate / whisper::N_FFT as f64;
166            let rising = (hz - low) / (middle - low);
167            let falling = (high - hz) / (high - middle);
168            filters.push((rising.min(falling).max(0.0) * area) as f32);
169        }
170    }
171    filters
172}
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177
178    #[test]
179    fn mel_filters_are_triangles_over_the_spectrum() {
180        let filters = mel_filters(80);
181        assert_eq!(filters.len(), 80 * 201);
182        for band in filters.chunks(201) {
183            assert!(band.iter().all(|weight| *weight >= 0.0));
184            assert!(band.iter().any(|weight| *weight > 0.0), "every band hears something");
185        }
186        // The first band starts at 0 Hz, on the second FFT bin, as in librosa.
187        assert_eq!(filters[0], 0.0);
188        assert!(filters[1] > 0.0);
189    }
190
191    /// The same numbers as the filters candle's example ships, when that file is given.
192    #[test]
193    fn mel_filters_match_the_published_ones() {
194        let Some(path) = std::env::var_os("WHISPER_MEL_FILTERS") else { return };
195        let bytes = std::fs::read(path).unwrap_or_default();
196        let published: Vec<f32> = bytes.as_chunks::<4>().0.iter().map(|b| f32::from_le_bytes(*b)).collect();
197        let ours = mel_filters(80);
198        assert_eq!(published.len(), ours.len());
199        let worst = published.iter().zip(&ours).map(|(a, b)| (a - b).abs()).fold(0.0, f32::max);
200        assert!(worst < 1e-6, "differs by {worst}");
201    }
202}
ravenrs_ai/hearing.rs

ravenrs_ai/
hearing.rs
