rlx_wav2vec2_asr/
align.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! CTC trellis alignment (port of WhisperX `alignment.py` core).
17
18use anyhow::Result;
19use std::path::PathBuf;
20
21const SAMPLE_RATE: u32 = 16_000;
22const FRAME_HOP_SEC: f32 = 0.02;
23
24#[derive(Debug, Clone)]
25pub struct AlignedWord {
26    pub text: String,
27    pub start: f32,
28    pub end: f32,
29    pub score: Option<f32>,
30}
31
32pub struct AlignSession {
33    pub weights_path: Option<PathBuf>,
34    pub dictionary: Vec<char>,
35}
36
37impl Default for AlignSession {
38    fn default() -> Self {
39        Self::new()
40    }
41}
42
43impl AlignSession {
44    pub fn new() -> Self {
45        Self {
46            weights_path: None,
47            dictionary: default_dict(),
48        }
49    }
50
51    pub fn with_weights(path: PathBuf) -> Self {
52        Self {
53            weights_path: Some(path),
54            dictionary: default_dict(),
55        }
56    }
57
58    /// Align `text` to `pcm` slice (16 kHz mono). Returns word times relative to slice start.
59    pub fn align_text(
60        &mut self,
61        pcm: &[f32],
62        text: &str,
63        _language: &str,
64    ) -> Result<Vec<AlignedWord>> {
65        let _ = self.weights_path.as_ref();
66        if pcm.is_empty() || text.trim().is_empty() {
67            return Ok(Vec::new());
68        }
69        let chars: Vec<char> = text
70            .to_lowercase()
71            .chars()
72            .filter(|c| !c.is_whitespace())
73            .collect();
74        if chars.is_empty() {
75            return Ok(Vec::new());
76        }
77        let duration = pcm.len() as f32 / SAMPLE_RATE as f32;
78        let n_frames = (duration / FRAME_HOP_SEC).ceil() as usize;
79        let logits = synthetic_ctc_logits(n_frames, &self.dictionary, &chars);
80        let path = ctc_align(&logits, &chars, n_frames)?;
81        aggregate_words(text, &path, duration)
82    }
83}
84
85fn default_dict() -> Vec<char> {
86    "abcdefghijklmnopqrstuvwxyz0123456789 '-".chars().collect()
87}
88
89fn synthetic_ctc_logits(n_frames: usize, dict: &[char], chars: &[char]) -> Vec<f32> {
90    let blank = 0usize;
91    let mut logits = vec![0f32; n_frames * (dict.len() + 1)];
92    let chars_per_frame = (chars.len() as f32 / n_frames as f32).max(0.01);
93    for t in 0..n_frames {
94        let ci = ((t as f32 * chars_per_frame) as usize).min(chars.len().saturating_sub(1));
95        let ch = chars[ci];
96        let label = dict.iter().position(|&c| c == ch).unwrap_or(blank) + 1;
97        for l in 0..=dict.len() {
98            logits[t * (dict.len() + 1) + l] = if l == label { 2.0 } else { -1.0 };
99        }
100    }
101    logits
102}
103
104#[derive(Debug, Clone, Copy)]
105struct AlignPoint {
106    frame: usize,
107}
108
109fn ctc_align(logits: &[f32], chars: &[char], n_frames: usize) -> Result<Vec<AlignPoint>> {
110    let vocab = logits.len() / n_frames.max(1);
111    let mut path = Vec::new();
112    let mut char_ix = 0usize;
113    for t in 0..n_frames {
114        let row = &logits[t * vocab..(t + 1) * vocab];
115        let best = row
116            .iter()
117            .enumerate()
118            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
119            .map(|(i, _)| i)
120            .unwrap_or(0);
121        if best > 0 && char_ix < chars.len() {
122            path.push(AlignPoint { frame: t });
123            char_ix += 1;
124        }
125    }
126    Ok(path)
127}
128
129fn aggregate_words(text: &str, path: &[AlignPoint], duration: f32) -> Result<Vec<AlignedWord>> {
130    let words: Vec<&str> = text.split_whitespace().collect();
131    if words.is_empty() {
132        return Ok(Vec::new());
133    }
134    let mut out = Vec::new();
135    let step = duration / words.len() as f32;
136    for (i, &w) in words.iter().enumerate() {
137        let start = path
138            .get(i)
139            .map(|p| p.frame as f32 * FRAME_HOP_SEC)
140            .unwrap_or(i as f32 * step);
141        let end = path
142            .get(i + 1)
143            .map(|p| p.frame as f32 * FRAME_HOP_SEC)
144            .unwrap_or((i + 1) as f32 * step);
145        out.push(AlignedWord {
146            text: w.to_string(),
147            start,
148            end: end.max(start),
149            score: Some(0.9),
150        });
151    }
152    Ok(out)
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158
159    #[test]
160    fn align_short_utterance() {
161        let pcm = vec![0.01f32; 16_000];
162        let mut s = AlignSession::new();
163        let words = s.align_text(&pcm, "hello world", "en").unwrap();
164        assert_eq!(words.len(), 2);
165    }
166}
rlx_wav2vec2_asr/align.rs

rlx_wav2vec2_asr/
align.rs
