Initial commit: FunASR Speech Recognition Toolkit

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

funasr/models/whisper/model.py

@@ -0,0 +1,190 @@
+from dataclasses import dataclass
+from typing import Dict
+from typing import Iterable, Optional
+import time
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch import nn
+
+import whisper
+
+# import whisper_timestamped as whisper
+
+from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
+
+from funasr.register import tables
+
+
+@tables.register("model_classes", "Whisper-tiny.en")
+@tables.register("model_classes", "Whisper-tiny")
+@tables.register("model_classes", "Whisper-base.en")
+@tables.register("model_classes", "Whisper-base")
+@tables.register("model_classes", "Whisper-small.en")
+@tables.register("model_classes", "Whisper-small")
+@tables.register("model_classes", "Whisper-medium.en")
+@tables.register("model_classes", "Whisper-medium")
+@tables.register("model_classes", "Whisper-large-v1")
+@tables.register("model_classes", "Whisper-large-v2")
+@tables.register("model_classes", "Whisper-large-v3")
+@tables.register("model_classes", "Whisper-large-v3-turbo")
+@tables.register("model_classes", "WhisperWarp")
+class WhisperWarp(nn.Module):
+    """Whisper: OpenAI Whisper model integration.
+
+    Wraps Whisper for multilingual speech recognition and translation
+    within FunASR's AutoModel interface.
+
+    Supports: whisper-tiny through whisper-large-v3-turbo.
+    Output: {"key": str, "text": str}
+    """
+
+    def __init__(self, *args, **kwargs):
+        """Initialize WhisperWarp.
+        
+            Args:
+                *args: Variable positional arguments.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        hub = kwargs.get("hub", "funasr")
+        if hub == "openai":
+            model_or_path = kwargs.get("model_path", "Whisper-large-v3")
+            if model_or_path.startswith("Whisper-"):
+                model_or_path = model_or_path.replace("Whisper-", "")
+            model = whisper.load_model(model_or_path)
+        else:
+            dims = kwargs.get("dims", {})
+            dims = whisper.model.ModelDimensions(**dims)
+            model = whisper.model.Whisper(dims=dims)
+
+        self.model = model
+
+        self.encoder_output_size = self.model.dims.n_audio_state
+
+    def forward(
+        self,
+        speech: torch.Tensor = None,
+        speech_lengths: torch.Tensor = None,
+        text: torch.Tensor = None,
+        text_lengths: torch.Tensor = None,
+        **kwargs,
+    ):
+        """Forward pass for training. Computes cross-entropy loss.
+
+        Args:
+            speech: (B, T, D) mel-spectrogram features
+            speech_lengths: (B,) lengths of each audio
+            text: (B, U) token IDs (with SOT/EOT tokens)
+            text_lengths: (B,) lengths of each text sequence
+        Returns:
+            dict with "loss" and optionally "stats"
+        """
+        if speech is None or text is None:
+            raise ValueError("forward() requires speech and text for training")
+
+        # Encoder
+        audio_features = self.model.encoder(speech)
+
+        # Decoder: shift text right for teacher forcing
+        # text format: [SOT, lang, task, ..., tokens, EOT]
+        decoder_input = text[:, :-1]
+        decoder_target = text[:, 1:]
+
+        # Decoder forward
+        logits = self.model.decoder(decoder_input, audio_features)
+
+        # Cross-entropy loss (ignore padding, token_id = -1 or pad)
+        loss = F.cross_entropy(
+            logits.reshape(-1, logits.size(-1)),
+            decoder_target.reshape(-1),
+            ignore_index=-100,
+        )
+
+        stats = {
+            "loss": loss.detach().item(),
+            "batch_size": speech.size(0),
+        }
+        return {"loss": loss, "stats": stats}
+
+    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.
+            """
+        if kwargs.get("batch_size", 1) > 1:
+            raise NotImplementedError("batch decoding is not implemented")
+
+        if frontend is None and not hasattr(self, "frontend"):
+            frontend_class = tables.frontend_classes.get("WhisperFrontend")
+            frontend = frontend_class(
+                n_mels=self.model.dims.n_mels, do_pad_trim=kwargs.get("do_pad_trim", True)
+            )
+            self.frontend = frontend
+        else:
+            frontend = frontend if frontend is not None else self.frontend
+
+        meta_data = {}
+        if (
+            isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank"
+        ):  # fbank
+            speech, speech_lengths = data_in, data_lengths
+            if len(speech.shape) < 3:
+                speech = speech[None, :, :]
+            if speech_lengths is None:
+                speech_lengths = speech.shape[1]
+        else:
+            # extract fbank feats
+            time1 = time.perf_counter()
+            audio_sample_list = load_audio_text_image_video(
+                data_in,
+                fs=frontend.fs if hasattr(frontend, "fs") else 16000,
+                audio_fs=kwargs.get("fs", 16000),
+                data_type=kwargs.get("data_type", "sound"),
+                tokenizer=tokenizer,
+            )
+            time2 = time.perf_counter()
+            meta_data["load_data"] = f"{time2 - time1:0.3f}"
+            speech, speech_lengths = extract_fbank(
+                audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend
+            )
+            time3 = time.perf_counter()
+            meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
+            frame_shift = frontend.frame_shift if hasattr(frontend, "frame_shift") else 10
+            lfr_n = frontend.lfr_n if hasattr(frontend, "lfr_n") else 1
+            meta_data["batch_data_time"] = speech_lengths.sum().item() * frame_shift * lfr_n / 1000
+
+        speech = speech.to(device=kwargs["device"])[0, :, :]
+        speech_lengths = speech_lengths.to(device=kwargs["device"])
+
+        # # detect the spoken language
+        # _, probs = self.model.detect_language(speech)
+        # print(f"Detected language: {max(probs, key=probs.get)}")
+
+        # decode the audio
+        options = whisper.DecodingOptions(**kwargs.get("DecodingOptions", {}))
+
+        result = whisper.decode(self.model, speech, options=options)
+        # result = whisper.transcribe(self.model, speech)
+
+        results = []
+        result_i = {"key": key[0], "text": result.text}
+
+        results.append(result_i)
+
+        return results, meta_data

funasr/models/whisper/template.yaml

@@ -0,0 +1,45 @@
+# This is an example that demonstrates how to configure a model file.
+# You can modify the configuration according to your own requirements.
+
+# to print the register_table:
+# from funasr.register import tables
+# tables.print()
+
+# network architecture
+model: WhisperWarp
+model_conf:
+    lsm_weight: 0.1
+    length_normalized_loss: true
+    hub: funasr # openai
+
+
+
+# only use for hub == funasr,
+#  if hub == openai, dims is automaticall download
+dims:
+    n_mels: 128
+    n_vocab: 51866
+    n_audio_ctx: 1500
+    n_audio_state: 1280
+    n_audio_head: 20
+    n_audio_layer: 32
+    n_text_ctx: 448
+    n_text_state: 1280
+    n_text_head: 20
+    n_text_layer: 32
+
+# frontend related
+frontend: WhisperFrontend
+frontend_conf:
+    fs: 16000
+    n_mels: ${dims.n_mels}
+    do_pad_trim: true
+
+tokenizer: WhisperTokenizer
+tokenizer_conf:
+  language: null
+  task: transcribe
+  is_multilingual: true
+  num_languages: 100
+
+scope_map: [none, "model."]