diff --git a/tests/test_client.py b/tests/test_client.py index 2808648..b636000 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -55,6 +55,7 @@ class TestClientCallbacks(BaseTestCase): "same_output_threshold": 10, "enable_translation": False, "target_language": "fr", + "hotwords": None, }) self.client.on_open(self.mock_ws_app) self.mock_ws_app.send.assert_called_with(expected_message) diff --git a/whisper_live/backend/faster_whisper_backend.py b/whisper_live/backend/faster_whisper_backend.py index 67463ca..76e498b 100644 --- a/whisper_live/backend/faster_whisper_backend.py +++ b/whisper_live/backend/faster_whisper_backend.py @@ -34,6 +34,7 @@ class ServeClientFasterWhisper(ServeClientBase): same_output_threshold=7, cache_path="~/.cache/whisper-live/", translation_queue=None, + hotwords=None, ): """ Initialize a ServeClient instance. @@ -78,6 +79,7 @@ class ServeClientFasterWhisper(ServeClientBase): self.task = task self.initial_prompt = initial_prompt self.vad_parameters = vad_parameters or {"threshold": 0.5} + self.hotwords = hotwords device = "cuda" if torch.cuda.is_available() else "cpu" if device == "cuda": @@ -231,7 +233,8 @@ class ServeClientFasterWhisper(ServeClientBase): language=self.language, task=self.task, vad_filter=self.use_vad, - vad_parameters=self.vad_parameters if self.use_vad else None) + vad_parameters=self.vad_parameters if self.use_vad else None, + hotwords=self.hotwords) if ServeClientFasterWhisper.SINGLE_MODEL: ServeClientFasterWhisper.SINGLE_MODEL_LOCK.release() diff --git a/whisper_live/client.py b/whisper_live/client.py index 5d78c13..14af63f 100644 --- a/whisper_live/client.py +++ b/whisper_live/client.py @@ -43,6 +43,7 @@ class Client: translation_srt_file_path="output_translated.srt", enable_timestamps=False, display_segments=4, + hotwords=None, ): """ Initializes a Client instance for audio recording and streaming to a server. @@ -101,6 +102,7 @@ class Client: self.task = "translate" self.enable_timestamps = enable_timestamps self.display_segments = display_segments + self.hotwords = hotwords self.audio_bytes = None @@ -299,6 +301,7 @@ class Client: "same_output_threshold": self.same_output_threshold, "enable_translation": self.enable_translation, "target_language": self.target_language, + "hotwords": self.hotwords, } ) ) @@ -820,6 +823,7 @@ class TranscriptionClient(TranscriptionTeeClient): translation_srt_file_path="./output_translated.srt", enable_timestamps=False, display_segments=4, + hotwords=None, ): self.client = Client( host, @@ -842,6 +846,7 @@ class TranscriptionClient(TranscriptionTeeClient): translation_srt_file_path=translation_srt_file_path, enable_timestamps=enable_timestamps, display_segments=display_segments, + hotwords=hotwords, ) if save_output_recording and not output_recording_filename.endswith(".wav"): diff --git a/whisper_live/server.py b/whisper_live/server.py index 54da64a..66bcd21 100644 --- a/whisper_live/server.py +++ b/whisper_live/server.py @@ -288,7 +288,8 @@ class TranscriptionServer: clip_audio=options.get("clip_audio", False), same_output_threshold=options.get("same_output_threshold", 10), cache_path=self.cache_path, - translation_queue=translation_queue + translation_queue=translation_queue, + hotwords=options.get("hotwords"), ) logging.info("Running faster_whisper backend.") @@ -533,7 +534,8 @@ class TranscriptionServer: include: Optional[List[str]] = Form(default=None), known_speaker_names: Optional[List[str]] = Form(default=None), known_speaker_references: Optional[List[str]] = Form(default=None), - stream: bool = Form(default=False) + stream: bool = Form(default=False), + hotwords: Optional[str] = Form(default=None), ): if stream: return JSONResponse({"error": "Streaming not supported in this backend."}, status_code=400) @@ -564,7 +566,8 @@ class TranscriptionServer: initial_prompt=prompt, temperature=temperature, vad_filter=False, - word_timestamps=(timestamp_granularities and "word" in timestamp_granularities) + word_timestamps=(timestamp_granularities and "word" in timestamp_granularities), + hotwords=hotwords, ) text = " ".join([s.text.strip() for s in segments])