Add word-level timestamps and confidence scores

- 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
This commit is contained in:
Aaron Boxer
2026-04-17 10:18:19 -04:00
committed by Aaron Boxer
parent 18de3eacc7
commit 4e31f8c61b
6 changed files with 162 additions and 22 deletions
+105
View File
@@ -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()
+18 -17
View File
@@ -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(
+28 -4
View File
@@ -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.
@@ -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()
+5
View File
@@ -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"):
+1
View File
@@ -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.")