Initial commit: FunASR Speech Recognition Toolkit

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

funasr/auto/auto_model_vllm.py

@@ -0,0 +1,348 @@
+#!/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)
+
+"""
+Generic vLLM inference wrapper for ALL LLM-based ASR models in FunASR.
+
+Applicable models (any model with audio_encoder + adaptor + LLM architecture):
+    - FunASRNano (Fun-ASR-Nano-2512, Fun-ASR-MLT-Nano-2512)
+    - LLMASR (Whisper + Qwen/Vicuna/LLaMA)
+    - GLMASR (GLM-ASR-Nano)
+
+NOT applicable (these models don't use autoregressive LLM decoding):
+    - Paraformer (non-autoregressive CIF predictor + attention decoder)
+    - SenseVoice (Whisper-like encoder-decoder, not LLM-based)
+    - Conformer/Transformer ASR (CTC/attention, no LLM)
+    - CT-Transformer (punctuation model, small transformer)
+    - Qwen3-ASR (uses external qwen-asr package with its own optimized inference)
+
+Usage:
+    from funasr.auto.auto_model_vllm import AutoModelVLLM
+
+    # Works for any LLM-based ASR model
+    model = AutoModelVLLM(
+        model="FunAudioLLM/Fun-ASR-Nano-2512",
+        tensor_parallel_size=2,
+    )
+    results = model.generate(["audio.wav"])
+
+    # Also works for LLMASR models
+    model = AutoModelVLLM(
+        model="/path/to/llm_asr_model",
+        tensor_parallel_size=4,
+    )
+"""
+
+import glob
+import json
+import logging
+import os
+import re
+import shutil
+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}
+
+# Models that use LLM and can benefit from vLLM
+_LLM_BASED_MODELS = {"FunASRNano", "LLMASR", "LLMASRNAR", "GLMASR", "QwenAudioWarp"}
+
+# Models that CANNOT use vLLM (no autoregressive LLM)
+_NON_LLM_MODELS = {
+    "Paraformer": "Non-autoregressive model using CIF predictor. No LLM decoding.",
+    "SenseVoice": "Whisper-like encoder-decoder. Not LLM-based.",
+    "CTTransformer": "Small punctuation model. No benefit from vLLM.",
+    "Conformer": "CTC/attention encoder-decoder. No LLM.",
+    "Qwen3ASR": "Uses external qwen-asr package with optimized inference.",
+}
+
+
+def check_vllm_applicable(model_name: str) -> bool:
+    """Check if a model can use vLLM inference.
+
+    Args:
+        model_name: The model class name from config.yaml.
+
+    Returns:
+        True if vLLM is applicable.
+
+    Raises:
+        ValueError: If model explicitly cannot use vLLM, with explanation.
+    """
+    if model_name in _LLM_BASED_MODELS:
+        return True
+    for non_llm, reason in _NON_LLM_MODELS.items():
+        if non_llm in model_name:
+            raise ValueError(
+                f"Model '{model_name}' cannot use vLLM: {reason}\n"
+                f"vLLM only accelerates autoregressive LLM decoding. "
+                f"Use the standard FunASR AutoModel for this model."
+            )
+    return False
+
+
+def prepare_vllm_weights(model_dir: str, output_dir: str = None) -> str:
+    """Extract LLM weights from model.pt into vLLM-compatible format.
+
+    Works for any model that stores LLM weights with 'llm.' prefix in model.pt
+    and has a config directory (e.g., Qwen3-0.6B/) with model config and tokenizer.
+
+    Args:
+        model_dir: Path to the FunASR model directory.
+        output_dir: Where to save extracted weights. Auto-detected if None.
+
+    Returns:
+        Path to vLLM-ready model directory.
+    """
+    if output_dir is None:
+        # Find the LLM config directory
+        from omegaconf import OmegaConf
+        config_path = os.path.join(model_dir, "config.yaml")
+        if os.path.exists(config_path):
+            config = OmegaConf.load(config_path)
+            llm_conf = OmegaConf.to_container(config.get("llm_conf", {}), resolve=True)
+            llm_path = llm_conf.get("init_param_path", "")
+            if llm_path and not os.path.isabs(llm_path):
+                llm_path = os.path.join(model_dir, llm_path)
+            if os.path.isdir(llm_path):
+                output_dir = llm_path + "-vllm"
+            else:
+                output_dir = os.path.join(model_dir, "llm-vllm")
+        else:
+            output_dir = os.path.join(model_dir, "llm-vllm")
+
+    # Check if already prepared
+    if glob.glob(os.path.join(output_dir, "*.safetensors")) or glob.glob(
+        os.path.join(output_dir, "model*.bin")
+    ):
+        logger.info(f"vLLM weights already at {output_dir}")
+        return output_dir
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Find and copy LLM config/tokenizer files
+    from omegaconf import OmegaConf
+    config = OmegaConf.load(os.path.join(model_dir, "config.yaml"))
+    llm_conf = OmegaConf.to_container(config.get("llm_conf", {}), resolve=True)
+    llm_config_dir = llm_conf.get("init_param_path", "")
+    if llm_config_dir and not os.path.isabs(llm_config_dir):
+        llm_config_dir = os.path.join(model_dir, llm_config_dir)
+
+    if os.path.isdir(llm_config_dir):
+        for fname in os.listdir(llm_config_dir):
+            src = os.path.join(llm_config_dir, fname)
+            dst = os.path.join(output_dir, fname)
+            if os.path.isfile(src) and not os.path.exists(dst):
+                shutil.copy2(src, dst)
+
+    # Extract LLM weights from model.pt
+    model_pt = os.path.join(model_dir, "model.pt")
+    if not os.path.exists(model_pt):
+        raise FileNotFoundError(f"model.pt not found at {model_pt}")
+
+    logger.info(f"Extracting LLM weights from {model_pt}...")
+    checkpoint = torch.load(model_pt, map_location="cpu")
+    state_dict = checkpoint.get("state_dict", checkpoint)
+
+    llm_state = {}
+    for key, value in state_dict.items():
+        if key.startswith("llm."):
+            llm_state[key[4:]] = value  # Remove 'llm.' prefix
+
+    if not llm_state:
+        raise RuntimeError("No LLM weights found (expected 'llm.*' prefix)")
+
+    logger.info(f"Extracted {len(llm_state)} LLM weight tensors")
+
+    try:
+        from safetensors.torch import save_file
+        save_path = os.path.join(output_dir, "model.safetensors")
+        save_file(llm_state, save_path)
+        index = {
+            "metadata": {"total_size": sum(v.numel() * v.element_size() for v in llm_state.values())},
+            "weight_map": {k: "model.safetensors" for k in llm_state.keys()},
+        }
+        with open(os.path.join(output_dir, "model.safetensors.index.json"), "w") as f:
+            json.dump(index, f, indent=2)
+    except ImportError:
+        torch.save(llm_state, os.path.join(output_dir, "model.bin"))
+
+    return output_dir
+
+
+class AutoModelVLLM:
+    """Generic vLLM wrapper for LLM-based ASR models.
+
+    Automatically detects model architecture, extracts LLM weights,
+    loads audio components in PyTorch, and uses vLLM for generation.
+
+    Works for: FunASRNano, LLMASR, GLMASR, and any model with
+    audio_encoder + audio_adaptor + LLM architecture.
+
+    Args:
+        model: Model name (hub) or local directory path.
+        hub: "ms" (ModelScope) or "hf" (HuggingFace).
+        device: Device for audio encoder/adaptor.
+        dtype: Compute dtype ("bf16", "fp16", "fp32").
+        tensor_parallel_size: GPUs for vLLM tensor parallelism.
+        gpu_memory_utilization: GPU memory fraction for vLLM.
+        max_model_len: Maximum sequence length.
+
+    Example:
+        >>> model = AutoModelVLLM(model="FunAudioLLM/Fun-ASR-Nano-2512")
+        >>> results = model.generate(["audio.wav"], language="中文")
+    """
+
+    def __init__(
+        self,
+        model: str,
+        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,
+    ):
+        # Resolve model directory
+        if os.path.isdir(model):
+            self.model_dir = model
+        else:
+            if hub in ("ms", "modelscope"):
+                from modelscope.hub.snapshot_download import snapshot_download
+                self.model_dir = snapshot_download(model, revision=kwargs.get("revision", "master"))
+            elif hub in ("hf", "huggingface"):
+                from huggingface_hub import snapshot_download
+                self.model_dir = snapshot_download(model)
+            else:
+                raise ValueError(f"Unsupported hub: {hub}")
+
+        # Check model type
+        from omegaconf import OmegaConf
+        config = OmegaConf.load(os.path.join(self.model_dir, "config.yaml"))
+        self.model_type = config.get("model", "unknown")
+        check_vllm_applicable(self.model_type)
+
+        self.device = device
+        self.dtype = dtype
+        self.torch_dtype = dtype_map.get(dtype, torch.bfloat16)
+
+        # Use the specialized implementation if available
+        if self.model_type == "FunASRNano":
+            from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM
+            self._engine = FunASRNanoVLLM(
+                model_dir=self.model_dir, 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,
+            )
+        elif self.model_type in ("GLMASR", "glmasr"):
+            from funasr.models.glm_asr.inference_vllm import GLMASRVLLMEngine
+            self._engine = GLMASRVLLMEngine(
+                model_dir=self.model_dir, device=device, dtype=dtype,
+                tensor_parallel_size=tensor_parallel_size,
+                gpu_memory_utilization=gpu_memory_utilization,
+                max_model_len=max_model_len,
+                **kwargs,
+            )
+        elif self.model_type in ("LLMASR", "LLMASRNAR"):
+            self._engine = self._build_llmasr_engine(
+                config, tensor_parallel_size, gpu_memory_utilization,
+                max_model_len, enforce_eager, **kwargs,
+            )
+        else:
+            # Generic fallback using the FunASRNano approach
+            # (works for any model with audio_encoder + adaptor + LLM)
+            from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM
+            self._engine = FunASRNanoVLLM(
+                model_dir=self.model_dir, 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,
+            )
+
+    def _build_llmasr_engine(self, config, tensor_parallel_size, gpu_memory_utilization,
+                              max_model_len, enforce_eager, **kwargs):
+        """Build vLLM engine for LLMASR models (Whisper + Qwen/Vicuna)."""
+        from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM
+
+        # LLMASR follows same pattern as FunASRNano
+        return FunASRNanoVLLM(
+            model_dir=self.model_dir, device=self.device, dtype=self.dtype,
+            tensor_parallel_size=tensor_parallel_size,
+            gpu_memory_utilization=gpu_memory_utilization,
+            max_model_len=max_model_len, enforce_eager=enforce_eager,
+            **kwargs,
+        )
+
+    def generate(self, inputs, **kwargs):
+        """Run ASR inference.
+
+        Args:
+            inputs: Audio file path(s), numpy arrays, or tensors.
+            **kwargs: Model-specific parameters (language, hotwords, etc.)
+
+        Returns:
+            List of result dicts with "key" and "text" fields.
+        """
+        self._warn_if_audio_too_long(inputs)
+        return self._engine.generate(inputs, **kwargs)
+
+    def _warn_if_audio_too_long(self, inputs, max_safe_sec=40.0):
+        """Warn (once) if a single audio input is long enough to be truncated.
+
+        Fun-ASR-Nano is a segment-level (LLM-)ASR model. Decoding very long
+        audio in a single pass can silently truncate or degrade the output -- the
+        decode hits ``max_new_tokens`` long before the audio ends, so the user
+        gets a partial transcript with no error. The right usage is to
+        pre-segment with VAD; this warning points users there instead of letting
+        them get a silently truncated result. It does not change the output.
+        """
+        if getattr(self, "_warned_audio_too_long", False):
+            return
+        items = inputs if isinstance(inputs, (list, tuple)) else [inputs]
+        for item in items:
+            duration = None
+            try:
+                if isinstance(item, str):
+                    import soundfile as sf
+
+                    duration = sf.info(item).duration
+                elif isinstance(item, np.ndarray):
+                    duration = item.shape[-1] / 16000.0
+                elif isinstance(item, torch.Tensor):
+                    duration = item.shape[-1] / 16000.0
+            except Exception:
+                continue
+            if duration is not None and duration > max_safe_sec:
+                logger.warning(
+                    "AutoModelVLLM received a %.0fs audio input. Fun-ASR-Nano is a "
+                    "segment-level model; decoding very long audio in a single pass can "
+                    "truncate or degrade the result. Pre-segment with VAD and pass the "
+                    "segments, or use the high-level `funasr.AutoModel(model=..., "
+                    'vad_model="fsmn-vad")`, which segments long audio automatically.',
+                    duration,
+                )
+                self._warned_audio_too_long = True
+                break
+
+    @classmethod
+    def supported_models(cls):
+        """Return dict of model types and their vLLM support status."""
+        info = {}
+        for m in _LLM_BASED_MODELS:
+            info[m] = {"supported": True, "reason": "LLM-based, autoregressive generation"}
+        for m, reason in _NON_LLM_MODELS.items():
+            info[m] = {"supported": False, "reason": reason}
+        return info