| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import math |
| 5 | from collections.abc import Iterable, Mapping, Sequence |
| 6 | from typing import Any, ClassVar |
| 7 | |
| 8 | import torch |
| 9 | import torch.nn.functional as F |
| 10 | from torch import nn |
| 11 | from transformers import PreTrainedConfig |
| 12 | |
| 13 | from vllm.compilation.decorators import support_torch_compile |
| 14 | from vllm.config import CacheConfig, ModelConfig, SpeechToTextConfig, VllmConfig |
| 15 | from vllm.config.multimodal import MultiModalDummyOptions |
| 16 | from vllm.config.speech_to_text import SpeechToTextParams |
| 17 | from vllm.distributed import get_tensor_model_parallel_world_size |
| 18 | from vllm.inputs import MultiModalDataDict, PromptType, TokensPrompt |
| 19 | from vllm.logger import init_logger |
| 20 | from vllm.model_executor.layers.activation import get_act_fn |
| 21 | from vllm.model_executor.layers.attention import ( |
| 22 | Attention, |
| 23 | CrossAttention, |
| 24 | ) |
| 25 | from vllm.model_executor.layers.linear import ( |
| 26 | ColumnParallelLinear, |
| 27 | QKVParallelLinear, |
| 28 | RowParallelLinear, |
| 29 | ) |
| 30 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 31 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 32 | from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead |
| 33 | from vllm.model_executor.model_loader.weight_utils import default_weight_loader |
| 34 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 35 | from vllm.multimodal.inputs import ( |
| 36 | MultiModalFieldConfig, |
| 37 | MultiModalKwargsItems, |
| 38 | ) |
| 39 | from vllm.multimodal.parse import ( |
| 40 | AudioProcessorItems, |
| 41 | MultiModalDataItems, |
| 42 | MultiModalDataParser, |
| 43 | ) |
| 44 | from vllm.multimodal.processing import ( |
| 45 | BaseDummyInputsBuilder, |
| 46 | BaseProcessingInfo, |
| 47 | EncDecMultiModalProcessor, |
| 48 | ProcessorInputs, |
| 49 | PromptReplacement, |
| 50 | PromptUpdate, |
| 51 | TimingContext, |
| 52 | ) |
| 53 | from vllm.multimodal.processing.processor import ( |
| 54 | HFMultiModalInputs, |
| 55 | MultiModalProcessingResult, |
| 56 | ) |
| 57 | from vllm.renderers import TokenizeParams |
| 58 | from vllm.tokenizers import cached_tokenizer_from_config |
| 59 | from vllm.transformers_utils.processors.cohere_asr import ( |
| 60 | INF_VAL, |
| 61 | CohereASRFeatureExtractor, |
| 62 | CohereASRProcessor, |
| 63 | ) |
| 64 | from vllm.utils.collection_utils import is_list_of |
| 65 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 66 | from vllm.v1.attention.backend import ( |
| 67 | AttentionType, |
| 68 | ) |
| 69 | |
| 70 | from .interfaces import ( |
| 71 | MultiModalEmbeddings, |
| 72 | SupportsMultiModal, |
| 73 | SupportsTranscription, |
| 74 | ) |
| 75 | from .utils import AutoWeightsLoader, WeightsMapper, make_layers, maybe_prefix |
| 76 | |
| 77 | logger = init_logger(__name__) |
| 78 | |
| 79 | # From https://platform.openai.com/docs/guides/speech-to-text/supported-languages |
| 80 | |
| 81 | ISO639_1_SUPPORTED_LANGS = { |
| 82 | "en": "English", |
| 83 | "fr": "French", |
| 84 | "de": "German", |
| 85 | "es": "Spanish", |
| 86 | "pt": "Portuguese", |
| 87 | "it": "Italian", |
| 88 | "nl": "Dutch", |
| 89 | "pl": "Polish", |
| 90 | "el": "Greek", |
| 91 | "ar": "Arabic", |
| 92 | "ko": "Korean", |
| 93 | "ja": "Japanese", |
| 94 | "vi": "Vietnamese", |
| 95 | "zh": "Chinese", |
| 96 | } |
| 97 | |
| 98 | |
| 99 | class CohereASRAttention(nn.Module): |
| 100 | def __init__( |
| 101 | self, |
| 102 | embed_dim: int, |
| 103 | num_heads: int, |
| 104 | bias: bool = True, |
| 105 | attn_type: AttentionType = AttentionType.DECODER, |
| 106 | cache_config: CacheConfig | None = None, |
| 107 | quant_config: QuantizationConfig | None = None, |
| 108 | prefix: str = "", |
| 109 | ): |
| 110 | super().__init__() |
| 111 | self.embed_dim = embed_dim |
| 112 | tp_size = get_tensor_model_parallel_world_size() |
| 113 | self.total_num_heads = num_heads |
| 114 | assert self.total_num_heads % tp_size == 0 |
| 115 | self.num_heads = self.total_num_heads // tp_size |
| 116 | if self.total_num_heads >= tp_size: |
| 117 | # Number of heads is greater than TP size, so we partition |
| 118 | # the KV heads across multiple tensor parallel GPUs. |
| 119 | assert self.total_num_heads % tp_size == 0 |
| 120 | else: |
| 121 | # Number of heads is less than TP size, so we replicate |
| 122 | # the KV heads across multiple tensor parallel GPUs. |
| 123 | assert tp_size % self.total_num_heads == 0 |
| 124 | self.num_kv_heads = max(1, self.total_num_heads // tp_size) |
| 125 | self.head_dim = self.embed_dim // self.total_num_heads |
| 126 | self.q_size = self.num_heads * self.head_dim |
| 127 | self.kv_size = self.num_kv_heads * self.head_dim |
| 128 | self.attn_type = attn_type |
| 129 | |
| 130 | if (self.head_dim * num_heads) != self.embed_dim: |
| 131 | raise ValueError( |
| 132 | f"embed_dim must be divisible by num_heads (got `embed_dim`: " |
| 133 | f"{self.embed_dim} and `num_heads`: {num_heads})." |
| 134 | ) |
| 135 | self.scaling = self.head_dim**-0.5 |
| 136 | |
| 137 | self._init_qkv(embed_dim, bias, quant_config, prefix=prefix) |
| 138 | |
| 139 | self.out_projection = RowParallelLinear( |
| 140 | input_size=embed_dim, |
| 141 | output_size=embed_dim, |
| 142 | bias=bias, |
| 143 | quant_config=quant_config, |
| 144 | prefix=f"{prefix}.out_projection", |
| 145 | ) |
| 146 | if attn_type == AttentionType.ENCODER: |
| 147 | raise NotImplementedError( |
| 148 | "CohereASRAttention does not support Encoder Self-Attention yet." |
| 149 | ) |
| 150 | |
| 151 | elif self.attn_type == AttentionType.ENCODER_DECODER: |
| 152 | self.attn = CrossAttention( |
| 153 | self.num_heads, |
| 154 | self.head_dim, |
| 155 | self.scaling, |
| 156 | num_kv_heads=self.num_kv_heads, |
| 157 | cache_config=cache_config, |
| 158 | quant_config=quant_config, |
| 159 | prefix=f"{prefix}.attn", |
| 160 | attn_type=self.attn_type, |
| 161 | ) |
| 162 | else: # AttentionType.DECODER (regular decoder self-attention) |
| 163 | self.attn = Attention( |
| 164 | self.num_heads, |
| 165 | self.head_dim, |
| 166 | self.scaling, |
| 167 | num_kv_heads=self.num_kv_heads, |
| 168 | cache_config=cache_config, |
| 169 | quant_config=quant_config, |
| 170 | prefix=f"{prefix}.attn", |
| 171 | attn_type=self.attn_type, |
| 172 | ) |
| 173 | |
| 174 | def _init_qkv( |
| 175 | self, |
| 176 | embed_dim: int, |
| 177 | bias: bool = True, |
| 178 | quant_config: QuantizationConfig | None = None, |
| 179 | prefix: str = "", |
| 180 | ) -> None: |
| 181 | self.qkv_proj = QKVParallelLinear( |
| 182 | hidden_size=embed_dim, |
| 183 | head_size=self.head_dim, |
| 184 | total_num_heads=self.total_num_heads, |
| 185 | total_num_kv_heads=self.total_num_heads, |
| 186 | bias=bias, |
| 187 | quant_config=quant_config, |
| 188 | prefix=f"{prefix}.qkv_proj", |
| 189 | ) |
| 190 | |
| 191 | def forward( |
| 192 | self, |
| 193 | hidden_states: torch.Tensor, |
| 194 | encoder_hidden_states: torch.Tensor | None = None, |
| 195 | ) -> torch.Tensor: |
| 196 | qkv, _ = self.qkv_proj(hidden_states) |
| 197 | q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) |
| 198 | |
| 199 | attn_output = self.attn(q, k, v) |
| 200 | |
| 201 | output, _ = self.out_projection(attn_output) |
| 202 | |
| 203 | return output |
| 204 | |
| 205 | |
| 206 | class CohereASRCrossAttention(CohereASRAttention): |
| 207 | def __init__( |
| 208 | self, |
| 209 | embed_dim: int, |
| 210 | num_heads: int, |
| 211 | bias: bool = True, |
| 212 | cache_config: CacheConfig | None = None, |
| 213 | quant_config: QuantizationConfig | None = None, |
| 214 | prefix: str = "", |
| 215 | ): |
| 216 | super().__init__( |
| 217 | embed_dim=embed_dim, |
| 218 | num_heads=num_heads, |
| 219 | bias=bias, |
| 220 | cache_config=cache_config, |
| 221 | quant_config=quant_config, |
| 222 | prefix=prefix, |
| 223 | attn_type=AttentionType.ENCODER_DECODER, |
| 224 | ) |
| 225 | |
| 226 | def _init_qkv( |
| 227 | self, |
| 228 | embed_dim: int, |
| 229 | bias: bool = True, |
| 230 | quant_config: QuantizationConfig | None = None, |
| 231 | prefix: str = "", |
| 232 | ) -> None: |
| 233 | self.q_proj = ColumnParallelLinear( |
| 234 | input_size=embed_dim, |
| 235 | output_size=embed_dim, |
| 236 | bias=bias, |
| 237 | quant_config=quant_config, |
| 238 | prefix=f"{prefix}.q_proj", |
| 239 | ) |
| 240 | self.kv_proj = QKVParallelLinear( |
| 241 | hidden_size=embed_dim, |
| 242 | head_size=self.head_dim, |
| 243 | total_num_heads=0, |
| 244 | total_num_kv_heads=self.total_num_heads, |
| 245 | bias=bias, |
| 246 | quant_config=quant_config, |
| 247 | prefix=f"{prefix}.kv_proj", |
| 248 | ) |
| 249 | |
| 250 | def forward( |
| 251 | self, |
| 252 | hidden_states: torch.Tensor, |
| 253 | encoder_hidden_states: torch.Tensor | None = None, |
| 254 | ) -> torch.Tensor: |
| 255 | q, _ = self.q_proj(hidden_states) |
| 256 | |
| 257 | # Encoder hidden states are only computed once during prefill phase. |
| 258 | # Afterwards, the keys and values should be available in the kv-cache. |
| 259 | if encoder_hidden_states is not None: |
| 260 | kv, _ = self.kv_proj(encoder_hidden_states) |
| 261 | k, v = kv.split([self.kv_size, self.kv_size], dim=-1) |
| 262 | else: |
| 263 | k = v = None |
| 264 | |
| 265 | attn_output = self.attn(q, k, v) |
| 266 | |
| 267 | output, _ = self.out_projection(attn_output) |
| 268 | |
| 269 | return output |
| 270 | |
| 271 | |
| 272 | # ----- Decoder START ----- |
| 273 | class CohereASRMLP(nn.Module): |
| 274 | def __init__( |
| 275 | self, |
| 276 | embed_dim: int, |
| 277 | ffn_dim: int, |
| 278 | act_fn: str, |
| 279 | quant_config: QuantizationConfig | None = None, |
| 280 | prefix: str = "", |
| 281 | ): |
| 282 | super().__init__() |
| 283 | |
| 284 | self.activation_fn = get_act_fn(act_fn) |
| 285 | self.dense_in = ColumnParallelLinear( |
| 286 | input_size=embed_dim, |
| 287 | output_size=ffn_dim, |
| 288 | quant_config=quant_config, |
| 289 | prefix=f"{prefix}.fc1", |
| 290 | ) |
| 291 | self.dense_out = RowParallelLinear( |
| 292 | input_size=ffn_dim, |
| 293 | output_size=embed_dim, |
| 294 | quant_config=quant_config, |
| 295 | prefix=f"{prefix}.fc2", |
| 296 | ) |
| 297 | |
| 298 | def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 299 | hidden_states, _ = self.dense_in(hidden_states) |
| 300 | hidden_states = self.activation_fn(hidden_states) |
| 301 | hidden_states, _ = self.dense_out(hidden_states) |
| 302 | return hidden_states |
| 303 | |
| 304 | |
| 305 | class FixedPositionalEncoding(nn.Module): |
| 306 | """Fixed positional encoding (embedding layer) from sine and cosine functions |
| 307 | of different frequencies according to https://arxiv.org/abs/1706.03762 |
| 308 | |
| 309 | Args: |
| 310 | hidden_size: size of the embeddings in the model, also known as d_model |
| 311 | max_sequence_length: maximum allowed length of the input sequence |
| 312 | |
| 313 | """ |
| 314 | |
| 315 | def __init__(self, hidden_size: int, max_sequence_length: int = 512) -> None: |
| 316 | super().__init__() |
| 317 | |
| 318 | self._hidden_size = hidden_size |
| 319 | self._max_sequence_length = max_sequence_length |
| 320 | self._build_pos_enc( |
| 321 | hidden_size=self._hidden_size, max_sequence_length=self._max_sequence_length |
| 322 | ) |
| 323 | |
| 324 | def _build_pos_enc(self, hidden_size: int, max_sequence_length: int) -> None: |
| 325 | """Builds/replaces pre-computed positional encoding.""" |
| 326 | pos_enc = torch.zeros(max_sequence_length, hidden_size) |
| 327 | position = torch.arange(0.0, max_sequence_length).unsqueeze(1) |
| 328 | coef = -math.log(10000.0) / hidden_size |
| 329 | div_term = torch.exp(coef * torch.arange(0.0, hidden_size, 2)) |
| 330 | pos_enc[:, 0::2] = torch.sin(position * div_term) |
| 331 | pos_enc[:, 1::2] = torch.cos(position * div_term) |
| 332 | pos_enc.div_(math.sqrt(hidden_size)) |
| 333 | self.register_buffer("pos_enc", pos_enc) |
| 334 | |
| 335 | def forward(self, position_ids: torch.Tensor) -> torch.Tensor: |
| 336 | embeddings = torch.embedding(self.pos_enc, position_ids) |
| 337 | return embeddings |
| 338 | |
| 339 | |
| 340 | class CohereASRDecoderLayer(nn.Module): |
| 341 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 342 | super().__init__() |
| 343 | config = vllm_config.model_config.hf_config.transf_decoder["config_dict"] |
| 344 | cache_config = vllm_config.cache_config |
| 345 | quant_config = vllm_config.quant_config |
| 346 | |
| 347 | self.hidden_dim = config.get("hidden_size") |
| 348 | self.ffn_dim = config.get("inner_size") |
| 349 | self.act_fn = config.get("hidden_act") |
| 350 | self.num_heads = config.get("num_attention_heads") |
| 351 | |
| 352 | # self_attn |
| 353 | self.layer_norm_1 = nn.LayerNorm(self.hidden_dim) |
| 354 | self.first_sub_layer = CohereASRAttention( |
| 355 | embed_dim=self.hidden_dim, |
| 356 | num_heads=self.num_heads, |
| 357 | attn_type=AttentionType.DECODER, |
| 358 | cache_config=cache_config, |
| 359 | quant_config=quant_config, |
| 360 | prefix=f"{prefix}.first_sub_layer", |
| 361 | ) |
| 362 | |
| 363 | # cross attn to attend to encoder |
| 364 | self.layer_norm_2 = nn.LayerNorm(self.hidden_dim) |
| 365 | self.second_sub_layer = CohereASRCrossAttention( |
| 366 | embed_dim=self.hidden_dim, |
| 367 | num_heads=self.num_heads, |
| 368 | cache_config=cache_config, |
| 369 | quant_config=quant_config, |
| 370 | prefix=f"{prefix}.second_sub_layer", |
| 371 | ) |
| 372 | |
| 373 | self.layer_norm_3 = nn.LayerNorm(self.hidden_dim) |
| 374 | self.third_sub_layer = CohereASRMLP( |
| 375 | embed_dim=self.hidden_dim, |
| 376 | ffn_dim=self.ffn_dim, |
| 377 | act_fn=self.act_fn, |
| 378 | quant_config=quant_config, |
| 379 | prefix=f"{prefix}.third_sub_layer", |
| 380 | ) |
| 381 | |
| 382 | def forward( |
| 383 | self, |
| 384 | hidden_states: torch.Tensor, |
| 385 | encoder_hidden_states: torch.Tensor | None, |
| 386 | ) -> torch.Tensor: |
| 387 | residual = hidden_states |
| 388 | hidden_states = self.layer_norm_1(hidden_states) |
| 389 | hidden_states = self.first_sub_layer(hidden_states=hidden_states) |
| 390 | |
| 391 | hidden_states = residual + hidden_states |
| 392 | residual = hidden_states |
| 393 | hidden_states = self.layer_norm_2(hidden_states) |
| 394 | hidden_states = self.second_sub_layer( |
| 395 | hidden_states=hidden_states, |
| 396 | encoder_hidden_states=encoder_hidden_states, |
| 397 | ) |
| 398 | |
| 399 | hidden_states = residual + hidden_states |
| 400 | residual = hidden_states |
| 401 | hidden_states = self.layer_norm_3(hidden_states) |
| 402 | hidden_states = self.third_sub_layer(hidden_states) |
| 403 | hidden_states = residual + hidden_states |
| 404 | |
| 405 | return hidden_states |
| 406 | |
| 407 | |
| 408 | class TransformerEmbedding(nn.Module): |
| 409 | def __init__( |
| 410 | self, |
| 411 | vocab_size: int, |
| 412 | hidden_size: int, |
| 413 | max_target_positions: int, |
| 414 | padding_idx: int, |
| 415 | ) -> None: |
| 416 | super().__init__() |
| 417 | self.token_embedding = nn.Embedding(vocab_size, hidden_size, padding_idx) |
| 418 | self.position_embedding = FixedPositionalEncoding( |
| 419 | hidden_size=hidden_size, |
| 420 | max_sequence_length=max_target_positions, |
| 421 | ) |
| 422 | self.layer_norm = nn.LayerNorm(hidden_size) |
| 423 | |
| 424 | def forward(self, input_ids: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: |
| 425 | inputs_embeds = self.token_embedding(input_ids) |
| 426 | positions = self.position_embedding(positions) |
| 427 | embeddings = inputs_embeds + positions |
| 428 | embeddings = self.layer_norm(embeddings) |
| 429 | return embeddings |
| 430 | |
| 431 | |
| 432 | @support_torch_compile(dynamic_arg_dims={"input_ids": 0, "positions": -1}) |
| 433 | class CohereASRDecoder(nn.Module): |
| 434 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 435 | super().__init__() |
| 436 | config = vllm_config.model_config.hf_config |
| 437 | self.padding_idx = 2 |
| 438 | config_dict = config.transf_decoder["config_dict"] |
| 439 | self.max_target_positions = config_dict.get("max_sequence_length") |
| 440 | self.hidden_size = config_dict.get("hidden_size") |
| 441 | self.num_decoder_layers = config_dict.get("num_layers") |
| 442 | self.vocab_size = config.head["num_classes"] |
| 443 | |
| 444 | self.embedding = TransformerEmbedding( |
| 445 | vocab_size=self.vocab_size, |
| 446 | hidden_size=self.hidden_size, |
| 447 | max_target_positions=self.max_target_positions, |
| 448 | padding_idx=self.padding_idx, |
| 449 | ) |
| 450 | |
| 451 | self.start_layer, self.end_layer, self.layers = make_layers( |
| 452 | self.num_decoder_layers, |
| 453 | lambda prefix: CohereASRDecoderLayer( |
| 454 | vllm_config=vllm_config, prefix=f"{prefix}.layers" |
| 455 | ), |
| 456 | prefix=f"{prefix}.layers", |
| 457 | ) |
| 458 | self.final_layer_norm = nn.LayerNorm(self.hidden_size) |
| 459 | |
| 460 | def forward( |
| 461 | self, |
| 462 | input_ids: torch.Tensor, |
| 463 | positions: torch.Tensor, |
| 464 | encoder_hidden_states: torch.Tensor | None, |
| 465 | ) -> torch.Tensor: |
| 466 | hidden_states = self.get_input_embeddings(input_ids, positions) |
| 467 | for decoder_layer in self.layers: |
| 468 | hidden_states = decoder_layer( |
| 469 | hidden_states, |
| 470 | encoder_hidden_states=encoder_hidden_states, |
| 471 | ) |
| 472 | |
| 473 | hidden_states = self.final_layer_norm(hidden_states) |
| 474 | return hidden_states |
| 475 | |
| 476 | def get_input_embeddings( |
| 477 | self, input_ids: torch.Tensor, positions: torch.Tensor |
| 478 | ) -> torch.Tensor: |
| 479 | return self.embedding(input_ids, positions) |
| 480 | |
| 481 | |
| 482 | # ----- Decoder END ----- |
| 483 | |
| 484 | |
| 485 | # ----- Encoder START ----- |
| 486 | class MaskedConvSequential(nn.Sequential): |
| 487 | def forward( |
| 488 | self, x: torch.Tensor, lengths: torch.Tensor |
| 489 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 490 | x = x.unsqueeze(1) # (batch, 1, time, features) |
| 491 | current_lengths = lengths.clone().float() |
| 492 | mask = self._create_mask(x, current_lengths.long()) |
| 493 | |
| 494 | # Process through each layer with mask propagation |
| 495 | for i, layer in enumerate(self): |
| 496 | # Apply current mask before layer |
| 497 | x = self.apply_channel_mask(x, mask) |
| 498 | |
| 499 | # Apply layer |
| 500 | x = layer(x) |
| 501 | |
| 502 | # Update lengths for stride operations with proper padding |
| 503 | if hasattr(layer, "stride") and layer.stride != (1, 1): |
| 504 | if hasattr(layer, "_left_padding"): |
| 505 | padding = ( |
| 506 | layer._left_padding, |
| 507 | layer._right_padding, |
| 508 | ) # CausalConv2D |
| 509 | else: |
| 510 | padding = layer.padding |
| 511 | current_lengths = self.calculate_conv_output_size( |
| 512 | current_lengths, layer.kernel_size[0], layer.stride[0], padding |
| 513 | ) |
| 514 | mask = self._create_mask(x, current_lengths.long()) |
| 515 | |
| 516 | # Final masking |
| 517 | x = self.apply_channel_mask(x, mask) |
| 518 | return x, current_lengths.long() |
| 519 | |
| 520 | def _create_mask(self, tensor: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor: |
| 521 | """Create broadcastable mask from per-sample lengths. |
| 522 | |
| 523 | Returns a (B, 1, T, 1) mask that broadcasts over channels and |
| 524 | features without materializing a full (B, C, T, F) tensor. |
| 525 | """ |
| 526 | batch_size, channels, time, features = tensor.shape |
| 527 | time_mask = torch.arange(time, device=tensor.device).expand( |
| 528 | batch_size, time |
| 529 | ) < lengths.unsqueeze(1) |
| 530 | return time_mask.to(tensor.dtype).unsqueeze(1).unsqueeze(-1) |
| 531 | |
| 532 | def apply_channel_mask( |
| 533 | self, tensor: torch.Tensor, mask: torch.Tensor |
| 534 | ) -> torch.Tensor: |
| 535 | """Apply mask in-place via broadcasting. |
| 536 | |
| 537 | tensor: (B, C, T, F), mask: (B, 1, T, 1) |
| 538 | """ |
| 539 | tensor.mul_(mask) |
| 540 | return tensor |
| 541 | |
| 542 | def calculate_conv_output_size( |
| 543 | self, |
| 544 | input_size: torch.Tensor, |
| 545 | kernel_size: int, |
| 546 | stride: int, |
| 547 | padding: tuple[int, int], |
| 548 | ): |
| 549 | """Calculate exact output size after convolution.""" |
| 550 | return (input_size + padding[0] + padding[1] - kernel_size) // stride + 1 |
| 551 | |
| 552 | |
| 553 | class ConvSubsampling(nn.Module): |
| 554 | def __init__( |
| 555 | self, |
| 556 | subsampling: str, |
| 557 | subsampling_factor: int, |
| 558 | feat_in: int, |
| 559 | feat_out: int, |
| 560 | conv_channels: int, |
| 561 | subsampling_conv_chunking_factor: int = 1, |
| 562 | activation: nn.Module | None = None, |
| 563 | is_causal: bool = False, |
| 564 | ) -> None: |
| 565 | super().__init__() |
| 566 | if activation is None: |
| 567 | activation = nn.ReLU() |
| 568 | |
| 569 | if subsampling_factor % 2 != 0: |
| 570 | raise ValueError("Sampling factor should be a multiply of 2!") |
| 571 | self._sampling_num = int(math.log(subsampling_factor, 2)) |
| 572 | |
| 573 | if ( |
| 574 | subsampling_conv_chunking_factor != -1 |
| 575 | and subsampling_conv_chunking_factor != 1 |
| 576 | and subsampling_conv_chunking_factor % 2 != 0 |
| 577 | ): |
| 578 | raise ValueError( |
| 579 | "subsampling_conv_chunking_factor should be -1, 1, or a power of 2" |
| 580 | ) |
| 581 | |
| 582 | in_channels = 1 |
| 583 | layers = [] |
| 584 | |
| 585 | assert subsampling == "dw_striding" |
| 586 | self._stride = 2 |
| 587 | self._kernel_size = 3 |
| 588 | self._ceil_mode = False |
| 589 | |
| 590 | assert not is_causal |
| 591 | |
| 592 | self._left_padding = (self._kernel_size - 1) // 2 |
| 593 | self._right_padding = (self._kernel_size - 1) // 2 |
| 594 | |
| 595 | # Layer 1 |
| 596 | # [1, T, num_melspec] -> [conv_channels, T//2, num_melspec//2] |
| 597 | layers.append( |
| 598 | torch.nn.Conv2d( |
| 599 | in_channels=in_channels, |
| 600 | out_channels=conv_channels, |
| 601 | kernel_size=self._kernel_size, |
| 602 | stride=self._stride, |
| 603 | padding=self._left_padding, |
| 604 | ) |
| 605 | ) |
| 606 | in_channels = conv_channels |
| 607 | layers.append(activation) |
| 608 | |
| 609 | for i in range(self._sampling_num - 1): |
| 610 | # [conv_channels, T//2^i, num_melspec//2^i] -> |
| 611 | # [conv_channels, T//2^(i+1), num_melspec//2^(i+1)] |
| 612 | # depthwise conv |
| 613 | layers.append( |
| 614 | torch.nn.Conv2d( |
| 615 | in_channels=in_channels, |
| 616 | out_channels=in_channels, |
| 617 | kernel_size=self._kernel_size, |
| 618 | stride=self._stride, |
| 619 | padding=self._left_padding, |
| 620 | groups=in_channels, |
| 621 | ) |
| 622 | ) |
| 623 | |
| 624 | # [conv_channels, T//2^(i+1), num_melspec//2^(i+1)] |
| 625 | # -> [conv_channels, T//2^(i+1), num_melspec//2^(i+1)] |
| 626 | # pointwise conv |
| 627 | layers.append( |
| 628 | torch.nn.Conv2d( |
| 629 | in_channels=in_channels, |
| 630 | out_channels=conv_channels, |
| 631 | kernel_size=1, |
| 632 | stride=1, |
| 633 | padding=0, |
| 634 | groups=1, |
| 635 | ) |
| 636 | ) |
| 637 | layers.append(activation) |
| 638 | in_channels = conv_channels |
| 639 | |
| 640 | in_length = torch.tensor(feat_in, dtype=torch.float) |
| 641 | out_length = self.calc_length( |
| 642 | lengths=in_length, |
| 643 | all_paddings=self._left_padding + self._right_padding, |
| 644 | kernel_size=self._kernel_size, |
| 645 | stride=self._stride, |
| 646 | ceil_mode=self._ceil_mode, |
| 647 | repeat_num=self._sampling_num, |
| 648 | ) |
| 649 | |
| 650 | # reshape: |
| 651 | # [conv_channels, T//sub_factor, num_melspec//sub_factor] |
| 652 | # -> [T//sub_factor, conv_channels * (num_melspec//sub_factor)] |
| 653 | # mlp: |
| 654 | # [T//sub_factor, conv_channels * (num_melspec//sub_factor)] |
| 655 | # -> [T//sub_factor, feat_out] |
| 656 | self.out = torch.nn.Linear(conv_channels * int(out_length), feat_out) |
| 657 | self.conv2d_subsampling = True |
| 658 | self.conv = MaskedConvSequential(*layers) |
| 659 | |
| 660 | def calc_length( |
| 661 | self, |
| 662 | lengths: torch.Tensor, |
| 663 | all_paddings: int, |
| 664 | kernel_size: int, |
| 665 | stride: int, |
| 666 | ceil_mode: bool, |
| 667 | repeat_num: int = 1, |
| 668 | ) -> torch.Tensor: |
| 669 | """Calculates the output length of a Tensor passed |
| 670 | through a convolution or max pooling layer""" |
| 671 | add_pad: float = all_paddings - kernel_size |
| 672 | one: float = 1.0 |
| 673 | for i in range(repeat_num): |
| 674 | lengths = torch.div(lengths.to(dtype=torch.float) + add_pad, stride) + one |
| 675 | lengths = torch.ceil(lengths) if ceil_mode else torch.floor(lengths) |
| 676 | return lengths.to(dtype=torch.int) |
| 677 | |
| 678 | def forward( |
| 679 | self, x: torch.Tensor, lengths: torch.Tensor |
| 680 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 681 | x, lengths = self.conv(x, lengths) |
| 682 | |
| 683 | if self.conv2d_subsampling: |
| 684 | b, c, t, f = x.size() |
| 685 | x = self.out(x.transpose(1, 2).reshape(b, t, -1)) |
| 686 | # Transpose to Channel Last mode |
| 687 | else: |
| 688 | x = x.transpose(1, 2) |
| 689 | |
| 690 | return x, lengths |
| 691 | |
| 692 | |
| 693 | class PositionalEncoding(torch.nn.Module): |
| 694 | """Fixed sinusoidal positional encoding. |
| 695 | |
| 696 | Args: |
| 697 | d_model (int): embedding dim |
| 698 | max_len (int): maximum input length |
| 699 | xscale (bool): whether to scale the input by sqrt(d_model) |
| 700 | |
| 701 | """ |
| 702 | |
| 703 | def __init__( |
| 704 | self, d_model: int, max_len: int = 5000, xscale: float | None = None |
| 705 | ) -> None: |
| 706 | super().__init__() |
| 707 | self.d_model = d_model |
| 708 | self.xscale = xscale |
| 709 | self.max_len = max_len |
| 710 | |
| 711 | def create_pe(self, positions: torch.Tensor, dtype: torch.dtype) -> None: |
| 712 | pos_length = positions.size(0) |
| 713 | pe = torch.zeros(pos_length, self.d_model, device=positions.device) |
| 714 | div_term = torch.exp( |
| 715 | torch.arange( |
| 716 | 0, self.d_model, 2, dtype=torch.float32, device=positions.device |
| 717 | ) |
| 718 | * -(math.log(10000.0) / self.d_model) |
| 719 | ) |
| 720 | pe[:, 0::2] = torch.sin(positions * div_term) |
| 721 | pe[:, 1::2] = torch.cos(positions * div_term) |
| 722 | pe = pe.unsqueeze(0).to(dtype) |
| 723 | if hasattr(self, "pe"): |
| 724 | self.pe = pe |
| 725 | else: |
| 726 | self.register_buffer("pe", pe, persistent=False) |
| 727 | |
| 728 | def forward( |
| 729 | self, x: torch.Tensor, cache_len: int = 0 |
| 730 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 731 | """Adds positional encoding. |
| 732 | |
| 733 | Args: |
| 734 | x (torch.Tensor): Input. Its shape is (batch, time, feature_size) |
| 735 | cache_len (int): the size of the cache which is used to shift positions |
| 736 | Returns: |
| 737 | x+pos_emb (torch.Tensor): Its shape is (batch, time, feature_size) |
| 738 | pos_emb (torch.Tensor): Its shape is (1, time, feature_size) |
| 739 | |
| 740 | """ |
| 741 | input_len = x.size(1) + cache_len |
| 742 | if self.xscale: |
| 743 | x = x * self.xscale |
| 744 | pos_emb = self.pe[:, :input_len] |
| 745 | x = x + pos_emb |
| 746 | return x, pos_emb |
| 747 | |
| 748 | |
| 749 | class RelPositionalEncoding(PositionalEncoding): |
| 750 | """Relative positional encoding for TransformerXL's layers |
| 751 | See : Appendix B in https://arxiv.org/abs/1901.02860 |
| 752 | Args: |
| 753 | d_model (int): embedding dim |
| 754 | max_len (int): maximum input length |
| 755 | xscale (bool): whether to scale the input by sqrt(d_model) |
| 756 | """ |
| 757 | |
| 758 | def extend_pe(self, length: int, device: torch.device, dtype: torch.dtype) -> None: |
| 759 | """Reset and extend the positional encodings if needed.""" |
| 760 | needed_size = 2 * length - 1 |
| 761 | if hasattr(self, "pe") and self.pe.size(1) >= needed_size: |
| 762 | return |
| 763 | positions = torch.arange( |
| 764 | length - 1, -length, -1, dtype=torch.float32, device=device |
| 765 | ).unsqueeze(1) |
| 766 | self.create_pe(positions=positions, dtype=dtype) |
| 767 | |
| 768 | def forward( |
| 769 | self, x: torch.Tensor, cache_len: int = 0 |
| 770 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 771 | """Compute positional encoding. |
| 772 | |
| 773 | Args: |
| 774 | x (torch.Tensor): Input. Its shape is (batch, time, feature_size) |
| 775 | cache_len (int): the size of the cache which is used to shift positions |
| 776 | Returns: |
| 777 | x (torch.Tensor): Its shape is (batch, time, feature_size) |
| 778 | pos_emb (torch.Tensor): Its shape is (1, time, feature_size) |
| 779 | |
| 780 | """ |
| 781 | if self.xscale: |
| 782 | x = x * self.xscale |
| 783 | |
| 784 | input_len = x.size(1) + cache_len |
| 785 | center_pos = self.pe.size(1) // 2 + 1 |
| 786 | start_pos = center_pos - input_len |
| 787 | end_pos = center_pos + input_len - 1 |
| 788 | pos_emb = self.pe[:, start_pos:end_pos] |
| 789 | |
| 790 | return x, pos_emb |
| 791 | |
| 792 | |
| 793 | class Swish(nn.SiLU): |
| 794 | """Swish activation function introduced in 'https://arxiv.org/abs/1710.05941' |
| 795 | Mathematically identical to SiLU. See note in nn.SiLU for references. |
| 796 | """ |
| 797 | |
| 798 | |
| 799 | class ConformerFeedForward(nn.Module): |
| 800 | """feed-forward module of Conformer model. |
| 801 | use_bias (bool): Apply bias to all Linear and Conv1d |
| 802 | layers to improve activation flow and stabilize |
| 803 | training of huge models. |
| 804 | """ |
| 805 | |
| 806 | def __init__( |
| 807 | self, |
| 808 | d_model: int, |
| 809 | d_ff: int, |
| 810 | activation: nn.Module | None = None, |
| 811 | use_bias: bool = True, |
| 812 | ) -> None: |
| 813 | super().__init__() |
| 814 | if activation is None: |
| 815 | activation = Swish() |
| 816 | self.linear1 = nn.Linear(d_model, d_ff, bias=use_bias) |
| 817 | self.activation = activation |
| 818 | self.linear2 = nn.Linear(d_ff, d_model, bias=use_bias) |
| 819 | |
| 820 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 821 | x = self.linear1(x) |
| 822 | x = self.activation(x) |
| 823 | x = self.linear2(x) |
| 824 | return x |
| 825 | |
| 826 | |
| 827 | class CausalConv1D(nn.Conv1d): |
| 828 | """A causal version of nn.Conv1d where each step would |
| 829 | have limited access to locations on its right or left. |
| 830 | All arguments are the same as nn.Conv1d except padding. |
| 831 | |
| 832 | If padding is set None, then paddings are set |
| 833 | automatically to make it a causal convolution where |
| 834 | each location would not see any steps on its right. |
| 835 | |
| 836 | If padding is set as a list (size of 2), then |
| 837 | padding[0] would be used as left padding and |
| 838 | padding[1] as right padding. It would make it possible |
| 839 | to control the number of steps to be accessible on the |
| 840 | right and left. This mode is not supported when |
| 841 | stride > 1. padding[0]+padding[1] should be equal to |
| 842 | (kernel_size - 1). |
| 843 | """ |
| 844 | |
| 845 | def __init__( |
| 846 | self, |
| 847 | in_channels: int, |
| 848 | out_channels: int, |
| 849 | kernel_size: int, |
| 850 | stride: int = 1, |
| 851 | padding: str | int | list[int] | None = 0, |
| 852 | dilation: int = 1, |
| 853 | groups: int = 1, |
| 854 | bias: bool = True, |
| 855 | padding_mode: str = "zeros", |
| 856 | device=None, |
| 857 | dtype=None, |
| 858 | ) -> None: |
| 859 | self._left_padding: int |
| 860 | self._right_padding: int |
| 861 | if padding is None: |
| 862 | self._left_padding = kernel_size - 1 |
| 863 | self._right_padding = stride - 1 |
| 864 | else: |
| 865 | if stride != 1 and padding != kernel_size - 1: |
| 866 | raise ValueError("No striding allowed for non-symmetric convolutions!") |
| 867 | if isinstance(padding, int): |
| 868 | self._left_padding = padding |
| 869 | self._right_padding = padding |
| 870 | elif ( |
| 871 | isinstance(padding, list) |
| 872 | and len(padding) == 2 |
| 873 | and padding[0] + padding[1] == kernel_size - 1 |
| 874 | ): |
| 875 | self._left_padding = padding[0] |
| 876 | self._right_padding = padding[1] |
| 877 | else: |
| 878 | raise ValueError(f"Invalid padding param: {padding}!") |
| 879 | |
| 880 | super().__init__( |
| 881 | in_channels=in_channels, |
| 882 | out_channels=out_channels, |
| 883 | kernel_size=kernel_size, |
| 884 | stride=stride, |
| 885 | padding=0, |
| 886 | dilation=dilation, |
| 887 | groups=groups, |
| 888 | bias=bias, |
| 889 | padding_mode=padding_mode, |
| 890 | device=device, |
| 891 | dtype=dtype, |
| 892 | ) |
| 893 | |
| 894 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 895 | x = F.pad(x, pad=(self._left_padding, self._right_padding)) |
| 896 | return super().forward(x) |
| 897 | |
| 898 | |
| 899 | class ConformerConvolution(nn.Module): |
| 900 | """The convolution module for the Conformer model. |
| 901 | |
| 902 | Args: |
| 903 | d_model (int): hidden dimension |
| 904 | kernel_size (int): kernel size for depthwise convolution |
| 905 | pointwise_activation (str): name of the activation |
| 906 | function to be used for the pointwise conv. |
| 907 | Note that Conformer uses a special key `glu_` |
| 908 | which is treated as the original default from |
| 909 | the paper. |
| 910 | use_bias (bool): Use bias in all Linear and Conv1d |
| 911 | layers to improve activation flow and stabilize |
| 912 | training of huge models. Defaults to True |
| 913 | |
| 914 | """ |
| 915 | |
| 916 | def __init__( |
| 917 | self, |
| 918 | d_model: int, |
| 919 | kernel_size: int, |
| 920 | norm_type: str = "batch_norm", |
| 921 | conv_context_size: int | list[int] | None = None, |
| 922 | pointwise_activation: str = "glu_", |
| 923 | use_bias: bool = True, |
| 924 | ) -> None: |
| 925 | super().__init__() |
| 926 | assert (kernel_size - 1) % 2 == 0 |
| 927 | |
| 928 | if conv_context_size is None: |
| 929 | conv_context_size = (kernel_size - 1) // 2 |
| 930 | |
| 931 | assert pointwise_activation == "glu_" |
| 932 | dw_conv_input_dim = d_model |
| 933 | |
| 934 | self.pointwise_conv1 = nn.Conv1d( |
| 935 | in_channels=d_model, |
| 936 | out_channels=d_model * 2, |
| 937 | kernel_size=1, |
| 938 | stride=1, |
| 939 | padding=0, |
| 940 | bias=use_bias, |
| 941 | ) |
| 942 | |
| 943 | self.depthwise_conv = CausalConv1D( |
| 944 | in_channels=dw_conv_input_dim, |
| 945 | out_channels=dw_conv_input_dim, |
| 946 | kernel_size=kernel_size, |
| 947 | stride=1, |
| 948 | padding=conv_context_size, |
| 949 | groups=dw_conv_input_dim, |
| 950 | bias=use_bias, |
| 951 | ) |
| 952 | |
| 953 | assert norm_type == "batch_norm" |
| 954 | self.batch_norm = nn.BatchNorm1d(dw_conv_input_dim) |
| 955 | |
| 956 | self.activation = Swish() |
| 957 | self.pointwise_conv2 = nn.Conv1d( |
| 958 | in_channels=dw_conv_input_dim, |
| 959 | out_channels=d_model, |
| 960 | kernel_size=1, |
| 961 | stride=1, |
| 962 | padding=0, |
| 963 | bias=use_bias, |
| 964 | ) |
| 965 | |
| 966 | def forward( |
| 967 | self, x: torch.Tensor, pad_mask: torch.Tensor | None = None |
| 968 | ) -> torch.Tensor: |
| 969 | x = x.transpose(1, 2) |
| 970 | x = self.pointwise_conv1(x) |
| 971 | |
| 972 | x = nn.functional.glu(x, dim=1) |
| 973 | |
| 974 | if pad_mask is not None: |
| 975 | x = x.masked_fill(pad_mask.unsqueeze(1), 0.0) |
| 976 | |
| 977 | x = self.depthwise_conv(x) |
| 978 | |
| 979 | x = self.batch_norm(x) |
| 980 | |
| 981 | x = self.activation(x) |
| 982 | x = self.pointwise_conv2(x) |
| 983 | x = x.transpose(1, 2) |
| 984 | return x |
| 985 | |
| 986 | |
| 987 | class CohereASRMultiHeadAttention(nn.Module): |
| 988 | """Multi-Head Attention layer of Transformer. |
| 989 | |
| 990 | Args: |
| 991 | n_head (int): number of heads |
| 992 | n_feat (int): size of the features |
| 993 | use_bias (bool): whether to remove bias in linear and conv layers |
| 994 | |
| 995 | """ |
| 996 | |
| 997 | def __init__( |
| 998 | self, |
| 999 | n_head: int, |
| 1000 | n_feat: int, |
| 1001 | use_bias: bool = True, |
| 1002 | ) -> None: |
| 1003 | """Construct an MultiHeadedAttention object.""" |
| 1004 | super().__init__() |
| 1005 | |
| 1006 | assert n_feat % n_head == 0 |
| 1007 | self.d_k = n_feat // n_head |
| 1008 | self.s_d_k = math.sqrt(self.d_k) |
| 1009 | self.h = n_head |
| 1010 | self.linear_q = nn.Linear(n_feat, n_feat, bias=use_bias) |
| 1011 | self.linear_k = nn.Linear(n_feat, n_feat, bias=use_bias) |
| 1012 | self.linear_v = nn.Linear(n_feat, n_feat, bias=use_bias) |
| 1013 | self.linear_out = nn.Linear(n_feat, n_feat, bias=use_bias) |
| 1014 | |
| 1015 | def forward_qkv( |
| 1016 | self, |
| 1017 | query: torch.Tensor, |
| 1018 | key: torch.Tensor, |
| 1019 | value: torch.Tensor, |
| 1020 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 1021 | """Transforms query, key and value. |
| 1022 | |
| 1023 | Args: |
| 1024 | query (torch.Tensor): (batch, time1, size) |
| 1025 | key (torch.Tensor): (batch, time2, size) |
| 1026 | value (torch.Tensor): (batch, time2, size) |
| 1027 | returns: |
| 1028 | q (torch.Tensor): (batch, head, time1, size) |
| 1029 | k (torch.Tensor): (batch, head, time2, size) |
| 1030 | v (torch.Tensor): (batch, head, time2, size) |
| 1031 | |
| 1032 | """ |
| 1033 | n_batch = query.size(0) |
| 1034 | q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k) |
| 1035 | k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k) |
| 1036 | v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k) |
| 1037 | q = q.transpose(1, 2) |
| 1038 | k = k.transpose(1, 2) |
| 1039 | v = v.transpose(1, 2) |
| 1040 | |
| 1041 | return q, k, v |
| 1042 | |
| 1043 | def forward_attention( |
| 1044 | self, |
| 1045 | value: torch.Tensor, |
| 1046 | scores: torch.Tensor, |
| 1047 | mask: torch.Tensor | None, |
| 1048 | ) -> torch.Tensor: |
| 1049 | """Compute attention context vector. |
| 1050 | |
| 1051 | Args: |
| 1052 | value (torch.Tensor): (batch, time2, size) |
| 1053 | scores(torch.Tensor): (batch, time1, time2) |
| 1054 | mask(torch.Tensor): (batch, time1, time2) |
| 1055 | returns: |
| 1056 | value (torch.Tensor): transformed `value` |
| 1057 | (batch, time2, d_model) weighted by the |
| 1058 | attention scores |
| 1059 | |
| 1060 | """ |
| 1061 | n_batch = value.size(0) |
| 1062 | if mask is not None: |
| 1063 | mask = mask.unsqueeze(1) # (batch, 1, time1, time2) |
| 1064 | scores = scores.masked_fill(mask, -INF_VAL) |
| 1065 | attn = torch.softmax(scores, dim=-1).masked_fill( |
| 1066 | mask, 0.0 |
| 1067 | ) # (batch, head, time1, time2) |
| 1068 | else: |
| 1069 | attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2) |
| 1070 | |
| 1071 | x = torch.matmul(attn, value) # (batch, head, time1, d_k) |
| 1072 | x = x.transpose(1, 2).reshape( |
| 1073 | n_batch, -1, self.h * self.d_k |
| 1074 | ) # (batch, time1, d_model) |
| 1075 | |
| 1076 | return self.linear_out(x) # (batch, time1, d_model) |
| 1077 | |
| 1078 | def forward( |
| 1079 | self, |
| 1080 | query: torch.Tensor, |
| 1081 | key: torch.Tensor, |
| 1082 | value: torch.Tensor, |
| 1083 | mask: torch.Tensor | None, |
| 1084 | pos_emb: torch.Tensor | None = None, |
| 1085 | ) -> torch.Tensor: |
| 1086 | """Compute 'Scaled Dot Product Attention'. |
| 1087 | |
| 1088 | Args: |
| 1089 | query (torch.Tensor): (batch, time1, size) |
| 1090 | key (torch.Tensor): (batch, time2, size) |
| 1091 | value(torch.Tensor): (batch, time2, size) |
| 1092 | mask (torch.Tensor): (batch, time1, time2) |
| 1093 | pos_emb (torch.Tensor): optional relative position embeddings |
| 1094 | |
| 1095 | Returns: |
| 1096 | output (torch.Tensor): transformed `value` |
| 1097 | (batch, time1, d_model) weighted by the |
| 1098 | query dot key attention |
| 1099 | |
| 1100 | """ |
| 1101 | q, k, v = self.forward_qkv(query, key, value) |
| 1102 | |
| 1103 | scores = torch.matmul(q, k.transpose(-2, -1)) / self.s_d_k |
| 1104 | return self.forward_attention(v, scores, mask) |
| 1105 | |
| 1106 | |
| 1107 | class RelPositionMultiHeadAttention(CohereASRMultiHeadAttention): |
| 1108 | """Multi-Head Attention layer of Transformer-XL with |
| 1109 | support of relative positional encoding. |
| 1110 | Paper: https://arxiv.org/abs/1901.02860 |
| 1111 | Args: |
| 1112 | n_head (int): number of heads |
| 1113 | n_feat (int): size of the features |
| 1114 | use_bias (bool): whether to apply bias in linear |
| 1115 | and conv layers of MultiHeadAttention |
| 1116 | """ |
| 1117 | |
| 1118 | def __init__( |
| 1119 | self, |
| 1120 | n_head: int, |
| 1121 | n_feat: int, |
| 1122 | pos_bias_u: nn.Parameter | torch.Tensor | None, |
| 1123 | pos_bias_v: nn.Parameter | torch.Tensor | None, |
| 1124 | use_bias: bool = True, |
| 1125 | ) -> None: |
| 1126 | """Construct an RelPositionMultiHeadedAttention object.""" |
| 1127 | super().__init__( |
| 1128 | n_head=n_head, |
| 1129 | n_feat=n_feat, |
| 1130 | use_bias=use_bias, |
| 1131 | ) |
| 1132 | # linear transformation for positional encoding |
| 1133 | self.linear_pos = nn.Linear(n_feat, n_feat, bias=False) |
| 1134 | # these two learnable biases are used in matrix c and matrix d |
| 1135 | # as described in https://arxiv.org/abs/1901.02860 Section 3.3 |
| 1136 | if pos_bias_u is None or pos_bias_v is None: |
| 1137 | self.pos_bias_u = nn.Parameter( |
| 1138 | torch.zeros(self.h, self.d_k), requires_grad=False |
| 1139 | ) |
| 1140 | self.pos_bias_v = nn.Parameter( |
| 1141 | torch.zeros(self.h, self.d_k), requires_grad=False |
| 1142 | ) |
| 1143 | else: |
| 1144 | self.pos_bias_u = pos_bias_u |
| 1145 | self.pos_bias_v = pos_bias_v |
| 1146 | |
| 1147 | def rel_shift(self, x: torch.Tensor) -> torch.Tensor: |
| 1148 | """Compute relative positional encoding. |
| 1149 | |
| 1150 | Args: |
| 1151 | x (torch.Tensor): (batch, nheads, time, 2*time-1) |
| 1152 | |
| 1153 | """ |
| 1154 | b, h, qlen, pos_len = x.size() # (b, h, t1, t2) |
| 1155 | # need to add a column of zeros on the left side of |
| 1156 | # last dimension to perform the relative shifting |
| 1157 | x = torch.nn.functional.pad(x, pad=(1, 0)) # (b, h, t1, t2+1) |
| 1158 | x = x.view(b, h, -1, qlen) # (b, h, t2+1, t1) |
| 1159 | # need to drop the first row |
| 1160 | x = x[:, :, 1:].view(b, h, qlen, pos_len) # (b, h, t1, t2) |
| 1161 | return x |
| 1162 | |
| 1163 | def forward( |
| 1164 | self, |
| 1165 | query: torch.Tensor, |
| 1166 | key: torch.Tensor, |
| 1167 | value: torch.Tensor, |
| 1168 | mask: torch.Tensor | None, |
| 1169 | pos_emb: torch.Tensor | None = None, |
| 1170 | ) -> torch.Tensor: |
| 1171 | """Compute 'Scaled Dot Product Attention' with rel. positional encoding. |
| 1172 | |
| 1173 | Args: |
| 1174 | query (torch.Tensor): (batch, time1, size) |
| 1175 | key (torch.Tensor): (batch, time2, size) |
| 1176 | value(torch.Tensor): (batch, time2, size) |
| 1177 | mask (torch.Tensor): (batch, time1, time2) |
| 1178 | pos_emb (torch.Tensor): (batch, time1, size) |
| 1179 | |
| 1180 | Returns: |
| 1181 | output (torch.Tensor): transformed `value` |
| 1182 | (batch, time1, d_model) weighted by the |
| 1183 | query dot key attention |
| 1184 | |
| 1185 | """ |
| 1186 | q, k, v = self.forward_qkv(query, key, value) |
| 1187 | q = q.transpose(1, 2) # (batch, time1, head, d_k) |
| 1188 | |
| 1189 | assert pos_emb is not None |
| 1190 | n_batch_pos = pos_emb.size(0) |
| 1191 | p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k) |
| 1192 | p = p.transpose(1, 2) # (batch, head, time1, d_k) |
| 1193 | |
| 1194 | # (batch, head, time1, d_k) |
| 1195 | q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2) |
| 1196 | # (batch, head, time1, d_k) |
| 1197 | q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2) |
| 1198 | |
| 1199 | # compute attention score |
| 1200 | # first compute matrix a and matrix c |
| 1201 | # as described in https://arxiv.org/abs/1901.02860 Section 3.3 |
| 1202 | # (batch, head, time1, time2) |
| 1203 | |
| 1204 | # compute matrix b and matrix d |
| 1205 | # (batch, head, time1, time2) |
| 1206 | matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) |
| 1207 | matrix_bd = self.rel_shift(matrix_bd) |
| 1208 | |
| 1209 | # drops extra elements in matrix_bd to match the key sequence length |
| 1210 | matrix_bd = matrix_bd[:, :, :, : k.size(-2)] |
| 1211 | batch, heads, query_len, key_len = matrix_bd.shape |
| 1212 | scale = 1.0 / self.s_d_k |
| 1213 | scores = torch.baddbmm( |
| 1214 | matrix_bd.flatten(0, 1), |
| 1215 | q_with_bias_u.flatten(0, 1), |
| 1216 | k.transpose(-2, -1).flatten(0, 1), |
| 1217 | beta=scale, |
| 1218 | alpha=scale, |
| 1219 | ).view(batch, heads, query_len, key_len) |
| 1220 | |
| 1221 | return self.forward_attention(v, scores, mask) |
| 1222 | |
| 1223 | |
| 1224 | class ConformerLayer(torch.nn.Module): |
| 1225 | """A single block of the Conformer encoder. |
| 1226 | |
| 1227 | Args: |
| 1228 | d_model (int): input dimension of |
| 1229 | MultiheadAttentionMechanism and |
| 1230 | PositionwiseFeedForward |
| 1231 | d_ff (int): hidden dimension of |
| 1232 | PositionwiseFeedForward |
| 1233 | self_attention_model (str): type of the attention |
| 1234 | layer and positional encoding |
| 1235 | n_heads (int): number of heads for multi-head |
| 1236 | attention |
| 1237 | conv_kernel_size (int): kernel size for depthwise |
| 1238 | convolution in convolution module |
| 1239 | use_bias (bool): Apply bias to all Linear and |
| 1240 | Conv1d layers from each ConformerLayer to |
| 1241 | improve activation flow and stabilize training |
| 1242 | of huge models. Defaults to True. |
| 1243 | |
| 1244 | """ |
| 1245 | |
| 1246 | def __init__( |
| 1247 | self, |
| 1248 | d_model: int, |
| 1249 | d_ff: int, |
| 1250 | self_attention_model: str = "rel_pos", |
| 1251 | n_heads: int = 4, |
| 1252 | conv_kernel_size: int = 31, |
| 1253 | conv_norm_type: str = "batch_norm", |
| 1254 | conv_context_size: int | list[int] | None = None, |
| 1255 | pos_bias_u: nn.Parameter | torch.Tensor | None = None, |
| 1256 | pos_bias_v: nn.Parameter | torch.Tensor | None = None, |
| 1257 | att_context_size: list[int] | None = None, |
| 1258 | use_bias: bool = True, |
| 1259 | ) -> None: |
| 1260 | super().__init__() |
| 1261 | if att_context_size is None: |
| 1262 | att_context_size = [-1, -1] |
| 1263 | |
| 1264 | self.self_attention_model = self_attention_model |
| 1265 | self.fc_factor = 0.5 |
| 1266 | |
| 1267 | # first feed forward module |
| 1268 | self.norm_feed_forward1 = nn.LayerNorm(d_model) |
| 1269 | self.feed_forward1 = ConformerFeedForward( |
| 1270 | d_model=d_model, d_ff=d_ff, use_bias=use_bias |
| 1271 | ) |
| 1272 | |
| 1273 | # convolution module |
| 1274 | self.norm_conv = nn.LayerNorm(d_model) |
| 1275 | self.conv = ConformerConvolution( |
| 1276 | d_model=d_model, |
| 1277 | kernel_size=conv_kernel_size, |
| 1278 | norm_type=conv_norm_type, |
| 1279 | conv_context_size=conv_context_size, |
| 1280 | use_bias=use_bias, |
| 1281 | ) |
| 1282 | |
| 1283 | # multi-headed self-attention module |
| 1284 | self.norm_self_att = nn.LayerNorm(d_model) |
| 1285 | |
| 1286 | assert self_attention_model == "rel_pos" |
| 1287 | |
| 1288 | self.self_attn = RelPositionMultiHeadAttention( |
| 1289 | n_head=n_heads, |
| 1290 | n_feat=d_model, |
| 1291 | pos_bias_u=pos_bias_u, |
| 1292 | pos_bias_v=pos_bias_v, |
| 1293 | use_bias=use_bias, |
| 1294 | ) |
| 1295 | |
| 1296 | # second feed forward module |
| 1297 | self.norm_feed_forward2 = nn.LayerNorm(d_model) |
| 1298 | self.feed_forward2 = ConformerFeedForward( |
| 1299 | d_model=d_model, d_ff=d_ff, use_bias=use_bias |
| 1300 | ) |
| 1301 | |
| 1302 | self.norm_out = nn.LayerNorm(d_model) |
| 1303 | |
| 1304 | def forward( |
| 1305 | self, |
| 1306 | x: torch.Tensor, |
| 1307 | att_mask: torch.Tensor | None = None, |
| 1308 | pos_emb: torch.Tensor | None = None, |
| 1309 | pad_mask: torch.Tensor | None = None, |
| 1310 | ) -> torch.Tensor: |
| 1311 | """Args: |
| 1312 | x (torch.Tensor): input signals (B, T, d_model) |
| 1313 | att_mask (torch.Tensor): attention masks(B, T, T) |
| 1314 | pos_emb (torch.Tensor): (L, 1, d_model) |
| 1315 | pad_mask (torch.tensor): padding mask |
| 1316 | Returns: |
| 1317 | x (torch.Tensor): (B, T, d_model) |
| 1318 | |
| 1319 | """ |
| 1320 | residual = x |
| 1321 | x = self.norm_feed_forward1(x) |
| 1322 | x = self.feed_forward1(x) |
| 1323 | residual = residual + x * self.fc_factor |
| 1324 | |
| 1325 | x = self.norm_self_att(residual) |
| 1326 | if self.self_attention_model == "rel_pos": |
| 1327 | x = self.self_attn( |
| 1328 | query=x, |
| 1329 | key=x, |
| 1330 | value=x, |
| 1331 | mask=att_mask, |
| 1332 | pos_emb=pos_emb, |
| 1333 | ) |
| 1334 | elif self.self_attention_model == "rel_pos_local_attn": |
| 1335 | x = self.self_attn( |
| 1336 | query=x, |
| 1337 | key=x, |
| 1338 | value=x, |
| 1339 | pad_mask=pad_mask, |
| 1340 | pos_emb=pos_emb, |
| 1341 | ) |
| 1342 | elif self.self_attention_model == "abs_pos": |
| 1343 | x = self.self_attn(query=x, key=x, value=x, mask=att_mask) |
| 1344 | else: |
| 1345 | x = None |
| 1346 | |
| 1347 | residual = residual + x |
| 1348 | |
| 1349 | x = self.norm_conv(residual) |
| 1350 | x = self.conv(x, pad_mask=pad_mask) |
| 1351 | residual = residual + x |
| 1352 | |
| 1353 | x = self.norm_feed_forward2(residual) |
| 1354 | x = self.feed_forward2(x) |
| 1355 | residual = residual + x * self.fc_factor |
| 1356 | |
| 1357 | x = self.norm_out(residual) |
| 1358 | |
| 1359 | return x |
| 1360 | |
| 1361 | |
| 1362 | class ConformerEncoder(nn.Module): |
| 1363 | """The encoder for ASR model of Conformer. |
| 1364 | Based on this paper: |
| 1365 | 'Conformer: Convolution-augmented Transformer for |
| 1366 | Speech Recognition' by Anmol Gulati et al. |
| 1367 | https://arxiv.org/abs/2005.08100 |
| 1368 | """ |
| 1369 | |
| 1370 | def __init__(self, *, vllm_config: VllmConfig): |
| 1371 | super().__init__() |
| 1372 | |
| 1373 | self.hf_config = vllm_config.model_config.hf_config |
| 1374 | |
| 1375 | feat_in = self.hf_config.encoder["feat_in"] |
| 1376 | n_layers = self.hf_config.encoder["n_layers"] |
| 1377 | d_model = self.hf_config.encoder["d_model"] |
| 1378 | feat_out = self.hf_config.encoder["feat_out"] |
| 1379 | causal_downsampling = self.hf_config.encoder["causal_downsampling"] |
| 1380 | subsampling = self.hf_config.encoder["subsampling"] |
| 1381 | subsampling_factor = self.hf_config.encoder["subsampling_factor"] |
| 1382 | subsampling_conv_chunking_factor = self.hf_config.encoder.get( |
| 1383 | "subsampling_conv_chunking_factor", 1 |
| 1384 | ) |
| 1385 | subsampling_conv_channels = self.hf_config.encoder["subsampling_conv_channels"] |
| 1386 | ff_expansion_factor = self.hf_config.encoder["ff_expansion_factor"] |
| 1387 | self_attention_model = self.hf_config.encoder["self_attention_model"] |
| 1388 | n_heads = self.hf_config.encoder["n_heads"] |
| 1389 | att_context_size = self.hf_config.encoder["att_context_size"] |
| 1390 | att_context_probs = self.hf_config.encoder.get("att_context_probs", None) |
| 1391 | att_context_style = self.hf_config.encoder.get("att_context_style", "regular") |
| 1392 | xscaling = self.hf_config.encoder["xscaling"] |
| 1393 | untie_biases = self.hf_config.encoder["untie_biases"] |
| 1394 | pos_emb_max_len = self.hf_config.encoder["pos_emb_max_len"] |
| 1395 | conv_kernel_size = self.hf_config.encoder["conv_kernel_size"] |
| 1396 | conv_norm_type = self.hf_config.encoder["conv_norm_type"] |
| 1397 | conv_context_size = self.hf_config.encoder["conv_context_size"] |
| 1398 | use_bias = self.hf_config.encoder.get("use_bias", True) |
| 1399 | |
| 1400 | d_ff = d_model * ff_expansion_factor |
| 1401 | self.d_model = d_model |
| 1402 | self._feat_in = feat_in |
| 1403 | self.att_context_style = att_context_style |
| 1404 | self.subsampling_factor = subsampling_factor |
| 1405 | |
| 1406 | self.self_attention_model = self_attention_model |
| 1407 | |
| 1408 | # Setting up the att_context_size |
| 1409 | ( |
| 1410 | _, |
| 1411 | self.att_context_size, |
| 1412 | _, |
| 1413 | self.conv_context_size, |
| 1414 | ) = self._calc_context_sizes( |
| 1415 | att_context_style=att_context_style, |
| 1416 | att_context_size=att_context_size, |
| 1417 | att_context_probs=att_context_probs, |
| 1418 | conv_context_size=conv_context_size, |
| 1419 | conv_kernel_size=conv_kernel_size, |
| 1420 | ) |
| 1421 | |
| 1422 | self.xscale: float | None |
| 1423 | if xscaling: |
| 1424 | self.xscale = math.sqrt(d_model) |
| 1425 | else: |
| 1426 | self.xscale = None |
| 1427 | |
| 1428 | # Subsampling |
| 1429 | if subsampling_conv_channels == -1: |
| 1430 | subsampling_conv_channels = d_model |
| 1431 | assert subsampling and subsampling_factor > 1 and subsampling == "dw_striding" |
| 1432 | |
| 1433 | self.pre_encode = ConvSubsampling( |
| 1434 | subsampling=subsampling, |
| 1435 | subsampling_factor=subsampling_factor, |
| 1436 | feat_in=feat_in, |
| 1437 | feat_out=d_model, |
| 1438 | conv_channels=subsampling_conv_channels, |
| 1439 | subsampling_conv_chunking_factor=subsampling_conv_chunking_factor, |
| 1440 | activation=nn.ReLU(True), |
| 1441 | is_causal=causal_downsampling, |
| 1442 | ) |
| 1443 | |
| 1444 | self._feat_out = d_model |
| 1445 | |
| 1446 | # Biases for relative positional encoding |
| 1447 | if not untie_biases and self_attention_model == "rel_pos": |
| 1448 | d_head = d_model // n_heads |
| 1449 | # Register as buffers instead of parameters since they're not trainable |
| 1450 | # and need to respect dtype during weight loading |
| 1451 | self.register_buffer( |
| 1452 | "pos_bias_u", torch.zeros(n_heads, d_head), persistent=True |
| 1453 | ) |
| 1454 | self.register_buffer( |
| 1455 | "pos_bias_v", torch.zeros(n_heads, d_head), persistent=True |
| 1456 | ) |
| 1457 | pos_bias_u = self.pos_bias_u |
| 1458 | pos_bias_v = self.pos_bias_v |
| 1459 | else: |
| 1460 | pos_bias_u = None |
| 1461 | pos_bias_v = None |
| 1462 | |
| 1463 | # Positional encodings |
| 1464 | self.pos_emb_max_len = pos_emb_max_len |
| 1465 | assert self_attention_model == "rel_pos" |
| 1466 | self.pos_enc = RelPositionalEncoding( |
| 1467 | d_model=d_model, |
| 1468 | max_len=pos_emb_max_len, |
| 1469 | xscale=self.xscale, |
| 1470 | ) |
| 1471 | |
| 1472 | self.layers = nn.ModuleList() |
| 1473 | for i in range(n_layers): |
| 1474 | layer = ConformerLayer( |
| 1475 | d_model=d_model, |
| 1476 | d_ff=d_ff, |
| 1477 | self_attention_model=self_attention_model, |
| 1478 | n_heads=n_heads, |
| 1479 | conv_kernel_size=conv_kernel_size, |
| 1480 | conv_norm_type=conv_norm_type, |
| 1481 | conv_context_size=self.conv_context_size, |
| 1482 | pos_bias_u=pos_bias_u, |
| 1483 | pos_bias_v=pos_bias_v, |
| 1484 | att_context_size=self.att_context_size, |
| 1485 | use_bias=use_bias, |
| 1486 | ) |
| 1487 | self.layers.append(layer) |
| 1488 | |
| 1489 | if feat_out > 0 and feat_out != self._feat_out: |
| 1490 | self.out_proj = nn.Linear(self._feat_out, feat_out) |
| 1491 | self._feat_out = feat_out |
| 1492 | else: |
| 1493 | self.out_proj = None |
| 1494 | self._feat_out = d_model |
| 1495 | self.set_max_audio_length(self.pos_emb_max_len) |
| 1496 | |
| 1497 | def get_num_encoder_cross_attn_tokens(self, num_encoder_input_tokens: int) -> int: |
| 1498 | num_encoder_cross_attn_tokens = math.ceil( |
| 1499 | num_encoder_input_tokens / self.subsampling_factor |
| 1500 | ) |
| 1501 | return num_encoder_cross_attn_tokens |
| 1502 | |
| 1503 | def set_max_audio_length(self, max_audio_length: int) -> None: |
| 1504 | """Sets maximum input length. |
| 1505 | Pre-calculates internal seq_range mask. |
| 1506 | |
| 1507 | Args: |
| 1508 | max_audio_length (int): New maximum sequence length. |
| 1509 | |
| 1510 | """ |
| 1511 | device = next(self.parameters()).device |
| 1512 | dtype = next(self.parameters()).dtype |
| 1513 | self.pos_enc.extend_pe(max_audio_length, device, dtype) |
| 1514 | |
| 1515 | def forward( |
| 1516 | self, |
| 1517 | audio_signal: torch.Tensor, |
| 1518 | length: torch.Tensor, |
| 1519 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1520 | if audio_signal.shape[-2] != self._feat_in: |
| 1521 | raise ValueError( |
| 1522 | f"audio_signal should have shape " |
| 1523 | f"(batch, {self._feat_in}, n_frame) but " |
| 1524 | f"got last dimension " |
| 1525 | f"{audio_signal.shape[-2]}." |
| 1526 | ) |
| 1527 | |
| 1528 | return self.forward_internal( |
| 1529 | audio_signal, |
| 1530 | length, |
| 1531 | ) |
| 1532 | |
| 1533 | def forward_internal( |
| 1534 | self, |
| 1535 | audio_signal: torch.Tensor, |
| 1536 | length: torch.Tensor | None = None, |
| 1537 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1538 | if length is None: |
| 1539 | length = audio_signal.new_full( |
| 1540 | (audio_signal.size(0),), |
| 1541 | audio_signal.size(-1), |
| 1542 | dtype=torch.int64, |
| 1543 | device=audio_signal.device, |
| 1544 | ) |
| 1545 | |
| 1546 | cur_att_context_size = self.att_context_size |
| 1547 | audio_signal = torch.transpose(audio_signal, 1, 2) |
| 1548 | |
| 1549 | audio_signal, length = self.pre_encode(x=audio_signal, lengths=length) |
| 1550 | length = length.to(torch.int64) |
| 1551 | |
| 1552 | max_audio_length = audio_signal.size(1) |
| 1553 | |
| 1554 | padding_length = length |
| 1555 | |
| 1556 | audio_signal, pos_emb = self.pos_enc(x=audio_signal, cache_len=0) |
| 1557 | |
| 1558 | pad_mask, att_mask = self._create_masks( |
| 1559 | att_context_size=cur_att_context_size, |
| 1560 | padding_length=padding_length, |
| 1561 | max_audio_length=max_audio_length, |
| 1562 | offset=None, |
| 1563 | device=audio_signal.device, |
| 1564 | ) |
| 1565 | |
| 1566 | for lth, layer in enumerate(self.layers): |
| 1567 | audio_signal = layer( |
| 1568 | x=audio_signal, |
| 1569 | att_mask=att_mask, |
| 1570 | pos_emb=pos_emb, |
| 1571 | pad_mask=pad_mask, |
| 1572 | ) |
| 1573 | |
| 1574 | if self.out_proj is not None: |
| 1575 | audio_signal = self.out_proj(audio_signal) |
| 1576 | |
| 1577 | audio_signal = torch.transpose(audio_signal, 1, 2) |
| 1578 | length = length.to(dtype=torch.int64) |
| 1579 | |
| 1580 | return audio_signal, length |
| 1581 | |
| 1582 | def _create_masks( |
| 1583 | self, |
| 1584 | att_context_size: list[int], |
| 1585 | padding_length: torch.Tensor, |
| 1586 | max_audio_length: int, |
| 1587 | offset: torch.Tensor | None, |
| 1588 | device: torch.device, |
| 1589 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 1590 | if self.self_attention_model != "rel_pos_local_attn": |
| 1591 | att_mask = torch.ones( |
| 1592 | 1, max_audio_length, max_audio_length, dtype=torch.bool, device=device |
| 1593 | ) |
| 1594 | |
| 1595 | if self.att_context_style == "regular": |
| 1596 | if att_context_size[0] >= 0: |
| 1597 | att_mask = att_mask.triu(diagonal=-att_context_size[0]) |
| 1598 | if att_context_size[1] >= 0: |
| 1599 | att_mask = att_mask.tril(diagonal=att_context_size[1]) |
| 1600 | elif self.att_context_style == "chunked_limited": |
| 1601 | # When right context is unlimited, just the |
| 1602 | # left side of masking needs to get updated |
| 1603 | if att_context_size[1] == -1: |
| 1604 | if att_context_size[0] >= 0: |
| 1605 | att_mask = att_mask.triu(diagonal=-att_context_size[0]) |
| 1606 | else: |
| 1607 | chunk_size = att_context_size[1] + 1 |
| 1608 | # left_chunks_num specifies the number |
| 1609 | # of chunks to be visible by each chunk |
| 1610 | # on the left side |
| 1611 | if att_context_size[0] >= 0: |
| 1612 | left_chunks_num = att_context_size[0] // chunk_size |
| 1613 | else: |
| 1614 | left_chunks_num = 10000 |
| 1615 | |
| 1616 | chunk_idx = torch.arange( |
| 1617 | 0, max_audio_length, dtype=torch.int, device=att_mask.device |
| 1618 | ) |
| 1619 | chunk_idx = torch.div(chunk_idx, chunk_size, rounding_mode="trunc") |
| 1620 | diff_chunks = chunk_idx.unsqueeze(1) - chunk_idx.unsqueeze(0) |
| 1621 | chunked_limited_mask = torch.logical_and( |
| 1622 | torch.le(diff_chunks, left_chunks_num), torch.ge(diff_chunks, 0) |
| 1623 | ) |
| 1624 | att_mask = torch.logical_and( |
| 1625 | att_mask, chunked_limited_mask.unsqueeze(0) |
| 1626 | ) |
| 1627 | else: |
| 1628 | att_mask = None |
| 1629 | |
| 1630 | # pad_mask is the masking to be used to ignore paddings |
| 1631 | pad_mask = torch.arange(0, max_audio_length, device=device).expand( |
| 1632 | padding_length.size(0), -1 |
| 1633 | ) < padding_length.unsqueeze(-1) |
| 1634 | |
| 1635 | if offset is not None: |
| 1636 | pad_mask_off = torch.arange(0, max_audio_length, device=device).expand( |
| 1637 | padding_length.size(0), -1 |
| 1638 | ) >= offset.unsqueeze(-1) |
| 1639 | pad_mask = pad_mask_off.logical_and(pad_mask) |
| 1640 | |
| 1641 | if att_mask is not None: |
| 1642 | # pad_mask_for_att_mask is the mask which helps to ignore paddings |
| 1643 | pad_mask_for_att_mask = pad_mask.unsqueeze(1).repeat( |
| 1644 | [1, max_audio_length, 1] |
| 1645 | ) |
| 1646 | pad_mask_for_att_mask = torch.logical_and( |
| 1647 | pad_mask_for_att_mask, pad_mask_for_att_mask.transpose(1, 2) |
| 1648 | ) |
| 1649 | # att_mask is the masking to be used by MHA |
| 1650 | # layers to ignore tokens not supposed to be |
| 1651 | # visible |
| 1652 | att_mask = att_mask[:, :max_audio_length, :max_audio_length] |
| 1653 | # paddings should also get ignored, so |
| 1654 | # pad_mask_for_att_mask is used to ignore their |
| 1655 | # corresponding scores |
| 1656 | att_mask = torch.logical_and( |
| 1657 | pad_mask_for_att_mask, att_mask.to(pad_mask_for_att_mask.device) |
| 1658 | ) |
| 1659 | att_mask = ~att_mask |
| 1660 | |
| 1661 | pad_mask = ~pad_mask |
| 1662 | return pad_mask, att_mask |
| 1663 | |
| 1664 | def _calc_context_sizes( |
| 1665 | self, |
| 1666 | att_context_size: list[int] | list[list[int]] | None, |
| 1667 | att_context_probs: list[float] | None, |
| 1668 | att_context_style: str, |
| 1669 | conv_context_size: list[int] | str | None, |
| 1670 | conv_kernel_size: int, |
| 1671 | ) -> tuple[list[list[int]], list[int], list[float], list[int]]: |
| 1672 | # convert att_context_size to a standard list of lists |
| 1673 | if att_context_size: |
| 1674 | if is_list_of(att_context_size, int, check="all"): |
| 1675 | att_context_size_all = [att_context_size] |
| 1676 | elif is_list_of(att_context_size, list, check="all"): |
| 1677 | att_context_size_all = att_context_size |
| 1678 | else: |
| 1679 | raise ValueError("att_context_size cannot mix nested and flat values") |
| 1680 | for i, att_cs in enumerate(att_context_size_all): |
| 1681 | if att_context_style == "chunked_limited": |
| 1682 | if att_cs[0] > 0 and att_cs[0] % (att_cs[1] + 1) > 0: |
| 1683 | raise ValueError( |
| 1684 | f"att_context_size[{i}][0] % " |
| 1685 | f"(att_context_size[{i}][1]" |
| 1686 | f" + 1) should be zero!" |
| 1687 | ) |
| 1688 | if att_cs[1] < 0 and len(att_context_size_all) <= 1: |
| 1689 | raise ValueError( |
| 1690 | f"Right context " |
| 1691 | f"(att_context_size[{i}][1])" |
| 1692 | f" can not be unlimited for" |
| 1693 | f" chunked_limited style!" |
| 1694 | ) |
| 1695 | else: |
| 1696 | att_context_size_all = [[-1, -1]] |
| 1697 | |
| 1698 | if att_context_probs: |
| 1699 | if len(att_context_probs) != len(att_context_size_all): |
| 1700 | raise ValueError( |
| 1701 | "The size of the att_context_probs " |
| 1702 | "should be the same as att_context_size." |
| 1703 | ) |
| 1704 | att_context_probs = list(att_context_probs) |
| 1705 | if sum(att_context_probs) != 1: |
| 1706 | raise ValueError( |
| 1707 | "The sum of numbers in " |
| 1708 | "att_context_probs should be equal " |
| 1709 | "to one to be a distribution." |
| 1710 | ) |
| 1711 | else: |
| 1712 | att_context_probs = [1.0 / len(att_context_size_all)] * len( |
| 1713 | att_context_size_all |
| 1714 | ) |
| 1715 | |
| 1716 | if conv_context_size is not None: |
| 1717 | if not isinstance(conv_context_size, list) and not isinstance( |
| 1718 | conv_context_size, str |
| 1719 | ): |
| 1720 | raise ValueError( |
| 1721 | "Invalid conv_context_size! It should " |
| 1722 | "be the string 'causal' or a list of " |
| 1723 | "two integers." |
| 1724 | ) |
| 1725 | if conv_context_size == "causal": |
| 1726 | conv_context_size = [conv_kernel_size - 1, 0] |
| 1727 | else: |
| 1728 | assert isinstance(conv_context_size, list) |
| 1729 | total = conv_context_size[0] + conv_context_size[1] + 1 |
| 1730 | if total != conv_kernel_size: |
| 1731 | raise ValueError( |
| 1732 | f"Invalid conv_context_size: {self.conv_context_size}!" |
| 1733 | ) |
| 1734 | else: |
| 1735 | conv_context_size = [ |
| 1736 | (conv_kernel_size - 1) // 2, |
| 1737 | (conv_kernel_size - 1) // 2, |
| 1738 | ] |
| 1739 | return ( |
| 1740 | att_context_size_all, |
| 1741 | att_context_size_all[0], |
| 1742 | att_context_probs, |
| 1743 | conv_context_size, |
| 1744 | ) |
| 1745 | |
| 1746 | |
| 1747 | # ----- Encoder END ----- |
| 1748 | |
| 1749 | |
| 1750 | # This subclass is specific to vLLM in order for |
| 1751 | # `_mark_composite_model` to target this module |
| 1752 | class CohereASRProjector(nn.Linear): |
| 1753 | pass |
| 1754 | |
| 1755 | |
| 1756 | class CohereASRModel(nn.Module): |
| 1757 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 1758 | super().__init__() |
| 1759 | self.encoder = ConformerEncoder(vllm_config=vllm_config) |
| 1760 | |
| 1761 | self.decoder = CohereASRDecoder( |
| 1762 | vllm_config=vllm_config, |
| 1763 | prefix=maybe_prefix(prefix, "decoder"), |
| 1764 | ) |
| 1765 | |
| 1766 | if self.encoder.d_model != self.decoder.hidden_size: |
| 1767 | self.encoder_decoder_proj = CohereASRProjector( |
| 1768 | self.encoder.d_model, self.decoder.hidden_size |
| 1769 | ) |
| 1770 | |
| 1771 | def forward( |
| 1772 | self, |
| 1773 | input_ids: torch.Tensor | None, |
| 1774 | positions: torch.Tensor, |
| 1775 | encoder_outputs: list[torch.Tensor], |
| 1776 | ) -> torch.Tensor: |
| 1777 | enc_states = torch.cat(encoder_outputs, dim=0) if len(encoder_outputs) else None |
| 1778 | decoder_outputs = self.decoder( |
| 1779 | input_ids=input_ids, |
| 1780 | positions=positions, |
| 1781 | encoder_hidden_states=enc_states, |
| 1782 | ) |
| 1783 | |
| 1784 | return decoder_outputs |
| 1785 | |
| 1786 | def get_encoder_outputs( |
| 1787 | self, |
| 1788 | input_features: torch.Tensor | list[torch.Tensor] | None, |
| 1789 | seq_lens: torch.Tensor | None, |
| 1790 | ) -> torch.Tensor | None: |
| 1791 | if input_features is None: |
| 1792 | return None |
| 1793 | |
| 1794 | if isinstance(input_features, torch.Tensor): |
| 1795 | encoder_input_length = seq_lens |
| 1796 | out, encoder_output_length = self.encoder( |
| 1797 | input_features, length=encoder_input_length |
| 1798 | ) # B x D x T |
| 1799 | out = out.permute(0, 2, 1) |
| 1800 | |
| 1801 | if hasattr(self, "encoder_decoder_proj"): |
| 1802 | out = self.encoder_decoder_proj(out) |
| 1803 | |
| 1804 | # Convert padded tensor to packed |
| 1805 | with gpu_sync_allowed(): |
| 1806 | lengths = encoder_output_length.tolist() |
| 1807 | outs = [feat[:length, :] for feat, length in zip(out, lengths)] |
| 1808 | |
| 1809 | return outs |
| 1810 | else: |
| 1811 | raise NotImplementedError("List input_features not supported") |
| 1812 | |
| 1813 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 1814 | stacked_params_mapping = [ |
| 1815 | # (param_name, shard_name, shard_id) |
| 1816 | (".first_sub_layer.qkv_proj", ".first_sub_layer.query_net", "q"), |
| 1817 | (".first_sub_layer.qkv_proj", ".first_sub_layer.key_net", "k"), |
| 1818 | (".first_sub_layer.qkv_proj", ".first_sub_layer.value_net", "v"), |
| 1819 | (".second_sub_layer.kv_proj", ".second_sub_layer.key_net", "k"), |
| 1820 | (".second_sub_layer.kv_proj", ".second_sub_layer.value_net", "v"), |
| 1821 | ] |
| 1822 | params_dict = dict(self.named_parameters()) |
| 1823 | buffers_dict = dict(self.named_buffers()) |
| 1824 | params_dict.update(buffers_dict) |
| 1825 | |
| 1826 | loaded_params: set[str] = set() |
| 1827 | for name, loaded_weight in weights: |
| 1828 | for param_name, weight_name, shard_id in stacked_params_mapping: |
| 1829 | if weight_name not in name: |
| 1830 | continue |
| 1831 | name = name.replace(weight_name, param_name) |
| 1832 | # Skip loading extra bias for GPTQ models. |
| 1833 | # if name.endswith(".bias") and name not in params_dict: |
| 1834 | # continue |
| 1835 | |
| 1836 | param = params_dict[name] |
| 1837 | weight_loader = param.weight_loader |
| 1838 | weight_loader(param, loaded_weight, shard_id) |
| 1839 | break |
| 1840 | else: |
| 1841 | # Skip loading extra bias for GPTQ models. |
| 1842 | if name.endswith(".bias") and name not in params_dict: |
| 1843 | continue |
| 1844 | |
| 1845 | param = params_dict[name] |
| 1846 | weight_loader = getattr(param, "weight_loader", default_weight_loader) |
| 1847 | |
| 1848 | # Convert buffer dtype to match loaded weight for pos_bias tensors |
| 1849 | if "pos_bias" in name and param.dtype != loaded_weight.dtype: |
| 1850 | logger.info( |
| 1851 | "Converting buffer %s dtype from %s to %s for loading.", |
| 1852 | name, |
| 1853 | param.dtype, |
| 1854 | loaded_weight.dtype, |
| 1855 | ) |
| 1856 | param.data = param.data.to(loaded_weight.dtype) |
| 1857 | |
| 1858 | weight_loader(param, loaded_weight) |
| 1859 | loaded_params.add(name) |
| 1860 | return loaded_params |
| 1861 | |
| 1862 | |
| 1863 | class CohereASRProcessingInfo(BaseProcessingInfo): |
| 1864 | def get_hf_config(self) -> PreTrainedConfig: |
| 1865 | return self.ctx.get_hf_config() |
| 1866 | |
| 1867 | def get_default_tok_params(self) -> TokenizeParams: |
| 1868 | # Special tokens should be provided by the user based on the |
| 1869 | # task and language of their request. Also needed to avoid |
| 1870 | # appending an EOS token to the prompt which disrupts generation. |
| 1871 | return super().get_default_tok_params().with_kwargs(add_special_tokens=False) |
| 1872 | |
| 1873 | def get_hf_processor(self, **kwargs: object) -> CohereASRProcessor: |
| 1874 | if not hasattr(self, "_cached_hf_processor"): |
| 1875 | hf_config = self.get_hf_config() |
| 1876 | preproc = hf_config.preprocessor |
| 1877 | |
| 1878 | sample_rate = preproc.get("sample_rate", 16000) |
| 1879 | window_size = preproc.get("window_size", 0.02) |
| 1880 | window_stride = preproc.get("window_stride", 0.01) |
| 1881 | |
| 1882 | feature_extractor = CohereASRFeatureExtractor( |
| 1883 | feature_size=preproc.get("features", 64), |
| 1884 | sampling_rate=sample_rate, |
| 1885 | padding_value=preproc.get("pad_value", 0.0), |
| 1886 | max_duration=hf_config.max_audio_clip_s, |
| 1887 | n_window_size=int(window_size * sample_rate), |
| 1888 | n_window_stride=int(window_stride * sample_rate), |
| 1889 | window=preproc.get("window", "hann"), |
| 1890 | normalize=preproc.get("normalize", "per_feature"), |
| 1891 | n_fft=preproc.get("n_fft", None), |
| 1892 | preemph=preproc.get("preemph", 0.97), |
| 1893 | lowfreq=preproc.get("lowfreq", 0), |
| 1894 | highfreq=preproc.get("highfreq", None), |
| 1895 | log=preproc.get("log", True), |
| 1896 | log_zero_guard_type=preproc.get("log_zero_guard_type", "add"), |
| 1897 | log_zero_guard_value=preproc.get("log_zero_guard_value", 2**-24), |
| 1898 | dither=preproc.get("dither", 1e-05), |
| 1899 | pad_to=preproc.get("pad_to", 16), |
| 1900 | frame_splicing=preproc.get("frame_splicing", 1), |
| 1901 | exact_pad=preproc.get("exact_pad", False), |
| 1902 | mag_power=preproc.get("mag_power", 2.0), |
| 1903 | mel_norm=preproc.get("mel_norm", "slaney"), |
| 1904 | stft_exact_pad=preproc.get("stft_exact_pad", False), |
| 1905 | stft_conv=preproc.get("stft_conv", False), |
| 1906 | device="cpu", |
| 1907 | ) |
| 1908 | |
| 1909 | tokenizer = self.ctx.tokenizer |
| 1910 | self._cached_hf_processor = CohereASRProcessor( |
| 1911 | feature_extractor=feature_extractor, |
| 1912 | tokenizer=tokenizer, |
| 1913 | ) |
| 1914 | return self._cached_hf_processor |
| 1915 | |
| 1916 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 1917 | return {"audio": 1} |
| 1918 | |
| 1919 | def get_data_parser(self) -> MultiModalDataParser: |
| 1920 | feature_extractor = self.get_feature_extractor() |
| 1921 | return MultiModalDataParser(target_sr=feature_extractor.sampling_rate) |
| 1922 | |
| 1923 | def get_feature_extractor(self, **kwargs: object) -> CohereASRFeatureExtractor: |
| 1924 | hf_processor = self.get_hf_processor(**kwargs) |
| 1925 | feature_extractor = hf_processor.feature_extractor |
| 1926 | assert isinstance(feature_extractor, CohereASRFeatureExtractor) |
| 1927 | return feature_extractor |
| 1928 | |
| 1929 | def get_num_audio_tokens(self, num_samples: int) -> int: |
| 1930 | num_tokens = self.get_feature_extractor().get_seq_len(num_samples) |
| 1931 | config = self.get_hf_config() |
| 1932 | subsampling_factor = config.encoder["subsampling_factor"] |
| 1933 | num_tokens = math.ceil(num_tokens / subsampling_factor) |
| 1934 | return num_tokens |
| 1935 | |
| 1936 | |
| 1937 | class CohereASRDummyInputsBuilder(BaseDummyInputsBuilder[CohereASRProcessingInfo]): |
| 1938 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 1939 | num_audios = mm_counts.get("audio", 0) |
| 1940 | |
| 1941 | return "<|startoftranscript|>" * num_audios |
| 1942 | |
| 1943 | def get_dummy_mm_data( |
| 1944 | self, |
| 1945 | seq_len: int, |
| 1946 | mm_counts: Mapping[str, int], |
| 1947 | mm_options: MultiModalDummyOptions, |
| 1948 | mm_processor_kwargs=None, |
| 1949 | ) -> MultiModalDataDict: |
| 1950 | feature_extractor = self.info.get_feature_extractor() |
| 1951 | |
| 1952 | sampling_rate = feature_extractor.sampling_rate |
| 1953 | audio_len = feature_extractor.max_duration * sampling_rate |
| 1954 | num_audios = mm_counts.get("audio", 0) |
| 1955 | |
| 1956 | return { |
| 1957 | "audio": self._get_dummy_audios(length=audio_len, num_audios=num_audios) |
| 1958 | } |
| 1959 | |
| 1960 | |
| 1961 | class CohereASRMultiModalProcessor(EncDecMultiModalProcessor[CohereASRProcessingInfo]): |
| 1962 | skip_decoder_start_token: bool = True |
| 1963 | |
| 1964 | @property |
| 1965 | def pad_dummy_encoder_prompt(self) -> bool: |
| 1966 | return True |
| 1967 | |
| 1968 | def create_encoder_prompt( |
| 1969 | self, |
| 1970 | prompt: list[int], |
| 1971 | mm_items: MultiModalDataItems, |
| 1972 | ) -> list[int]: |
| 1973 | return [0] |
| 1974 | |
| 1975 | def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: |
| 1976 | return self.dummy_inputs.get_dummy_text(mm_counts) |
| 1977 | |
| 1978 | def _get_hf_mm_inputs( |
| 1979 | self, |
| 1980 | mm_items: MultiModalDataItems, |
| 1981 | hf_kwargs: Mapping[str, object], |
| 1982 | ) -> HFMultiModalInputs: |
| 1983 | hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) |
| 1984 | |
| 1985 | feature_extractor = self.info.get_feature_extractor(**hf_kwargs) |
| 1986 | return hf_inputs._replace( |
| 1987 | hf_kwargs=dict( |
| 1988 | hf_inputs.hf_kwargs, |
| 1989 | sampling_rate=feature_extractor.sampling_rate, |
| 1990 | ) |
| 1991 | ) |
| 1992 | |
| 1993 | def _cached_apply_hf_processor( |
| 1994 | self, |
| 1995 | inputs: ProcessorInputs, |
| 1996 | timing_ctx: TimingContext, |
| 1997 | ) -> MultiModalProcessingResult: |
| 1998 | # Dithering injects noise into the extracted features, so the |
| 1999 | # feature extractor is not a pure function of its input. Since the |
| 2000 | # processing cache assumes that processor outputs are invariant |
| 2001 | # across calls, bypass the cache when dithering is active. |
| 2002 | preproc = self.info.get_hf_config().preprocessor |
| 2003 | if preproc.get("dither", 1e-05) > 0: |
| 2004 | return self._apply_hf_processor(inputs, timing_ctx) |
| 2005 | |
| 2006 | return super()._cached_apply_hf_processor(inputs, timing_ctx) |
| 2007 | |
| 2008 | def _get_mm_fields_config( |
| 2009 | self, |
| 2010 | hf_inputs, |
| 2011 | hf_processor_mm_kwargs: Mapping[str, object], |
| 2012 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 2013 | return dict( |
| 2014 | input_features=MultiModalFieldConfig.batched("audio"), |
| 2015 | length=MultiModalFieldConfig.batched("audio"), |
| 2016 | ) |
| 2017 | |
| 2018 | def _get_prompt_updates( |
| 2019 | self, |
| 2020 | mm_items: MultiModalDataItems, |
| 2021 | hf_processor_mm_kwargs: Mapping[str, object], |
| 2022 | out_mm_kwargs: MultiModalKwargsItems, |
| 2023 | ) -> Sequence[PromptUpdate]: |
| 2024 | def get_audio_replacement_cohere_asr(item_idx: int): |
| 2025 | audios = mm_items.get_items("audio", AudioProcessorItems) |
| 2026 | audio_len = audios.get_audio_length(item_idx) |
| 2027 | num_tokens = self.info.get_num_audio_tokens(num_samples=audio_len) |
| 2028 | return [0] * num_tokens |
| 2029 | |
| 2030 | return [ |
| 2031 | PromptReplacement( |
| 2032 | modality="audio", |
| 2033 | target=[0], |
| 2034 | replacement=get_audio_replacement_cohere_asr, |
| 2035 | ) |
| 2036 | ] |
| 2037 | |
| 2038 | |
| 2039 | @MULTIMODAL_REGISTRY.register_processor( |
| 2040 | CohereASRMultiModalProcessor, |
| 2041 | info=CohereASRProcessingInfo, |
| 2042 | dummy_inputs=CohereASRDummyInputsBuilder, |
| 2043 | ) |
| 2044 | class CohereAsrForConditionalGeneration( |
| 2045 | nn.Module, SupportsTranscription, SupportsMultiModal |
| 2046 | ): |
| 2047 | packed_modules_mapping = { |
| 2048 | "self_attn.qkv_proj": [ |
| 2049 | "self_attn.q_proj", |
| 2050 | "self_attn.k_proj", |
| 2051 | "self_attn.v_proj", |
| 2052 | ], |
| 2053 | "encoder_attn.kv_proj": ["encoder_attn.k_proj", "encoder_attn.v_proj"], |
| 2054 | } |
| 2055 | |
| 2056 | hf_to_vllm_mapper = WeightsMapper( |
| 2057 | orig_to_new_substr={ |
| 2058 | ".fc1.": ".mlp.fc1.", |
| 2059 | ".fc2.": ".mlp.fc2.", |
| 2060 | "model.conv.batch_norm.num_batches_tracked": None, |
| 2061 | }, |
| 2062 | orig_to_new_prefix={ |
| 2063 | "model.preprocessor.featurizer.fb": None, |
| 2064 | "model.preprocessor.featurizer.window": None, |
| 2065 | }, |
| 2066 | ) |
| 2067 | |
| 2068 | supports_transcription_only = True |
| 2069 | supported_languages = ISO639_1_SUPPORTED_LANGS |
| 2070 | skip_warmup_audio_preprocessing = True |
| 2071 | no_space_languages = {"ja", "zh"} |
| 2072 | _default_prompt_token_ids_cache: ClassVar[ |
| 2073 | dict[tuple[str | None, str | None, str], tuple[int, ...]] |
| 2074 | ] = {} |
| 2075 | |
| 2076 | @classmethod |
| 2077 | def validate_language(cls, language: str | None) -> str | None: |
| 2078 | if language is None: |
| 2079 | logger.warning( |
| 2080 | "Defaulting to language='en'. If you wish to transcribe " |
| 2081 | "audio in a different language, pass the `language` field " |
| 2082 | "in the TranscriptionRequest." |
| 2083 | ) |
| 2084 | language = "en" |
| 2085 | return super().validate_language(language) |
| 2086 | |
| 2087 | @classmethod |
| 2088 | def get_generation_prompt(cls, stt_params: SpeechToTextParams) -> PromptType: |
| 2089 | audio = stt_params.audio |
| 2090 | stt_config = stt_params.stt_config |
| 2091 | language = stt_params.language |
| 2092 | model_config = stt_params.model_config |
| 2093 | |
| 2094 | if language is None: |
| 2095 | raise ValueError( |
| 2096 | "Language must be specified when creating the CohereASR prompt" |
| 2097 | ) |
| 2098 | |
| 2099 | tokenizer = cached_tokenizer_from_config(model_config) |
| 2100 | |
| 2101 | # prompt_text is None because CoherASR uses fast implementation of |
| 2102 | # sentencepiece tokenizer which needs "▁" as the first token |
| 2103 | # (which is different from "_") and encode("▁ABC") ignores the first token |
| 2104 | # so the prompt_text is unreliable. However, prompt_token_ids can be used |
| 2105 | # to get prompt_text but it wont have the first token "▁". |
| 2106 | prompt_token_ids = cls._get_default_prompt_token_ids( |
| 2107 | tokenizer, |
| 2108 | model_config, |
| 2109 | language, |
| 2110 | ) |
| 2111 | |
| 2112 | return TokensPrompt( |
| 2113 | prompt_token_ids=prompt_token_ids, |
| 2114 | multi_modal_data={"audio": (audio, stt_config.sample_rate)}, |
| 2115 | ) |
| 2116 | |
| 2117 | @classmethod |
| 2118 | def _get_default_prompt_tokens(cls, language: str) -> tuple[str, ...]: |
| 2119 | # Use token-level control tags so fast tokenizers do not have to parse |
| 2120 | # the raw string form of the decoder prefix. |
| 2121 | return ( |
| 2122 | "▁", |
| 2123 | "<|startofcontext|>", |
| 2124 | "<|startoftranscript|>", |
| 2125 | "<|emo:undefined|>", |
| 2126 | f"<|{language}|>", |
| 2127 | f"<|{language}|>", |
| 2128 | "<|pnc|>", |
| 2129 | "<|noitn|>", |
| 2130 | "<|notimestamp|>", |
| 2131 | "<|nodiarize|>", |
| 2132 | ) |
| 2133 | |
| 2134 | @classmethod |
| 2135 | def _get_default_prompt_token_ids( |
| 2136 | cls, |
| 2137 | tokenizer: Any, |
| 2138 | model_config: ModelConfig, |
| 2139 | language: str, |
| 2140 | ) -> list[int]: |
| 2141 | cache_key = ( |
| 2142 | getattr(model_config, "tokenizer", None), |
| 2143 | getattr(model_config, "tokenizer_revision", None), |
| 2144 | language, |
| 2145 | ) |
| 2146 | prompt_token_ids = cls._default_prompt_token_ids_cache.get(cache_key) |
| 2147 | if prompt_token_ids is None: |
| 2148 | prompt_tokens = list(cls._get_default_prompt_tokens(language)) |
| 2149 | token_ids = tokenizer.convert_tokens_to_ids(prompt_tokens) |
| 2150 | if not isinstance(token_ids, list): |
| 2151 | token_ids = [token_ids] |
| 2152 | unk_token_id = getattr(tokenizer, "unk_token_id", None) |
| 2153 | if unk_token_id is not None and any( |
| 2154 | token_id == unk_token_id for token_id in token_ids |
| 2155 | ): |
| 2156 | raise ValueError( |
| 2157 | "Failed to resolve the CohereASR decoder control tokens " |
| 2158 | "with the configured tokenizer." |
| 2159 | ) |
| 2160 | prompt_token_ids = tuple(int(token_id) for token_id in token_ids) |
| 2161 | cls._default_prompt_token_ids_cache[cache_key] = prompt_token_ids |
| 2162 | |
| 2163 | return list(prompt_token_ids) |
| 2164 | |
| 2165 | @classmethod |
| 2166 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 2167 | # Required as part of SupportsMultiModal interface. |
| 2168 | if modality.startswith("audio"): |
| 2169 | return None |
| 2170 | |
| 2171 | raise ValueError("Only audio modality is supported") |
| 2172 | |
| 2173 | @classmethod |
| 2174 | def get_speech_to_text_config( |
| 2175 | cls, model_config: ModelConfig, task_type: str |
| 2176 | ) -> SpeechToTextConfig: |
| 2177 | sampling_rate = model_config.hf_config.sample_rate |
| 2178 | assert sampling_rate == 16000 |
| 2179 | max_audio_clip_s = model_config.hf_config.max_audio_clip_s |
| 2180 | overlap_chunk_second = model_config.hf_config.overlap_chunk_second |
| 2181 | |
| 2182 | return SpeechToTextConfig( |
| 2183 | max_audio_clip_s=max_audio_clip_s, |
| 2184 | overlap_chunk_second=overlap_chunk_second, |
| 2185 | sample_rate=sampling_rate, |
| 2186 | ) |
| 2187 | |
| 2188 | @classmethod |
| 2189 | def get_num_audio_tokens( |
| 2190 | cls, |
| 2191 | audio_duration_s: float, |
| 2192 | stt_config: SpeechToTextConfig, |
| 2193 | model_config: ModelConfig, |
| 2194 | ) -> int | None: |
| 2195 | hf_config = model_config.hf_config |
| 2196 | preprocessor = hf_config.preprocessor |
| 2197 | sample_rate = preprocessor["sample_rate"] |
| 2198 | window_stride = preprocessor["window_stride"] |
| 2199 | hop_length = int(window_stride * sample_rate) |
| 2200 | # Floor-divide to match get_seq_len. |
| 2201 | num_frames = int(audio_duration_s * sample_rate) // hop_length |
| 2202 | subsampling_factor = hf_config.encoder["subsampling_factor"] |
| 2203 | return math.ceil(num_frames / subsampling_factor) |
| 2204 | |
| 2205 | def get_num_encoder_cross_attn_tokens(self, num_encoder_input_tokens: int) -> int: |
| 2206 | return self.model.encoder.get_num_encoder_cross_attn_tokens( |
| 2207 | num_encoder_input_tokens |
| 2208 | ) |
| 2209 | |
| 2210 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 2211 | super().__init__() |
| 2212 | config = vllm_config.model_config.hf_config |
| 2213 | quant_config = vllm_config.quant_config |
| 2214 | self.config = config |
| 2215 | self.dtype = vllm_config.model_config.dtype |
| 2216 | |
| 2217 | with self._mark_composite_model( |
| 2218 | vllm_config, |
| 2219 | language_targets=CohereASRDecoder, |
| 2220 | tower_targets={"audio": (ConformerEncoder, CohereASRProjector)}, |
| 2221 | ): |
| 2222 | self.model = CohereASRModel(vllm_config=vllm_config, prefix=prefix) |
| 2223 | |
| 2224 | head_config = config.head |
| 2225 | |
| 2226 | self.proj_out = ParallelLMHead( |
| 2227 | head_config["num_classes"], |
| 2228 | head_config["hidden_size"], |
| 2229 | quant_config=quant_config, |
| 2230 | bias=True, |
| 2231 | ) # NOTE: bias is True |
| 2232 | |
| 2233 | logit_scale = getattr(head_config, "logit_scale", 1.0) |
| 2234 | self.logits_processor = LogitsProcessor( |
| 2235 | head_config["num_classes"], scale=logit_scale |
| 2236 | ) |
| 2237 | |
| 2238 | def forward( |
| 2239 | self, |
| 2240 | input_ids: torch.Tensor, |
| 2241 | positions: torch.Tensor, |
| 2242 | encoder_outputs: list[torch.Tensor] | None = None, |
| 2243 | **kwargs, |
| 2244 | ) -> torch.Tensor: |
| 2245 | if encoder_outputs is None: |
| 2246 | encoder_outputs = [] |
| 2247 | decoder_outputs = self.model( |
| 2248 | input_ids=input_ids, |
| 2249 | positions=positions, |
| 2250 | encoder_outputs=encoder_outputs, |
| 2251 | ) |
| 2252 | |
| 2253 | return decoder_outputs |
| 2254 | |
| 2255 | def get_language_model(self) -> torch.nn.Module: |
| 2256 | # Required as part of SupportsMultiModal interface. |
| 2257 | return self.model.decoder |
| 2258 | |
| 2259 | def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: |
| 2260 | # Required as part of SupportsMultiModal interface. |
| 2261 | audio_input, seq_lens = self._parse_and_validate_audio_input(**kwargs) |
| 2262 | |
| 2263 | if hasattr(audio_input, "input_features"): |
| 2264 | out = self.model.get_encoder_outputs( |
| 2265 | audio_input["input_features"], seq_lens |
| 2266 | ) |
| 2267 | else: |
| 2268 | out = self.model.get_encoder_outputs(audio_input, seq_lens) |
| 2269 | |
| 2270 | return out |
| 2271 | |
| 2272 | def _parse_and_validate_audio_input( |
| 2273 | self, **kwargs: object |
| 2274 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 2275 | input_features = kwargs.pop("input_features", None) |
| 2276 | length = kwargs.pop("length", None) |
| 2277 | |
| 2278 | if input_features is None: |
| 2279 | raise ValueError("Audio features are required for CohereASR model.") |
| 2280 | |
| 2281 | if not isinstance(input_features, (torch.Tensor, list)): |
| 2282 | raise ValueError( |
| 2283 | f"Incorrect type of audio features. Got type: {type(input_features)}" |
| 2284 | ) |
| 2285 | |
| 2286 | if not isinstance(length, torch.Tensor): |
| 2287 | raise ValueError("Audio feature lengths must be a tensor.") |
| 2288 | |
| 2289 | if isinstance(input_features, torch.Tensor): |
| 2290 | seq_lens = length.reshape(-1) |
| 2291 | else: |
| 2292 | if not is_list_of(input_features, torch.Tensor, check="all"): |
| 2293 | raise ValueError("All audio features must be tensors.") |
| 2294 | input_features = [ |
| 2295 | feature.to(self.dtype).squeeze(0).transpose(1, 0) |
| 2296 | for feature in input_features |
| 2297 | ] |
| 2298 | seq_lens = length.reshape(-1) |
| 2299 | input_features = torch.nn.utils.rnn.pad_sequence( |
| 2300 | input_features, batch_first=True, padding_value=0.0 |
| 2301 | ) |
| 2302 | input_features = input_features.transpose(1, 2) |
| 2303 | |
| 2304 | return input_features, seq_lens |
| 2305 | |
| 2306 | def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 2307 | logits = self.logits_processor(self.proj_out, hidden_states, self.proj_out.bias) |
| 2308 | return logits |
| 2309 | |
| 2310 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 2311 | def transform(inputs): |
| 2312 | name, loaded_weight = inputs |
| 2313 | |
| 2314 | if name.startswith("transf_decoder._decoder"): |
| 2315 | name = name.replace("transf_decoder._decoder", "decoder") |
| 2316 | if name.startswith("transf_decoder._embedding"): |
| 2317 | name = name.replace("transf_decoder._embedding", "decoder.embedding") |
| 2318 | if "second_sub_layer.query_net" in name: |
| 2319 | name = name.replace( |
| 2320 | "second_sub_layer.query_net", "second_sub_layer.q_proj" |
| 2321 | ) |
| 2322 | |
| 2323 | if name in ["log_softmax.mlp.layer0.weight", "log_softmax.mlp.layer0.bias"]: |
| 2324 | name = name.replace("log_softmax.mlp.layer0", "proj_out") |
| 2325 | else: |
| 2326 | name = "model." + name |
| 2327 | |
| 2328 | return name, loaded_weight |
| 2329 | |
| 2330 | loader = AutoWeightsLoader(self) |
| 2331 | |
| 2332 | return loader.load_weights( |
| 2333 | map(transform, weights), mapper=self.hf_to_vllm_mapper |
| 2334 | ) |