diff --git a/examples/doa/README.md b/examples/doa/README.md new file mode 100644 index 0000000..e836c90 --- /dev/null +++ b/examples/doa/README.md @@ -0,0 +1,164 @@ +# Decoder-Only Attention (DOA) Policy + +The [**Decoder-Only Attention (DOA)** policy](https://arxiv.org/abs/2605.31432) extends the +[encoder-decoder **StreamAtt** policy](https://aclanthology.org/2024.acl-long.202/) to SpeechLLMs +that have no cross-attention mechanism. Instead of relying on encoder-decoder cross-attention, DOA +builds a *proxy* cross-attention matrix by extracting the self-attention weights between the audio +tokens and the text tokens from the decoder layers. The resulting matrix is then used by the +[AlignAtt policy](https://www.isca-archive.org/interspeech_2023/papi23_interspeech.html) to decide +which generated tokens can be safely emitted at each step. + +## Supported models + +| Class | HuggingFace model | +|---|---| +| `simulstream.server.speech_processors.phi4multimodal_doa.Phi4MultimodalDOA` | `microsoft/Phi-4-multimodal-instruct` | +| `simulstream.server.speech_processors.qwenomni_doa.Qwen3OmniDOA` | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | + +## Configuration + +A DOA config file is a YAML file passed via `--speech-processor-config`. Below is a full annotated +example, elements between `{}` brackets should be replaced as mentioned below: + +```yaml +type: "{MODEL_CLASS}" +text_history: + type: "simulstream.server.speech_processors.base_streamatt.{TEXT_HISTORY}" +audio_history_max_duration: 180 +text_history_max_len: 128 +speech_chunk_size: 1 # seconds of audio fed to the model at each step +max_new_tokens: 32 # max tokens generated per chunk +hf_model_name: "{MODEL_NAME}" +detokenizer_type: "hf" +word_level_postprocess: {WORD_POSTPROCESS} +bow_prefix: "{BOW_MARKER}" +prompt: "{PROMPT}" + +# --- DOA parameters --- +attn_layer: {ATTN_LAYER} +attn_head: {ATTN_HEAD} +average_attn_over_layers: {ATTN_AVG} +cutoff_frame_num: {CUTOFF_FRAME} +``` + +Parameters to be replaced: +- `{MODEL_CLASS}` from the [supported models Class](#supported-models) +(e.g., `simulstream.server.speech_processors.phi4multimodal_doa.Phi4MultimodalDOA`) +- `{TEXT_HISTORY}` among: + - `FixedWordsTextHistory`: Retains the last *N* complete words. Recommended for space-separated +languages (English, Italian, …). In this case, the number of `history_words` should be added to +the config, for instance: + ```yaml + text_history: + type: "simulstream.server.speech_processors.base_streamatt.FixedWordsTextHistory" + history_words: 10 + ``` + - `FixedCharsTextHistory`: Retains the last *N* characters. Recommended for character-level +languages (Chinese, Japanese) where `FixedWordsTextHistory` is ineffective because spaces are +sparse. n this case, the number of `history_chars` should be added to the config, for instance: + ```yaml + text_history: + type: "simulstream.server.speech_processors.base_streamatt.FixedWordsTextHistory" + history_chars: 20 + ``` + - `PunctuationTextHistory`: Retains the text from the last strong punctuation mark (`.`, `!`, +`?`, `:`, `;`, `。`). Works for both space-separated and character-level languages. +- `{MODEL_NAME}` from the [supported models HuggingFace model name](#supported-models) (e.g., +`microsoft/Phi-4-multimodal-instruct`) +- `{WORD_POSTPROCESS}`: When `true`, the output is trimmed to complete words before emission. +Set to `false` for character-level languages (Chinese, Japanese). +- `{BOW_MARKER}`: The beginning-of-word (BOW) marker used by the model's tokenizer (e.g., +Phi-4-multimodal and Qwen3-Omni use a plain space `" "`). +- `{PROMPT}`: User prompt. The default is `"Translate the audio to {tgt_lang}:"`. It can be +overridden with the `prompt` key in the yaml. The placeholders `{src_lang}` and `{tgt_lang}` are +filled in automatically from the language codes passed at inference time. +- `{ATTN_LAYER}`: Decoder layer to extract self-attention from (int, 0-indexed). +- `{ATTN_HEAD}`: Attention head to extract self-attention from (int, 0-indexed). `null` averages +over all heads. +- `{ATTN_AVG}`: If `true` (default), average the selected per-layer attention view across all +layers; `attn_layer` is used only when this is `false`. +- `{CUTOFF_FRAME}`: Cutoff frame of the AlignAtt policy. Tokens whose attention peak falls in the +last *N* audio frames are withheld. Higher values add more latency but reduce the risk of cutting +correct tokens. + +## Configurations of DOA's paper + +The configurations used to report the final results in Figure 3 are reported below: + +### Phi4-Multimodal +```yaml +type: "simulstream.server.speech_processors.phi4multimodal_doa.Phi4MultimodalDOA" +text_history: + type: "simulstream.server.speech_processors.base_streamatt.PunctuationTextHistory" +audio_history_max_duration: 180 +text_history_max_len: 128 +speech_chunk_size: 1 +detokenizer_type: "hf" +hf_model_name: "microsoft/Phi-4-multimodal-instruct" +word_level_postprocess: True +max_new_tokens: 32 +bow_prefix: " " +attn_layer: 0 +attn_head: null +average_attn_over_layers: True +cutoff_frame_num: __FRAME__ +``` + +### Qwen3-Omni +```yaml +type: "simulstream.server.speech_processors.qwenomni_doa.Qwen3OmniDOA" +text_history: + type: "simulstream.server.speech_processors.base_streamatt.PunctuationTextHistory" +audio_history_max_duration: 60 +text_history_max_len: 128 +speech_chunk_size: 1 +detokenizer_type: "hf" +hf_model_name: "Qwen/Qwen3-Omni-30B-A3B-Instruct" +word_level_postprocess: True +max_new_tokens: 32 +bow_prefix: " " +prompt: "You are a professional {src_lang}-to-{tgt_lang} translator. Your goal is to accurately \ +convey the meaning and nuances of the original {src_lang} speech while adhering to {tgt_lang} \ +grammar, vocabulary, and cultural sensitivities. Use precise terminology and a tone appropriate \ +for academic or instructional materials. Produce only the {tgt_lang} translation, without any \ +additional explanations or commentary. Please translate the provided {src_lang} speech into \ +{tgt_lang}:" +attn_layer: 0 +attn_head: null +average_attn_over_layers: True +cutoff_frame_num: __FRAME__ +``` +To run the inference, `__FRAME__` should be replaced with `sed`: +```bash +simulstream_inference --speech-processor-config <(sed "s/__FRAME__/${FRAME}/g" ${CONFIG_YAML}) \ + --wav-list-file ${AUDIOPATH_LIST} \ + --tgt-lang $TGTLANG --src-lang en \ + --metrics-log-file ${OUTLOG} +``` +where `${FRAME}` is `5`, `10`, or `15` following the paper, `${CONFIG_YAML}` is the path to the +aforementioned configuration yaml file, `${AUDIOPATH_LIST}` is the list of test audio files path, +and `${OUTLOG}` is the path to the jsonl output log file. + +## Adding a new model + +To support a new SpeechLLM, subclass `DecoderOnlyAttention` +(`simulstream.server.speech_processors.base_doa.DecoderOnlyAttention`) and implement: + +- `load_model(config)` — load the model and processor. +- `build_prompt()` — return the text prompt string. +- `build_processor_inputs(waveform)` — build processor inputs from the rolling audio history. +- `_do_generate(inputs)` — run generation and return `(new_tokens, attentions)`. +- `_find_audio_positions(input_ids)` — return the positions of audio tokens in the input sequence. + +## Citation +If you use DOA in your work, please cite: + +```bibtex +@article{papi-2026-doa, + title = {{DOA}: Training-Free Decoder-Only Attention Policy for Long-Form + Simultaneous Translation with {SpeechLLMs}}, + author = {Papi, Sara and Bentivogli, Luisa}, + journal = {arXiv preprint arXiv:2605.31432}, + year = {2026}, +} +``` \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 45a6581..e4912cb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,7 +14,8 @@ dependencies = [ "pyyaml>6.0", "websockets", "torch", - "librosa" + "librosa", + "pycountry" ] dynamic = ["version"] diff --git a/simulstream/server/speech_processors/base_doa.py b/simulstream/server/speech_processors/base_doa.py new file mode 100644 index 0000000..a2ec6e1 --- /dev/null +++ b/simulstream/server/speech_processors/base_doa.py @@ -0,0 +1,230 @@ +# Copyright 2026 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 logging +from abc import abstractmethod +from types import SimpleNamespace +from typing import List, Tuple, Any, Dict + +import numpy as np +import pycountry +import torch + +from simulstream.server.speech_processors import SAMPLE_RATE +from simulstream.server.speech_processors.base_streamatt import BaseStreamAtt + + +logger = logging.getLogger(__name__) + + +def get_language_name(code: str) -> str: + """Return the language name for an ISO 639-1 code, falling back to the code itself.""" + lang = pycountry.languages.get(alpha_2=code) + if lang is not None: + return lang.name + else: + logger.warning(f"Language code '{code}' not found in the language list. Using language " + f"code directly in the prompt, but this can lead to unexpected behavior.") + return code + + +class DecoderOnlyAttention(BaseStreamAtt): + """ + Generic Decoder-only Attention-based policy for SpeechLLMs. + + The class handles: + - Raw-waveform history accumulation. + - Greedy generation with ``output_attentions=True``. + - Building the proxy cross-attention matrix from self-attention weights. + - Applying the StreamAtt-based policy on the proxy cross-attention matrix. + + The derived class should implement the following methods: + - **load_model**: Loads the model and processor. + - **build_prompt**: Builds the text prompt to use with audio inputs. + - **build_processor_inputs**: Builds model inputs from the rolling audio history. + - **_do_generate**: Returns newly-generated tokens and self-attention scores. + - **_find_audio_positions**: Returns the indices of audio tokens. + + Args: + config (SimpleNamespace): Configuration object. The following additional attributes are + expected: + - **attn_layer (int)**: Layer from which to extract attention scores. Defaults to 0. + - **attn_head (int)**: Attention head to use. If not set, attention scores are averaged + over all heads. + - **average_attn_over_layers (bool)**: Whether to average the selected attention view + over all decoder layers. Defaults to True. + - **audio_history_max_duration (int)**: Maximum raw waveform length to keep in the + rolling history, in seconds. Defaults to 180. + - **max_new_tokens (int)**: Maximum tokens to generate per chunk. Defaults to 32. + """ + + def __init__(self, config: SimpleNamespace): + super().__init__(config) + self.cross_attn_layer = getattr(self.config, "attn_layer", 0) + self.cross_attn_head = getattr(self.config, "attn_head", None) + self.average_attn_over_layers = getattr(self.config, "average_attn_over_layers", True) + self.audio_history_max_duration = getattr(self.config, "audio_history_max_duration", 180) + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.max_new_tokens = getattr(self.config, "max_new_tokens", 32) + self.prompt = getattr(self.config, "prompt", "Translate the audio to {tgt_lang}:") + logger.debug("Prompt:\n%s", self.prompt) + + @property + def audio_max_len(self) -> int: + """Maximum raw-waveform samples to keep in the rolling audio history.""" + return self.audio_history_max_duration * SAMPLE_RATE + + @abstractmethod + def load_model(self, config: SimpleNamespace) -> None: + """ + Load the model and processor from *config* and assign them to ``self.model`` and + ``self.processor``. + + The model **must** be loaded with ``output_attentions=True`` (or the equivalent flag for + the architecture) and ``_attn_implementation="eager"``. + """ + ... + + @abstractmethod + def build_prompt(self) -> str: + """ + Return the prompt string to be used with audio tokens. + """ + ... + + @abstractmethod + def build_processor_inputs(self, waveform: np.ndarray) -> Any: + """ + Build processor inputs from the entire rolling waveform history (float32, 16 kHz). + """ + ... + + @abstractmethod + def _do_generate(self, inputs: Dict[str, Any]) -> Tuple[List[str], List[torch.Tensor]]: + """ + Runs the actual generation from the underlying model and returns the generated tokens and + the corresponding self-attention scores. + """ + ... + + @abstractmethod + def _find_audio_positions(self, input_ids: torch.Tensor) -> torch.Tensor: + """Return token positions corresponding to the encoded audio span.""" + ... + + def _generate(self, waveform: np.ndarray) -> Tuple[List[str], torch.Tensor]: + """ + Generate tokens from the given inputs together with the self-attention scores. + + Returns: + Tuple[List[str], torch.Tensor]: + List[str]: A list of generated tokens. + torch.Tensor: Self-attention scores between speech and text with dimension + (token_len, audio_len). + """ + inputs = self.build_processor_inputs(waveform).to(self.device) + input_ids = inputs["input_ids"] # (1, input_len) + input_len = input_ids.shape[1] + + audio_positions = self._find_audio_positions(input_ids) + audio_len = audio_positions.shape[0] + + # Run the actual generate on the underlying model + new_tokens, attentions = self._do_generate(inputs) + + # Build proxy cross-attention for the hypothesis (prefix + new_tokens) + prefill_attn = self.average_attn(attentions[0]) + prefix_len = len(self.text_history) if self.text_history else 0 + if prefix_len > 0: + # Prefix rows come from the prefill pass + prefix_rows = prefill_attn[input_len - prefix_len:, :][:, audio_positions] + else: + prefix_rows = torch.zeros(0, audio_len, device=self.device) + + if new_tokens: + # The prefill pass predicts the first generated token, so its last row corresponds to + # the first generated token's proxy audio-attention + first_new_row = prefill_attn[-1:, audio_positions] + # Other tokens' attention is present in each generation step + new_rows = [ + self.average_attn(step_attn).squeeze(0)[audio_positions] + for step_attn in attentions[1:] + ] + subsequent_new_attn = torch.stack(new_rows, dim=0) + new_attn = torch.cat([first_new_row, subsequent_new_attn], dim=0) + else: + new_attn = torch.zeros(0, audio_len, device=self.device) + + cross_attn = torch.cat([prefix_rows, new_attn], dim=0) + cross_attn = self.normalize_attn(cross_attn) + return new_tokens, cross_attn + + def set_target_language(self, language: str) -> None: + self.tgt_lang = language + + def set_source_language(self, language: str) -> None: + self.src_lang = language + + def build_raw_text_prefix(self) -> str: + return "".join(self.text_history) if self.text_history else "" + + def _select_attn_from_layer(self, layer_attn: torch.Tensor) -> torch.Tensor: + # Generation runs one stream at a time, so remove the singleton batch dimension. + layer_attn = layer_attn.squeeze(0) + if self.cross_attn_head is None: + # Default behavior: average over all heads for this layer. + return layer_attn.mean(dim=0) + + num_heads = layer_attn.shape[0] + if self.cross_attn_head < 0 or self.cross_attn_head >= num_heads: + raise ValueError( + f"Invalid attn_head={self.cross_attn_head}. Layer has {num_heads} heads." + ) + return layer_attn[self.cross_attn_head] + + def average_attn(self, attn) -> torch.Tensor: + """ + Average or select attentions according to ``attn_layer``, ``attn_head``, and + ``average_attn_over_layers``. + + If ``attn_head`` is not set, attention is averaged over heads. If + ``average_attn_over_layers`` is set, the selected per-layer attention view is also + averaged across layers; otherwise only ``attn_layer`` is used. + """ + if self.average_attn_over_layers: + # Average the per-layer attention view selected by _select_attn_from_layer. + return torch.stack( + [self._select_attn_from_layer(layer_attn) for layer_attn in attn], + dim=0, + ).mean(dim=0) + return self._select_attn_from_layer(attn[self.cross_attn_layer]) + + def _preprocess(self, waveform: np.float32) -> np.ndarray: + """ + Append *waveform* to ``self.audio_history`` and enforce the maximum length. + """ + if self.audio_history is None: + self.audio_history = waveform + else: + self.audio_history = np.concatenate([self.audio_history, waveform]) + + if len(self.audio_history) > self.audio_max_len: + logger.warning("Audio history exceeded %d samples; trimming.", self.audio_max_len) + self.audio_history = self.audio_history[-self.audio_max_len:] + + return self.audio_history + + def tokens_to_string(self, tokens: List[str]) -> str: + """Convert a list of decoded tokens to a plain output string.""" + return "".join(tokens) diff --git a/simulstream/server/speech_processors/base_streamatt.py b/simulstream/server/speech_processors/base_streamatt.py index 3d31ed8..b5c5128 100644 --- a/simulstream/server/speech_processors/base_streamatt.py +++ b/simulstream/server/speech_processors/base_streamatt.py @@ -75,7 +75,8 @@ def __init__(self, config: SimpleNamespace): 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.bow_prefix = getattr(self.config, "bow_prefix", BOW_PREFIX) + self.text_history_method = text_history_cls(text_history_config, self.bow_prefix) 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) @@ -182,8 +183,7 @@ def _update_speech_history(self, discarded_text: int, cross_attn: torch.Tensor) # Check audio history not exceeding maximum allowed length self._cut_audio_exceeding_maxlen() - @staticmethod - def _strip_incomplete_words(tokens: List[str]) -> List[str]: + def _strip_incomplete_words(self, tokens: List[str]) -> List[str]: """ Remove last incomplete word(s) from the new hypothesis. @@ -198,7 +198,7 @@ def _strip_incomplete_words(tokens: List[str]) -> List[str]: num_tokens_incomplete = 0 for tok in reversed(tokens): num_tokens_incomplete += 1 - if tok.startswith(BOW_PREFIX): + if tok.startswith(self.bow_prefix): # slice off the trailing incomplete tokens tokens_to_write = tokens[:-num_tokens_incomplete] break @@ -273,11 +273,10 @@ 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.bow_prefix = getattr(config, "bow_prefix", BOW_PREFIX) self.config = config def select_text_history(self, text_history: List[str]): @@ -285,9 +284,9 @@ def select_text_history(self, text_history: List[str]): 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: + # Check if bow_prefix is contained in the token, meaning that we reached + # the beginning of the word that should be counted + if self.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 diff --git a/simulstream/server/speech_processors/phi4multimodal_doa.py b/simulstream/server/speech_processors/phi4multimodal_doa.py new file mode 100644 index 0000000..f045d95 --- /dev/null +++ b/simulstream/server/speech_processors/phi4multimodal_doa.py @@ -0,0 +1,97 @@ +# Copyright 2026 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, Any, Dict, Tuple + +import numpy as np +import torch +from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig + +from simulstream.server.speech_processors import SAMPLE_RATE +from simulstream.server.speech_processors.base_doa import DecoderOnlyAttention, get_language_name + + +class Phi4MultimodalDOA(DecoderOnlyAttention): + """ + Decoder-Only Attention agent for Phi4-Multimodal. + """ + + # Phi-4 special tokens + _USER_START = "<|user|>" + _AUDIO_TOKEN = "<|audio_1|>" + _END_TOKEN = "<|end|>" + _ASST_START = "<|assistant|>" + + ENCODER_SUBSAMPLING_FACTOR = 8 + HOP_LENGTH = 160 # 10ms at 16kHz + AUDIO_SPECIAL_TOKEN_ID = 200011 # _AUDIO_SPECIAL_TOKEN_ID in modeling_phi4mm.py + + def __init__(self, config: SimpleNamespace): + super().__init__(config) + self.audio_subsampling_factor = self.ENCODER_SUBSAMPLING_FACTOR * self.HOP_LENGTH + + @classmethod + def load_model(cls, config: SimpleNamespace) -> None: + model_path = getattr( + config, + "hf_model_name", + getattr(config, "model_path", "microsoft/Phi-4-multimodal-instruct"), + ) + + cls.processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True) + cls.model = AutoModelForCausalLM.from_pretrained( + model_path, + device_map="cuda", + torch_dtype="auto", + trust_remote_code=True, + _attn_implementation="eager", + ) + cls.model.eval() + cls.generation_config = GenerationConfig.from_pretrained(model_path) + + def build_prompt(self) -> str: + filled_prompt = self.prompt.replace("{tgt_lang}", get_language_name(self.tgt_lang)) + raw_prefix = self.build_raw_text_prefix() + prompt = ( + f"{self._USER_START}{self._AUDIO_TOKEN}" + f"{filled_prompt}{self._END_TOKEN}" + f"{self._ASST_START}{raw_prefix}" + ) + return prompt + + def build_processor_inputs(self, waveform: np.ndarray) -> dict: + return self.processor( + text=self.build_prompt(), + audios=[(waveform, SAMPLE_RATE)], + return_tensors="pt", + ) + + def _do_generate(self, inputs: Dict[str, Any]) -> Tuple[List[str], List[torch.Tensor]]: + input_len = inputs["input_ids"].shape[1] + output = self.model.generate( + **inputs, + max_new_tokens=self.max_new_tokens, + generation_config=self.generation_config, + num_logits_to_keep=1, + output_attentions=True, + return_dict_in_generate=True, + do_sample=False, + ) + new_tokens = self.processor.tokenizer.convert_ids_to_tokens( + output.sequences[0, input_len:], skip_special_tokens=True) + return new_tokens, output.attentions + + def _find_audio_positions(self, input_ids: torch.Tensor) -> torch.Tensor: + return (input_ids[0] == self.AUDIO_SPECIAL_TOKEN_ID).nonzero(as_tuple=True)[0] diff --git a/simulstream/server/speech_processors/qwenomni_doa.py b/simulstream/server/speech_processors/qwenomni_doa.py new file mode 100644 index 0000000..ba8648b --- /dev/null +++ b/simulstream/server/speech_processors/qwenomni_doa.py @@ -0,0 +1,152 @@ +# Copyright 2026 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, Dict, Any, Tuple + +import numpy as np +import torch +from qwen_omni_utils import process_mm_info +from transformers import Qwen3OmniMoeForConditionalGeneration, Qwen3OmniMoeProcessor + +from simulstream.server.speech_processors import SAMPLE_RATE +from simulstream.server.speech_processors.base_doa import DecoderOnlyAttention, get_language_name + + +class Qwen3OmniDOA(DecoderOnlyAttention): + """ + Decoder-Only Attention agent for Qwen3-Omni. + + Args: + config (SimpleNamespace): Configuration object. The following additional attributes are + expected: + - **repetition_penalty (float)**: Repetition penalty for text generation. + Default: ``1.05``. + - **temperature (float)**: Temperature parameter. Default: ``1.0``. + - **no_repeat_ngram_size (int)**: Ngram size for text generation. Default: ``5``. + """ + + AUDIO_TOKEN_STRIDE = 640 + AUDIO_TOKEN_INDEX = 151675 # <|audio_pad|> + AUDIO_START_TOKEN_ID = 151669 # <|audio_start|> + AUDIO_END_TOKEN_ID = 151670 # <|audio_end|> + SYSTEM_PROMPT = ( + "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of " + "perceiving auditory and visual inputs, as well as generating text and speech." + ) + + def __init__(self, config: SimpleNamespace): + super().__init__(config) + self.audio_subsampling_factor = self.AUDIO_TOKEN_STRIDE + self.repetition_penalty = getattr(self.config, "repetition_penalty", 1.05) + self.temperature = getattr(self.config, "temperature", 1.0) + self.no_repeat_ngram_size = getattr(self.config, "no_repeat_ngram_size", 5) + + @classmethod + def load_model(cls, config: SimpleNamespace) -> None: + """Load the Qwen3-Omni model and processor.""" + model_name = getattr( + config, + "hf_model_name", + getattr(config, "model_path", "Qwen/Qwen3-Omni-30B-A3B-Instruct"), + ) + attn_impl = getattr(config, "attn_implementation", "eager") + + cls.processor = Qwen3OmniMoeProcessor.from_pretrained(model_name) + cls.model = Qwen3OmniMoeForConditionalGeneration.from_pretrained( + model_name, + torch_dtype="auto", + device_map="auto", + attn_implementation=attn_impl, + enable_audio_output=False, + ) + cls.model.eval() + + def build_prompt(self) -> str: + """Build the translation instruction used alongside the audio input.""" + return ( + self.prompt + .replace("{src_lang}", get_language_name(self.src_lang)) + .replace("{tgt_lang}", get_language_name(self.tgt_lang)) + ) + + def build_processor_inputs(self, waveform: np.ndarray) -> dict: + """Build multimodal processor inputs from the rolling audio history.""" + prompt_text = self.build_prompt() + prefix = self.build_raw_text_prefix() + + conversation = [ + { + "role": "system", + "content": [{"type": "text", "text": self.SYSTEM_PROMPT}], + }, + { + "role": "user", + "content": [ + {"type": "audio", "audio": waveform}, + {"type": "text", "text": prompt_text}, + ], + }, + ] + + prompt = self.processor.apply_chat_template( + conversation, + add_generation_prompt=True, + tokenize=False, + ) + + audios, images, videos = process_mm_info(conversation, use_audio_in_video=True) + + inputs = self.processor( + text=f"{prompt}{prefix}", + audio=audios, + images=images, + videos=videos, + sampling_rate=SAMPLE_RATE, + return_tensors="pt", + padding=True, + use_audio_in_video=True, + ) + return inputs.to(self.device).to(self.model.dtype) + + def _find_audio_positions(self, input_ids: torch.Tensor) -> torch.Tensor: + """Return token positions corresponding to the encoded audio span.""" + audio_positions = (input_ids[0] == self.AUDIO_TOKEN_INDEX).nonzero(as_tuple=True)[0] + if audio_positions.numel() > 0: + return audio_positions + + start_pos = (input_ids[0] == self.AUDIO_START_TOKEN_ID).nonzero(as_tuple=True)[0][0] + end_pos = (input_ids[0] == self.AUDIO_END_TOKEN_ID).nonzero(as_tuple=True)[0] + end_pos = end_pos[end_pos > start_pos][0] + return torch.arange(start_pos + 1, end_pos, device=input_ids.device) + + def _do_generate(self, inputs: Dict[str, Any]) -> Tuple[List[str], List[torch.Tensor]]: + input_len = inputs["input_ids"].shape[1] + output = self.model.generate( + **inputs, + use_audio_in_video=True, + return_audio=False, + thinker_max_new_tokens=self.max_new_tokens, + thinker_repetition_penalty=self.repetition_penalty, + thinker_no_repeat_ngram_size=self.no_repeat_ngram_size, + thinker_output_attentions=True, + thinker_return_dict_in_generate=True, + thinker_do_sample=False, + temperature=self.temperature, + ) + if isinstance(output, tuple): + output = output[0] + new_tokens = self.processor.tokenizer.convert_ids_to_tokens( + output.sequences[0, input_len:], skip_special_tokens=True) + return new_tokens, output.attentions