Initial commit: FunASR Speech Recognition Toolkit

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

funasr/models/rwkv_bat/rwkv_subsampling.py

@@ -0,0 +1,250 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import math
+import torch
+from typing import Optional, Tuple, Union
+from funasr.models.transformer.utils.nets_utils import pad_to_len
+
+
+class TooShortUttError(Exception):
+    """Raised when the utt is too short for subsampling.
+
+    Args:
+        message (str): Message for error catch
+        actual_size (int): the short size that cannot pass the subsampling
+        limit (int): the limit size for subsampling
+
+    """
+
+    def __init__(self, message, actual_size, limit):
+        """Construct a TooShortUttError for error handler."""
+        super().__init__(message)
+        self.actual_size = actual_size
+        self.limit = limit
+
+
+def check_short_utt(ins, size):
+    """Check if the utterance is too short for subsampling."""
+    if isinstance(ins, Conv2dSubsampling2) and size < 3:
+        return True, 3
+    if isinstance(ins, Conv2dSubsampling) and size < 7:
+        return True, 7
+    if isinstance(ins, Conv2dSubsampling6) and size < 11:
+        return True, 11
+    if isinstance(ins, Conv2dSubsampling8) and size < 15:
+        return True, 15
+    return False, -1
+
+
+class RWKVConvInput(torch.nn.Module):
+    """Streaming ConvInput module definition.
+    Args:
+        input_size: Input size.
+        conv_size: Convolution size.
+        subsampling_factor: Subsampling factor.
+        output_size: Block output dimension.
+    """
+
+    def __init__(
+        self,
+        input_size: int,
+        conv_size: Union[int, Tuple],
+        subsampling_factor: int = 4,
+        conv_kernel_size: int = 3,
+        output_size: Optional[int] = None,
+    ) -> None:
+        """Construct a ConvInput object."""
+        super().__init__()
+        if subsampling_factor == 1:
+            conv_size1, conv_size2, conv_size3 = conv_size
+
+            self.conv = torch.nn.Sequential(
+                torch.nn.Conv2d(
+                    1, conv_size1, conv_kernel_size, stride=1, padding=(conv_kernel_size - 1) // 2
+                ),
+                torch.nn.ReLU(),
+                torch.nn.Conv2d(
+                    conv_size1,
+                    conv_size1,
+                    conv_kernel_size,
+                    stride=[1, 2],
+                    padding=(conv_kernel_size - 1) // 2,
+                ),
+                torch.nn.ReLU(),
+                torch.nn.Conv2d(
+                    conv_size1,
+                    conv_size2,
+                    conv_kernel_size,
+                    stride=1,
+                    padding=(conv_kernel_size - 1) // 2,
+                ),
+                torch.nn.ReLU(),
+                torch.nn.Conv2d(
+                    conv_size2,
+                    conv_size2,
+                    conv_kernel_size,
+                    stride=[1, 2],
+                    padding=(conv_kernel_size - 1) // 2,
+                ),
+                torch.nn.ReLU(),
+                torch.nn.Conv2d(
+                    conv_size2,
+                    conv_size3,
+                    conv_kernel_size,
+                    stride=1,
+                    padding=(conv_kernel_size - 1) // 2,
+                ),
+                torch.nn.ReLU(),
+                torch.nn.Conv2d(
+                    conv_size3,
+                    conv_size3,
+                    conv_kernel_size,
+                    stride=[1, 2],
+                    padding=(conv_kernel_size - 1) // 2,
+                ),
+                torch.nn.ReLU(),
+            )
+
+            output_proj = conv_size3 * ((input_size // 2) // 2)
+
+            self.subsampling_factor = 1
+
+            self.stride_1 = 1
+
+            self.create_new_mask = self.create_new_vgg_mask
+
+        else:
+            conv_size1, conv_size2, conv_size3 = conv_size
+
+            kernel_1 = int(subsampling_factor / 2)
+
+            self.conv = torch.nn.Sequential(
+                torch.nn.Conv2d(
+                    1, conv_size1, conv_kernel_size, stride=1, padding=(conv_kernel_size - 1) // 2
+                ),
+                torch.nn.ReLU(),
+                torch.nn.Conv2d(
+                    conv_size1,
+                    conv_size1,
+                    conv_kernel_size,
+                    stride=[kernel_1, 2],
+                    padding=(conv_kernel_size - 1) // 2,
+                ),
+                torch.nn.ReLU(),
+                torch.nn.Conv2d(
+                    conv_size1,
+                    conv_size2,
+                    conv_kernel_size,
+                    stride=1,
+                    padding=(conv_kernel_size - 1) // 2,
+                ),
+                torch.nn.ReLU(),
+                torch.nn.Conv2d(
+                    conv_size2,
+                    conv_size2,
+                    conv_kernel_size,
+                    stride=[2, 2],
+                    padding=(conv_kernel_size - 1) // 2,
+                ),
+                torch.nn.ReLU(),
+                torch.nn.Conv2d(
+                    conv_size2,
+                    conv_size3,
+                    conv_kernel_size,
+                    stride=1,
+                    padding=(conv_kernel_size - 1) // 2,
+                ),
+                torch.nn.ReLU(),
+                torch.nn.Conv2d(
+                    conv_size3,
+                    conv_size3,
+                    conv_kernel_size,
+                    stride=1,
+                    padding=(conv_kernel_size - 1) // 2,
+                ),
+                torch.nn.ReLU(),
+            )
+
+            output_proj = conv_size3 * ((input_size // 2) // 2)
+
+            self.subsampling_factor = subsampling_factor
+
+            self.create_new_mask = self.create_new_vgg_mask
+
+            self.stride_1 = kernel_1
+
+        self.min_frame_length = 7
+
+        if output_size is not None:
+            self.output = torch.nn.Linear(output_proj, output_size)
+            self.output_size = output_size
+        else:
+            self.output = None
+            self.output_size = output_proj
+
+    def forward(
+        self, x: torch.Tensor, mask: Optional[torch.Tensor], chunk_size: Optional[torch.Tensor]
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Encode input sequences.
+        Args:
+            x: ConvInput input sequences. (B, T, D_feats)
+            mask: Mask of input sequences. (B, 1, T)
+        Returns:
+            x: ConvInput output sequences. (B, sub(T), D_out)
+            mask: Mask of output sequences. (B, 1, sub(T))
+        """
+        if mask is not None:
+            mask = self.create_new_mask(mask)
+            olens = max(mask.eq(0).sum(1))
+
+        b, t, f = x.size()
+        x = x.unsqueeze(1)  # (b. 1. t. f)
+
+        if chunk_size is not None:
+            max_input_length = int(
+                chunk_size
+                * self.subsampling_factor
+                * (math.ceil(float(t) / (chunk_size * self.subsampling_factor)))
+            )
+            x = map(lambda inputs: pad_to_len(inputs, max_input_length, 1), x)
+            x = list(x)
+            x = torch.stack(x, dim=0)
+            N_chunks = max_input_length // (chunk_size * self.subsampling_factor)
+            x = x.view(b * N_chunks, 1, chunk_size * self.subsampling_factor, f)
+
+        x = self.conv(x)
+
+        _, c, _, f = x.size()
+        if chunk_size is not None:
+            x = x.transpose(1, 2).contiguous().view(b, -1, c * f)[:, :olens, :]
+        else:
+            x = x.transpose(1, 2).contiguous().view(b, -1, c * f)
+
+        if self.output is not None:
+            x = self.output(x)
+
+        return x, mask[:, :olens][:, : x.size(1)]
+
+    def create_new_vgg_mask(self, mask: torch.Tensor) -> torch.Tensor:
+        """Create a new mask for VGG output sequences.
+        Args:
+            mask: Mask of input sequences. (B, T)
+        Returns:
+            mask: Mask of output sequences. (B, sub(T))
+        """
+        if self.subsampling_factor > 1:
+            return mask[:, ::2][:, :: self.stride_1]
+        else:
+            return mask
+
+    def get_size_before_subsampling(self, size: int) -> int:
+        """Return the original size before subsampling for a given size.
+        Args:
+            size: Number of frames after subsampling.
+        Returns:
+            : Number of frames before subsampling.
+        """
+        return size * self.subsampling_factor