Refactor ASR module; add Paraformer support

config/interface/options.yaml

@@ -7,10 +7,8 @@ asr_models:
     name: Whisper medium
   - id: openai/whisper-small
     name: Whisper small
-  - id: sanchit-gandhi/whisper-small-dv
-    name: Whisper small-dv
-  - id: facebook/wav2vec2-base-960h
-    name: Wav2Vec2-Base-960h
+  - id: funasr/paraformer-zh
+    name: Paraformer-zh
 
 llm_models:
   - id: gemini-2.5-flash

modules/asr/__init__.py

@@ -0,0 +1,11 @@
+from .base import AbstractASRModel
+from .registry import ASR_MODEL_REGISTRY, get_asr_model, register_asr_model
+from .whisper import WhisperASR
+from .paraformer import ParaformerASR
+
+__all__ = [
+    "AbstractASRModel",
+    "get_asr_model",
+    "register_asr_model",
+    "ASR_MODEL_REGISTRY",
+]

modules/asr/base.py

@@ -0,0 +1,24 @@
+from abc import ABC, abstractmethod
+from typing import Optional
+
+import numpy as np
+
+
+class AbstractASRModel(ABC):
+    def __init__(
+        self, model_id: str, device: str = "cpu", cache_dir: str = "cache", **kwargs
+    ):
+        print(f"Loading ASR model {model_id}...")
+        self.model_id = model_id
+        self.device = device
+        self.cache_dir = cache_dir
+
+    @abstractmethod
+    def transcribe(
+        self,
+        audio: np.ndarray,
+        audio_sample_rate: int,
+        language: Optional[str] = None,
+        **kwargs,
+    ) -> str:
+        pass

modules/asr/paraformer.py

@@ -0,0 +1,86 @@
+import os
+import tempfile
+from typing import Optional
+
+import numpy as np
+import soundfile as sf
+
+try:
+    from funasr import AutoModel
+except ImportError:
+    AutoModel = None
+
+from .base import AbstractASRModel
+from .registry import register_asr_model
+
+
+@register_asr_model("funasr/paraformer-zh")
+class ParaformerASR(AbstractASRModel):
+    def __init__(
+        self, model_id: str, device: str = "cpu", cache_dir: str = "cache", **kwargs
+    ):
+        super().__init__(model_id, device, cache_dir, **kwargs)
+
+        if AutoModel is None:
+            raise ImportError(
+                "funasr is not installed. Please install it with: pip3 install -U funasr"
+            )
+
+        model_name = model_id.replace("funasr/", "")
+        language = model_name.split("-")[1]
+        if language == "zh":
+            self.language = "mandarin"
+        elif language == "en":
+            self.language = "english"
+        else:
+            raise ValueError(
+                f"Language cannot be determined. {model_id} is not supported"
+            )
+
+        try:
+            original_cache_dir = os.getenv("MODELSCOPE_CACHE")
+            os.makedirs(cache_dir, exist_ok=True)
+            os.environ["MODELSCOPE_CACHE"] = cache_dir
+            self.model = AutoModel(
+                model=model_name,
+                model_revision="v2.0.4",
+                vad_model="fsmn-vad",
+                vad_model_revision="v2.0.4",
+                punc_model="ct-punc-c",
+                punc_model_revision="v2.0.4",
+                device=device,
+            )
+            if original_cache_dir:
+                os.environ["MODELSCOPE_CACHE"] = original_cache_dir
+            else:
+                del os.environ["MODELSCOPE_CACHE"]
+
+        except Exception as e:
+            raise ValueError(f"Error loading Paraformer model: {e}")
+
+    def transcribe(
+        self,
+        audio: np.ndarray,
+        audio_sample_rate: int,
+        language: Optional[str] = None,
+        **kwargs,
+    ) -> str:
+        if language and language != self.language:
+            raise ValueError(
+                f"Paraformer model {self.model_id} only supports {self.language} language, but {language} was requested"
+            )
+
+        try:
+            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+                sf.write(f.name, audio, audio_sample_rate)
+                temp_file = f.name
+
+            result = self.model.generate(input=temp_file, batch_size_s=300, **kwargs)
+
+            os.unlink(temp_file)
+
+            print(f"Transcription result: {result}, type: {type(result)}")
+
+            return result[0]["text"]
+        except Exception as e:
+            raise ValueError(f"Error during transcription: {e}")

