simulstream 0.1.0__py3-none-any.whl

simulstream/server/speech_processors/sliding_window_retranslation.py

@@ -0,0 +1,135 @@
+# Copyright 2025 FBK
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+from difflib import SequenceMatcher
+from types import SimpleNamespace
+from typing import List
+
+import torch
+
+from simulstream.server.speech_processors import SAMPLE_RATE
+from simulstream.server.speech_processors.base import BaseSpeechProcessor
+from simulstream.server.speech_processors.incremental_output import IncrementalOutput
+
+
+class SlidingWindowRetranslator(BaseSpeechProcessor):
+    """
+    A speech processor that applies a fixed-length sliding window retranslation with
+    deduplication to mitigate overlapping outputs when processing unsegmented audio streams.
+
+    This class implements the algorithm introduced in:
+
+       S. Sen, et al. 2025. *"Simultaneous Translation for Unsegmented Input:
+       A Sliding Window Approach"* (https://arxiv.org/pdf/2210.09754)
+
+    The approach relies on detecting the **longest common subsequence** between the current window
+    and the previous one, in order to prevent repeating tokens caused by overlapping audio windows.
+
+    Args:
+       config (SimpleNamespace): Configuration object. The following attributes are expected:
+
+           - **window_len (int)**: Length of the sliding window (in seconds).
+           - **matching_threshold (float, optional)**: Minimum fraction of the current tokens that
+             must match the previous history to be considered aligned. Default = ``0.1``.
+           - **override_on_failed_match (bool, optional)**: If ``True``, the previous history is
+             deleted from the output when no sufficient match is found. Otherwise, previous history
+             is kept and the new output is appended to the end of the previous history.
+             Default = ``False``.
+           - **max_tokens_per_second (int, optional)**: Maximum output tokens allowed per second of
+             audio. Default = ``10``.
+    """
+
+    def __init__(self, config: SimpleNamespace):
+        super().__init__(config)
+        self.window_len = self.config.window_len * SAMPLE_RATE
+        self.matching_threshold = getattr(self.config, "matching_threshold", 0.1)
+        self.override_on_failed_match = getattr(self.config, "override_on_failed_match", False)
+        self.max_tokens_per_second = getattr(self.config, "max_tokens_per_second", 10)
+        self.within_first_window = True
+
+    def _build_incremental_outputs(self, generated_tokens: List[str]) -> IncrementalOutput:
+        """
+        Deduplicates the output stream of overlapping windows using the algorithm introduced in
+        `S. Sen, et al. 2025. "Simultaneous Translation for Unsegmented Input:
+        A Sliding Window Approach" <https://arxiv.org/pdf/2210.09754>`_
+
+        This algorithm is based on the longest matching substring between the current and previous
+        window. We use tokens instead of string to match, though, as we have empirically observed
+        that tokenization is mostly consistent across generations of the same word.
+        """
+        if self.text_history is None or len(self.text_history) == 0:
+            self.text_history = generated_tokens
+            generated_string = self.tokens_to_string(generated_tokens)
+            return IncrementalOutput(
+                new_tokens=generated_tokens,
+                new_string=generated_string,
+                deleted_tokens=[],
+                deleted_string=""
+            )
+        seq_matcher = SequenceMatcher(None, self.text_history, generated_tokens, autojunk=False)
+        longest_match = seq_matcher.find_longest_match()
+        if longest_match.size >= self.matching_threshold * len(generated_tokens):
+            new_tokens = generated_tokens[longest_match.b + longest_match.size:]
+            deleted_tokens = self.text_history[longest_match.a + longest_match.size:]
+            new_string = self.tokens_to_string(new_tokens)
+            deleted_string = self.tokens_to_string(deleted_tokens)
+            # we take the matching part and the last part of the generated string as part of
+            # the history. Then we take from the history the tokens corresponding to the amount
+            # generated in this step, to ensure we have a sufficiently wide window
+            matching_and_last_tokens = generated_tokens[longest_match.b:]
+            initial_discarded_tokens = len(generated_tokens) - len(matching_and_last_tokens)
+            history_tokens_discarded = self.text_history[longest_match.a:]
+            history_initial_tokens = len(self.text_history) - len(history_tokens_discarded)
+            new_history_initial_tokens = self.text_history[
+                max(history_initial_tokens - initial_discarded_tokens, 0):history_initial_tokens]
+            self.text_history = new_history_initial_tokens + matching_and_last_tokens
+        else:
+            if self.within_first_window or self.override_on_failed_match:
+                deleted_tokens = self.text_history
+                deleted_string = self.tokens_to_string(self.text_history)
+                self.text_history = None
+            else:
+                deleted_tokens = []
+                deleted_string = ""
+            new_tokens = generated_tokens
+            new_string = self.tokens_to_string(generated_tokens)
+            self.text_history = generated_tokens
+        return IncrementalOutput(
+            new_tokens=new_tokens,
+            new_string=new_string,
+            deleted_tokens=deleted_tokens,
+            deleted_string=deleted_string,
+        )
+
+    def _update_speech_history(
+            self,
+            new_speech: torch.Tensor,
+            generated_tokens: List[str],
+            new_output: IncrementalOutput) -> None:
+        if self.within_first_window and len(self.audio_history) >= self.window_len:
+            self.within_first_window = False
+
+    def _update_text_history(
+            self,
+            new_speech: torch.Tensor,
+            generated_tokens: List[str],
+            new_output: IncrementalOutput) -> None:
+        pass
+
+    def end_of_stream(self) -> IncrementalOutput:
+        return IncrementalOutput([], "", [], "")
+
+    def clear(self) -> None:
+        super().clear()
+        self.within_first_window = True

