wavo 0.3.1

Speech-to-text in pure Rust on the GPU (Metal, Vulkan, DX12): Parakeet, GigaAM, Whisper and Moonshine, 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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151

// Copyright (c) vladkens | MIT License | https://github.com/vladkens/wavo
//! Greedy decoding on the CPU. RNN-T: embedding + one LSTM layer as the predictor, joint network
//! `out(relu(enc + pred(h)))`; the encoder side of the joint is projected on the GPU. CTC: argmax
//! per frame over logits computed on the GPU, then collapse.

use crate::cpu::{matmul, rows4};
use crate::error::Result;
use crate::gguf::Gguf;

/// At most this many tokens per encoder frame.
const MAX_SYMBOLS: usize = 10;
/// Frames scored at once while the predictor output stays fixed.
const SPAN: usize = 2;

pub struct Decoder {
  hidden: usize,
  joint: usize,
  blank: usize,
  /// `[vocab][hidden]`
  embed: Vec<f32>,
  /// `[4·hidden][hidden]`, gates i, f, g, o.
  wx: Vec<f32>,
  wh: Vec<f32>,
  /// `bias_ih + bias_hh`
  bias: Vec<f32>,
  /// `[joint][hidden]`, rows padded to a multiple of 4.
  pred_w: Vec<f32>,
  pred_b: Vec<f32>,
  /// `[vocab][joint]`, rows padded to a multiple of 4.
  out_w: Vec<f32>,
  out_b: Vec<f32>,
}

struct State {
  h: Vec<f32>,
  c: Vec<f32>,
}

impl Decoder {
  pub fn new(g: &Gguf, hidden: usize, joint: usize, vocab: usize, blank: usize) -> Result<Self> {
    let t = |name: &str, dims: &[usize]| g.tensor(name, dims).and_then(|t| t.to_f32());
    Ok(Self {
      hidden,
      joint,
      blank,
      embed: t("pred.embed.weight", &[hidden, vocab])?,
      wx: t("pred.lstm.0.Wx", &[hidden, 4 * hidden])?,
      wh: t("pred.lstm.0.Wh", &[hidden, 4 * hidden])?,
      bias: t("pred.lstm.0.bias", &[4 * hidden])?,
      pred_w: rows4(t("joint.pred.weight", &[hidden, joint])?, hidden),
      pred_b: t("joint.pred.bias", &[joint])?,
      out_w: rows4(t("joint.out.weight", &[joint, vocab])?, joint),
      out_b: t("joint.out.bias", &[vocab])?,
    })
  }

  /// Greedy search over `enc` `[frames][joint]` (encoder projection with bias). Returns
  /// `(token, frame)` pairs. The first predictor input is zero; the predictor state advances
  /// only when a token is emitted, so frames are scored `SPAN` at a time against one predictor
  /// output, and the frames after a token are scored again.
  pub fn decode(&self, enc: &[f32]) -> Vec<(u32, u32)> {
    let (n, j) = (self.hidden, self.joint);
    let zero = State { h: vec![0.0; n], c: vec![0.0; n] };
    let mut next = self.lstm(&vec![0.0; n], &zero);
    let mut pred = self.predict(&next.h);
    let mut out = Vec::new();
    let frames = enc.len() / j;
    let mut t = 0;
    while t < frames {
      let span = self.joint(&enc[t * j..(t + SPAN).min(frames) * j], &pred);
      let Some(skip) = span.iter().position(|&token| token != self.blank) else {
        t += span.len();
        continue;
      };
      t += skip;
      let mut token = span[skip];
      for _ in 0..MAX_SYMBOLS {
        out.push((token as u32, t as u32));
        next = self.lstm(&self.embed[token * n..][..n], &next);
        pred = self.predict(&next.h);
        token = self.joint(&enc[t * j..][..j], &pred)[0];
        if token == self.blank {
          break;
        }
      }
      t += 1;
    }
    out
  }

  fn lstm(&self, x: &[f32], s: &State) -> State {
    let n = self.hidden;
    let (gx, gh) = (matmul(&self.wx, x, n), matmul(&self.wh, &s.h, n));
    let gate = |i: usize| self.bias[i] + gx[i] + gh[i];
    let sigmoid = |v: f32| 1.0 / (1.0 + (-v).exp());
    let mut next = State { h: vec![0.0; n], c: vec![0.0; n] };
    for j in 0..n {
      let (i, f, g, o) = (gate(j), gate(n + j), gate(2 * n + j), gate(3 * n + j));
      next.c[j] = sigmoid(f) * s.c[j] + sigmoid(i) * g.tanh();
      next.h[j] = sigmoid(o) * next.c[j].tanh();
    }
    next
  }

  fn predict(&self, h: &[f32]) -> Vec<f32> {
    let d = matmul(&self.pred_w, h, self.hidden);
    d.iter().zip(&self.pred_b).map(|(d, b)| b + d).collect()
  }

  /// Argmax of the joint logits for each frame in `enc`; ties go to the lowest id.
  fn joint(&self, enc: &[f32], pred: &[f32]) -> Vec<usize> {
    let z: Vec<f32> = (enc.chunks_exact(self.joint))
      .flat_map(|e| e.iter().zip(pred).map(|(e, p)| (e + p).max(0.0)))
      .collect();
    let logits = matmul(&self.out_w, &z, self.joint);
    (logits.chunks_exact(self.out_w.len() / self.joint))
      .map(|l| {
        let mut best = (0, f32::NEG_INFINITY);
        for (v, (d, b)) in l.iter().zip(&self.out_b).enumerate() {
          let logit = d + b;
          if logit > best.1 {
            best = (v, logit);
          }
        }
        best.0
      })
      .collect()
  }
}

/// Greedy CTC over `logits` `[frames][classes]`: per-frame argmax (ties go to the lowest id; the
/// reference takes it after a log-softmax, which only shifts each row), repeats of the previous
/// frame's label dropped, then blanks. Returns `(token, frame)` pairs; a token's frame is the
/// first of its run.
pub fn ctc(logits: &[f32], classes: usize, blank: usize) -> Vec<(u32, u32)> {
  let mut out = Vec::new();
  let mut prev = None;
  for (t, row) in logits.chunks_exact(classes).enumerate() {
    let mut best = (0, row[0]);
    for (v, &x) in row.iter().enumerate().skip(1) {
      if x > best.1 {
        best = (v, x);
      }
    }
    if prev != Some(best.0) && best.0 != blank {
      out.push((best.0 as u32, t as u32));
    }
    prev = Some(best.0);
  }
  out
}