From 4e31f8c61b0c59ce9160c0466523cc32f406e76f Mon Sep 17 00:00:00 2001 From: Aaron Boxer Date: Fri, 17 Apr 2026 10:18:19 -0400 Subject: [PATCH] Add word-level timestamps and confidence scores MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - New word_timestamps option (default False) in client handshake - When enabled, each segment includes 'words' array with per-word start/end times and probability scores - Wired through entire pipeline: client → server → backend → transcribe() - Words include timestamp_offset for accurate absolute times - REST API already supported word timestamps; now WebSocket does too - Added 9 unit tests for word timestamp extraction and formatting --- tests/test_base_backend.py | 105 ++++++++++++++++++ tests/test_client.py | 35 +++--- whisper_live/backend/base.py | 32 +++++- .../backend/faster_whisper_backend.py | 6 +- whisper_live/client.py | 5 + whisper_live/server.py | 1 + 6 files changed, 162 insertions(+), 22 deletions(-) diff --git a/tests/test_base_backend.py b/tests/test_base_backend.py index 624ec7f..bdb6f95 100644 --- a/tests/test_base_backend.py +++ b/tests/test_base_backend.py @@ -410,5 +410,110 @@ class TestGetSegmentHelpers(unittest.TestCase): self.assertAlmostEqual(self.client.get_segment_start(seg), 2.0) +class TestWordTimestamps(unittest.TestCase): + """Tests for word-level timestamp extraction.""" + + def _make_client(self, word_timestamps=False): + ws = MagicMock() + return ConcreteServeClient( + client_uid="wt-uid", websocket=ws, word_timestamps=word_timestamps + ) + + def _make_word(self, word, start, end, prob): + w = MagicMock() + w.word = word + w.start = start + w.end = end + w.probability = prob + return w + + def _make_segment(self, text, start, end, no_speech_prob=0.0, words=None): + seg = MagicMock() + seg.text = text + seg.start = start + seg.end = end + seg.no_speech_prob = no_speech_prob + seg.words = words + return seg + + def test_word_timestamps_disabled_by_default(self): + client = self._make_client() + self.assertFalse(client.word_timestamps) + + def test_word_timestamps_enabled(self): + client = self._make_client(word_timestamps=True) + self.assertTrue(client.word_timestamps) + + def test_extract_words_when_disabled(self): + client = self._make_client(word_timestamps=False) + seg = self._make_segment("hello", 0.0, 1.0, words=[self._make_word("hello", 0.0, 0.5, 0.99)]) + result = client._extract_words(seg, 0.0) + self.assertIsNone(result) + + def test_extract_words_when_enabled(self): + client = self._make_client(word_timestamps=True) + words = [ + self._make_word("hello", 0.0, 0.3, 0.95), + self._make_word("world", 0.4, 0.8, 0.88), + ] + seg = self._make_segment("hello world", 0.0, 1.0, words=words) + result = client._extract_words(seg, 10.0) + self.assertEqual(len(result), 2) + self.assertEqual(result[0]["word"], "hello") + self.assertEqual(result[0]["start"], "10.000") + self.assertEqual(result[0]["end"], "10.300") + self.assertEqual(result[0]["probability"], 0.95) + self.assertEqual(result[1]["word"], "world") + self.assertEqual(result[1]["start"], "10.400") + + def test_extract_words_no_words_on_segment(self): + client = self._make_client(word_timestamps=True) + seg = self._make_segment("hello", 0.0, 1.0, words=None) + result = client._extract_words(seg, 0.0) + self.assertIsNone(result) + + def test_format_segment_without_words(self): + client = self._make_client() + seg = client.format_segment(0.0, 1.0, "hello") + self.assertNotIn("words", seg) + + def test_format_segment_with_words(self): + client = self._make_client(word_timestamps=True) + words = [{"word": "hello", "start": "0.000", "end": "0.500", "probability": 0.95}] + seg = client.format_segment(0.0, 1.0, "hello", words=words) + self.assertIn("words", seg) + self.assertEqual(len(seg["words"]), 1) + self.assertEqual(seg["words"][0]["word"], "hello") + + def test_update_segments_includes_words(self): + client = self._make_client(word_timestamps=True) + words1 = [self._make_word("hello", 0.0, 0.5, 0.9)] + words2 = [self._make_word("world", 0.6, 1.0, 0.85)] + segments = [ + self._make_segment(" hello", 0.0, 0.5, words=words1), + self._make_segment(" world", 0.6, 1.0, words=words2), + ] + last = client.update_segments(segments, 2.0) + # First segment should be completed (in transcript) with words + self.assertTrue(len(client.transcript) > 0) + self.assertIn("words", client.transcript[-1]) + # Last segment should be in-progress with words + self.assertIsNotNone(last) + self.assertIn("words", last) + + def test_update_segments_no_words_when_disabled(self): + client = self._make_client(word_timestamps=False) + words1 = [self._make_word("hello", 0.0, 0.5, 0.9)] + words2 = [self._make_word("world", 0.6, 1.0, 0.85)] + segments = [ + self._make_segment(" hello", 0.0, 0.5, words=words1), + self._make_segment(" world", 0.6, 1.0, words=words2), + ] + last = client.update_segments(segments, 2.0) + self.assertTrue(len(client.transcript) > 0) + self.assertNotIn("words", client.transcript[-1]) + self.assertNotIn("words", last) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_client.py b/tests/test_client.py index d55bf6d..9838668 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -43,24 +43,25 @@ class TestClientWebSocketCommunication(BaseTestCase): class TestClientCallbacks(BaseTestCase): def test_on_open(self): - expected_message = json.dumps({ - "uid": self.client.uid, - "language": self.client.language, - "task": self.client.task, - "model": self.client.model, - "use_vad": True, - "send_last_n_segments": 10, - "no_speech_thresh": 0.45, - "clip_audio": False, - "same_output_threshold": 10, - "enable_translation": False, - "target_language": "fr", - "hotwords": None, - "enable_diarization": False, - "max_speakers": 10, - }) + self.client.on_open(self.mock_ws_app) - self.mock_ws_app.send.assert_called_with(expected_message) + self.mock_ws_app.send.assert_called_once() + sent_message = json.loads(self.mock_ws_app.send.call_args[0][0]) + self.assertEqual(sent_message["uid"], self.client.uid) + self.assertEqual(sent_message["language"], self.client.language) + self.assertEqual(sent_message["task"], self.client.task) + self.assertEqual(sent_message["model"], self.client.model) + self.assertTrue(sent_message["use_vad"]) + self.assertEqual(sent_message["send_last_n_segments"], 10) + self.assertAlmostEqual(sent_message["no_speech_thresh"], 0.45) + self.assertFalse(sent_message["clip_audio"]) + self.assertEqual(sent_message["same_output_threshold"], 10) + self.assertFalse(sent_message["enable_translation"]) + self.assertEqual(sent_message["target_language"], "fr") + self.assertIsNone(sent_message["hotwords"]) + self.assertFalse(sent_message["enable_diarization"]) + self.assertEqual(sent_message["max_speakers"], 10) + self.assertFalse(sent_message["word_timestamps"]) def test_on_message(self): message = json.dumps( diff --git a/whisper_live/backend/base.py b/whisper_live/backend/base.py index 1926730..c6e6a7e 100644 --- a/whisper_live/backend/base.py +++ b/whisper_live/backend/base.py @@ -39,6 +39,7 @@ class ServeClientBase(object): same_output_threshold=10, translation_queue=None, diarization=None, + word_timestamps=False, ): self.client_uid = client_uid self.websocket = websocket @@ -47,6 +48,7 @@ class ServeClientBase(object): self.clip_audio = clip_audio self.same_output_threshold = same_output_threshold self.diarization = diarization + self.word_timestamps = word_timestamps self.frames = b"" self.timestamp_offset = 0.0 @@ -118,7 +120,7 @@ class ServeClientBase(object): def handle_transcription_output(self, result, duration): raise NotImplementedError - def format_segment(self, start, end, text, completed=False, speaker=None): + def format_segment(self, start, end, text, completed=False, speaker=None, words=None): """ Formats a transcription segment with precise start and end times alongside the transcribed text. @@ -127,6 +129,7 @@ class ServeClientBase(object): end (float): The end time of the transcription segment in seconds. text (str): The transcribed text corresponding to the segment. speaker (str, optional): Speaker label from diarization. + words (list, optional): Word-level timestamps and probabilities. Returns: dict: A dictionary representing the formatted transcription segment, including @@ -141,6 +144,8 @@ class ServeClientBase(object): } if speaker is not None: seg['speaker'] = speaker + if words is not None: + seg['words'] = words return seg def add_frames(self, frame_np): @@ -311,7 +316,6 @@ class ServeClientBase(object): seg_end = self.get_segment_end(segment) start_sample = int(seg_start * self.RATE) end_sample = int(seg_end * self.RATE) - # Extract audio relative to the current buffer samples_offset = max(0, int((self.timestamp_offset - self.frames_offset) * self.RATE)) audio_slice = self.frames_np[samples_offset + start_sample:samples_offset + end_sample] if len(audio_slice) < self.RATE * 0.3: @@ -321,6 +325,23 @@ class ServeClientBase(object): logging.error(f"Diarization error: {e}") return None + def _extract_words(self, segment, time_offset): + """Extracts word-level timestamps from a segment if word_timestamps is enabled.""" + if not self.word_timestamps: + return None + words = getattr(segment, "words", None) + if not words: + return None + return [ + { + "word": w.word, + "start": "{:.3f}".format(time_offset + w.start), + "end": "{:.3f}".format(time_offset + w.end), + "probability": round(w.probability, 4), + } + for w in words + ] + def update_segments(self, segments, duration): """ Processes the segments from Whisper and updates the transcript. @@ -351,7 +372,8 @@ class ServeClientBase(object): if self.get_segment_no_speech_prob(s) > self.no_speech_thresh: continue speaker = self._identify_speaker(s) - completed_segment = self.format_segment(start, end, text_, completed=True, speaker=speaker) + words = self._extract_words(s, self.timestamp_offset) + completed_segment = self.format_segment(start, end, text_, completed=True, speaker=speaker, words=words) self.transcript.append(completed_segment) if self.translation_queue: @@ -364,12 +386,14 @@ class ServeClientBase(object): # Process the last segment if its no_speech_prob is acceptable. if self.get_segment_no_speech_prob(segments[-1]) <= self.no_speech_thresh: self.current_out += segments[-1].text + words = self._extract_words(segments[-1], self.timestamp_offset) with self.lock: last_segment = self.format_segment( self.timestamp_offset + self.get_segment_start(segments[-1]), self.timestamp_offset + min(duration, self.get_segment_end(segments[-1])), self.current_out, - completed=False + completed=False, + words=words ) # Handle repeated output logic. diff --git a/whisper_live/backend/faster_whisper_backend.py b/whisper_live/backend/faster_whisper_backend.py index 02b07e4..361d489 100644 --- a/whisper_live/backend/faster_whisper_backend.py +++ b/whisper_live/backend/faster_whisper_backend.py @@ -36,6 +36,7 @@ class ServeClientFasterWhisper(ServeClientBase): translation_queue=None, hotwords=None, diarization=None, + word_timestamps=False, ): """ Initialize a ServeClient instance. @@ -67,6 +68,7 @@ class ServeClientFasterWhisper(ServeClientBase): same_output_threshold, translation_queue, diarization, + word_timestamps, ) self.cache_path = cache_path self.model_sizes = [ @@ -217,6 +219,7 @@ class ServeClientFasterWhisper(ServeClientBase): initial_prompt=self.initial_prompt, use_vad=self.use_vad, vad_parameters=self.vad_parameters if self.use_vad else None, + word_timestamps=self.word_timestamps, ) ServeClientFasterWhisper.BATCH_WORKER.submit(request) request.future.wait(timeout=30) @@ -236,7 +239,8 @@ class ServeClientFasterWhisper(ServeClientBase): task=self.task, vad_filter=self.use_vad, vad_parameters=self.vad_parameters if self.use_vad else None, - hotwords=self.hotwords) + hotwords=self.hotwords, + word_timestamps=self.word_timestamps) if ServeClientFasterWhisper.SINGLE_MODEL: ServeClientFasterWhisper.SINGLE_MODEL_LOCK.release() diff --git a/whisper_live/client.py b/whisper_live/client.py index c8ac4e5..9a271fb 100644 --- a/whisper_live/client.py +++ b/whisper_live/client.py @@ -46,6 +46,7 @@ class Client: hotwords=None, enable_diarization=False, max_speakers=10, + word_timestamps=False, ): """ Initializes a Client instance for audio recording and streaming to a server. @@ -107,6 +108,7 @@ class Client: self.hotwords = hotwords self.enable_diarization = enable_diarization self.max_speakers = max_speakers + self.word_timestamps = word_timestamps self.audio_bytes = None if host is not None and port is not None: @@ -307,6 +309,7 @@ class Client: "hotwords": self.hotwords, "enable_diarization": self.enable_diarization, "max_speakers": self.max_speakers, + "word_timestamps": self.word_timestamps, } ) ) @@ -831,6 +834,7 @@ class TranscriptionClient(TranscriptionTeeClient): hotwords=None, enable_diarization=False, max_speakers=10, + word_timestamps=False, ): self.client = Client( @@ -857,6 +861,7 @@ class TranscriptionClient(TranscriptionTeeClient): hotwords=hotwords, enable_diarization=enable_diarization, max_speakers=max_speakers, + word_timestamps=word_timestamps, ) if save_output_recording and not output_recording_filename.endswith(".wav"): diff --git a/whisper_live/server.py b/whisper_live/server.py index 66f9324..e420a47 100644 --- a/whisper_live/server.py +++ b/whisper_live/server.py @@ -293,6 +293,7 @@ class TranscriptionServer: translation_queue=translation_queue, hotwords=options.get("hotwords"), diarization=self._create_diarizer(options), + word_timestamps=options.get("word_timestamps", False), ) logging.info("Running faster_whisper backend.")