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:
@@ -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
@@ -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(
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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"):
|
||||
|
||||
@@ -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.")
|
||||
|
||||
Reference in New Issue
Block a user