| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Modules below used for the audio encoder component in: models/nano_nemotron_vl.py.""" |
| 4 | |
| 5 | from collections.abc import Iterable |
| 6 | from functools import cache |
| 7 | from typing import Any |
| 8 | |
| 9 | import numpy as np |
| 10 | import torch |
| 11 | import torch.nn as nn |
| 12 | from transformers import ParakeetEncoder as HFParakeetEncoder |
| 13 | from transformers import PreTrainedConfig |
| 14 | from transformers.audio_utils import mel_filter_bank |
| 15 | |
| 16 | from vllm.logger import init_logger |
| 17 | from vllm.model_executor.layers.activation import ReLUSquaredActivation |
| 18 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 19 | from vllm.model_executor.model_loader.weight_utils import default_weight_loader |
| 20 | from vllm.platforms import current_platform |
| 21 | from vllm.transformers_utils.configs.parakeet import ExtractorConfig, ParakeetConfig |
| 22 | |
| 23 | logger = init_logger(__name__) |
| 24 | |
| 25 | |
| 26 | class ParakeetProjection(nn.Module): |
| 27 | def __init__(self, config: ParakeetConfig) -> None: |
| 28 | super().__init__() |
| 29 | sound_hidden_size = config.hidden_size |
| 30 | proj_hidden_size = config.projection_hidden_size |
| 31 | llm_hidden_size = config.llm_hidden_size |
| 32 | bias = config.projection_bias |
| 33 | |
| 34 | self.norm = RMSNorm(sound_hidden_size, eps=config.projection_eps) |
| 35 | self.linear1 = nn.Linear(sound_hidden_size, proj_hidden_size, bias=bias) |
| 36 | self.activation = ReLUSquaredActivation() |
| 37 | self.linear2 = nn.Linear(proj_hidden_size, llm_hidden_size, bias=bias) |
| 38 | |
| 39 | def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 40 | hidden_states = self.norm(hidden_states) |
| 41 | hidden_states = self.linear1(hidden_states) |
| 42 | hidden_states = self.activation(hidden_states) |
| 43 | hidden_states = self.linear2(hidden_states) |
| 44 | return hidden_states |
| 45 | |
| 46 | |
| 47 | class ProjectedParakeet(nn.Module): |
| 48 | def __init__( |
| 49 | self, |
| 50 | config: PreTrainedConfig, |
| 51 | *, |
| 52 | dtype: torch.dtype, |
| 53 | llm_hidden_size: int, |
| 54 | max_model_len: int, |
| 55 | ) -> None: |
| 56 | super().__init__() |
| 57 | self.config = ParakeetConfig.from_hf_config( |
| 58 | config, llm_hidden_size=llm_hidden_size, max_model_len=max_model_len |
| 59 | ) |
| 60 | self.encoder = HFParakeetEncoder(self.config) |
| 61 | self.encoder = self.encoder.to(dtype) |
| 62 | self.projection = ParakeetProjection(self.config) |
| 63 | self.projection = self.projection.to(dtype) |
| 64 | |
| 65 | def forward( |
| 66 | self, input_features: torch.Tensor, attention_mask: torch.Tensor | None = None |
| 67 | ) -> torch.Tensor: |
| 68 | outputs = self.encoder( |
| 69 | input_features=input_features, attention_mask=attention_mask |
| 70 | ) |
| 71 | outputs = outputs.last_hidden_state |
| 72 | outputs = outputs.to(dtype=torch.bfloat16) |
| 73 | outputs = self.projection(outputs) |
| 74 | return outputs |
| 75 | |
| 76 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 77 | loaded_params: set[str] = set() |
| 78 | params_dict = dict(self.named_parameters()) |
| 79 | buffers_dict = dict(self.named_buffers()) |
| 80 | |
| 81 | if isinstance(weights, dict): |
| 82 | weights_list = list(weights.items()) |
| 83 | else: |
| 84 | weights_list = list(weights) |
| 85 | |
| 86 | for name, weight in weights_list: |
| 87 | if name.startswith("sound_encoder.encoder.feature_extractor."): |
| 88 | # Feature extractor buffers are handled outside the encoder. |
| 89 | continue |
| 90 | if name.startswith("sound_encoder."): |
| 91 | target_name = name[len("sound_encoder.") :] |
| 92 | elif name.startswith("sound_projection."): |
| 93 | target_name = f"projection.{name[len('sound_projection.') :]}" |
| 94 | else: |
| 95 | continue |
| 96 | |
| 97 | target = params_dict.get(target_name) |
| 98 | if target is None: |
| 99 | target = buffers_dict.get(target_name) |
| 100 | if target is None: |
| 101 | if self._can_skip_missing_named_param(target_name): |
| 102 | continue |
| 103 | raise ValueError(f"Unknown weight: {name}") |
| 104 | weight_loader = getattr(target, "weight_loader", default_weight_loader) |
| 105 | with torch.no_grad(): |
| 106 | weight_loader(target, weight) |
| 107 | loaded_params.add(target_name) |
| 108 | |
| 109 | return loaded_params |
| 110 | |
| 111 | def _can_skip_missing_named_param(self, target_name: str) -> bool: |
| 112 | if self.config.convolution_bias: |
| 113 | return False |
| 114 | |
| 115 | # In transformers v5 (not v4), `convolution_bias=False` is |
| 116 | # propagated from parakeet config. If `False`, torch.conv1d will |
| 117 | # *skip registering the param*, thus it will be missing in the |
| 118 | # module's named params. *If* you happen to also have the bias |
| 119 | # tensors in the weights, it will cause a mismatch between the |
| 120 | # weights and the params. |
| 121 | # This allows us to have `convolution_bias=False` in the sound config, |
| 122 | # but still allow for the weights to exist. |
| 123 | |
| 124 | return target_name.endswith( |
| 125 | ( |
| 126 | ".conv.pointwise_conv1.bias", |
| 127 | ".conv.depthwise_conv.bias", |
| 128 | ".conv.pointwise_conv2.bias", |
| 129 | ) |
| 130 | ) |
| 131 | |
| 132 | |
| 133 | EPSILON = 1e-5 |
| 134 | LOG_ZERO_GUARD_VALUE = 2**-24 |
| 135 | |
| 136 | |
| 137 | class ParakeetExtractor: |
| 138 | def __init__(self, config: PreTrainedConfig) -> None: |
| 139 | self.config = ExtractorConfig.from_hf_config(config) |
| 140 | """`config` is named *exactly* for `._get_subsampling_output_length` below""" |
| 141 | self._clip_target_samples = int( |
| 142 | round(self.config.clip_duration_s * self.config.sampling_rate) |
| 143 | ) |
| 144 | self._tail_min_samples = int( |
| 145 | round(self.config.clip_min_duration_s * self.config.sampling_rate) |
| 146 | ) |
| 147 | |
| 148 | @staticmethod |
| 149 | @cache |
| 150 | def _get_window(win_length: int, device: str) -> torch.Tensor: |
| 151 | return torch.hann_window(win_length, periodic=False, device=device) |
| 152 | |
| 153 | @staticmethod |
| 154 | @cache |
| 155 | def _get_mel_filters( |
| 156 | feature_size: int, sampling_rate: int, n_fft: int, device: str |
| 157 | ) -> torch.Tensor: |
| 158 | filter_bank = mel_filter_bank( |
| 159 | num_frequency_bins=n_fft // 2 + 1, |
| 160 | num_mel_filters=feature_size, |
| 161 | min_frequency=0.0, |
| 162 | max_frequency=sampling_rate / 2, |
| 163 | sampling_rate=sampling_rate, |
| 164 | norm="slaney", |
| 165 | mel_scale="slaney", |
| 166 | ) |
| 167 | return torch.from_numpy(filter_bank.T).to(device=device, dtype=torch.float32) |
| 168 | |
| 169 | def _torch_extract_fbank_features(self, waveform: torch.Tensor, device: str): |
| 170 | # spectrogram |
| 171 | device = str(torch.device(device)) |
| 172 | cfg = self.config |
| 173 | window = self._get_window(cfg.win_length, device) |
| 174 | stft = torch.stft( |
| 175 | waveform, |
| 176 | self.config.n_fft, |
| 177 | hop_length=cfg.hop_length, |
| 178 | win_length=cfg.win_length, |
| 179 | window=window, |
| 180 | return_complex=True, |
| 181 | pad_mode="constant", |
| 182 | ) |
| 183 | mel_filters = self._get_mel_filters( |
| 184 | cfg.feature_size, cfg.sampling_rate, cfg.n_fft, device |
| 185 | ) |
| 186 | return self._apply_mel_filters(stft, mel_filters) |
| 187 | |
| 188 | @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) |
| 189 | def _apply_mel_filters( |
| 190 | self, stft_output: torch.Tensor, mel_filters: torch.Tensor |
| 191 | ) -> torch.Tensor: |
| 192 | magnitudes = stft_output.real.square() + stft_output.imag.square() |
| 193 | mel_spec = mel_filters @ magnitudes |
| 194 | mel_spec = torch.log(mel_spec + LOG_ZERO_GUARD_VALUE) |
| 195 | return mel_spec.permute(0, 2, 1) |
| 196 | |
| 197 | @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) |
| 198 | def _apply_preemphasis( |
| 199 | self, input_features: torch.Tensor, audio_lengths: torch.Tensor |
| 200 | ) -> torch.Tensor: |
| 201 | timemask = torch.arange( |
| 202 | input_features.shape[1], device=input_features.device |
| 203 | ).unsqueeze(0) < audio_lengths.unsqueeze(1) |
| 204 | input_features = torch.cat( |
| 205 | [ |
| 206 | input_features[:, :1], |
| 207 | input_features[:, 1:] |
| 208 | - self.config.preemphasis * input_features[:, :-1], |
| 209 | ], |
| 210 | dim=1, |
| 211 | ) |
| 212 | input_features = input_features.masked_fill(~timemask, 0.0) |
| 213 | return input_features |
| 214 | |
| 215 | @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) |
| 216 | def _normalize_mel_features( |
| 217 | self, mel_features: torch.Tensor, audio_lengths: torch.Tensor |
| 218 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 219 | features_lengths = torch.floor_divide( |
| 220 | audio_lengths + self.config.n_fft // 2 * 2 - self.config.n_fft, |
| 221 | self.config.hop_length, |
| 222 | ) |
| 223 | attention_mask = ( |
| 224 | torch.arange(mel_features.shape[1], device=mel_features.device)[None, :] |
| 225 | < features_lengths[:, None] |
| 226 | ) |
| 227 | mask = attention_mask.unsqueeze(-1) |
| 228 | lengths = attention_mask.sum(dim=1) |
| 229 | mel_features_masked = mel_features * mask |
| 230 | mean = (mel_features_masked.sum(dim=1) / lengths.unsqueeze(-1)).unsqueeze(1) |
| 231 | variance = ((mel_features_masked - mean) ** 2 * mask).sum(dim=1) / ( |
| 232 | lengths - 1 |
| 233 | ).unsqueeze(-1) |
| 234 | std = torch.sqrt(variance).unsqueeze(1) |
| 235 | return (mel_features - mean) / (std + EPSILON) * mask, attention_mask |
| 236 | |
| 237 | def _pad_raw_speech( |
| 238 | self, raw_speech: list[torch.Tensor], max_len: int, device: str |
| 239 | ) -> torch.Tensor: |
| 240 | output = torch.full( |
| 241 | (len(raw_speech), max_len), |
| 242 | self.config.padding_value, |
| 243 | device=device, |
| 244 | dtype=torch.float32, |
| 245 | ) |
| 246 | dsts = [output[i, : raw_speech[i].shape[0]] for i in range(len(raw_speech))] |
| 247 | srcs = [s.squeeze(-1) for s in raw_speech] |
| 248 | # single kernel horizontal fusion |
| 249 | torch._foreach_copy_(dsts, srcs) |
| 250 | return output |
| 251 | |
| 252 | def _clip_sizes(self, audio_len: int) -> list[int]: |
| 253 | audio_len = max(audio_len, self._tail_min_samples) |
| 254 | num_full_clips, remainder = divmod(audio_len, self._clip_target_samples) |
| 255 | clip_sizes = [self._clip_target_samples] * num_full_clips |
| 256 | if remainder > 0: |
| 257 | clip_sizes.append(max(remainder, self._tail_min_samples)) |
| 258 | return clip_sizes |
| 259 | |
| 260 | def audio_token_count(self, audio_len: int) -> int: |
| 261 | total_tokens = 0 |
| 262 | for clip_size in self._clip_sizes(audio_len): |
| 263 | num_frames = clip_size // self.config.hop_length |
| 264 | n_tokens = HFParakeetEncoder._get_subsampling_output_length( |
| 265 | self, torch.tensor([num_frames], dtype=torch.float) |
| 266 | ) |
| 267 | total_tokens += int(n_tokens.item()) |
| 268 | return max(1, total_tokens) |
| 269 | |
| 270 | def split_audio_into_clips(self, audio: torch.Tensor) -> list[torch.Tensor]: |
| 271 | assert audio.ndim == 1 |
| 272 | audio_len = int(audio.shape[0]) |
| 273 | clip_sizes = self._clip_sizes(audio_len) |
| 274 | target_len = sum(clip_sizes) |
| 275 | if audio_len < target_len: |
| 276 | audio = torch.nn.functional.pad(audio, (0, target_len - audio_len)) |
| 277 | |
| 278 | clips = list[torch.Tensor]() |
| 279 | offset = 0 |
| 280 | for clip_size in clip_sizes: |
| 281 | clips.append(audio[offset : offset + clip_size]) |
| 282 | offset += clip_size |
| 283 | return clips |
| 284 | |
| 285 | def __call__( |
| 286 | self, |
| 287 | raw_speech: list[np.ndarray], |
| 288 | *, |
| 289 | device: str = "cpu", |
| 290 | ) -> dict[str, Any]: |
| 291 | raw_speech = [ |
| 292 | torch.as_tensor(speech, device=device, dtype=torch.float32) |
| 293 | for speech in raw_speech |
| 294 | ] |
| 295 | |
| 296 | for i, speech in enumerate(raw_speech): |
| 297 | if len(speech.shape) > 1: |
| 298 | logger.warning( |
| 299 | "Only mono-channel audio is supported for input to %s. " |
| 300 | "We will take the mean of the channels to convert to mono.", |
| 301 | self.__class__.__name__, |
| 302 | ) |
| 303 | raw_speech[i] = speech.mean(-1) |
| 304 | |
| 305 | audio_clips = list[torch.Tensor]() |
| 306 | audio_num_clips = list[int]() |
| 307 | for audio in raw_speech: |
| 308 | clips = self.split_audio_into_clips(audio) |
| 309 | audio_clips.extend(clips) |
| 310 | audio_num_clips.append(len(clips)) |
| 311 | raw_speech = audio_clips |
| 312 | |
| 313 | audio_lengths = torch.tensor( |
| 314 | [len(speech) for speech in raw_speech], dtype=torch.long, device=device |
| 315 | ) |
| 316 | |
| 317 | max_length = max(len(speech) for speech in raw_speech) |
| 318 | input_features = self._pad_raw_speech(raw_speech, max_length, device) |
| 319 | input_features = self._apply_preemphasis(input_features, audio_lengths) |
| 320 | input_features = self._torch_extract_fbank_features(input_features, device) |
| 321 | input_features, attention_mask = self._normalize_mel_features( |
| 322 | input_features, audio_lengths |
| 323 | ) |
| 324 | |
| 325 | return { |
| 326 | "input_audio_features": input_features, |
| 327 | "feature_attention_mask": attention_mask, |
| 328 | "audio_num_clips": audio_num_clips, |
| 329 | } |
| 330 | |
| 331 | @staticmethod |
| 332 | def audio_length(raw_config: PreTrainedConfig, audio_tokens: int) -> int: |
| 333 | config = ExtractorConfig.from_hf_config(raw_config) |
| 334 | return int(audio_tokens * config.subsampling_factor * config.hop_length) |