freedakgmail
6116b1f3c6
Update API Documentation / build-api-docs (push) Has been cancelled
Add complete FunASR codebase including models, runtime, and documentation.
729 lines
29 KiB
Python
729 lines
29 KiB
Python
#!/usr/bin/env python3
|
||
# -*- encoding: utf-8 -*-
|
||
# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
|
||
# MIT License (https://opensource.org/licenses/MIT)
|
||
|
||
"""
|
||
Fun-ASR-Nano vLLM inference engine.
|
||
|
||
Uses vLLM for high-throughput LLM decoding while keeping the audio encoder
|
||
and adaptor in PyTorch. Supports batch inference and tensor-parallel for
|
||
multi-GPU acceleration.
|
||
|
||
Usage:
|
||
from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM
|
||
|
||
engine = FunASRNanoVLLM.from_pretrained(
|
||
model="FunAudioLLM/Fun-ASR-Nano-2512",
|
||
tensor_parallel_size=2,
|
||
)
|
||
results = engine.generate(["audio1.wav", "audio2.wav"])
|
||
"""
|
||
|
||
import glob
|
||
import json
|
||
import logging
|
||
import os
|
||
import re
|
||
import shutil
|
||
import time
|
||
from typing import List, Optional, Union
|
||
|
||
import numpy as np
|
||
import torch
|
||
import torch.nn as nn
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}
|
||
|
||
|
||
def prepare_vllm_model_dir(model_dir: str, output_dir: str = None) -> str:
|
||
"""Extract LLM weights from Fun-ASR-Nano model.pt and save in HuggingFace format.
|
||
|
||
Fun-ASR-Nano stores all weights (audio encoder + adaptor + LLM) in a single
|
||
model.pt file. vLLM needs the LLM weights in standard HuggingFace format.
|
||
This function extracts LLM weights and saves them alongside the config/tokenizer
|
||
files from the Qwen3-0.6B subdirectory.
|
||
|
||
Args:
|
||
model_dir: Path to the Fun-ASR-Nano model directory.
|
||
output_dir: Where to save the extracted LLM. Defaults to model_dir/Qwen3-0.6B-vllm.
|
||
|
||
Returns:
|
||
Path to the directory containing the vLLM-ready LLM model.
|
||
"""
|
||
if output_dir is None:
|
||
output_dir = os.path.join(model_dir, "Qwen3-0.6B-vllm")
|
||
|
||
# Check if already prepared
|
||
safetensors_files = glob.glob(os.path.join(output_dir, "*.safetensors"))
|
||
bin_files = glob.glob(os.path.join(output_dir, "model*.bin"))
|
||
if safetensors_files or bin_files:
|
||
logger.info(f"vLLM model already prepared at {output_dir}")
|
||
return output_dir
|
||
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
# Copy config and tokenizer from Qwen3-0.6B
|
||
qwen_dir = os.path.join(model_dir, "Qwen3-0.6B")
|
||
if not os.path.isdir(qwen_dir):
|
||
raise FileNotFoundError(f"Qwen3-0.6B config directory not found at {qwen_dir}")
|
||
|
||
for fname in os.listdir(qwen_dir):
|
||
src = os.path.join(qwen_dir, fname)
|
||
dst = os.path.join(output_dir, fname)
|
||
if os.path.isfile(src) and not os.path.exists(dst):
|
||
shutil.copy2(src, dst)
|
||
|
||
# Load model.pt and extract LLM weights
|
||
model_pt = os.path.join(model_dir, "model.pt")
|
||
if not os.path.exists(model_pt):
|
||
raise FileNotFoundError(
|
||
f"model.pt not found at {model_pt}. Make sure the model is fully downloaded."
|
||
)
|
||
|
||
logger.info(f"Loading model.pt from {model_pt}...")
|
||
checkpoint = torch.load(model_pt, map_location="cpu")
|
||
if "state_dict" in checkpoint:
|
||
state_dict = checkpoint["state_dict"]
|
||
else:
|
||
state_dict = checkpoint
|
||
|
||
# Extract LLM weights (prefixed with "llm.")
|
||
llm_state = {}
|
||
for key, value in state_dict.items():
|
||
if key.startswith("llm."):
|
||
new_key = key[len("llm."):]
|
||
llm_state[new_key] = value
|
||
|
||
if not llm_state:
|
||
raise RuntimeError("No LLM weights found in model.pt (expected prefix 'llm.')")
|
||
|
||
logger.info(f"Extracted {len(llm_state)} LLM weight tensors")
|
||
|
||
# Save in safetensors format (preferred by vLLM)
|
||
try:
|
||
from safetensors.torch import save_file
|
||
|
||
save_path = os.path.join(output_dir, "model.safetensors")
|
||
save_file(llm_state, save_path)
|
||
logger.info(f"Saved LLM weights to {save_path}")
|
||
|
||
# Create model index
|
||
index = {
|
||
"metadata": {"total_size": sum(v.numel() * v.element_size() for v in llm_state.values())},
|
||
"weight_map": {k: "model.safetensors" for k in llm_state.keys()},
|
||
}
|
||
with open(os.path.join(output_dir, "model.safetensors.index.json"), "w") as f:
|
||
json.dump(index, f, indent=2)
|
||
except ImportError:
|
||
save_path = os.path.join(output_dir, "model.bin")
|
||
torch.save(llm_state, save_path)
|
||
logger.info(f"Saved LLM weights to {save_path} (install safetensors for faster loading)")
|
||
|
||
return output_dir
|
||
|
||
|
||
class FunASRNanoVLLM:
|
||
"""Fun-ASR-Nano with vLLM backend for high-throughput inference.
|
||
|
||
Architecture:
|
||
Audio -> WavFrontend -> SenseVoiceEncoder -> AudioAdaptor -> audio embeddings
|
||
Text tokens -> LLM embedding layer -> text embeddings
|
||
Combined embeddings -> vLLM (Qwen3-0.6B) -> generated text
|
||
|
||
The audio encoder and adaptor run in PyTorch on a single GPU,
|
||
while vLLM handles the LLM inference with optional tensor parallelism.
|
||
|
||
Args:
|
||
model_dir: Path to the Fun-ASR-Nano model directory.
|
||
device: Device for audio encoder/adaptor (e.g. "cuda:0").
|
||
dtype: Dtype for audio processing ("bf16", "fp16", "fp32").
|
||
tensor_parallel_size: Number of GPUs for vLLM tensor parallelism.
|
||
gpu_memory_utilization: Fraction of GPU memory for vLLM KV cache.
|
||
max_model_len: Maximum sequence length for vLLM.
|
||
enforce_eager: Disable CUDA graph for debugging.
|
||
|
||
Example:
|
||
>>> engine = FunASRNanoVLLM(
|
||
... model_dir="/path/to/Fun-ASR-Nano-2512",
|
||
... tensor_parallel_size=2,
|
||
... )
|
||
>>> results = engine.generate(["audio1.wav", "audio2.wav"])
|
||
>>> for r in results:
|
||
... print(r["text"])
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
model_dir: str,
|
||
device: str = "cuda:0",
|
||
dtype: str = "bf16",
|
||
tensor_parallel_size: int = 1,
|
||
gpu_memory_utilization: float = 0.8,
|
||
max_model_len: int = 2048,
|
||
enforce_eager: bool = False,
|
||
**kwargs,
|
||
):
|
||
from vllm import LLM, SamplingParams
|
||
try:
|
||
from vllm.inputs import EmbedsPrompt
|
||
except ImportError:
|
||
from vllm.inputs.data import EmbedsPrompt
|
||
|
||
self.device = device
|
||
self.dtype = dtype
|
||
self.torch_dtype = dtype_map.get(dtype, torch.bfloat16)
|
||
if self.torch_dtype == torch.float16:
|
||
logger.warning(
|
||
"dtype='fp16' can produce degraded or garbage transcription for "
|
||
"Fun-ASR-Nano (numerical overflow in the audio embedding path). "
|
||
"Use dtype='bf16' (recommended) or dtype='fp32'. On GPUs without "
|
||
"bfloat16 support (e.g. NVIDIA V100), use 'fp32'."
|
||
)
|
||
self.model_dir = model_dir
|
||
|
||
# Step 1: Prepare LLM weights for vLLM (extract from model.pt if needed)
|
||
vllm_model_dir = prepare_vllm_model_dir(model_dir)
|
||
|
||
# Step 2: Load audio components (encoder + adaptor + frontend)
|
||
self._load_audio_components(model_dir, **kwargs)
|
||
|
||
# Step 3: Initialize vLLM engine
|
||
logger.info(f"Initializing vLLM with model: {vllm_model_dir}")
|
||
logger.info(f" tensor_parallel_size={tensor_parallel_size}")
|
||
logger.info(f" gpu_memory_utilization={gpu_memory_utilization}")
|
||
|
||
vllm_kwargs = kwargs.get("vllm_kwargs", {})
|
||
self.vllm_engine = LLM(
|
||
enable_prompt_embeds=True,
|
||
model=vllm_model_dir,
|
||
tensor_parallel_size=tensor_parallel_size,
|
||
gpu_memory_utilization=gpu_memory_utilization,
|
||
max_model_len=max_model_len,
|
||
enforce_eager=enforce_eager,
|
||
dtype={"bf16": "bfloat16", "fp16": "float16", "fp32": "auto"}.get(dtype, dtype),
|
||
trust_remote_code=True,
|
||
**vllm_kwargs,
|
||
)
|
||
|
||
# Step 4: Get tokenizer and LLM embedding layer
|
||
self.tokenizer = self.vllm_engine.get_tokenizer()
|
||
self._load_embedding_layer(model_dir)
|
||
|
||
def _load_audio_components(self, model_dir: str, **kwargs):
|
||
"""Load audio encoder, adaptor, frontend, and CTC from checkpoint."""
|
||
from omegaconf import OmegaConf
|
||
from funasr.register import tables
|
||
|
||
config_path = os.path.join(model_dir, "config.yaml")
|
||
config = OmegaConf.load(config_path)
|
||
self._config = OmegaConf.to_container(config, resolve=True)
|
||
|
||
# --- Frontend ---
|
||
frontend_class = tables.frontend_classes.get(config["frontend"])
|
||
frontend_conf = OmegaConf.to_container(config.get("frontend_conf", {}), resolve=True)
|
||
cmvn_file = frontend_conf.get("cmvn_file")
|
||
if cmvn_file and not os.path.isabs(cmvn_file):
|
||
frontend_conf["cmvn_file"] = os.path.join(model_dir, cmvn_file)
|
||
self.frontend = frontend_class(**frontend_conf)
|
||
self.frontend.eval()
|
||
|
||
# --- Audio Encoder ---
|
||
encoder_conf = OmegaConf.to_container(config.get("audio_encoder_conf", {}), resolve=True)
|
||
hub = encoder_conf.get("hub", None)
|
||
if hub == "ms":
|
||
from funasr import AutoModel as FunAutoModel
|
||
|
||
enc_model = FunAutoModel(
|
||
model=config["audio_encoder"], model_revision="master", disable_update=True
|
||
)
|
||
self.audio_encoder_output_size = (
|
||
enc_model.model.encoder_output_size
|
||
if hasattr(enc_model.model, "encoder_output_size")
|
||
else -1
|
||
)
|
||
self.audio_encoder = (
|
||
enc_model.model.model.encoder
|
||
if hasattr(enc_model.model, "model")
|
||
else enc_model.model.encoder
|
||
)
|
||
else:
|
||
encoder_class = tables.encoder_classes.get(config["audio_encoder"])
|
||
input_size = self.frontend.output_size()
|
||
self.audio_encoder = encoder_class(input_size=input_size, **encoder_conf)
|
||
self.audio_encoder_output_size = self.audio_encoder.output_size()
|
||
|
||
self.audio_encoder.eval()
|
||
for p in self.audio_encoder.parameters():
|
||
p.requires_grad = False
|
||
|
||
# --- Audio Adaptor ---
|
||
adaptor_conf = OmegaConf.to_container(config.get("audio_adaptor_conf", {}), resolve=True)
|
||
adaptor_class = tables.adaptor_classes.get(config["audio_adaptor"])
|
||
if self.audio_encoder_output_size > 0:
|
||
adaptor_conf["encoder_dim"] = self.audio_encoder_output_size
|
||
self.audio_adaptor = adaptor_class(**adaptor_conf)
|
||
self.audio_adaptor.eval()
|
||
for p in self.audio_adaptor.parameters():
|
||
p.requires_grad = False
|
||
self.use_low_frame_rate = adaptor_conf.get("use_low_frame_rate", False)
|
||
|
||
# --- CTC Decoder (optional, for timestamps) ---
|
||
self.ctc_decoder = None
|
||
self.ctc = None
|
||
self.ctc_tokenizer = None
|
||
self.blank_id = None
|
||
|
||
ctc_decoder_name = self._config.get("ctc_decoder", None)
|
||
if ctc_decoder_name:
|
||
ctc_decoder_class = tables.adaptor_classes.get(ctc_decoder_name)
|
||
ctc_decoder_conf = self._config.get("ctc_decoder_conf", {})
|
||
if self.audio_encoder_output_size > 0:
|
||
ctc_decoder_conf["encoder_dim"] = self.audio_encoder_output_size
|
||
self.ctc_decoder = ctc_decoder_class(**ctc_decoder_conf)
|
||
self.ctc_decoder.eval()
|
||
for p in self.ctc_decoder.parameters():
|
||
p.requires_grad = False
|
||
|
||
from funasr.models.fun_asr_nano.ctc import CTC
|
||
|
||
ctc_conf = self._config.get("ctc_conf", {})
|
||
ctc_vocab_size = self._config.get("ctc_vocab_size", 60515)
|
||
self.blank_id = ctc_conf.get("blank_id", ctc_vocab_size - 1)
|
||
self.ctc = CTC(
|
||
odim=ctc_vocab_size,
|
||
encoder_output_size=self.audio_encoder_output_size,
|
||
blank_id=self.blank_id,
|
||
**ctc_conf,
|
||
)
|
||
|
||
# CTC tokenizer
|
||
ds_conf = self._config.get("dataset_conf", {})
|
||
ctc_tokenizer_name = ds_conf.get("ctc_tokenizer", None)
|
||
ctc_tokenizer_conf = ds_conf.get("ctc_tokenizer_conf", {})
|
||
if ctc_tokenizer_name:
|
||
ctc_tokenizer_class = tables.tokenizer_classes.get(ctc_tokenizer_name)
|
||
vocab_path = ctc_tokenizer_conf.get("vocab_path")
|
||
if vocab_path is None or not os.path.isabs(vocab_path):
|
||
multilingual_path = os.path.join(model_dir, "multilingual.tiktoken")
|
||
if os.path.exists(multilingual_path):
|
||
ctc_tokenizer_conf["vocab_path"] = multilingual_path
|
||
elif vocab_path and not os.path.isabs(vocab_path):
|
||
ctc_tokenizer_conf["vocab_path"] = os.path.join(model_dir, vocab_path)
|
||
self.ctc_tokenizer = ctc_tokenizer_class(**ctc_tokenizer_conf)
|
||
|
||
# --- Load weights from model.pt ---
|
||
model_pt = os.path.join(model_dir, "model.pt")
|
||
if os.path.exists(model_pt):
|
||
logger.info(f"Loading audio component weights from {model_pt}")
|
||
checkpoint = torch.load(model_pt, map_location="cpu")
|
||
state_dict = checkpoint.get("state_dict", checkpoint)
|
||
|
||
# Audio encoder
|
||
enc_state = {
|
||
k[len("audio_encoder."):]: v
|
||
for k, v in state_dict.items()
|
||
if k.startswith("audio_encoder.")
|
||
}
|
||
if enc_state:
|
||
self.audio_encoder.load_state_dict(enc_state, strict=False)
|
||
logger.info(f" Loaded audio_encoder: {len(enc_state)} params")
|
||
|
||
# Audio adaptor
|
||
adp_state = {
|
||
k[len("audio_adaptor."):]: v
|
||
for k, v in state_dict.items()
|
||
if k.startswith("audio_adaptor.")
|
||
}
|
||
if adp_state:
|
||
self.audio_adaptor.load_state_dict(adp_state, strict=False)
|
||
logger.info(f" Loaded audio_adaptor: {len(adp_state)} params")
|
||
|
||
# CTC decoder
|
||
if self.ctc_decoder is not None:
|
||
ctc_dec_state = {
|
||
k[len("ctc_decoder."):]: v
|
||
for k, v in state_dict.items()
|
||
if k.startswith("ctc_decoder.")
|
||
}
|
||
if ctc_dec_state:
|
||
self.ctc_decoder.load_state_dict(ctc_dec_state, strict=False)
|
||
ctc_state = {
|
||
k[len("ctc."):]: v
|
||
for k, v in state_dict.items()
|
||
if k.startswith("ctc.") and not k.startswith("ctc_decoder.")
|
||
}
|
||
if ctc_state:
|
||
self.ctc.load_state_dict(ctc_state, strict=False)
|
||
|
||
# Move to device
|
||
self.audio_encoder = self.audio_encoder.to(self.device, dtype=torch.float32)
|
||
self.audio_adaptor = self.audio_adaptor.to(self.device, dtype=self.torch_dtype)
|
||
if self.ctc_decoder is not None:
|
||
self.ctc_decoder = self.ctc_decoder.to(self.device, dtype=torch.float32)
|
||
self.ctc = self.ctc.to(self.device, dtype=torch.float32)
|
||
|
||
def _load_embedding_layer(self, model_dir: str):
|
||
"""Load the LLM embedding layer for text token embedding computation."""
|
||
model_pt = os.path.join(model_dir, "model.pt")
|
||
checkpoint = torch.load(model_pt, map_location="cpu")
|
||
state_dict = checkpoint.get("state_dict", checkpoint)
|
||
|
||
# Look for embedding weights
|
||
embed_key = None
|
||
for key in state_dict.keys():
|
||
if "embed_tokens.weight" in key and key.startswith("llm."):
|
||
embed_key = key
|
||
break
|
||
|
||
if embed_key is None:
|
||
raise RuntimeError("Could not find LLM embedding weights in model.pt")
|
||
|
||
embed_weight = state_dict[embed_key]
|
||
self.embed_tokens = nn.Embedding.from_pretrained(embed_weight, freeze=True)
|
||
self.embed_tokens = self.embed_tokens.to(self.device, dtype=self.torch_dtype)
|
||
logger.info(f"Loaded embedding layer: {embed_weight.shape}")
|
||
|
||
@torch.no_grad()
|
||
def _encode_audio(self, audio_input: Union[str, torch.Tensor, np.ndarray]):
|
||
"""Encode audio through frontend -> encoder -> adaptor.
|
||
|
||
Returns:
|
||
adaptor_out: (1, T', D_llm) audio embeddings for LLM input
|
||
adaptor_out_lens: (1,) lengths
|
||
encoder_out: (1, T, D_enc) encoder output for CTC
|
||
encoder_out_lens: (1,) encoder output lengths
|
||
"""
|
||
from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
|
||
|
||
if isinstance(audio_input, str):
|
||
data_src = load_audio_text_image_video(audio_input, fs=self.frontend.fs)
|
||
elif isinstance(audio_input, np.ndarray):
|
||
data_src = torch.from_numpy(audio_input).float()
|
||
elif isinstance(audio_input, torch.Tensor):
|
||
data_src = audio_input.float()
|
||
else:
|
||
raise ValueError(f"Unsupported audio input type: {type(audio_input)}")
|
||
|
||
speech, speech_lengths = extract_fbank(
|
||
data_src, data_type="sound", frontend=self.frontend, is_final=True
|
||
)
|
||
speech = speech.to(self.device, dtype=torch.float32)
|
||
speech_lengths = speech_lengths.to(self.device)
|
||
|
||
encoder_out, encoder_out_lens = self.audio_encoder(speech, speech_lengths)
|
||
encoder_out_for_adaptor = encoder_out.to(dtype=self.torch_dtype)
|
||
adaptor_out, adaptor_out_lens = self.audio_adaptor(encoder_out_for_adaptor, encoder_out_lens)
|
||
|
||
# Apply low frame rate: compute effective token count from fbank length
|
||
# Matches PyTorch model.py data_load_speech formula exactly
|
||
if self.use_low_frame_rate:
|
||
for i in range(adaptor_out.shape[0]):
|
||
fbank_len = speech_lengths[i].item()
|
||
olens = 1 + (fbank_len - 3 + 2 * 1) // 2
|
||
olens = 1 + (olens - 3 + 2 * 1) // 2
|
||
fake_token_len = (olens - 1) // 2 + 1
|
||
adaptor_out_lens[i] = fake_token_len
|
||
|
||
return adaptor_out, adaptor_out_lens, encoder_out, encoder_out_lens
|
||
|
||
def _build_prompt_text(
|
||
self,
|
||
hotwords: List[str] = None,
|
||
language: str = None,
|
||
itn: bool = True,
|
||
) -> str:
|
||
"""Build the ASR prompt string."""
|
||
hotwords = hotwords or []
|
||
if len(hotwords) > 0:
|
||
hotwords_str = ", ".join(hotwords)
|
||
prompt = (
|
||
"请结合上下文信息，更加准确地完成语音转写任务。"
|
||
"如果没有相关信息，我们会留空。\n\n\n**上下文信息：**\n\n\n"
|
||
)
|
||
prompt += f"热词列表：[{hotwords_str}]\n"
|
||
else:
|
||
prompt = ""
|
||
if language is None:
|
||
prompt += "语音转写"
|
||
else:
|
||
prompt += f"语音转写成{language}"
|
||
if not itn:
|
||
prompt += "，不进行文本规整"
|
||
return prompt + "："
|
||
|
||
@torch.no_grad()
|
||
def _build_input_embeds(
|
||
self,
|
||
audio_embeds: torch.Tensor,
|
||
audio_embed_lens: torch.Tensor,
|
||
hotwords: List[str] = None,
|
||
language: str = None,
|
||
itn: bool = True,
|
||
system_prompt: str = "You are a helpful assistant.",
|
||
) -> torch.Tensor:
|
||
"""Build the full input embedding sequence with audio inserted.
|
||
|
||
Returns:
|
||
Tensor of shape (seq_len, D_llm)
|
||
"""
|
||
prompt = self._build_prompt_text(hotwords, language, itn)
|
||
|
||
# ChatML format with speech markers and thinking prefix
|
||
prefix_text = (
|
||
f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
|
||
f"<|im_start|>user\n{prompt}<|startofspeech|>"
|
||
)
|
||
suffix_text = "<|endofspeech|><|im_end|>\n<|im_start|>assistant\n"
|
||
|
||
# Tokenize
|
||
prefix_ids = self.tokenizer.encode(prefix_text, add_special_tokens=False)
|
||
suffix_ids = self.tokenizer.encode(suffix_text, add_special_tokens=False)
|
||
|
||
# Embed text tokens
|
||
prefix_tensor = torch.tensor(prefix_ids, dtype=torch.long, device=self.device)
|
||
suffix_tensor = torch.tensor(suffix_ids, dtype=torch.long, device=self.device)
|
||
prefix_embeds = self.embed_tokens(prefix_tensor)
|
||
suffix_embeds = self.embed_tokens(suffix_tensor)
|
||
|
||
# Audio embeddings
|
||
audio_len = audio_embed_lens[0].item()
|
||
audio_emb = audio_embeds[0, :audio_len, :]
|
||
|
||
# Concat: [prefix_text_emb | audio_emb | suffix_text_emb]
|
||
inputs_embeds = torch.cat([prefix_embeds, audio_emb, suffix_embeds], dim=0)
|
||
return inputs_embeds
|
||
|
||
def generate(
|
||
self,
|
||
inputs: Union[str, List[str], np.ndarray, torch.Tensor, List],
|
||
hotwords: List[str] = None,
|
||
language: str = None,
|
||
itn: bool = True,
|
||
max_new_tokens: int = 512,
|
||
temperature: float = 0.0,
|
||
top_p: float = 1.0,
|
||
top_k: int = -1,
|
||
repetition_penalty: float = 1.0,
|
||
**kwargs,
|
||
) -> List[dict]:
|
||
"""Run batch ASR inference using vLLM.
|
||
|
||
Args:
|
||
inputs: Audio input(s). Accepts:
|
||
- str: single file path
|
||
- List[str]: batch of file paths
|
||
- np.ndarray / torch.Tensor: raw audio samples (16kHz)
|
||
hotwords: Keywords to boost recognition accuracy.
|
||
language: Language hint (e.g. "中文", "英文", "日文").
|
||
itn: Apply inverse text normalization (default True).
|
||
max_new_tokens: Maximum tokens to generate per sample.
|
||
temperature: Sampling temperature (0 = greedy decoding).
|
||
top_p: Nucleus sampling parameter.
|
||
top_k: Top-k sampling (-1 = disabled).
|
||
repetition_penalty: Repetition penalty factor.
|
||
|
||
Returns:
|
||
List of result dicts: [{"key": str, "text": str, "timestamps": [...]}]
|
||
"""
|
||
from vllm import SamplingParams
|
||
try:
|
||
from vllm.inputs import EmbedsPrompt
|
||
except ImportError:
|
||
from vllm.inputs.data import EmbedsPrompt
|
||
|
||
from funasr.models.fun_asr_nano.vllm_utils import resolve_repetition_penalty
|
||
|
||
if isinstance(inputs, (str, np.ndarray, torch.Tensor)):
|
||
inputs = [inputs]
|
||
|
||
sampling_params = SamplingParams(
|
||
max_tokens=max_new_tokens,
|
||
temperature=temperature,
|
||
top_p=top_p,
|
||
top_k=top_k if top_k > 0 else -1,
|
||
# Prompt-embeds mode has no token IDs to penalize; see #2948.
|
||
repetition_penalty=resolve_repetition_penalty(repetition_penalty),
|
||
skip_special_tokens=True,
|
||
)
|
||
|
||
# Batch encode audio and build embedding prompts
|
||
prompts = []
|
||
encoder_outputs = []
|
||
|
||
t0 = time.perf_counter()
|
||
|
||
# Pre-compute text embeddings (shared across batch)
|
||
prompt_text = self._build_prompt_text(hotwords, language, itn)
|
||
prefix_text = f"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n{prompt_text}"
|
||
suffix_text = "<|im_end|>\n<|im_start|>assistant\n"
|
||
prefix_ids = self.tokenizer.encode(prefix_text, add_special_tokens=False)
|
||
suffix_ids = self.tokenizer.encode(suffix_text, add_special_tokens=False)
|
||
prefix_emb = self.embed_tokens(torch.tensor(prefix_ids, dtype=torch.long, device=self.device))
|
||
suffix_emb = self.embed_tokens(torch.tensor(suffix_ids, dtype=torch.long, device=self.device))
|
||
|
||
# Batch encode audio (groups of 8 for memory efficiency)
|
||
batch_size_enc = 8
|
||
all_adaptor_outs = []
|
||
all_adaptor_lens = []
|
||
for i in range(0, len(inputs), batch_size_enc):
|
||
batch_inputs = inputs[i:i+batch_size_enc]
|
||
# Load and extract fbank for batch
|
||
from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
|
||
audio_tensors = []
|
||
for audio_input in batch_inputs:
|
||
if isinstance(audio_input, str):
|
||
data_src = load_audio_text_image_video(audio_input, fs=self.frontend.fs)
|
||
elif isinstance(audio_input, np.ndarray):
|
||
data_src = torch.from_numpy(audio_input).float()
|
||
elif isinstance(audio_input, torch.Tensor):
|
||
data_src = audio_input.float()
|
||
else:
|
||
raise ValueError(f"Unsupported audio input type: {type(audio_input)}")
|
||
audio_tensors.append(data_src)
|
||
|
||
speech, speech_lengths = extract_fbank(
|
||
audio_tensors, data_type="sound", frontend=self.frontend, is_final=True
|
||
)
|
||
speech = speech.to(self.device, dtype=torch.float32)
|
||
speech_lengths = speech_lengths.to(self.device)
|
||
|
||
with torch.no_grad():
|
||
enc_out, enc_lens = self.audio_encoder(speech, speech_lengths)
|
||
adp_out, adp_lens = self.audio_adaptor(enc_out.to(dtype=self.torch_dtype), enc_lens)
|
||
|
||
# Apply low frame rate token length correction
|
||
if self.use_low_frame_rate:
|
||
for j in range(len(batch_inputs)):
|
||
fbank_len = speech_lengths[j].item()
|
||
olens = 1 + (fbank_len - 3 + 2 * 1) // 2
|
||
olens = 1 + (olens - 3 + 2 * 1) // 2
|
||
adp_lens[j] = (olens - 1) // 2 + 1
|
||
|
||
for j in range(len(batch_inputs)):
|
||
all_adaptor_outs.append(adp_out[j, :adp_lens[j], :])
|
||
all_adaptor_lens.append(adp_lens[j])
|
||
encoder_outputs.append((enc_out[j:j+1, :enc_lens[j], :], enc_lens[j:j+1]))
|
||
|
||
# Build prompts
|
||
for audio_emb in all_adaptor_outs:
|
||
input_embeds = torch.cat([prefix_emb, audio_emb, suffix_emb], dim=0)
|
||
prompts.append(EmbedsPrompt(prompt_embeds=input_embeds.float()))
|
||
|
||
t1 = time.perf_counter()
|
||
logger.info(f"Audio encoding: {len(inputs)} samples in {t1 - t0:.3f}s")
|
||
|
||
# vLLM batch generation
|
||
outputs = self.vllm_engine.generate(prompts, sampling_params, use_tqdm=len(inputs) > 1)
|
||
|
||
t2 = time.perf_counter()
|
||
logger.info(f"vLLM generation: {t2 - t1:.3f}s")
|
||
|
||
# Process results
|
||
results = []
|
||
for i, output in enumerate(outputs):
|
||
token_ids = list(output.outputs[0].token_ids)
|
||
text = self.tokenizer.decode(token_ids, skip_special_tokens=True)
|
||
# Clean vLLM artifacts: remove garbage prefix/tags
|
||
text = re.sub(r'<[^>]*>', '', text)
|
||
text = re.sub(r'\[[^\]]*\]', '', text)
|
||
text = re.sub(r'endofpatch|/sil|FFFF|</strong>', '', text)
|
||
# Strip non-CJK/non-alnum prefix garbage
|
||
text = re.sub(r'^[^\w一-鿿]+', '', text)
|
||
text_clean = re.sub(r"\s+", " ", text).strip()
|
||
|
||
key = (
|
||
os.path.splitext(os.path.basename(inputs[i]))[0]
|
||
if isinstance(inputs[i], str)
|
||
else f"sample_{i}"
|
||
)
|
||
result = {"key": key, "text": text_clean}
|
||
|
||
# Timestamps via CTC forced alignment
|
||
if self.ctc_decoder is not None and self.ctc_tokenizer is not None:
|
||
try:
|
||
timestamps = self._compute_timestamps(
|
||
encoder_outputs[i][0], encoder_outputs[i][1], text_clean
|
||
)
|
||
if timestamps:
|
||
result["timestamps"] = timestamps
|
||
except Exception as e:
|
||
logger.debug(f"Timestamp computation failed for {key}: {e}")
|
||
|
||
results.append(result)
|
||
|
||
return results
|
||
|
||
@torch.no_grad()
|
||
def _compute_timestamps(self, encoder_out, encoder_out_lens, text):
|
||
"""CTC forced alignment for character-level timestamps."""
|
||
from funasr.models.fun_asr_nano.tools.utils import forced_align
|
||
|
||
decoder_out, decoder_out_lens = self.ctc_decoder(encoder_out, encoder_out_lens)
|
||
ctc_logits = self.ctc.log_softmax(decoder_out)
|
||
x = ctc_logits[0, : encoder_out_lens[0].item(), :]
|
||
|
||
target_ids = torch.tensor(self.ctc_tokenizer.encode(text), dtype=torch.int64)
|
||
if len(target_ids) == 0:
|
||
return []
|
||
|
||
timestamps = forced_align(x, target_ids, self.blank_id)
|
||
for ts in timestamps:
|
||
ts["token"] = self.ctc_tokenizer.decode([ts["token"]])
|
||
ts["start_time"] = ts["start_time"] * 6 * 10 / 1000
|
||
ts["end_time"] = ts["end_time"] * 6 * 10 / 1000
|
||
return timestamps
|
||
|
||
@classmethod
|
||
def from_pretrained(
|
||
cls,
|
||
model: str = "FunAudioLLM/Fun-ASR-Nano-2512",
|
||
hub: str = "ms",
|
||
device: str = "cuda:0",
|
||
dtype: str = "bf16",
|
||
tensor_parallel_size: int = 1,
|
||
gpu_memory_utilization: float = 0.8,
|
||
max_model_len: int = 2048,
|
||
**kwargs,
|
||
) -> "FunASRNanoVLLM":
|
||
"""Load model from hub or local path.
|
||
|
||
Args:
|
||
model: Model name or local directory path.
|
||
hub: "ms" (ModelScope) or "hf" (HuggingFace).
|
||
device: Device for audio encoder/adaptor.
|
||
dtype: Compute dtype ("bf16", "fp16", "fp32").
|
||
tensor_parallel_size: GPUs for vLLM tensor parallel.
|
||
gpu_memory_utilization: GPU memory fraction for vLLM.
|
||
max_model_len: Maximum sequence length.
|
||
|
||
Returns:
|
||
Initialized FunASRNanoVLLM engine.
|
||
"""
|
||
if os.path.isdir(model):
|
||
model_dir = model
|
||
else:
|
||
if hub in ("ms", "modelscope"):
|
||
from modelscope.hub.snapshot_download import snapshot_download
|
||
|
||
model_dir = snapshot_download(model, revision=kwargs.pop("revision", "master"))
|
||
elif hub in ("hf", "huggingface"):
|
||
from huggingface_hub import snapshot_download
|
||
|
||
model_dir = snapshot_download(model)
|
||
else:
|
||
raise ValueError(f"Unsupported hub: {hub}. Use 'ms' or 'hf'.")
|
||
|
||
logger.info(f"Model directory: {model_dir}")
|
||
return cls(
|
||
model_dir=model_dir,
|
||
device=device,
|
||
dtype=dtype,
|
||
tensor_parallel_size=tensor_parallel_size,
|
||
gpu_memory_utilization=gpu_memory_utilization,
|
||
max_model_len=max_model_len,
|
||
**kwargs,
|
||
)
|