studio_worker/stt_stream/
session.rs

1//! One streaming speech-to-text session: the `/transcribe` protocol over
2//! any streaming transcriber.  Pure (no sockets), so every rule is tested.
3//!
4//! Client: binary frames of 16 kHz mono s16le PCM; text `end` finalises;
5//! text `cancel` closes without a final.  Server: `{"partial":true,"text"}`
6//! as the transcript grows, `{"final":true,"text"}` once, `{"error"}`.
7
8use super::vad::{Vad, SAMPLE_RATE};
9
10/// Silence appended on finalise so the model emits its last words (a
11/// streaming model lags its input by up to a chunk plus look-ahead).
12/// Measured: 1.5 s flushed the last word in the spike.  Safe range 0.5..=3 s.
13pub const FLUSH_SILENCE_MS: usize = 1500;
14
15/// A streaming speech model with per-utterance state.
16pub trait StreamingTranscriber {
17    /// Samples the model takes per step (e.g. 2560 = 160 ms).
18    fn chunk_samples(&self) -> usize;
19    /// Feed one chunk; returns the text it adds (may be empty).
20    fn step(&mut self, chunk: &[f32]) -> anyhow::Result<String>;
21    /// Forget the previous utterance.
22    fn reset(&mut self);
23}
24
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub enum ClientFrame {
27    /// 16 kHz mono s16le PCM.
28    Audio(Vec<u8>),
29    End,
30    Cancel,
31}
32
33impl ClientFrame {
34    /// A text frame by name; `None` for anything unknown.
35    pub fn from_text(text: &str) -> Option<Self> {
36        match text.trim() {
37            "end" => Some(Self::End),
38            "cancel" => Some(Self::Cancel),
39            _ => None,
40        }
41    }
42}
43
44#[derive(Debug, Clone, PartialEq, Eq)]
45pub enum ServerFrame {
46    Partial(String),
47    Final(String),
48    Error(String),
49}
50
51impl ServerFrame {
52    pub fn to_json(&self) -> serde_json::Value {
53        match self {
54            Self::Partial(text) => serde_json::json!({ "partial": true, "text": text }),
55            Self::Final(text) => serde_json::json!({ "final": true, "text": text }),
56            Self::Error(error) => serde_json::json!({ "error": error }),
57        }
58    }
59}
60
61/// Keep the socket open, or close it after sending.
62#[derive(Debug, Clone, Copy, PartialEq, Eq)]
63pub enum Next {
64    Continue,
65    Close,
66}
67
68/// SentencePiece word markers to spaces, whitespace collapsed.
69pub fn normalise(raw: &str) -> String {
70    raw.replace('\u{2581}', " ")
71        .split_whitespace()
72        .collect::<Vec<_>>()
73        .join(" ")
74}
75
76/// One utterance over a transcriber.
77pub struct StreamSession<'a> {
78    transcriber: &'a mut dyn StreamingTranscriber,
79    pending: Vec<f32>,
80    odd_byte: Option<u8>,
81    raw: String,
82    sent: String,
83    vad: Vad,
84}
85
86impl<'a> StreamSession<'a> {
87    pub fn new(transcriber: &'a mut dyn StreamingTranscriber) -> Self {
88        transcriber.reset();
89        Self {
90            transcriber,
91            pending: Vec::new(),
92            odd_byte: None,
93            raw: String::new(),
94            sent: String::new(),
95            vad: Vad::default(),
96        }
97    }
98
99    /// Handle one client frame: the frames to send, then whether to close.
100    pub fn handle(&mut self, frame: ClientFrame) -> (Vec<ServerFrame>, Next) {
101        match frame {
102            ClientFrame::Audio(bytes) => {
103                let samples = self.decode(&bytes);
104                let ended = self.vad.feed(&samples);
105                self.pending.extend_from_slice(&samples);
106                let mut out = Vec::new();
107                if let Err(err) = self.drain(&mut out, true) {
108                    out.push(failure(&err));
109                    return (out, Next::Close);
110                }
111                if ended {
112                    return self.finish(out);
113                }
114                (out, Next::Continue)
115            }
116            ClientFrame::End => self.finish(Vec::new()),
117            ClientFrame::Cancel => (Vec::new(), Next::Close),
118        }
119    }
120
121    fn decode(&mut self, bytes: &[u8]) -> Vec<f32> {
122        let mut joined = Vec::with_capacity(bytes.len() + 1);
123        joined.extend(self.odd_byte.take());
124        joined.extend_from_slice(bytes);
125        if joined.len() % 2 == 1 {
126            self.odd_byte = joined.pop();
127        }
128        joined
129            .as_chunks::<2>()
130            .0
131            .iter()
132            .map(|b| i16::from_le_bytes(*b) as f32 / 32768.0)
133            .collect()
134    }
135
136    /// Step every whole chunk; send a partial when the transcript changed.
137    fn drain(&mut self, out: &mut Vec<ServerFrame>, partials: bool) -> anyhow::Result<()> {
138        let chunk = self.transcriber.chunk_samples().max(1);
139        while self.pending.len() >= chunk {
140            let piece: Vec<f32> = self.pending.drain(..chunk).collect();
141            self.raw.push_str(&self.transcriber.step(&piece)?);
142            let text = normalise(&self.raw);
143            if partials && text != self.sent {
144                out.push(ServerFrame::Partial(text.clone()));
145                self.sent = text;
146            }
147        }
148        Ok(())
149    }
150
151    /// Flush the tail with silence, then send the final transcript.
152    fn finish(&mut self, mut out: Vec<ServerFrame>) -> (Vec<ServerFrame>, Next) {
153        let chunk = self.transcriber.chunk_samples().max(1);
154        let flush = SAMPLE_RATE * FLUSH_SILENCE_MS / 1000;
155        let target = (self.pending.len() + flush).div_ceil(chunk) * chunk;
156        self.pending.resize(target, 0.0);
157        match self.drain(&mut out, false) {
158            Ok(()) => out.push(ServerFrame::Final(normalise(&self.raw))),
159            Err(err) => out.push(failure(&err)),
160        }
161        (out, Next::Close)
162    }
163}
164
165fn failure(err: &anyhow::Error) -> ServerFrame {
166    ServerFrame::Error(format!("transcription failed: {err:#}"))
167}
168
169#[cfg(test)]
170mod tests {
171    use super::*;
172
173    /// Emits one word per chunk from a script; counts resets.
174    struct Scripted {
175        words: Vec<&'static str>,
176        next: usize,
177        resets: usize,
178        chunk: usize,
179        fail_at: Option<usize>,
180    }
181
182    impl Scripted {
183        fn new(words: &[&'static str]) -> Self {
184            Self {
185                words: words.to_vec(),
186                next: 0,
187                resets: 0,
188                chunk: 1600,
189                fail_at: None,
190            }
191        }
192    }
193
194    impl StreamingTranscriber for Scripted {
195        fn chunk_samples(&self) -> usize {
196            self.chunk
197        }
198        fn step(&mut self, _chunk: &[f32]) -> anyhow::Result<String> {
199            if self.fail_at == Some(self.next) {
200                anyhow::bail!("cuda went away");
201            }
202            let piece = self.words.get(self.next).copied().unwrap_or("");
203            self.next += 1;
204            Ok(piece.to_string())
205        }
206        fn reset(&mut self) {
207            self.resets += 1;
208            self.next = 0;
209        }
210    }
211
212    /// `ms` of PCM at a constant level, as s16le bytes.
213    fn pcm(ms: usize, level: i16) -> Vec<u8> {
214        (0..ms * 16)
215            .flat_map(|i| (if i % 2 == 0 { level } else { -level }).to_le_bytes())
216            .collect()
217    }
218
219    fn texts(frames: &[ServerFrame]) -> Vec<String> {
220        frames
221            .iter()
222            .map(|f| match f {
223                ServerFrame::Partial(t) => format!("partial:{t}"),
224                ServerFrame::Final(t) => format!("final:{t}"),
225                ServerFrame::Error(e) => format!("error:{e}"),
226            })
227            .collect()
228    }
229
230    #[test]
231    fn a_session_starts_from_a_clean_transcriber() {
232        let mut t = Scripted::new(&[]);
233        StreamSession::new(&mut t);
234        assert_eq!(t.resets, 1);
235    }
236
237    #[test]
238    fn partials_grow_as_chunks_arrive() {
239        let mut t = Scripted::new(&["\u{2581}hello", "\u{2581}runa"]);
240        let mut s = StreamSession::new(&mut t);
241        let (out, next) = s.handle(ClientFrame::Audio(pcm(200, 8000)));
242        assert_eq!(next, Next::Continue);
243        assert_eq!(texts(&out), ["partial:hello", "partial:hello runa"]);
244    }
245
246    #[test]
247    fn an_unchanged_transcript_sends_no_partial() {
248        let mut t = Scripted::new(&["hi", "", ""]);
249        let mut s = StreamSession::new(&mut t);
250        let (out, _) = s.handle(ClientFrame::Audio(pcm(300, 8000)));
251        assert_eq!(texts(&out), ["partial:hi"]);
252    }
253
254    #[test]
255    fn audio_is_buffered_until_a_whole_chunk_arrives() {
256        let mut t = Scripted::new(&["one"]);
257        let mut s = StreamSession::new(&mut t);
258        let (out, _) = s.handle(ClientFrame::Audio(pcm(50, 8000)));
259        assert!(out.is_empty(), "50 ms is less than one 100 ms chunk");
260        let (out, _) = s.handle(ClientFrame::Audio(pcm(50, 8000)));
261        assert_eq!(texts(&out), ["partial:one"]);
262    }
263
264    #[test]
265    fn an_odd_byte_carries_over_to_the_next_frame() {
266        let mut t = Scripted::new(&["one"]);
267        let mut s = StreamSession::new(&mut t);
268        let bytes = pcm(100, 8000);
269        let (a, b) = bytes.split_at(1601);
270        assert!(s.handle(ClientFrame::Audio(a.to_vec())).0.is_empty());
271        let (out, _) = s.handle(ClientFrame::Audio(b.to_vec()));
272        assert_eq!(texts(&out), ["partial:one"]);
273    }
274
275    #[test]
276    fn end_flushes_and_sends_the_final_then_closes() {
277        let mut t = Scripted::new(&["hello", " there", " friend"]);
278        let mut s = StreamSession::new(&mut t);
279        s.handle(ClientFrame::Audio(pcm(100, 8000)));
280        let (out, next) = s.handle(ClientFrame::End);
281        assert_eq!(next, Next::Close);
282        assert_eq!(texts(&out).last().unwrap(), "final:hello there friend");
283    }
284
285    #[test]
286    fn cancel_closes_without_a_final() {
287        let mut t = Scripted::new(&["x"]);
288        let mut s = StreamSession::new(&mut t);
289        let (out, next) = s.handle(ClientFrame::Cancel);
290        assert_eq!(next, Next::Close);
291        assert!(out.is_empty());
292    }
293
294    #[test]
295    fn silence_after_speech_finalises_hands_free() {
296        let mut t = Scripted::new(&["hi"]);
297        let mut s = StreamSession::new(&mut t);
298        s.handle(ClientFrame::Audio(pcm(400, 8000)));
299        let (out, next) = s.handle(ClientFrame::Audio(pcm(1600, 0)));
300        assert_eq!(next, Next::Close);
301        assert_eq!(texts(&out).last().unwrap(), "final:hi");
302    }
303
304    #[test]
305    fn a_transcriber_failure_is_reported_and_closes() {
306        let mut t = Scripted::new(&["a", "b"]);
307        t.fail_at = Some(1);
308        let mut s = StreamSession::new(&mut t);
309        let (out, next) = s.handle(ClientFrame::Audio(pcm(200, 8000)));
310        assert_eq!(next, Next::Close);
311        assert_eq!(
312            texts(&out),
313            ["partial:a", "error:transcription failed: cuda went away"]
314        );
315    }
316
317    #[test]
318    fn frames_serialise_to_the_wire_protocol() {
319        assert_eq!(
320            ServerFrame::Partial("hi".into()).to_json(),
321            serde_json::json!({ "partial": true, "text": "hi" })
322        );
323        assert_eq!(
324            ServerFrame::Final("hi".into()).to_json(),
325            serde_json::json!({ "final": true, "text": "hi" })
326        );
327        assert_eq!(
328            ServerFrame::Error("x".into()).to_json(),
329            serde_json::json!({ "error": "x" })
330        );
331    }
332
333    #[test]
334    fn client_text_frames_parse_by_name() {
335        assert_eq!(ClientFrame::from_text("end"), Some(ClientFrame::End));
336        assert_eq!(
337            ClientFrame::from_text(" cancel "),
338            Some(ClientFrame::Cancel)
339        );
340        assert_eq!(ClientFrame::from_text("hello"), None);
341    }
342
343    #[test]
344    fn pieces_normalise_to_plain_text() {
345        assert_eq!(
346            normalise("\u{2581}hello\u{2581}\u{2581}world  "),
347            "hello world"
348        );
349        assert_eq!(normalise(" Hello Runa, please "), "Hello Runa, please");
350    }
351}
studio_worker/stt_stream/session.rs

studio_worker/stt_stream/
session.rs
