fix: address PR review comments on StreamingTranscriptionClient
This commit is contained in:
@@ -265,7 +265,7 @@ Audio is automatically normalized to float32 range [-1.0, 1.0]. Clients can also
|
|||||||
|
|
||||||
`StreamingTranscriptionClient` lets you push raw PCM audio bytes from any source — a live microphone capture loop, a network stream, an audio pipeline — and receive transcripts via callbacks as speech is detected. Unlike `TranscriptionClient`, it does not manage audio capture internally; you control when and how audio is fed.
|
`StreamingTranscriptionClient` lets you push raw PCM audio bytes from any source — a live microphone capture loop, a network stream, an audio pipeline — and receive transcripts via callbacks as speech is detected. Unlike `TranscriptionClient`, it does not manage audio capture internally; you control when and how audio is fed.
|
||||||
|
|
||||||
A runnable example that reads from an audio file and stream the chunks is at [`examples/manual_audio_chunking.py`](examples/manual_audio_chunking.py):
|
A runnable example that reads from an audio file and streams the chunks is at [`examples/manual_audio_chunking.py`](examples/manual_audio_chunking.py):
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
python examples/manual_audio_chunking.py --file assets/jfk.flac
|
python examples/manual_audio_chunking.py --file assets/jfk.flac
|
||||||
|
|||||||
@@ -69,8 +69,8 @@ def main():
|
|||||||
print("\n[INFO] Final transcript:")
|
print("\n[INFO] Final transcript:")
|
||||||
for seg in client.transcript:
|
for seg in client.transcript:
|
||||||
print(f" [{float(seg['start']):.2f}s → {float(seg['end']):.2f}s] {seg['text'].strip()}")
|
print(f" [{float(seg['start']):.2f}s → {float(seg['end']):.2f}s] {seg['text'].strip()}")
|
||||||
if client.last_partial:
|
seg = client.last_partial
|
||||||
seg = client.last_segment
|
if seg:
|
||||||
print(f" [{float(seg['start']):.2f}s → {float(seg['end']):.2f}s] {seg['text'].strip()} (partial)")
|
print(f" [{float(seg['start']):.2f}s → {float(seg['end']):.2f}s] {seg['text'].strip()} (partial)")
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,164 @@
|
|||||||
|
import json
|
||||||
|
import unittest
|
||||||
|
from unittest.mock import patch, MagicMock
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
from whisper_live.client import Client, StreamingTranscriptionClient
|
||||||
|
|
||||||
|
|
||||||
|
class StreamingClientTestCase(unittest.TestCase):
|
||||||
|
@patch('whisper_live.client.websocket.WebSocketApp')
|
||||||
|
def setUp(self, mock_websocket):
|
||||||
|
self.mock_websocket = mock_websocket
|
||||||
|
self.mock_ws_app = mock_websocket.return_value
|
||||||
|
self.mock_ws_app.send = MagicMock()
|
||||||
|
|
||||||
|
self.committed = []
|
||||||
|
self.partials = []
|
||||||
|
self.session_started = []
|
||||||
|
|
||||||
|
self.client = StreamingTranscriptionClient(
|
||||||
|
host='localhost',
|
||||||
|
port=9090,
|
||||||
|
lang="en",
|
||||||
|
on_session_started=lambda: self.session_started.append(True),
|
||||||
|
on_committed_transcript=lambda text, segs: self.committed.append((text, segs)),
|
||||||
|
on_partial_transcript=lambda text, segs: self.partials.append((text, segs)),
|
||||||
|
)
|
||||||
|
self._inner = self.client._client
|
||||||
|
|
||||||
|
def tearDown(self):
|
||||||
|
self._inner.close_websocket()
|
||||||
|
self.mock_websocket.stop()
|
||||||
|
|
||||||
|
def _server_ready(self, backend="faster_whisper"):
|
||||||
|
self._inner.on_message(self.mock_ws_app, json.dumps({
|
||||||
|
"uid": self._inner.uid,
|
||||||
|
"message": "SERVER_READY",
|
||||||
|
"backend": backend,
|
||||||
|
}))
|
||||||
|
|
||||||
|
def _send_segments(self, segments):
|
||||||
|
self._inner.on_message(self.mock_ws_app, json.dumps({
|
||||||
|
"uid": self._inner.uid,
|
||||||
|
"segments": segments,
|
||||||
|
}))
|
||||||
|
|
||||||
|
|
||||||
|
class TestPcmFormatConversion(StreamingClientTestCase):
|
||||||
|
def test_int16_is_normalized_to_float32(self):
|
||||||
|
self._server_ready()
|
||||||
|
raw = np.array([0, 16384, -32768], dtype=np.int16).tobytes()
|
||||||
|
with patch.object(self._inner, 'send_packet_to_server') as mock_send:
|
||||||
|
self.client.send(raw, pcm_format="int16")
|
||||||
|
sent = np.frombuffer(mock_send.call_args[0][0], dtype=np.float32)
|
||||||
|
np.testing.assert_allclose(sent, [0.0, 0.5, -1.0], atol=1e-4)
|
||||||
|
|
||||||
|
def test_float32_passes_through(self):
|
||||||
|
self._server_ready()
|
||||||
|
raw = np.array([0.1, -0.2], dtype=np.float32).tobytes()
|
||||||
|
with patch.object(self._inner, 'send_packet_to_server') as mock_send:
|
||||||
|
self.client.send(raw, pcm_format="float32")
|
||||||
|
self.assertEqual(mock_send.call_args[0][0], raw)
|
||||||
|
|
||||||
|
def test_default_format_is_int16(self):
|
||||||
|
self._server_ready()
|
||||||
|
raw = np.array([32767], dtype=np.int16).tobytes()
|
||||||
|
with patch.object(self._inner, 'send_packet_to_server') as mock_send:
|
||||||
|
self.client.send(raw)
|
||||||
|
sent = np.frombuffer(mock_send.call_args[0][0], dtype=np.float32)
|
||||||
|
self.assertAlmostEqual(float(sent[0]), 32767 / 32768.0, places=4)
|
||||||
|
|
||||||
|
def test_unsupported_format_raises(self):
|
||||||
|
self._server_ready()
|
||||||
|
with self.assertRaises(ValueError):
|
||||||
|
self.client.send(b"\x00\x00", pcm_format="int8")
|
||||||
|
|
||||||
|
def test_send_after_close_raises(self):
|
||||||
|
self.client._closed = True
|
||||||
|
with self.assertRaises(RuntimeError):
|
||||||
|
self.client.send(b"\x00\x00", pcm_format="int16")
|
||||||
|
|
||||||
|
def test_send_array_normalizes_integers(self):
|
||||||
|
with patch.object(self._inner, 'send_packet_to_server') as mock_send:
|
||||||
|
self.client.send_array(np.array([0, 16384, -32768], dtype=np.int16))
|
||||||
|
sent = np.frombuffer(mock_send.call_args[0][0], dtype=np.float32)
|
||||||
|
np.testing.assert_allclose(sent, [0.0, 0.5, -1.0], atol=1e-4)
|
||||||
|
|
||||||
|
|
||||||
|
class TestTranscriptDispatch(StreamingClientTestCase):
|
||||||
|
def test_partial_then_committed(self):
|
||||||
|
self._server_ready()
|
||||||
|
self._send_segments([{"start": 0, "end": 1, "text": "hello", "completed": False}])
|
||||||
|
self.assertEqual(len(self.partials), 1)
|
||||||
|
self.assertEqual(self.partials[0][0], "hello")
|
||||||
|
self.assertEqual(len(self.committed), 0)
|
||||||
|
|
||||||
|
self._send_segments([{"start": 0, "end": 1, "text": "hello world", "completed": True}])
|
||||||
|
self.assertEqual(len(self.committed), 1)
|
||||||
|
self.assertEqual(self.committed[0][0], "hello world")
|
||||||
|
self.assertEqual(len(self.client.transcript), 1)
|
||||||
|
|
||||||
|
def test_committed_deduplicated(self):
|
||||||
|
self._server_ready()
|
||||||
|
seg = {"start": 0, "end": 1, "text": "hi", "completed": True}
|
||||||
|
self._send_segments([seg])
|
||||||
|
self._send_segments([seg])
|
||||||
|
self.assertEqual(len(self.committed), 1)
|
||||||
|
self.assertEqual(len(self.client.transcript), 1)
|
||||||
|
|
||||||
|
def test_committed_backend_agnostic(self):
|
||||||
|
"""Committed dispatch must work for non-faster_whisper backends."""
|
||||||
|
self._server_ready(backend="tensorrt")
|
||||||
|
self._send_segments([{"start": 0, "end": 1, "text": "trt seg", "completed": True}])
|
||||||
|
self.assertEqual(len(self.committed), 1)
|
||||||
|
self.assertEqual(len(self.client.transcript), 1)
|
||||||
|
|
||||||
|
def test_last_partial_alias(self):
|
||||||
|
self._server_ready()
|
||||||
|
self._send_segments([{"start": 0, "end": 1, "text": "pending", "completed": False}])
|
||||||
|
self.assertIsNotNone(self.client.last_partial)
|
||||||
|
self.assertIs(self.client.last_partial, self.client.last_segment)
|
||||||
|
|
||||||
|
|
||||||
|
class TestConnectLifecycle(StreamingClientTestCase):
|
||||||
|
def test_connect_returns_after_ready(self):
|
||||||
|
self._server_ready()
|
||||||
|
self.assertIs(self.client.connect(), self.client)
|
||||||
|
self.assertEqual(len(self.session_started), 1)
|
||||||
|
|
||||||
|
def test_connect_times_out(self):
|
||||||
|
self.client._ready_timeout = 0.1
|
||||||
|
with self.assertRaises(TimeoutError):
|
||||||
|
self.client.connect()
|
||||||
|
|
||||||
|
def test_connect_raises_on_server_error(self):
|
||||||
|
self._inner.on_message(self.mock_ws_app, json.dumps({
|
||||||
|
"uid": self._inner.uid,
|
||||||
|
"status": "ERROR",
|
||||||
|
"message": "boom",
|
||||||
|
}))
|
||||||
|
with self.assertRaises(RuntimeError):
|
||||||
|
self.client.connect()
|
||||||
|
|
||||||
|
def test_connect_raises_when_server_full(self):
|
||||||
|
self._inner.on_message(self.mock_ws_app, json.dumps({
|
||||||
|
"uid": self._inner.uid,
|
||||||
|
"status": "WAIT",
|
||||||
|
"message": 5,
|
||||||
|
}))
|
||||||
|
with self.assertRaises(RuntimeError):
|
||||||
|
self.client.connect()
|
||||||
|
|
||||||
|
def test_close_sends_end_of_audio(self):
|
||||||
|
self._server_ready()
|
||||||
|
with patch.object(self._inner, 'send_packet_to_server') as mock_send, \
|
||||||
|
patch.object(self._inner, 'close_websocket') as mock_close:
|
||||||
|
self.client.close(drain_seconds=0)
|
||||||
|
mock_send.assert_called_once_with(Client.END_OF_AUDIO.encode("utf-8"))
|
||||||
|
mock_close.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
unittest.main()
|
||||||
+18
-14
@@ -927,10 +927,10 @@ class _HookedClient(Client):
|
|||||||
super().__init__(*args, **kwargs)
|
super().__init__(*args, **kwargs)
|
||||||
|
|
||||||
def on_message(self, ws, message):
|
def on_message(self, ws, message):
|
||||||
data = json.loads(message)
|
was_recording = self.recording
|
||||||
if data.get("message") == "SERVER_READY" and self._on_session_started:
|
|
||||||
self._on_session_started()
|
|
||||||
super().on_message(ws, message)
|
super().on_message(ws, message)
|
||||||
|
if not was_recording and self.recording and self._on_session_started:
|
||||||
|
self._on_session_started()
|
||||||
|
|
||||||
def on_error(self, ws, error):
|
def on_error(self, ws, error):
|
||||||
if self._on_error_hook:
|
if self._on_error_hook:
|
||||||
@@ -1007,7 +1007,8 @@ class StreamingTranscriptionClient:
|
|||||||
self._on_committed_transcript = on_committed_transcript
|
self._on_committed_transcript = on_committed_transcript
|
||||||
self._ready_timeout = ready_timeout
|
self._ready_timeout = ready_timeout
|
||||||
self._closed = False
|
self._closed = False
|
||||||
self._last_committed_count = 0
|
self._transcript = []
|
||||||
|
self._committed_keys = set()
|
||||||
|
|
||||||
self._client = _HookedClient(
|
self._client = _HookedClient(
|
||||||
host=host,
|
host=host,
|
||||||
@@ -1031,11 +1032,16 @@ class StreamingTranscriptionClient:
|
|||||||
)
|
)
|
||||||
|
|
||||||
def _dispatch_transcript(self, text: str, segments: list) -> None:
|
def _dispatch_transcript(self, text: str, segments: list) -> None:
|
||||||
new_committed = self._client.transcript[self._last_committed_count:]
|
for seg in segments:
|
||||||
for seg in new_committed:
|
if not seg.get("completed", False):
|
||||||
|
continue
|
||||||
|
key = (seg.get("start"), seg.get("end"), seg.get("text"))
|
||||||
|
if key in self._committed_keys:
|
||||||
|
continue
|
||||||
|
self._committed_keys.add(key)
|
||||||
|
self._transcript.append(seg)
|
||||||
if self._on_committed_transcript:
|
if self._on_committed_transcript:
|
||||||
self._on_committed_transcript(seg["text"].strip(), [seg])
|
self._on_committed_transcript(seg["text"].strip(), [seg])
|
||||||
self._last_committed_count = len(self._client.transcript)
|
|
||||||
|
|
||||||
last = segments[-1] if segments else None
|
last = segments[-1] if segments else None
|
||||||
if last and not last.get("completed", False) and self._on_partial_transcript:
|
if last and not last.get("completed", False) and self._on_partial_transcript:
|
||||||
@@ -1054,12 +1060,12 @@ class StreamingTranscriptionClient:
|
|||||||
time.sleep(0.05)
|
time.sleep(0.05)
|
||||||
return self
|
return self
|
||||||
|
|
||||||
def send(self, audio_bytes: bytes, pcm_format: PcmFormat = "float32") -> None:
|
def send(self, audio_bytes: bytes, pcm_format: PcmFormat = "int16") -> None:
|
||||||
"""Send one PCM chunk. Any chunk size is fine; must be mono 16 kHz.
|
"""Send one PCM chunk. Any chunk size is fine; must be mono 16 kHz.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
audio_bytes: Raw PCM payload.
|
audio_bytes: Raw PCM payload.
|
||||||
pcm_format: ``"float32"`` passes through; ``"int16"`` is normalized to float32.
|
pcm_format: ``"int16"`` is normalized to float32; ``"float32"`` passes through.
|
||||||
"""
|
"""
|
||||||
if self._closed:
|
if self._closed:
|
||||||
raise RuntimeError("Client is already closed.")
|
raise RuntimeError("Client is already closed.")
|
||||||
@@ -1092,17 +1098,15 @@ class StreamingTranscriptionClient:
|
|||||||
@property
|
@property
|
||||||
def transcript(self) -> list:
|
def transcript(self) -> list:
|
||||||
"""All committed segments received so far."""
|
"""All committed segments received so far."""
|
||||||
return self._client.transcript
|
return self._transcript
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def last_partial(self) -> Optional[dict]:
|
def last_partial(self) -> Optional[dict]:
|
||||||
"""The most recent in-progress segment, or ``None`` if none pending."""
|
"""The most recent in-progress segment, or ``None`` if none pending."""
|
||||||
return self._client.last_segment
|
return self._client.last_segment
|
||||||
|
|
||||||
@property
|
# Alias for ``last_partial``; kept for readability at call sites.
|
||||||
def last_segment(self) -> Optional[dict]:
|
last_segment = last_partial
|
||||||
"""The most recent in-progress segment, or ``None`` if none pending."""
|
|
||||||
return self._client.last_segment
|
|
||||||
|
|
||||||
def close(self, drain_seconds: float = 2.0) -> None:
|
def close(self, drain_seconds: float = 2.0) -> None:
|
||||||
"""Signal end-of-stream, wait briefly for final transcripts, then close.
|
"""Signal end-of-stream, wait briefly for final transcripts, then close.
|
||||||
|
|||||||
Reference in New Issue
Block a user