Initial commit: FunASR Speech Recognition Toolkit

Add complete FunASR codebase including models, runtime, and documentation.
2026-07-09 22:38:58 +08:00
commit 6116b1f3c6
3683 changed files with 990984 additions and 0 deletions
@@ -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