wavo 0.1.0

Speech-to-text in pure Rust on the GPU (Metal, Vulkan): Parakeet, GigaAM and Whisper, as a library and an ollama-like CLI.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79

// Copyright (c) vladkens | MIT License | https://github.com/vladkens/wavo
//! Speech-to-text for the GGUF ASR models published by transcribe.cpp. The encoder runs on the GPU
//! through wgpu; the frontend runs on the CPU, and the decoder on the CPU (or for Whisper on the
//! GPU).

mod conformer;
mod cpu;
mod error;
mod fft;
mod gguf;
mod gigaam;
mod gpu;
mod parakeet;
mod whisper;

use std::path::Path;

use error::bail;
pub use error::{Error, Result};

pub struct Model(Family);

enum Family {
  Gigaam(Box<gigaam::Gigaam>),
  Parakeet(Box<parakeet::Parakeet>),
  Whisper(Box<whisper::Whisper>),
}

#[derive(Debug, Clone, Default, PartialEq)]
pub struct Transcript {
  pub text: String,
  pub tokens: Vec<Token>,
}

#[derive(Debug, Clone, PartialEq)]
pub struct Token {
  pub id: u32,
  /// Vocabulary piece: a SentencePiece piece (`▁` marks a word start), a character, or for
  /// Whisper the text the token completes (its bytes may end inside a character).
  pub piece: String,
  /// Start time in milliseconds: the encoder frame where the token was emitted times the frame
  /// length (40 ms for GigaAM, 80 ms for Parakeet); for Whisper the start of its segment.
  pub start_ms: u32,
}

impl Model {
  pub fn load(path: impl AsRef<Path>) -> Result<Self> {
    let g = gguf::Gguf::open(path.as_ref())?;
    Ok(Self(match g.get::<&str>("general.architecture")? {
      "gigaam" => Family::Gigaam(Box::new(gigaam::Gigaam::load(&g)?)),
      "parakeet" => Family::Parakeet(Box::new(parakeet::Parakeet::load(&g)?)),
      "whisper" => Family::Whisper(Box::new(whisper::Whisper::load(&g)?)),
      arch => {
        bail!("general.architecture is {arch:?}, only gigaam, parakeet and whisper are supported")
      }
    }))
  }

  /// Transcribes 16 kHz mono PCM samples in [-1, 1].
  pub fn transcribe(&self, pcm: &[f32]) -> Result<Transcript> {
    match &self.0 {
      Family::Gigaam(m) => m.transcribe(pcm),
      Family::Parakeet(m) => m.transcribe(pcm),
      Family::Whisper(m) => m.transcribe(pcm),
    }
  }

  /// The longest audio, in milliseconds, the model was trained on. Longer audio still runs, but
  /// accuracy degrades, so split it into segments up to this length. `None`: no such window.
  pub fn max_audio_ms(&self) -> Option<u32> {
    match &self.0 {
      // No GGUF key holds GigaAM's window; transcribe.cpp hardcodes the same 25 s.
      Family::Gigaam(_) => Some(25_000),
      // Its window; `transcribe` still takes any length in 30 s windows, as transcribe.cpp does.
      Family::Whisper(_) => Some(30_000),
      Family::Parakeet(_) => None,
    }
  }
}