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