rustlavel_ai/transcribe/
transcript.rs

1//! What comes back: the words, and when they were said.
2
3use rustlavel_core::Json;
4
5/// One word, with the moment it was spoken.
6///
7/// Times are seconds from the start of the audio. A clip is cut on these, so
8/// they are `f64`, not milliseconds rounded somewhere on the way in.
9#[derive(Debug, Clone, PartialEq)]
10pub struct Word {
11    pub text: String,
12    pub start: f64,
13    pub end: f64,
14    /// Zero-based speaker index, when the provider was asked to tell speakers
15    /// apart and could.
16    pub speaker: Option<u32>,
17    /// 0–1, when the provider reports one.
18    pub confidence: Option<f64>,
19}
20
21/// A run of speech: a sentence, an utterance, one speaker's turn.
22#[derive(Debug, Clone, PartialEq)]
23pub struct Segment {
24    pub text: String,
25    pub start: f64,
26    pub end: f64,
27    pub speaker: Option<u32>,
28}
29
30#[derive(Debug, Clone, PartialEq)]
31pub struct Transcript {
32    pub text: String,
33    /// ISO 639-1 where the provider says (`"id"`, `"en"`); `None` when it
34    /// does not.
35    pub language: Option<String>,
36    /// Length of the audio in seconds, when reported.
37    pub duration: Option<f64>,
38    pub segments: Vec<Segment>,
39    pub words: Vec<Word>,
40    /// Whether speakers were told apart. `false` means every `speaker` is
41    /// `None` because nobody looked, not because there was one speaker.
42    pub diarized: bool,
43    pub model: String,
44    pub provider: String,
45}
46
47impl Transcript {
48    /// How many distinct speakers were heard. Zero when not diarized.
49    pub fn speakers(&self) -> usize {
50        let mut seen: Vec<u32> = self
51            .words
52            .iter()
53            .filter_map(|w| w.speaker)
54            .chain(self.segments.iter().filter_map(|s| s.speaker))
55            .collect();
56        seen.sort_unstable();
57        seen.dedup();
58        seen.len()
59    }
60
61    /// The words spoken between two moments — what a clip from `start` to
62    /// `end` contains. A word counts when any part of it is inside.
63    pub fn words_between(&self, start: f64, end: f64) -> Vec<&Word> {
64        self.words.iter().filter(|word| word.end > start && word.start < end).collect()
65    }
66
67    /// The segments that overlap a window, for cutting on sentence edges.
68    pub fn segments_between(&self, start: f64, end: f64) -> Vec<&Segment> {
69        self.segments.iter().filter(|segment| segment.end > start && segment.start < end).collect()
70    }
71
72    /// SubRip subtitles, one cue per segment. `ffmpeg -vf subtitles=` burns
73    /// this in.
74    pub fn to_srt(&self) -> String {
75        let mut out = String::new();
76        for (index, segment) in self.segments.iter().enumerate() {
77            out.push_str(&format!(
78                "{}\n{} --> {}\n{}\n\n",
79                index + 1,
80                timestamp(segment.start, ','),
81                timestamp(segment.end, ','),
82                labelled(segment)
83            ));
84        }
85        out
86    }
87
88    /// WebVTT, for a `<track>` on a web player.
89    pub fn to_vtt(&self) -> String {
90        let mut out = String::from("WEBVTT\n\n");
91        for segment in &self.segments {
92            out.push_str(&format!(
93                "{} --> {}\n{}\n\n",
94                timestamp(segment.start, '.'),
95                timestamp(segment.end, '.'),
96                labelled(segment)
97            ));
98        }
99        out
100    }
101
102    pub fn to_json(&self) -> Json {
103        Json::object([
104            ("text", Json::from(self.text.as_str())),
105            ("language", self.language.as_deref().map_or(Json::Null, Json::from)),
106            ("duration", self.duration.map_or(Json::Null, Json::from)),
107            ("diarized", Json::from(self.diarized)),
108            ("model", Json::from(self.model.as_str())),
109            ("provider", Json::from(self.provider.as_str())),
110            (
111                "segments",
112                Json::Array(
113                    self.segments
114                        .iter()
115                        .map(|s| {
116                            Json::object([
117                                ("text", Json::from(s.text.as_str())),
118                                ("start", Json::from(s.start)),
119                                ("end", Json::from(s.end)),
120                                ("speaker", s.speaker.map_or(Json::Null, |n| Json::from(n as i64))),
121                            ])
122                        })
123                        .collect(),
124                ),
125            ),
126            (
127                "words",
128                Json::Array(
129                    self.words
130                        .iter()
131                        .map(|w| {
132                            Json::object([
133                                ("text", Json::from(w.text.as_str())),
134                                ("start", Json::from(w.start)),
135                                ("end", Json::from(w.end)),
136                                ("speaker", w.speaker.map_or(Json::Null, |n| Json::from(n as i64))),
137                                ("confidence", w.confidence.map_or(Json::Null, Json::from)),
138                            ])
139                        })
140                        .collect(),
141                ),
142            ),
143        ])
144    }
145}
146
147/// Segments built from words when a provider gives words but no segments:
148/// a new one at each change of speaker, at sentence-ending punctuation, or
149/// after a pause of a second.
150pub(crate) fn segments_from_words(words: &[Word]) -> Vec<Segment> {
151    let mut segments: Vec<Segment> = Vec::new();
152    let mut current: Option<Segment> = None;
153    let mut previous_end = 0.0;
154
155    for word in words {
156        let breaks = match &current {
157            None => true,
158            Some(segment) => segment.speaker != word.speaker || word.start - previous_end > 1.0,
159        };
160        if breaks && let Some(done) = current.take() {
161            segments.push(done);
162        }
163        match &mut current {
164            Some(segment) => {
165                segment.text.push(' ');
166                segment.text.push_str(&word.text);
167                segment.end = word.end;
168            }
169            None => {
170                current = Some(Segment { text: word.text.clone(), start: word.start, end: word.end, speaker: word.speaker });
171            }
172        }
173        previous_end = word.end;
174        if word.text.ends_with(['.', '?', '!']) && let Some(done) = current.take() {
175            segments.push(done);
176        }
177    }
178    if let Some(done) = current {
179        segments.push(done);
180    }
181    segments
182}
183
184fn labelled(segment: &Segment) -> String {
185    match segment.speaker {
186        Some(speaker) => format!("[Speaker {}] {}", speaker + 1, segment.text),
187        None => segment.text.clone(),
188    }
189}
190
191/// `HH:MM:SS,mmm` (SRT) or `HH:MM:SS.mmm` (VTT).
192fn timestamp(seconds: f64, separator: char) -> String {
193    let millis = (seconds.max(0.0) * 1000.0).round() as u64;
194    format!(
195        "{:02}:{:02}:{:02}{separator}{:03}",
196        millis / 3_600_000,
197        millis / 60_000 % 60,
198        millis / 1000 % 60,
199        millis % 1000
200    )
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206
207    fn word(text: &str, start: f64, end: f64, speaker: Option<u32>) -> Word {
208        Word { text: text.into(), start, end, speaker, confidence: None }
209    }
210
211    #[test]
212    fn timestamps_are_subtitle_shaped() {
213        assert_eq!(timestamp(0.0, ','), "00:00:00,000");
214        assert_eq!(timestamp(61.5, ','), "00:01:01,500");
215        assert_eq!(timestamp(3725.042, '.'), "01:02:05.042");
216    }
217
218    #[test]
219    fn segments_break_on_speaker_punctuation_and_pause() {
220        let words = vec![
221            word("Halo", 0.0, 0.4, Some(0)),
222            word("semua.", 0.5, 0.9, Some(0)),
223            word("Apa", 1.0, 1.2, Some(0)),
224            word("kabar", 1.3, 1.6, Some(1)),
225            word("hari", 3.0, 3.3, Some(1)),
226            word("ini", 3.4, 3.6, Some(1)),
227        ];
228        let segments = segments_from_words(&words);
229        let texts: Vec<&str> = segments.iter().map(|s| s.text.as_str()).collect();
230        assert_eq!(texts, ["Halo semua.", "Apa", "kabar", "hari ini"]);
231        assert_eq!(segments[3].start, 3.0);
232        assert_eq!(segments[3].end, 3.6);
233        assert_eq!(segments[3].speaker, Some(1));
234    }
235
236    #[test]
237    fn a_clip_window_keeps_words_that_touch_it() {
238        let transcript = Transcript {
239            text: String::new(),
240            language: None,
241            duration: None,
242            segments: vec![],
243            words: vec![word("a", 0.0, 1.0, None), word("b", 1.0, 2.0, None), word("c", 2.5, 3.0, None)],
244            diarized: false,
245            model: String::new(),
246            provider: String::new(),
247        };
248        let inside: Vec<&str> = transcript.words_between(0.9, 2.2).iter().map(|w| w.text.as_str()).collect();
249        assert_eq!(inside, ["a", "b"]);
250        assert_eq!(transcript.speakers(), 0);
251    }
252
253    #[test]
254    fn srt_and_vtt_label_speakers_when_there_are_any() {
255        let transcript = Transcript {
256            text: "Halo semua. Apa kabar".into(),
257            language: Some("id".into()),
258            duration: Some(2.0),
259            segments: vec![
260                Segment { text: "Halo semua.".into(), start: 0.0, end: 0.9, speaker: Some(0) },
261                Segment { text: "Apa kabar".into(), start: 1.0, end: 1.6, speaker: Some(1) },
262            ],
263            words: vec![],
264            diarized: true,
265            model: "m".into(),
266            provider: "p".into(),
267        };
268        let srt = transcript.to_srt();
269        assert!(srt.starts_with("1\n00:00:00,000 --> 00:00:00,900\n[Speaker 1] Halo semua.\n\n2\n"), "{srt}");
270        let vtt = transcript.to_vtt();
271        assert!(vtt.starts_with("WEBVTT\n\n00:00:00.000 --> 00:00:00.900\n[Speaker 1] Halo semua.\n"), "{vtt}");
272        assert_eq!(transcript.speakers(), 2);
273    }
274}
rustlavel_ai/transcribe/transcript.rs

rustlavel_ai/transcribe/
transcript.rs
