@ondeinference / onde
Code Commits Issues Pulls Cloud Agent CI/CD
feature/diffusion-vision
copilot/assess-inference-engine-complexity copilot/fix-onde-npm-sdk-v1-1-2-release development feature/diffusion-vision main
rs 931 lines 29.6 KB
Raw
1 //! onde-wasm — Whisper speech-to-text compiled to WebAssembly.
2 //!
3 //! This crate exposes a single [`WhisperDecoder`] class to JavaScript via
4 //! `wasm-bindgen`. It is backed by `candle-transformers`' pure-Rust Whisper
5 //! implementation, which compiles cleanly to `wasm32-unknown-unknown` without
6 //! any C FFI or platform intrinsics.
7 //!
8 //! # JavaScript API
9 //!
10 //! ```js
11 //! import init, { WhisperDecoder } from "./onde_wasm.js";
12 //!
13 //! await init();
14 //!
15 //! // All four byte arrays are fetched from HuggingFace (see build.sh).
16 //! const decoder = new WhisperDecoder(
17 //! weightsBytes, // model.safetensors — or model-tiny-en-q80.gguf
18 //! tokenizerBytes, // tokenizer.json
19 //! melFiltersBytes, // mel_filters.safetensors
20 //! configBytes, // config.json
21 //! quantized, // bool — true for .gguf weights
22 //! isMultilingual, // bool — false for *.en models
23 //! timestamps, // bool — include per-segment timestamps
24 //! task, // "transcribe" | "translate" | null
25 //! language, // BCP-47 string | null (null = auto-detect)
26 //! );
27 //!
28 //! // wavBytes is a Uint8Array of a 16 kHz mono WAV file.
29 //! const json = decoder.decode(wavBytes);
30 //!
31 //! // json is a string with the shape:
32 //! // {
33 //! // "text": "Full transcript…",
34 //! // "segments": [
35 //! // { "start": 0.0, "end": 3.5, "text": " Hello world" },
36 //! // …
37 //! // ]
38 //! // }
39 //! const result = JSON.parse(json);
40 //! ```
41 //!
42 //! # Result JSON shape
43 //!
44 //! The JSON string returned by [`WhisperDecoder::decode`] mirrors the
45 //! `WhisperResult` / `WhisperSegment` types in `onde::whisper` so that the
46 //! same TypeScript interface can be reused across both the native Tauri app
47 //! and the browser connected-devices app.
48 //!
49 //! ```ts
50 //! interface OndeSegment { start: number; end: number; text: string; }
51 //! interface OndeResult { text: string; segments: OndeSegment[]; }
52 //! ```
53 //!
54 //! # Model assets
55 //!
56 //! | File | Source |
57 //! |---|---|
58 //! | `model.safetensors` | `openai/whisper-tiny` (or tiny.en, base, …) |
59 //! | `model-tiny-en-q80.gguf` | `lmz/candle-whisper` (quantized) |
60 //! | `tokenizer.json` | same HF repo |
61 //! | `config.json` | same HF repo |
62 //! | `mel_filters.safetensors` | `huggingface/candle-whisper` space assets |
63 //!
64 //! # Build
65 //!
66 //! ```sh
67 //! wasm-pack build --target web --out-dir pkg
68 //! ```
69 //!
70 //! See `onde/crates/onde-wasm/build.sh` for the full reproducible build
71 //! script with asset download instructions.
72
73 use anyhow::anyhow;
74 use candle_core::{safetensors::Load, Device, IndexOp, Tensor, D};
75 use candle_nn::{ops::softmax, VarBuilder};
76 use candle_transformers::models::whisper::{self as m, Config};
77 use rand::{distr::Distribution, rngs::StdRng, SeedableRng};
78 use serde::{Deserialize, Serialize};
79 use tokenizers::Tokenizer;
80 use wasm_bindgen::prelude::*;
81
82 // ---------------------------------------------------------------------------
83 // Panic hook
84 // ---------------------------------------------------------------------------
85
86 fn set_panic_hook() {
87 #[cfg(feature = "console_error_panic_hook")]
88 console_error_panic_hook::set_once();
89 }
90
91 // ---------------------------------------------------------------------------
92 // Browser console bridge
93 // ---------------------------------------------------------------------------
94
95 #[wasm_bindgen]
96 extern "C" {
97 #[wasm_bindgen(js_namespace = console, js_name = log)]
98 fn console_log(s: &str);
99
100 #[wasm_bindgen(js_namespace = console, js_name = warn)]
101 fn console_warn(s: &str);
102 }
103
104 macro_rules! clog {
105 ($($t:tt)*) => { console_log(&format!("[onde-wasm] {}", format_args!($($t)*))) }
106 }
107
108 macro_rules! cwarn {
109 ($($t:tt)*) => { console_warn(&format!("[onde-wasm] {}", format_args!($($t)*))) }
110 }
111
112 // ---------------------------------------------------------------------------
113 // Language table
114 //
115 // Copied verbatim from candle-wasm-examples/whisper/src/languages.rs at
116 // rev c3bb5bf. This constant is NOT re-exported by candle-transformers
117 // itself at that revision.
118 // ---------------------------------------------------------------------------
119
120 const LANGUAGES: [(&str, &str); 99] = [
121 ("en", "english"),
122 ("zh", "chinese"),
123 ("de", "german"),
124 ("es", "spanish"),
125 ("ru", "russian"),
126 ("ko", "korean"),
127 ("fr", "french"),
128 ("ja", "japanese"),
129 ("pt", "portuguese"),
130 ("tr", "turkish"),
131 ("pl", "polish"),
132 ("ca", "catalan"),
133 ("nl", "dutch"),
134 ("ar", "arabic"),
135 ("sv", "swedish"),
136 ("it", "italian"),
137 ("id", "indonesian"),
138 ("hi", "hindi"),
139 ("fi", "finnish"),
140 ("vi", "vietnamese"),
141 ("he", "hebrew"),
142 ("uk", "ukrainian"),
143 ("el", "greek"),
144 ("ms", "malay"),
145 ("cs", "czech"),
146 ("ro", "romanian"),
147 ("da", "danish"),
148 ("hu", "hungarian"),
149 ("ta", "tamil"),
150 ("no", "norwegian"),
151 ("th", "thai"),
152 ("ur", "urdu"),
153 ("hr", "croatian"),
154 ("bg", "bulgarian"),
155 ("lt", "lithuanian"),
156 ("la", "latin"),
157 ("mi", "maori"),
158 ("ml", "malayalam"),
159 ("cy", "welsh"),
160 ("sk", "slovak"),
161 ("te", "telugu"),
162 ("fa", "persian"),
163 ("lv", "latvian"),
164 ("bn", "bengali"),
165 ("sr", "serbian"),
166 ("az", "azerbaijani"),
167 ("sl", "slovenian"),
168 ("kn", "kannada"),
169 ("et", "estonian"),
170 ("mk", "macedonian"),
171 ("br", "breton"),
172 ("eu", "basque"),
173 ("is", "icelandic"),
174 ("hy", "armenian"),
175 ("ne", "nepali"),
176 ("mn", "mongolian"),
177 ("bs", "bosnian"),
178 ("kk", "kazakh"),
179 ("sq", "albanian"),
180 ("sw", "swahili"),
181 ("gl", "galician"),
182 ("mr", "marathi"),
183 ("pa", "punjabi"),
184 ("si", "sinhala"),
185 ("km", "khmer"),
186 ("sn", "shona"),
187 ("yo", "yoruba"),
188 ("so", "somali"),
189 ("af", "afrikaans"),
190 ("oc", "occitan"),
191 ("ka", "georgian"),
192 ("be", "belarusian"),
193 ("tg", "tajik"),
194 ("sd", "sindhi"),
195 ("gu", "gujarati"),
196 ("am", "amharic"),
197 ("yi", "yiddish"),
198 ("lo", "lao"),
199 ("uz", "uzbek"),
200 ("fo", "faroese"),
201 ("ht", "haitian creole"),
202 ("ps", "pashto"),
203 ("tk", "turkmen"),
204 ("nn", "nynorsk"),
205 ("mt", "maltese"),
206 ("sa", "sanskrit"),
207 ("lb", "luxembourgish"),
208 ("my", "myanmar"),
209 ("bo", "tibetan"),
210 ("tl", "tagalog"),
211 ("mg", "malagasy"),
212 ("as", "assamese"),
213 ("tt", "tatar"),
214 ("haw", "hawaiian"),
215 ("ln", "lingala"),
216 ("ha", "hausa"),
217 ("ba", "bashkir"),
218 ("jw", "javanese"),
219 ("su", "sundanese"),
220 ];
221
222 // ---------------------------------------------------------------------------
223 // Audio processing
224 //
225 // Ported from candle-wasm-examples/whisper/src/audio.rs (rev c3bb5bf).
226 // Pure-Rust FFT + mel-spectrogram — no C dependencies, WASM-safe.
227 // ---------------------------------------------------------------------------
228
229 trait Float: num_traits::Float + num_traits::FloatConst + num_traits::NumAssign {}
230 impl Float for f32 {}
231 impl Float for f64 {}
232
233 fn fft<T: Float>(inp: &[T]) -> Vec<T> {
234 let n = inp.len();
235 let zero = T::zero();
236 if n == 1 {
237 return vec![inp[0], zero];
238 }
239 if n % 2 == 1 {
240 return dft(inp);
241 }
242 let mut out = vec![zero; n * 2];
243 let mut even = Vec::with_capacity(n / 2);
244 let mut odd = Vec::with_capacity(n / 2);
245 for (i, &v) in inp.iter().enumerate() {
246 if i % 2 == 0 { even.push(v) } else { odd.push(v) }
247 }
248 let even_fft = fft(&even);
249 let odd_fft = fft(&odd);
250 let two_pi = T::PI() + T::PI();
251 let n_t = T::from(n).unwrap();
252 for k in 0..n / 2 {
253 let k_t = T::from(k).unwrap();
254 let theta = two_pi * k_t / n_t;
255 let re = theta.cos();
256 let im = -theta.sin();
257 let re_odd = odd_fft[2 * k];
258 let im_odd = odd_fft[2 * k + 1];
259 out[2 * k] = even_fft[2 * k] + re * re_odd - im * im_odd;
260 out[2 * k + 1] = even_fft[2 * k + 1] + re * im_odd + im * re_odd;
261 out[2 * (k + n / 2)] = even_fft[2 * k] - re * re_odd + im * im_odd;
262 out[2 * (k + n / 2) + 1] = even_fft[2 * k + 1] - re * im_odd - im * re_odd;
263 }
264 out
265 }
266
267 fn dft<T: Float>(inp: &[T]) -> Vec<T> {
268 let zero = T::zero();
269 let n = inp.len();
270 let two_pi = T::PI() + T::PI();
271 let mut out = Vec::with_capacity(2 * n);
272 let n_t = T::from(n).unwrap();
273 for k in 0..n {
274 let k_t = T::from(k).unwrap();
275 let mut re = zero;
276 let mut im = zero;
277 for (j, &v) in inp.iter().enumerate() {
278 let j_t = T::from(j).unwrap();
279 let angle = two_pi * k_t * j_t / n_t;
280 re += v * angle.cos();
281 im -= v * angle.sin();
282 }
283 out.push(re);
284 out.push(im);
285 }
286 out
287 }
288
289 #[allow(clippy::too_many_arguments)]
290 fn log_mel_spectrogram_w<T: Float>(
291 ith: usize,
292 hann: &[T],
293 samples: &[T],
294 filters: &[T],
295 fft_size: usize,
296 fft_step: usize,
297 speed_up: bool,
298 n_len: usize,
299 n_mel: usize,
300 n_threads: usize,
301 ) -> Vec<T> {
302 let n_fft = if speed_up { 1 + fft_size / 4 } else { 1 + fft_size / 2 };
303 let zero = T::zero();
304 let half = T::from(0.5).unwrap();
305 let mut fft_in = vec![zero; fft_size];
306 let mut mel = vec![zero; n_len * n_mel];
307 for i in (ith..n_len).step_by(n_threads) {
308 let offset = i * fft_step;
309 for j in 0..fft_size {
310 fft_in[j] = if offset + j < samples.len() {
311 hann[j] * samples[offset + j]
312 } else {
313 zero
314 };
315 }
316 let mut fft_out: Vec<T> = fft(&fft_in);
317 for j in 0..fft_size {
318 fft_out[j] = fft_out[2 * j] * fft_out[2 * j]
319 + fft_out[2 * j + 1] * fft_out[2 * j + 1];
320 }
321 for j in 1..fft_size / 2 {
322 let v = fft_out[fft_size - j];
323 fft_out[j] += v;
324 }
325 if speed_up {
326 for j in 0..n_fft {
327 fft_out[j] = half * (fft_out[2 * j] + fft_out[2 * j + 1]);
328 }
329 }
330 for j in 0..n_mel {
331 let mut sum = zero;
332 for k in 0..n_fft {
333 sum += fft_out[k] * filters[j * n_fft + k];
334 }
335 mel[j * n_len + i] = T::max(sum, T::from(1e-10).unwrap()).log10();
336 }
337 }
338 mel
339 }
340
341 fn log_mel_spectrogram_<T: Float>(
342 samples: &[T],
343 filters: &[T],
344 fft_size: usize,
345 fft_step: usize,
346 n_mel: usize,
347 speed_up: bool,
348 ) -> Vec<T> {
349 let zero = T::zero();
350 let two_pi = T::PI() + T::PI();
351 let half = T::from(0.5).unwrap();
352 let one = T::from(1.0).unwrap();
353 let four = T::from(4.0).unwrap();
354 let fft_size_t = T::from(fft_size).unwrap();
355 let hann: Vec<T> = (0..fft_size)
356 .map(|i| half * (one - ((two_pi * T::from(i).unwrap()) / fft_size_t).cos()))
357 .collect();
358 let n_len = samples.len() / fft_step;
359 let pad = 100 * m::CHUNK_LENGTH / 2;
360 let n_len = if n_len % pad != 0 {
361 (n_len / pad + 1) * pad
362 } else {
363 n_len
364 };
365 let n_len = n_len + pad;
366 let samples = {
367 let mut padded = samples.to_vec();
368 let to_add = n_len * fft_step - samples.len();
369 padded.extend(std::iter::repeat_n(zero, to_add));
370 padded
371 };
372 let mut mel = log_mel_spectrogram_w(
373 0, &hann, &samples, filters, fft_size, fft_step, speed_up, n_len, n_mel, 1,
374 );
375 let mmax = mel
376 .iter()
377 .max_by(|u, v| u.partial_cmp(v).unwrap_or(std::cmp::Ordering::Greater))
378 .copied()
379 .unwrap_or(zero)
380 - T::from(8).unwrap();
381 for m in mel.iter_mut() {
382 let v = T::max(*m, mmax);
383 *m = v / four + one;
384 }
385 mel
386 }
387
388 fn pcm_to_mel(cfg: &Config, samples: &[f32], filters: &[f32]) -> anyhow::Result<Vec<f32>> {
389 Ok(log_mel_spectrogram_(
390 samples,
391 filters,
392 m::N_FFT,
393 m::HOP_LENGTH,
394 cfg.num_mel_bins,
395 false,
396 ))
397 }
398
399 // ---------------------------------------------------------------------------
400 // Model enum — wraps normal (f32 safetensors) and quantized (GGUF) variants
401 // ---------------------------------------------------------------------------
402
403 enum Model {
404 Normal(m::model::Whisper),
405 Quantized(m::quantized_model::Whisper),
406 }
407
408 impl Model {
409 fn config(&self) -> &Config {
410 match self {
411 Self::Normal(m) => &m.config,
412 Self::Quantized(m) => &m.config,
413 }
414 }
415
416 fn encoder_forward(&mut self, x: &Tensor, flush: bool) -> candle_core::Result<Tensor> {
417 match self {
418 Self::Normal(m) => m.encoder.forward(x, flush),
419 Self::Quantized(m) => m.encoder.forward(x, flush),
420 }
421 }
422
423 fn decoder_forward(
424 &mut self,
425 x: &Tensor,
426 xa: &Tensor,
427 flush: bool,
428 ) -> candle_core::Result<Tensor> {
429 match self {
430 Self::Normal(m) => m.decoder.forward(x, xa, flush),
431 Self::Quantized(m) => m.decoder.forward(x, xa, flush),
432 }
433 }
434
435 fn decoder_final_linear(&self, x: &Tensor) -> candle_core::Result<Tensor> {
436 match self {
437 Self::Normal(m) => m.decoder.final_linear(x),
438 Self::Quantized(m) => m.decoder.final_linear(x),
439 }
440 }
441 }
442
443 // ---------------------------------------------------------------------------
444 // Token helpers
445 // ---------------------------------------------------------------------------
446
447 fn token_id(tokenizer: &Tokenizer, token: &str) -> anyhow::Result<u32> {
448 tokenizer
449 .token_to_id(token)
450 .ok_or_else(|| anyhow!("no token-id for {token}"))
451 }
452
453 fn detect_language(
454 model: &mut Model,
455 tokenizer: &Tokenizer,
456 mel: &Tensor,
457 ) -> anyhow::Result<u32> {
458 clog!("detecting language");
459 let (_bsize, _, seq_len) = mel.dims3()?;
460 let mel = mel.narrow(
461 2,
462 0,
463 usize::min(seq_len, model.config().max_source_positions),
464 )?;
465 let device = mel.device();
466 let language_token_ids = LANGUAGES
467 .iter()
468 .map(|(t, _)| token_id(tokenizer, &format!("<|{t}|>")))
469 .collect::<anyhow::Result<Vec<_>>>()?;
470 let sot = token_id(tokenizer, m::SOT_TOKEN)?;
471 let audio_features = model.encoder_forward(&mel, true)?;
472 let tokens = Tensor::new(&[[sot]], device)?;
473 let language_token_ids_t = Tensor::new(language_token_ids.as_slice(), device)?;
474 let ys = model.decoder_forward(&tokens, &audio_features, true)?;
475 let logits = model.decoder_final_linear(&ys.i(..1)?)?.i(0)?.i(0)?;
476 let logits = logits.index_select(&language_token_ids_t, 0)?;
477 let probs = softmax(&logits, D::Minus1)?.to_vec1::<f32>()?;
478 let mut ranked: Vec<_> = LANGUAGES.iter().zip(probs.iter()).collect();
479 ranked.sort_by(|(_, p1), (_, p2)| p2.total_cmp(p1));
480 let best_code = ranked[0].0 .0;
481 let token_str = format!("<|{best_code}|>");
482 let language_id = token_id(tokenizer, &token_str)?;
483 clog!("detected language: {best_code} ({token_str})");
484 Ok(language_id)
485 }
486
487 // ---------------------------------------------------------------------------
488 // Internal decoding state
489 // ---------------------------------------------------------------------------
490
491 #[derive(Clone)]
492 struct DecodingResult {
493 tokens: Vec<u32>,
494 text: String,
495 avg_logprob: f64,
496 no_speech_prob: f64,
497 compression_ratio: f64,
498 }
499
500 struct Decoder {
501 model: Model,
502 rng: StdRng,
503 tokenizer: Tokenizer,
504 mel_filters: Vec<f32>,
505 language: Option<String>,
506 is_multilingual: bool,
507 timestamps: bool,
508 task_tokens: Vec<u32>, // [task_token] or [task_token, no_timestamps_token]
509 suppress_tokens: Tensor,
510 sot_token: u32,
511 eot_token: u32,
512 no_speech_token: u32,
513 no_timestamps_token: u32,
514 }
515
516 impl Decoder {
517 #[allow(clippy::too_many_arguments)]
518 fn new(
519 model: Model,
520 tokenizer: Tokenizer,
521 mel_filters: Vec<f32>,
522 device: &Device,
523 task: Option<&str>,
524 language: Option<String>,
525 is_multilingual: bool,
526 timestamps: bool,
527 ) -> anyhow::Result<Self> {
528 let suppress_tokens_vec: Vec<f32> = (0..model.config().vocab_size as u32)
529 .map(|i| {
530 if model.config().suppress_tokens.contains(&i) {
531 f32::NEG_INFINITY
532 } else {
533 0f32
534 }
535 })
536 .collect();
537 let suppress_tokens = Tensor::new(suppress_tokens_vec.as_slice(), device)?;
538
539 let sot_token = token_id(&tokenizer, m::SOT_TOKEN)?;
540 let transcribe_token = token_id(&tokenizer, m::TRANSCRIBE_TOKEN)?;
541 let translate_token = token_id(&tokenizer, m::TRANSLATE_TOKEN)?;
542 let eot_token = token_id(&tokenizer, m::EOT_TOKEN)?;
543 let no_timestamps_token = token_id(&tokenizer, m::NO_TIMESTAMPS_TOKEN)?;
544 let no_speech_token = m::NO_SPEECH_TOKENS
545 .iter()
546 .find_map(|t| token_id(&tokenizer, t).ok())
547 .ok_or_else(|| anyhow!("no no-speech token found"))?;
548
549 let task_token = match task {
550 Some("translate") => translate_token,
551 _ => transcribe_token,
552 };
553
554 let mut task_tokens = vec![task_token];
555 if !timestamps {
556 task_tokens.push(no_timestamps_token);
557 }
558
559 Ok(Self {
560 model,
561 rng: StdRng::seed_from_u64(299_792_458),
562 tokenizer,
563 mel_filters,
564 language,
565 is_multilingual,
566 timestamps,
567 task_tokens,
568 suppress_tokens,
569 sot_token,
570 eot_token,
571 no_speech_token,
572 no_timestamps_token,
573 })
574 }
575
576 fn decode(&mut self, mel: &Tensor, temperature: f64) -> anyhow::Result<DecodingResult> {
577 let language_token = match (self.is_multilingual, &self.language) {
578 (true, None) => Some(detect_language(&mut self.model, &self.tokenizer, mel)?),
579 (false, None) => None,
580 (true, Some(lang)) => {
581 let tok = format!("<|{lang}|>");
582 Some(token_id(&self.tokenizer, &tok).map_err(|_| {
583 anyhow!("language '{lang}' is not supported by this model")
584 })?)
585 }
586 (false, Some(_)) => {
587 anyhow::bail!("language cannot be set for non-multilingual models")
588 }
589 };
590
591 let audio_features = self.model.encoder_forward(mel, true)?;
592 let sample_len = self.model.config().max_target_positions / 2;
593
594 // Build the initial prompt token sequence:
595 // [SOT, (language), task, (no_timestamps)]
596 let mut tokens = vec![self.sot_token];
597 if let Some(lang_tok) = language_token {
598 tokens.push(lang_tok);
599 }
600 tokens.extend_from_slice(&self.task_tokens);
601
602 let mut sum_logprob = 0f64;
603 let mut no_speech_prob = f64::NAN;
604
605 for i in 0..sample_len {
606 let tokens_t = Tensor::new(tokens.as_slice(), mel.device())?.unsqueeze(0)?;
607 let ys = self.model.decoder_forward(&tokens_t, &audio_features, i == 0)?;
608
609 // Extract no-speech probability from the first decoder step.
610 if i == 0 {
611 let logits = self.model.decoder_final_linear(&ys.i(..1)?)?.i(0)?.i(0)?;
612 no_speech_prob = softmax(&logits, 0)?
613 .i(self.no_speech_token as usize)?
614 .to_scalar::<f32>()? as f64;
615 }
616
617 let (_, seq_len, _) = ys.dims3()?;
618 let logits = self
619 .model
620 .decoder_final_linear(&ys.i((..1, seq_len - 1..))?)?
621 .i(0)?
622 .i(0)?;
623 let logits = logits.broadcast_add(&self.suppress_tokens)?;
624
625 let next_token = if temperature > 0.0 {
626 let prs = softmax(&(&logits / temperature)?, 0)?;
627 let logits_v: Vec<f32> = prs.to_vec1()?;
628 let distr = rand::distr::weighted::WeightedIndex::new(&logits_v)?;
629 distr.sample(&mut self.rng) as u32
630 } else {
631 let logits_v: Vec<f32> = logits.to_vec1()?;
632 logits_v
633 .iter()
634 .enumerate()
635 .max_by(|(_, a), (_, b)| a.total_cmp(b))
636 .map(|(i, _)| i as u32)
637 .unwrap_or(self.eot_token)
638 };
639
640 if next_token == self.eot_token
641 || tokens.len() > self.model.config().max_target_positions
642 {
643 break;
644 }
645
646 let prob = softmax(&logits, D::Minus1)?
647 .i(next_token as usize)?
648 .to_scalar::<f32>()? as f64;
649 sum_logprob += prob.ln();
650 tokens.push(next_token);
651 }
652
653 let text = self
654 .tokenizer
655 .decode(&tokens, true)
656 .map_err(|e| anyhow!("{e}"))?;
657 let avg_logprob = sum_logprob / tokens.len().max(1) as f64;
658
659 Ok(DecodingResult {
660 tokens,
661 text,
662 avg_logprob,
663 no_speech_prob,
664 compression_ratio: f64::NAN,
665 })
666 }
667
668 fn decode_with_fallback(&mut self, segment: &Tensor) -> anyhow::Result<DecodingResult> {
669 for (i, &t) in m::TEMPERATURES.iter().enumerate() {
670 let result = self.decode(segment, t);
671 if i == m::TEMPERATURES.len() - 1 {
672 return result;
673 }
674 match result {
675 Ok(dr) => {
676 let needs_fallback = dr.compression_ratio > m::COMPRESSION_RATIO_THRESHOLD
677 || dr.avg_logprob < m::LOGPROB_THRESHOLD;
678 if !needs_fallback || dr.no_speech_prob > m::NO_SPEECH_THRESHOLD {
679 return Ok(dr);
680 }
681 }
682 Err(e) => cwarn!("decode at temperature {t} failed: {e}"),
683 }
684 }
685 unreachable!()
686 }
687
688 fn run(&mut self, mel: &Tensor) -> anyhow::Result<Vec<Segment>> {
689 let (_, _, content_frames) = mel.dims3()?;
690 let mut seek = 0;
691 let mut segments = vec![];
692 while seek < content_frames {
693 let time_offset = (seek * m::HOP_LENGTH) as f64 / m::SAMPLE_RATE as f64;
694 let segment_size = usize::min(content_frames - seek, m::N_FRAMES);
695 let mel_segment = mel.narrow(2, seek, segment_size)?;
696 let segment_duration =
697 (segment_size * m::HOP_LENGTH) as f64 / m::SAMPLE_RATE as f64;
698 let dr = self.decode_with_fallback(&mel_segment)?;
699 seek += segment_size;
700 if dr.no_speech_prob > m::NO_SPEECH_THRESHOLD
701 && dr.avg_logprob < m::LOGPROB_THRESHOLD
702 {
703 clog!("skipping silent segment at seek={seek}");
704 continue;
705 }
706 clog!("segment [{:.1}s] {}", time_offset, dr.text.trim());
707 segments.push(Segment {
708 start: time_offset,
709 end: time_offset + segment_duration,
710 dr,
711 });
712 }
713 Ok(segments)
714 }
715 }
716
717 // ---------------------------------------------------------------------------
718 // Segment — intermediate result per 30-second chunk
719 // ---------------------------------------------------------------------------
720
721 struct Segment {
722 start: f64,
723 end: f64,
724 dr: DecodingResult,
725 }
726
727 // ---------------------------------------------------------------------------
728 // Public JSON output types
729 // ---------------------------------------------------------------------------
730
731 #[derive(Serialize, Deserialize)]
732 struct OndeSegment {
733 start: f64,
734 end: f64,
735 text: String,
736 }
737
738 #[derive(Serialize, Deserialize)]
739 struct OndeResult {
740 text: String,
741 segments: Vec<OndeSegment>,
742 }
743
744 // ---------------------------------------------------------------------------
745 // WhisperDecoder — the public wasm-bindgen export
746 // ---------------------------------------------------------------------------
747
748 /// Whisper decoder bound to a pre-loaded set of model weights.
749 ///
750 /// Construct once, then call `decode` for each audio clip.
751 #[wasm_bindgen]
752 pub struct WhisperDecoder {
753 decoder: Decoder,
754 }
755
756 #[wasm_bindgen]
757 impl WhisperDecoder {
758 /// Construct a new `WhisperDecoder` from raw byte buffers.
759 ///
760 /// All four byte slices must stay alive for the duration of this call;
761 /// they are consumed (weights are moved into the model, everything else
762 /// is copied into small owned structures) and the originals are no
763 /// longer needed once the constructor returns.
764 ///
765 /// # Arguments
766 ///
767 /// - `weights` — `model.safetensors` bytes, OR `.gguf` bytes when `quantized = true`
768 /// - `tokenizer` — `tokenizer.json` bytes
769 /// - `mel_filters` — `mel_filters.safetensors` bytes
770 /// - `config` — `config.json` bytes
771 /// - `quantized` — `true` when `weights` is a `.gguf` file
772 /// - `is_multilingual` — `false` for `*.en` models
773 /// - `timestamps` — include per-segment timestamps in the output
774 /// - `task` — `"transcribe"` | `"translate"` | `null`
775 /// - `language` — BCP-47 code (`"id"`, `"en"`, …) or `null` for auto-detect
776 #[wasm_bindgen(constructor)]
777 pub fn new(
778 weights: &[u8],
779 tokenizer: &[u8],
780 mel_filters: &[u8],
781 config: &[u8],
782 quantized: bool,
783 is_multilingual: bool,
784 timestamps: bool,
785 task: Option<String>,
786 language: Option<String>,
787 ) -> Result<WhisperDecoder, JsError> {
788 set_panic_hook();
789 Self::build(
790 weights,
791 tokenizer,
792 mel_filters,
793 config,
794 quantized,
795 is_multilingual,
796 timestamps,
797 task.as_deref(),
798 language,
799 )
800 .map_err(|e| JsError::new(&e.to_string()))
801 }
802
803 fn build(
804 weights: &[u8],
805 tokenizer_bytes: &[u8],
806 mel_filters_bytes: &[u8],
807 config_bytes: &[u8],
808 quantized: bool,
809 is_multilingual: bool,
810 timestamps: bool,
811 task: Option<&str>,
812 language: Option<String>,
813 ) -> anyhow::Result<WhisperDecoder> {
814 let device = Device::Cpu;
815
816 // Tokenizer
817 let tokenizer = Tokenizer::from_bytes(tokenizer_bytes)
818 .map_err(|e| anyhow!("tokenizer load failed: {e}"))?;
819 clog!("tokenizer loaded");
820
821 // Mel filters — shape (n_mels, n_fft/2+1) stored as "mel_80" tensor
822 let mel_st = safetensors::SafeTensors::deserialize(mel_filters_bytes)
823 .map_err(|e| anyhow!("mel_filters load failed: {e}"))?;
824 let mel_tensor = mel_st
825 .tensor("mel_80")
826 .map_err(|e| anyhow!("mel_80 tensor missing: {e}"))?
827 .load(&device)?;
828 let mel_filters = mel_tensor.flatten_all()?.to_vec1::<f32>()?;
829 clog!("mel filters loaded: {:?}", mel_tensor.shape());
830
831 // Config
832 let config: Config = serde_json::from_slice(config_bytes)
833 .map_err(|e| anyhow!("config.json parse failed: {e}"))?;
834
835 // Model weights
836 let model = if quantized {
837 let vb =
838 candle_transformers::quantized_var_builder::VarBuilder::from_gguf_buffer(
839 weights,
840 &device,
841 )?;
842 Model::Quantized(m::quantized_model::Whisper::load(&vb, config)?)
843 } else {
844 let vb = VarBuilder::from_buffered_safetensors(
845 weights.to_vec(),
846 m::DTYPE,
847 &device,
848 )?;
849 Model::Normal(m::model::Whisper::load(&vb, config)?)
850 };
851 clog!("model loaded (quantized={quantized})");
852
853 let decoder = Decoder::new(
854 model,
855 tokenizer,
856 mel_filters,
857 &device,
858 task,
859 language,
860 is_multilingual,
861 timestamps,
862 )?;
863
864 Ok(WhisperDecoder { decoder })
865 }
866
867 /// Decode a 16 kHz mono PCM WAV file.
868 ///
869 /// Returns a JSON string with shape `{ text, segments: [{start, end, text}] }`.
870 pub fn decode(&mut self, wav_bytes: &[u8]) -> Result<String, JsError> {
871 self.decode_inner(wav_bytes)
872 .map_err(|e| JsError::new(&e.to_string()))
873 }
874
875 fn decode_inner(&mut self, wav_bytes: &[u8]) -> anyhow::Result<String> {
876 let device = Device::Cpu;
877
878 // Decode WAV
879 let mut cursor = std::io::Cursor::new(wav_bytes);
880 let reader = hound::WavReader::new(&mut cursor)?;
881 let spec = reader.spec();
882 clog!("wav: {spec:?}");
883 if spec.sample_rate != m::SAMPLE_RATE as u32 {
884 anyhow::bail!(
885 "expected {}Hz WAV, got {}Hz — please resample before passing to decode()",
886 m::SAMPLE_RATE,
887 spec.sample_rate
888 );
889 }
890
891 // Read samples, take only the first channel if stereo
892 let raw: Vec<_> = reader.into_samples::<i16>().collect::<hound::Result<Vec<_>>>()?;
893 let mono: Vec<f32> = raw
894 .chunks(spec.channels as usize)
895 .map(|ch| ch[0] as f32 / 32768.0)
896 .collect();
897 clog!("pcm samples: {}", mono.len());
898
899 // Build mel spectrogram tensor (shape: [1, n_mels, n_frames])
900 let cfg = self.decoder.model.config();
901 let mel = pcm_to_mel(cfg, &mono, &self.decoder.mel_filters)?;
902 let mel_len = mel.len();
903 let n_mels = cfg.num_mel_bins;
904 let mel_t = Tensor::from_vec(mel, (1, n_mels, mel_len / n_mels), &device)?;
905 clog!("mel tensor: {:?}", mel_t.dims());
906
907 // Run the sliding-window decoder
908 let segments = self.decoder.run(&mel_t)?;
909
910 // Build output
911 let mut full_text = String::new();
912 let mut out_segments = Vec::with_capacity(segments.len());
913 for seg in &segments {
914 if !full_text.is_empty() && !seg.dr.text.starts_with(' ') {
915 full_text.push(' ');
916 }
917 full_text.push_str(seg.dr.text.trim());
918 out_segments.push(OndeSegment {
919 start: seg.start,
920 end: seg.end,
921 text: seg.dr.text.trim().to_string(),
922 });
923 }
924
925 let result = OndeResult {
926 text: full_text,
927 segments: out_segments,
928 };
929 Ok(serde_json::to_string(&result)?)
930 }
931 }