Initial commit: FunASR Speech Recognition Toolkit

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

funasr/models/glm_asr/inference_vllm.py

@@ -0,0 +1,351 @@
+#!/usr/bin/env python3
+"""GLM-ASR vLLM inference engine.
+
+Architecture: audio_tower (Whisper-like) + multi_modal_projector + language_model (Llama)
+Strategy: audio_tower + projector in PyTorch, language_model in vLLM via EmbedsPrompt.
+
+Usage:
+    from funasr.models.glm_asr.inference_vllm import GLMASRVLLMEngine
+
+    engine = GLMASRVLLMEngine.from_pretrained("zai-org/GLM-ASR-Nano-2512")
+    results = engine.generate(inputs=["audio.wav"])
+    print(results[0]["text"])
+"""
+
+import glob
+import json
+import logging
+import os
+import re
+import shutil
+import time
+
+import numpy as np
+import torch
+
+logger = logging.getLogger(__name__)
+dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}
+
+# Warn only once per process so batch loops do not spam the log.
+_warned_rep_penalty = False
+
+
+def _safe_repetition_penalty(repetition_penalty):
+    """Force ``repetition_penalty`` to the neutral value for prompt-embeds mode.
+
+    GLM-ASR feeds vLLM precomputed embeddings (``enable_prompt_embeds=True``), so
+    a request carries no prompt token IDs. vLLM applies ``repetition_penalty`` by
+    scattering over those IDs, so any value other than 1.0 indexes an empty
+    token-id tensor and aborts the engine with a CUDA
+    ``scatter gather index out of bounds`` assertion (issue #2948). We therefore
+    warn once and fall back to the neutral value of 1.0.
+    """
+    global _warned_rep_penalty
+
+    if repetition_penalty is None or repetition_penalty == 1.0:
+        return 1.0
+
+    if not _warned_rep_penalty:
+        logger.warning(
+            "repetition_penalty=%s is not supported in vLLM prompt-embeds mode "
+            "(no prompt token IDs to penalize) and would trigger a CUDA scatter "
+            "index-out-of-bounds crash; using repetition_penalty=1.0 instead. "
+            "See https://github.com/modelscope/FunASR/issues/2948.",
+            repetition_penalty,
+        )
+        _warned_rep_penalty = True
+    return 1.0
+
+
+def prepare_glmasr_vllm_dir(model_dir: str) -> str:
+    """Extract language_model weights into vLLM-compatible Llama format."""
+    output_dir = os.path.join(model_dir, "language_model_vllm")
+
+    if glob.glob(os.path.join(output_dir, "*.safetensors")):
+        logger.info(f"vLLM LM weights already at {output_dir}")
+        return output_dir
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    from safetensors import safe_open
+    from safetensors.torch import save_file
+
+    st_files = sorted(glob.glob(os.path.join(model_dir, "*.safetensors")))
+    lm_state = {}
+    for st_file in st_files:
+        with safe_open(st_file, framework="pt") as f:
+            for key in f.keys():
+                if key.startswith("language_model."):
+                    lm_state[key[len("language_model."):]] = f.get_tensor(key)
+
+    if not lm_state:
+        raise RuntimeError("No language_model weights found in safetensors")
+
+    logger.info(f"Extracted {len(lm_state)} LM tensors")
+    save_file(lm_state, os.path.join(output_dir, "model.safetensors"))
+
+    with open(os.path.join(model_dir, "config.json")) as f:
+        full_config = json.load(f)
+    text_config = full_config["text_config"]
+    text_config["architectures"] = ["LlamaForCausalLM"]
+    text_config["model_type"] = "llama"
+    with open(os.path.join(output_dir, "config.json"), "w") as f:
+        json.dump(text_config, f, indent=2)
+
+    for fname in os.listdir(model_dir):
+        if "tokenizer" in fname or fname == "generation_config.json":
+            src = os.path.join(model_dir, fname)
+            dst = os.path.join(output_dir, fname)
+            if os.path.isfile(src) and not os.path.exists(dst):
+                shutil.copy2(src, dst)
+
+    index = {
+        "metadata": {"total_size": sum(v.numel() * v.element_size() for v in lm_state.values())},
+        "weight_map": {k: "model.safetensors" for k in lm_state.keys()},
+    }
+    with open(os.path.join(output_dir, "model.safetensors.index.json"), "w") as f:
+        json.dump(index, f, indent=2)
+
+    logger.info(f"Saved vLLM LM to {output_dir}")
+    return output_dir
+
+
+# Warn only once per process so batch loops do not spam the log.
+_warned_dup_keys = False
+
+
+def _dedup_keys(keys):
+    """Make result keys unique while preserving order and first-occurrence names.
+
+    Each result key is derived from the audio file basename
+    (``os.path.splitext(os.path.basename(path))[0]``), so two inputs that live
+    in different directories but share a basename -- e.g. ``spk1/segment.wav``
+    and ``spk2/segment.wav`` -- both map to ``"segment"``. A downstream
+    ``{r["key"]: r["text"]}`` mapping (the canonical FunASR result shape) would
+    then silently drop all but the last colliding entry, returning fewer
+    transcripts than inputs with no error. Appending a deterministic ``_N``
+    suffix to later collisions keeps every transcript addressable.
+
+    Args:
+        keys: Result keys in input order.
+
+    Returns:
+        A new list of unique keys, same length and order as ``keys``. The first
+        occurrence of each key is preserved unchanged; the n-th repeat becomes
+        ``"<key>_<n-1>"`` (e.g. ``"seg"`` -> ``"seg"``, ``"seg_1"``, ``"seg_2"``).
+    """
+    global _warned_dup_keys
+
+    seen = set()
+    out = []
+    collided = False
+    for key in keys:
+        if key not in seen:
+            seen.add(key)
+            out.append(key)
+            continue
+        # Find the first free "<key>_<n>" so a suffixed key cannot itself clash
+        # with an existing one (e.g. inputs "seg", "seg_1", "seg").
+        collided = True
+        n = 1
+        candidate = f"{key}_{n}"
+        while candidate in seen:
+            n += 1
+            candidate = f"{key}_{n}"
+        seen.add(candidate)
+        out.append(candidate)
+
+    if collided and not _warned_dup_keys:
+        logger.warning(
+            "Duplicate result keys from audio basenames were made unique with "
+            "'_N' suffixes (e.g. two files named 'segment.wav' in different "
+            "directories map to the same key); pass distinct filenames if you "
+            "rely on the basename as the result key."
+        )
+        _warned_dup_keys = True
+
+    return out
+
+
+class GLMASRVLLMEngine:
+    """GLM-ASR with vLLM backend.
+
+    Audio tower + projector run in PyTorch on a single device.
+    Language model is decoded by vLLM with PagedAttention for high throughput.
+
+    Args:
+        model_dir: Path to GLM-ASR model directory.
+        device: Device for audio encoder.
+        dtype: Compute dtype ("bf16", "fp16", "fp32").
+        tensor_parallel_size: GPUs for vLLM tensor parallelism.
+        gpu_memory_utilization: GPU memory fraction for vLLM KV cache.
+        max_model_len: Maximum sequence length for vLLM.
+    """
+
+    def __init__(self, model_dir, device="cuda:0", dtype="bf16",
+                 tensor_parallel_size=1, gpu_memory_utilization=0.5,
+                 max_model_len=4096, **kwargs):
+        from vllm import LLM
+        from transformers import AutoProcessor, AutoConfig, AutoModel as HFAutoModel
+
+        from funasr.models.glm_asr.vllm_utils import warn_if_degraded_dtype
+
+        self.device = device
+        self.torch_dtype = dtype_map.get(warn_if_degraded_dtype(dtype), torch.bfloat16)
+        self.model_dir = model_dir
+
+        logger.info(f"Loading GLM-ASR audio components from {model_dir}")
+        full_model = HFAutoModel.from_pretrained(
+            model_dir, dtype=self.torch_dtype, device_map=device, trust_remote_code=True
+        )
+        full_model.eval()
+
+        self.audio_tower = full_model.audio_tower
+        self.multi_modal_projector = full_model.multi_modal_projector
+        self.get_audio_features = full_model.get_audio_features
+        self.embed_tokens = full_model.language_model.get_input_embeddings()
+        self._full_model_config = full_model.config
+
+        # Free LM weights from GPU (vLLM loads its own copy)
+        del full_model.language_model
+        torch.cuda.empty_cache()
+
+        self.processor = AutoProcessor.from_pretrained(model_dir, trust_remote_code=True)
+
+        # Prepare and load vLLM engine
+        vllm_dir = prepare_glmasr_vllm_dir(model_dir)
+        logger.info(f"Initializing vLLM LM from {vllm_dir}")
+        self.vllm_engine = LLM(
+            model=vllm_dir,
+            enable_prompt_embeds=True,
+            tensor_parallel_size=tensor_parallel_size,
+            gpu_memory_utilization=gpu_memory_utilization,
+            max_model_len=max_model_len,
+            dtype={"bf16": "bfloat16", "fp16": "float16", "fp32": "auto"}.get(dtype, dtype),
+            trust_remote_code=True,
+        )
+        self.tokenizer = self.vllm_engine.get_tokenizer()
+        logger.info("GLM-ASR vLLM engine ready")
+
+    @torch.no_grad()
+    def _encode_audio(self, audio_input):
+        """Encode a single audio input through audio_tower + projector.
+
+        Returns:
+            audio_embeds: (1, T, hidden_size) tensor
+        """
+        import librosa
+
+        if isinstance(audio_input, str):
+            audio, _ = librosa.load(audio_input, sr=16000)
+        elif isinstance(audio_input, np.ndarray):
+            audio = audio_input.astype(np.float32)
+        elif isinstance(audio_input, torch.Tensor):
+            audio = audio_input.cpu().numpy().astype(np.float32)
+        else:
+            raise ValueError(f"Unsupported audio type: {type(audio_input)}")
+
+        inputs = self.processor.feature_extractor(audio, sampling_rate=16000, return_tensors="pt")
+        input_features = inputs["input_features"].to(self.device, dtype=self.torch_dtype)
+        feat_len = input_features.shape[-1]
+        input_features_mask = torch.ones(1, feat_len, dtype=torch.long, device=self.device)
+
+        audio_outputs = self.get_audio_features(
+            input_features, input_features_mask, return_dict=True
+        )
+        audio_embeds = audio_outputs.pooler_output
+        return audio_embeds.unsqueeze(0)
+
+    def _build_prompt_embeds(self, audio_embeds, prompt="转录以下音频内容"):
+        """Build [prefix_text_emb | audio_emb | suffix_text_emb]."""
+        prefix_text = "<|user|>\n<|begin_of_audio|>"
+        suffix_text = f"<|end_of_audio|><|user|>\n{prompt}<|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))
+
+        audio_emb = audio_embeds[0] if audio_embeds.dim() == 3 else audio_embeds
+        return torch.cat([prefix_emb, audio_emb, suffix_emb], dim=0)
+
+    def generate(self, inputs, prompt="转录以下音频内容", max_new_tokens=500,
+                 temperature=0.0, top_p=1.0, top_k=-1, repetition_penalty=1.0,
+                 **kwargs):
+        """Run batch ASR inference.
+
+        Args:
+            inputs: Audio file path(s), numpy arrays, or tensors.
+            prompt: Instruction prompt for ASR.
+            max_new_tokens: Maximum tokens to generate per sample.
+            temperature: Sampling temperature (0 = greedy decoding).
+            top_p: Nucleus sampling parameter.
+            top_k: Top-k sampling (-1 = disabled).
+            repetition_penalty: Repetition penalty factor. Non-neutral values are
+                forced back to 1.0 here because this engine feeds vLLM precomputed
+                embeddings (``enable_prompt_embeds=True``); see
+                ``resolve_repetition_penalty`` and issue #2948.
+
+        Returns:
+            List of {"key": str, "text": str}
+        """
+        from vllm import SamplingParams
+        try:
+            from vllm.inputs import EmbedsPrompt
+        except ImportError:
+            from vllm.inputs.data import EmbedsPrompt
+
+        if isinstance(inputs, (str, np.ndarray, torch.Tensor)):
+            inputs = [inputs]
+
+        sampling_params = SamplingParams(
+            max_tokens=max_new_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k if top_k > 0 else -1,
+            repetition_penalty=_safe_repetition_penalty(repetition_penalty),
+            skip_special_tokens=True,
+        )
+
+        t0 = time.perf_counter()
+        prompts = []
+        for audio_input in inputs:
+            audio_embeds = self._encode_audio(audio_input)
+            full_embeds = self._build_prompt_embeds(audio_embeds, prompt=prompt)
+            prompts.append(EmbedsPrompt(prompt_embeds=full_embeds.float()))
+        t1 = time.perf_counter()
+        logger.info(f"Audio encoding: {len(inputs)} samples in {t1-t0:.3f}s")
+
+        outputs = self.vllm_engine.generate(prompts, sampling_params, use_tqdm=False)
+        t2 = time.perf_counter()
+        logger.info(f"vLLM generation: {t2-t1:.3f}s")
+
+        raw_keys = [
+            os.path.splitext(os.path.basename(x))[0] if isinstance(x, str) else f"sample_{i}"
+            for i, x in enumerate(inputs)
+        ]
+        keys = _dedup_keys(raw_keys)
+
+        results = []
+        for i, output in enumerate(outputs):
+            token_ids = list(output.outputs[0].token_ids)
+            text = self.tokenizer.decode(token_ids, skip_special_tokens=True)
+            text = re.sub(r'\s+', ' ', text).strip()
+            results.append({"key": keys[i], "text": text})
+
+        return results
+
+    @classmethod
+    def from_pretrained(cls, model="zai-org/GLM-ASR-Nano-2512", hub="ms",
+                        device="cuda:0", dtype="bf16", **kwargs):
+        """Load from hub or local path."""
+        if os.path.isdir(model):
+            model_dir = model
+        else:
+            if hub in ("ms", "modelscope"):
+                from modelscope.hub.snapshot_download import snapshot_download
+                model_dir = snapshot_download(model, revision=kwargs.pop("revision", "master"))
+            else:
+                from huggingface_hub import snapshot_download
+                model_dir = snapshot_download(model)
+        return cls(model_dir=model_dir, device=device, dtype=dtype, **kwargs)

