@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 21dfb842d7
Fork 0

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 8, 12:30 UTC

[model] support FunASR model (#33247)

Signed-off-by: zixiao <shunli.dsl@alibaba-inc.com> Co-authored-by: zixiao <shunli.dsl@alibaba-inc.com>

AllenDou committed Feb 11, 2026 at 15:37 UTC 21dfb842d76c61204d44f6f1dd1e99f55a9b2cf4
7 files changed +1585 -3
docs/models/supported_models.md
+1
@@ -790,6 +790,7 @@ Speech2Text models trained specifically for Automatic Speech Recognition.
790
791 | Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
792 |--------------|--------|-------------------|----------------------|---------------------------|
793 +| `FunASRForConditionalGeneration` | FunASR | `allendou/Fun-ASR-Nano-2512-vllm`, etc. | | |
794 | `Gemma3nForConditionalGeneration` | Gemma3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
795 | `GlmAsrForConditionalGeneration` | GLM-ASR | `zai-org/GLM-ASR-Nano-2512` | ✅︎ | ✅︎ |
796 | `GraniteSpeechForConditionalGeneration` | Granite Speech | `ibm-granite/granite-speech-3.3-2b`, `ibm-granite/granite-speech-3.3-8b`, etc. | ✅︎ | ✅︎ |
examples/online_serving/openai_transcription_client.py
+16 -3
@@ -26,7 +26,9 @@ from openai import AsyncOpenAI, OpenAI
26 from vllm.assets.audio import AudioAsset
27
28
29 -def sync_openai(audio_path: str, client: OpenAI, model: str):
29 +def sync_openai(
30 + audio_path: str, client: OpenAI, model: str, *, repetition_penalty: float = 1.3
31 +):
32 """
33 Perform synchronous transcription using OpenAI-compatible API.
34 """
@@ -40,7 +42,7 @@ def sync_openai(audio_path: str, client: OpenAI, model: str):
42 # Additional sampling params not provided by OpenAI API.
43 extra_body=dict(
44 seed=4419,
43 - repetition_penalty=1.3,
45 + repetition_penalty=repetition_penalty,
46 ),
47 )
48 print("transcription result [sync]:", transcription.text)
@@ -129,7 +131,12 @@ def main(args):
131 print(f"Using model: {model}")
132
133 # Run the synchronous function
132 - sync_openai(args.audio_path if args.audio_path else mary_had_lamb, client, model)
134 + sync_openai(
135 + audio_path=args.audio_path if args.audio_path else mary_had_lamb,
136 + client=client,
137 + model=model,
138 + repetition_penalty=args.repetition_penalty,
139 + )
140
141 # Run the asynchronous function
142 if "openai" in model:
@@ -161,5 +168,11 @@ if __name__ == "__main__":
168 default=None,
169 help="The path to the audio file to transcribe.",
170 )
171 + parser.add_argument(
172 + "--repetition_penalty",
173 + type=float,
174 + default=1.3,
175 + help="repetition penalty",
176 + )
177 args = parser.parse_args()
178 main(args)
tests/models/registry.py
+4
@@ -713,6 +713,10 @@ _MULTIMODAL_EXAMPLE_MODELS = {
713 "baidu/ERNIE-4.5-VL-28B-A3B-PT",
714 trust_remote_code=True,
715 ),
716 + "FunASRForConditionalGeneration": _HfExamplesInfo(
717 + "allendou/Fun-ASR-Nano-2512-vllm",
718 + is_available_online=False,
719 + ),
720 "FunAudioChatForConditionalGeneration": _HfExamplesInfo(
721 "funaudiochat", is_available_online=False
722 ),
vllm/model_executor/models/funasr.py new
+1057
@@ -0,0 +1,1057 @@
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 Annotated, 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 (
13 + BatchFeature,
14 + Qwen3Config,
15 +)
16 +
17 +from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
18 +from vllm.config.multimodal import BaseDummyOptions
19 +from vllm.distributed import get_tensor_model_parallel_world_size
20 +from vllm.inputs.data import PromptType
21 +from vllm.logger import init_logger
22 +from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY
23 +from vllm.model_executor.layers.attention.mm_encoder_attention import (
24 + MMEncoderAttention,
25 +)
26 +from vllm.model_executor.layers.linear import (
27 + ColumnParallelLinear,
28 + QKVParallelLinear,
29 + ReplicatedLinear,
30 + RowParallelLinear,
31 +)
32 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
33 +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
34 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
35 +from vllm.model_executor.models.whisper_utils import (
36 + ISO639_1_SUPPORTED_LANGS,
37 +)
38 +from vllm.multimodal import MULTIMODAL_REGISTRY
39 +from vllm.multimodal.inputs import (
40 + MultiModalDataDict,
41 + MultiModalFieldConfig,
42 + MultiModalKwargsItems,
43 +)
44 +from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
45 +from vllm.multimodal.processing import (
46 + BaseDummyInputsBuilder,
47 + BaseMultiModalProcessor,
48 + BaseProcessingInfo,
49 + PromptReplacement,
50 + PromptUpdate,
51 + PromptUpdateDetails,
52 +)
53 +from vllm.transformers_utils.processor import cached_processor_from_config
54 +from vllm.transformers_utils.processors.funasr_processor import FunASRFeatureExtractor
55 +from vllm.utils.jsontree import json_map_leaves
56 +from vllm.utils.tensor_schema import TensorSchema, TensorShape
57 +
58 +from .interfaces import (
59 + MultiModalEmbeddings,
60 + SupportsMultiModal,
61 + SupportsTranscription,
62 + _require_is_multimodal,
63 +)
64 +from .qwen3 import Qwen3Model
65 +from .utils import (
66 + AutoWeightsLoader,
67 + WeightsMapper,
68 + _merge_multimodal_embeddings,
69 + maybe_prefix,
70 +)
71 +
72 +logger = init_logger(__name__)
73 +
74 +
75 +def sequence_mask(lengths, maxlen=None, dtype=torch.float32, device=None):
76 + if maxlen is None:
77 + maxlen = lengths.max()
78 + row_vector = torch.arange(0, maxlen, 1).to(lengths.device)
79 + matrix = torch.unsqueeze(lengths, dim=-1)
80 + mask = row_vector < matrix
81 + mask = mask.detach()
82 +
83 + return mask.type(dtype).to(device) if device is not None else mask.type(dtype)
84 +
85 +
86 +class LayerNorm(torch.nn.LayerNorm):
87 + def __init__(self, nout, dim=-1):
88 + super().__init__(nout, eps=1e-12)
89 + self.dim = dim
90 +
91 + def forward(self, x: torch.Tensor):
92 + if self.dim == -1:
93 + return super().forward(x)
94 + return super().forward(x.transpose(self.dim, -1)).transpose(self.dim, -1)
95 +
96 +
97 +class EncoderLayerSANM(nn.Module):
98 + def __init__(
99 + self,
100 + in_size: int,
101 + size: int,
102 + self_attn: nn.Module,
103 + feed_forward: nn.Module,
104 + normalize_before=True,
105 + ):
106 + super().__init__()
107 + self.self_attn = self_attn
108 + self.feed_forward = feed_forward
109 + self.norm1 = LayerNorm(in_size)
110 + self.norm2 = LayerNorm(size)
111 + self.in_size = in_size
112 + self.size = size
113 + self.normalize_before = normalize_before
114 +
115 + def forward(
116 + self,
117 + hidden_states: torch.Tensor,
118 + mask: torch.Tensor | None = None,
119 + cache=None,
120 + mask_shfit_chunk=None,
121 + mask_att_chunk_encoder=None,
122 + ):
123 + residual = hidden_states
124 + hidden_states = self.norm1(hidden_states)
125 +
126 + if self.in_size == self.size:
127 + hidden_states = residual + self.self_attn(
128 + hidden_states,
129 + mask,
130 + mask_shfit_chunk=mask_shfit_chunk,
131 + mask_att_chunk_encoder=mask_att_chunk_encoder,
132 + )
133 + else:
134 + hidden_states = self.self_attn(
135 + hidden_states,
136 + mask,
137 + mask_shfit_chunk=mask_shfit_chunk,
138 + mask_att_chunk_encoder=mask_att_chunk_encoder,
139 + )
140 +
141 + residual = hidden_states
142 + hidden_states = self.norm2(hidden_states)
143 + hidden_states = residual + self.feed_forward(hidden_states)
144 +
145 + return hidden_states, mask, cache, mask_shfit_chunk, mask_att_chunk_encoder
146 +
147 +
148 +class MultiHeadedAttentionSANM(nn.Module):
149 + def __init__(
150 + self,
151 + n_head: int,
152 + in_feat: int,
153 + n_feat: int,
154 + kernel_size: int,
155 + sanm_shift: int = 0,
156 + ):
157 + super().__init__()
158 + assert n_feat % n_head == 0
159 + # We assume d_v always equals d_k
160 + self.d_k = n_feat // n_head
161 + self.h = n_head
162 + self.out_proj = ReplicatedLinear(
163 + input_size=n_feat,
164 + output_size=n_feat,
165 + bias=True,
166 + )
167 + self.linear_q_k_v = ReplicatedLinear(
168 + input_size=in_feat,
169 + output_size=n_feat * 3,
170 + bias=True,
171 + )
172 + self.attn = None
173 +
174 + self.fsmn_block = nn.Conv1d(
175 + n_feat, n_feat, kernel_size, stride=1, padding=0, groups=n_feat, bias=False
176 + )
177 + # padding
178 + left_padding = (kernel_size - 1) // 2
179 + if sanm_shift > 0:
180 + left_padding = left_padding + sanm_shift
181 + right_padding = kernel_size - 1 - left_padding
182 + self.pad_fn = nn.ConstantPad1d((left_padding, right_padding), 0.0)
183 +
184 + def forward_fsmn(
185 + self,
186 + inputs: torch.Tensor,
187 + mask: torch.Tensor,
188 + mask_shfit_chunk: torch.Tensor = None,
189 + ):
190 + b, t, d = inputs.size()
191 + if mask is not None:
192 + mask = torch.reshape(mask, (b, -1, 1))
193 + if mask_shfit_chunk is not None:
194 + mask = mask * mask_shfit_chunk
195 + inputs = inputs * mask
196 +
197 + x = inputs.transpose(1, 2)
198 + x = self.pad_fn(x)
199 + x = self.fsmn_block(x)
200 + x = x.transpose(1, 2)
201 + x += inputs
202 + if mask is not None:
203 + x = x * mask
204 + return x
205 +
206 + def forward_qkv(self, x: torch.Tensor):
207 + b, t, d = x.size()
208 +
209 + q_k_v, _ = self.linear_q_k_v(x)
210 + q, k, v = torch.split(q_k_v, int(self.h * self.d_k), dim=-1)
211 + q_h = torch.reshape(q, (b, t, self.h, self.d_k)).transpose(1, 2)
212 + k_h = torch.reshape(k, (b, t, self.h, self.d_k)).transpose(1, 2)
213 + v_h = torch.reshape(v, (b, t, self.h, self.d_k)).transpose(1, 2)
214 +
215 + return q_h, k_h, v_h, v
216 +
217 + def forward_attention(
218 + self,
219 + value: torch.Tensor,
220 + scores: torch.Tensor,
221 + mask: torch.Tensor,
222 + mask_att_chunk_encoder: torch.Tensor = None,
223 + ):
224 + n_batch = value.size(0)
225 + if mask is not None:
226 + if mask_att_chunk_encoder is not None:
227 + mask = mask * mask_att_chunk_encoder
228 +
229 + mask = mask.unsqueeze(1).eq(0)
230 +
231 + min_value = -float("inf")
232 + scores = scores.masked_fill(mask, min_value)
233 + attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0)
234 + else:
235 + attn = torch.softmax(scores, dim=-1)
236 +
237 + p_attn = attn
238 + x = torch.matmul(p_attn, value)
239 + x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)
240 +
241 + out, _ = self.out_proj(x)
242 + return out
243 +
244 + def forward(
245 + self,
246 + hidden_states: torch.Tensor,
247 + mask: torch.Tensor,
248 + mask_shfit_chunk: torch.Tensor = None,
249 + mask_att_chunk_encoder: torch.Tensor = None,
250 + ):
251 + q_h, k_h, v_h, v = self.forward_qkv(hidden_states)
252 + fsmn_memory = self.forward_fsmn(v, mask, mask_shfit_chunk)
253 + q_h = q_h * self.d_k ** (-0.5)
254 + scores = torch.matmul(q_h, k_h.transpose(-2, -1))
255 + att_outs = self.forward_attention(v_h, scores, mask, mask_att_chunk_encoder)
256 + return att_outs + fsmn_memory
257 +
258 +
259 +class SinusoidalPositionEncoder(torch.nn.Module):
260 + def __init__(self, d_model=80):
261 + super().__init__()
262 +
263 + def encode(
264 + self,
265 + positions: torch.Tensor = None,
266 + depth: int = None,
267 + dtype: torch.dtype = torch.float32,
268 + ):
269 + batch_size = positions.size(0)
270 + positions = positions.type(dtype)
271 + device = positions.device
272 + log_timescale_increment = torch.log(
273 + torch.tensor([10000], dtype=dtype, device=device)
274 + ) / (depth / 2 - 1)
275 + inv_timescales = torch.exp(
276 + torch.arange(depth / 2, device=device).type(dtype)
277 + * (-log_timescale_increment)
278 + )
279 + inv_timescales = torch.reshape(inv_timescales, [batch_size, -1])
280 + scaled_time = torch.reshape(positions, [1, -1, 1]) * torch.reshape(
281 + inv_timescales, [1, 1, -1]
282 + )
283 + encoding = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=2)
284 + return encoding.type(dtype)
285 +
286 + def forward(self, hidden_states: torch.Tensor):
287 + batch_size, timesteps, input_dim = hidden_states.size()
288 + positions = torch.arange(1, timesteps + 1, device=hidden_states.device)[None, :]
289 + position_encoding = self.encode(positions, input_dim, hidden_states.dtype).to(
290 + hidden_states.device
291 + )
292 +
293 + return hidden_states + position_encoding
294 +
295 +
296 +class SenseVoiceEncoderSmall(nn.Module):
297 + def __init__(
298 + self,
299 + input_size: int,
300 + output_size: int = 256,
301 + attention_heads: int = 4,
302 + linear_units: int = 2048,
303 + num_blocks: int = 6,
304 + tp_blocks: int = 0,
305 + attention_dropout_rate: float = 0.0,
306 + normalize_before: bool = True,
307 + kernel_size: int = 11,
308 + sanm_shift: int = 0,
309 + **kwargs,
310 + ):
311 + super().__init__()
312 + self._output_size = output_size
313 + self.embed = SinusoidalPositionEncoder()
314 +
315 + self.normalize_before = normalize_before
316 +
317 + positionwise_layer = PositionwiseFeedForward
318 + positionwise_layer_args = (
319 + output_size,
320 + linear_units,
321 + )
322 +
323 + encoder_selfattn_layer = MultiHeadedAttentionSANM
324 + encoder_selfattn_layer_args0 = (
325 + attention_heads,
326 + input_size,
327 + output_size,
328 + kernel_size,
329 + sanm_shift,
330 + )
331 + encoder_selfattn_layer_args = (
332 + attention_heads,
333 + output_size,
334 + output_size,
335 + kernel_size,
336 + sanm_shift,
337 + )
338 +
339 + self.encoders0 = nn.ModuleList(
340 + [
341 + EncoderLayerSANM(
342 + input_size,
343 + output_size,
344 + encoder_selfattn_layer(*encoder_selfattn_layer_args0),
345 + positionwise_layer(*positionwise_layer_args),
346 + )
347 + for i in range(1)
348 + ]
349 + )
350 + self.encoders = nn.ModuleList(
351 + [
352 + EncoderLayerSANM(
353 + output_size,
354 + output_size,
355 + encoder_selfattn_layer(*encoder_selfattn_layer_args),
356 + positionwise_layer(*positionwise_layer_args),
357 + )
358 + for i in range(num_blocks - 1)
359 + ]
360 + )
361 +
362 + self.tp_encoders = nn.ModuleList(
363 + [
364 + EncoderLayerSANM(
365 + output_size,
366 + output_size,
367 + encoder_selfattn_layer(*encoder_selfattn_layer_args),
368 + positionwise_layer(*positionwise_layer_args),
369 + )
370 + for i in range(tp_blocks)
371 + ]
372 + )
373 +
374 + self.after_norm = LayerNorm(output_size)
375 +
376 + self.tp_norm = LayerNorm(output_size)
377 +
378 + def output_size(self) -> int:
379 + return self._output_size
380 +
381 + def forward(
382 + self,
383 + xs_pad: torch.Tensor,
384 + ilens: torch.Tensor,
385 + ):
386 + maxlen = xs_pad.shape[1]
387 + masks = sequence_mask(
388 + ilens, maxlen=maxlen, dtype=ilens.dtype, device=ilens.device
389 + )[:, None, :]
390 +
391 + xs_pad *= self.output_size() ** 0.5
392 +
393 + xs_pad = self.embed(xs_pad)
394 +
395 + for layer_idx, encoder_layer in enumerate(self.encoders0):
396 + encoder_outs = encoder_layer(xs_pad, masks)
397 + xs_pad, masks = encoder_outs[0], encoder_outs[1]
398 +
399 + for layer_idx, encoder_layer in enumerate(self.encoders):
400 + encoder_outs = encoder_layer(xs_pad, masks)
401 + xs_pad, masks = encoder_outs[0], encoder_outs[1]
402 +
403 + xs_pad = self.after_norm(xs_pad)
404 +
405 + olens = masks.squeeze(1).sum(1).int()
406 +
407 + for layer_idx, encoder_layer in enumerate(self.tp_encoders):
408 + encoder_outs = encoder_layer(xs_pad, masks)
409 + xs_pad, masks = encoder_outs[0], encoder_outs[1]
410 +
411 + xs_pad = self.tp_norm(xs_pad)
412 + return xs_pad, olens
413 +
414 +
415 +class PositionwiseFeedForward(nn.Module):
416 + def __init__(self, idim: int, hidden_units: int):
417 + super().__init__()
418 + self.w_1 = ColumnParallelLinear(
419 + input_size=idim,
420 + output_size=hidden_units,
421 + bias=True,
422 + )
423 + self.w_2 = RowParallelLinear(
424 + input_size=hidden_units,
425 + output_size=idim,
426 + bias=True,
427 + )
428 + self.activation = _ACTIVATION_REGISTRY["relu"]
429 +
430 + def forward(self, hidden_states: torch.Tensor):
431 + hidden_states, _ = self.w_1(hidden_states)
432 + hidden_states = self.activation(hidden_states)
433 + hidden_states, _ = self.w_2(hidden_states)
434 + return hidden_states
435 +
436 +
437 +class EncoderLayer(nn.Module):
438 + def __init__(
439 + self,
440 + size: int,
441 + self_attn: nn.Module,
442 + feed_forward: nn.Module,
443 + ):
444 + super().__init__()
445 + self.self_attn = self_attn
446 + self.feed_forward = feed_forward
447 + self.norm1 = LayerNorm(size)
448 + self.norm2 = LayerNorm(size)
449 +
450 + def forward(self, hidden_states: torch.Tensor):
451 + residual = hidden_states
452 + hidden_states = self.norm1(hidden_states)
453 + hidden_states = residual + self.self_attn(hidden_states, None, None)
454 + residual = hidden_states
455 + hidden_states = self.norm2(hidden_states)
456 + hidden_states = residual + self.feed_forward(hidden_states)
457 +
458 + return hidden_states
459 +
460 +
461 +class FunASRAudioAttention(nn.Module):
462 + def __init__(
463 + self,
464 + num_heads: int,
465 + embed_dim: int,
466 + prefix: str = "",
467 + ):
468 + super().__init__()
469 + self.embed_dim = embed_dim
470 + self.num_heads = num_heads
471 + self.head_dim = self.embed_dim // self.num_heads
472 + tp_size = get_tensor_model_parallel_world_size()
473 + self.num_local_heads = self.num_heads // tp_size
474 +
475 + if (self.head_dim * self.num_heads) != self.embed_dim:
476 + raise ValueError(
477 + f"embed_dim must be divisible by num_heads (got `embed_dim`: "
478 + f"{self.embed_dim} and `num_heads`: {self.num_heads})."
479 + )
480 +
481 + self.scaling = self.head_dim**-0.5
482 +
483 + self.qkv = QKVParallelLinear(
484 + hidden_size=self.embed_dim,
485 + head_size=self.head_dim,
486 + total_num_heads=self.num_heads,
487 + total_num_kv_heads=self.num_heads,
488 + bias=True,
489 + prefix=f"{prefix}.qkv",
490 + )
491 +
492 + self.out_proj = RowParallelLinear(
493 + input_size=self.embed_dim,
494 + output_size=self.embed_dim,
495 + bias=True,
496 + prefix=f"{prefix}.out_proj",
497 + )
498 +
499 + self.attn = MMEncoderAttention(
500 + num_heads=self.num_local_heads,
501 + head_size=self.head_dim,
502 + scale=self.scaling,
503 + )
504 +
505 + def forward(
506 + self,
507 + hidden_states: torch.Tensor,
508 + cu_seqlens: torch.Tensor,
509 + max_seqlen: torch.Tensor | None,
510 + ) -> torch.Tensor:
511 + bs, seq_length, _ = hidden_states.size()
512 + qkv, _ = self.qkv(hidden_states)
513 + q, k, v = qkv.chunk(3, dim=-1)
514 + q = q.view(bs, seq_length, -1, self.head_dim)
515 + k = k.view(bs, seq_length, -1, self.head_dim)
516 + v = v.view(bs, seq_length, -1, self.head_dim)
517 +
518 + attn_output = self.attn(
519 + query=q,
520 + key=k,
521 + value=v,
522 + cu_seqlens=cu_seqlens,
523 + max_seqlen=max_seqlen,
524 + )
525 +
526 + attn_output = attn_output.view(bs, seq_length, -1)
527 + output, _ = self.out_proj(attn_output)
528 + return output
529 +
530 +
531 +class Transformer(nn.Module):
532 + def __init__(
533 + self,
534 + downsample_rate=2,
535 + encoder_dim=1280,
536 + llm_dim=4096,
537 + ffn_dim: int = 2048,
538 + prefix: str = "",
539 + **kwargs,
540 + ):
541 + super().__init__()
542 + self.k = downsample_rate
543 + self.encoder_dim = encoder_dim
544 + self.llm_dim = llm_dim
545 + self.linear1 = ColumnParallelLinear(
546 + input_size=self.encoder_dim * self.k,
547 + output_size=ffn_dim,
548 + bias=True,
549 + )
550 + self.relu = nn.ReLU()
551 + self.linear2 = RowParallelLinear(
552 + input_size=ffn_dim,
553 + output_size=self.llm_dim,
554 + bias=True,
555 + )
556 +
557 + self.blocks = None
558 + if kwargs.get("n_layer", 2) > 0:
559 + self.blocks = nn.ModuleList(
560 + [
561 + EncoderLayer(
562 + llm_dim,
563 + FunASRAudioAttention(
564 + kwargs.get("attention_heads", 8),
565 + llm_dim,
566 + prefix=f"{prefix}.self_attn",
567 + ),
568 + PositionwiseFeedForward(
569 + llm_dim,
570 + llm_dim // 4,
571 + ),
572 + )
573 + for _ in range(kwargs.get("n_layer", 2))
574 + ]
575 + )
576 +
577 + def forward(self, hidden_states: torch.Tensor, ilens: int = 0):
578 + batch_size, seq_len, dim = hidden_states.size()
579 + chunk_num = (seq_len - 1) // self.k + 1
580 + pad_num = chunk_num * self.k - seq_len
581 + hidden_states = F.pad(hidden_states, (0, 0, 0, pad_num, 0, 0), value=0.0)
582 + seq_len = hidden_states.size(1)
583 +
584 + hidden_states = hidden_states.contiguous()
585 + hidden_states = hidden_states.view(batch_size, chunk_num, dim * self.k)
586 + hidden_states, _ = self.linear1(hidden_states)
587 + hidden_states = self.relu(hidden_states)
588 + hidden_states, _ = self.linear2(hidden_states)
589 +
590 + olens = None
591 + olens = (ilens - 1) // self.k + 1
592 +
593 + if self.blocks is not None:
594 + for layer, block in enumerate(self.blocks):
595 + hidden_states = block(hidden_states)
596 + return hidden_states, olens
597 +
598 +
599 +class FunASRAudioInputs(TensorSchema):
600 + """
601 + Dimensions:
602 + - b: Batch size
603 + - nmb: Number of mel bins
604 + - t: Time frames (M)
605 + """
606 +
607 + input_features: Annotated[
608 + list[torch.Tensor] | None,
609 + TensorShape("b", "nmb", "t"),
610 + ]
611 + speech_lengths: Annotated[
612 + list[torch.Tensor] | None,
613 + TensorShape("b"),
614 + ]
615 +
616 +
617 +class FunASREncoder(nn.Module):
618 + def __init__(
619 + self, *, vllm_config: VllmConfig, prefix: str = "", init_in_fp32: bool = False
620 + ):
621 + super().__init__()
622 + self.audio_encoder = SenseVoiceEncoderSmall(
623 + input_size=560, **vllm_config.model_config.hf_config.audio_encoder_conf
624 + )
625 + self.audio_adaptor = Transformer(
626 + downsample_rate=1,
627 + use_low_frame_rate=True,
628 + ffn_dim=2048,
629 + llm_dim=1024,
630 + encoder_dim=512,
631 + n_layer=2,
632 + freeze=True,
633 + prefix=maybe_prefix(prefix, "audio_encoder"),
634 + )
635 +
636 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
637 + """Load weights with mapping from HuggingFace format."""
638 + stacked_params_mapping = [
639 + # (param_name, shard_name, shard_id)
640 + ("self_attn.qkv.", "self_attn.q_proj.", "q"),
641 + ("self_attn.qkv.", "self_attn.k_proj.", "k"),
642 + ("self_attn.qkv.", "self_attn.v_proj.", "v"),
643 + ]
644 + params_dict = dict(self.named_parameters(remove_duplicate=False))
645 + loaded_params: set[str] = set()
646 +
647 + for name, loaded_weight in weights:
648 + for param_name, weight_name, shard_id in stacked_params_mapping:
649 + if weight_name not in name:
650 + continue
651 + name = name.replace(weight_name, param_name)
652 +
653 + param = params_dict[name]
654 + weight_loader = param.weight_loader
655 + weight_loader(param, loaded_weight, shard_id)
656 + break
657 + else:
658 + param = params_dict.get(name)
659 + if param is not None:
660 + weight_loader = getattr(
661 + param, "weight_loader", default_weight_loader
662 + )
663 + weight_loader(param, loaded_weight)
664 + loaded_params.add(name)
665 + return loaded_params
666 +
667 +
668 +class FunASRModel(nn.Module):
669 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
670 + super().__init__()
671 + self.encoder = FunASREncoder(
672 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "encoder")
673 + )
674 + self.decoder = Qwen3Model(
675 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "decoder")
676 + )
677 +
678 + def forward(
679 + self,
680 + input_ids: torch.Tensor | None,
681 + positions: torch.Tensor,
682 + inputs_embeds: torch.Tensor | None = None,
683 + ) -> torch.Tensor:
684 + decoder_outputs = self.decoder(
685 + input_ids=input_ids,
686 + positions=positions,
687 + inputs_embeds=inputs_embeds,
688 + )
689 + return decoder_outputs
690 +
691 + def get_encoder_outputs(
692 + self,
693 + speech: torch.Tensor | list[torch.Tensor] | None,
694 + speech_lengths: torch.Tensor | list[torch.Tensor] | None,
695 + ) -> torch.Tensor | None:
696 + self.feat_permute = False
697 +
698 + if self.feat_permute:
699 + encoder_out, encoder_out_lens = self.encoder.audio_encoder(
700 + speech.permute(0, 2, 1), speech_lengths
701 + )
702 + else:
703 + encoder_out, encoder_out_lens = self.encoder.audio_encoder(
704 + speech, speech_lengths
705 + )
706 +
707 + encoder_out, encoder_out_lens = self.encoder.audio_adaptor(
708 + encoder_out, encoder_out_lens
709 + )
710 + return encoder_out
711 +
712 +
713 +class FunASRProcessingInfo(BaseProcessingInfo):
714 + def get_hf_config(self) -> Qwen3Config:
715 + return self.ctx.get_hf_config(Qwen3Config)
716 +
717 + @property
718 + def skip_prompt_length_check(self) -> bool:
719 + return True # Because the encoder prompt is padded
720 +
721 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:
722 + return {"audio": 1}
723 +
724 + def get_feature_extractor(self, **kwargs: object) -> FunASRFeatureExtractor:
725 + hf_processor = self.get_hf_processor(**kwargs)
726 + feature_extractor = hf_processor.feature_extractor # type: ignore
727 + assert isinstance(feature_extractor, FunASRFeatureExtractor)
728 + return feature_extractor
729 +
730 + def get_target_channels(self) -> int:
731 + return 1
732 +
733 + def get_num_audio_tokens(self) -> int:
734 + return self.get_hf_config().max_source_positions
735 +
736 +
737 +class FunASRDummyInputsBuilder(BaseDummyInputsBuilder[FunASRProcessingInfo]):
738 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
739 + num_audios = mm_counts.get("audio", 0)
740 +
741 + return "<|AUDIO|>" * num_audios
742 +
743 + def get_dummy_mm_data(
744 + self,
745 + seq_len: int,
746 + mm_counts: Mapping[str, int],
747 + mm_options: Mapping[str, BaseDummyOptions] | None = None,
748 + ) -> MultiModalDataDict:
749 + feature_extractor = self.info.get_feature_extractor()
750 +
751 + sampling_rate = feature_extractor.sampling_rate
752 + audio_len = feature_extractor.chunk_length * sampling_rate
753 + num_audios = mm_counts.get("audio", 0)
754 +
755 + audio_overrides = mm_options.get("audio") if mm_options else None
756 +
757 + return {
758 + "audio": self._get_dummy_audios(
759 + length=audio_len, num_audios=num_audios, overrides=audio_overrides
760 + )
761 + }
762 +
763 +
764 +class FunASRMultiModalProcessor(BaseMultiModalProcessor[FunASRProcessingInfo]):
765 + def _get_data_parser(self) -> MultiModalDataParser:
766 + feature_extractor = self.info.get_feature_extractor()
767 + return MultiModalDataParser(
768 + target_sr=feature_extractor.sampling_rate,
769 + target_channels=self.info.get_target_channels(),
770 + )
771 +
772 + def _call_hf_processor(
773 + self,
774 + prompt: str,
775 + mm_data: Mapping[str, object],
776 + mm_kwargs: Mapping[str, object],
777 + tok_kwargs: Mapping[str, object],
778 + ) -> BatchFeature:
779 + if mm_data:
780 + feature_extractor = self.info.get_feature_extractor(**mm_kwargs)
781 + mm_data = dict(audio=mm_data.pop("audios"))
782 + mm_kwargs = dict(
783 + **mm_kwargs,
784 + sampling_rate=feature_extractor.sampling_rate,
785 + )
786 + processed_outputs = super()._call_hf_processor(
787 + prompt=prompt,
788 + mm_data=mm_data,
789 + mm_kwargs=mm_kwargs,
790 + tok_kwargs=tok_kwargs,
791 + )
792 + if "labels" in processed_outputs:
793 + processed_outputs["input_ids"] = processed_outputs.pop("labels")
794 + return processed_outputs
795 +
796 + def _get_mm_fields_config(
797 + self,
798 + hf_inputs: BatchFeature,
799 + hf_processor_mm_kwargs: Mapping[str, object],
800 + ) -> Mapping[str, MultiModalFieldConfig]:
801 + return dict(
802 + input_features=MultiModalFieldConfig.batched("audio"),
803 + speech_lengths=MultiModalFieldConfig.batched("audio"),
804 + fake_token_len=MultiModalFieldConfig.batched("audio"),
805 + )
806 +
807 + def _get_prompt_updates(
808 + self,
809 + mm_items: MultiModalDataItems,
810 + hf_processor_mm_kwargs: Mapping[str, object],
811 + out_mm_kwargs: MultiModalKwargsItems,
812 + ) -> Sequence[PromptUpdate]:
813 + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
814 + tokenizer = self.info.get_tokenizer()
815 + vocab = tokenizer.get_vocab()
816 +
817 + # Use getattr with default to be compatible with transformers<4.48
818 + audio_token = getattr(processor, "audio_token", "<|AUDIO|>")
819 +
820 + audio_token_id = vocab[audio_token]
821 +
822 + out_mm_data = out_mm_kwargs.get_data()
823 +
824 + fake_token_len = out_mm_data.get("fake_token_len")
825 + if fake_token_len is None:
826 + audio_output_lengths = []
827 + else:
828 + assert isinstance(fake_token_len, torch.Tensor)
829 +
830 + audio_output_lengths = fake_token_len.tolist()
831 +
832 + def get_replacement_qwen2_audio(item_idx: int):
833 + if audio_output_lengths:
834 + num_features = audio_output_lengths[item_idx]
835 + else:
836 + audio_embeds = out_mm_data["audio_embeds"][item_idx]
837 + assert len(audio_embeds.shape) == 2, "audio_embeds must be a 2D tensor"
838 + num_features = audio_embeds.shape[0]
839 +
840 + audio_tokens = [audio_token_id] * num_features
841 +
842 + return PromptUpdateDetails.select_token_id(
843 + audio_tokens,
844 + embed_token_id=audio_token_id,
845 + )
846 +
847 + return [
848 + PromptReplacement(
849 + modality="audio",
850 + target=audio_token,
851 + replacement=get_replacement_qwen2_audio,
852 + )
853 + ]
854 +
855 +
856 +@MULTIMODAL_REGISTRY.register_processor(
857 + FunASRMultiModalProcessor,
858 + info=FunASRProcessingInfo,
859 + dummy_inputs=FunASRDummyInputsBuilder,
860 +)
861 +class FunASRForConditionalGeneration(
862 + nn.Module, SupportsTranscription, SupportsMultiModal
863 +):
864 + packed_modules_mapping = {
865 + "self_attn.qkv_proj": [
866 + "self_attn.q_proj",
867 + "self_attn.k_proj",
868 + "self_attn.v_proj",
869 + ],
870 + "encoder_attn.kv_proj": ["encoder_attn.k_proj", "encoder_attn.v_proj"],
871 + }
872 +
873 + hf_to_vllm_mapper = WeightsMapper(
874 + orig_to_new_substr={
875 + "linear_q.": "q_proj.",
876 + "linear_k.": "k_proj.",
877 + "linear_v.": "v_proj.",
878 + "linear_out.": "out_proj.",
879 + }
880 + )
881 +
882 + supports_transcription_only = True
883 + supports_segment_timestamp = True
884 + supported_languages = ISO639_1_SUPPORTED_LANGS
885 +
886 + @classmethod
887 + def validate_language(cls, language: str | None) -> str | None:
888 + if language is None:
889 + # TODO language should be optional and can be guessed.
890 + # For now we default to en. See
891 + # https://github.com/huggingface/transformers/blob/main/src/transformers/models/whisper/generation_whisper.py#L1520
892 + logger.warning(
893 + "Defaulting to language='en'. If you wish to transcribe "
894 + "audio in a different language, pass the `language` field "
895 + "in the TranscriptionRequest."
896 + )
897 + language = "en"
898 + return super().validate_language(language)
899 +
900 + @classmethod
901 + def get_generation_prompt(
902 + cls,
903 + audio: np.ndarray,
904 + model_config: ModelConfig, # not needed here
905 + stt_config: SpeechToTextConfig,
906 + language: str | None,
907 + task_type: Literal["transcribe", "translate"],
908 + request_prompt: str,
909 + to_language: str | None,
910 + ) -> PromptType:
911 + if language is None:
912 + raise ValueError(
913 + "Language must be specified when creating the funasr prompt"
914 + )
915 +
916 + funasr_prompt = "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n语音转写：<|AUDIO|><|im_end|>\n<|im_start|>assistant\n" # noqa: E501
917 + prompt = {
918 + "prompt": funasr_prompt,
919 + "multi_modal_data": {
920 + "audio": (audio, stt_config.sample_rate),
921 + },
922 + }
923 + return cast(PromptType, prompt)
924 +
925 + @classmethod
926 + def get_speech_to_text_config(
927 + cls, model_config: ModelConfig, task_type: str
928 + ) -> SpeechToTextConfig:
929 + processor = cached_processor_from_config(model_config)
930 +
931 + return SpeechToTextConfig(
932 + max_audio_clip_s=processor.feature_extractor.chunk_length,
933 + sample_rate=processor.feature_extractor.sampling_rate,
934 + )
935 +
936 + @classmethod
937 + def get_num_audio_tokens(
938 + cls,
939 + audio_duration_s: float,
940 + stt_config: SpeechToTextConfig,
941 + model_config: ModelConfig,
942 + ) -> int | None:
943 + processor = cached_processor_from_config(model_config)
944 + hop_length = processor.feature_extractor.hop_length
945 + assert hop_length is not None
946 + return math.ceil(audio_duration_s * stt_config.sample_rate / hop_length)
947 +
948 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
949 + super().__init__()
950 + config = vllm_config.model_config.hf_config
951 + quant_config = vllm_config.quant_config
952 + self.config = config
953 + self.dtype = vllm_config.model_config.dtype
954 +
955 + self.model = FunASRModel(
956 + vllm_config=vllm_config,
957 + prefix=maybe_prefix(prefix, "model"),
958 + )
959 + logit_scale = getattr(config, "logit_scale", 1.0)
960 +
961 + if config.tie_word_embeddings:
962 + self.lm_head = self.model.decoder.embed_tokens
963 + else:
964 + self.lm_head = ParallelLMHead(
965 + config.vocab_size,
966 + config.hidden_size,
967 + quant_config=quant_config,
968 + prefix=maybe_prefix(prefix, "lm_head"),
969 + )
970 + self.logits_processor = LogitsProcessor(config.vocab_size, scale=logit_scale)
971 +
972 + def forward(
973 + self,
974 + input_ids: torch.Tensor,
975 + positions: torch.Tensor,
976 + inputs_embeds: torch.Tensor | None = None,
977 + **kwargs,
978 + ) -> torch.Tensor:
979 + decoder_outputs = self.model(
980 + input_ids=input_ids,
981 + positions=positions,
982 + inputs_embeds=inputs_embeds,
983 + )
984 + return decoder_outputs
985 +
986 + def get_language_model(self) -> torch.nn.Module:
987 + return self.model.decoder
988 +
989 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
990 + audio_input = self._parse_and_validate_audio_input(**kwargs)
991 +
992 + speech = audio_input["input_features"]
993 + speech_lengths = audio_input["speech_lengths"]
994 + enc_output = self.model.get_encoder_outputs(
995 + speech=speech, speech_lengths=speech_lengths
996 + )
997 +
998 + return enc_output
999 +
1000 + def embed_input_ids(
1001 + self,
1002 + input_ids: torch.Tensor,
1003 + multimodal_embeddings: MultiModalEmbeddings | None = None,
1004 + *,
1005 + is_multimodal: torch.Tensor | None = None,
1006 + handle_oov_mm_token: bool = False,
1007 + ) -> torch.Tensor:
1008 + inputs_embeds = self.model.decoder.embed_input_ids(input_ids)
1009 +
1010 + return _merge_multimodal_embeddings(
1011 + inputs_embeds=inputs_embeds,
1012 + multimodal_embeddings=multimodal_embeddings,
1013 + is_multimodal=_require_is_multimodal(is_multimodal),
1014 + )
1015 +
1016 + def _parse_and_validate_audio_input(self, **kwargs: object) -> FunASRAudioInputs:
1017 + input_features = kwargs.pop("input_features", None)
1018 + speech_lengths = kwargs.pop("speech_lengths", None)
1019 +
1020 + if input_features is not None:
1021 + input_features = json_map_leaves(lambda x: x.to(self.dtype), input_features)
1022 +
1023 + if speech_lengths is not None:
1024 + speech_lengths = json_map_leaves(lambda x: x.to(self.dtype), speech_lengths)
1025 +
1026 + return FunASRAudioInputs(
1027 + input_features=input_features, speech_lengths=speech_lengths
1028 + )
1029 +
1030 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
1031 + logits = self.logits_processor(self.lm_head, hidden_states)
1032 + return logits
1033 +
1034 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1035 + loader = AutoWeightsLoader(
1036 + self,
1037 + )
1038 +
1039 + # add fake zeros bias for k_proj to state_dict
1040 + weights = _create_fake_bias_for_k_proj(weights)
1041 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1042 +
1043 +
1044 +def _create_fake_bias_for_k_proj(
1045 + weights: Iterable[tuple[str, torch.Tensor]],
1046 +) -> Iterable[tuple[str, torch.Tensor]]:
1047 + """
1048 + Create full zeros bias for k_proj weight in self-attn and x-attn layers.
1049 + So that the bias for k_proj in qkv_proj can be initialized with zeros.
1050 + """
1051 + for name, weight in weights:
1052 + if name.endswith(".k_proj.weight"):
1053 + bias = torch.zeros(weight.size(0))
1054 + bias_name = name.replace("weight", "bias")
1055 + yield from [(name, weight), (bias_name, bias)]
1056 + else:
1057 + yield name, weight
vllm/model_executor/models/registry.py
+1
@@ -325,6 +325,7 @@ _MULTIMODAL_MODELS = {
325 "ernie45_vl",
326 "Ernie4_5_VLMoeForConditionalGeneration",
327 ),
328 + "FunASRForConditionalGeneration": ("funasr", "FunASRForConditionalGeneration"), # noqa: E501
329 "FunAudioChatForConditionalGeneration": (
330 "funaudiochat",
331 "FunAudioChatForConditionalGeneration",
vllm/transformers_utils/processors/__init__.py
+2
@@ -10,6 +10,7 @@ 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.funasr_processor import FunASRProcessor
14 from vllm.transformers_utils.processors.hunyuan_vl import HunYuanVLProcessor
15 from vllm.transformers_utils.processors.hunyuan_vl_image import HunYuanVLImageProcessor
16 from vllm.transformers_utils.processors.ovis import OvisProcessor
@@ -18,6 +19,7 @@ from vllm.transformers_utils.processors.ovis2_5 import Ovis2_5Processor
19 __all__ = [
20 "BagelProcessor",
21 "DeepseekVLV2Processor",
22 + "FunASRProcessor",
23 "HunYuanVLProcessor",
24 "HunYuanVLImageProcessor",
25 "OvisProcessor",
vllm/transformers_utils/processors/funasr_processor.py new
+504
@@ -0,0 +1,504 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +
5 +import numpy as np
6 +import torch
7 +import torch.nn as nn
8 +import torchaudio.compliance.kaldi as kaldi
9 +from torch.nn.utils.rnn import pad_sequence
10 +from transformers import (
11 + AutoFeatureExtractor,
12 + AutoProcessor,
13 + BatchFeature,
14 +)
15 +from transformers.feature_extraction_sequence_utils import SequenceFeatureExtractor
16 +from transformers.processing_utils import ProcessorMixin
17 +from transformers.utils import TensorType
18 +
19 +from vllm.logger import init_logger
20 +
21 +logger = init_logger(__name__)
22 +
23 +
24 +def apply_cmvn(inputs, cmvn): # noqa
25 + """
26 + Apply CMVN with mvn data
27 + """
28 +
29 + device = inputs.device
30 + # dtype = inputs.dtype
31 + frame, dim = inputs.shape
32 +
33 + means = cmvn[0:1, :dim]
34 + vars = cmvn[1:2, :dim]
35 + inputs += means.to(device)
36 + inputs *= vars.to(device)
37 +
38 + return inputs.type(torch.float32)
39 +
40 +
41 +def apply_lfr(inputs, lfr_m, lfr_n):
42 + # LFR_inputs = []
43 + T = inputs.shape[0]
44 + T_lfr = int(np.ceil(T / lfr_n))
45 + left_padding = inputs[0].repeat((lfr_m - 1) // 2, 1)
46 + inputs = torch.vstack((left_padding, inputs))
47 + T = T + (lfr_m - 1) // 2
48 + feat_dim = inputs.shape[-1]
49 + strides = (lfr_n * feat_dim, 1)
50 + sizes = (T_lfr, lfr_m * feat_dim)
51 + last_idx = (T - lfr_m) // lfr_n + 1
52 + num_padding = lfr_m - (T - last_idx * lfr_n)
53 + if num_padding > 0:
54 + num_padding = (
55 + (2 * lfr_m - 2 * T + (T_lfr - 1 + last_idx) * lfr_n)
56 + / 2
57 + * (T_lfr - last_idx)
58 + )
59 + inputs = torch.vstack([inputs] + [inputs[-1:]] * int(num_padding))
60 + LFR_outputs = inputs.as_strided(sizes, strides)
61 + return LFR_outputs.clone().type(torch.float32)
62 +
63 +
64 +def load_cmvn(cmvn_file):
65 + with open(cmvn_file, encoding="utf-8") as f:
66 + lines = f.readlines()
67 + means_list = []
68 + vars_list = []
69 + for i in range(len(lines)):
70 + line_item = lines[i].split()
71 + if line_item[0] == "<AddShift>":
72 + line_item = lines[i + 1].split()
73 + if line_item[0] == "<LearnRateCoef>":
74 + add_shift_line = line_item[3 : (len(line_item) - 1)]
75 + means_list = list(add_shift_line)
76 + continue
77 + elif line_item[0] == "<Rescale>":
78 + line_item = lines[i + 1].split()
79 + if line_item[0] == "<LearnRateCoef>":
80 + rescale_line = line_item[3 : (len(line_item) - 1)]
81 + vars_list = list(rescale_line)
82 + continue
83 + means = np.array(means_list).astype(np.float32)
84 + vars = np.array(vars_list).astype(np.float32)
85 + cmvn = np.array([means, vars])
86 + cmvn = torch.as_tensor(cmvn, dtype=torch.float32)
87 + return cmvn
88 +
89 +
90 +class WavFrontend(nn.Module):
91 + """Conventional frontend structure for ASR."""
92 +
93 + def __init__(
94 + self,
95 + cmvn_file: str = "null",
96 + fs: int = 16000,
97 + window: str = "hamming",
98 + n_mels: int = 80,
99 + frame_length: int = 25,
100 + frame_shift: int = 10,
101 + filter_length_min: int = -1,
102 + filter_length_max: int = -1,
103 + lfr_m: int = 1,
104 + lfr_n: int = 1,
105 + dither: float = 1.0,
106 + snip_edges: bool = True,
107 + upsacle_samples: bool = True,
108 + **kwargs,
109 + ):
110 + super().__init__()
111 + self.fs = fs
112 + self.window = window
113 + self.n_mels = n_mels
114 + self.frame_length = frame_length
115 + self.frame_shift = frame_shift
116 + self.filter_length_min = filter_length_min
117 + self.filter_length_max = filter_length_max
118 + self.lfr_m = lfr_m
119 + self.lfr_n = lfr_n
120 + self.cmvn_file = cmvn_file
121 + self.dither = dither
122 + self.snip_edges = snip_edges
123 + self.upsacle_samples = upsacle_samples
124 + self.cmvn = None if self.cmvn_file is None else load_cmvn(self.cmvn_file)
125 +
126 + def output_size(self) -> int:
127 + return self.n_mels * self.lfr_m
128 +
129 + def forward(
130 + self,
131 + input: torch.Tensor,
132 + input_lengths,
133 + **kwargs,
134 + ) -> tuple[torch.Tensor, torch.Tensor]:
135 + batch_size = input.size(0)
136 + feats = []
137 + feats_lens = []
138 + for i in range(batch_size):
139 + waveform_length = input_lengths[i]
140 + waveform = input[i][:waveform_length]
141 + if self.upsacle_samples:
142 + waveform = waveform * (1 << 15)
143 + waveform = waveform.unsqueeze(0)
144 + mat = kaldi.fbank(
145 + waveform,
146 + num_mel_bins=self.n_mels,
147 + frame_length=min(self.frame_length, waveform_length / self.fs * 1000),
148 + frame_shift=self.frame_shift,
149 + dither=self.dither,
150 + energy_floor=0.0,
151 + window_type=self.window,
152 + sample_frequency=self.fs,
153 + snip_edges=self.snip_edges,
154 + )
155 +
156 + if self.lfr_m != 1 or self.lfr_n != 1:
157 + mat = apply_lfr(mat, self.lfr_m, self.lfr_n)
158 + if self.cmvn is not None:
159 + mat = apply_cmvn(mat, self.cmvn)
160 + feat_length = mat.size(0)
161 + feats.append(mat)
162 + feats_lens.append(feat_length)
163 +
164 + feats_lens = torch.as_tensor(feats_lens)
165 + if batch_size == 1:
166 + feats_pad = feats[0][None, :, :]
167 + else:
168 + feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0)
169 + return feats_pad, feats_lens
170 +
171 + def forward_fbank(
172 + self, input: torch.Tensor, input_lengths: torch.Tensor
173 + ) -> tuple[torch.Tensor, torch.Tensor]:
174 + batch_size = input.size(0)
175 + feats = []
176 + feats_lens = []
177 + for i in range(batch_size):
178 + waveform_length = input_lengths[i]
179 + waveform = input[i][:waveform_length]
180 + waveform = waveform * (1 << 15)
181 + waveform = waveform.unsqueeze(0)
182 + mat = kaldi.fbank(
183 + waveform,
184 + num_mel_bins=self.n_mels,
185 + frame_length=self.frame_length,
186 + frame_shift=self.frame_shift,
187 + dither=self.dither,
188 + energy_floor=0.0,
189 + window_type=self.window,
190 + sample_frequency=self.fs,
191 + )
192 +
193 + feat_length = mat.size(0)
194 + feats.append(mat)
195 + feats_lens.append(feat_length)
196 +
197 + feats_lens = torch.as_tensor(feats_lens)
198 + feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0)
199 + return feats_pad, feats_lens
200 +
201 + def forward_lfr_cmvn(
202 + self, input: torch.Tensor, input_lengths: torch.Tensor
203 + ) -> tuple[torch.Tensor, torch.Tensor]:
204 + batch_size = input.size(0)
205 + feats = []
206 + feats_lens = []
207 + for i in range(batch_size):
208 + mat = input[i, : input_lengths[i], :]
209 + if self.lfr_m != 1 or self.lfr_n != 1:
210 + mat = apply_lfr(mat, self.lfr_m, self.lfr_n)
211 + if self.cmvn is not None:
212 + mat = apply_cmvn(mat, self.cmvn)
213 + feat_length = mat.size(0)
214 + feats.append(mat)
215 + feats_lens.append(feat_length)
216 +
217 + feats_lens = torch.as_tensor(feats_lens)
218 + feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0)
219 + return feats_pad, feats_lens
220 +
221 +
222 +class FunASRFeatureExtractor(SequenceFeatureExtractor):
223 + r"""
224 + Constructs a FunASR feature extractor.
225 +
226 + This feature extractor inherits from [`~feature_extraction_sequence_
227 + utils.SequenceFeatureExtractor`] which contains most of the main
228 + methods. Users should refer to this superclass for more information
229 + regarding those methods.
230 +
231 + This class extracts mel-filter bank features from raw speech using a custom
232 + numpy implementation of the `Short Time Fourier Transform` which should
233 + match pytorch's `torch.stft` equivalent.
234 +
235 + Args:
236 + feature_size (`int`, *optional*, defaults to 80):
237 + The feature dimension of the extracted features.
238 + sampling_rate (`int`, *optional*, defaults to 16000):
239 + The sampling rate at which the audio files should be digitalized
240 + expressed in hertz (Hz).
241 + hop_length (`int`, *optional*, defaults to 160):
242 + Length of the overlapping windows for the STFT used to obtain the
243 + Mel Frequency coefficients.
244 + chunk_length (`int`, *optional*, defaults to 30):
245 + The maximum number of chunks of `sampling_rate` samples used to
246 + trim and pad longer or shorter audio sequences.
247 + n_fft (`int`, *optional*, defaults to 400):
248 + Size of the Fourier transform.
249 + padding_value (`float`, *optional*, defaults to 0.0):
250 + Padding value used to pad the audio. Should correspond to silences.
251 + dither (`float`, *optional*, defaults to 0.0):
252 + Adds dithering. In other words, adds a small Gaussian noise to each frame.
253 + E.g. use 0.0001 to add dithering with a normal distribution centered
254 + around 0.0 with standard deviation 0.0001 (assuming [-1,+1] range
255 + of raw_speech). The value 0.0 means no dithering.
256 + Dithering has similar effect as `spectrogram(mel_floor=...)`. It reduces
257 + the high log_mel_fbank values for signals with hard-zero sections,
258 + when VAD cutoff is present in the signal.
259 + """
260 +
261 + model_input_names = ["input_features"]
262 +
263 + def __init__(
264 + self,
265 + feature_size=80,
266 + sampling_rate=16000,
267 + hop_length=160,
268 + chunk_length=30,
269 + n_fft=400,
270 + padding_value=0.0,
271 + dither=0.0,
272 + return_attention_mask=False,
273 + **kwargs,
274 + ):
275 + super().__init__(
276 + feature_size=feature_size,
277 + sampling_rate=sampling_rate,
278 + padding_value=padding_value,
279 + return_attention_mask=return_attention_mask,
280 + **kwargs,
281 + )
282 + self.frontend_conf = kwargs.get("frontend_conf", {})
283 + self.n_fft = n_fft
284 + self.hop_length = hop_length
285 + self.chunk_length = chunk_length
286 + self.n_samples = chunk_length * sampling_rate
287 + self.nb_max_frames = self.n_samples // hop_length
288 + self.sampling_rate = sampling_rate
289 + self.dither = dither
290 +
291 + def extract_fbank(
292 + self, data, data_len=None, data_type: str = "sound", frontend=None, **kwargs
293 + ):
294 + if isinstance(data, np.ndarray):
295 + data = torch.from_numpy(data)
296 + if len(data.shape) < 2:
297 + data = data[None, :] # data: [batch, N]
298 + data_len = [data.shape[1]] if data_len is None else data_len
299 + elif isinstance(data, torch.Tensor):
300 + if len(data.shape) < 2:
301 + data = data[None, :] # data: [batch, N]
302 + data_len = [data.shape[1]] if data_len is None else data_len
303 + elif isinstance(data, (list, tuple)):
304 + data_list, data_len = [], []
305 + for data_i in data:
306 + if isinstance(data_i, np.ndarray):
307 + data_i = torch.from_numpy(data_i)
308 + data_list.append(data_i)
309 + data_len.append(data_i.shape[0])
310 + data = pad_sequence(data_list, batch_first=True)
311 +
312 + data, data_len = frontend(data, data_len, **kwargs)
313 +
314 + if isinstance(data_len, (list, tuple)):
315 + data_len = torch.tensor([data_len])
316 + return data.to(torch.float32), data_len.to(torch.int32)
317 +
318 + def __call__(
319 + self,
320 + raw_speech: np.ndarray | list[float] | list[np.ndarray] | list[list[float]],
321 + truncation: bool = True,
322 + pad_to_multiple_of: int | None = None,
323 + return_tensors: str | TensorType | None = None,
324 + return_attention_mask: bool | None = None,
325 + padding: str | None = "max_length",
326 + max_length: int | None = None,
327 + sampling_rate: int | None = None,
328 + do_normalize: bool | None = None,
329 + device: str | None = "cpu",
330 + return_token_timestamps: bool | None = None,
331 + **kwargs,
332 + ) -> BatchFeature:
333 + is_batched = isinstance(raw_speech, (list, tuple)) and (
334 + isinstance(raw_speech[0], (np.ndarray, tuple, list))
335 + )
336 +
337 + if is_batched:
338 + raw_speech = [
339 + np.asarray([speech], dtype=np.float32).T for speech in raw_speech
340 + ]
341 + elif not is_batched and not isinstance(raw_speech, np.ndarray):
342 + raw_speech = np.asarray(raw_speech, dtype=np.float32)
343 + elif isinstance(raw_speech, np.ndarray) and raw_speech.dtype is np.dtype(
344 + np.float64
345 + ):
346 + raw_speech = raw_speech.astype(np.float32)
347 +
348 + if not is_batched:
349 + raw_speech = [np.asarray([raw_speech]).T]
350 +
351 + batched_speech = BatchFeature({"input_features": raw_speech})
352 +
353 + padded_inputs = self.pad(
354 + batched_speech,
355 + padding=padding,
356 + max_length=max_length if max_length else self.n_samples,
357 + truncation=truncation,
358 + pad_to_multiple_of=pad_to_multiple_of,
359 + return_attention_mask=return_attention_mask or do_normalize,
360 + )
361 +
362 + input_features = padded_inputs.get("input_features").transpose(2, 0, 1)
363 +
364 + self.frontend = WavFrontend(**self.frontend_conf)
365 + input_features, speech_lengths = self.extract_fbank(
366 + input_features[0],
367 + data_type=kwargs.get("data_type", "sound"),
368 + frontend=self.frontend,
369 + is_final=True,
370 + )
371 + olens = 1 + (speech_lengths - 3 + 2 * 1) // 2
372 + olens = 1 + (olens - 3 + 2 * 1) // 2
373 + fake_token_len = (olens - 1) // 2 + 1
374 + if isinstance(input_features[0], list):
375 + padded_inputs["input_features"] = [
376 + np.asarray(feature, dtype=np.float32) for feature in input_features
377 + ]
378 +
379 + else:
380 + padded_inputs["input_features"] = input_features
381 +
382 + if return_tensors is not None:
383 + padded_inputs = padded_inputs.convert_to_tensors(return_tensors)
384 +
385 + padded_inputs["speech_lengths"] = speech_lengths
386 + padded_inputs["fake_token_len"] = fake_token_len
387 +
388 + return padded_inputs
389 +
390 +
391 +class FunASRProcessor(ProcessorMixin):
392 + r"""
393 + Constructs a FunASR processor which wraps a FunASR feature extractor and
394 + a FunASR tokenizer into a single processor.
395 +
396 + [`FunASRProcessor`] offers all the functionalities of
397 + [`FunASRFeatureExtractor`] and [`Qwen2Tokenizer`]. See the
398 + [`~FunASRProcessor.__call__`] and [`~FunASRProcessor.decode`] for more
399 + information.
400 +
401 + Args:
402 + feature_extractor (`FunASRFeatureExtractor`): An instance of
403 + [`FunASRFeatureExtractor`].
404 + The feature extractor is a required input.
405 + tokenizer (`Qwen2Tokenizer`):
406 + An instance of [`Qwen2Tokenizer`]. The tokenizer is a required
407 + input.
408 + """
409 +
410 + feature_extractor_class = "FunASRFeatureExtractor"
411 + tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")
412 +
413 + def __init__(
414 + self,
415 + feature_extractor,
416 + tokenizer,
417 + audio_token="<|AUDIO|>",
418 + ):
419 + super().__init__(feature_extractor, tokenizer)
420 + self.current_processor = self.feature_extractor
421 + self._in_target_context_manager = False
422 + self.audio_token = (
423 + tokenizer.audio_token if hasattr(tokenizer, "audio_token") else audio_token
424 + )
425 + self.audio_token_id = tokenizer.convert_tokens_to_ids(self.audio_token)
426 +
427 + def get_decoder_prompt_ids(self, task=None, language=None, no_timestamps=True):
428 + return self.tokenizer.get_decoder_prompt_ids(
429 + task=task, language=language, no_timestamps=no_timestamps
430 + )
431 +
432 + def __call__(self, *args, **kwargs):
433 + """
434 + Forwards the `audio` argument to FunASRFeatureExtractor's
435 + [`~FunASRFeatureExtractor.__call__`] and the `text` argument to
436 + [`~Qwen2Tokenizer.__call__`]. Please refer to the docstring of the
437 + above two methods for more information.
438 + """
439 + if self._in_target_context_manager:
440 + return self.current_processor(*args, **kwargs)
441 +
442 + audio = kwargs.pop("audio", None)
443 + sampling_rate = kwargs.pop("sampling_rate", None)
444 + text = kwargs.pop("text", None)
445 + if len(args) > 0:
446 + audio = args[0]
447 + args = args[1:]
448 +
449 + if text is None:
450 + raise ValueError("You need to specify `text` input to process.")
451 + elif isinstance(text, str):
452 + text = [text]
453 + elif not isinstance(text, list) and not isinstance(text[0], str):
454 + raise ValueError(
455 + "Invalid input text. Please provide a string, or a list of strings"
456 + )
457 +
458 + if audio is not None:
459 + # ensure we have as much audios as audio tokens
460 + num_audio_tokens = sum(sample.count(self.audio_token) for sample in text)
461 + num_audios = 1 if type(audio) is np.ndarray else len(audio)
462 + if num_audio_tokens != num_audios:
463 + raise ValueError(
464 + 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
465 + )
466 + inputs = self.feature_extractor(
467 + audio, *args, sampling_rate=sampling_rate, **kwargs
468 + )
469 +
470 + expanded_text = []
471 + for sample in text:
472 + replace_str = []
473 + while self.audio_token in sample:
474 + num_audio_tokens = inputs["fake_token_len"].item()
475 +
476 + expanded_audio_token = self.audio_token * num_audio_tokens
477 +
478 + replace_str.append(expanded_audio_token)
479 + sample = sample.replace(self.audio_token, "<placeholder>", 1)
480 +
481 + while "<placeholder>" in sample:
482 + sample = sample.replace("<placeholder>", replace_str.pop(0), 1)
483 + expanded_text.append(sample)
484 + text = expanded_text
485 +
486 + if text is not None:
487 + encodings = self.tokenizer(text, **kwargs)
488 +
489 + if text is None:
490 + return inputs
491 +
492 + elif audio is None:
493 + return encodings
494 + else:
495 + inputs["labels"] = encodings["input_ids"]
496 +
497 + return inputs
498 +
499 + def get_prompt_ids(self, text: str, return_tensors="np"):
500 + return self.tokenizer.get_prompt_ids(text, return_tensors=return_tensors)
501 +
502 +
503 +AutoFeatureExtractor.register("FunASRFeatureExtractor", FunASRFeatureExtractor)
504 +AutoProcessor.register("FunASRProcessor", FunASRProcessor)