| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # Copyright (c) Microsoft Corporation. |
| 4 | # Licensed under the MIT license. |
| 5 | # Code copied from Microsoft/MoE by Jacob Platin (jacobplatin@microsoft.com) |
| 6 | # but implemented by the Phi-Speech team |
| 7 | #!/usr/bin/env python3 |
| 8 | import abc |
| 9 | import math |
| 10 | import typing |
| 11 | from typing import Any, Literal |
| 12 | |
| 13 | import numpy as np |
| 14 | import torch |
| 15 | import torch.nn.functional as F |
| 16 | from torch import Tensor, nn |
| 17 | from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( |
| 18 | CheckpointWrapper, |
| 19 | ) |
| 20 | from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel |
| 21 | from transformers import PreTrainedConfig |
| 22 | |
| 23 | from vllm.model_executor.models.phi4mm_utils import ( |
| 24 | AbsolutePositionalEncoding, |
| 25 | ConvModule, |
| 26 | FeedForward, |
| 27 | MeanVarianceNormLayer, |
| 28 | MultiHeadedAttention, |
| 29 | MultiSequential, |
| 30 | NemoConvSubsampling, |
| 31 | T5RelativeAttentionLogitBias, |
| 32 | adaptive_enc_mask, |
| 33 | get_offset, |
| 34 | unfold_tensor, |
| 35 | ) |
| 36 | from vllm.utils.torch_utils import async_tensor_h2d |
| 37 | |
| 38 | |
| 39 | class ConformerEncoderLayer(nn.Module): |
| 40 | """ConformerEncoder Layer module. |
| 41 | for more details see conformer paper: |
| 42 | https://arxiv.org/abs/2005.08100 |
| 43 | This module implement the Conformer block layer. |
| 44 | |
| 45 | Args: |
| 46 | d_model: int |
| 47 | attention dim. |
| 48 | ext_pw_out_channel: int |
| 49 | if > 0, ext_pw_out_channel is a dim channel size |
| 50 | for the last pointwise conv after swish activation. |
| 51 | depthwise_seperable_out_channel: int |
| 52 | if set different to 0, the number of |
| 53 | depthwise_seperable_out_channel will be used as a |
| 54 | channel_out of the second conv1d layer. |
| 55 | otherwise, it equals to 0, the second conv1d layer is skipped. |
| 56 | depthwise_multiplier: int |
| 57 | number of input_dim channels duplication. this value |
| 58 | will be used to compute the hidden channels of the Conv1D. |
| 59 | n_head: int |
| 60 | the number of heads for multihead attention module. |
| 61 | d_ffn: int |
| 62 | output size of the feed_forward blocks. |
| 63 | ext_pw_kernel_size: int |
| 64 | kernel size of the conv pointwise of the conformer. |
| 65 | kernel_size: int |
| 66 | kernel size. |
| 67 | dropout_rate: float |
| 68 | dropout rate. |
| 69 | causal: bool, optional |
| 70 | if set to True, convolution have no access |
| 71 | to future frames. default False. |
| 72 | batch_norm: bool, optional |
| 73 | if set to True, apply batchnorm before activation |
| 74 | in ConvModule layer of the conformer. |
| 75 | default False |
| 76 | activation: str, optional |
| 77 | activation function name, |
| 78 | one of ["relu", "swish", "sigmoid"], |
| 79 | sigmoid activation is only used with "glu_in_fnn=True", |
| 80 | default "relu". |
| 81 | chunk_se: int, optional |
| 82 | 0 for offline SE. |
| 83 | 1 for streaming SE, where mean is computed |
| 84 | by accumulated history until current chunk_se. |
| 85 | 2 for streaming SE, where mean is computed |
| 86 | by only the current chunk. |
| 87 | default 0. |
| 88 | chunk_size: int, optional |
| 89 | chunk_size for cnn. default 18 |
| 90 | conv_activation: str, optional |
| 91 | activation function used in ConvModule part |
| 92 | of the conformer, default "relu". |
| 93 | conv_glu_type: str, optional |
| 94 | activation function used for the glu inside |
| 95 | the ConvModule part of the conformer. |
| 96 | default: "sigmoid". |
| 97 | bias_in_glu: bool, optional |
| 98 | if set to True, use additive bias in the weight module |
| 99 | before GLU. |
| 100 | linear_glu_in_convm: bool, optional |
| 101 | if set to True, use GLULinear module, |
| 102 | otherwise, used GLUPointWiseConv module. |
| 103 | default to False. |
| 104 | attention_inner_dim: int, optional |
| 105 | if equal to -1, attention dim for linears k/q/v is |
| 106 | equal to d_model. otherwise attention_inner_dim is used. |
| 107 | default -1. |
| 108 | attention_glu_type: str, optional |
| 109 | activation function for glu used in the multihead attention, |
| 110 | default "swish". |
| 111 | activation_checkpointing: str, optional |
| 112 | a dictionary of {"module","interval","offload"}, where |
| 113 | "module": str |
| 114 | accept ["transformer", "attention"] to select |
| 115 | which module should do activation checkpointing. |
| 116 | "interval": int, default 1, |
| 117 | interval of applying activation checkpointing, |
| 118 | interval = 1 means that we apply checkpointing |
| 119 | on every layer (if activation), otherwise, |
| 120 | we apply it every x interval. |
| 121 | "offload": bool, default False, |
| 122 | if set to True, we offload activation to cpu and |
| 123 | reload it during backward, otherwise, |
| 124 | we recalculate activation in backward. |
| 125 | default "". |
| 126 | export: bool, optional |
| 127 | if set to True, it removes the padding from convolutional layers |
| 128 | and allow the onnx conversion for inference. |
| 129 | default False. |
| 130 | use_pt_scaled_dot_product_attention: bool, optional |
| 131 | if set to True, use pytorch's scaled dot product attention |
| 132 | implementation in training. |
| 133 | attn_group_sizes: int, optional |
| 134 | the number of groups to use for attention, default 1 |
| 135 | (Multi-Head Attention), |
| 136 | 1 = typical Multi-Head Attention, |
| 137 | 1 < attn_group_sizes < attention_heads = Grouped-Query Attention |
| 138 | attn_group_sizes = attention_heads = Multi-Query Attention |
| 139 | |
| 140 | """ |
| 141 | |
| 142 | def __init__( |
| 143 | self, |
| 144 | d_model: int = 512, |
| 145 | ext_pw_out_channel: int = 0, |
| 146 | depthwise_seperable_out_channel: int = 256, |
| 147 | depthwise_multiplier: int = 1, |
| 148 | n_head: int = 4, |
| 149 | d_ffn: int = 2048, |
| 150 | ext_pw_kernel_size: int = 1, |
| 151 | kernel_size: int = 3, |
| 152 | dropout_rate: float = 0.1, |
| 153 | causal: bool = False, |
| 154 | batch_norm: bool = False, |
| 155 | activation: str = "relu", |
| 156 | chunk_se: int = 0, |
| 157 | chunk_size: int | list[int] = 18, |
| 158 | conv_activation: str = "relu", |
| 159 | conv_glu_type: str = "sigmoid", |
| 160 | bias_in_glu: bool = True, |
| 161 | linear_glu_in_convm: bool = False, |
| 162 | attention_inner_dim: int = -1, |
| 163 | attention_glu_type: str = "swish", |
| 164 | activation_checkpointing: str = "", |
| 165 | export: bool = False, |
| 166 | use_pt_scaled_dot_product_attention: bool = False, |
| 167 | attn_group_sizes: int = 1, |
| 168 | ) -> None: |
| 169 | super().__init__() |
| 170 | |
| 171 | self.feed_forward_in = FeedForward( |
| 172 | d_model=d_model, |
| 173 | d_inner=d_ffn, |
| 174 | dropout_rate=dropout_rate, |
| 175 | activation=activation, |
| 176 | bias_in_glu=bias_in_glu, |
| 177 | ) |
| 178 | |
| 179 | self.self_attn = MultiHeadedAttention( |
| 180 | n_head, |
| 181 | d_model, |
| 182 | dropout_rate, |
| 183 | attention_inner_dim, |
| 184 | attention_glu_type, |
| 185 | bias_in_glu, |
| 186 | use_pt_scaled_dot_product_attention=use_pt_scaled_dot_product_attention, |
| 187 | group_size=attn_group_sizes, |
| 188 | ) |
| 189 | self.conv = ConvModule( |
| 190 | d_model, |
| 191 | ext_pw_out_channel, |
| 192 | depthwise_seperable_out_channel, |
| 193 | ext_pw_kernel_size, |
| 194 | kernel_size, |
| 195 | depthwise_multiplier, |
| 196 | dropout_rate, |
| 197 | causal, |
| 198 | batch_norm, |
| 199 | chunk_se, |
| 200 | chunk_size, |
| 201 | conv_activation, |
| 202 | conv_glu_type, |
| 203 | bias_in_glu, |
| 204 | linear_glu_in_convm, |
| 205 | export=export, |
| 206 | ) |
| 207 | |
| 208 | self.feed_forward_out = FeedForward( |
| 209 | d_model=d_model, |
| 210 | d_inner=d_ffn, |
| 211 | dropout_rate=dropout_rate, |
| 212 | activation=activation, |
| 213 | bias_in_glu=bias_in_glu, |
| 214 | ) |
| 215 | |
| 216 | self.layer_norm_att = nn.LayerNorm(d_model) |
| 217 | self.layer_norm = nn.LayerNorm(d_model) |
| 218 | |
| 219 | def forward( |
| 220 | self, |
| 221 | x: torch.Tensor, |
| 222 | pos_k: torch.Tensor, |
| 223 | pos_v: torch.Tensor, |
| 224 | mask: torch.Tensor, |
| 225 | relative_attention_bias: Tensor | None = None, |
| 226 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 227 | """ConformerEncoder forward. |
| 228 | |
| 229 | Args: |
| 230 | x: input feature of shape (batch, max_time_in, size) |
| 231 | pos_k: positional key embedding. |
| 232 | pos_v: positional value embedding. |
| 233 | mask: mask for x (batch, max_time_in) |
| 234 | relative_attention_bias: bias added to attention logits w.r.t. |
| 235 | relative positions (1, n_head, time1, time2) |
| 236 | |
| 237 | """ |
| 238 | x = x + 0.5 * self.feed_forward_in(x) |
| 239 | norm_x = self.layer_norm_att(x) |
| 240 | |
| 241 | x = x + self.self_attn( |
| 242 | norm_x, |
| 243 | norm_x, |
| 244 | norm_x, |
| 245 | pos_k, |
| 246 | pos_v, |
| 247 | mask, |
| 248 | relative_attention_bias=relative_attention_bias, |
| 249 | ) |
| 250 | x = x + self.conv(x) |
| 251 | x = x + 0.5 * self.feed_forward_out(x) |
| 252 | |
| 253 | out = self.layer_norm(x) |
| 254 | |
| 255 | return out, pos_k, pos_v, mask |
| 256 | |
| 257 | |
| 258 | class TransformerEncoderBase(abc.ABC, nn.Module): |
| 259 | """The Base class for Transformer based encoders. |
| 260 | |
| 261 | Please set causal = True in streaming model |
| 262 | Args: |
| 263 | input_size: int |
| 264 | input feature dimension. |
| 265 | chunk_size: int, list(int) |
| 266 | Number of frames for each chunk |
| 267 | This variable can take 2 forms: |
| 268 | int: Used for inference, or single chunk size training |
| 269 | list(int) : Used only for variable chunk size training |
| 270 | Some examples for the 2 cases: |
| 271 | chunk_size = 12 |
| 272 | chunk_size = [6, 8, 12, 24] |
| 273 | left_chunk: int, list(int) |
| 274 | Number of chunks used for masking in streaming mode. |
| 275 | This variable can take 2 forms: |
| 276 | int: Used for inference, or single chunk size training |
| 277 | list(int) : Used only for variable chunk size training. When |
| 278 | chunk_size is a list, left_chunk must be a list with same length. |
| 279 | Some examples for the 2 cases: |
| 280 | left_chunk = 6 |
| 281 | left_chunk = [12, 9, 6, 3] |
| 282 | attention_dim: int, optional |
| 283 | attention dimension. default 256. |
| 284 | attention_heads: int, optional |
| 285 | the number of heads. default 4 |
| 286 | input_layer: str, optional |
| 287 | input layer type before Conformer, |
| 288 | one of ["linear", "conv2d", "custom", "vgg2l", "embed"], |
| 289 | default "conv2d" |
| 290 | cnn_out: int, optional |
| 291 | the number of CNN channels before Conformer. |
| 292 | default -1. |
| 293 | cnn_layer_norm: bool, optional |
| 294 | layer norm between Conformer and the first CNN. |
| 295 | default False. |
| 296 | time_reduction: int, optional |
| 297 | time reduction factor |
| 298 | default 4 |
| 299 | dropout_rate: float, optional |
| 300 | dropout rate. default 0.1 |
| 301 | padding_idx: int, optional |
| 302 | padding index for input_layer=embed |
| 303 | default -1 |
| 304 | relative_attention_bias_args: dict, optional |
| 305 | use more efficient scalar bias-based relative multihead attention |
| 306 | (Q*K^T + B) implemented in cmb.basics.embedding. |
| 307 | [T5/ALiBi]RelativeAttentionLogitBias |
| 308 | usage: relative_attention_bias_args={"type": t5/alibi} |
| 309 | additional method-specific arguments can be provided (see |
| 310 | transformer_base.py) |
| 311 | positional_dropout_rate: float, optional |
| 312 | dropout rate after positional encoding. default 0.0 |
| 313 | nemo_conv_settings: dict, optional |
| 314 | A dictionary of settings for NeMo Subsampling. |
| 315 | default None |
| 316 | conv2d_extra_padding: str, optional |
| 317 | Add extra padding in conv2d subsampling layers. Choices are |
| 318 | (feat, feat_time, none, True). |
| 319 | if True or feat_time, the extra padding is added into non full |
| 320 | supraframe utts in batch. |
| 321 | Default: none |
| 322 | attention_group_size: int, optional |
| 323 | the number of groups to use for attention, default 1 |
| 324 | (Multi-Head Attention), |
| 325 | 1 = typical Multi-Head Attention, |
| 326 | 1 < attention_group_size < attention_heads = Grouped-Query |
| 327 | Attention |
| 328 | attention_group_size = attention_heads = Multi-Query Attention |
| 329 | """ |
| 330 | |
| 331 | def __init__( |
| 332 | self, |
| 333 | input_size: int, |
| 334 | chunk_size: int | list[int], |
| 335 | left_chunk: int | list[int], |
| 336 | attention_dim: int = 256, |
| 337 | attention_heads: int = 4, |
| 338 | input_layer: str = "nemo_conv", |
| 339 | cnn_out: int = -1, |
| 340 | cnn_layer_norm: bool = False, |
| 341 | time_reduction: int = 4, |
| 342 | dropout_rate: float = 0.0, |
| 343 | padding_idx: int = -1, |
| 344 | relative_attention_bias_args: dict[str, Any] | None = None, |
| 345 | positional_dropout_rate: float = 0.0, |
| 346 | nemo_conv_settings: dict[str, Any] | None = None, |
| 347 | conv2d_extra_padding: Literal["feat", "feat_time", "none", True] = "none", |
| 348 | attention_group_size: int = 1, |
| 349 | encoder_embedding_config: dict[str, Any] | None = None, |
| 350 | ) -> None: |
| 351 | super().__init__() |
| 352 | self.input_size = input_size |
| 353 | self.input_layer = input_layer |
| 354 | self.chunk_size = chunk_size |
| 355 | self.left_chunk = left_chunk |
| 356 | self.attention_dim = attention_dim |
| 357 | self.num_heads = attention_heads |
| 358 | self.attention_group_size = attention_group_size |
| 359 | self.time_reduction = time_reduction |
| 360 | self.nemo_conv_settings = nemo_conv_settings |
| 361 | self.encoder_embedding_config = encoder_embedding_config |
| 362 | |
| 363 | if self.input_layer == "nemo_conv": |
| 364 | default_nemo_conv_settings = { |
| 365 | "subsampling": "dw_striding", |
| 366 | "subsampling_factor": self.time_reduction, |
| 367 | "feat_in": input_size, |
| 368 | "feat_out": attention_dim, |
| 369 | "conv_channels": 256, |
| 370 | "subsampling_conv_chunking_factor": 1, |
| 371 | "activation": nn.ReLU(), |
| 372 | "is_causal": False, |
| 373 | } |
| 374 | # Override any of the defaults with the incoming, user settings |
| 375 | if nemo_conv_settings: |
| 376 | default_nemo_conv_settings.update(nemo_conv_settings) |
| 377 | for i in ["subsampling_factor", "feat_in", "feat_out"]: |
| 378 | assert i not in nemo_conv_settings, ( |
| 379 | "{i} should be specified outside of the NeMo dictionary" |
| 380 | ) |
| 381 | |
| 382 | self.embed = NemoConvSubsampling( |
| 383 | **default_nemo_conv_settings, |
| 384 | ) |
| 385 | else: |
| 386 | raise ValueError("unknown input_layer: " + input_layer) |
| 387 | |
| 388 | self.pos_emb = AbsolutePositionalEncoding( |
| 389 | attention_dim, positional_dropout_rate |
| 390 | ) |
| 391 | |
| 392 | self.relative_attention_bias_type = ( |
| 393 | relative_attention_bias_args.get("type") |
| 394 | if relative_attention_bias_args |
| 395 | else None |
| 396 | ) |
| 397 | if self.relative_attention_bias_type == "t5": |
| 398 | assert relative_attention_bias_args is not None |
| 399 | assert self.num_heads % self.attention_group_size == 0, ( |
| 400 | "attention_group_size must divide n_head" |
| 401 | ) |
| 402 | self.relative_attention_bias_layer = T5RelativeAttentionLogitBias( |
| 403 | self.num_heads // self.attention_group_size, |
| 404 | max_distance=relative_attention_bias_args.get( |
| 405 | "t5_bias_max_distance", 1000 |
| 406 | ), |
| 407 | symmetric=relative_attention_bias_args.get("t5_bias_symmetric", False), |
| 408 | ) |
| 409 | else: |
| 410 | raise NotImplementedError |
| 411 | |
| 412 | assert self.encoder_embedding_config is not None |
| 413 | self.encoder_embedding = MeanVarianceNormLayer( |
| 414 | self.encoder_embedding_config["input_size"] |
| 415 | ) |
| 416 | |
| 417 | def compute_lens_change( |
| 418 | self, feature_lens: int | torch.Tensor |
| 419 | ) -> int | torch.Tensor: |
| 420 | """feature_lens: int |
| 421 | return updated feature lens. |
| 422 | |
| 423 | This used to return a different lambda function for each case that |
| 424 | computed the right thing. That does not work within Torchscript. |
| 425 | If you really need this to be faster, create nn.Module()-s for all |
| 426 | the cases and return one of them. Torchscript does support that. |
| 427 | """ |
| 428 | if self.input_layer == "nemo_conv": |
| 429 | assert self.nemo_conv_settings is not None |
| 430 | # Handle the special causal case |
| 431 | subsampling_causal_cond = self.nemo_conv_settings.get( |
| 432 | "subsampling", "dw_striding" |
| 433 | ) in [ |
| 434 | "dw_striding", |
| 435 | "striding", |
| 436 | "striding_conv1d", |
| 437 | ] |
| 438 | is_causal = self.nemo_conv_settings.get("is_causal", False) |
| 439 | if is_causal and subsampling_causal_cond: |
| 440 | feature_lens_remainder = feature_lens % self.time_reduction |
| 441 | if isinstance(feature_lens, Tensor): |
| 442 | lens_change = torch.ceil(feature_lens / self.time_reduction).long() |
| 443 | lens_change[feature_lens_remainder != 1] += 1 |
| 444 | return lens_change |
| 445 | |
| 446 | lens_change = math.ceil(feature_lens / self.time_reduction) |
| 447 | if feature_lens_remainder != 1: |
| 448 | lens_change += 1 |
| 449 | return lens_change |
| 450 | ceil_func = math.ceil if isinstance(feature_lens, int) else torch.ceil |
| 451 | return ceil_func(feature_lens / self.time_reduction) |
| 452 | raise ValueError("unknown input_layer: " + self.input_layer) |
| 453 | |
| 454 | @abc.abstractmethod |
| 455 | def forward(self) -> Any: |
| 456 | """Abstract forward method implementation.""" |
| 457 | |
| 458 | def _chunk_size_selection( |
| 459 | self, |
| 460 | chunk_size: int | list[int] | None = None, |
| 461 | left_chunk: int | list[int] | None = None, |
| 462 | ) -> tuple[int, int]: |
| 463 | """If chunk size is a list, we will randomly select a chunk size.""" |
| 464 | if chunk_size is None: |
| 465 | chunk_size = self.chunk_size |
| 466 | if left_chunk is None: |
| 467 | left_chunk = self.left_chunk |
| 468 | if isinstance(chunk_size, list): |
| 469 | # Variable chunk size during training |
| 470 | chunk_size_index = int( |
| 471 | torch.randint(low=0, high=len(chunk_size), size=(1,)) |
| 472 | ) |
| 473 | chunk_size_train_eff = chunk_size[chunk_size_index] |
| 474 | if not isinstance(left_chunk, list): |
| 475 | raise ValueError( |
| 476 | "Since chunk_size is a list, left_chunk must be a list" |
| 477 | ) |
| 478 | if len(left_chunk) != len(chunk_size): |
| 479 | raise ValueError( |
| 480 | "The length of left_chunk must be the same as length of chunk_size." |
| 481 | ) |
| 482 | left_chunk_train_eff = left_chunk[chunk_size_index] |
| 483 | else: |
| 484 | assert chunk_size is not None |
| 485 | assert isinstance(left_chunk, int) |
| 486 | chunk_size_train_eff = chunk_size |
| 487 | left_chunk_train_eff = left_chunk |
| 488 | |
| 489 | return chunk_size_train_eff, left_chunk_train_eff |
| 490 | |
| 491 | def _get_embed_class(self, embed: nn.Module) -> nn.Module: |
| 492 | # pylint: disable=protected-access |
| 493 | is_embed_using_act_chkpt = isinstance(embed, CheckpointWrapper) |
| 494 | is_embed_fsdp_wrapped = isinstance(embed, FullyShardedDataParallel) |
| 495 | embed_class = embed |
| 496 | if is_embed_using_act_chkpt: |
| 497 | embed_class = embed._checkpoint_wrapped_module |
| 498 | if is_embed_fsdp_wrapped: |
| 499 | embed_class = embed.module |
| 500 | return embed_class |
| 501 | |
| 502 | def _forward_embeddings_core( |
| 503 | self, input_tensor: torch.Tensor, masks: torch.Tensor |
| 504 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 505 | embed_class = self._get_embed_class(self.embed) |
| 506 | assert isinstance(embed_class, NemoConvSubsampling) |
| 507 | input_tensor, masks = self.embed(input_tensor, masks) |
| 508 | return input_tensor, masks |
| 509 | |
| 510 | def _position_embedding( |
| 511 | self, input_tensor: torch.Tensor |
| 512 | ) -> tuple[torch.Tensor | None, torch.Tensor | None]: |
| 513 | pos_k = None |
| 514 | pos_v = None |
| 515 | if self.relative_attention_bias_layer is None: |
| 516 | input_tensor = self.pos_emb( |
| 517 | input_tensor |
| 518 | ) # default to add abs sinusoid embedding |
| 519 | return pos_k, pos_v |
| 520 | |
| 521 | def _streaming_mask( |
| 522 | self, |
| 523 | seq_len: int, |
| 524 | batch_size: int, |
| 525 | chunk_size: int | list[int], |
| 526 | left_chunk: int | list[int] | None, |
| 527 | ) -> torch.Tensor: |
| 528 | chunk_size_train_eff, left_chunk_train_eff = self._chunk_size_selection( |
| 529 | chunk_size, left_chunk |
| 530 | ) |
| 531 | |
| 532 | # Create mask matrix for streaming |
| 533 | # S stores start index. if chunksize is 18, s is [0,18,36,....] |
| 534 | chunk_start_idx = np.arange(0, seq_len, chunk_size_train_eff) |
| 535 | |
| 536 | enc_streaming_mask = ( |
| 537 | adaptive_enc_mask( |
| 538 | seq_len, chunk_start_idx, left_window=left_chunk_train_eff |
| 539 | ) |
| 540 | .unsqueeze(0) |
| 541 | .expand([batch_size, -1, -1]) |
| 542 | ) |
| 543 | return enc_streaming_mask |
| 544 | |
| 545 | @typing.overload |
| 546 | def forward_embeddings( |
| 547 | self, |
| 548 | xs_pad: torch.Tensor, |
| 549 | masks: torch.Tensor, |
| 550 | chunk_size_nc: None = None, |
| 551 | left_chunk_nc: int | list[int] | None = None, |
| 552 | ) -> tuple[ |
| 553 | torch.Tensor, |
| 554 | torch.Tensor | None, |
| 555 | torch.Tensor | None, |
| 556 | torch.Tensor, |
| 557 | torch.Tensor, |
| 558 | ]: ... |
| 559 | |
| 560 | @typing.overload |
| 561 | def forward_embeddings( |
| 562 | self, |
| 563 | xs_pad: torch.Tensor, |
| 564 | masks: torch.Tensor, |
| 565 | chunk_size_nc: int | list[int], |
| 566 | left_chunk_nc: int | list[int] | None = None, |
| 567 | ) -> tuple[ |
| 568 | torch.Tensor, |
| 569 | torch.Tensor | None, |
| 570 | torch.Tensor | None, |
| 571 | torch.Tensor, |
| 572 | torch.Tensor, |
| 573 | torch.Tensor, |
| 574 | ]: ... |
| 575 | |
| 576 | def forward_embeddings( |
| 577 | self, |
| 578 | xs_pad: torch.Tensor, |
| 579 | masks: torch.Tensor, |
| 580 | chunk_size_nc: int | list[int] | None = None, |
| 581 | left_chunk_nc: int | list[int] | None = None, |
| 582 | ) -> ( |
| 583 | tuple[ |
| 584 | torch.Tensor, |
| 585 | torch.Tensor | None, |
| 586 | torch.Tensor | None, |
| 587 | torch.Tensor, |
| 588 | torch.Tensor, |
| 589 | ] |
| 590 | | tuple[ |
| 591 | torch.Tensor, |
| 592 | torch.Tensor | None, |
| 593 | torch.Tensor | None, |
| 594 | torch.Tensor, |
| 595 | torch.Tensor, |
| 596 | torch.Tensor, |
| 597 | ] |
| 598 | ): |
| 599 | """Forwarding the inputs through the top embedding layers. |
| 600 | |
| 601 | Args: |
| 602 | xs_pad: torch.Tensor |
| 603 | input tensor |
| 604 | masks: torch.Tensor |
| 605 | input mask |
| 606 | chunk_size_nc: (optional, default is None) chunk size for |
| 607 | non-causal layers |
| 608 | left_chunk_nc: (optional, default is None) # of left chunks for |
| 609 | non-causal layers |
| 610 | |
| 611 | """ |
| 612 | # pylint: disable=R0915 |
| 613 | # get new lens. |
| 614 | seq_len = int(self.compute_lens_change(xs_pad.shape[1])) |
| 615 | if seq_len <= 0: |
| 616 | raise ValueError( |
| 617 | f"""The sequence length after time reduction is invalid: |
| 618 | {seq_len}. Your input feature is too short. Consider |
| 619 | filtering out the very short sentence from data |
| 620 | loader""", |
| 621 | ) |
| 622 | |
| 623 | batch_size = xs_pad.shape[0] |
| 624 | |
| 625 | enc_streaming_mask = self._streaming_mask( |
| 626 | seq_len, batch_size, self.chunk_size, self.left_chunk |
| 627 | ) |
| 628 | device = xs_pad.device |
| 629 | enc_streaming_mask = async_tensor_h2d(enc_streaming_mask.contiguous(), device) |
| 630 | xs_pad = xs_pad.to(device) |
| 631 | |
| 632 | input_tensor = xs_pad |
| 633 | input_tensor, masks = self._forward_embeddings_core(input_tensor, masks) |
| 634 | |
| 635 | streaming_mask = enc_streaming_mask |
| 636 | if streaming_mask is not None and masks is not None: |
| 637 | hs_mask = masks & streaming_mask |
| 638 | elif masks is not None: |
| 639 | hs_mask = masks |
| 640 | else: |
| 641 | hs_mask = streaming_mask |
| 642 | |
| 643 | if chunk_size_nc is not None: |
| 644 | enc_streaming_mask_nc = self._streaming_mask( |
| 645 | seq_len, batch_size, chunk_size_nc, left_chunk_nc |
| 646 | ) |
| 647 | if device.type != "cpu": |
| 648 | enc_streaming_mask_nc = async_tensor_h2d( |
| 649 | enc_streaming_mask_nc.contiguous(), device |
| 650 | ) |
| 651 | if masks is not None: |
| 652 | hs_mask_nc = masks & enc_streaming_mask_nc |
| 653 | else: |
| 654 | hs_mask_nc = enc_streaming_mask_nc |
| 655 | else: |
| 656 | hs_mask_nc = None |
| 657 | |
| 658 | pos_k, pos_v = self._position_embedding(input_tensor) |
| 659 | |
| 660 | if chunk_size_nc is None: |
| 661 | return input_tensor, pos_k, pos_v, hs_mask, masks |
| 662 | return input_tensor, pos_k, pos_v, hs_mask, masks, hs_mask_nc |
| 663 | |
| 664 | def get_offset(self) -> int: |
| 665 | """Returns offset used when retaining inputs for decoding. |
| 666 | |
| 667 | This is essentially, how many additional frames have to be added to |
| 668 | the front-end CNN input to ensure it can produce a single output. |
| 669 | So if the "padding" parameter is 0, typically offset will be > 0. |
| 670 | """ |
| 671 | return get_offset(self.input_layer, self.time_reduction) |
| 672 | |
| 673 | |
| 674 | class ConformerEncoder(TransformerEncoderBase): |
| 675 | """ConformerEncoder module. |
| 676 | see original paper for more details: |
| 677 | https://arxiv.org/abs/2005.08100 |
| 678 | |
| 679 | Please set causal = True in streaming model |
| 680 | Args: |
| 681 | input_size: int |
| 682 | input feature dimension. |
| 683 | chunk_size: int, list(int) |
| 684 | Number of frames for each chunk |
| 685 | This variable can take 2 forms: |
| 686 | int: Used for inference, or single chunk size training |
| 687 | list(int) : Used only for variable chunk size training |
| 688 | Some examples for the 2 cases: |
| 689 | chunk_size = 12 |
| 690 | chunk_size = [6, 8, 12, 24] |
| 691 | left_chunk: int, list(int) |
| 692 | Number of chunks used for masking in streaming mode. |
| 693 | This variable can take 2 forms: |
| 694 | int: Used for inference, or single chunk size training |
| 695 | list(int) : Used only for variable chunk size training. When |
| 696 | chunk_size is a list, left_chunk must be a list with same length. |
| 697 | Some examples for the 2 cases: |
| 698 | left_chunk = 6 |
| 699 | left_chunk = [12, 9, 6, 3] |
| 700 | num_lang: int |
| 701 | This parameter is used to store the number of languages in the |
| 702 | lang_dict, only used for multiseed/multilingual models. |
| 703 | default None. |
| 704 | attention_dim: int, optional |
| 705 | attention dimension. default 256. |
| 706 | attention_heads: int, optional |
| 707 | the number of heads. default 4 |
| 708 | linear_units: |
| 709 | the number of units of position-wise feed forward. |
| 710 | default 2048 |
| 711 | num_block: |
| 712 | number of Transformer layer. default 6 |
| 713 | dropout_rate: float, optional |
| 714 | dropout rate. default 0.1 |
| 715 | input_layer: str, optional |
| 716 | input layer type before Conformer, |
| 717 | one of ["linear", "conv2d", "custom", "vgg2l", "embed"], |
| 718 | default "conv2d" |
| 719 | causal: bool, optional |
| 720 | if set to True, convolution have no access |
| 721 | to future frames. default False. |
| 722 | batch_norm: bool, optional |
| 723 | if set to True, apply batchnorm before activation |
| 724 | in ConvModule layer of the conformer. |
| 725 | default False |
| 726 | cnn_out: int, optional |
| 727 | the number of CNN channels before Conformer. |
| 728 | default -1. |
| 729 | cnn_layer_norm: bool, optional |
| 730 | layer norm between Conformer and the first CNN. |
| 731 | default False. |
| 732 | ext_pw_out_channel: int, optional |
| 733 | the number of channel for CNN |
| 734 | before depthwise_separable_CNN. |
| 735 | If 0 then use linear. default 0. |
| 736 | ext_pw_kernel_size: int, optional |
| 737 | kernel size of N before depthwise_separable_CNN. |
| 738 | only work for ext_pw_out_channel > 0. |
| 739 | default 1 |
| 740 | depthwise_seperable_out_channel: int, optional |
| 741 | the number of channel for |
| 742 | depthwise_separable_CNN. |
| 743 | default 256. |
| 744 | depthwise_multiplier: int, optional |
| 745 | the number of multiplier for |
| 746 | depthwise_separable_CNN. |
| 747 | default 1. |
| 748 | chunk_se: int, optional |
| 749 | 0 for offline SE. |
| 750 | 1 for streaming SE, where mean is computed |
| 751 | by accumulated history until current chunk_se. |
| 752 | 2 for streaming SE, where mean is computed |
| 753 | by only the current chunk. |
| 754 | default 0. |
| 755 | kernel_size: int, optional |
| 756 | the number of kernels for depthwise_separable_CNN. |
| 757 | default 3. |
| 758 | activation: str, optional |
| 759 | FeedForward block activation. |
| 760 | one of ["relu", "swish", "sigmoid"] |
| 761 | default "relu". |
| 762 | conv_activation: str, optional |
| 763 | activation function used in ConvModule part |
| 764 | of the conformer, default "relu". |
| 765 | conv_glu_type: str, optional |
| 766 | activation used use glu in depthwise_separable_CNN, |
| 767 | default "sigmoid" |
| 768 | bias_in_glu: bool, optional |
| 769 | if set to True, use additive bias in the weight module |
| 770 | before GLU. default True |
| 771 | linear_glu_in_convm: bool, optional |
| 772 | if set to True, use GLULinear module, |
| 773 | otherwise, used GLUPointWiseConv module. |
| 774 | default to False. |
| 775 | attention_glu_type: str |
| 776 | only work for glu_in_attention !=0 |
| 777 | default "swish". |
| 778 | export: bool, optional |
| 779 | if set to True, it removes the padding from convolutional layers |
| 780 | and allow the onnx conversion for inference. |
| 781 | default False. |
| 782 | activation_checkpointing: str, optional |
| 783 | a dictionarry of {"module","interval","offload"}, where |
| 784 | "module": str |
| 785 | accept ["transformer", "attention"] to select |
| 786 | which module should do activation checkpointing. |
| 787 | "interval": int, default 1, |
| 788 | interval of applying activation checkpointing, |
| 789 | interval = 1 means that we apply checkpointing |
| 790 | on every layer (if activation), otherwise, |
| 791 | we apply it every x interval. |
| 792 | "offload": bool, default False, |
| 793 | if set to True, we offload activation to cpu and |
| 794 | reload it during backward, otherwise, |
| 795 | we recalculate activation in backward. |
| 796 | default "". |
| 797 | extra_layer_output_idx: int |
| 798 | the layer index to be exposed. |
| 799 | relative_attention_bias_args: dict, optional |
| 800 | use more efficient scalar bias-based relative multihead attention |
| 801 | (Q*K^T + B) implemented in cmb.basics.embedding. |
| 802 | [T5/ALiBi]RelativeAttentionLogitBias |
| 803 | usage: relative_attention_bias_args={"type": t5/alibi} |
| 804 | additional method-specific arguments can be provided (see |
| 805 | transformer_base.py) |
| 806 | time_reduction: int optional |
| 807 | time reduction factor |
| 808 | default 4 |
| 809 | use_pt_scaled_dot_product_attention: whether to use pytorch scaled |
| 810 | dot product attention in training. |
| 811 | Default: False |
| 812 | nemo_conv_settings: dict, optional |
| 813 | A dictionary of settings for NeMo Subsampling. |
| 814 | default: None |
| 815 | usage: nemo_conv_settings= |
| 816 | { |
| 817 | "subsampling": |
| 818 | dw_striding/striding/dw_striding_conv1d/striding_conv1d, |
| 819 | "conv_channels": int, |
| 820 | "subsampling_conv_chunking_factor": int, |
| 821 | "is_causal": True/False |
| 822 | } |
| 823 | conv2d_extra_padding: str, optional |
| 824 | Add extra padding in conv2d subsampling layers. Choices are |
| 825 | (feat, feat_time, none, True) |
| 826 | Default: none |
| 827 | replication_pad_for_subsample_embedding: For batched-streaming |
| 828 | decoding, use "replication" padding for the cache at start of |
| 829 | utterance. |
| 830 | Default: False |
| 831 | attention_group_size: int, optional |
| 832 | the number of groups to use for attention, default 1 |
| 833 | (Multi-Head Attention), |
| 834 | 1 = typical Multi-Head Attention, |
| 835 | 1 < attention_group_size < attention_heads = Grouped-Query |
| 836 | Attention |
| 837 | attention_group_size = attention_heads = Multi-Query Attention |
| 838 | """ |
| 839 | |
| 840 | extra_multi_layer_output_idxs: list[int] |
| 841 | |
| 842 | def __init__( # pylint: disable-all |
| 843 | self, |
| 844 | input_size: int, |
| 845 | chunk_size: int | list[int], |
| 846 | left_chunk: int | list[int], |
| 847 | num_lang: int | None = None, |
| 848 | attention_dim: int = 256, |
| 849 | attention_heads: int = 4, |
| 850 | linear_units: int = 2048, |
| 851 | num_blocks: int = 6, |
| 852 | dropout_rate: float = 0.1, |
| 853 | input_layer: str = "nemo_conv", |
| 854 | causal: bool = True, |
| 855 | batch_norm: bool = False, |
| 856 | cnn_out: int = -1, |
| 857 | cnn_layer_norm: bool = False, |
| 858 | ext_pw_out_channel: int = 0, |
| 859 | ext_pw_kernel_size: int = 1, |
| 860 | depthwise_seperable_out_channel: int = 256, |
| 861 | depthwise_multiplier: int = 1, |
| 862 | chunk_se: int = 0, |
| 863 | kernel_size: int = 3, |
| 864 | activation: str = "relu", |
| 865 | conv_activation: str = "relu", |
| 866 | conv_glu_type: str = "sigmoid", |
| 867 | bias_in_glu: bool = True, |
| 868 | linear_glu_in_convm: bool = False, |
| 869 | attention_glu_type: str = "swish", |
| 870 | export: bool = False, |
| 871 | extra_layer_output_idx: int = -1, |
| 872 | extra_multi_layer_output_idxs: list[int] = [], # noqa |
| 873 | activation_checkpointing: str = "", |
| 874 | relative_attention_bias_args: dict[str, Any] | None = None, |
| 875 | time_reduction: int = 4, |
| 876 | use_pt_scaled_dot_product_attention: bool = False, |
| 877 | nemo_conv_settings: dict[str, Any] | None = None, |
| 878 | conv2d_extra_padding: Literal["feat", "feat_time", "none", True] = "none", |
| 879 | replication_pad_for_subsample_embedding: bool = False, |
| 880 | attention_group_size: int = 1, |
| 881 | encoder_embedding_config: dict[str, Any] | None = None, |
| 882 | ) -> None: |
| 883 | super().__init__( |
| 884 | input_size, |
| 885 | chunk_size, |
| 886 | left_chunk, |
| 887 | attention_dim, |
| 888 | attention_heads, |
| 889 | input_layer, |
| 890 | cnn_out, |
| 891 | cnn_layer_norm, |
| 892 | time_reduction, |
| 893 | dropout_rate=dropout_rate, |
| 894 | relative_attention_bias_args=relative_attention_bias_args, |
| 895 | positional_dropout_rate=0.0, |
| 896 | nemo_conv_settings=nemo_conv_settings, |
| 897 | conv2d_extra_padding=conv2d_extra_padding, |
| 898 | attention_group_size=attention_group_size, |
| 899 | encoder_embedding_config=encoder_embedding_config, |
| 900 | ) |
| 901 | self.num_blocks = num_blocks |
| 902 | self.num_lang = num_lang |
| 903 | self.kernel_size = kernel_size |
| 904 | self.replication_pad_for_subsample_embedding: bool = ( |
| 905 | replication_pad_for_subsample_embedding |
| 906 | ) |
| 907 | assert self.num_heads % attention_group_size == 0, ( |
| 908 | "attention_group_size must divide n_head" |
| 909 | ) |
| 910 | self.num_heads_k = self.num_heads // attention_group_size |
| 911 | |
| 912 | self.encoders = MultiSequential( |
| 913 | *[ |
| 914 | ConformerEncoderLayer( |
| 915 | d_model=attention_dim, |
| 916 | ext_pw_out_channel=ext_pw_out_channel, |
| 917 | depthwise_seperable_out_channel=depthwise_seperable_out_channel, |
| 918 | depthwise_multiplier=depthwise_multiplier, |
| 919 | n_head=attention_heads, |
| 920 | d_ffn=linear_units, |
| 921 | ext_pw_kernel_size=ext_pw_kernel_size, |
| 922 | kernel_size=kernel_size, |
| 923 | dropout_rate=dropout_rate, |
| 924 | causal=causal, |
| 925 | batch_norm=batch_norm, |
| 926 | activation=activation, |
| 927 | chunk_se=chunk_se, |
| 928 | chunk_size=chunk_size, |
| 929 | conv_activation=conv_activation, |
| 930 | conv_glu_type=conv_glu_type, |
| 931 | bias_in_glu=bias_in_glu, |
| 932 | linear_glu_in_convm=linear_glu_in_convm, |
| 933 | attention_glu_type=attention_glu_type, |
| 934 | activation_checkpointing=activation_checkpointing, |
| 935 | export=export, |
| 936 | use_pt_scaled_dot_product_attention=use_pt_scaled_dot_product_attention, |
| 937 | attn_group_sizes=attention_group_size, |
| 938 | ) |
| 939 | for _ in range(num_blocks) |
| 940 | ] |
| 941 | ) |
| 942 | self.extra_layer_output_idx = extra_layer_output_idx |
| 943 | self.extra_multi_layer_output_idxs = extra_multi_layer_output_idxs |
| 944 | # Make a zeros scalar we can use in get_initial_state to determine |
| 945 | # the device and the needed dtype: |
| 946 | self.register_buffer("dev_type", torch.zeros(()), persistent=False) |
| 947 | |
| 948 | def init_relative_attention_bias( |
| 949 | self, input_tensor: torch.Tensor |
| 950 | ) -> torch.Tensor | None: |
| 951 | if self.relative_attention_bias_layer: |
| 952 | return self.relative_attention_bias_layer(input_tensor) |
| 953 | return None |
| 954 | |
| 955 | def calculate_hs_mask( |
| 956 | self, xs_pad: torch.Tensor, device: torch.device, mask: torch.Tensor | None |
| 957 | ) -> torch.Tensor: |
| 958 | max_audio_length = xs_pad.shape[1] |
| 959 | batch_size = xs_pad.shape[0] |
| 960 | enc_streaming_mask = self._streaming_mask( |
| 961 | max_audio_length, batch_size, self.chunk_size, self.left_chunk |
| 962 | ) |
| 963 | enc_streaming_mask = async_tensor_h2d(enc_streaming_mask.contiguous(), device) |
| 964 | if mask is None: |
| 965 | return enc_streaming_mask |
| 966 | |
| 967 | feature_lens = mask.sum(1) |
| 968 | padding_length = feature_lens |
| 969 | pad_mask = torch.arange(0, max_audio_length, device=device).expand( |
| 970 | padding_length.size(0), -1 |
| 971 | ) < padding_length.unsqueeze(1) |
| 972 | pad_mask = pad_mask.unsqueeze(1) |
| 973 | pad_mask = pad_mask & enc_streaming_mask |
| 974 | return pad_mask |
| 975 | |
| 976 | @torch.jit.ignore |
| 977 | def forward( |
| 978 | self, xs_pad: torch.Tensor, masks: torch.Tensor |
| 979 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 980 | """Conformer Forward function. |
| 981 | |
| 982 | Args: |
| 983 | xs_pad: torch.Tensor |
| 984 | input tensor |
| 985 | masks: torch.Tensor |
| 986 | post-embedding input lengths |
| 987 | |
| 988 | """ |
| 989 | xs_pad = self.encoder_embedding(xs_pad) |
| 990 | input_tensor, pos_k, pos_v, hs_mask, masks = self.forward_embeddings( |
| 991 | xs_pad, masks |
| 992 | ) |
| 993 | |
| 994 | unfolded = False |
| 995 | ori_bz, seq_len, D = input_tensor.shape |
| 996 | max_seq_len = 500 # maximum position for absolute positional encoding |
| 997 | if seq_len > max_seq_len: |
| 998 | # audio sequence is longer than max_seq_len, unfold it into chunks |
| 999 | # of max_seq_len |
| 1000 | unfolded = True |
| 1001 | # the unfold op will drop residual frames, pad it to the multiple |
| 1002 | # of max_seq_len |
| 1003 | if seq_len % max_seq_len > 0: |
| 1004 | chunk_pad_size = max_seq_len - (seq_len % max_seq_len) |
| 1005 | else: |
| 1006 | chunk_pad_size = 0 |
| 1007 | if chunk_pad_size > 0: |
| 1008 | input_tensor_pad = F.pad( |
| 1009 | input_tensor, (0, 0, 0, chunk_pad_size), "constant", 0 |
| 1010 | ) |
| 1011 | input_tensor = input_tensor_pad.to(input_tensor.device) |
| 1012 | input_tensor = unfold_tensor(input_tensor, max_seq_len) |
| 1013 | if masks is not None: |
| 1014 | # revise hs_mask here because the previous calculated hs_mask |
| 1015 | # did not consider extra pad |
| 1016 | subsampled_pad_mask = masks.squeeze( |
| 1017 | 1 |
| 1018 | ) # [bz, subsampled_unmask_seq_len] |
| 1019 | extra_padded_subsamlped_pad_mask = F.pad( |
| 1020 | subsampled_pad_mask, (0, chunk_pad_size), "constant", False |
| 1021 | ) # extra padding to the pad mask |
| 1022 | extra_padded_subsamlped_pad_mask = ( |
| 1023 | extra_padded_subsamlped_pad_mask.unsqueeze(-1).float() |
| 1024 | ) |
| 1025 | masks_unfold = unfold_tensor( |
| 1026 | extra_padded_subsamlped_pad_mask, max_seq_len |
| 1027 | ) # unfold the pad mask like we did to the input tensor |
| 1028 | masks_unfold = masks_unfold.squeeze( |
| 1029 | -1 |
| 1030 | ).bool() # unfold op does not support bool tensor |
| 1031 | else: |
| 1032 | masks_unfold = None |
| 1033 | hs_mask = self.calculate_hs_mask( |
| 1034 | input_tensor, input_tensor.device, masks_unfold |
| 1035 | ) # calculate hs_mask based on the unfolded pad mask |
| 1036 | |
| 1037 | # layer_emb = None |
| 1038 | |
| 1039 | relative_attention_bias = self.init_relative_attention_bias(input_tensor) |
| 1040 | |
| 1041 | _simplified_path = ( |
| 1042 | self.extra_layer_output_idx == -1 and relative_attention_bias is None |
| 1043 | ) |
| 1044 | |
| 1045 | if _simplified_path: |
| 1046 | input_tensor, *_ = self.encoders(input_tensor, pos_k, pos_v, hs_mask) |
| 1047 | else: |
| 1048 | for i, layer in enumerate(self.encoders): |
| 1049 | input_tensor, _, _, _ = layer( |
| 1050 | input_tensor, |
| 1051 | pos_k, |
| 1052 | pos_v, |
| 1053 | hs_mask, |
| 1054 | relative_attention_bias=relative_attention_bias, |
| 1055 | ) |
| 1056 | |
| 1057 | # if i == self.extra_layer_output_idx: |
| 1058 | # layer_emb = input_tensor |
| 1059 | |
| 1060 | if unfolded: |
| 1061 | embed_dim = input_tensor.shape[-1] |
| 1062 | input_tensor = input_tensor.reshape(ori_bz, -1, embed_dim) |
| 1063 | # if we ever padded before unfolding, we need to remove the padding |
| 1064 | if chunk_pad_size > 0: |
| 1065 | input_tensor = input_tensor[:, :-chunk_pad_size, :] |
| 1066 | |
| 1067 | return input_tensor, masks # , layer_emb |
| 1068 | |
| 1069 | |
| 1070 | class WindowQformer(nn.Module): |
| 1071 | """Window-level Qformer.""" |
| 1072 | |
| 1073 | def __init__( |
| 1074 | self, |
| 1075 | window_size: int = 8, |
| 1076 | num_queries: int = 1, |
| 1077 | num_blocks: int = 2, |
| 1078 | attention_dim: int = 512, |
| 1079 | attention_heads: int = 8, |
| 1080 | linear_units: int = 2048, |
| 1081 | dropout_rate: float = 0.0, |
| 1082 | normalize_before: bool = True, |
| 1083 | ): |
| 1084 | super().__init__() |
| 1085 | |
| 1086 | self.decoders = nn.ModuleList( |
| 1087 | [ |
| 1088 | nn.TransformerDecoderLayer( |
| 1089 | d_model=attention_dim, |
| 1090 | nhead=attention_heads, |
| 1091 | dim_feedforward=linear_units, |
| 1092 | dropout=dropout_rate, |
| 1093 | activation="relu", |
| 1094 | batch_first=True, |
| 1095 | norm_first=normalize_before, # TODO need to verify |
| 1096 | ) |
| 1097 | for _ in range(num_blocks) |
| 1098 | ] |
| 1099 | ) |
| 1100 | |
| 1101 | self.queries = nn.Parameter(torch.zeros(1, num_queries, attention_dim)) |
| 1102 | self.after_norm = ( |
| 1103 | nn.LayerNorm(attention_dim, eps=1e-12) if normalize_before else None |
| 1104 | ) |
| 1105 | self.window_size = window_size |
| 1106 | |
| 1107 | def forward( |
| 1108 | self, |
| 1109 | audio_embed: torch.Tensor, |
| 1110 | mask: torch.Tensor | None, |
| 1111 | embed_len: int | None = None, |
| 1112 | ) -> tuple[torch.Tensor, int | None]: |
| 1113 | """Forward decoder.""" |
| 1114 | # audio_embed: N x T x D => N x D x T |
| 1115 | |
| 1116 | audio_embed = audio_embed.transpose(1, 2) |
| 1117 | # audio_embed: N x D x 1 x T => N x DK x T' |
| 1118 | padding = audio_embed.shape[-1] % self.window_size |
| 1119 | if padding > 0: |
| 1120 | audio_embed = F.pad( |
| 1121 | audio_embed, (0, self.window_size - padding), "constant", 0 |
| 1122 | ) |
| 1123 | |
| 1124 | embed_chunk = F.unfold( |
| 1125 | audio_embed[..., None, :], |
| 1126 | kernel_size=(1, self.window_size), |
| 1127 | stride=(1, self.window_size), |
| 1128 | ) |
| 1129 | bsz, _, slen = embed_chunk.shape |
| 1130 | # N x D x K x T' |
| 1131 | embed_chunk = embed_chunk.view(bsz, -1, self.window_size, slen) |
| 1132 | # N x T' x K x D |
| 1133 | embed_chunk = embed_chunk.transpose(1, 3).contiguous() |
| 1134 | # NT' x K x D |
| 1135 | embed_chunk = embed_chunk.view(bsz * slen, self.window_size, -1) |
| 1136 | # NT' x 1 x D |
| 1137 | q = self.queries.expand(bsz * slen, -1, -1) |
| 1138 | for layer in self.decoders: |
| 1139 | q = layer(tgt=q, memory=embed_chunk, tgt_mask=None, memory_mask=mask) |
| 1140 | |
| 1141 | if self.after_norm is not None: |
| 1142 | q = self.after_norm(q) |
| 1143 | |
| 1144 | if embed_len is not None: |
| 1145 | embed_len = embed_len // self.window_size |
| 1146 | # N x T' x D |
| 1147 | out = q.view(bsz, slen, -1) |
| 1148 | |
| 1149 | return out, embed_len |
| 1150 | |
| 1151 | |
| 1152 | class AudioEmbedding(nn.Module): |
| 1153 | """Image embedding.""" |
| 1154 | |
| 1155 | def __init__(self, config: PreTrainedConfig, **kwargs: Any) -> None: |
| 1156 | super().__init__() |
| 1157 | self.config = config |
| 1158 | # n_embed or hidden_size for text LM |
| 1159 | hidden_size = config.n_embd if hasattr(config, "n_embd") else config.hidden_size |
| 1160 | |
| 1161 | # self.wte = nn.Embedding(config.vocab_size, hidden_size) |
| 1162 | |
| 1163 | audio_dim_out = ( |
| 1164 | None # Set this variable according to the actual audio processor |
| 1165 | ) |
| 1166 | self.layer_idx = -2 |
| 1167 | |
| 1168 | if ( |
| 1169 | isinstance(config.audio_processor, dict) |
| 1170 | and config.audio_processor.get("name", None) == "cascades" |
| 1171 | ): |
| 1172 | encoder_config = config.audio_processor.get("config", None) |
| 1173 | assert encoder_config is not None |
| 1174 | self.encoder = ConformerEncoder(**encoder_config) |
| 1175 | |
| 1176 | audio_dim_out = encoder_config["attention_dim"] |
| 1177 | n_mels = encoder_config["input_size"] |
| 1178 | else: |
| 1179 | raise NotImplementedError("") |
| 1180 | |
| 1181 | assert audio_dim_out is not None, "Remember to set values for audio_dim_out" |
| 1182 | self.audio_dim_out = audio_dim_out |
| 1183 | self.audio_dim_in = n_mels |
| 1184 | |
| 1185 | self.freeze_audio_processor = kwargs.get("freeze_audio_processor", False) |
| 1186 | |
| 1187 | self.downsample_rate = kwargs.get("downsample_rate", 1) |
| 1188 | |
| 1189 | if kwargs.get("use_qformer", False): |
| 1190 | qformer_config = kwargs.get("qformer_config", {}) |
| 1191 | qformer_config["attention_dim"] = audio_dim_out |
| 1192 | self.qformer: WindowQformer | None = WindowQformer(**qformer_config) |
| 1193 | else: |
| 1194 | self.qformer = None |
| 1195 | |
| 1196 | if kwargs.get("use_conv_downsample", False): |
| 1197 | assert self.qformer is None, ( |
| 1198 | "don't support use qformer and conv downsample together" |
| 1199 | ) |
| 1200 | nemo_conv_settings = kwargs.get("nemo_conv_settings", {}) |
| 1201 | default_nemo_conv_settings = { |
| 1202 | "subsampling": "dw_striding", |
| 1203 | "subsampling_factor": self.downsample_rate, |
| 1204 | "feat_in": audio_dim_out, |
| 1205 | "feat_out": audio_dim_out, |
| 1206 | "conv_channels": 256, |
| 1207 | "subsampling_conv_chunking_factor": 1, |
| 1208 | "activation": nn.ReLU(), |
| 1209 | "is_causal": False, |
| 1210 | } |
| 1211 | # Override any of the defaults with the incoming, user settings |
| 1212 | if nemo_conv_settings: |
| 1213 | default_nemo_conv_settings.update(nemo_conv_settings) |
| 1214 | for i in ["subsampling_factor", "feat_in", "feat_out"]: |
| 1215 | assert i not in nemo_conv_settings, ( |
| 1216 | "{i} should be specified outside of the NeMo dictionary" |
| 1217 | ) |
| 1218 | |
| 1219 | self.conv_ds: NemoConvSubsampling | None = NemoConvSubsampling( |
| 1220 | **default_nemo_conv_settings, |
| 1221 | ) |
| 1222 | else: |
| 1223 | self.conv_ds = None |
| 1224 | |
| 1225 | projection_cls = kwargs.get("projection_cls", "linear") |
| 1226 | if projection_cls == "linear": |
| 1227 | self.audio_projection = nn.Linear(audio_dim_out, hidden_size) |
| 1228 | elif projection_cls == "mlp": |
| 1229 | # follow llava-v1.5's implementation |
| 1230 | # (do not use image_projection and image_proj_norm) |
| 1231 | dim_projection = hidden_size |
| 1232 | depth = 2 |
| 1233 | self.linear_downsample_rate = ( |
| 1234 | 1 if (self.qformer or self.conv_ds) else self.downsample_rate |
| 1235 | ) |
| 1236 | layers = [ |
| 1237 | nn.Linear(audio_dim_out * self.linear_downsample_rate, dim_projection) |
| 1238 | ] |
| 1239 | for _ in range(1, depth): |
| 1240 | layers.extend([nn.GELU(), nn.Linear(dim_projection, dim_projection)]) |
| 1241 | self.audio_projection = nn.Sequential(*layers) |
| 1242 | # NOTE vision-speech tasks use a separate projection layer |
| 1243 | layers = [ |
| 1244 | nn.Linear(audio_dim_out * self.linear_downsample_rate, dim_projection) |
| 1245 | ] |
| 1246 | for _ in range(1, depth): |
| 1247 | layers.extend([nn.GELU(), nn.Linear(dim_projection, dim_projection)]) |
| 1248 | self.audio_projection_for_vision = nn.Sequential(*layers) |
| 1249 | else: |
| 1250 | raise NotImplementedError( |
| 1251 | f"projection_cls = {projection_cls}, not implemented" |
| 1252 | ) |
| 1253 | |
| 1254 | # TODO: audio sequence compression - Qformer |
| 1255 | self.vocab_size = config.vocab_size |
| 1256 | self.input_embeds = None |
| 1257 | self.audio_embed_sizes = None |
| 1258 | |
| 1259 | def set_audio_embeds(self, input_embeds: torch.Tensor) -> None: |
| 1260 | self.input_embeds = input_embeds |
| 1261 | |
| 1262 | def set_audio_embed_sizes(self, audio_embed_sizes: torch.Tensor) -> None: |
| 1263 | self.audio_embed_sizes = audio_embed_sizes |
| 1264 | |
| 1265 | def get_audio_features( |
| 1266 | self, |
| 1267 | input_embeds: torch.Tensor, |
| 1268 | audio_attention_mask: torch.Tensor | None = None, |
| 1269 | audio_projection_mode: str = "speech", |
| 1270 | ) -> torch.Tensor: |
| 1271 | """arguments: |
| 1272 | input_embeds: audio features (B, T, D) B: num audios in a sequence |
| 1273 | |
| 1274 | """ |
| 1275 | if self.freeze_audio_processor: |
| 1276 | with torch.no_grad(): |
| 1277 | audio_features, masks = self.encoder(input_embeds, audio_attention_mask) |
| 1278 | else: |
| 1279 | audio_features, masks = self.encoder(input_embeds, audio_attention_mask) |
| 1280 | |
| 1281 | if self.qformer is not None: |
| 1282 | audio_features, _ = self.qformer(audio_features, mask=None) |
| 1283 | |
| 1284 | if self.conv_ds is not None: |
| 1285 | if masks is not None: |
| 1286 | masks = masks.squeeze(1) |
| 1287 | |
| 1288 | audio_features, masks = self.conv_ds(audio_features, mask=masks) |
| 1289 | |
| 1290 | if self.linear_downsample_rate != 1: |
| 1291 | bs, seq_len, feat_dim = audio_features.size() |
| 1292 | padding = seq_len % self.linear_downsample_rate |
| 1293 | if padding > 0: |
| 1294 | audio_features = F.pad( |
| 1295 | audio_features, |
| 1296 | (0, 0, 0, self.linear_downsample_rate - padding), |
| 1297 | "constant", |
| 1298 | 0, |
| 1299 | ) |
| 1300 | |
| 1301 | seq_len = audio_features.size(1) |
| 1302 | audio_features = audio_features.view( |
| 1303 | bs, |
| 1304 | seq_len // self.linear_downsample_rate, |
| 1305 | feat_dim * self.linear_downsample_rate, |
| 1306 | ) |
| 1307 | |
| 1308 | if audio_projection_mode == "speech": |
| 1309 | audio_set_tensor = self.audio_projection(audio_features) |
| 1310 | elif audio_projection_mode == "vision": |
| 1311 | audio_set_tensor = self.audio_projection_for_vision(audio_features) |
| 1312 | else: |
| 1313 | raise ValueError( |
| 1314 | f"audio_projection_mode = {audio_projection_mode} not implemented" |
| 1315 | ) |
| 1316 | |
| 1317 | return audio_set_tensor |
| 1318 | |
| 1319 | def forward( |
| 1320 | self, |
| 1321 | audio_features: torch.Tensor, |
| 1322 | audio_attention_mask: torch.Tensor | None = None, |
| 1323 | audio_projection_mode: str = "speech", |
| 1324 | ) -> torch.Tensor: |
| 1325 | """arguments: |
| 1326 | audio_features: audio features (T, D) |
| 1327 | |
| 1328 | Returns: |
| 1329 | audio_embeds: audio embeddings (num_audio_tokens, hidden_dim) |
| 1330 | |
| 1331 | """ |
| 1332 | audio_embeds = self.get_audio_features( |
| 1333 | audio_features.unsqueeze(0), |
| 1334 | audio_attention_mask=audio_attention_mask, |
| 1335 | audio_projection_mode=audio_projection_mode, |
| 1336 | ) |
| 1337 | return audio_embeds.squeeze(0) |