Skip to main content

rlx_whisper/
subtitles.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//! SRT / VTT / TSV export for [`crate::transcript::WhisperTranscript`].
17
18use crate::transcript::WhisperTranscript;
19use anyhow::Result;
20
21fn fmt_srt_time(sec: f32) -> String {
22    let ms = (sec * 1000.0).round() as u64;
23    let h = ms / 3_600_000;
24    let m = (ms % 3_600_000) / 60_000;
25    let s = (ms % 60_000) / 1000;
26    let f = ms % 1000;
27    format!("{h:02}:{m:02}:{s:02},{f:03}")
28}
29
30fn fmt_vtt_time(sec: f32) -> String {
31    let ms = (sec * 1000.0).round() as u64;
32    let h = ms / 3_600_000;
33    let m = (ms % 3_600_000) / 60_000;
34    let s = (ms % 60_000) / 1000;
35    let f = ms % 1000;
36    format!("{h:02}:{m:02}:{s:02}.{f:03}")
37}
38
39pub fn to_srt(t: &WhisperTranscript) -> String {
40    let mut out = String::new();
41    for (i, seg) in t.segments.iter().enumerate() {
42        if seg.text.trim().is_empty() {
43            continue;
44        }
45        let label = seg
46            .speaker
47            .as_ref()
48            .map(|s| format!("[{s}] "))
49            .unwrap_or_default();
50        out.push_str(&(i + 1).to_string());
51        out.push('\n');
52        out.push_str(&fmt_srt_time(seg.start));
53        out.push_str(" --> ");
54        out.push_str(&fmt_srt_time(seg.end));
55        out.push('\n');
56        out.push_str(&label);
57        out.push_str(seg.text.trim());
58        out.push_str("\n\n");
59    }
60    out
61}
62
63pub fn to_vtt(t: &WhisperTranscript) -> String {
64    let mut out = String::from("WEBVTT\n\n");
65    for seg in &t.segments {
66        if seg.text.trim().is_empty() {
67            continue;
68        }
69        out.push_str(&fmt_vtt_time(seg.start));
70        out.push_str(" --> ");
71        out.push_str(&fmt_vtt_time(seg.end));
72        out.push('\n');
73        if let Some(spk) = &seg.speaker {
74            out.push_str(&format!("<v {spk}>"));
75        }
76        out.push_str(seg.text.trim());
77        out.push_str("\n\n");
78    }
79    out
80}
81
82pub fn to_tsv(t: &WhisperTranscript) -> String {
83    let mut out = String::from("start\tend\ttext\tspeaker\n");
84    for seg in &t.segments {
85        let spk = seg.speaker.as_deref().unwrap_or("");
86        out.push_str(&format!(
87            "{:.3}\t{:.3}\t{}\t{spk}\n",
88            seg.start,
89            seg.end,
90            seg.text.replace('\t', " ").trim()
91        ));
92    }
93    out
94}
95
96pub fn to_json_pretty(t: &WhisperTranscript) -> Result<String> {
97    Ok(serde_json::to_string_pretty(t)?)
98}
99
100#[derive(Debug, Clone, Copy, PartialEq, Eq)]
101pub enum SubtitleFormat {
102    Srt,
103    Vtt,
104    Tsv,
105    Json,
106}
107
108impl SubtitleFormat {
109    pub fn parse(s: &str) -> Option<Self> {
110        match s.to_ascii_lowercase().as_str() {
111            "srt" => Some(Self::Srt),
112            "vtt" | "webvtt" => Some(Self::Vtt),
113            "tsv" => Some(Self::Tsv),
114            "json" => Some(Self::Json),
115            _ => None,
116        }
117    }
118
119    pub fn render(self, t: &WhisperTranscript) -> Result<String> {
120        match self {
121            Self::Srt => Ok(to_srt(t)),
122            Self::Vtt => Ok(to_vtt(t)),
123            Self::Tsv => Ok(to_tsv(t)),
124            Self::Json => to_json_pretty(t),
125        }
126    }
127}