diff --git a/whisper_live/backend/base.py b/whisper_live/backend/base.py index c6e6a7e..4d926dc 100644 --- a/whisper_live/backend/base.py +++ b/whisper_live/backend/base.py @@ -63,6 +63,13 @@ class ServeClientBase(object): self.end_time_for_same_output = None self.translation_queue = translation_queue + # Optional post-processing callable for segments. + # If set, called with a segment dict and must return a segment dict. + # Allows external projects to plug in custom post-processing + # (e.g. PII redaction, formatting, diarization) without modifying + # WhisperLive's core code. + self.segment_post_processor = None + # threading self.lock = threading.Lock() @@ -254,9 +261,23 @@ class ServeClientBase(object): This method formats the transcription segments into a JSON object and attempts to send this object to the client. If an error occurs during the send operation, it logs the error. + If a ``segment_post_processor`` callable is set, each segment is passed through it + before sending. The callable receives a segment dict and must return a segment dict. + Returns: segments (list): A list of transcription segments to be sent to the client. """ + if self.segment_post_processor is not None: + processed = [] + for seg in segments: + try: + result = self.segment_post_processor(seg) + processed.append(result if result is not None else seg) + except Exception as e: + logging.error(f"[ERROR]: segment_post_processor failed: {e}") + processed.append(seg) + segments = processed + try: self.websocket.send( json.dumps({ diff --git a/whisper_live/server.py b/whisper_live/server.py index e420a47..9798618 100644 --- a/whisper_live/server.py +++ b/whisper_live/server.py @@ -176,6 +176,7 @@ class TranscriptionServer: self.single_model = False self.batch_config = None self.raw_pcm_input = False + self.segment_post_processor = None def initialize_client( self, websocket, options, faster_whisper_custom_model_path, @@ -316,6 +317,10 @@ class TranscriptionServer: if client is None: raise ValueError(f"Backend type {self.backend.value} not recognised or not handled.") + # Attach segment post-processor if configured + if self.segment_post_processor is not None: + client.segment_post_processor = self.segment_post_processor + if translation_client: client.translation_client = translation_client client.translation_thread = translation_thread @@ -477,7 +482,8 @@ class TranscriptionServer: batch_max_size=8, batch_window_ms=50, raw_pcm_input=False, - metrics_port: int = 0): + metrics_port: int = 0, + segment_post_processor=None): """ Run the transcription server. @@ -493,6 +499,11 @@ class TranscriptionServer: batch_window_ms (int): Maximum time in milliseconds to wait for the batch to fill after the first request arrives. Defaults to 50. + segment_post_processor (callable, optional): A callable that receives + a transcription segment dict and returns a modified segment dict. + Applied to every segment before sending to the client. Useful for + plugging in custom post-processing (e.g. formatting, redaction). + Defaults to None. """ self.cache_path = cache_path self.raw_pcm_input = raw_pcm_input @@ -506,6 +517,7 @@ class TranscriptionServer: if batch_enabled and batch_window_ms < 0: raise ValueError(f"batch_window_ms must be >= 0, got {batch_window_ms}") + self.segment_post_processor = segment_post_processor self.client_manager = ClientManager(max_clients, max_connection_time) if faster_whisper_custom_model_path is not None and not os.path.exists(faster_whisper_custom_model_path): if "/" not in faster_whisper_custom_model_path: