simulstream 0.1.0__py3-none-any.whl

simulstream/server/speech_processors/base.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 abc import abstractmethod
+from types import SimpleNamespace
+from typing import List, Union, Dict
+
+import numpy as np
+import torch
+
+from simulstream.server.speech_processors import SpeechProcessor
+from simulstream.server.speech_processors.incremental_output import IncrementalOutput
+
+
+class BaseSpeechProcessor(SpeechProcessor):
+    """
+    A partial implementation of :class:`SpeechProcessor` that provides
+    common logic for handling incremental speech-to-text processing.
+
+    This class defines the high-level workflow of processing an incoming
+    audio chunk (preprocessing, generation, updating history, building
+    outputs), while leaving the model-specific details to subclasses.
+
+    Subclasses must implement the abstract helper methods to define how
+    audio is preprocessed, tokens are generated, and histories are updated.
+    """
+
+    def __init__(self, config: SimpleNamespace):
+        super().__init__(config)
+        self.tgt_lang_tag = None
+        self.src_lang_tag = None
+        self.audio_history = None
+        self.text_history = None
+
+    @abstractmethod
+    def _preprocess(self, waveform: np.float32) -> Union[Dict[str, torch.Tensor], torch.Tensor]:
+        """
+        Convert a raw audio waveform into model-ready features.
+
+        Args:
+            waveform (np.ndarray): A 1D NumPy array of type ``float32``
+                containing normalized audio samples in the range ``[-1.0, 1.0]``. This is the new
+                incoming chunk of audio.
+
+        Returns:
+            Union[Dict[str, torch.Tensor], torch.Tensor]:
+                A tensor or a dictionary of tensors ready for model inference.
+        """
+        ...
+
+    @abstractmethod
+    def _update_speech_history(
+            self,
+            new_speech: torch.Tensor,
+            generated_tokens: List[str],
+            new_output: IncrementalOutput) -> None:
+        """
+        Update the internal audio history.
+
+        Args:
+            new_speech (torch.Tensor): The newly preprocessed speech features.
+            generated_tokens (List[str]): Tokens generated from the new input.
+            new_output (IncrementalOutput): Incremental output object containing
+                new and deleted tokens/strings.
+        """
+        ...
+
+    @abstractmethod
+    def _update_text_history(
+            self,
+            new_speech: torch.Tensor,
+            generated_tokens: List[str],
+            new_output: IncrementalOutput) -> None:
+        """
+        Update the internal text history with new generated tokens.
+
+        Args:
+            new_speech (torch.Tensor): The newly preprocessed speech features.
+            generated_tokens (List[str]): Tokens generated from the new input.
+            new_output (IncrementalOutput): Incremental output object containing
+                new and deleted tokens/strings.
+        """
+        ...
+
+    @abstractmethod
+    def _generate(self, speech: Union[Dict[str, torch.Tensor], torch.Tensor]) -> List[str]:
+        """
+        Generate tokens from the given speech features.
+
+        Args:
+            speech (Union[Dict[str, torch.Tensor], torch.Tensor]):
+                Model-ready speech features as produced by :meth:`_preprocess`.
+
+        Returns:
+            List[str]: A list of generated tokens.
+        """
+        ...
+
+    @abstractmethod
+    def _build_incremental_outputs(self, generated_tokens: List[str]) -> IncrementalOutput:
+        """
+        Build an :class:`IncrementalOutput` object from generated tokens.
+
+        Args:
+            generated_tokens (List[str]): Tokens generated by :meth:`_generate`.
+
+        Returns:
+            IncrementalOutput: The structured incremental output.
+        """
+        ...
+
+    def process_chunk(self, waveform: np.float32) -> IncrementalOutput:
+        speech = self._preprocess(waveform)
+        generated_tokens = self._generate(speech)
+        new_output = self._build_incremental_outputs(generated_tokens)
+        self._update_speech_history(speech, generated_tokens, new_output)
+        self._update_text_history(speech, generated_tokens, new_output)
+        return new_output
+
+    def clear(self) -> None:
+        self.text_history = None
+        self.audio_history = None
+        self.src_lang_tag = None
+        self.tgt_lang_tag = None

simulstream/server/speech_processors/base_streamatt.py

@@ -0,0 +1,320 @@
+# 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 torch
+import logging
+import numpy as np
+
+from types import SimpleNamespace
+from abc import abstractmethod
+from typing import List, Tuple
+
+from simulstream.server.speech_processors import class_load
+from simulstream.server.speech_processors.base import BaseSpeechProcessor
+from simulstream.server.speech_processors.incremental_output import IncrementalOutput
+
+
+BOW_PREFIX = "\u2581"
+
+
+logger = logging.getLogger(__name__)
+
+
+class BaseStreamAtt(BaseSpeechProcessor):
+    """
+    A partial implementation of :class:`BaseSpeechProcessor` that provides common logic for the
+    StreamAtt policy, introduced in:
+
+       S. Papi, et al. 2024. *"StreamAtt: Direct Streaming Speech-to-Text Translation with
+       Attention-based Audio History Selection"* (https://aclanthology.org/2024.acl-long.202/)
+
+    The approach relies on selecting the audio history based on the cross-attention mechanism.
+    Specifically, the history for the next decoding step is defined as follows:
+     - First, the new textual history is selected by the **text_history_method**, which is in
+     charge of selecting the tokens to retain;
+     - Second, the new audio history is selected according to cross-attention scores between the
+     audio features and the retained textual history by discarding past features that are not
+     attended by any tokens of the textual history.
+
+     The derived class should implement the following methods:
+        - **audio_max_len**: Returns the maximum audio feature length.
+        - **load_model**: Loads the model to device.
+        - **_preprocess**: Preprocess the audio features before feeding them into the model.
+        - **_generate**: Generate that also returns cross attention scores.
+
+    Args:
+       config (SimpleNamespace): Configuration object. The following attributes are expected:
+           - **text_history (str)**: config (SimpleNamespace) with the following attribute:
+               - **type (str)**: Name of the class to use to determine the text history to keep as
+                 context for next predictions.
+           - **audio_subsampling_factor (int)**: Subsampling factor of the model, if any.
+             Defaults to 1.
+           - **text_history_max_len (int)**: The maximum length of the textual history after which
+             the current content is cut. Defaults to 128.
+           - **cross_attention_layer (int)**: Layer from which to extract the cross-attention from.
+           - **cutoff_frame_num (int)**: Number of last frames that cannot be attended by tokens
+             in the AlignAtt policy.
+           - **word_level_postprocess (bool)**: Whether to postprocess the generated tokens to keep
+             only complete words in the emitted hypothesis. To be disabled when operating with
+             character-level languages. Defaults to True.
+    """
+
+    def __init__(self, config: SimpleNamespace):
+        super().__init__(config)
+        self.config = config
+        text_history_config = self.config.text_history
+        text_history_cls = class_load(text_history_config.type)
+        self.text_history_method = text_history_cls(text_history_config)
+        self.audio_subsampling_factor = getattr(self.config, "audio_subsampling_factor", 1)
+        self.text_history_max_len = getattr(self.config, "text_history_max_len", 128)
+        self.cross_attn_layer = getattr(self.config, "cross_attention_layer", 3)
+        self.cutoff_frame_num = getattr(self.config, "cutoff_frame_num", 2)
+        self.word_level_postprocess = getattr(self.config, "word_level_postprocess", True)
+        self.unselected_tokens = []
+
+    @property
+    @abstractmethod
+    def audio_max_len(self) -> float:
+        """
+        Return the maximum allowed length of the audio features, beyond which the audio is cut off.
+        """
+        ...
+
+    @abstractmethod
+    def _generate(self, speech: torch.Tensor) -> Tuple[List[str], torch.Tensor]:
+        """
+        Generate tokens from the given speech features together with the cross-attention scores.
+
+        Args:
+            speech (torch.Tensor): Model-ready speech features as produced by :meth:`_preprocess`.
+
+        Returns:
+            Tuple[List[str], torch.Tensor]:
+                List[str]: A list of generated tokens.
+                torch.Tensor: Cross-attention scores with dimension (generated_tokens,
+                input_length).
+        """
+        ...
+
+    @staticmethod
+    def normalize_attn(attn):
+        """
+        Normalize the cross-attention scores along the frame dimension to avoid attention sinks.
+        """
+        std = attn.std(axis=0)
+        std[std == 0.] = 1.0
+        mean = attn.mean(axis=0)
+        return (attn - mean) / std
+
+    def _update_text_history(self, new_output: List[str]) -> int:
+        if self.text_history:
+            self.text_history += new_output
+        else:
+            self.text_history = new_output
+        new_history = self.text_history_method.select_text_history(self.text_history)
+        discarded_text = len(self.text_history) - len(new_history)
+        self.text_history = new_history
+
+        # Ensure not exceeding max text history length
+        if self.text_history and len(self.text_history) > self.text_history_max_len:
+            if logger.isEnabledFor(logging.WARNING):
+                logger.warning(
+                    f"The textual history has hit the maximum predefined length of "
+                    f"{self.text_history_max_len}")
+            self.text_history = self.text_history[-self.text_history_max_len:]
+        return discarded_text
+
+    def _cut_audio_exceeding_maxlen(self):
+        # Ensure not exceeding max audio history length
+        if len(self.audio_history) > self.audio_max_len:
+            if logger.isEnabledFor(logging.WARNING):
+                logger.warning(
+                    f"The audio history has hit the maximum predefined length of "
+                    f"{self.audio_max_len}")
+            self.audio_history = self.audio_history[-self.audio_max_len:]
+
+    def _update_speech_history(self, discarded_text: int, cross_attn: torch.Tensor) -> None:
+        # If no history is discarded, no need for attention-based audio trimming
+        if discarded_text == 0:
+            # Check audio history not exceeding maximum allowed length
+            self._cut_audio_exceeding_maxlen()
+            return
+
+        # Trim the cross-attention by excluding the discarded new generated tokens and the
+        # discarded textual history. Output shape: (text_history_len, n_audio_features)
+        cross_attn = cross_attn[discarded_text:discarded_text + len(self.text_history), :]
+
+        # Compute the frame to which each token of the textual history mostly attends to
+        most_attended_idxs = torch.argmax(cross_attn.float(), dim=1)
+
+        # Find the first feature that is attended
+        if most_attended_idxs.shape[0] > 1:
+            # Multiple tokens: sort and get the earliest attended frame
+            sorted_idxs = torch.sort(most_attended_idxs)[0]
+            earliest_attended_idx = sorted_idxs[0]
+        else:
+            # Only one token: use the unique most attended frame
+            earliest_attended_idx = most_attended_idxs[0]
+
+        # Multiply by the subsampling factor to recover the original number of frames
+        frames_to_cut = earliest_attended_idx * self.audio_subsampling_factor
+
+        # Cut the unattended audio features
+        self.audio_history = self.audio_history[frames_to_cut:]
+
+        # Check audio history not exceeding maximum allowed length
+        self._cut_audio_exceeding_maxlen()
+
+    @staticmethod
+    def _strip_incomplete_words(tokens: List[str]) -> List[str]:
+        """
+        Remove last incomplete word(s) from the new hypothesis.
+
+        Args:
+            tokens (List[str]): selected tokens, possibly containing partial words to be removed.
+
+        Returns:
+            List[str]: A list of generated tokens from which partial words are removed.
+        """
+        tokens_to_write = []
+        # iterate from the end and count how many trailing tokens to drop
+        num_tokens_incomplete = 0
+        for tok in reversed(tokens):
+            num_tokens_incomplete += 1
+            if tok.startswith(BOW_PREFIX):
+                # slice off the trailing incomplete tokens
+                tokens_to_write = tokens[:-num_tokens_incomplete]
+                break
+        return tokens_to_write
+
+    def alignatt_policy(self, generated_tokens, cross_attn) -> List[str]:
+        """
+        Apply the AlignAtt policy by cutting off tokens whose attention falls
+        beyond the allowed frame range.
+        The AlignAtt policy was introduced in:
+            S. Papi, et al. 2023. *"AlignAtt: Using Attention-based Audio-Translation
+            Alignments as a Guide for Simultaneous Speech Translation"*
+            (https://www.isca-archive.org/interspeech_2023/papi23_interspeech.html)
+        """
+        # Select attention scores corresponding to the new generated tokens
+        cross_attn = cross_attn[-len(generated_tokens):, :]
+        selected_tokens = generated_tokens
+
+        # Find the frame to which each token mostly attends to
+        most_attended_frames = torch.argmax(cross_attn, dim=1)
+        cutoff = cross_attn.size(1) - self.cutoff_frame_num
+
+        # Find the first token that attends beyond the cutoff frame
+        invalid_tok_ids = torch.where(most_attended_frames >= cutoff)[0]
+
+        # Truncate tokens up to the first invalid alignment (if any)
+        if len(invalid_tok_ids) > 0:
+            selected_tokens = selected_tokens[:invalid_tok_ids[0]]
+
+        if self.word_level_postprocess:
+            selected_tokens = self._strip_incomplete_words(selected_tokens)
+
+        # Store unselected tokens, to be used in the case of end of stream
+        self.unselected_tokens = generated_tokens[len(selected_tokens):]
+
+        return selected_tokens
+
+    def _build_incremental_outputs(self, generated_tokens: List[str]) -> IncrementalOutput:
+        return IncrementalOutput(
+            new_tokens=generated_tokens,
+            new_string=self.tokens_to_string(generated_tokens),
+            deleted_tokens=[],
+            deleted_string="",
+        )
+
+    def process_chunk(self, waveform: np.float32) -> IncrementalOutput:
+        speech = self._preprocess(waveform)
+        # Generate new hypothesis with its corresponding cross-attention scores (no prefix)
+        generated_tokens, cross_attn = self._generate(speech)
+        # Select the part of the new hypothesis to be emitted, and trim cross-attention accordingly
+        selected_output = self.alignatt_policy(generated_tokens, cross_attn)
+        incremental_output = self._build_incremental_outputs(selected_output)
+        # Discard textual history, if needed
+        discarded_text = self._update_text_history(selected_output)
+        # Trim audio corresponding to the discarded textual history
+        self._update_speech_history(discarded_text, cross_attn)
+        return incremental_output
+
+    def end_of_stream(self) -> IncrementalOutput:
+        last_output = self._build_incremental_outputs(self.unselected_tokens)
+        self.unselected_tokens = []
+        return last_output
+
+    def clear(self) -> None:
+        super().clear()
+        self.text_history = None
+        self.audio_history = None
+        self.unselected_tokens = []
+
+
+class FixedWordsTextHistory:
+    """
+    Fixed Words textual history selection method that retains a pre-defined
+    number of words in the history (*history_words*).
+
+    The current implementation supports only SentencePiece.
+    """
+    def __init__(self, config: SimpleNamespace):
+        self.history_words = getattr(config, "history_words", 20)
+        self.config = config
+
+    def select_text_history(self, text_history: List[str]):
+        words_to_keep = self.history_words
+        new_history = []
+        for token in reversed(text_history):
+            new_history.append(token)
+            # Check if 'BOW_PREFIX' (space in SentencePiece) is contained in the token,
+            # meaning that we reached the beginning of the word that should be counted
+            if BOW_PREFIX in token:
+                words_to_keep -= 1
+                # When all the words to keep are consumed, the accumulation is stopped
+                # and the prefix is returned
+                if words_to_keep == 0:
+                    break
+        # Reverse the list
+        return new_history[::-1]
+
+
+class PunctuationTextHistory:
+    """
+    Punctuation textual history selection method that retains the sentence
+    before the last strong punctuation character.
+
+    The current implementation supports only SentencePiece.
+    """
+
+    STRONG_PUNCTUATION = [".", "!", "?", ":", ";"]
+
+    def __init__(self, config: SimpleNamespace):
+        self.config = config
+
+    def select_text_history(self, text_history):
+        new_history = []
+        for token in reversed(text_history):
+            prefix_token = token
+            contains_punctuation = False
+            for punct in self.STRONG_PUNCTUATION:
+                if punct in prefix_token:
+                    contains_punctuation = True
+                    break
+            if contains_punctuation:
+                break
+            new_history.append(token)
+        # Reverse the list
+        return new_history[::-1]

simulstream/server/speech_processors/canary_sliding_window_retranslation.py

@@ -0,0 +1,73 @@
+# 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
+import torch
+from nemo.collections.asr.models import ASRModel
+
+from simulstream.server.speech_processors import SAMPLE_RATE
+from simulstream.server.speech_processors.sliding_window_retranslation import \
+    SlidingWindowRetranslator
+
+
+class CanarySlidingWindowRetranslator(SlidingWindowRetranslator):
+    """
+    Perform Sliding Window Retranslation with Canary.
+    """
+
+    @classmethod
+    def load_model(cls, config: SimpleNamespace):
+        if not hasattr(cls, "model") or cls.model is None:
+            cls.model = ASRModel.from_pretrained(model_name=config.model_name)
+            cls.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            assert cls.model.preprocessor._sample_rate == SAMPLE_RATE
+            cls.model.to(cls.device)
+
+    def _generate(self, speech: torch.Tensor) -> List[str]:
+        output = self.model.transcribe(
+            speech, source_lang=self.src_lang_tag, target_lang=self.tgt_lang_tag)
+        return self.model.tokenizer.ids_to_tokens(output[0].y_sequence)
+
+    def tokens_to_string(self, tokens: List[str]) -> str:
+        # avoid that the initial space, if it is there, get removed in the detokenization
+        check_for_init_space = self.text_history is not None and len(self.text_history) > 0
+        if check_for_init_space:
+            tokens = [' '] + tokens
+        text = self.model.tokenizer.tokens_to_text(tokens)
+        if check_for_init_space:
+            text = text[1:]
+        return text
+
+    def _preprocess(self, waveform: np.float32) -> torch.Tensor:
+        """
+        Extracts the filter-bank features from the input waveform and appends them to the audio
+        history. Returns the concatenated audio history and new frames, taking the last
+        `self.window_len` frames, and returns it after storing it in the audio history.
+        """
+        if self.audio_history is not None:
+            waveform = np.concatenate((self.audio_history, waveform))
+        new_speech_len = len(waveform)
+        if new_speech_len > self.window_len:
+            waveform = waveform[-self.window_len:]
+        self.audio_history = waveform
+        return torch.tensor(waveform).to(self.device)
+
+    def set_target_language(self, language: str) -> None:
+        self.tgt_lang_tag = language
+
+    def set_source_language(self, language: str) -> None:
+        self.src_lang_tag = language

simulstream/server/speech_processors/hf_sliding_window_retranslation.py

@@ -0,0 +1,87 @@
+# 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
+import torch
+from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq
+
+from simulstream.server.speech_processors import SAMPLE_RATE
+from simulstream.server.speech_processors.sliding_window_retranslation import \
+    SlidingWindowRetranslator
+
+
+class HFSlidingWindowRetranslator(SlidingWindowRetranslator):
+    """
+    Perform Sliding Window Retranslation with a Huggingface speech-to-text model.
+    """
+
+    @classmethod
+    def load_model(cls, config: SimpleNamespace):
+        if not hasattr(cls, "model") or cls.model is None:
+            lang_tags = None
+            if hasattr(config, "supported_langs") and config.supported_langs is not None:
+                lang_tags = [
+                    config.lang_tag_template.format(lang) for lang in config.supported_langs]
+            cls.processor = AutoProcessor.from_pretrained(
+                config.hf_model_name,
+                additional_special_tokens=lang_tags)
+            cls.model = AutoModelForSpeechSeq2Seq.from_pretrained(
+                config.hf_model_name, trust_remote_code=True)
+            cls.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            cls.model.to(cls.device)
+
+    def _generate(self, speech: torch.Tensor) -> List[str]:
+        speech_seconds = speech.shape[1] / 100  # 1 frame every 10 ms
+        extra_kwargs = {
+            "max_new_tokens": int(max(self.max_tokens_per_second * speech_seconds, 10))}
+        if self.tgt_lang_tag is not None:
+            extra_kwargs["forced_bos_token_id"] = self.tgt_lang_tag
+        generated_ids = self.model.generate(speech, **extra_kwargs)[0]
+        return self.processor.tokenizer.convert_ids_to_tokens(
+            generated_ids, skip_special_tokens=True)
+
+    def tokens_to_string(self, tokens: List[str]) -> str:
+        # avoid that the initial space, if it is there, get removed in the detokenization
+        if self.text_history is not None and len(self.text_history) > 0:
+            tokens = [''] + tokens
+        return self.processor.tokenizer.convert_tokens_to_string(tokens)
+
+    def _preprocess(self, waveform: np.float32) -> torch.Tensor:
+        """
+        Extracts the filter-bank features from the input waveform and appends them to the audio
+        history. Returns the concatenated audio history and new frames, taking the last
+        `self.window_len` frames, and returns it after storing it in the audio history.
+        """
+        if self.audio_history is not None:
+            waveform = np.concatenate((self.audio_history, waveform))
+        new_speech_len = len(waveform)
+        if new_speech_len > self.window_len:
+            waveform = waveform[-self.window_len:]
+        self.audio_history = waveform
+        new_speech = self.processor(
+            waveform,
+            sampling_rate=SAMPLE_RATE,
+            return_tensors="pt")["input_features"]
+        return new_speech.to(self.device)
+
+    def set_target_language(self, language: str) -> None:
+        lang_tag_id = self.processor.tokenizer.convert_tokens_to_ids(
+            self.config.lang_tag_template.format(language))
+        self.tgt_lang_tag = torch.tensor(lang_tag_id, dtype=torch.int, device=self.device)
+
+    def set_source_language(self, language: str) -> None:
+        pass