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 7, 07:45 UTC
[model] support FireRedASR2 (#35727)
Signed-off-by: zixiao <shunli.dsl@alibaba-inc.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: zixiao <shunli.dsl@alibaba-inc.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
AllenDou committed
Mar 4, 2026 at 11:41 UTC
c1d963403c4f09cc0d5a25573c45d7405cd09abb
7 files changed
+1183
docs/models/supported_models.md
+1
@@ -793,6 +793,7 @@ Speech2Text models trained specifically for Automatic Speech Recognition.
793
794
| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
795
|--------------|--------|-------------------|----------------------|---------------------------|
796
+| `FireRedASR2ForConditionalGeneration` | FireRedASR2 | `allendou/FireRedASR2-LLM-vllm`, etc. | | |
797
| `FunASRForConditionalGeneration` | FunASR | `allendou/Fun-ASR-Nano-2512-vllm`, etc. | | |
798
| `Gemma3nForConditionalGeneration` | Gemma3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
799
| `GlmAsrForConditionalGeneration` | GLM-ASR | `zai-org/GLM-ASR-Nano-2512` | ✅︎ | ✅︎ |
requirements/common.txt
+1
@@ -57,3 +57,4 @@ opentelemetry-sdk >= 1.27.0
57
opentelemetry-api >= 1.27.0
58
opentelemetry-exporter-otlp >= 1.27.0
59
opentelemetry-semantic-conventions-ai >= 0.4.1
60
+kaldi-native-fbank >= 1.18.7
tests/models/registry.py
+3
@@ -743,6 +743,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {
743
"baidu/ERNIE-4.5-VL-28B-A3B-PT",
744
trust_remote_code=True,
745
),
746
+ "FireRedASR2ForConditionalGeneration": _HfExamplesInfo(
747
+ "allendou/FireRedASR2-LLM-vllm",
748
+ ),
749
"FunASRForConditionalGeneration": _HfExamplesInfo(
750
"allendou/Fun-ASR-Nano-2512-vllm",
751
),
vllm/model_executor/models/fireredasr2.py
new
+829
@@ -0,0 +1,829 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import math
4
+from collections.abc import Iterable, Mapping, Sequence
5
+from typing import Annotated, Literal, cast
6
+
7
+import numpy as np
8
+import torch
9
+import torch.nn.functional as F
10
+from torch import nn
11
+from transformers import (
12
+ BatchFeature,
13
+ Qwen2Config,
14
+)
15
+
16
+from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
17
+from vllm.config.multimodal import BaseDummyOptions
18
+from vllm.inputs.data import PromptType
19
+from vllm.logger import init_logger
20
+from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY
21
+from vllm.model_executor.layers.linear import (
22
+ ReplicatedLinear,
23
+)
24
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
25
+from vllm.model_executor.models.whisper_utils import (
26
+ ISO639_1_SUPPORTED_LANGS,
27
+)
28
+from vllm.multimodal import MULTIMODAL_REGISTRY
29
+from vllm.multimodal.inputs import (
30
+ MultiModalDataDict,
31
+ MultiModalFieldConfig,
32
+ MultiModalKwargsItems,
33
+)
34
+from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
35
+from vllm.multimodal.processing import (
36
+ BaseDummyInputsBuilder,
37
+ BaseMultiModalProcessor,
38
+ BaseProcessingInfo,
39
+ PromptReplacement,
40
+ PromptUpdate,
41
+ PromptUpdateDetails,
42
+)
43
+from vllm.transformers_utils.processor import cached_processor_from_config
44
+from vllm.transformers_utils.processors.fireredasr2_processor import (
45
+ FireRedASR2FeatureExtractor,
46
+)
47
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
48
+
49
+from .interfaces import (
50
+ MultiModalEmbeddings,
51
+ SupportsMultiModal,
52
+ SupportsTranscription,
53
+ _require_is_multimodal,
54
+)
55
+from .qwen2 import Qwen2ForCausalLM
56
+from .utils import (
57
+ AutoWeightsLoader,
58
+ WeightsMapper,
59
+ _merge_multimodal_embeddings,
60
+ maybe_prefix,
61
+)
62
+
63
+logger = init_logger(__name__)
64
+
65
+
66
+class FireRedASR2AudioInputs(TensorSchema):
67
+ """
68
+ Dimensions:
69
+ - b: Batch size
70
+ - nmb: Number of mel bins
71
+ - t: Time frames (M)
72
+ """
73
+
74
+ input_features: Annotated[
75
+ list[torch.Tensor] | None,
76
+ TensorShape("b", "nmb", "t"),
77
+ ]
78
+ speech_lengths: Annotated[
79
+ list[torch.Tensor] | None,
80
+ TensorShape("b"),
81
+ ]
82
+ fake_token_lengths: Annotated[
83
+ list[torch.Tensor] | None,
84
+ TensorShape("b"),
85
+ ]
86
+
87
+
88
+class Swish(nn.Module):
89
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
90
+ return x * torch.sigmoid(x)
91
+
92
+
93
+class Conv2dSubsampling(nn.Module):
94
+ def __init__(self, idim: int, d_model: int, out_channels: int = 32):
95
+ super().__init__()
96
+ self.conv = nn.Sequential(
97
+ nn.Conv2d(1, out_channels, 3, 2),
98
+ nn.ReLU(),
99
+ nn.Conv2d(out_channels, out_channels, 3, 2),
100
+ nn.ReLU(),
101
+ )
102
+ subsample_idim = ((idim - 1) // 2 - 1) // 2
103
+ self.out = ReplicatedLinear(
104
+ input_size=out_channels * subsample_idim,
105
+ output_size=d_model,
106
+ bias=True,
107
+ )
108
+
109
+ self.subsampling = 4
110
+ left_context = right_context = 3 # both exclude currect frame
111
+ self.context = left_context + 1 + right_context # 7
112
+
113
+ def forward(
114
+ self, x: torch.Tensor, x_mask: torch.Tensor
115
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
116
+ x = x.unsqueeze(1)
117
+ x = self.conv(x)
118
+ N, C, T, D = x.size()
119
+ x, _ = self.out(x.transpose(1, 2).contiguous().view(N, T, C * D))
120
+ mask = x_mask[:, :, :-2:2][:, :, :-2:2]
121
+ input_lengths = mask[:, -1, :].sum(dim=-1)
122
+ return x, input_lengths, mask
123
+
124
+
125
+class RelPositionalEncoding(nn.Module):
126
+ def __init__(self, d_model: int, max_len: int = 5000):
127
+ super().__init__()
128
+ pe_positive = torch.zeros(max_len, d_model, requires_grad=False)
129
+ pe_negative = torch.zeros(max_len, d_model, requires_grad=False)
130
+ position = torch.arange(0, max_len).unsqueeze(1).float()
131
+ div_term = torch.exp(
132
+ torch.arange(0, d_model, 2).float()
133
+ * -(torch.log(torch.tensor(10000.0)).item() / d_model)
134
+ )
135
+ pe_positive[:, 0::2] = torch.sin(position * div_term)
136
+ pe_positive[:, 1::2] = torch.cos(position * div_term)
137
+ pe_negative[:, 0::2] = torch.sin(-1 * position * div_term)
138
+ pe_negative[:, 1::2] = torch.cos(-1 * position * div_term)
139
+
140
+ pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0)
141
+ pe_negative = pe_negative[1:].unsqueeze(0)
142
+ self.pe = torch.cat([pe_positive, pe_negative], dim=1)
143
+
144
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
145
+ # Tmax = 2 * max_len - 1
146
+ Tmax, T = self.pe.size(1), x.size(1)
147
+ pos_emb = self.pe[:, Tmax // 2 - T + 1 : Tmax // 2 + T].clone().detach()
148
+ return pos_emb
149
+
150
+
151
+class ConformerFeedForward(nn.Module):
152
+ def __init__(self, d_model: int):
153
+ super().__init__()
154
+ self.pre_layer_norm = nn.LayerNorm(d_model)
155
+ self.linear_expand = ReplicatedLinear(
156
+ input_size=d_model,
157
+ output_size=d_model * 4,
158
+ bias=True,
159
+ )
160
+ self.nonlinear = Swish()
161
+ self.linear_project = ReplicatedLinear(
162
+ input_size=d_model * 4,
163
+ output_size=d_model,
164
+ bias=True,
165
+ )
166
+
167
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
168
+ residual = x
169
+ x = self.pre_layer_norm(x)
170
+ x, _ = self.linear_expand(x)
171
+ x = self.nonlinear(x)
172
+ x, _ = self.linear_project(x)
173
+ output = x + residual
174
+ return output
175
+
176
+
177
+class EncoderMultiHeadAttention(nn.Module):
178
+ def __init__(self, n_head: int, d_model: int):
179
+ super().__init__()
180
+ assert d_model % n_head == 0
181
+ self.n_head = n_head
182
+ self.d_k = d_model // n_head
183
+ self.d_v = self.d_k
184
+
185
+ self.w_qs = ReplicatedLinear(
186
+ input_size=d_model, output_size=n_head * self.d_k, bias=False
187
+ )
188
+ self.w_ks = ReplicatedLinear(
189
+ input_size=d_model, output_size=n_head * self.d_k, bias=False
190
+ )
191
+ self.w_vs = ReplicatedLinear(
192
+ input_size=d_model, output_size=n_head * self.d_v, bias=False
193
+ )
194
+
195
+ self.layer_norm_q = nn.LayerNorm(d_model)
196
+ self.layer_norm_k = nn.LayerNorm(d_model)
197
+ self.layer_norm_v = nn.LayerNorm(d_model)
198
+
199
+ self.fc = ReplicatedLinear(
200
+ input_size=n_head * self.d_v, output_size=d_model, bias=False
201
+ )
202
+
203
+ def forward_qkv(
204
+ self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor
205
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
206
+ d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
207
+ sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)
208
+
209
+ q = self.layer_norm_q(q)
210
+ k = self.layer_norm_k(k)
211
+ v = self.layer_norm_v(v)
212
+
213
+ q = self.w_qs(q)[0].view(sz_b, len_q, n_head, d_k)
214
+ k = self.w_ks(k)[0].view(sz_b, len_k, n_head, d_k)
215
+ v = self.w_vs(v)[0].view(sz_b, len_v, n_head, d_v)
216
+ q = q.transpose(1, 2)
217
+ k = k.transpose(1, 2)
218
+ v = v.transpose(1, 2)
219
+ return q, k, v
220
+
221
+ def forward_output(
222
+ self, output: torch.Tensor, residual: torch.Tensor, sz_b: int, len_q: int
223
+ ) -> torch.Tensor:
224
+ output = output.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
225
+ fc_out, _ = self.fc(output)
226
+ output = fc_out
227
+ output = output + residual
228
+ return output
229
+
230
+ def forward_attention(
231
+ self, attn: torch.Tensor, v: torch.Tensor, mask: torch.Tensor | None = None
232
+ ) -> tuple[torch.Tensor, torch.Tensor]:
233
+ if mask is not None:
234
+ mask = mask.unsqueeze(1)
235
+ mask = mask.eq(0)
236
+ attn = attn.masked_fill(mask, -float("inf"))
237
+ attn = torch.softmax(attn, dim=-1).masked_fill(mask, 0.0)
238
+ else:
239
+ attn = torch.softmax(attn, dim=-1)
240
+
241
+ d_attn = attn
242
+ output = torch.matmul(d_attn, v)
243
+
244
+ return output, attn
245
+
246
+
247
+class RelPosMultiHeadAttention(EncoderMultiHeadAttention):
248
+ def __init__(self, n_head: int, d_model: int):
249
+ super().__init__(n_head, d_model)
250
+ d_k = d_model // n_head
251
+ self.scale = 1.0 / (d_k**0.5)
252
+ self.linear_pos = ReplicatedLinear(
253
+ input_size=d_model, output_size=n_head * d_k, bias=False
254
+ )
255
+ self.pos_bias_u = nn.Parameter(torch.empty([n_head, d_k]))
256
+ self.pos_bias_v = nn.Parameter(torch.empty([n_head, d_k]))
257
+
258
+ def _rel_shift(self, x):
259
+ N, H, T1, T2 = x.size()
260
+ zero_pad = torch.zeros((N, H, T1, 1), device=x.device, dtype=x.dtype)
261
+ x_padded = torch.cat([zero_pad, x], dim=-1)
262
+
263
+ x_padded = x_padded.view(N, H, T2 + 1, T1)
264
+ x = x_padded[:, :, 1:].view_as(x)
265
+ x = x[:, :, :, : x.size(-1) // 2 + 1]
266
+ return x
267
+
268
+ def forward(
269
+ self,
270
+ q: torch.Tensor,
271
+ k: torch.Tensor,
272
+ v: torch.Tensor,
273
+ pos_emb: torch.Tensor,
274
+ mask: torch.Tensor | None = None,
275
+ ) -> tuple[torch.Tensor, torch.Tensor]:
276
+ sz_b, len_q = q.size(0), q.size(1)
277
+
278
+ residual = q
279
+ q, k, v = self.forward_qkv(q, k, v)
280
+
281
+ q = q.transpose(1, 2)
282
+ n_batch_pos = pos_emb.size(0)
283
+ p = self.linear_pos(pos_emb)[0].view(n_batch_pos, -1, self.n_head, self.d_k)
284
+ p = p.transpose(1, 2)
285
+
286
+ q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2)
287
+ q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2)
288
+
289
+ matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
290
+
291
+ matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
292
+ matrix_bd = self._rel_shift(matrix_bd)
293
+
294
+ attn_scores = matrix_ac + matrix_bd
295
+ attn_scores.mul_(self.scale)
296
+
297
+ output, attn = self.forward_attention(attn_scores, v, mask=mask)
298
+
299
+ output = self.forward_output(output, residual, sz_b, len_q)
300
+ return output, attn
301
+
302
+
303
+class ConformerConvolution(nn.Module):
304
+ def __init__(self, d_model: int, kernel_size: int = 33):
305
+ super().__init__()
306
+ assert kernel_size % 2 == 1
307
+ self.pre_layer_norm = nn.LayerNorm(d_model)
308
+ self.pointwise_conv1 = nn.Conv1d(
309
+ d_model, d_model * 4, kernel_size=1, bias=False
310
+ )
311
+ self.padding = (kernel_size - 1) // 2
312
+ self.depthwise_conv = nn.Conv1d(
313
+ d_model * 2,
314
+ d_model * 2,
315
+ kernel_size,
316
+ stride=1,
317
+ padding=self.padding,
318
+ groups=d_model * 2,
319
+ bias=False,
320
+ )
321
+ self.batch_norm = nn.LayerNorm(d_model * 2)
322
+ self.swish = Swish()
323
+ self.pointwise_conv2 = nn.Conv1d(
324
+ d_model * 2, d_model, kernel_size=1, bias=False
325
+ )
326
+
327
+ def forward(
328
+ self, x: torch.Tensor, mask: torch.Tensor | None = None
329
+ ) -> torch.Tensor:
330
+ residual = x
331
+ out = self.pre_layer_norm(x)
332
+ out = out.transpose(1, 2)
333
+ if mask is not None:
334
+ out.masked_fill_(mask.ne(1), 0.0)
335
+ out = self.pointwise_conv1(out)
336
+ out = F.glu(out, dim=1)
337
+ out = self.depthwise_conv(out)
338
+
339
+ out = out.transpose(1, 2)
340
+ out = self.swish(self.batch_norm(out))
341
+ out = out.transpose(1, 2)
342
+
343
+ out = self.pointwise_conv2(out)
344
+ if mask is not None:
345
+ out.masked_fill_(mask.ne(1), 0.0)
346
+ out = out.transpose(1, 2)
347
+ return out + residual
348
+
349
+
350
+class RelPosEmbConformerBlock(nn.Module):
351
+ def __init__(self, d_model, n_head, kernel_size=33):
352
+ super().__init__()
353
+ self.ffn1 = ConformerFeedForward(d_model)
354
+ self.mhsa = RelPosMultiHeadAttention(n_head, d_model)
355
+ self.conv = ConformerConvolution(d_model, kernel_size)
356
+ self.ffn2 = ConformerFeedForward(d_model)
357
+ self.layer_norm = nn.LayerNorm(d_model)
358
+
359
+ def forward(
360
+ self,
361
+ x: torch.Tensor,
362
+ pos_emb: torch.Tensor,
363
+ slf_attn_mask: torch.Tensor | None = None,
364
+ pad_mask: torch.Tensor | None = None,
365
+ ) -> torch.Tensor:
366
+ out = 0.5 * x + 0.5 * self.ffn1(x)
367
+ out = self.mhsa(out, out, out, pos_emb, mask=slf_attn_mask)[0]
368
+ out = self.conv(out, pad_mask)
369
+ out = 0.5 * out + 0.5 * self.ffn2(out)
370
+ out = self.layer_norm(out)
371
+ return out
372
+
373
+
374
+class ConformerEncoder(nn.Module):
375
+ def __init__(
376
+ self,
377
+ idim: int,
378
+ n_layers_enc: int,
379
+ n_head: int,
380
+ d_model: int,
381
+ kernel_size: int = 33,
382
+ pe_maxlen: int = 5000,
383
+ ):
384
+ super().__init__()
385
+ self.odim = d_model
386
+
387
+ self.input_preprocessor = Conv2dSubsampling(idim, d_model)
388
+ self.positional_encoding = RelPositionalEncoding(d_model)
389
+
390
+ self.layer_stack = nn.ModuleList()
391
+ for _ in range(n_layers_enc):
392
+ block = RelPosEmbConformerBlock(d_model, n_head, kernel_size)
393
+ self.layer_stack.append(block)
394
+
395
+ def forward(
396
+ self, padded_input: torch.Tensor, input_lengths: torch.Tensor, pad: bool = True
397
+ ):
398
+ if pad:
399
+ padded_input = F.pad(
400
+ padded_input,
401
+ (0, 0, 0, self.input_preprocessor.context - 1),
402
+ "constant",
403
+ 0.0,
404
+ )
405
+ src_mask = self.padding_position_is_0(padded_input, input_lengths)
406
+
407
+ embed_output, input_lengths, src_mask = self.input_preprocessor(
408
+ padded_input, src_mask
409
+ )
410
+ enc_output = embed_output
411
+
412
+ pos_emb = self.positional_encoding(embed_output)
413
+
414
+ enc_outputs = []
415
+ for enc_layer in self.layer_stack:
416
+ enc_output = enc_layer(
417
+ enc_output, pos_emb, slf_attn_mask=src_mask, pad_mask=src_mask
418
+ )
419
+ enc_outputs.append(enc_output)
420
+
421
+ return enc_output, input_lengths, src_mask
422
+
423
+ def padding_position_is_0(
424
+ self, padded_input: torch.Tensor, input_lengths: torch.Tensor
425
+ ) -> torch.Tensor:
426
+ N, T = padded_input.size()[:2]
427
+ mask = torch.ones((N, T)).to(padded_input.device)
428
+ for i in range(N):
429
+ mask[i, input_lengths[i] :] = 0
430
+ mask = mask.unsqueeze(dim=1)
431
+ return mask.to(torch.uint8)
432
+
433
+
434
+class FireRedASR2Adapter(nn.Module):
435
+ def __init__(self, encoder_dim: int, llm_dim: int, downsample_rate: int = 2):
436
+ super().__init__()
437
+ self.ds = downsample_rate
438
+ self.linear1 = ReplicatedLinear(
439
+ input_size=encoder_dim * downsample_rate,
440
+ output_size=llm_dim,
441
+ bias=True,
442
+ )
443
+ self.relu = _ACTIVATION_REGISTRY["relu"]
444
+ self.linear2 = ReplicatedLinear(
445
+ input_size=llm_dim,
446
+ output_size=llm_dim,
447
+ bias=True,
448
+ )
449
+
450
+ def forward(self, x, x_lens):
451
+ batch_size, seq_len, feat_dim = x.size()
452
+ num_frames_to_discard = seq_len % self.ds
453
+ if num_frames_to_discard > 0:
454
+ x = x[:, :-num_frames_to_discard, :]
455
+ seq_len = x.size(1)
456
+
457
+ x = x.contiguous()
458
+ x = x.view(batch_size, seq_len // self.ds, feat_dim * self.ds)
459
+
460
+ x, _ = self.linear1(x)
461
+ x = self.relu(x)
462
+ x, _ = self.linear2(x)
463
+
464
+ new_x_lens = torch.clamp(x_lens, max=seq_len) // self.ds
465
+ return x, new_x_lens
466
+
467
+
468
+class FireRedASR2Encoder(nn.Module):
469
+ def __init__(
470
+ self,
471
+ *,
472
+ vllm_config: VllmConfig,
473
+ ):
474
+ super().__init__()
475
+ self.audio_encoder = ConformerEncoder(
476
+ **vllm_config.model_config.hf_config.audio_encoder_conf
477
+ )
478
+
479
+
480
+class FireRedASR2Model(nn.Module):
481
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
482
+ super().__init__()
483
+ self.encoder = FireRedASR2Encoder(
484
+ vllm_config=vllm_config,
485
+ )
486
+ encoder_dim = self.encoder.audio_encoder.odim
487
+ llm_dim = vllm_config.model_config.hf_config.hidden_size
488
+ self.encoder_projector = FireRedASR2Adapter(
489
+ encoder_dim,
490
+ llm_dim,
491
+ vllm_config.model_config.hf_config.encoder_downsample_rate,
492
+ )
493
+
494
+ self.decoder = Qwen2ForCausalLM(
495
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "decoder")
496
+ )
497
+
498
+ def forward(
499
+ self,
500
+ input_ids: torch.Tensor | None,
501
+ positions: torch.Tensor,
502
+ inputs_embeds: torch.Tensor | None = None,
503
+ ) -> torch.Tensor:
504
+ decoder_outputs = self.decoder(
505
+ input_ids=input_ids,
506
+ positions=positions,
507
+ inputs_embeds=inputs_embeds,
508
+ )
509
+ return decoder_outputs
510
+
511
+ def get_encoder_outputs(
512
+ self,
513
+ speech: torch.Tensor | list[torch.Tensor] | None,
514
+ speech_lengths: torch.Tensor | list[torch.Tensor] | None,
515
+ ) -> torch.Tensor | None:
516
+ encoder_outs, enc_lengths, enc_mask = self.encoder.audio_encoder(
517
+ speech, speech_lengths
518
+ )
519
+ speech_features, speech_lens = self.encoder_projector(encoder_outs, enc_lengths)
520
+ return speech_features
521
+
522
+
523
+class FireRedASR2ProcessingInfo(BaseProcessingInfo):
524
+ def get_hf_config(self) -> Qwen2Config:
525
+ return self.ctx.get_hf_config(Qwen2Config)
526
+
527
+ def get_supported_mm_limits(self) -> Mapping[str, int | None]:
528
+ return {"audio": 1}
529
+
530
+ def get_feature_extractor(self, **kwargs: object) -> FireRedASR2FeatureExtractor:
531
+ hf_processor = self.get_hf_processor(**kwargs)
532
+ feature_extractor = hf_processor.feature_extractor # type: ignore
533
+ assert isinstance(feature_extractor, FireRedASR2FeatureExtractor)
534
+ return feature_extractor
535
+
536
+ def get_data_parser(self) -> MultiModalDataParser:
537
+ feature_extractor = self.get_feature_extractor()
538
+ return MultiModalDataParser(
539
+ target_sr=feature_extractor.sampling_rate,
540
+ target_channels=self.get_target_channels(),
541
+ )
542
+
543
+ def get_target_channels(self) -> int:
544
+ return 1
545
+
546
+
547
+class FireRedASR2DummyInputsBuilder(BaseDummyInputsBuilder[FireRedASR2ProcessingInfo]):
548
+ def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
549
+ num_audios = mm_counts.get("audio", 0)
550
+
551
+ return "<|AUDIO|>" * num_audios
552
+
553
+ def get_dummy_mm_data(
554
+ self,
555
+ seq_len: int,
556
+ mm_counts: Mapping[str, int],
557
+ mm_options: Mapping[str, BaseDummyOptions],
558
+ ) -> MultiModalDataDict:
559
+ feature_extractor = self.info.get_feature_extractor()
560
+
561
+ sampling_rate = feature_extractor.sampling_rate
562
+ audio_len = feature_extractor.chunk_length * sampling_rate
563
+ num_audios = mm_counts.get("audio", 0)
564
+
565
+ audio_overrides = mm_options.get("audio")
566
+
567
+ ret = {
568
+ "audio": self._get_dummy_audios(
569
+ length=audio_len, num_audios=num_audios, overrides=audio_overrides
570
+ )
571
+ }
572
+ return ret
573
+
574
+
575
+class FireRedASR2MultiModalProcessor(
576
+ BaseMultiModalProcessor[FireRedASR2ProcessingInfo]
577
+):
578
+ def _call_hf_processor(
579
+ self,
580
+ prompt: str,
581
+ mm_data: Mapping[str, object],
582
+ mm_kwargs: Mapping[str, object],
583
+ tok_kwargs: Mapping[str, object],
584
+ ) -> BatchFeature:
585
+ if mm_data:
586
+ feature_extractor = self.info.get_feature_extractor(**mm_kwargs)
587
+ mm_data = dict(audio=mm_data.pop("audios"))
588
+ mm_kwargs = dict(
589
+ **mm_kwargs,
590
+ sampling_rate=feature_extractor.sampling_rate,
591
+ )
592
+ processed_outputs = super()._call_hf_processor(
593
+ prompt=prompt,
594
+ mm_data=mm_data,
595
+ mm_kwargs=mm_kwargs,
596
+ tok_kwargs=tok_kwargs,
597
+ )
598
+ if "labels" in processed_outputs:
599
+ processed_outputs["input_ids"] = processed_outputs.pop("labels")
600
+ return processed_outputs
601
+
602
+ def _get_mm_fields_config(
603
+ self,
604
+ hf_inputs: BatchFeature,
605
+ hf_processor_mm_kwargs: Mapping[str, object],
606
+ ) -> Mapping[str, MultiModalFieldConfig]:
607
+ return dict(
608
+ input_features=MultiModalFieldConfig.batched("audio"),
609
+ speech_lengths=MultiModalFieldConfig.batched("audio"),
610
+ fake_token_lengths=MultiModalFieldConfig.batched("audio"),
611
+ )
612
+
613
+ def _get_prompt_updates(
614
+ self,
615
+ mm_items: MultiModalDataItems,
616
+ hf_processor_mm_kwargs: Mapping[str, object],
617
+ out_mm_kwargs: MultiModalKwargsItems,
618
+ ) -> Sequence[PromptUpdate]:
619
+ processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
620
+ tokenizer = self.info.get_tokenizer()
621
+ vocab = tokenizer.get_vocab()
622
+
623
+ audio_token = getattr(processor, "audio_token", "<|AUDIO|>")
624
+
625
+ audio_token_id = vocab[audio_token]
626
+
627
+ out_mm_data = out_mm_kwargs.get_data()
628
+
629
+ fake_token_lengths = out_mm_data.get("fake_token_lengths")
630
+
631
+ if fake_token_lengths is None:
632
+ audio_output_lengths = []
633
+ else:
634
+ assert isinstance(fake_token_lengths, torch.Tensor)
635
+
636
+ audio_output_lengths = fake_token_lengths.tolist()
637
+
638
+ def get_replacement_fireredasr2_audio(item_idx: int):
639
+ num_features = audio_output_lengths[item_idx]
640
+
641
+ audio_tokens = [audio_token_id] * int(num_features)
642
+
643
+ return PromptUpdateDetails.select_token_id(
644
+ audio_tokens,
645
+ embed_token_id=audio_token_id,
646
+ )
647
+
648
+ return [
649
+ PromptReplacement(
650
+ modality="audio",
651
+ target=[audio_token_id],
652
+ replacement=get_replacement_fireredasr2_audio,
653
+ )
654
+ ]
655
+
656
+
657
+@MULTIMODAL_REGISTRY.register_processor(
658
+ FireRedASR2MultiModalProcessor,
659
+ info=FireRedASR2ProcessingInfo,
660
+ dummy_inputs=FireRedASR2DummyInputsBuilder,
661
+)
662
+class FireRedASR2ForConditionalGeneration(
663
+ nn.Module, SupportsTranscription, SupportsMultiModal
664
+):
665
+ packed_modules_mapping = {
666
+ "self_attn.qkv_proj": [
667
+ "self_attn.q_proj",
668
+ "self_attn.k_proj",
669
+ "self_attn.v_proj",
670
+ ],
671
+ "encoder_attn.kv_proj": ["encoder_attn.k_proj", "encoder_attn.v_proj"],
672
+ }
673
+
674
+ hf_to_vllm_mapper = WeightsMapper(
675
+ orig_to_new_substr={
676
+ "llm.": "model.decoder.",
677
+ "encoder.": "model.encoder.audio_encoder.",
678
+ "encoder_projector.": "model.encoder_projector.",
679
+ "net.0": "pre_layer_norm",
680
+ "net.1": "linear_expand",
681
+ "net.4": "linear_project",
682
+ }
683
+ )
684
+
685
+ supports_transcription_only = True
686
+ supports_segment_timestamp = True
687
+ supported_languages = ISO639_1_SUPPORTED_LANGS
688
+
689
+ @classmethod
690
+ def validate_language(cls, language: str | None) -> str | None:
691
+ if language is None:
692
+ # TODO language should be optional and can be guessed.
693
+ # For now we default to en. See
694
+ # https://github.com/huggingface/transformers/blob/main/src/transformers/models/whisper/generation_whisper.py#L1520
695
+ logger.warning(
696
+ "Defaulting to language='en'. If you wish to transcribe "
697
+ "audio in a different language, pass the `language` field "
698
+ "in the TranscriptionRequest."
699
+ )
700
+ language = "en"
701
+ return super().validate_language(language)
702
+
703
+ @classmethod
704
+ def get_generation_prompt(
705
+ cls,
706
+ audio: np.ndarray,
707
+ model_config: ModelConfig, # not needed here
708
+ stt_config: SpeechToTextConfig,
709
+ language: str | None,
710
+ task_type: Literal["transcribe", "translate"],
711
+ request_prompt: str,
712
+ to_language: str | None,
713
+ ) -> PromptType:
714
+ if language is None:
715
+ raise ValueError(
716
+ "Language must be specified when creating the fireredasr2 prompt"
717
+ )
718
+
719
+ prompt_str = "<|im_start|>user\n<|AUDIO|>请转写音频为文字<|im_end|>\n<|im_start|>assistant\n" # noqa: E501
720
+ prompt = {
721
+ "prompt": prompt_str,
722
+ "multi_modal_data": {
723
+ "audio": (audio, stt_config.sample_rate),
724
+ },
725
+ }
726
+ return cast(PromptType, prompt)
727
+
728
+ @classmethod
729
+ def get_speech_to_text_config(
730
+ cls, model_config: ModelConfig, task_type: str
731
+ ) -> SpeechToTextConfig:
732
+ processor = cached_processor_from_config(model_config)
733
+
734
+ return SpeechToTextConfig(
735
+ max_audio_clip_s=processor.feature_extractor.chunk_length,
736
+ sample_rate=processor.feature_extractor.sampling_rate,
737
+ )
738
+
739
+ @classmethod
740
+ def get_num_audio_tokens(
741
+ cls,
742
+ audio_duration_s: float,
743
+ stt_config: SpeechToTextConfig,
744
+ model_config: ModelConfig,
745
+ ) -> int | None:
746
+ processor = cached_processor_from_config(model_config)
747
+ hop_length = processor.feature_extractor.hop_length
748
+ assert hop_length is not None
749
+ return math.ceil(audio_duration_s * stt_config.sample_rate / hop_length)
750
+
751
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
752
+ super().__init__()
753
+ config = vllm_config.model_config.hf_config
754
+ self.config = config
755
+ self.dtype = vllm_config.model_config.dtype
756
+
757
+ self.model = FireRedASR2Model(
758
+ vllm_config=vllm_config,
759
+ prefix=maybe_prefix(prefix, "model"),
760
+ )
761
+ logit_scale = getattr(config, "logit_scale", 1.0)
762
+
763
+ self.logits_processor = LogitsProcessor(config.vocab_size, scale=logit_scale)
764
+
765
+ def forward(
766
+ self,
767
+ input_ids: torch.Tensor,
768
+ positions: torch.Tensor,
769
+ inputs_embeds: torch.Tensor | None = None,
770
+ **kwargs,
771
+ ) -> torch.Tensor:
772
+ decoder_outputs = self.model(
773
+ input_ids=input_ids,
774
+ positions=positions,
775
+ inputs_embeds=inputs_embeds,
776
+ )
777
+ return decoder_outputs
778
+
779
+ def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
780
+ audio_input = self._parse_and_validate_audio_input(**kwargs)
781
+
782
+ speech = audio_input["input_features"]
783
+ speech_lengths = audio_input["speech_lengths"].to(torch.int32)
784
+ enc_output = self.model.get_encoder_outputs(
785
+ speech=speech, speech_lengths=speech_lengths
786
+ )
787
+
788
+ return enc_output
789
+
790
+ def embed_input_ids(
791
+ self,
792
+ input_ids: torch.Tensor,
793
+ multimodal_embeddings: MultiModalEmbeddings | None = None,
794
+ *,
795
+ is_multimodal: torch.Tensor | None = None,
796
+ handle_oov_mm_token: bool = False,
797
+ ) -> torch.Tensor:
798
+ inputs_embeds = self.model.decoder.embed_input_ids(input_ids)
799
+
800
+ ret = _merge_multimodal_embeddings(
801
+ inputs_embeds=inputs_embeds,
802
+ multimodal_embeddings=multimodal_embeddings,
803
+ is_multimodal=_require_is_multimodal(is_multimodal),
804
+ )
805
+ return ret
806
+
807
+ def _parse_and_validate_audio_input(
808
+ self, **kwargs: object
809
+ ) -> FireRedASR2AudioInputs:
810
+ input_features = kwargs.pop("input_features", None)
811
+ speech_lengths = kwargs.pop("speech_lengths", None)
812
+ fake_token_lengths = kwargs.pop("fake_token_lengths", None)
813
+
814
+ return FireRedASR2AudioInputs(
815
+ input_features=input_features,
816
+ speech_lengths=speech_lengths,
817
+ fake_token_lengths=fake_token_lengths,
818
+ )
819
+
820
+ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
821
+ logits = self.logits_processor(self.model.decoder.lm_head, hidden_states)
822
+ return logits
823
+
824
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
825
+ loader = AutoWeightsLoader(
826
+ self, skip_prefixes=["model.encoder.audio_encoder.positional_encoding.pe"]
827
+ )
828
+
829
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/registry.py
+4
@@ -341,6 +341,10 @@ _MULTIMODAL_MODELS = {
341
"ernie45_vl",
342
"Ernie4_5_VLMoeForConditionalGeneration",
343
),
344
+ "FireRedASR2ForConditionalGeneration": (
345
+ "fireredasr2",
346
+ "FireRedASR2ForConditionalGeneration",
347
+ ),
348
"FunASRForConditionalGeneration": ("funasr", "FunASRForConditionalGeneration"), # noqa: E501
349
"FunAudioChatForConditionalGeneration": (
350
"funaudiochat",
vllm/transformers_utils/processors/__init__.py
+4
@@ -10,6 +10,9 @@ reasons:
10
11
from vllm.transformers_utils.processors.bagel import BagelProcessor
12
from vllm.transformers_utils.processors.deepseek_vl2 import DeepseekVLV2Processor
13
+from vllm.transformers_utils.processors.fireredasr2_processor import (
14
+ FireRedASR2Processor,
15
+)
16
from vllm.transformers_utils.processors.funasr_processor import FunASRProcessor
17
from vllm.transformers_utils.processors.hunyuan_vl import HunYuanVLProcessor
18
from vllm.transformers_utils.processors.hunyuan_vl_image import HunYuanVLImageProcessor
@@ -19,6 +22,7 @@ from vllm.transformers_utils.processors.ovis2_5 import Ovis2_5Processor
22
__all__ = [
23
"BagelProcessor",
24
"DeepseekVLV2Processor",
25
+ "FireRedASR2Processor",
26
"FunASRProcessor",
27
"HunYuanVLProcessor",
28
"HunYuanVLImageProcessor",
vllm/transformers_utils/processors/fireredasr2_processor.py
new
+341
@@ -0,0 +1,341 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+
5
+import kaldi_native_fbank as knf
6
+import numpy as np
7
+import torch
8
+import torch.nn.functional as F
9
+from transformers import (
10
+ AutoFeatureExtractor,
11
+ AutoProcessor,
12
+ BatchFeature,
13
+)
14
+from transformers.feature_extraction_sequence_utils import SequenceFeatureExtractor
15
+from transformers.processing_utils import ProcessorMixin
16
+from transformers.utils import TensorType
17
+
18
+from vllm.logger import init_logger
19
+
20
+logger = init_logger(__name__)
21
+
22
+
23
+class CMVN:
24
+ def __init__(self, dim, means, inverse_std_variences):
25
+ self.dim, self.means, self.inverse_std_variences = (
26
+ dim,
27
+ np.array(means),
28
+ np.array(inverse_std_variences),
29
+ )
30
+
31
+ def __call__(self, x):
32
+ assert x.shape[-1] == self.dim, "CMVN dim mismatch"
33
+ out = x - self.means
34
+ out = out * self.inverse_std_variences
35
+ return out
36
+
37
+
38
+class KaldifeatFbank:
39
+ def __init__(self, num_mel_bins=80, frame_length=25, frame_shift=10, dither=1.0):
40
+ self.dither = dither
41
+ opts = knf.FbankOptions()
42
+ opts.frame_opts.dither = dither
43
+ opts.mel_opts.num_bins = num_mel_bins
44
+ opts.frame_opts.snip_edges = True
45
+ opts.mel_opts.debug_mel = False
46
+ self.opts = opts
47
+
48
+ def __call__(self, sample_rate, wav_np, is_train=False):
49
+ dither = self.dither if is_train else 0.0
50
+ self.opts.frame_opts.dither = dither
51
+ fbank = knf.OnlineFbank(self.opts)
52
+
53
+ fbank.accept_waveform(sample_rate, wav_np.tolist())
54
+ feat = []
55
+ for i in range(fbank.num_frames_ready):
56
+ feat.append(fbank.get_frame(i))
57
+ if len(feat) == 0:
58
+ print("Check data, len(feat) == 0", wav_np, flush=True)
59
+ return np.zeros((0, self.opts.mel_opts.num_bins))
60
+ feat = np.vstack(feat)
61
+ return feat
62
+
63
+
64
+class FireRedASR2FeatureExtractor(SequenceFeatureExtractor):
65
+ r"""
66
+ Constructs a FireRedASR2 feature extractor.
67
+
68
+ This feature extractor inherits from [`~feature_extraction_sequence_
69
+ utils.SequenceFeatureExtractor`] which contains most of the main
70
+ methods. Users should refer to this superclass for more information
71
+ regarding those methods.
72
+
73
+ This class extracts mel-filter bank features from raw speech using a custom
74
+ numpy implementation of the `Short Time Fourier Transform` which should
75
+ match pytorch's `torch.stft` equivalent.
76
+
77
+ Args:
78
+ feature_size (`int`, *optional*, defaults to 80):
79
+ The feature dimension of the extracted features.
80
+ sampling_rate (`int`, *optional*, defaults to 16000):
81
+ The sampling rate at which the audio files should be digitalized
82
+ expressed in hertz (Hz).
83
+ chunk_length (`int`, *optional*, defaults to 30):
84
+ The maximum number of chunks of `sampling_rate` samples used to
85
+ trim and pad longer or shorter audio sequences.
86
+ padding_value (`float`, *optional*, defaults to 0.0):
87
+ Padding value used to pad the audio. Should correspond to silences.
88
+ dither (`float`, *optional*, defaults to 0.0):
89
+ Adds dithering. In other words, adds a small Gaussian noise to each frame.
90
+ E.g. use 0.0001 to add dithering with a normal distribution centered
91
+ around 0.0 with standard deviation 0.0001 (assuming [-1,+1] range
92
+ of raw_speech). The value 0.0 means no dithering.
93
+ Dithering has similar effect as `spectrogram(mel_floor=...)`. It reduces
94
+ the high log_mel_fbank values for signals with hard-zero sections,
95
+ when VAD cutoff is present in the signal.
96
+ """
97
+
98
+ model_input_names = ["input_features"]
99
+
100
+ def __init__(
101
+ self,
102
+ feature_size=80,
103
+ sampling_rate=16000,
104
+ chunk_length=30,
105
+ padding_value=0.0,
106
+ return_attention_mask=False,
107
+ dim=80,
108
+ means=None,
109
+ inverse_std_variences=None,
110
+ num_mel_bins=80,
111
+ frame_length=25,
112
+ frame_shift=10,
113
+ dither=0.0,
114
+ max_length=3000,
115
+ downsample_rate=2,
116
+ left_context=3,
117
+ right_context=3,
118
+ **kwargs,
119
+ ):
120
+ super().__init__(
121
+ feature_size=feature_size,
122
+ sampling_rate=sampling_rate,
123
+ padding_value=padding_value,
124
+ return_attention_mask=return_attention_mask,
125
+ **kwargs,
126
+ )
127
+ self.chunk_length = chunk_length
128
+ self.max_length = max_length
129
+ self.dim = dim
130
+ self.means = means
131
+ self.inverse_std_variences = inverse_std_variences
132
+ self.num_mel_bins = num_mel_bins
133
+ self.frame_length = frame_length
134
+ self.frame_shift = frame_shift
135
+ self.dither = dither
136
+ self.sampling_rate = sampling_rate
137
+ self.downsample_rate = downsample_rate
138
+ self.context = left_context + 1 + right_context
139
+
140
+ def __call__(
141
+ self,
142
+ raw_speech: np.ndarray | list[float] | list[np.ndarray] | list[list[float]],
143
+ truncation: bool = True,
144
+ pad_to_multiple_of: int | None = None,
145
+ return_tensors: str | TensorType | None = None,
146
+ return_attention_mask: bool | None = None,
147
+ padding: str | None = "max_length",
148
+ max_length: int | None = None,
149
+ sampling_rate: int | None = None,
150
+ do_normalize: bool | None = None,
151
+ **kwargs,
152
+ ) -> BatchFeature:
153
+ if sampling_rate != self.sampling_rate:
154
+ raise ValueError(
155
+ f"The model corresponding to this feature extractor: "
156
+ f"{self.__class__.__name__} was trained using a sampling "
157
+ f"rate of {self.sampling_rate}. Please make sure that the "
158
+ f"provided `raw_speech` input was sampled with "
159
+ f"{self.sampling_rate} and not {sampling_rate}."
160
+ )
161
+
162
+ def padding_position_is_0(padded_input, input_lengths):
163
+ N, T = padded_input.size()[:2]
164
+ mask = torch.ones((N, T)).to(padded_input.device)
165
+ for i in range(N):
166
+ mask[i, input_lengths[i] :] = 0
167
+ mask = mask.unsqueeze(dim=1)
168
+ return mask.to(torch.uint8)
169
+
170
+ # initialize the CMVN and Fbank objects
171
+ self.cmvn = CMVN(self.dim, self.means, self.inverse_std_variences)
172
+ self.fbank = KaldifeatFbank(
173
+ num_mel_bins=self.num_mel_bins,
174
+ frame_length=self.frame_length,
175
+ frame_shift=self.frame_shift,
176
+ dither=self.dither,
177
+ )
178
+
179
+ feats = []
180
+ speech_lengths = []
181
+ fake_token_lengths = []
182
+ for speech in raw_speech:
183
+ """
184
+ We must multiply by 32768 here because FireRedASR2 loads audio data
185
+ using kaldiio.load_mat, while vLLM loads audio data using librosa.
186
+ """
187
+ speech = speech * 32768
188
+ fbank = self.fbank(sampling_rate, speech)
189
+ fbank = self.cmvn(fbank)
190
+ fbank = torch.from_numpy(fbank).float()
191
+ length = fbank.size(0)
192
+ feats.append(fbank)
193
+ speech_lengths.append(length)
194
+ padded_input2 = fbank
195
+ padded_input2 = F.pad(
196
+ padded_input2, (0, 0, 0, self.context - 1), "constant", 0.0
197
+ )
198
+ src_mask = padding_position_is_0(
199
+ padded_input2[None, :, :], torch.tensor([length], dtype=torch.int32)
200
+ )
201
+ x_mask = src_mask
202
+ mask = x_mask[:, :, :-2:2][:, :, :-2:2]
203
+ input_lengths = mask[:, -1, :].sum(dim=-1)
204
+ input_lengths = input_lengths // self.downsample_rate
205
+ fake_token_len = torch.clamp(input_lengths, min=1)
206
+ fake_token_lengths.append(fake_token_len)
207
+
208
+ feats = torch.stack(feats, dim=0)
209
+ batched_speech = self.pad(
210
+ BatchFeature({"input_features": feats}),
211
+ padding=padding,
212
+ max_length=max_length if max_length else self.max_length,
213
+ truncation=truncation,
214
+ pad_to_multiple_of=pad_to_multiple_of,
215
+ return_attention_mask=return_attention_mask or do_normalize,
216
+ )
217
+
218
+ if return_tensors is not None:
219
+ batched_speech = batched_speech.convert_to_tensors(return_tensors)
220
+
221
+ batched_speech["speech_lengths"] = torch.tensor(speech_lengths)
222
+ batched_speech["fake_token_lengths"] = torch.concat(fake_token_lengths)
223
+ return batched_speech
224
+
225
+
226
+class FireRedASR2Processor(ProcessorMixin):
227
+ r"""
228
+ Constructs a FireRedASR2 processor which wraps a FireRedASR2 feature extractor and
229
+ a FireRedASR2 tokenizer into a single processor.
230
+
231
+ [`FireRedASR2Processor`] offers all the functionalities of
232
+ [`FireRedASR2FeatureExtractor`] and [`Qwen2Tokenizer`]. See the
233
+ [`~FireRedASR2Processor.__call__`] and [`~FireRedASR2Processor.decode`] for more
234
+ information.
235
+
236
+ Args:
237
+ feature_extractor (`FireRedASR2FeatureExtractor`): An instance of
238
+ [`FireRedASR2FeatureExtractor`].
239
+ The feature extractor is a required input.
240
+ tokenizer (`Qwen2Tokenizer`):
241
+ An instance of [`Qwen2Tokenizer`]. The tokenizer is a required
242
+ input.
243
+ """
244
+
245
+ feature_extractor_class = "FireRedASR2FeatureExtractor"
246
+ tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")
247
+
248
+ def __init__(
249
+ self,
250
+ feature_extractor,
251
+ tokenizer,
252
+ audio_token="<|AUDIO|>",
253
+ ):
254
+ super().__init__(feature_extractor, tokenizer)
255
+ self.current_processor = self.feature_extractor
256
+ self._in_target_context_manager = False
257
+ self.audio_token = (
258
+ tokenizer.audio_token if hasattr(tokenizer, "audio_token") else audio_token
259
+ )
260
+ self.audio_token_id = tokenizer.convert_tokens_to_ids(self.audio_token)
261
+
262
+ def get_decoder_prompt_ids(self, task=None, language=None, no_timestamps=True):
263
+ return self.tokenizer.get_decoder_prompt_ids(
264
+ task=task, language=language, no_timestamps=no_timestamps
265
+ )
266
+
267
+ def __call__(self, *args, **kwargs):
268
+ """
269
+ Forwards the `audio` argument to FireRedASR2FeatureExtractor's
270
+ [`~FireRedASR2FeatureExtractor.__call__`] and the `text` argument to
271
+ [`~Qwen2Tokenizer.__call__`]. Please refer to the docstring of the
272
+ above two methods for more information.
273
+ """
274
+ if self._in_target_context_manager:
275
+ return self.current_processor(*args, **kwargs)
276
+
277
+ audio = kwargs.pop("audio", None)
278
+ sampling_rate = kwargs.pop("sampling_rate", None)
279
+ text = kwargs.pop("text", None)
280
+ if len(args) > 0:
281
+ audio = args[0]
282
+ args = args[1:]
283
+
284
+ if text is None:
285
+ raise ValueError("You need to specify `text` input to process.")
286
+ elif isinstance(text, str):
287
+ text = [text]
288
+ elif not isinstance(text, list) and not isinstance(text[0], str):
289
+ raise ValueError(
290
+ "Invalid input text. Please provide a string, or a list of strings"
291
+ )
292
+
293
+ if audio is not None:
294
+ # ensure we have as much audios as audio tokens
295
+ num_audio_tokens = sum(sample.count(self.audio_token) for sample in text)
296
+ num_audios = 1 if type(audio) is np.ndarray else len(audio)
297
+ if num_audio_tokens != num_audios:
298
+ raise ValueError(
299
+ f"Found {num_audio_tokens} {self.audio_token} token{'s' if num_audio_tokens > 1 else ''} in provided text but received {num_audios} audio{'s' if num_audios > 1 else ''}" # noqa: E501
300
+ )
301
+ inputs = self.feature_extractor(
302
+ audio, *args, sampling_rate=sampling_rate, **kwargs
303
+ )
304
+
305
+ expanded_text = []
306
+ for sample in text:
307
+ replace_str = []
308
+ while self.audio_token in sample:
309
+ num_audio_tokens = int(inputs["fake_token_lengths"].item())
310
+
311
+ expanded_audio_token = self.audio_token * num_audio_tokens
312
+
313
+ replace_str.append(expanded_audio_token)
314
+ sample = sample.replace(self.audio_token, "<placeholder>", 1)
315
+
316
+ while "<placeholder>" in sample:
317
+ sample = sample.replace("<placeholder>", replace_str.pop(0), 1)
318
+ expanded_text.append(sample)
319
+ text = expanded_text
320
+
321
+ if text is not None:
322
+ encodings = self.tokenizer(text, **kwargs)
323
+
324
+ if text is None:
325
+ return inputs
326
+
327
+ elif audio is None:
328
+ return encodings
329
+ else:
330
+ inputs["labels"] = encodings["input_ids"]
331
+
332
+ return inputs
333
+
334
+ def get_prompt_ids(self, text: str, return_tensors="np"):
335
+ return self.tokenizer.get_prompt_ids(text, return_tensors=return_tensors)
336
+
337
+
338
+AutoFeatureExtractor.register(
339
+ "FireRedASR2FeatureExtractor", FireRedASR2FeatureExtractor
340
+)
341
+AutoProcessor.register("FireRedASR2Processor", FireRedASR2Processor)