simulstream/server/speech_processors/vad_wrapper.py

@@ -0,0 +1,180 @@
+# Copyright 2025 FBK
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+from types import SimpleNamespace
+from typing import List
+
+import numpy as np
+from silero_vad import load_silero_vad, VADIterator
+
+from simulstream.server.speech_processors import SAMPLE_RATE, SpeechProcessor, \
+    speech_processor_class_load
+from simulstream.server.speech_processors.incremental_output import merge_incremental_outputs, \
+    IncrementalOutput
+
+
+class VADWrapperSpeechProcessor(SpeechProcessor):
+    """
+    A speech processor that integrates **Voice Activity Detection (VAD)** to filter and split
+    continuous audio streams into meaningful speech chunks before processing them with an
+    underlying speech processor.
+
+    This class wraps a :class:`SpeechProcessor` implementation (defined by in the configuration via
+    the attribute `base_speech_processor_class`) with a Silero VAD-based iterator that detects the
+    start and end of speech segments. Audio outside of speech is ignored, and each detected segment
+    is passed to the underlying speech processor.
+
+    Args:
+        config (SimpleNamespace): Configuration object. The following attributes are used:
+
+            - **base_speech_processor_class (str)**: full name of the underlying speech processor
+              class to use.
+            - **vad_threshold (float, optional)**: VAD probability threshold. Default = ``0.5``.
+            - **vad_min_silence_duration_ms (int, optional)**: Minimum silence duration
+              (milliseconds) to consider the end of a speech segment. Default = ``100``.
+            - **vad_speech_pad_ms (int, optional)**: Padding (milliseconds) to include before and
+              after detected speech. Default = ``30``.
+            - **min_speech_size (int, optional)**: Minimum segment size in seconds; shorter
+              segments are ignored. Default = ``1``.
+            - Any additional attributes required by the subclass :py:attr:`speech_processor_class`.
+    """
+
+    @classmethod
+    def speech_processor_class(cls, config: SimpleNamespace) -> type[SpeechProcessor]:
+        return speech_processor_class_load(config.base_speech_processor_class)
+
+    @classmethod
+    def load_model(cls, config: SimpleNamespace):
+        super().load_model(config)
+        if not hasattr(cls, "vad_model") or cls.vad_model is None:
+            cls.vad_model = load_silero_vad()
+        cls.speech_processor_class(config).load_model(config)
+
+    def __init__(self, config: SimpleNamespace):
+        super().__init__(config)
+        self.speech_processor = self.speech_processor_class(self.config)(self.config)
+        self.min_speech_size = getattr(self.config, "min_speech_size", 1) * SAMPLE_RATE
+        self.vad_iterator = VADIterator(
+            self.vad_model,
+            threshold=getattr(self.config, "vad_threshold", 0.5),
+            sampling_rate=SAMPLE_RATE,
+            min_silence_duration_ms=getattr(self.config, "vad_min_silence_duration_ms", 100),
+            speech_pad_ms=getattr(self.config, "vad_speech_pad_ms", 30),
+        )
+        self.residual_prev_audio = None
+        self.speech_buffer = None
+        self.audio_cursor_position = 0
+        self.in_speech = False
+        assert SAMPLE_RATE == 16000, \
+            "SileroHFSlidingWindowRetranslator supports only 16kHz sampling rate"
+        self.window_size_samples = 512  # assuming 16kHz
+        self.previous_audio_chunk = None  # needed as VAD uses padding before start
+
+    def clear(self) -> None:
+        super().clear()
+        self.residual_prev_audio = None
+        self.speech_buffer = None
+        self.audio_cursor_position = 0
+        self.in_speech = False
+        self.previous_audio_chunk = None
+        self.vad_iterator.reset_states()
+        self.speech_processor.clear()
+
+    def process_chunk(self, waveform: np.float32) -> IncrementalOutput:
+        if self.residual_prev_audio is not None:
+            waveform = np.concatenate((self.residual_prev_audio, waveform))
+            self.residual_prev_audio = None
+        # we can have more than one generate if there are multiple speech segments in the current
+        # chunk
+        outputs = []
+
+        for i in range(0, len(waveform), self.window_size_samples):
+            chunk = waveform[i: i + self.window_size_samples]
+            if len(chunk) < self.window_size_samples:
+                # process tailing audio with the next waveform chunk
+                self.residual_prev_audio = chunk
+                break
+            speech_dict = self.vad_iterator(chunk, return_seconds=False)
+            if speech_dict:
+                # if a VAD event happens, update the status accordingly
+                assert not ('start' in speech_dict and 'end' in speech_dict)
+                if 'start' in speech_dict:
+                    assert not self.in_speech, \
+                        "Cannot start a new segment when current one is being processed. " \
+                        "This means there is a bug in the implementation."
+                    start_offset = self.audio_cursor_position - len(self.previous_audio_chunk)
+                    chunk_start_position = speech_dict['start'] - start_offset
+                    assert chunk_start_position >= 0
+                    self.speech_buffer = np.concatenate(
+                        (self.previous_audio_chunk, chunk))[chunk_start_position:]
+                    self.in_speech = True
+                if 'end' in speech_dict:
+                    assert self.in_speech, \
+                        "Cannot end a segment when no current segment is being processed. " \
+                        "This means there is a bug in the implementation."
+                    speech_buffer_len = len(self.speech_buffer) \
+                        if self.speech_buffer is not None else 0
+                    speech_buffer_offset = self.audio_cursor_position - speech_buffer_len
+                    chunk_end_position = speech_dict['end'] - speech_buffer_offset
+                    # In case we already processed more audio than needed (i.e., we already
+                    # processed beyond the end by VAD, we skip further processing; otherwise, we
+                    # process the remaining unhandled audio).
+                    # We can process more audio after the end as the VAD takes ~100ms (see the
+                    # min_silence_duration_ms parameter of the VAD) to emit the end signal, so if
+                    # we already processed the partial speech buffer (see min_speech_size in this
+                    # processor), we may have processed up to extra 100ms.
+                    if chunk_end_position >= 0:
+                        self.append_to_speech_buffer(chunk)
+                        self.speech_buffer = self.speech_buffer[:chunk_end_position]
+                        outputs.append(self.speech_processor.process_chunk(self.speech_buffer))
+                    self.in_speech = False
+                    self.speech_buffer = None
+                    outputs.append(self.speech_processor.end_of_stream())
+                    # reset history at the end of a segment
+                    if hasattr(self.speech_processor, 'text_history'):
+                        self.speech_processor.text_history = None
+                    if hasattr(self.speech_processor, 'audio_history'):
+                        self.speech_processor.audio_history = None
+            else:
+                # if no VAD event happens, we just ignore the audio if we are outside speech and
+                # update the buffer in case we are in speech
+                if self.in_speech:
+                    self.append_to_speech_buffer(chunk)
+            # update cursor position
+            self.audio_cursor_position += self.window_size_samples
+            self.previous_audio_chunk = chunk
+
+        if self.in_speech and len(self.speech_buffer) > self.min_speech_size:
+            outputs.append(self.speech_processor.process_chunk(self.speech_buffer))
+            self.speech_buffer = None
+
+        return merge_incremental_outputs(outputs, self.tokens_to_string)
+
+    def append_to_speech_buffer(self, audio_chunk: np.float32) -> None:
+        if self.speech_buffer is None:
+            self.speech_buffer = audio_chunk
+        else:
+            self.speech_buffer = np.concatenate((self.speech_buffer, audio_chunk))
+
+    def set_source_language(self, language: str) -> None:
+        self.speech_processor.set_source_language(language)
+
+    def set_target_language(self, language: str) -> None:
+        self.speech_processor.set_target_language(language)
+
+    def tokens_to_string(self, tokens: List[str]) -> str:
+        return self.speech_processor.tokens_to_string(tokens)
+
+    def end_of_stream(self) -> IncrementalOutput:
+        return self.speech_processor.end_of_stream()

