diff --git a/examples/industrial_data_pretraining/fun_asr_nano/benchmark_vllm.py b/examples/industrial_data_pretraining/fun_asr_nano/benchmark_vllm.py new file mode 100644 index 000000000..8da6de601 --- /dev/null +++ b/examples/industrial_data_pretraining/fun_asr_nano/benchmark_vllm.py @@ -0,0 +1,186 @@ +#!/usr/bin/env python3 +"""Fun-ASR-Nano vLLM Benchmark: Speed + CER comparison. + +Usage: + # Full benchmark (PyTorch + vLLM + CER) + CUDA_VISIBLE_DEVICES=0 python benchmark_vllm.py \ + --audio-dir /path/to/benchmark_audio \ + --label-json /path/to/benchmark_testset.json + + # Quick test (first 20 files) + CUDA_VISIBLE_DEVICES=0 python benchmark_vllm.py \ + --audio-dir /path/to/benchmark_audio \ + --label-json /path/to/benchmark_testset.json \ + --max-files 20 +""" + +import argparse +import json +import os +import re +import time + +import kaldialign +import numpy as np +import soundfile as sf +import torch + + +def normalize_zh(text): + """Remove punctuation for CER calculation.""" + text = re.sub(r'[^\w一-鿿]', '', text) + return text.upper() + + +def compute_cer(refs, hyps): + """Compute CER between reference and hypothesis lists.""" + total_ref = 0 + total_errs = 0 + for ref, hyp in zip(refs, hyps): + r = list(normalize_zh(ref)) + h = list(normalize_zh(hyp)) + total_ref += len(r) + ali = kaldialign.align(r, h, '*') + total_errs += sum(1 for a, b in ali if a != b) + return total_errs / total_ref * 100 if total_ref > 0 else 0 + + +def vad_segment(files, device="cuda:0"): + """Pre-segment all files with fsmn-vad.""" + from funasr import AutoModel + + vad_model = AutoModel(model="fsmn-vad", device=device, disable_update=True) + all_segments = [] + for fi, wav_path in enumerate(files): + audio, sr = sf.read(wav_path) + if audio.ndim > 1: + audio = audio[:, 0] + audio = audio.astype(np.float32) + res = vad_model.generate(input=wav_path, dynamic_silence=False) + segs = res[0]["value"] + for seg in segs: + s0 = int(seg[0] * sr / 1000) + s1 = int(seg[1] * sr / 1000) + seg_audio = audio[s0:s1] + if len(seg_audio) > sr * 0.5: + all_segments.append((fi, seg[0], seg[1], seg_audio)) + return all_segments + + +def concat_results(all_segments, seg_texts, n_files): + """Concatenate segment texts back to per-file results.""" + file_texts = {} + for (fi, _, _, _), text in zip(all_segments, seg_texts): + if fi not in file_texts: + file_texts[fi] = [] + file_texts[fi].append(text) + return ["".join(file_texts.get(fi, [])) for fi in range(n_files)] + + +def run_pytorch(seg_files, device="cuda:0"): + """PyTorch native inference.""" + from funasr import AutoModel + + model = AutoModel( + model="FunAudioLLM/Fun-ASR-Nano-2512", trust_remote_code=True, + remote_code=os.path.join(os.path.dirname(__file__), "model.py"), + device=device, disable_update=True, + ) + model.generate(input=seg_files[0]) # warmup + + t0 = time.perf_counter() + texts = [] + for f in seg_files: + res = model.generate(input=f) + texts.append(res[0]["text"]) + t1 = time.perf_counter() + return t1 - t0, texts + + +def run_vllm(seg_files, device="cuda:0"): + """Our FunASRNanoVLLM batch inference.""" + from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM + + engine = FunASRNanoVLLM.from_pretrained( + model="FunAudioLLM/Fun-ASR-Nano-2512", device=device, dtype="bf16", + max_model_len=4096, gpu_memory_utilization=0.5, + ) + engine.generate(inputs=[seg_files[0]], language="中文") # warmup + + t0 = time.perf_counter() + results = engine.generate(inputs=seg_files, language="中文", max_new_tokens=500) + t1 = time.perf_counter() + texts = [r["text"] for r in results] + return t1 - t0, texts + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description="Fun-ASR-Nano vLLM Benchmark") + parser.add_argument("--audio-dir", type=str, required=True) + parser.add_argument("--label-json", type=str, required=True) + parser.add_argument("--device", type=str, default="cuda:0") + parser.add_argument("--max-files", type=int, default=0, help="Limit files (0=all)") + parser.add_argument("--skip-pytorch", action="store_true") + args = parser.parse_args() + + with open(args.label_json) as f: + dataset = json.load(f) + + files = [] + refs = [] + for item in dataset: + wav_path = os.path.join(args.audio_dir, f"{item['id']:03d}.wav") + if os.path.exists(wav_path): + files.append(wav_path) + refs.append(item["ref"]) + if args.max_files > 0: + files = files[:args.max_files] + refs = refs[:args.max_files] + + total_audio = sum(sf.info(f).duration for f in files) + print(f"Dataset: {len(files)} files, {total_audio:.0f}s audio") + + # VAD pre-segment + print("\n>>> VAD pre-segmenting...") + all_segments = vad_segment(files, device=args.device) + print(f" {len(all_segments)} segments (avg {sum(len(s[3])/16000 for s in all_segments)/len(all_segments):.1f}s)") + + # Save segments + os.makedirs("/tmp/benchmark_segs", exist_ok=True) + seg_files = [] + for i, (fi, s, e, audio) in enumerate(all_segments): + path = f"/tmp/benchmark_segs/{i:04d}.wav" + sf.write(path, audio, 16000) + seg_files.append(path) + + # PyTorch + if not args.skip_pytorch: + print("\n>>> PyTorch native...") + pt_time, pt_seg_texts = run_pytorch(seg_files, device=args.device) + pt_texts = concat_results(all_segments, pt_seg_texts, len(files)) + cer_pt = compute_cer(refs, pt_texts) + print(f" Time: {pt_time:.1f}s | RTFx: {total_audio/pt_time:.1f} | CER: {cer_pt:.2f}%") + + del torch.cuda.memory_allocated + torch.cuda.empty_cache() + + # vLLM + print("\n>>> vLLM (FunASRNanoVLLM batch)...") + vllm_time, vllm_seg_texts = run_vllm(seg_files, device=args.device) + vllm_texts = concat_results(all_segments, vllm_seg_texts, len(files)) + cer_vllm = compute_cer(refs, vllm_texts) + print(f" Time: {vllm_time:.1f}s | RTFx: {total_audio/vllm_time:.1f} | CER: {cer_vllm:.2f}%") + + # Summary + print(f"\n{'='*60}") + print(f"BENCHMARK RESULTS ({len(files)} files, {total_audio:.0f}s)") + print(f"{'-'*60}") + print(f"{'Method':<20} {'Time':<10} {'RTFx':<10} {'CER'}") + print(f"{'-'*60}") + if not args.skip_pytorch: + print(f"{'PyTorch':<20} {pt_time:<10.1f} {total_audio/pt_time:<10.1f} {cer_pt:.2f}%") + print(f"{'vLLM (ours)':<20} {vllm_time:<10.1f} {total_audio/vllm_time:<10.1f} {cer_vllm:.2f}%") + if not args.skip_pytorch: + print(f"{'-'*60}") + print(f"Speedup: {total_audio/vllm_time / (total_audio/pt_time):.1f}x | CER diff: {cer_vllm - cer_pt:+.2f}%") + print(f"{'='*60}") diff --git a/funasr/models/fun_asr_nano/inference_vllm.py b/funasr/models/fun_asr_nano/inference_vllm.py index 927da7b5c..38384b509 100644 --- a/funasr/models/fun_asr_nano/inference_vllm.py +++ b/funasr/models/fun_asr_nano/inference_vllm.py @@ -406,6 +406,17 @@ class FunASRNanoVLLM: 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( @@ -450,12 +461,12 @@ class FunASRNanoVLLM: """ prompt = self._build_prompt_text(hotwords, language, itn) - # ChatML format + # 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\n\n\n\n" + suffix_text = "<|endofspeech|><|im_end|>\n<|im_start|>assistant\n" # Tokenize prefix_ids = self.tokenizer.encode(prefix_text, add_special_tokens=False) @@ -522,21 +533,67 @@ class FunASRNanoVLLM: skip_special_tokens=True, ) - # Encode all audio samples and build embedding prompts + # Batch encode audio and build embedding prompts prompts = [] encoder_outputs = [] t0 = time.perf_counter() - for audio_input in inputs: - adaptor_out, adaptor_out_lens, encoder_out, encoder_out_lens = self._encode_audio( - audio_input - ) - encoder_outputs.append((encoder_out, encoder_out_lens)) - input_embeds = self._build_input_embeds( - adaptor_out, adaptor_out_lens, hotwords, language, itn + # 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 ) - # vLLM EmbedsPrompt expects float32 or the model's dtype + 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() @@ -551,10 +608,15 @@ class FunASRNanoVLLM: # Process results results = [] for i, output in enumerate(outputs): - text = output.outputs[0].text - if not text and output.outputs[0].token_ids: - text = self.tokenizer.decode(list(output.outputs[0].token_ids), skip_special_tokens=True) - text_clean = re.sub(r"\s+", " ", text.replace("/sil", " ")).strip() + 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|', '', 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]