Mercurial
diff dictation/session.py @ 275:78699f810817
Add Qwen3-VL and WebRTC dictation services
Add Bazel targets for the CUDA-backed Qwen3-VL server and a local WebRTC faster-whisper dictation service.
Co-authored-by: Copilot <[email protected]>
Copilot-Session: e3d8cb06-6c95-4ae0-9757-651d3796ab00
| author | MrJuneJune <me@mrjunejune.com> |
|---|---|
| date | Mon, 17 Aug 2026 10:58:47 -0700 |
| parents | |
| children | 49e9e591c9bb |
line wrap: on
line diff
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/dictation/session.py Mon Aug 17 10:58:47 2026 -0700 @@ -0,0 +1,177 @@ +from __future__ import annotations + +import asyncio +import json +from typing import Any + +from aiortc.mediastreams import MediaStreamError +from av.audio.resampler import AudioResampler +import numpy as np + +from dictation.audio import AudioEventKind, AudioSegmenter +from dictation.config import DictationConfig +from dictation.transcriber import Transcriber + + +class DictationSession: + def __init__( + self, + session_id: str, + peer: Any, + transcriber: Transcriber, + config: DictationConfig, + ) -> None: + self.session_id = session_id + self.peer = peer + self.transcriber = transcriber + self.segmenter = AudioSegmenter( + speech_threshold=config.speech_threshold, + silence_ms=config.silence_ms, + partial_interval_ms=config.partial_interval_ms, + max_utterance_seconds=config.max_utterance_seconds, + ) + self.channel = None + self.audio_task: asyncio.Task | None = None + self.partial_task: asyncio.Task | None = None + self.final_tasks: set[asyncio.Task] = set() + self.closed = False + self.stopping = False + self._resampler = AudioResampler( + format="s16", + layout="mono", + rate=16000, + ) + + def attach_channel(self, channel: Any) -> None: + if channel.label != "transcripts": + channel.close() + return + self.channel = channel + + @channel.on("open") + def on_open() -> None: + self.send("ready", sessionId=self.session_id) + + @channel.on("message") + def on_message(message: Any) -> None: + if message == "stop": + self.stopping = True + event = self.segmenter.flush() + if event and event.samples is not None: + self._start_final(event.samples) + else: + self.send("speech.ended") + + if channel.readyState == "open": + self.send("ready", sessionId=self.session_id) + + def attach_track(self, track: Any) -> None: + if track.kind != "audio" or self.audio_task is not None: + return + self.audio_task = asyncio.create_task(self._consume_audio(track)) + + async def _consume_audio(self, track: Any) -> None: + try: + while not self.closed: + frame = await track.recv() + for converted in self._resampler.resample(frame): + if self.stopping: + continue + samples = converted.to_ndarray().reshape(-1) + normalized = samples.astype(np.float32) / 32768.0 + for event in self.segmenter.feed(normalized): + if event.kind == AudioEventKind.SPEECH_STARTED: + self.send("speech.started") + elif ( + event.kind == AudioEventKind.PARTIAL_READY + and event.samples is not None + ): + self._start_partial(event.samples) + elif ( + event.kind == AudioEventKind.FINAL_READY + and event.samples is not None + ): + self._start_final(event.samples) + except MediaStreamError: + if not self.stopping: + event = self.segmenter.flush() + if event and event.samples is not None: + self._start_final(event.samples) + else: + self.send("speech.ended") + except asyncio.CancelledError: + raise + except Exception as error: + if not self.closed: + self.send( + "error", + code="audio_track_failed", + message=str(error), + ) + + def _start_partial(self, samples: np.ndarray) -> None: + if self.partial_task and not self.partial_task.done(): + return + self.partial_task = asyncio.create_task( + self._transcribe(samples, final=False) + ) + + def _start_final(self, samples: np.ndarray) -> None: + task = asyncio.create_task(self._transcribe(samples, final=True)) + self.final_tasks.add(task) + task.add_done_callback(self.final_tasks.discard) + + async def _transcribe(self, samples: np.ndarray, *, final: bool) -> None: + try: + transcript = await self.transcriber.transcribe( + samples, + final=final, + ) + if transcript.text: + self.send( + "transcript.final" if final else "transcript.partial", + text=transcript.text, + language=transcript.language, + probability=transcript.probability, + ) + if final: + self.send("speech.ended") + except asyncio.CancelledError: + raise + except Exception as error: + self.send( + "error", + code="transcription_failed", + message=str(error), + ) + + def send(self, event_type: str, **payload: Any) -> None: + if not self.channel or self.channel.readyState != "open": + return + self.channel.send( + json.dumps( + {"type": event_type, **payload}, + separators=(",", ":"), + ) + ) + + async def close(self) -> None: + if self.closed: + return + self.closed = True + tasks = [ + task + for task in [ + self.audio_task, + self.partial_task, + *self.final_tasks, + ] + if task is not None and not task.done() + ] + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + if self.channel and self.channel.readyState != "closed": + self.channel.close() + await self.peer.close()