rlx-whisper 0.2.11

OpenAI Whisper ASR for RLX
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

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! Whisper configuration — HuggingFace `config.json` / OpenAI dims.

use serde::Deserialize;
use std::path::Path;

/// Whisper model dimensions (HF field names; see OpenAI `model.py` comments in candle).
#[derive(Debug, Clone, Deserialize)]
pub struct WhisperConfig {
    pub num_mel_bins: usize,
    pub max_source_positions: usize,
    pub d_model: usize,
    pub encoder_attention_heads: usize,
    pub encoder_layers: usize,
    pub vocab_size: usize,
    pub max_target_positions: usize,
    pub decoder_attention_heads: usize,
    pub decoder_layers: usize,
    #[serde(default)]
    pub suppress_tokens: Vec<u32>,
    /// Suppressed on the first generated token only (HF `begin_suppress_tokens`).
    #[serde(default)]
    pub begin_suppress_tokens: Vec<u32>,
}

impl WhisperConfig {
    pub fn from_file(path: &Path) -> anyhow::Result<Self> {
        let data = std::fs::read_to_string(path)?;
        Ok(serde_json::from_str(&data)?)
    }

    pub fn head_dim(&self) -> usize {
        self.d_model / self.encoder_attention_heads
    }

    pub fn decoder_head_dim(&self) -> usize {
        self.d_model / self.decoder_attention_heads
    }

    /// Encoder sequence length after the two conv layers (stride-2 on the second).
    pub fn encoder_seq_len(&self, mel_frames: usize) -> usize {
        let after_conv1 = mel_frames;
        let pad = 1usize;
        let k = 3usize;
        let stride2 = 2usize;
        (after_conv1 + 2 * pad - k) / stride2 + 1
    }

    /// Tiny config for compile tests (no real weights required).
    pub fn tiny_synthetic() -> Self {
        Self {
            num_mel_bins: 4,
            max_source_positions: 16,
            d_model: 8,
            encoder_attention_heads: 2,
            encoder_layers: 1,
            vocab_size: 32,
            max_target_positions: 16,
            decoder_attention_heads: 2,
            decoder_layers: 1,
            suppress_tokens: vec![],
            begin_suppress_tokens: vec![],
        }
    }

    /// `openai/whisper-tiny` dimensions.
    pub fn tiny() -> Self {
        Self {
            num_mel_bins: 80,
            max_source_positions: 1500,
            d_model: 384,
            encoder_attention_heads: 6,
            encoder_layers: 4,
            vocab_size: 51865,
            max_target_positions: 448,
            decoder_attention_heads: 6,
            decoder_layers: 4,
            suppress_tokens: vec![],
            begin_suppress_tokens: vec![220, 50257],
        }
    }
}