diff --git a/README.md b/README.md index 81fc92e..dc88e93 100644 --- a/README.md +++ b/README.md @@ -145,6 +145,10 @@ if scores["hey_livekit"] > 0.5: print("Wake word detected!") ``` +`WakeWordModel.predict()` is stateless. For a sequence of overlapping full +windows, `model.create_stream()` reuses only speech embeddings whose mel inputs +match exactly. `WakeWordListener` enables this optimization automatically. + **Async listener with microphone:** ```python diff --git a/docs/export-and-inference.md b/docs/export-and-inference.md index 67d26ab..e487db7 100644 --- a/docs/export-and-inference.md +++ b/docs/export-and-inference.md @@ -107,8 +107,29 @@ Feature extraction models (`melspectrogram.onnx`, `embedding_model.onnx`) are bu | Method | Returns | Description | |--------|---------|-------------| | `predict(audio_chunk)` | `dict[str, float]` | Scores for each loaded model (0-1) | +| `create_stream()` | `WakeWordStream` | Isolated cache for overlapping full audio windows | | `load_model(path, name)` | `None` | Load additional wake word model | +#### Streaming Prediction + +`WakeWordModel` remains stateless. For one continuous audio source, create a +separate stream and pass it the same complete ~2-second windows, advanced by +80 ms (1,280 samples), that you would otherwise pass to `model.predict()`: + +```python +stream = model.create_stream() + +scores = stream.predict(first_audio_window) +scores = stream.predict(next_overlapping_window) +stream.reset() # Clear cached mel windows and speech embeddings. +``` + +The stream still computes the full mel spectrogram because mel output can depend +on the complete audio window. It reuses a speech embedding only when its 76-frame +mel input exactly matches the shifted input from the previous window, then batches +all changed windows into one embedding call. Use one stream per audio source and +call it serially. `WakeWordListener` creates and resets its own stream automatically. + #### Audio Input - **Format:** 16kHz mono, int16 or float32 @@ -118,6 +139,8 @@ Feature extraction models (`melspectrogram.onnx`, `embedding_model.onnx`) are bu ### WakeWordListener The `WakeWordListener` class provides async microphone detection with debouncing. +It also uses `WakeWordModel.create_stream()` to avoid recomputing unchanged speech +embeddings between overlapping microphone windows. ```python import asyncio diff --git a/src/livekit/wakeword/__init__.py b/src/livekit/wakeword/__init__.py index 0c633c0..28a6e7d 100644 --- a/src/livekit/wakeword/__init__.py +++ b/src/livekit/wakeword/__init__.py @@ -1,7 +1,7 @@ """livekit-wakeword — Wake word detection for voice-enabled applications.""" from .inference.listener import Detection, WakeWordListener -from .inference.model import WakeWordModel +from .inference.model import WakeWordModel, WakeWordStream __version__ = "0.1.0" @@ -35,6 +35,7 @@ def __getattr__(name: str) -> object: "ExportFormat", "WakeWordListener", "WakeWordModel", + "WakeWordStream", "Detection", "load_config", "run_augment", diff --git a/src/livekit/wakeword/inference/__init__.py b/src/livekit/wakeword/inference/__init__.py index e954182..8ca7b2f 100644 --- a/src/livekit/wakeword/inference/__init__.py +++ b/src/livekit/wakeword/inference/__init__.py @@ -1,6 +1,6 @@ """Real-time inference engine.""" from .listener import Detection, WakeWordListener -from .model import WakeWordModel +from .model import WakeWordModel, WakeWordStream -__all__ = ["Detection", "WakeWordListener", "WakeWordModel"] +__all__ = ["Detection", "WakeWordListener", "WakeWordModel", "WakeWordStream"] diff --git a/src/livekit/wakeword/inference/listener.py b/src/livekit/wakeword/inference/listener.py index a7b68dd..9685b93 100644 --- a/src/livekit/wakeword/inference/listener.py +++ b/src/livekit/wakeword/inference/listener.py @@ -8,6 +8,7 @@ import time from collections import deque from dataclasses import dataclass +from typing import Protocol import numpy as np @@ -18,10 +19,14 @@ SAMPLE_RATE = 16000 FRAME_SAMPLES = 1280 # 80ms per frame CHUNK_SECONDS = 2.0 -# Number of frames that fill a ~2-second chunk (25 × 80ms = 2000ms) +# Number of frames that fill a ~2-second chunk (25 * 80ms = 2000ms) CHUNK_FRAMES = int(CHUNK_SECONDS * SAMPLE_RATE / FRAME_SAMPLES) +class _Predictor(Protocol): + def predict(self, audio_chunk: np.ndarray) -> dict[str, float]: ... + + @dataclass class Detection: """Wake word detection result.""" @@ -35,9 +40,10 @@ class WakeWordListener: """Async wake word listener that handles audio capture. The listener owns the audio buffer and passes fixed ~2-second chunks - to the stateless ``WakeWordModel.predict()``. After a detection the - loop pauses automatically and resumes when the consumer calls - ``wait_for_detection()`` again. + through an isolated stream predictor. The predictor reuses unchanged + speech embeddings while ``WakeWordModel`` remains stateless. After a + detection the loop pauses automatically and resumes when the consumer + calls ``wait_for_detection()`` again. Example: from livekit.wakeword import WakeWordModel, WakeWordListener @@ -64,6 +70,7 @@ def __init__( debounce: Minimum seconds between detections. """ self._model = model + self._predictor: _Predictor = model self._threshold = threshold self._debounce = debounce @@ -95,6 +102,8 @@ async def __aenter__(self) -> WakeWordListener: """Start audio capture.""" import pyaudio + create_stream = getattr(self._model, "create_stream", None) + self._predictor = create_stream() if callable(create_stream) else self._model self._pa = pyaudio.PyAudio() self._stream = self._pa.open( format=pyaudio.paInt16, @@ -124,7 +133,7 @@ async def __aexit__(self, *_: object) -> None: if self._task: try: await asyncio.wait_for(self._task, timeout=2.0) - except (asyncio.TimeoutError, asyncio.CancelledError): + except (TimeoutError, asyncio.CancelledError): self._task.cancel() try: await self._task @@ -175,7 +184,7 @@ async def _audio_loop(self) -> None: chunk = np.concatenate(list(self._frame_buffer)) scores = await loop.run_in_executor( self._executor, - self._model.predict, + self._predictor.predict, chunk, ) if not self._running: @@ -193,11 +202,12 @@ async def _audio_loop(self) -> None: # the detection. self._listening.clear() self._frame_buffer.clear() + reset = getattr(self._predictor, "reset", None) + if callable(reset): + await loop.run_in_executor(self._executor, reset) await self._detection_queue.put( - Detection( - name=name, confidence=score, timestamp=now - ) + Detection(name=name, confidence=score, timestamp=now) ) break # one detection per iteration except asyncio.CancelledError: @@ -243,8 +253,6 @@ async def wait_for_detection(self) -> Detection: return self._detection_queue.get_nowait() if self._error is not None: - raise RuntimeError( - f"Audio loop crashed: {self._error}" - ) from self._error + raise RuntimeError(f"Audio loop crashed: {self._error}") from self._error raise RuntimeError("Audio loop ended unexpectedly") diff --git a/src/livekit/wakeword/inference/model.py b/src/livekit/wakeword/inference/model.py index 4f16e02..6b6e4f0 100644 --- a/src/livekit/wakeword/inference/model.py +++ b/src/livekit/wakeword/inference/model.py @@ -1,4 +1,4 @@ -"""Stateless wake word detection model.""" +"""Wake word detection model with optional per-stream caching.""" from __future__ import annotations @@ -61,7 +61,10 @@ def __init__( ) self._mel_frontend = MelSpectrogramFrontend(onnx_path=mel_path, sess_options=sess_options) - self._speech_embedding = SpeechEmbedding(onnx_path=embedding_path, sess_options=sess_options) + self._speech_embedding = SpeechEmbedding( + onnx_path=embedding_path, + sess_options=sess_options, + ) # name -> (onnx_session, input_name) self._classifiers: dict[str, tuple] = {} @@ -70,7 +73,12 @@ def __init__( for model_path in models: self.load_model(model_path, sess_options=sess_options) - def load_model(self, model_path: str | Path, model_name: str | None = None, sess_options: SessionOptions = None ) -> None: + def load_model( + self, + model_path: str | Path, + model_name: str | None = None, + sess_options: SessionOptions | None = None, + ) -> None: """Load a wake word classifier model. Args: @@ -89,12 +97,20 @@ def load_model(self, model_path: str | Path, model_name: str | None = None, sess session = ort.InferenceSession( str(model_path), providers=["CPUExecutionProvider"], - sess_options=sess_options + sess_options=sess_options, ) input_name = session.get_inputs()[0].name self._classifiers[model_name] = (session, input_name) logger.info(f"Loaded wake word model '{model_name}' from {model_path}") + def create_stream(self) -> WakeWordStream: + """Create an isolated predictor for overlapping audio windows. + + The returned stream caches speech embeddings. The model itself remains + stateless, and each stream owns independent cache state. + """ + return WakeWordStream(self) + def predict(self, audio_chunk: np.ndarray) -> dict[str, float]: """Get wake word predictions for an audio chunk. @@ -112,38 +128,104 @@ def predict(self, audio_chunk: np.ndarray) -> dict[str, float]: if not self._classifiers: return {} - # Convert int16 to float32 if needed + windows = self._mel_windows(audio_chunk) + if windows is None: + return self._zero_scores() + + embeddings = np.stack( + [self._speech_embedding(window[np.newaxis, :, :])[0] for window in windows], + axis=0, + ) + return self._classify(embeddings) + + def _mel_windows(self, audio_chunk: np.ndarray) -> np.ndarray | None: if audio_chunk.dtype == np.int16: audio_chunk = audio_chunk.astype(np.float32) / 32768.0 - audio_chunk = audio_chunk.flatten() - - # Mel spectrogram over the full chunk - all_mel = self._mel_frontend(audio_chunk) + all_mel = self._mel_frontend(audio_chunk.flatten()) if all_mel.ndim == 3: all_mel = all_mel[0] - if all_mel.shape[0] < EMBEDDING_WINDOW: - return {name: 0.0 for name in self._classifiers} + starts = range( + 0, + all_mel.shape[0] - EMBEDDING_WINDOW + 1, + EMBEDDING_STRIDE, + ) + windows = [all_mel[start : start + EMBEDDING_WINDOW] for start in starts] + if len(windows) < MIN_EMBEDDINGS: + return None + return np.stack(windows[-MIN_EMBEDDINGS:], axis=0) + + def _classify(self, embeddings: np.ndarray) -> dict[str, float]: + emb_input = embeddings[np.newaxis, :, :].astype(np.float32) + predictions: dict[str, float] = {} + for name, (session, input_name) in self._classifiers.items(): + outputs = session.run(None, {input_name: emb_input}) + predictions[name] = float(outputs[0][0, 0]) + return predictions + + def _zero_scores(self) -> dict[str, float]: + return {name: 0.0 for name in self._classifiers} - # Extract embeddings: 76-frame windows, stride 8 - embeddings = [] - for start in range(0, all_mel.shape[0] - EMBEDDING_WINDOW + 1, EMBEDDING_STRIDE): - window = all_mel[start : start + EMBEDDING_WINDOW] - emb = self._speech_embedding(window[np.newaxis, :, :]) - embeddings.append(emb[0]) - if len(embeddings) < MIN_EMBEDDINGS: - return {name: 0.0 for name in self._classifiers} +class WakeWordStream: + """Prediction state for a sequence of overlapping full audio windows. - # Use last 16 embeddings - emb_sequence = np.stack(embeddings[-MIN_EMBEDDINGS:], axis=0) - emb_input = emb_sequence[np.newaxis, :, :].astype(np.float32) + Each call still computes the full mel spectrogram because that model's + output can depend on the complete audio window. Speech embeddings are reused + only when their 76-frame mel inputs exactly match the shifted prior inputs. + """ - predictions = {} - for name, (session, input_name) in self._classifiers.items(): - outputs = session.run(None, {input_name: emb_input}) - score = float(outputs[0][0, 0]) - predictions[name] = score + def __init__(self, model: WakeWordModel): + self._model = model + self._mel_windows: np.ndarray | None = None + self._embeddings: np.ndarray | None = None - return predictions + def reset(self) -> None: + """Clear cached windows and embeddings.""" + self._mel_windows = None + self._embeddings = None + + def predict(self, audio_chunk: np.ndarray) -> dict[str, float]: + """Predict from the next complete audio window in the stream.""" + if not self._model._classifiers: + return {} + + windows = self._model._mel_windows(audio_chunk) + if windows is None: + self.reset() + return self._model._zero_scores() + + embeddings = self._embedding_sequence(windows) + self._mel_windows = windows + self._embeddings = embeddings + return self._model._classify(embeddings) + + def _embedding_sequence(self, windows: np.ndarray) -> np.ndarray: + changed_indexes: list[int] = [] + changed_windows: list[np.ndarray] = [] + reused: dict[int, np.ndarray] = {} + + for index, window in enumerate(windows): + previous_index = index + 1 + if ( + self._mel_windows is not None + and self._embeddings is not None + and previous_index < len(self._mel_windows) + and np.array_equal(window, self._mel_windows[previous_index]) + ): + reused[index] = self._embeddings[previous_index] + else: + changed_indexes.append(index) + changed_windows.append(window) + + fresh = self._model._speech_embedding(np.stack(changed_windows, axis=0)) + embeddings = np.empty( + (len(windows), fresh.shape[1]), + dtype=fresh.dtype, + ) + for index, embedding in reused.items(): + embeddings[index] = embedding + for index, embedding in zip(changed_indexes, fresh, strict=True): + embeddings[index] = embedding + return embeddings diff --git a/tests/test_listener.py b/tests/test_listener.py index 16151e4..803d6bc 100644 --- a/tests/test_listener.py +++ b/tests/test_listener.py @@ -11,12 +11,10 @@ from livekit.wakeword.inference.listener import ( CHUNK_FRAMES, - Detection, - WakeWordListener, FRAME_SAMPLES, + WakeWordListener, ) - # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -43,6 +41,31 @@ def predict(self, audio_chunk: np.ndarray) -> dict[str, float]: return scores +class FakeStreamingPredictor(FakeModel): + def __init__(self, scores_sequence: list[dict[str, float]]): + super().__init__(scores_sequence) + self.chunks: list[np.ndarray] = [] + self.reset_count = 0 + + def predict(self, audio_chunk: np.ndarray) -> dict[str, float]: + self.chunks.append(audio_chunk.copy()) + return super().predict(audio_chunk) + + def reset(self) -> None: + self.reset_count += 1 + + +class FakeStreamingModel: + def __init__(self, scores_sequence: list[dict[str, float]]): + self._scores_sequence = scores_sequence + self.streams: list[FakeStreamingPredictor] = [] + + def create_stream(self) -> FakeStreamingPredictor: + stream = FakeStreamingPredictor(self._scores_sequence) + self.streams.append(stream) + return stream + + class FakeStream: """Mock PyAudio stream that returns silence.""" @@ -53,7 +76,7 @@ def __init__(self, *, error_after: int | None = None): def read(self, num_frames: int, exception_on_overflow: bool = False) -> bytes: self._read_count += 1 if self._error_after is not None and self._read_count > self._error_after: - raise IOError("Simulated stream read error") + raise OSError("Simulated stream read error") return np.zeros(num_frames, dtype=np.int16).tobytes() def stop_stream(self) -> None: @@ -63,6 +86,14 @@ def close(self) -> None: pass +class CountingFakeStream(FakeStream): + """Mock stream whose frame values expose the sliding-window overlap.""" + + def read(self, num_frames: int, exception_on_overflow: bool = False) -> bytes: + self._read_count += 1 + return np.full(num_frames, self._read_count, dtype=np.int16).tobytes() + + class FakePyAudio: """Mock PyAudio that returns a FakeStream.""" @@ -90,6 +121,51 @@ def _patch_pyaudio(stream: FakeStream): # --------------------------------------------------------------------------- +@pytest.mark.asyncio +async def test_listener_uses_stream_predictor_for_overlapping_windows(): + model = FakeStreamingModel([{"test": 0.0}, {"test": 0.9}]) + audio = CountingFakeStream() + + with _patch_pyaudio(audio): + async with WakeWordListener(model, threshold=0.5, debounce=0.0) as listener: + await asyncio.wait_for(listener.wait_for_detection(), timeout=5.0) + + predictor = model.streams[0] + assert len(predictor.chunks) == 2 + assert np.array_equal( + predictor.chunks[0][FRAME_SAMPLES:], + predictor.chunks[1][:-FRAME_SAMPLES], + ) + + +@pytest.mark.asyncio +async def test_detection_resets_stream_cache(): + model = FakeStreamingModel([{"test": 0.9}]) + audio = FakeStream() + + with _patch_pyaudio(audio): + async with WakeWordListener(model, threshold=0.5, debounce=0.0) as listener: + await asyncio.wait_for(listener.wait_for_detection(), timeout=5.0) + + assert model.streams[0].reset_count == 1 + + +@pytest.mark.asyncio +async def test_context_reentry_creates_fresh_stream_state(): + model = FakeStreamingModel([{"test": 0.9}]) + audio = FakeStream() + listener = WakeWordListener(model, threshold=0.5, debounce=0.0) + + with _patch_pyaudio(audio): + async with listener: + await asyncio.wait_for(listener.wait_for_detection(), timeout=5.0) + async with listener: + await asyncio.wait_for(listener.wait_for_detection(), timeout=5.0) + + assert len(model.streams) == 2 + assert all(stream._call_count == 1 for stream in model.streams) + + @pytest.mark.asyncio async def test_predict_runs_in_executor(): """predict() should execute on a non-main thread (in the executor).""" @@ -104,7 +180,7 @@ def predict(self, audio_chunk: np.ndarray) -> dict[str, float]: stream = FakeStream() with _patch_pyaudio(stream): - async with WakeWordListener(model, threshold=0.5) as listener: + async with WakeWordListener(model, threshold=0.5): await asyncio.sleep(0.5) assert len(predict_threads) > 0 @@ -153,9 +229,7 @@ async def test_buffer_cleared_after_detection(): with _patch_pyaudio(stream): async with WakeWordListener(model, threshold=0.5, debounce=0.0) as listener: - detection = await asyncio.wait_for( - listener.wait_for_detection(), timeout=5.0 - ) + detection = await asyncio.wait_for(listener.wait_for_detection(), timeout=5.0) assert detection.name == "test" assert detection.confidence == pytest.approx(0.9) # Buffer should be empty after detection @@ -188,9 +262,7 @@ async def test_loop_pauses_after_detection(): with _patch_pyaudio(stream): async with WakeWordListener(model, threshold=0.5, debounce=0.0) as listener: - det1 = await asyncio.wait_for( - listener.wait_for_detection(), timeout=5.0 - ) + det1 = await asyncio.wait_for(listener.wait_for_detection(), timeout=5.0) # Give time for the loop to potentially queue more (it shouldn't) await asyncio.sleep(0.3) @@ -199,9 +271,7 @@ async def test_loop_pauses_after_detection(): ) # Second call resumes listening and gets next detection - det2 = await asyncio.wait_for( - listener.wait_for_detection(), timeout=5.0 - ) + det2 = await asyncio.wait_for(listener.wait_for_detection(), timeout=5.0) assert det2.timestamp > det1.timestamp @@ -227,6 +297,6 @@ async def test_shutdown_during_active_listening(): stream = FakeStream() with _patch_pyaudio(stream): - async with WakeWordListener(model, threshold=0.5) as listener: + async with WakeWordListener(model, threshold=0.5): await asyncio.sleep(0.2) # let loop run a few iterations # If we reach here without hanging, shutdown is clean. diff --git a/tests/test_stream.py b/tests/test_stream.py new file mode 100644 index 0000000..287c7c9 --- /dev/null +++ b/tests/test_stream.py @@ -0,0 +1,157 @@ +"""Tests for cached prediction over overlapping wake-word windows.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import onnxruntime as ort + +from livekit.wakeword import WakeWordModel, WakeWordStream + + +class FakeMelFrontend: + def __init__(self, outputs: list[np.ndarray]): + self._outputs = iter(outputs) + self.call_count = 0 + + def __call__(self, _audio: np.ndarray) -> np.ndarray: + self.call_count += 1 + return next(self._outputs)[np.newaxis, :, :] + + +class FakeSpeechEmbedding: + def __init__(self) -> None: + self.batch_sizes: list[int] = [] + + def __call__(self, windows: np.ndarray) -> np.ndarray: + self.batch_sizes.append(windows.shape[0]) + values = windows[:, 0, 0].astype(np.float32) + return np.repeat(values[:, np.newaxis], 96, axis=1) + + +@dataclass +class FakeInput: + name: str = "features" + + +class FakeClassifier: + def __init__(self) -> None: + self.call_count = 0 + + def get_inputs(self) -> list[FakeInput]: + return [FakeInput()] + + def run(self, _outputs: object, feeds: dict[str, np.ndarray]) -> list[np.ndarray]: + self.call_count += 1 + score = float(feeds["features"].mean()) + return [np.array([[score]], dtype=np.float32)] + + +def mel_frames() -> np.ndarray: + return np.arange(197 * 32, dtype=np.float32).reshape(197, 32) + + +def shifted_mel(previous: np.ndarray) -> np.ndarray: + new_frames = previous[-8:] + 100_000 + return np.vstack((previous[8:], new_frames)) + + +def fake_model(mels: list[np.ndarray]) -> tuple[WakeWordModel, FakeSpeechEmbedding, FakeClassifier]: + model = WakeWordModel.__new__(WakeWordModel) + model._mel_frontend = FakeMelFrontend(mels) + embedding = FakeSpeechEmbedding() + model._speech_embedding = embedding + classifier = FakeClassifier() + model._classifiers = {"test": (classifier, "features")} + return model, embedding, classifier + + +def test_first_stream_prediction_embeds_all_windows_in_one_batch(): + model, embedding, classifier = fake_model([mel_frames()]) + + stream = model.create_stream() + scores = stream.predict(np.zeros(32_000, dtype=np.int16)) + + assert isinstance(stream, WakeWordStream) + assert embedding.batch_sizes == [16] + assert classifier.call_count == 1 + assert set(scores) == {"test"} + + +def test_shifted_prediction_reuses_fifteen_embeddings(): + first = mel_frames() + model, embedding, classifier = fake_model([first, shifted_mel(first)]) + stream = model.create_stream() + + stream.predict(np.zeros(32_000, dtype=np.int16)) + stream.predict(np.zeros(32_000, dtype=np.int16)) + + assert model._mel_frontend.call_count == 2 + assert embedding.batch_sizes == [16, 1] + assert classifier.call_count == 2 + + +def test_changed_overlapping_window_is_recomputed(): + first = mel_frames() + second = shifted_mel(first) + second[0, 0] += 1 + model, embedding, _classifier = fake_model([first, second]) + stream = model.create_stream() + + stream.predict(np.zeros(32_000, dtype=np.int16)) + stream.predict(np.zeros(32_000, dtype=np.int16)) + + # The modified first window and the newly arrived final window both change. + assert embedding.batch_sizes == [16, 2] + + +def test_streams_from_one_model_do_not_share_cache_state(): + first = mel_frames() + model, embedding, _classifier = fake_model([first, first]) + + model.create_stream().predict(np.zeros(32_000, dtype=np.int16)) + model.create_stream().predict(np.zeros(32_000, dtype=np.int16)) + + assert embedding.batch_sizes == [16, 16] + + +def test_reset_forces_the_next_prediction_to_rebuild_all_embeddings(): + first = mel_frames() + second = shifted_mel(first) + model, embedding, _classifier = fake_model([first, second, second]) + stream = model.create_stream() + + stream.predict(np.zeros(32_000, dtype=np.int16)) + stream.predict(np.zeros(32_000, dtype=np.int16)) + stream.reset() + stream.predict(np.zeros(32_000, dtype=np.int16)) + + assert embedding.batch_sizes == [16, 1, 16] + + +def test_stateless_prediction_keeps_individual_embedding_calls(): + model, embedding, classifier = fake_model([mel_frames()]) + + model.predict(np.zeros(32_000, dtype=np.int16)) + + assert embedding.batch_sizes == [1] * 16 + assert classifier.call_count == 1 + + +def test_stream_matches_stateless_scores_with_real_feature_models(): + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + options.inter_op_num_threads = 1 + model = WakeWordModel(sess_options=options) + classifier = FakeClassifier() + model._classifiers = {"test": (classifier, "features")} + stream = model.create_stream() + + rng = np.random.default_rng(2820) + audio = np.zeros(32_000 + 4 * 1_280, dtype=np.int16) + audio[30_000:34_000] = rng.integers(-20_000, 20_000, 4_000, dtype=np.int16) + chunks = [audio[start : start + 32_000] for start in range(0, 4 * 1_280, 1_280)] + + for chunk in chunks: + assert stream.predict(chunk) == model.predict(chunk)