| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Sequence |
| 5 | from typing import cast |
| 6 | |
| 7 | import torch |
| 8 | import torch.nn as nn |
| 9 | |
| 10 | DEFAULT_MAX_AUDIO_LEN_S = 655 |
| 11 | DEFAULT_MERGE_FACTOR = 4 |
| 12 | # Default convolution parameters: (padding, kernel_size, stride) |
| 13 | # These correspond to the two conv layers in GlmAsrEncoder |
| 14 | DEFAULT_CONV_PARAMS = [(1, 3, 1), (1, 3, 2)] |
| 15 | |
| 16 | |
| 17 | def _calculate_conv_output_length( |
| 18 | input_length: torch.Tensor, padding: int, kernel_size: int, stride: int |
| 19 | ) -> torch.Tensor: |
| 20 | """Calculate Conv1d output length using standard formula.""" |
| 21 | # in sync with `hf_processor._get_audio_token_length` |
| 22 | return (input_length + 2 * padding - (kernel_size - 1) - 1) // stride + 1 |
| 23 | |
| 24 | |
| 25 | def _as_list_chunk_counts( |
| 26 | chunk_counts: torch.Tensor | list[int] | list[torch.Tensor], |
| 27 | ) -> list[int]: |
| 28 | if isinstance(chunk_counts, torch.Tensor): |
| 29 | return chunk_counts.tolist() |
| 30 | if chunk_counts and isinstance(chunk_counts[0], torch.Tensor): |
| 31 | tensor_counts = cast(list[torch.Tensor], chunk_counts) |
| 32 | return [int(c.item()) for c in tensor_counts] |
| 33 | return [int(c) for c in chunk_counts] |
| 34 | |
| 35 | |
| 36 | def _normalize_chunk_counts( |
| 37 | chunk_counts: torch.Tensor | list[int] | list[torch.Tensor] | None, |
| 38 | num_chunks: int, |
| 39 | ) -> list[int]: |
| 40 | if chunk_counts is None: |
| 41 | return [1] * num_chunks |
| 42 | return _as_list_chunk_counts(chunk_counts) |
| 43 | |
| 44 | |
| 45 | def _get_audio_output_lengths_from_lengths( |
| 46 | audio_lengths: torch.Tensor, |
| 47 | merge_factor: int, |
| 48 | conv_params: list[tuple[int, int, int]], |
| 49 | ) -> torch.Tensor: |
| 50 | for padding, kernel_size, stride in conv_params: |
| 51 | audio_lengths = _calculate_conv_output_length( |
| 52 | audio_lengths, padding, kernel_size, stride |
| 53 | ) |
| 54 | return (audio_lengths - merge_factor) // merge_factor + 1 |
| 55 | |
| 56 | |
| 57 | def _get_audio_output_lengths_from_mask( |
| 58 | mask: torch.Tensor, |
| 59 | merge_factor: int, |
| 60 | conv_params: list[tuple[int, int, int]], |
| 61 | ) -> torch.Tensor: |
| 62 | audio_lengths = mask.sum(-1) |
| 63 | return _get_audio_output_lengths_from_lengths( |
| 64 | audio_lengths, merge_factor, conv_params |
| 65 | ) |
| 66 | |
| 67 | |
| 68 | def _get_audio_output_lengths_for_tower( |
| 69 | audio_tower: nn.Module, |
| 70 | audio_lengths: torch.Tensor, |
| 71 | merge_factor: int, |
| 72 | conv_params: list[tuple[int, int, int]], |
| 73 | ) -> torch.Tensor: |
| 74 | """Calculate the output lengths after audio processing. |
| 75 | |
| 76 | The output length accounts for: |
| 77 | 1. Convolution layers (downsampling) |
| 78 | 2. Merge factor (further downsampling during projection) |
| 79 | |
| 80 | Args: |
| 81 | audio_tower: The audio encoder module |
| 82 | audio_lengths: Input feature lengths [batch_size] |
| 83 | merge_factor: Factor for merging adjacent features |
| 84 | conv_params: List of (padding, kernel_size, stride) for each conv layer |
| 85 | |
| 86 | Returns: |
| 87 | Output lengths after all processing [batch_size] |
| 88 | |
| 89 | """ |
| 90 | # First, calculate the output length after convolutions |
| 91 | if hasattr(audio_tower, "_get_feat_extract_output_lengths"): |
| 92 | _, conv_output_lengths = audio_tower._get_feat_extract_output_lengths( |
| 93 | audio_lengths |
| 94 | ) |
| 95 | else: |
| 96 | conv_output_lengths = audio_lengths |
| 97 | for padding, kernel_size, stride in conv_params: |
| 98 | conv_output_lengths = _calculate_conv_output_length( |
| 99 | conv_output_lengths, padding, kernel_size, stride |
| 100 | ) |
| 101 | |
| 102 | # Then, apply merge_factor to get final output length |
| 103 | # Formula: (conv_output_lengths - merge_factor) // merge_factor + 1 |
| 104 | return (conv_output_lengths - merge_factor) // merge_factor + 1 |
| 105 | |
| 106 | |
| 107 | def _flatten_audio_features_by_length( |
| 108 | audio_features: torch.Tensor, |
| 109 | audio_output_lengths: torch.Tensor, |
| 110 | ) -> torch.Tensor: |
| 111 | num_chunks, max_audio_tokens, embed_dim = audio_features.shape |
| 112 | audio_output_lengths = audio_output_lengths.unsqueeze(1) |
| 113 | audio_features_mask = ( |
| 114 | torch.arange(max_audio_tokens) |
| 115 | .expand(num_chunks, max_audio_tokens) |
| 116 | .to(audio_output_lengths.device) |
| 117 | < audio_output_lengths |
| 118 | ) |
| 119 | return audio_features[audio_features_mask].view(-1, embed_dim) |
| 120 | |
| 121 | |
| 122 | def _group_audio_embeddings( |
| 123 | chunk_embeddings: Sequence[torch.Tensor], |
| 124 | chunk_counts: Sequence[int], |
| 125 | ) -> tuple[torch.Tensor, ...]: |
| 126 | grouped_embeddings = [] |
| 127 | current_idx = 0 |
| 128 | for count in chunk_counts: |
| 129 | audio_chunks = chunk_embeddings[current_idx : current_idx + count] |
| 130 | grouped_embeddings.append(torch.cat(audio_chunks, dim=0)) |
| 131 | current_idx += count |
| 132 | return tuple(grouped_embeddings) |