simulstream/server/websocket_server.py

@@ -0,0 +1,236 @@
+# Copyright 2025 FBK
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+import argparse
+import asyncio
+import logging
+import time
+from contextlib import asynccontextmanager
+from types import SimpleNamespace
+from typing import Callable, Awaitable
+
+import websockets
+from websockets.asyncio.server import serve, ServerConnection
+import json
+
+import simulstream
+from simulstream.config import yaml_config
+from simulstream.metrics.logger import setup_metrics_logger, METRICS_LOGGER
+from simulstream.server.message_processor import MessageProcessor
+from simulstream.server.speech_processors import build_speech_processor
+
+
+logging.basicConfig(
+    format='%(asctime)s | %(levelname)s | %(name)s | %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S',
+    level=logging.INFO,
+)
+LOGGER = logging.getLogger('simulstream.websocket_server')
+
+
+class SpeechProcessorPool:
+    """
+    A pool of speech processors initialized at startup and made available to clients.
+
+    Args:
+        speech_processor_config (SimpleNamespace): configuration for the speech processors to
+           create.
+        size (int): number of speech processors to have in the pool.
+        acquire_timeout (int): timeout (in seconds) for waiting the availability of a speech
+           processor.
+    """
+    def __init__(
+            self,
+            speech_processor_config: SimpleNamespace,
+            size: int,
+            acquire_timeout: int):
+        self.size = size
+        self.acquire_timeout = acquire_timeout
+        self.available = asyncio.Queue(maxsize=size)
+        for _ in range(size):
+            self.available.put_nowait(build_speech_processor(speech_processor_config))
+
+    @asynccontextmanager
+    async def acquire(self):
+        """
+        Acquire one process from the pool and release it automatically.
+
+        Returns:
+            SpeechProcessor: a speech processor available for usage.
+
+        Raises:
+            TimeoutError: if no speech processor is available within the configured timeout.
+        """
+        speech_processor = await asyncio.wait_for(
+            self.available.get(), timeout=self.acquire_timeout)
+        try:
+            yield speech_processor
+        finally:
+            # Return worker to pool
+            self.available.put_nowait(speech_processor)
+
+
+def connection_handler_factory(
+        speech_processor_pool: SpeechProcessorPool
+) -> Callable[[ServerConnection], Awaitable[None]]:
+    """
+    Factory function that creates a connection handler for the WebSocket server.
+
+    The returned connection handler routine will process audio and metadata messages sent by a
+    single client over WebSocket.
+
+    The handler receives client data (raw audio chunks and textual metadata) and forwards it to a
+    message processor using a speech processor retrieved form the pool of the available ones.
+    If no speech processor is available for the configured waiting time, the client connection is
+    closed.
+    The handler also sends incremental processing results back to the client in JSON format.
+
+    Args:
+        speech_processor_pool (SpeechProcessorPool): Pool of speech processors to use to handle
+            client connections.
+
+    Returns:
+        Callable[[websockets.asyncio.server.ServerConnection], Awaitable[None]]: An asynchronous
+            WebSocket connection handler coroutine.
+    """
+
+    async def handle_connection(websocket: ServerConnection) -> None:
+        """
+        Handles a single client WebSocket connection.
+
+        This is the coroutine that processes incoming messages from a client:
+
+        - If the message is binary (``bytes``), it is interpreted as raw audio data and
+          buffered until a full chunk is ready for processing.
+        - If the message is text (``str``), it is parsed as JSON metadata and can:
+
+          - Set the input sample rate.
+          - Set source and target languages for translation.
+          - Log custom metadata to the metrics logger.
+          - Indicate the end of the audio stream.
+
+        At the end of the stream, any remaining audio is processed, the processor state is cleared,
+        and an ``end_of_processing`` message is sent to the client.
+
+        Args:
+            websocket (websockets.asyncio.server.ServerConnection): The WebSocket connection for
+                the client.
+        """
+        loop = asyncio.get_running_loop()
+        client_id = id(websocket)
+        LOGGER.info(f"Client {client_id} connected")
+        try:
+            async with speech_processor_pool.acquire() as speech_processor:
+                message_processor = MessageProcessor(client_id, speech_processor)
+
+                try:
+                    async for message in websocket:
+                        if isinstance(message, bytes):
+                            # in this case we are processing an audio chunk
+                            incremental_output = await loop.run_in_executor(
+                                None, message_processor.process_speech, message)
+                            if incremental_output is not None:
+                                await websocket.send(incremental_output.strings_to_json())
+                        elif isinstance(message, str):
+                            # textual message are used to handle metadata
+                            try:
+                                data = json.loads(message)
+                                if 'end_of_stream' in data:
+                                    incremental_output = await loop.run_in_executor(
+                                        None, message_processor.end_of_stream)
+                                    await websocket.send(incremental_output.strings_to_json())
+                                    await websocket.send(json.dumps({'end_of_processing': True}))
+                                else:
+                                    message_processor.process_metadata(data)
+                            except Exception as e:
+                                LOGGER.error(
+                                    f"Invalid string message: {message}. Error: {e}. Ignoring it.")
+                except websockets.exceptions.ConnectionClosed:
+                    LOGGER.info(f"Client {client_id} disconnected.")
+                except Exception as e:
+                    LOGGER.exception(f"Error: {e}")
+                finally:
+                    message_processor.clear()
+        except TimeoutError:
+            LOGGER.error(f"Timeout waiting for a new processor for client {client_id}")
+            await websocket.close()
+
+    return handle_connection
+
+
+async def main(args: argparse.Namespace):
+    """
+    Main entry point for running the WebSocket speech server.
+
+    This function loads the server and speech processor configurations from YAML,
+    initializes logging (including metrics logging), and starts the WebSocket server
+    on the configured host and port.
+
+    Args:
+        args (argparse.Namespace): parsed command-line arguments with configuration file paths.
+    """
+    LOGGER.info(f"Loading server configuration from {args.server_config}")
+    server_config = yaml_config(args.server_config)
+    LOGGER.info(
+        f"Metric logging is{'' if server_config.metrics.enabled else ' NOT'} enabled at "
+        f"{server_config.metrics.filename}")
+    setup_metrics_logger(server_config.metrics)
+    LOGGER.info(f"Loading speech processor from {args.speech_processor_config}")
+    speech_processor_config = yaml_config(args.speech_processor_config)
+    LOGGER.info(f"Using as speech processor: {speech_processor_config.type}")
+    speech_processor_loading_time = time.time()
+    speech_processors_pool = SpeechProcessorPool(
+        speech_processor_config, server_config.pool_size, server_config.acquire_timeout)
+    speech_processor_loading_time = time.time() - speech_processor_loading_time
+    LOGGER.info(f"Loaded speech processor in {speech_processor_loading_time:.3f} seconds")
+    METRICS_LOGGER.info(json.dumps({
+        "model_loading_time": speech_processor_loading_time,
+    }))
+    LOGGER.info(f"Serving websocket server at {server_config.hostname}:{server_config.port}")
+    async with serve(
+            connection_handler_factory(speech_processors_pool),
+            server_config.hostname,
+            server_config.port,
+            ping_timeout=None) as server:
+        await server.serve_forever()
+
+
+def cli_main():
+    """
+    Simulstream WebSocket server command-line interface (CLI) entry point.
+
+    This function parses command-line arguments and starts the asynchronous :func:`main` routine.
+
+    Example usage::
+
+        $ python websocket_server.py --server-config config/server.yaml \\
+              --speech-processor-config config/speech.yaml
+
+    Command-line arguments:
+
+    - ``--server-config`` (str, optional): Path to the server configuration file
+      (default: ``config/server.yaml``).
+    - ``--speech-processor-config`` (str, required): Path to the speech processor configuration
+      file.
+    """
+    LOGGER.info(f"Websocket server version: {simulstream.__version__}")
+    parser = argparse.ArgumentParser("websocket_simul_server")
+    parser.add_argument("--server-config", type=str, default="config/server.yaml")
+    parser.add_argument("--speech-processor-config", type=str, required=True)
+    args = parser.parse_args()
+    asyncio.run(main(args))
+
+
+if __name__ == "__main__":
+    cli_main()