|
| 1 | +#!/usr/bin/env python3 |
| 2 | +"""FunASR vLLM Benchmark: unified speed + CER comparison for all supported models. |
| 3 | +
|
| 4 | +Supports: Fun-ASR-Nano, GLM-ASR-Nano (and any model via AutoModelVLLM). |
| 5 | +
|
| 6 | +Usage: |
| 7 | + # Fun-ASR-Nano |
| 8 | + CUDA_VISIBLE_DEVICES=0 python benchmark_vllm.py \ |
| 9 | + --model FunAudioLLM/Fun-ASR-Nano-2512 \ |
| 10 | + --audio-dir /path/to/benchmark_audio \ |
| 11 | + --label-json /path/to/benchmark_testset.json |
| 12 | +
|
| 13 | + # GLM-ASR-Nano |
| 14 | + CUDA_VISIBLE_DEVICES=0 python benchmark_vllm.py \ |
| 15 | + --model zai-org/GLM-ASR-Nano-2512 \ |
| 16 | + --audio-dir /path/to/benchmark_audio \ |
| 17 | + --label-json /path/to/benchmark_testset.json |
| 18 | +
|
| 19 | + # Quick test (first N files) |
| 20 | + CUDA_VISIBLE_DEVICES=0 python benchmark_vllm.py \ |
| 21 | + --model FunAudioLLM/Fun-ASR-Nano-2512 \ |
| 22 | + --audio-dir /path/to/benchmark_audio \ |
| 23 | + --label-json /path/to/benchmark_testset.json \ |
| 24 | + --max-files 20 |
| 25 | +
|
| 26 | + # Skip PyTorch (only test vLLM) |
| 27 | + CUDA_VISIBLE_DEVICES=0 python benchmark_vllm.py \ |
| 28 | + --model FunAudioLLM/Fun-ASR-Nano-2512 \ |
| 29 | + --skip-pytorch \ |
| 30 | + --audio-dir /path/to/benchmark_audio \ |
| 31 | + --label-json /path/to/benchmark_testset.json |
| 32 | +""" |
| 33 | + |
| 34 | +import argparse |
| 35 | +import json |
| 36 | +import os |
| 37 | +import re |
| 38 | +import time |
| 39 | + |
| 40 | +import kaldialign |
| 41 | +import numpy as np |
| 42 | +import soundfile as sf |
| 43 | +import torch |
| 44 | + |
| 45 | + |
| 46 | +def normalize_zh(text): |
| 47 | + text = re.sub(r'[^\w一-鿿]', '', text) |
| 48 | + return text.upper() |
| 49 | + |
| 50 | + |
| 51 | +def compute_cer(refs, hyps): |
| 52 | + total_ref = 0 |
| 53 | + total_errs = 0 |
| 54 | + for ref, hyp in zip(refs, hyps): |
| 55 | + r = list(normalize_zh(ref)) |
| 56 | + h = list(normalize_zh(hyp)) |
| 57 | + total_ref += len(r) |
| 58 | + ali = kaldialign.align(r, h, '*') |
| 59 | + total_errs += sum(1 for a, b in ali if a != b) |
| 60 | + return total_errs / total_ref * 100 if total_ref > 0 else 0 |
| 61 | + |
| 62 | + |
| 63 | +def vad_segment(files, device="cuda:0"): |
| 64 | + from funasr import AutoModel |
| 65 | + vad_model = AutoModel(model="fsmn-vad", device=device, disable_update=True) |
| 66 | + all_segments = [] |
| 67 | + for fi, wav_path in enumerate(files): |
| 68 | + audio, sr = sf.read(wav_path) |
| 69 | + if audio.ndim > 1: |
| 70 | + audio = audio[:, 0] |
| 71 | + audio = audio.astype(np.float32) |
| 72 | + res = vad_model.generate(input=wav_path, dynamic_silence=False) |
| 73 | + for seg in res[0]["value"]: |
| 74 | + s0 = int(seg[0] * sr / 1000) |
| 75 | + s1 = int(seg[1] * sr / 1000) |
| 76 | + seg_audio = audio[s0:s1] |
| 77 | + if len(seg_audio) > sr * 0.5: |
| 78 | + all_segments.append((fi, seg_audio)) |
| 79 | + return all_segments |
| 80 | + |
| 81 | + |
| 82 | +def concat_results(all_segments, seg_texts, n_files): |
| 83 | + file_texts = {} |
| 84 | + for (fi, _), text in zip(all_segments, seg_texts): |
| 85 | + file_texts.setdefault(fi, []).append(text) |
| 86 | + return ["".join(file_texts.get(fi, [])) for fi in range(n_files)] |
| 87 | + |
| 88 | + |
| 89 | +def run_pytorch(model_name, seg_files, device="cuda:0"): |
| 90 | + from funasr import AutoModel |
| 91 | + |
| 92 | + kwargs = {"model": model_name, "device": device, "disable_update": True} |
| 93 | + if "Fun-ASR-Nano" in model_name: |
| 94 | + kwargs["trust_remote_code"] = True |
| 95 | + kwargs["remote_code"] = os.path.join( |
| 96 | + os.path.dirname(__file__), |
| 97 | + "examples/industrial_data_pretraining/fun_asr_nano/model.py" |
| 98 | + ) |
| 99 | + |
| 100 | + model = AutoModel(**kwargs) |
| 101 | + model.generate(input=seg_files[0]) # warmup |
| 102 | + |
| 103 | + t0 = time.perf_counter() |
| 104 | + texts = [] |
| 105 | + for f in seg_files: |
| 106 | + res = model.generate(input=f) |
| 107 | + texts.append(res[0]["text"]) |
| 108 | + t1 = time.perf_counter() |
| 109 | + return t1 - t0, texts |
| 110 | + |
| 111 | + |
| 112 | +def run_vllm(model_name, seg_files, device="cuda:0", hub="ms"): |
| 113 | + if "Fun-ASR-Nano" in model_name: |
| 114 | + from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM |
| 115 | + engine = FunASRNanoVLLM.from_pretrained( |
| 116 | + model=model_name, hub=hub, device=device, dtype="bf16", |
| 117 | + max_model_len=4096, gpu_memory_utilization=0.5) |
| 118 | + engine.generate(inputs=[seg_files[0]], language="中文") # warmup |
| 119 | + t0 = time.perf_counter() |
| 120 | + results = engine.generate(inputs=seg_files, language="中文", max_new_tokens=500) |
| 121 | + t1 = time.perf_counter() |
| 122 | + texts = [r["text"] for r in results] |
| 123 | + |
| 124 | + elif "GLM-ASR" in model_name: |
| 125 | + from funasr.models.glm_asr.inference_vllm import GLMASRVLLMEngine |
| 126 | + engine = GLMASRVLLMEngine.from_pretrained( |
| 127 | + model=model_name, hub=hub, device=device, dtype="bf16", |
| 128 | + gpu_memory_utilization=0.4, max_model_len=4096) |
| 129 | + engine.generate(inputs=[seg_files[0]]) # warmup |
| 130 | + t0 = time.perf_counter() |
| 131 | + results = engine.generate(inputs=seg_files, max_new_tokens=500) |
| 132 | + t1 = time.perf_counter() |
| 133 | + texts = [r["text"] for r in results] |
| 134 | + |
| 135 | + else: |
| 136 | + from funasr.auto.auto_model_vllm import AutoModelVLLM |
| 137 | + engine = AutoModelVLLM(model=model_name, hub=hub, device=device) |
| 138 | + engine.generate(inputs=[seg_files[0]]) |
| 139 | + t0 = time.perf_counter() |
| 140 | + results = engine.generate(inputs=seg_files, max_new_tokens=500) |
| 141 | + t1 = time.perf_counter() |
| 142 | + texts = [r["text"] for r in results] |
| 143 | + |
| 144 | + return t1 - t0, texts |
| 145 | + |
| 146 | + |
| 147 | +if __name__ == '__main__': |
| 148 | + parser = argparse.ArgumentParser(description="FunASR vLLM Benchmark") |
| 149 | + parser.add_argument("--model", type=str, required=True, help="Model name or path") |
| 150 | + parser.add_argument("--hub", type=str, default="ms", choices=["ms", "hf"]) |
| 151 | + parser.add_argument("--audio-dir", type=str, required=True) |
| 152 | + parser.add_argument("--label-json", type=str, required=True) |
| 153 | + parser.add_argument("--device", type=str, default="cuda:0") |
| 154 | + parser.add_argument("--max-files", type=int, default=0) |
| 155 | + parser.add_argument("--skip-pytorch", action="store_true") |
| 156 | + args = parser.parse_args() |
| 157 | + |
| 158 | + with open(args.label_json) as f: |
| 159 | + dataset = json.load(f) |
| 160 | + |
| 161 | + files = [] |
| 162 | + refs = [] |
| 163 | + for item in dataset: |
| 164 | + wav_path = os.path.join(args.audio_dir, f"{item['id']:03d}.wav") |
| 165 | + if os.path.exists(wav_path): |
| 166 | + files.append(wav_path) |
| 167 | + refs.append(item["ref"]) |
| 168 | + if args.max_files > 0: |
| 169 | + files = files[:args.max_files] |
| 170 | + refs = refs[:args.max_files] |
| 171 | + |
| 172 | + total_audio = sum(sf.info(f).duration for f in files) |
| 173 | + print(f"{'='*60}") |
| 174 | + print(f"FunASR vLLM Benchmark") |
| 175 | + print(f"{'='*60}") |
| 176 | + print(f"Model: {args.model}") |
| 177 | + print(f"Dataset: {len(files)} files, {total_audio:.0f}s audio") |
| 178 | + |
| 179 | + # VAD |
| 180 | + print(f"\n>>> VAD pre-segmenting...") |
| 181 | + all_segments = vad_segment(files, device=args.device) |
| 182 | + print(f" {len(all_segments)} segments") |
| 183 | + |
| 184 | + os.makedirs("/tmp/benchmark_vllm_segs", exist_ok=True) |
| 185 | + seg_files = [] |
| 186 | + for i, (fi, audio) in enumerate(all_segments): |
| 187 | + path = f"/tmp/benchmark_vllm_segs/{i:04d}.wav" |
| 188 | + sf.write(path, audio, 16000) |
| 189 | + seg_files.append(path) |
| 190 | + |
| 191 | + # PyTorch |
| 192 | + cer_pt = None |
| 193 | + pt_time = None |
| 194 | + if not args.skip_pytorch: |
| 195 | + print(f"\n>>> PyTorch native...") |
| 196 | + pt_time, pt_seg_texts = run_pytorch(args.model, seg_files, device=args.device) |
| 197 | + pt_texts = concat_results(all_segments, pt_seg_texts, len(files)) |
| 198 | + cer_pt = compute_cer(refs, pt_texts) |
| 199 | + print(f" Time: {pt_time:.1f}s | RTFx: {total_audio/pt_time:.1f} | CER: {cer_pt:.2f}%") |
| 200 | + |
| 201 | + del torch.cuda.memory_allocated |
| 202 | + torch.cuda.empty_cache() |
| 203 | + |
| 204 | + # vLLM |
| 205 | + print(f"\n>>> vLLM...") |
| 206 | + vllm_time, vllm_seg_texts = run_vllm(args.model, seg_files, device=args.device, hub=args.hub) |
| 207 | + vllm_texts = concat_results(all_segments, vllm_seg_texts, len(files)) |
| 208 | + cer_vllm = compute_cer(refs, vllm_texts) |
| 209 | + print(f" Time: {vllm_time:.1f}s | RTFx: {total_audio/vllm_time:.1f} | CER: {cer_vllm:.2f}%") |
| 210 | + |
| 211 | + # Summary |
| 212 | + print(f"\n{'='*60}") |
| 213 | + print(f"RESULTS") |
| 214 | + print(f"{'-'*60}") |
| 215 | + print(f"{'Method':<20} {'Time':<10} {'RTFx':<10} {'CER'}") |
| 216 | + print(f"{'-'*60}") |
| 217 | + if not args.skip_pytorch: |
| 218 | + print(f"{'PyTorch':<20} {pt_time:<10.1f} {total_audio/pt_time:<10.1f} {cer_pt:.2f}%") |
| 219 | + print(f"{'vLLM':<20} {vllm_time:<10.1f} {total_audio/vllm_time:<10.1f} {cer_vllm:.2f}%") |
| 220 | + if not args.skip_pytorch: |
| 221 | + print(f"{'-'*60}") |
| 222 | + speedup = (total_audio/vllm_time) / (total_audio/pt_time) |
| 223 | + print(f"Speedup: {speedup:.1f}x | CER diff: {cer_vllm - cer_pt:+.2f}%") |
| 224 | + print(f"{'='*60}") |
0 commit comments