whisper_captions/
lib.rs

1//! Turn whisper.cpp output into clean, caption-ready segments.
2//!
3//! whisper.cpp can emit a JSON file (`-oj` / `-ojf`) or plain stdout lines like
4//! `[00:00:01.000 --> 00:00:03.500]  Hello there`. This crate normalizes both
5//! into [`Segment`]s with millisecond timestamps and, when token-level data is
6//! present, whole-word [`Word`] timings.
7//!
8//! whisper.cpp's special tokens (`[_BEG_]`, `[_TT_150]`, `<|endoftext|>`, …)
9//! are dropped, and [`to_srt`] / [`to_vtt`] turn segments into subtitle files.
10//!
11//! No whisper dependency: it works on output files from any whisper.cpp
12//! version, and [`merge_bpe_tokens`] also accepts tokens from `whisper-rs`.
13//!
14//! ```
15//! let stdout = "[00:00:01.000 --> 00:00:02.500]  Hello   world\n";
16//! let segments = whisper_captions::parse_whisper_stdout(stdout);
17//! assert_eq!(segments[0].text, "Hello world");
18//! assert_eq!(segments[0].start_ms, 1000);
19//! assert_eq!(segments[0].duration_ms, 1500);
20//!
21//! let srt = whisper_captions::to_srt(&segments);
22//! assert!(srt.starts_with("1\n00:00:01,000 --> 00:00:02,500\nHello world\n"));
23//! ```
24
25mod js;
26
27use std::sync::OnceLock;
28
29use regex::Regex;
30use serde::{Deserialize, Serialize};
31use serde_json::Value;
32
33use js::{
34    get, get_path, is_falsy, is_js_space, js_round, js_string, js_string_or_empty, nullish,
35    sanitize, trim_start_js,
36};
37
38/// A whole word with its timing, reassembled from whisper's BPE tokens.
39#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
40#[serde(rename_all = "camelCase")]
41pub struct Word {
42    pub text: String,
43    pub start_ms: i64,
44    pub end_ms: i64,
45}
46
47/// One caption-ready transcript segment.
48#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
49#[serde(rename_all = "camelCase")]
50pub struct Segment {
51    pub start_ms: i64,
52    pub end_ms: i64,
53    pub duration_ms: i64,
54    pub text: String,
55    /// Empty when the source had no usable word/token timings.
56    #[serde(default, skip_serializing_if = "Vec::is_empty")]
57    pub words: Vec<Word>,
58}
59
60/// A raw whisper token before BPE merging. The text is kept untouched so the
61/// leading-space "starts a new word" signal survives until merging.
62#[derive(Debug, Clone, PartialEq)]
63pub struct RawToken {
64    pub raw_text: String,
65    pub start_ms: f64,
66    pub end_ms: f64,
67}
68
69/// Parse a whisper timestamp into milliseconds.
70///
71/// * numbers above 100000 are treated as milliseconds, smaller ones as seconds
72/// * strings accept `HH:MM:SS.mmm`, `MM:SS,mmm`, `H:MM:SS` and similar
73pub fn parse_timestamp_ms(value: &Value) -> Option<f64> {
74    match value {
75        Value::Number(n) => {
76            let v = n.as_f64()?;
77            Some(if v > 100000.0 {
78                js_round(v)
79            } else {
80                js_round(v * 1000.0)
81            })
82        }
83        Value::String(s) => parse_timestamp_str(s),
84        _ => None,
85    }
86}
87
88fn parse_timestamp_str(s: &str) -> Option<f64> {
89    static RE: OnceLock<Regex> = OnceLock::new();
90    let re = RE.get_or_init(|| {
91        Regex::new(r"^(?:([0-9]{1,2}):)?([0-9]{1,2}):([0-9]{1,2})(?:[.,]([0-9]{1,3}))?$").unwrap()
92    });
93    let trimmed = s.trim_matches(is_js_space);
94    let caps = re.captures(trimmed)?;
95    let num = |i: usize| {
96        caps.get(i)
97            .map_or(0.0, |m| m.as_str().parse::<f64>().unwrap_or(0.0))
98    };
99    let ms = caps.get(4).map_or(0.0, |m| {
100        format!("{:0<3}", m.as_str()).parse::<f64>().unwrap_or(0.0)
101    });
102    Some(num(1) * 3_600_000.0 + num(2) * 60_000.0 + num(3) * 1000.0 + ms)
103}
104
105/// Build a segment, or `None` if its timing is invalid or it has no content.
106fn normalize_segment(
107    start: Option<f64>,
108    end: Option<f64>,
109    text: &str,
110    words: Vec<Word>,
111) -> Option<Segment> {
112    let (start, end) = (start?, end?);
113    if !start.is_finite() || !end.is_finite() || end <= start {
114        return None;
115    }
116    let safe_text = sanitize(text);
117    if safe_text.is_empty() && words.is_empty() {
118        return None;
119    }
120    let text = if safe_text.is_empty() {
121        words
122            .iter()
123            .map(|w| w.text.as_str())
124            .collect::<Vec<_>>()
125            .join(" ")
126    } else {
127        safe_text
128    };
129    Some(Segment {
130        start_ms: js_round(start).max(0.0) as i64,
131        end_ms: js_round(end).max(0.0) as i64,
132        duration_ms: js_round(end - start) as i64,
133        text,
134        words,
135    })
136}
137
138/// Read a start or end time from a whisper segment/token object, trying the
139/// field names used by the different whisper.cpp versions and output modes.
140fn read_time(obj: &Value, which: Edge) -> Option<f64> {
141    let (plain, ms_key, short) = match which {
142        Edge::Start => ("start", "start_ms", "from"),
143        Edge::End => ("end", "end_ms", "to"),
144    };
145    let value = [
146        get(obj, plain),
147        get(obj, ms_key),
148        get(obj, short),
149        get_path(obj, &["offsets", short]),
150        get_path(obj, &["timestamps", short]),
151    ]
152    .into_iter()
153    .find_map(nullish)?;
154
155    // Explicit millisecond fields are used as-is; everything else is parsed.
156    let explicit_ms = nullish(get(obj, ms_key)).is_some()
157        || nullish(get_path(obj, &["offsets", short])).is_some();
158    match value {
159        Value::Number(n) if explicit_ms => n.as_f64(),
160        other => parse_timestamp_ms(other),
161    }
162}
163
164#[derive(Clone, Copy)]
165enum Edge {
166    Start,
167    End,
168}
169
170/// Normalize one raw whisper word/token, keeping its text untouched.
171pub fn normalize_raw_token(token: &Value) -> Option<RawToken> {
172    let start_ms = read_time(token, Edge::Start);
173    let end_ms = read_time(token, Edge::End);
174    let raw = [get(token, "text"), get(token, "word"), get(token, "token")]
175        .into_iter()
176        .find_map(nullish)?;
177    if is_falsy(raw) {
178        return None;
179    }
180    let raw_text = js_string(raw);
181    let (start_ms, end_ms) = (start_ms?, end_ms?);
182    if !start_ms.is_finite() || !end_ms.is_finite() || end_ms <= start_ms {
183        return None;
184    }
185    Some(RawToken {
186        raw_text,
187        start_ms,
188        end_ms,
189    })
190}
191
192/// Is this a whisper control token rather than speech?
193///
194/// Covers whisper.cpp's `[_BEG_]`, `[_TT_150]`, `[_SOT_]`, `[_EOT_]`, … and
195/// the `<|endoftext|>` / `<|0.00|>` style used by OpenAI / Hugging Face.
196pub fn is_special_token(text: &str) -> bool {
197    let t = text.trim_matches(is_js_space);
198    (t.starts_with("[_") && t.ends_with(']')) || (t.starts_with("<|") && t.ends_with("|>"))
199}
200
201/// Reassemble whisper's BPE tokens into whole words.
202///
203/// A token that starts a new word carries a leading space (`" Mul"`), a
204/// continuation does not (`"do"`, `"on"`). Sanitizing tokens before merging
205/// destroys that signal and produces `"M uld oon"` or `"friend ."`.
206/// Special tokens (see [`is_special_token`]) are skipped.
207///
208/// ```
209/// use whisper_captions::{merge_bpe_tokens, RawToken};
210/// let t = |text: &str, start_ms: f64, end_ms: f64| RawToken { raw_text: text.into(), start_ms, end_ms };
211/// let words = merge_bpe_tokens([t("[_BEG_]", 0.0, 10.0), t(" Mul", 0.0, 300.0), t("do", 300.0, 500.0), t("on", 500.0, 800.0)]);
212/// assert_eq!(words[0].text, "Muldoon");
213/// assert_eq!((words[0].start_ms, words[0].end_ms), (0, 800));
214/// ```
215pub fn merge_bpe_tokens(tokens: impl IntoIterator<Item = RawToken>) -> Vec<Word> {
216    let mut merged: Vec<RawToken> = Vec::new();
217    for token in tokens {
218        if is_special_token(&token.raw_text) {
219            continue;
220        }
221        let starts_with_space = token.raw_text.chars().next().is_some_and(is_js_space);
222        match merged.last_mut() {
223            Some(previous) if !starts_with_space => {
224                previous.raw_text.push_str(&token.raw_text);
225                previous.end_ms = token.end_ms;
226            }
227            _ => merged.push(RawToken {
228                raw_text: trim_start_js(&token.raw_text).to_string(),
229                ..token
230            }),
231        }
232    }
233
234    merged
235        .into_iter()
236        .filter_map(|w| {
237            let text = sanitize(&w.raw_text);
238            let start_ms = js_round(w.start_ms).max(0.0);
239            let end_ms = js_round(w.end_ms).max(0.0);
240            (!text.is_empty() && end_ms > start_ms).then_some(Word {
241                text,
242                start_ms: start_ms as i64,
243                end_ms: end_ms as i64,
244            })
245        })
246        .collect()
247}
248
249fn normalize_json_segment(segment: &Value) -> Option<Segment> {
250    let tokens = [
251        get(segment, "words"),
252        get(segment, "tokens"),
253        get_path(segment, &["result", "words"]),
254    ]
255    .into_iter()
256    .flatten()
257    .find_map(Value::as_array);
258    let words = merge_bpe_tokens(tokens.into_iter().flatten().filter_map(normalize_raw_token));
259
260    let start =
261        read_time(segment, Edge::Start).or_else(|| words.first().map(|w| w.start_ms as f64));
262    let end = read_time(segment, Edge::End).or_else(|| words.last().map(|w| w.end_ms as f64));
263    let text = js_string_or_empty(get(segment, "text"));
264
265    normalize_segment(start, end, &text, words)
266}
267
268/// Parse a whisper.cpp JSON result (any of the known layouts).
269pub fn parse_whisper_json(json: &Value) -> Vec<Segment> {
270    if is_falsy(json) {
271        return Vec::new();
272    }
273    let candidates = [
274        get(json, "transcription"),
275        get(json, "segments"),
276        get_path(json, &["result", "segments"]),
277        get(json, "results"),
278    ]
279    .into_iter()
280    .flatten()
281    .find_map(Value::as_array);
282
283    candidates
284        .map(|list| list.iter().filter_map(normalize_json_segment).collect())
285        .unwrap_or_default()
286}
287
288/// Parse a whisper.cpp JSON string. Invalid JSON yields an error.
289pub fn parse_whisper_json_str(json: &str) -> Result<Vec<Segment>, serde_json::Error> {
290    Ok(parse_whisper_json(&serde_json::from_str(json)?))
291}
292
293/// Parse whisper.cpp's human-readable stdout (`[start --> end] text` lines).
294pub fn parse_whisper_stdout(stdout: &str) -> Vec<Segment> {
295    static RE: OnceLock<Regex> = OnceLock::new();
296    let re = RE.get_or_init(|| {
297        // `S` mirrors JavaScript's `\s` exactly (Rust's `\s` differs on U+0085 / U+FEFF).
298        const S: &str = r"[\t\n\x0B\x0C\r \x{A0}\x{1680}\x{2000}-\x{200A}\x{2028}\x{2029}\x{202F}\x{205F}\x{3000}\x{FEFF}]";
299        const TS: &str = r"[0-9]{1,2}:[0-9]{2}:[0-9]{2}[.,][0-9]{1,3}";
300        Regex::new(&format!(r"\[({TS}){S}*-->{S}*({TS})\]{S}*([^\n\r]+)")).unwrap()
301    });
302
303    re.captures_iter(stdout)
304        .filter_map(|c| {
305            normalize_segment(
306                parse_timestamp_str(&c[1]),
307                parse_timestamp_str(&c[2]),
308                &c[3],
309                Vec::new(),
310            )
311        })
312        .collect()
313}
314
315/// Prefer JSON segments; fall back to stdout when JSON is missing or empty.
316pub fn parse_transcript(json: Option<&Value>, stdout: Option<&str>) -> Vec<Segment> {
317    let from_json = json.map(parse_whisper_json).unwrap_or_default();
318    if !from_json.is_empty() {
319        return from_json;
320    }
321    stdout.map(parse_whisper_stdout).unwrap_or_default()
322}
323
324fn format_timestamp(ms: i64, separator: char) -> String {
325    let ms = ms.max(0);
326    format!(
327        "{:02}:{:02}:{:02}{separator}{:03}",
328        ms / 3_600_000,
329        ms / 60_000 % 60,
330        ms / 1000 % 60,
331        ms % 1000
332    )
333}
334
335/// Render segments as SubRip (`.srt`).
336pub fn to_srt(segments: &[Segment]) -> String {
337    segments
338        .iter()
339        .enumerate()
340        .map(|(i, s)| {
341            format!(
342                "{}\n{} --> {}\n{}\n",
343                i + 1,
344                format_timestamp(s.start_ms, ','),
345                format_timestamp(s.end_ms, ','),
346                s.text
347            )
348        })
349        .collect::<Vec<_>>()
350        .join("\n")
351}
352
353/// Render segments as WebVTT (`.vtt`).
354pub fn to_vtt(segments: &[Segment]) -> String {
355    let mut out = String::from("WEBVTT\n");
356    for s in segments {
357        out.push_str(&format!(
358            "\n{} --> {}\n{}\n",
359            format_timestamp(s.start_ms, '.'),
360            format_timestamp(s.end_ms, '.'),
361            s.text
362        ));
363    }
364    out
365}
366
367#[cfg(test)]
368mod tests;
369
370/// Compiles and runs the README examples as doctests.
371#[cfg(doctest)]
372#[doc = include_str!("../README.md")]
373pub struct ReadmeDoctests;
whisper_captions/lib.rs

whisper_captions/
lib.rs
