Initial commit: FunASR Speech Recognition Toolkit

Add complete FunASR codebase including models, runtime, and documentation.

benchmark_vllm.py

@@ -0,0 +1,224 @@
+#!/usr/bin/env python3
+"""FunASR vLLM Benchmark: unified speed + CER comparison for all supported models.
+
+Supports: Fun-ASR-Nano, GLM-ASR-Nano (and any model via AutoModelVLLM).
+
+Usage:
+    # Fun-ASR-Nano
+    CUDA_VISIBLE_DEVICES=0 python benchmark_vllm.py \
+        --model FunAudioLLM/Fun-ASR-Nano-2512 \
+        --audio-dir /path/to/benchmark_audio \
+        --label-json /path/to/benchmark_testset.json
+
+    # GLM-ASR-Nano
+    CUDA_VISIBLE_DEVICES=0 python benchmark_vllm.py \
+        --model zai-org/GLM-ASR-Nano-2512 \
+        --audio-dir /path/to/benchmark_audio \
+        --label-json /path/to/benchmark_testset.json
+
+    # Quick test (first N files)
+    CUDA_VISIBLE_DEVICES=0 python benchmark_vllm.py \
+        --model FunAudioLLM/Fun-ASR-Nano-2512 \
+        --audio-dir /path/to/benchmark_audio \
+        --label-json /path/to/benchmark_testset.json \
+        --max-files 20
+
+    # Skip PyTorch (only test vLLM)
+    CUDA_VISIBLE_DEVICES=0 python benchmark_vllm.py \
+        --model FunAudioLLM/Fun-ASR-Nano-2512 \
+        --skip-pytorch \
+        --audio-dir /path/to/benchmark_audio \
+        --label-json /path/to/benchmark_testset.json
+"""
+
+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):
+    text = re.sub(r'[^\w一-鿿]', '', text)
+    return text.upper()
+
+
+def compute_cer(refs, hyps):
+    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"):
+    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)
+        for seg in res[0]["value"]:
+            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_audio))
+    return all_segments
+
+
+def concat_results(all_segments, seg_texts, n_files):
+    file_texts = {}
+    for (fi, _), text in zip(all_segments, seg_texts):
+        file_texts.setdefault(fi, []).append(text)
+    return ["".join(file_texts.get(fi, [])) for fi in range(n_files)]
+
+
+def run_pytorch(model_name, seg_files, device="cuda:0"):
+    from funasr import AutoModel
+
+    kwargs = {"model": model_name, "device": device, "disable_update": True}
+    if "Fun-ASR-Nano" in model_name:
+        kwargs["trust_remote_code"] = True
+        kwargs["remote_code"] = os.path.join(
+            os.path.dirname(__file__),
+            "examples/industrial_data_pretraining/fun_asr_nano/model.py"
+        )
+
+    model = AutoModel(**kwargs)
+    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(model_name, seg_files, device="cuda:0", hub="ms"):
+    if "Fun-ASR-Nano" in model_name:
+        from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM
+        engine = FunASRNanoVLLM.from_pretrained(
+            model=model_name, hub=hub, 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]
+
+    elif "GLM-ASR" in model_name:
+        from funasr.models.glm_asr.inference_vllm import GLMASRVLLMEngine
+        engine = GLMASRVLLMEngine.from_pretrained(
+            model=model_name, hub=hub, device=device, dtype="bf16",
+            gpu_memory_utilization=0.4, max_model_len=4096)
+        engine.generate(inputs=[seg_files[0]])  # warmup
+        t0 = time.perf_counter()
+        results = engine.generate(inputs=seg_files, max_new_tokens=500)
+        t1 = time.perf_counter()
+        texts = [r["text"] for r in results]
+
+    else:
+        from funasr.auto.auto_model_vllm import AutoModelVLLM
+        engine = AutoModelVLLM(model=model_name, hub=hub, device=device)
+        engine.generate(inputs=[seg_files[0]])
+        t0 = time.perf_counter()
+        results = engine.generate(inputs=seg_files, 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="FunASR vLLM Benchmark")
+    parser.add_argument("--model", type=str, required=True, help="Model name or path")
+    parser.add_argument("--hub", type=str, default="ms", choices=["ms", "hf"])
+    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)
+    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"{'='*60}")
+    print(f"FunASR vLLM Benchmark")
+    print(f"{'='*60}")
+    print(f"Model: {args.model}")
+    print(f"Dataset: {len(files)} files, {total_audio:.0f}s audio")
+
+    # VAD
+    print(f"\n>>> VAD pre-segmenting...")
+    all_segments = vad_segment(files, device=args.device)
+    print(f"    {len(all_segments)} segments")
+
+    os.makedirs("/tmp/benchmark_vllm_segs", exist_ok=True)
+    seg_files = []
+    for i, (fi, audio) in enumerate(all_segments):
+        path = f"/tmp/benchmark_vllm_segs/{i:04d}.wav"
+        sf.write(path, audio, 16000)
+        seg_files.append(path)
+
+    # PyTorch
+    cer_pt = None
+    pt_time = None
+    if not args.skip_pytorch:
+        print(f"\n>>> PyTorch native...")
+        pt_time, pt_seg_texts = run_pytorch(args.model, 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(f"\n>>> vLLM...")
+    vllm_time, vllm_seg_texts = run_vllm(args.model, seg_files, device=args.device, hub=args.hub)
+    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"RESULTS")
+    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':<20} {vllm_time:<10.1f} {total_audio/vllm_time:<10.1f} {cer_vllm:.2f}%")
+    if not args.skip_pytorch:
+        print(f"{'-'*60}")
+        speedup = (total_audio/vllm_time) / (total_audio/pt_time)
+        print(f"Speedup: {speedup:.1f}x | CER diff: {cer_vllm - cer_pt:+.2f}%")
+    print(f"{'='*60}")