| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Lightweight utility functions for Gemma3n audio processing. |
| 4 | |
| 5 | This module is separate from gemma3n_mm.py to avoid heavy CUDA dependencies, |
| 6 | making it testable without a full vLLM build. |
| 7 | """ |
| 8 | |
| 9 | import torch |
| 10 | |
| 11 | |
| 12 | def adjust_audio_features_to_expected_length( |
| 13 | audio_features: torch.Tensor, |
| 14 | expected_tokens: int, |
| 15 | audio_padding_embs: torch.Tensor, |
| 16 | ) -> tuple[torch.Tensor, int]: |
| 17 | """Adjust audio features to expected token length via padding or truncation. |
| 18 | |
| 19 | The Gemma3nProcessor expects all audio will be ~30s in length and inserts |
| 20 | a fixed number of audio soft tokens into the text. However, the audio |
| 21 | preprocessing and encoder do not guarantee they will produce exactly that |
| 22 | many soft tokens; they may produce fewer tokens (for shorter audio) or more |
| 23 | tokens (for longer audio or due to BOA/EOA special tokens). |
| 24 | |
| 25 | This function handles both cases: |
| 26 | - If fewer tokens: pad with the provided padding embeddings |
| 27 | - If more tokens: truncate to the expected count |
| 28 | |
| 29 | Args: |
| 30 | audio_features: Audio embeddings tensor of shape |
| 31 | (batch_size, seq_len, embed_dim) |
| 32 | expected_tokens: The expected number of audio tokens (e.g., 188) |
| 33 | audio_padding_embs: Padding embeddings tensor of shape (1, 1, embed_dim) |
| 34 | |
| 35 | Returns: |
| 36 | Tuple of: |
| 37 | - adjusted_features: Audio features adjusted to expected_tokens length |
| 38 | - tokens_truncated: Number of tokens truncated (0 if padding was applied) |
| 39 | |
| 40 | """ |
| 41 | audio_batch_size, audio_seq_len, audio_embed_dim = audio_features.shape |
| 42 | tokens_truncated = 0 |
| 43 | |
| 44 | if audio_seq_len < expected_tokens: |
| 45 | # Pad to expected length with padding embeddings |
| 46 | extra_padding_tokens = expected_tokens - audio_seq_len |
| 47 | extra_padding_features = audio_padding_embs.expand( |
| 48 | audio_batch_size, extra_padding_tokens, audio_embed_dim |
| 49 | ) |
| 50 | audio_features = torch.cat((audio_features, extra_padding_features), dim=1) |
| 51 | elif audio_seq_len > expected_tokens: |
| 52 | # Truncate to expected length (audio encoder produced more tokens |
| 53 | # than expected, e.g., due to longer audio or placeholder mismatch) |
| 54 | tokens_truncated = audio_seq_len - expected_tokens |
| 55 | audio_features = audio_features[:, :expected_tokens, :] |
| 56 | |
| 57 | return audio_features, tokens_truncated |