merge vllm: fix use_low_frame_rate (CER 8.91% matches PyTorch), add benchmark

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}")

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<think>\n\n</think>\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|</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]