funasr/models/glm_asr/model.py

@@ -0,0 +1,154 @@
+import logging
+import os
+import time
+
+import torch
+import torch.nn as nn
+
+from funasr.register import tables
+
+
+@tables.register("model_classes", "GLMASR")
+@tables.register("model_classes", "zai-org/GLM-ASR-Nano-2512")
+@tables.register("model_classes", "ZhipuAI/GLM-ASR-Nano-2512")
+class GLMASR(nn.Module):
+
+    def __init__(self, **kwargs):
+        """Initialize GLMASR.
+        
+            Args:
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        model_path = kwargs.get("model_path", kwargs.get("model", "zai-org/GLM-ASR-Nano-2512"))
+        device = kwargs.get("device", "cuda:0")
+        dtype = kwargs.get("dtype", "bf16")
+        hub = kwargs.get("hub", "ms")
+        self._max_new_tokens = kwargs.get("max_new_tokens", 512)
+
+        self._dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}
+        self._device = device
+        self._torch_dtype = self._dtype_map.get(dtype, torch.bfloat16)
+        self._placeholder = nn.Parameter(torch.empty(0))
+
+        model_path = self._resolve_model_path(model_path, hub, kwargs)
+        self.model_path = model_path
+
+        from transformers import AutoModel as HFAutoModel
+        from transformers import AutoProcessor
+
+        self.processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
+        self.glm_model = HFAutoModel.from_pretrained(
+            model_path,
+            dtype=self._torch_dtype,
+            device_map=device,
+            trust_remote_code=True,
+        )
+        self.glm_model.eval()
+        logging.info(f"GLM-ASR model loaded from {model_path}")
+
+    def _resolve_model_path(self, model_path, hub, kwargs):
+        """Internal: resolve model path.
+        
+            Args:
+                model_path: TODO.
+                hub: TODO.
+                kwargs: Additional keyword arguments.
+            """
+        if os.path.exists(model_path):
+            return model_path
+
+        if hub in ("ms", "modelscope"):
+            try:
+                from modelscope.hub.snapshot_download import snapshot_download
+                model_revision = kwargs.get("model_revision", "master")
+                local_path = snapshot_download(model_path, revision=model_revision)
+                logging.info(f"Downloaded from ModelScope: {model_path} -> {local_path}")
+                return local_path
+            except Exception as e:
+                logging.warning(f"ModelScope download failed: {e}, falling back to HuggingFace path")
+
+        return model_path
+
+    def forward(self, **kwargs):
+        """Forward pass for training.
+        
+            Args:
+                **kwargs: Additional keyword arguments.
+            """
+        raise NotImplementedError("GLMASR only supports inference mode")
+
+    def inference(
+        self,
+        data_in,
+        data_lengths=None,
+        key: list = None,
+        tokenizer=None,
+        frontend=None,
+        **kwargs,
+    ):
+        """Run inference on input data.
+        
+            Args:
+                data_in: Input data (audio samples, file paths, or text).
+                data_lengths: Lengths of each input sample in the batch.
+                key: Sample identifiers.
+                tokenizer: Tokenizer instance for text encoding/decoding.
+                frontend: Audio frontend for feature extraction.
+                **kwargs: Additional keyword arguments.
+            """
+        meta_data = {}
+        time1 = time.perf_counter()
+
+        prompt = kwargs.get("prompt", "Please transcribe this audio into text")
+
+        if isinstance(data_in, (list, tuple)):
+            audio_list = list(data_in)
+        elif isinstance(data_in, str):
+            audio_list = [data_in]
+        else:
+            audio_list = [data_in]
+
+        time2 = time.perf_counter()
+        meta_data["load_data"] = f"{time2 - time1:0.3f}"
+
+        output = []
+        for i, audio_input in enumerate(audio_list):
+            messages = [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "audio", "url": audio_input},
+                        {"type": "text", "text": prompt},
+                    ],
+                }
+            ]
+
+            inputs = self.processor.apply_chat_template(
+                messages,
+                tokenize=True,
+                add_generation_prompt=True,
+                return_dict=True,
+                return_tensors="pt",
+            )
+            inputs = inputs.to(self._device, dtype=self._torch_dtype)
+
+            with torch.inference_mode():
+                generated = self.glm_model.generate(
+                    **inputs,
+                    max_new_tokens=self._max_new_tokens,
+                    do_sample=False,
+                )
+
+            text = self.processor.batch_decode(
+                generated[:, inputs["input_ids"].shape[1]:],
+                skip_special_tokens=True,
+            )[0].strip()
+
+            k = key[i] if key and i < len(key) else f"sample_{i}"
+            output.append({"key": k, "text": text})
+
+        time3 = time.perf_counter()
+        meta_data["batch_data_time"] = time3 - time2
+
+        return output, meta_data

funasr/models/glm_asr/vllm_utils.py

@@ -0,0 +1,43 @@
+"""Helpers for the GLM-ASR vLLM serving path.
+
+Kept dependency-free (standard library only) so the dtype guard can be unit
+tested without a CUDA device, a torch build, or a vLLM installation.
+"""
+
+import logging
+
+logger = logging.getLogger("funasr.glm_asr.vllm")
+
+# Compute dtype that is known to degrade GLM-ASR transcription quality.
+DEGRADED_DTYPE = "fp16"
+
+# Warn only once per process so batch loops do not spam the log.
+_warned_fp16 = False
+
+
+def warn_if_degraded_dtype(dtype):
+    """Warn once when a compute dtype is known to degrade GLM-ASR output.
+
+    ``fp16`` can produce degraded or garbage transcription for GLM-ASR
+    (numerical overflow in the audio embedding path), matching the documented
+    Fun-ASR-Nano behaviour. The value is still honoured -- some GPUs only
+    support fp16 -- but the caller is warned once about why output may be poor.
+
+    Args:
+        dtype: Requested compute dtype string ("bf16", "fp16", "fp32").
+
+    Returns:
+        ``dtype`` unchanged, so callers can wrap the value inline.
+    """
+    global _warned_fp16
+
+    if dtype == DEGRADED_DTYPE and not _warned_fp16:
+        logger.warning(
+            "dtype='fp16' can produce degraded or garbage transcription for "
+            "GLM-ASR (numerical overflow in the audio embedding path). "
+            "Use dtype='bf16' (recommended) or dtype='fp32'. On GPUs without "
+            "bfloat16 support (e.g. NVIDIA V100), use 'fp32'."
+        )
+        _warned_fp16 = True
+
+    return dtype