Initial commit: FunASR Speech Recognition Toolkit

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

funasr/models/fun_asr_nano/inference_vllm_pipeline.py

@@ -0,0 +1,372 @@
+#!/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 Pipeline: VAD + ASR(vLLM) + Speaker Diarization.
+
+Replicates AutoModel's inference_with_vad pipeline but uses vLLM for
+the LLM decoding step, enabling batch processing of all VAD segments
+in a single generate() call.
+
+Usage:
+    from funasr.models.fun_asr_nano.inference_vllm_pipeline import FunASRNanoVLLMPipeline
+
+    model = FunASRNanoVLLMPipeline(
+        model="FunAudioLLM/Fun-ASR-Nano-2512",
+        vad_model="fsmn-vad",
+        spk_model="cam++",
+        tensor_parallel_size=2,
+    )
+    results = model.generate("long_meeting.wav", language="中文")
+"""
+
+import logging
+import os
+import re
+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 _clean_text(text: str) -> str:
+    """Remove tags, fillers, and garbage from output."""
+    text = re.sub(r"<[^>]*>|</[^>]*>", "", text)
+    text = re.sub(r"(>.{2,8}?)\1{3,}", "", text)
+    text = re.sub(r"\[breath\]|\[noise\]|/sil|endofbreak|FFFF", "", text)
+    text = re.sub(r"\s+", " ", text)
+    text = text.replace("�", "").lstrip(">")
+    return text.strip()
+
+
+class FunASRNanoVLLMPipeline:
+    """VAD + ASR(vLLM) + Speaker pipeline.
+
+    Pipeline:
+        1. VAD: segment long audio into speech regions (torch)
+        2. ASR: batch ALL segments through vLLM in single generate() call
+        3. Speaker: extract embeddings per segment, cluster (torch)
+        4. Combine: merge text + timestamps + speaker labels
+
+    Args:
+        model: Fun-ASR-Nano model name or path.
+        vad_model: VAD model name (e.g. "fsmn-vad"). None to disable.
+        vad_kwargs: VAD config (e.g. {"max_single_segment_time": 30000}).
+        spk_model: Speaker model name (e.g. "cam++"). None to disable.
+        hub: "ms" or "hf".
+        device: Device for audio encoder + VAD + speaker.
+        dtype: Compute dtype for ASR.
+        tensor_parallel_size: GPUs for vLLM.
+        gpu_memory_utilization: GPU memory fraction for vLLM.
+        max_model_len: Maximum sequence length for vLLM.
+    """
+
+    def __init__(
+        self,
+        model: str = "FunAudioLLM/Fun-ASR-Nano-2512",
+        vad_model: str = None,
+        vad_kwargs: dict = None,
+        spk_model: str = None,
+        spk_kwargs: dict = None,
+        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 = 4096,
+        enforce_eager: bool = False,
+        **kwargs,
+    ):
+        from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM
+
+        # ASR engine (vLLM)
+        self.asr_engine = FunASRNanoVLLM.from_pretrained(
+            model=model, hub=hub, device=device, dtype=dtype,
+            tensor_parallel_size=tensor_parallel_size,
+            gpu_memory_utilization=gpu_memory_utilization,
+            max_model_len=max_model_len, enforce_eager=enforce_eager,
+            **kwargs,
+        )
+
+        # VAD model (torch)
+        self.vad_model = None
+        if vad_model is not None:
+            from funasr import AutoModel
+            vad_kw = vad_kwargs or {}
+            self.vad_model = AutoModel(
+                model=vad_model, device=device, disable_update=True, **vad_kw
+            )
+
+        # Speaker model (torch)
+        self.spk_model = None
+        self.cb_model = None
+        if spk_model is not None:
+            from funasr import AutoModel
+            from funasr.models.campplus.cluster_backend import ClusterBackend
+            spk_kw = spk_kwargs or {}
+            self.spk_model = AutoModel(
+                model=spk_model, device=device, disable_update=True, **spk_kw
+            )
+            cb_kwargs = spk_kw.get("cb_kwargs", {})
+            self.cb_model = ClusterBackend(**cb_kwargs).to(device)
+
+        self.device = device
+        self.sample_rate = 16000
+
+    def generate(
+        self,
+        input: Union[str, List[str]],
+        hotwords: List[str] = None,
+        language: str = None,
+        itn: bool = True,
+        max_new_tokens: int = 512,
+        batch_size_s: int = 300,
+        return_spk_res: bool = True,
+        **kwargs,
+    ) -> List[dict]:
+        """Run the full pipeline: VAD → ASR(vLLM) → Speaker.
+
+        Args:
+            input: Audio file path(s).
+            hotwords: Hotwords for ASR.
+            language: Language hint.
+            itn: Inverse text normalization.
+            max_new_tokens: Max tokens per segment.
+            batch_size_s: Max batch duration in seconds (for memory control).
+            return_spk_res: Whether to return speaker info.
+
+        Returns:
+            List of dicts: [{"key", "text", "timestamp", "sentence_info"}]
+        """
+        if isinstance(input, str):
+            input = [input]
+
+        results_all = []
+        for audio_path in input:
+            result = self._process_one(
+                audio_path, hotwords=hotwords, language=language,
+                itn=itn, max_new_tokens=max_new_tokens,
+                batch_size_s=batch_size_s, return_spk_res=return_spk_res,
+                **kwargs,
+            )
+            results_all.append(result)
+        return results_all
+
+    def _process_one(self, audio_path, **kwargs):
+        """Process a single audio file through the full pipeline."""
+        from funasr.utils.load_utils import load_audio_text_image_video
+        from funasr.utils.vad_utils import slice_padding_audio_samples
+
+        key = os.path.splitext(os.path.basename(audio_path))[0]
+
+        # Load audio
+        audio_data = load_audio_text_image_video(audio_path, fs=self.sample_rate)
+        if isinstance(audio_data, torch.Tensor):
+            audio_np = audio_data.numpy()
+        else:
+            audio_np = np.array(audio_data)
+        speech_length = len(audio_np)
+
+        # Step 1: VAD
+        if self.vad_model is not None:
+            vad_res = self.vad_model.generate(input=audio_path, cache={}, is_final=True)
+            vad_segments = vad_res[0]["value"]  # [[start_ms, end_ms], ...]
+        else:
+            vad_segments = [[0, int(speech_length / self.sample_rate * 1000)]]
+
+        if not vad_segments:
+            return {"key": key, "text": "", "timestamp": []}
+
+        n_segments = len(vad_segments)
+        logger.info(f"VAD: {n_segments} segments for {key}")
+
+        # Step 2: Slice audio by VAD segments and encode
+        segment_audios = []
+        for seg in vad_segments:
+            start_sample = int(seg[0] * self.sample_rate / 1000)
+            end_sample = int(seg[1] * self.sample_rate / 1000)
+            end_sample = min(end_sample, speech_length)
+            segment_audios.append(audio_np[start_sample:end_sample])
+
+        # Step 3: Batch ASR via vLLM
+        # Encode all segments and build prompts
+        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
+
+        prompts = []
+        for seg_audio in segment_audios:
+            seg_tensor = torch.from_numpy(seg_audio).float()
+            adaptor_out, adaptor_out_lens, _, _ = self.asr_engine._encode_audio(seg_tensor)
+            input_embeds = self.asr_engine._build_input_embeds(
+                adaptor_out, adaptor_out_lens,
+                hotwords=kwargs.get("hotwords"),
+                language=kwargs.get("language"),
+                itn=kwargs.get("itn", True),
+            )
+            prompts.append(EmbedsPrompt(prompt_embeds=input_embeds.float()))
+
+        params = SamplingParams(
+            max_tokens=kwargs.get("max_new_tokens", 512),
+            temperature=0.0,
+            # Prompt-embeds mode has no token IDs to penalize; see #2948.
+            repetition_penalty=resolve_repetition_penalty(
+                kwargs.get("repetition_penalty", 1.0)
+            ),
+            skip_special_tokens=True,
+        )
+
+        # Single batch generate for ALL segments
+        t0 = time.perf_counter()
+        outputs = self.asr_engine.vllm_engine.generate(prompts, params, use_tqdm=False)
+        t1 = time.perf_counter()
+        logger.info(f"vLLM batch ASR: {n_segments} segments in {t1-t0:.3f}s")
+
+        # Decode results
+        asr_results = []
+        for output in outputs:
+            text = output.outputs[0].text
+            if not text and output.outputs[0].token_ids:
+                text = self.asr_engine.tokenizer.decode(
+                    list(output.outputs[0].token_ids), skip_special_tokens=True
+                )
+            text = _clean_text(text)
+            asr_results.append(text)
+
+        # Step 4: Speaker embeddings (if spk_model configured)
+        spk_embeddings = None
+        if self.spk_model is not None and kwargs.get("return_spk_res", True):
+            from funasr.models.campplus.utils import sv_chunk, postprocess, distribute_spk
+
+            all_segments = []
+            all_spk_embs = []
+            for i, seg_audio in enumerate(segment_audios):
+                vad_seg = [
+                    [vad_segments[i][0] / 1000.0, vad_segments[i][1] / 1000.0, seg_audio]
+                ]
+                chunks = sv_chunk(vad_seg)
+                all_segments.extend(chunks)
+                speech_chunks = [c[2] for c in chunks]
+                spk_res = self.spk_model.generate(input=speech_chunks, cache={}, is_final=True)
+                embs = torch.cat([r["spk_embedding"] for r in spk_res], dim=0)
+                all_spk_embs.append(embs)
+
+            if all_spk_embs:
+                spk_embeddings = torch.cat(all_spk_embs, dim=0)
+
+        # Step 5: Combine results
+        # Merge text with timestamps
+        full_text = ""
+        all_timestamps = []
+        for i, (seg, text) in enumerate(zip(vad_segments, asr_results)):
+            if text:
+                if full_text:
+                    full_text += " "
+                full_text += text
+                # Simple word-level timestamp from VAD boundaries
+                all_timestamps.append([int(seg[0]), int(seg[1])])
+
+        result = {"key": key, "text": full_text}
+
+        # Add timestamps if available from CTC
+        if self.asr_engine.ctc_decoder is not None:
+            try:
+                detailed_timestamps = self._compute_all_timestamps(
+                    segment_audios, vad_segments, asr_results
+                )
+                if detailed_timestamps:
+                    result["timestamp"] = detailed_timestamps
+            except Exception as e:
+                logger.debug(f"Timestamp computation failed: {e}")
+                result["timestamp"] = all_timestamps
+        else:
+            result["timestamp"] = all_timestamps
+
+        # Add speaker info
+        if spk_embeddings is not None and self.cb_model is not None:
+            from funasr.models.campplus.utils import postprocess, distribute_spk
+
+            all_segments_sorted = sorted(all_segments, key=lambda x: x[0])
+            labels = self.cb_model(
+                spk_embeddings.cpu(),
+                oracle_num=kwargs.get("preset_spk_num", None),
+            )
+            sv_output = postprocess(all_segments_sorted, None, labels, spk_embeddings.cpu())
+
+            # Build sentence_info
+            sentence_list = []
+            for i, (seg, text) in enumerate(zip(vad_segments, asr_results)):
+                if text:
+                    sentence_list.append({
+                        "start": seg[0],
+                        "end": seg[1],
+                        "text": text,
+                        "timestamp": [[int(seg[0]), int(seg[1])]],
+                    })
+            distribute_spk(sentence_list, sv_output)
+            result["sentence_info"] = sentence_list
+
+        return result
+
+    def _compute_all_timestamps(self, segment_audios, vad_segments, asr_results):
+        """Compute CTC timestamps for all segments with VAD offsets."""
+        from funasr.models.fun_asr_nano.tools.utils import forced_align
+
+        all_timestamps = []
+        for seg_audio, vad_seg, text in zip(segment_audios, vad_segments, asr_results):
+            if not text:
+                continue
+            try:
+                seg_tensor = torch.from_numpy(seg_audio).float()
+                from funasr.utils.load_utils import extract_fbank
+                speech, speech_lengths = extract_fbank(
+                    seg_tensor, data_type="sound",
+                    frontend=self.asr_engine.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.asr_engine.audio_encoder(speech, speech_lengths)
+                    dec_out, dec_lens = self.asr_engine.ctc_decoder(enc_out, enc_lens)
+                    ctc_logits = self.asr_engine.ctc.log_softmax(dec_out)
+
+                x = ctc_logits[0, :enc_lens[0].item(), :]
+                target_ids = torch.tensor(
+                    self.asr_engine.ctc_tokenizer.encode(text), dtype=torch.int64
+                )
+                if len(target_ids) == 0:
+                    continue
+
+                timestamps = forced_align(x, target_ids, self.asr_engine.blank_id)
+                vad_offset_ms = int(vad_seg[0])
+                for ts in timestamps:
+                    ts["token"] = self.asr_engine.ctc_tokenizer.decode([ts["token"]])
+                    ts["start_time"] = ts["start_time"] * 6 * 10 / 1000 + vad_offset_ms / 1000
+                    ts["end_time"] = ts["end_time"] * 6 * 10 / 1000 + vad_offset_ms / 1000
+                all_timestamps.extend(timestamps)
+            except Exception as e:
+                logger.debug(f"Timestamp failed for segment: {e}")
+                all_timestamps.append({
+                    "start_time": vad_seg[0] / 1000,
+                    "end_time": vad_seg[1] / 1000,
+                    "token": text,
+                })
+        return all_timestamps
+
+    @classmethod
+    def from_pretrained(cls, model="FunAudioLLM/Fun-ASR-Nano-2512", **kwargs):
+        """Convenience constructor."""
+        return cls(model=model, **kwargs)