Initial commit: FunASR Speech Recognition Toolkit

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

funasr/models/whisper_lid/encoder.py

@@ -0,0 +1,137 @@
+# Copyright ESPnet (https://github.com/espnet/espnet). All Rights Reserved.
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+import copy
+from typing import Optional, Tuple, Union
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import whisper
+
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+from funasr.models.specaug.specaug import SpecAug
+from funasr.register import tables
+
+
+@tables.register("encoder_classes", "OpenAIWhisperEncoderWarp")
+class OpenAIWhisperEncoderWarp(nn.Module):
+    """Transformer-based Speech Encoder from OpenAI's Whisper Model:
+
+    URL: https://github.com/openai/whisper
+    """
+
+    def __init__(
+        self,
+        dropout_rate: float = 0.0,
+        whisper_model: str = "small",
+        download_dir: str = None,
+        use_specaug: bool = False,
+        use_padmask: bool = False,
+        specaug_conf: Union[dict, None] = None,
+    ):
+        """Initialize OpenAIWhisperEncoderWarp.
+        
+            Args:
+                dropout_rate: TODO.
+                whisper_model: Whisper Model instance.
+                download_dir: TODO.
+                use_specaug: TODO.
+                use_padmask: TODO.
+                specaug_conf: Configuration dict for specaug.
+            """
+        super().__init__()
+
+        # note that originally Whisper doesn't use dropouts
+        self.dropout = torch.nn.Dropout(dropout_rate)
+
+        assert whisper_model in whisper.available_models()
+        _model = whisper.load_model(whisper_model, download_root=download_dir, device="cpu")
+        self.encoders = copy.deepcopy(_model.encoder)
+        self.encoders.train()
+
+        del _model
+
+        if use_specaug:
+            self.specaug = SpecAug(**specaug_conf)
+        else:
+            self.specaug = None
+        self.use_padmask = use_padmask
+
+    def whisper_encode(
+        self,
+        input: torch.Tensor,
+        ilens: torch.Tensor = None,
+    ) -> torch.Tensor:
+        """Whisper encode.
+        
+            Args:
+                input: Input audio/text data.
+                ilens: TODO.
+            """
+        x = F.gelu(self.encoders.conv1(input))
+        x = F.gelu(self.encoders.conv2(x))
+        x = x.permute(0, 2, 1)
+
+        n_frames = x.size(1)
+        max_pos = self.encoders.positional_embedding.size(0)
+        if n_frames <= max_pos:
+            x = (x + self.encoders.positional_embedding[: x.size(1), :]).to(x.dtype)
+        else:
+            # due to positional encoding, audios >30 sec won't be accepted
+            x = x[:, :max_pos, :] + self.encoders.positional_embedding
+
+        if ilens is not None:
+            olens = (
+                1
+                + (ilens - self.encoders.conv2.kernel_size[0] + 2 * self.encoders.conv2.padding[0])
+                // self.encoders.conv2.stride[0]
+            )
+            olens = torch.clamp(olens, max=max_pos)
+        else:
+            olens = None
+
+        if self.use_padmask:
+            padding_mask = (~make_pad_mask(olens)[:, None, :]).to(x.device)
+        else:
+            padding_mask = None
+
+        x = self.dropout(x)
+
+        for layer, block in enumerate(self.encoders.blocks):
+            x = block(x)
+            if layer < len(self.encoders.blocks) - 1:
+                x = self.dropout(x)
+
+        x = self.encoders.ln_post(x)
+
+        return x, olens
+
+    def output_size(self) -> int:
+        # dummy output size
+        """Output size."""
+        return self.encoders.conv2.weight.shape[0]
+
+    def forward(
+        self,
+        xs_pad: torch.Tensor,
+        ilens: torch.Tensor,
+        prev_states: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        """Forward pass for training.
+        
+            Args:
+                xs_pad: TODO.
+                ilens: TODO.
+                prev_states: TODO.
+            """
+        feats, feats_lens = xs_pad, ilens
+
+        if self.specaug is not None and self.encoders.training:
+            feats = torch.transpose(feats, 1, 2)
+            feats, feats_lens = self.specaug(feats, feats_lens)
+            feats = torch.transpose(feats, 1, 2)
+
+        xs_pad, olens = self.whisper_encode(feats, feats_lens)
+
+        return xs_pad, olens, None