modules/asr/registry.py

@@ -0,0 +1,19 @@
+from .base import AbstractASRModel
+
+ASR_MODEL_REGISTRY = {}
+
+
+def register_asr_model(prefix: str):
+    def wrapper(cls):
+        assert issubclass(cls, AbstractASRModel), f"{cls} must inherit AbstractASRModel"
+        ASR_MODEL_REGISTRY[prefix] = cls
+        return cls
+
+    return wrapper
+
+
+def get_asr_model(model_id: str, device="cpu", **kwargs) -> AbstractASRModel:
+    for prefix, cls in ASR_MODEL_REGISTRY.items():
+        if model_id.startswith(prefix):
+            return cls(model_id, device=device, **kwargs)
+    raise ValueError(f"No ASR wrapper found for model: {model_id}")

modules/asr/whisper.py

@@ -0,0 +1,82 @@
+import os
+from typing import Optional
+
+import librosa
+import numpy as np
+from transformers.pipelines import pipeline
+
+from .base import AbstractASRModel
+from .registry import register_asr_model
+
+hf_token = os.getenv("HF_TOKEN")
+
+
+@register_asr_model("openai/whisper")
+class WhisperASR(AbstractASRModel):
+    def __init__(
+        self, model_id: str, device: str = "cpu", cache_dir: str = "cache", **kwargs
+    ):
+        super().__init__(model_id, device, cache_dir, **kwargs)
+        model_kwargs = kwargs.setdefault("model_kwargs", {})
+        model_kwargs["cache_dir"] = cache_dir
+        self.pipe = pipeline(
+            "automatic-speech-recognition",
+            model=model_id,
+            device=0 if device == "cuda" else -1,
+            token=hf_token,
+            **kwargs,
+        )
+
+    def transcribe(
+        self,
+        audio: np.ndarray,
+        audio_sample_rate: int,
+        language: Optional[str] = None,
+        **kwargs,
+    ) -> str:
+        """
+        Transcribe audio using Whisper model
+
+        Args:
+            audio: Audio numpy array
+            audio_sample_rate: Sample rate of the audio
+            language: Language hint (optional)
+
+        Returns:
+            Transcribed text as string
+        """
+        try:
+            # Resample to 16kHz if needed
+            if audio_sample_rate != 16000:
+                audio = librosa.resample(
+                    audio, orig_sr=audio_sample_rate, target_sr=16000
+                )
+
+            # Generate transcription
+            generate_kwargs = {}
+            if language:
+                generate_kwargs["language"] = language
+
+            result = self.pipe(
+                audio,
+                generate_kwargs=generate_kwargs,
+                return_timestamps=False,
+                **kwargs,
+            )
+
+            # Extract text from result
+            if isinstance(result, dict) and "text" in result:
+                return result["text"]
+            elif isinstance(result, list) and len(result) > 0:
+                # Handle list of results
+                first_result = result[0]
+                if isinstance(first_result, dict):
+                    return first_result.get("text", str(first_result))
+                else:
+                    return str(first_result)
+            else:
+                return str(result)
+
+        except Exception as e:
+            print(f"Error during Whisper transcription: {e}")
+            return ""

tests/test_asr_infer.py

@@ -0,0 +1,19 @@
+from modules.asr import get_asr_model
+import librosa
+
+if __name__ == "__main__":
+    supported_asrs = [
+        "funasr/paraformer-zh",
+        "openai/whisper-large-v3-turbo",
+    ]
+    for model_id in supported_asrs:
+        try:
+            print(f"Loading model: {model_id}")
+            asr_model = get_asr_model(model_id, device="cpu", cache_dir=".cache")
+            audio, sample_rate = librosa.load("tests/audio/hello.wav", sr=None)
+            result = asr_model.transcribe(audio, sample_rate, language="mandarin")
+            print(result)
+        except Exception as e:
+            print(f"Failed to load model {model_id}: {e}")
+            breakpoint()
+            continue