Initial commit: FunASR Speech Recognition Toolkit

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

funasr/frontends/abs_frontend.py

@@ -0,0 +1,12 @@
+from abc import ABC, abstractmethod
+import torch
+
+
+class AbsFrontend(ABC, torch.nn.Module):
+    @abstractmethod
+    def output_size(self) -> int:
+        raise NotImplementedError
+
+    @abstractmethod
+    def forward(self, input, input_lengths):
+        raise NotImplementedError

funasr/frontends/default.py

@@ -0,0 +1,416 @@
+import copy
+from typing import Optional
+from typing import Tuple
+from typing import Union
+import logging
+import numpy as np
+import torch
+import torch.nn as nn
+
+try:
+    from torch_complex.tensor import ComplexTensor
+except:
+    print("Please install torch_complex firstly")
+
+from funasr.frontends.utils.log_mel import LogMel
+from funasr.frontends.utils.stft import Stft
+from funasr.frontends.utils.frontend import Frontend
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+from funasr.register import tables
+
+
+@tables.register("frontend_classes", "DefaultFrontend")
+@tables.register("frontend_classes", "EspnetFrontend")
+class DefaultFrontend(nn.Module):
+    """Conventional frontend structure for ASR.
+    Stft -> WPE -> MVDR-Beamformer -> Power-spec -> Mel-Fbank -> CMVN
+    """
+
+    def __init__(
+        self,
+        fs: int = 16000,
+        n_fft: int = 512,
+        win_length: int = None,
+        hop_length: int = 128,
+        window: Optional[str] = "hann",
+        center: bool = True,
+        normalized: bool = False,
+        onesided: bool = True,
+        n_mels: int = 80,
+        fmin: int = None,
+        fmax: int = None,
+        htk: bool = False,
+        frontend_conf: Optional[dict] = None,
+        apply_stft: bool = True,
+        use_channel: int = None,
+        **kwargs,
+    ):
+        """Initialize DefaultFrontend.
+        
+            Args:
+                fs: TODO.
+                n_fft: TODO.
+                win_length: TODO.
+                hop_length: TODO.
+                window: TODO.
+                center: TODO.
+                normalized: TODO.
+                onesided: TODO.
+                n_mels: TODO.
+                fmin: TODO.
+                fmax: TODO.
+                htk: TODO.
+                frontend_conf: Configuration dict for frontend.
+                apply_stft: TODO.
+                use_channel: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+
+        # Deepcopy (In general, dict shouldn't be used as default arg)
+        frontend_conf = copy.deepcopy(frontend_conf)
+        self.hop_length = hop_length
+        self.fs = fs
+
+        if apply_stft:
+            self.stft = Stft(
+                n_fft=n_fft,
+                win_length=win_length,
+                hop_length=hop_length,
+                center=center,
+                window=window,
+                normalized=normalized,
+                onesided=onesided,
+            )
+        else:
+            self.stft = None
+        self.apply_stft = apply_stft
+
+        if frontend_conf is not None:
+            self.frontend = Frontend(idim=n_fft // 2 + 1, **frontend_conf)
+        else:
+            self.frontend = None
+
+        self.logmel = LogMel(
+            fs=fs,
+            n_fft=n_fft,
+            n_mels=n_mels,
+            fmin=fmin,
+            fmax=fmax,
+            htk=htk,
+        )
+        self.n_mels = n_mels
+        self.use_channel = use_channel
+        self.frontend_type = "default"
+
+    def output_size(self) -> int:
+        """Output size."""
+        return self.n_mels
+
+    def forward(
+        self, input: torch.Tensor, input_lengths: Union[torch.Tensor, list]
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass for training.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+            """
+        if isinstance(input_lengths, list):
+            input_lengths = torch.tensor(input_lengths)
+        if input.dtype == torch.float64:
+            input = input.float()
+        # 1. Domain-conversion: e.g. Stft: time -> time-freq
+        if self.stft is not None:
+            input_stft, feats_lens = self._compute_stft(input, input_lengths)
+        else:
+            input_stft = ComplexTensor(input[..., 0], input[..., 1])
+            feats_lens = input_lengths
+        # 2. [Option] Speech enhancement
+        if self.frontend is not None:
+            assert isinstance(input_stft, ComplexTensor), type(input_stft)
+            # input_stft: (Batch, Length, [Channel], Freq)
+            input_stft, _, mask = self.frontend(input_stft, feats_lens)
+
+        # 3. [Multi channel case]: Select a channel
+        if input_stft.dim() == 4:
+            # h: (B, T, C, F) -> h: (B, T, F)
+            if self.training:
+                if self.use_channel is not None:
+                    input_stft = input_stft[:, :, self.use_channel, :]
+                else:
+                    # Select 1ch randomly
+                    ch = np.random.randint(input_stft.size(2))
+                    input_stft = input_stft[:, :, ch, :]
+            else:
+                # Use the first channel
+                input_stft = input_stft[:, :, 0, :]
+
+        # 4. STFT -> Power spectrum
+        # h: ComplexTensor(B, T, F) -> torch.Tensor(B, T, F)
+        input_power = input_stft.real**2 + input_stft.imag**2
+
+        # 5. Feature transform e.g. Stft -> Log-Mel-Fbank
+        # input_power: (Batch, [Channel,] Length, Freq)
+        #       -> input_feats: (Batch, Length, Dim)
+        input_feats, _ = self.logmel(input_power, feats_lens)
+
+        return input_feats, feats_lens
+
+    def _compute_stft(self, input: torch.Tensor, input_lengths: torch.Tensor) -> torch.Tensor:
+        """Internal: compute stft.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+            """
+        input_stft, feats_lens = self.stft(input, input_lengths)
+
+        assert input_stft.dim() >= 4, input_stft.shape
+        # "2" refers to the real/imag parts of Complex
+        assert input_stft.shape[-1] == 2, input_stft.shape
+
+        # Change torch.Tensor to ComplexTensor
+        # input_stft: (..., F, 2) -> (..., F)
+        input_stft = ComplexTensor(input_stft[..., 0], input_stft[..., 1])
+        return input_stft, feats_lens
+
+
+class MultiChannelFrontend(nn.Module):
+    """Conventional frontend structure for ASR.
+    Stft -> WPE -> MVDR-Beamformer -> Power-spec -> Mel-Fbank -> CMVN
+    """
+
+    def __init__(
+        self,
+        fs: int = 16000,
+        n_fft: int = 512,
+        win_length: int = None,
+        hop_length: int = None,
+        frame_length: int = None,
+        frame_shift: int = None,
+        window: Optional[str] = "hann",
+        center: bool = True,
+        normalized: bool = False,
+        onesided: bool = True,
+        n_mels: int = 80,
+        fmin: int = None,
+        fmax: int = None,
+        htk: bool = False,
+        frontend_conf: Optional[dict] = None,
+        apply_stft: bool = True,
+        use_channel: int = None,
+        lfr_m: int = 1,
+        lfr_n: int = 1,
+        cmvn_file: str = None,
+        mc: bool = True,
+    ):
+        """Initialize MultiChannelFrontend.
+        
+            Args:
+                fs: TODO.
+                n_fft: TODO.
+                win_length: TODO.
+                hop_length: TODO.
+                frame_length: TODO.
+                frame_shift: TODO.
+                window: TODO.
+                center: TODO.
+                normalized: TODO.
+                onesided: TODO.
+                n_mels: TODO.
+                fmin: TODO.
+                fmax: TODO.
+                htk: TODO.
+                frontend_conf: Configuration dict for frontend.
+                apply_stft: TODO.
+                use_channel: TODO.
+                lfr_m: TODO.
+                lfr_n: TODO.
+                cmvn_file: TODO.
+                mc: TODO.
+            """
+        super().__init__()
+        # Deepcopy (In general, dict shouldn't be used as default arg)
+        frontend_conf = copy.deepcopy(frontend_conf)
+        if win_length is None and hop_length is None:
+            self.win_length = frame_length * 16
+            self.hop_length = frame_shift * 16
+        elif frame_length is None and frame_shift is None:
+            self.win_length = self.win_length
+            self.hop_length = self.hop_length
+        else:
+            logging.error(
+                "Only one of (win_length, hop_length) and (frame_length, frame_shift)" "can be set."
+            )
+            exit(1)
+
+        if apply_stft:
+            self.stft = Stft(
+                n_fft=n_fft,
+                win_length=self.win_length,
+                hop_length=self.hop_length,
+                center=center,
+                window=window,
+                normalized=normalized,
+                onesided=onesided,
+            )
+        else:
+            self.stft = None
+        self.apply_stft = apply_stft
+
+        if frontend_conf is not None:
+            self.frontend = Frontend(idim=n_fft // 2 + 1, **frontend_conf)
+        else:
+            self.frontend = None
+
+        self.logmel = LogMel(
+            fs=fs,
+            n_fft=n_fft,
+            n_mels=n_mels,
+            fmin=fmin,
+            fmax=fmax,
+            htk=htk,
+        )
+        self.n_mels = n_mels
+        self.use_channel = use_channel
+        self.mc = mc
+        if not self.mc:
+            if self.use_channel is not None:
+                logging.info("use the channel %d" % (self.use_channel))
+            else:
+                logging.info("random select channel")
+            self.cmvn_file = cmvn_file
+            if self.cmvn_file is not None:
+                mean, std = self._load_cmvn(self.cmvn_file)
+                self.register_buffer("mean", torch.from_numpy(mean))
+                self.register_buffer("std", torch.from_numpy(std))
+        self.frontend_type = "multichannelfrontend"
+
+    def output_size(self) -> int:
+        """Output size."""
+        return self.n_mels
+
+    def forward(
+        self, input: torch.Tensor, input_lengths: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # 1. Domain-conversion: e.g. Stft: time -> time-freq
+        """Forward pass for training.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+            """
+        if self.stft is not None:
+            input_stft, feats_lens = self._compute_stft(input, input_lengths)
+        else:
+            input_stft = ComplexTensor(input[..., 0], input[..., 1])
+            feats_lens = input_lengths
+        # 2. [Option] Speech enhancement
+        if self.frontend is not None:
+            assert isinstance(input_stft, ComplexTensor), type(input_stft)
+            # input_stft: (Batch, Length, [Channel], Freq)
+            input_stft, _, mask = self.frontend(input_stft, feats_lens)
+
+        # 3. [Multi channel case]: Select a channel(sa_asr)
+        if input_stft.dim() == 4 and not self.mc:
+            # h: (B, T, C, F) -> h: (B, T, F)
+            if self.training:
+                if self.use_channel is not None:
+                    input_stft = input_stft[:, :, self.use_channel, :]
+
+                else:
+                    # Select 1ch randomly
+                    ch = np.random.randint(input_stft.size(2))
+                    input_stft = input_stft[:, :, ch, :]
+            else:
+                # Use the first channel
+                input_stft = input_stft[:, :, 0, :]
+
+        # 4. STFT -> Power spectrum
+        # h: ComplexTensor(B, T, F) -> torch.Tensor(B, T, F)
+        input_power = input_stft.real**2 + input_stft.imag**2
+
+        # 5. Feature transform e.g. Stft -> Log-Mel-Fbank
+        # input_power: (Batch, [Channel,] Length, Freq)
+        #       -> input_feats: (Batch, Length, Dim)
+        input_feats, _ = self.logmel(input_power, feats_lens)
+        if self.mc:
+            # MFCCA
+            if input_feats.dim() == 4:
+                bt = input_feats.size(0)
+                channel_size = input_feats.size(2)
+                input_feats = (
+                    input_feats.transpose(1, 2).reshape(bt * channel_size, -1, 80).contiguous()
+                )
+                feats_lens = feats_lens.repeat(1, channel_size).squeeze()
+            else:
+                channel_size = 1
+            return input_feats, feats_lens, channel_size
+        else:
+            # 6. Apply CMVN
+            if self.cmvn_file is not None:
+                if feats_lens is None:
+                    feats_lens = input_feats.new_full([input_feats.size(0)], input_feats.size(1))
+                self.mean = self.mean.to(input_feats.device, input_feats.dtype)
+                self.std = self.std.to(input_feats.device, input_feats.dtype)
+                mask = make_pad_mask(feats_lens, input_feats, 1)
+
+                if input_feats.requires_grad:
+                    input_feats = input_feats + self.mean
+                else:
+                    input_feats += self.mean
+                if input_feats.requires_grad:
+                    input_feats = input_feats.masked_fill(mask, 0.0)
+                else:
+                    input_feats.masked_fill_(mask, 0.0)
+
+                input_feats *= self.std
+
+            return input_feats, feats_lens
+
+    def _compute_stft(self, input: torch.Tensor, input_lengths: torch.Tensor) -> torch.Tensor:
+        """Internal: compute stft.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+            """
+        input_stft, feats_lens = self.stft(input, input_lengths)
+
+        assert input_stft.dim() >= 4, input_stft.shape
+        # "2" refers to the real/imag parts of Complex
+        assert input_stft.shape[-1] == 2, input_stft.shape
+
+        # Change torch.Tensor to ComplexTensor
+        # input_stft: (..., F, 2) -> (..., F)
+        input_stft = ComplexTensor(input_stft[..., 0], input_stft[..., 1])
+        return input_stft, feats_lens
+
+    def _load_cmvn(self, cmvn_file):
+        """Internal: load cmvn.
+        
+            Args:
+                cmvn_file: TODO.
+            """
+        with open(cmvn_file, "r", encoding="utf-8") as f:
+            lines = f.readlines()
+        means_list = []
+        vars_list = []
+        for i in range(len(lines)):
+            line_item = lines[i].split()
+            if line_item[0] == "<AddShift>":
+                line_item = lines[i + 1].split()
+                if line_item[0] == "<LearnRateCoef>":
+                    add_shift_line = line_item[3 : (len(line_item) - 1)]
+                    means_list = list(add_shift_line)
+                    continue
+            elif line_item[0] == "<Rescale>":
+                line_item = lines[i + 1].split()
+                if line_item[0] == "<LearnRateCoef>":
+                    rescale_line = line_item[3 : (len(line_item) - 1)]
+                    vars_list = list(rescale_line)
+                    continue
+        means = np.array(means_list).astype(np.float)
+        vars = np.array(vars_list).astype(np.float)
+        return means, vars

funasr/frontends/eend_ola_feature.py

@@ -0,0 +1,73 @@
+# Copyright 2019 Hitachi, Ltd. (author: Yusuke Fujita)
+# Licensed under the MIT license.
+#
+# This module is for computing audio features
+
+import librosa
+import numpy as np
+
+
+def transform(Y, dtype=np.float32):
+    """Transform.
+    
+        Args:
+            Y: TODO.
+            dtype: TODO.
+        """
+    Y = np.abs(Y)
+    n_fft = 2 * (Y.shape[1] - 1)
+    sr = 8000
+    n_mels = 23
+    mel_basis = librosa.filters.mel(sr, n_fft, n_mels)
+    Y = np.dot(Y**2, mel_basis.T)
+    Y = np.log10(np.maximum(Y, 1e-10))
+    mean = np.mean(Y, axis=0)
+    Y = Y - mean
+    return Y.astype(dtype)
+
+
+def subsample(Y, T, subsampling=1):
+    """Subsample.
+    
+        Args:
+            Y: TODO.
+            T: TODO.
+            subsampling: TODO.
+        """
+    Y_ss = Y[::subsampling]
+    T_ss = T[::subsampling]
+    return Y_ss, T_ss
+
+
+def splice(Y, context_size=0):
+    """Splice.
+    
+        Args:
+            Y: TODO.
+            context_size: Size/dimension parameter.
+        """
+    Y_pad = np.pad(Y, [(context_size, context_size), (0, 0)], "constant")
+    Y_spliced = np.lib.stride_tricks.as_strided(
+        np.ascontiguousarray(Y_pad),
+        (Y.shape[0], Y.shape[1] * (2 * context_size + 1)),
+        (Y.itemsize * Y.shape[1], Y.itemsize),
+        writeable=False,
+    )
+    return Y_spliced
+
+
+def stft(data, frame_size=1024, frame_shift=256):
+    """Stft.
+    
+        Args:
+            data: TODO.
+            frame_size: Size/dimension parameter.
+            frame_shift: TODO.
+        """
+    fft_size = 1 << (frame_size - 1).bit_length()
+    if len(data) % frame_shift == 0:
+        return librosa.stft(data, n_fft=fft_size, win_length=frame_size, hop_length=frame_shift).T[
+            :-1
+        ]
+    else:
+        return librosa.stft(data, n_fft=fft_size, win_length=frame_size, hop_length=frame_shift).T

funasr/frontends/fused.py

@@ -0,0 +1,157 @@
+from funasr.frontends.default import DefaultFrontend
+from funasr.frontends.s3prl import S3prlFrontend
+import numpy as np
+import torch
+import torch.nn as nn
+from typing import Tuple
+
+
+class FusedFrontends(nn.Module):
+    def __init__(self, frontends=None, align_method="linear_projection", proj_dim=100, fs=16000):
+
+        """Initialize FusedFrontends.
+        
+            Args:
+                frontends: TODO.
+                align_method: TODO.
+                proj_dim: Size/dimension parameter.
+                fs: TODO.
+            """
+        super().__init__()
+        self.align_method = align_method  # fusing method : linear_projection only for now
+        self.proj_dim = proj_dim  # dim of the projection done on each frontend
+        self.frontends = []  # list of the frontends to combine
+
+        for i, frontend in enumerate(frontends):
+            frontend_type = frontend["frontend_type"]
+            if frontend_type == "default":
+                n_mels, fs, n_fft, win_length, hop_length = (
+                    frontend.get("n_mels", 80),
+                    fs,
+                    frontend.get("n_fft", 512),
+                    frontend.get("win_length"),
+                    frontend.get("hop_length", 128),
+                )
+                window, center, normalized, onesided = (
+                    frontend.get("window", "hann"),
+                    frontend.get("center", True),
+                    frontend.get("normalized", False),
+                    frontend.get("onesided", True),
+                )
+                fmin, fmax, htk, apply_stft = (
+                    frontend.get("fmin", None),
+                    frontend.get("fmax", None),
+                    frontend.get("htk", False),
+                    frontend.get("apply_stft", True),
+                )
+
+                self.frontends.append(
+                    DefaultFrontend(
+                        n_mels=n_mels,
+                        n_fft=n_fft,
+                        fs=fs,
+                        win_length=win_length,
+                        hop_length=hop_length,
+                        window=window,
+                        center=center,
+                        normalized=normalized,
+                        onesided=onesided,
+                        fmin=fmin,
+                        fmax=fmax,
+                        htk=htk,
+                        apply_stft=apply_stft,
+                    )
+                )
+            elif frontend_type == "s3prl":
+                frontend_conf, download_dir, multilayer_feature = (
+                    frontend.get("frontend_conf"),
+                    frontend.get("download_dir"),
+                    frontend.get("multilayer_feature"),
+                )
+                self.frontends.append(
+                    S3prlFrontend(
+                        fs=fs,
+                        frontend_conf=frontend_conf,
+                        download_dir=download_dir,
+                        multilayer_feature=multilayer_feature,
+                    )
+                )
+
+            else:
+                raise NotImplementedError  # frontends are only default or s3prl
+
+        self.frontends = torch.nn.ModuleList(self.frontends)
+
+        self.gcd = np.gcd.reduce([frontend.hop_length for frontend in self.frontends])
+        self.factors = [frontend.hop_length // self.gcd for frontend in self.frontends]
+        if torch.cuda.is_available():
+            dev = "cuda"
+        elif torch.xpu.is_available():
+            dev = "xpu"
+        elif torch.backends.mps.is_available():
+            dev = "mps"
+        else:
+            dev = "cpu"
+        if self.align_method == "linear_projection":
+            self.projection_layers = [
+                torch.nn.Linear(
+                    in_features=frontend.output_size(),
+                    out_features=self.factors[i] * self.proj_dim,
+                )
+                for i, frontend in enumerate(self.frontends)
+            ]
+            self.projection_layers = torch.nn.ModuleList(self.projection_layers)
+            self.projection_layers = self.projection_layers.to(torch.device(dev))
+
+    def output_size(self) -> int:
+        """Output size."""
+        return len(self.frontends) * self.proj_dim
+
+    def forward(
+        self, input: torch.Tensor, input_lengths: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        # step 0 : get all frontends features
+        """Forward pass for training.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+            """
+        self.feats = []
+        for frontend in self.frontends:
+            with torch.no_grad():
+                input_feats, feats_lens = frontend.forward(input, input_lengths)
+            self.feats.append([input_feats, feats_lens])
+
+        if self.align_method == "linear_projection":  # TODO(Dan): to add other align methods
+
+            # first step : projections
+            self.feats_proj = []
+            for i, frontend in enumerate(self.frontends):
+                input_feats = self.feats[i][0]
+                self.feats_proj.append(self.projection_layers[i](input_feats))
+
+            # 2nd step : reshape
+            self.feats_reshaped = []
+            for i, frontend in enumerate(self.frontends):
+                input_feats_proj = self.feats_proj[i]
+                bs, nf, dim = input_feats_proj.shape
+                input_feats_reshaped = torch.reshape(
+                    input_feats_proj, (bs, nf * self.factors[i], dim // self.factors[i])
+                )
+                self.feats_reshaped.append(input_feats_reshaped)
+
+            # 3rd step : drop the few last frames
+            m = min([x.shape[1] for x in self.feats_reshaped])
+            self.feats_final = [x[:, :m, :] for x in self.feats_reshaped]
+
+            input_feats = torch.cat(
+                self.feats_final, dim=-1
+            )  # change the input size of the preencoder : proj_dim * n_frontends
+            feats_lens = torch.ones_like(self.feats[0][1]) * (m)
+
+        else:
+            raise NotImplementedError
+
+        return input_feats, feats_lens

funasr/frontends/s3prl.py

@@ -0,0 +1,166 @@
+import copy
+import logging
+import os
+from argparse import Namespace
+from typing import Optional
+from typing import Tuple
+from typing import Union
+
+try:
+    import humanfriendly
+except ImportError:
+    humanfriendly = None
+import torch
+import torch.nn as nn
+
+from funasr.frontends.utils.frontend import Frontend
+from funasr.models.transformer.utils.nets_utils import pad_list
+
+
+def base_s3prl_setup(args):
+    """Base s3prl setup.
+    
+        Args:
+            args: TODO.
+        """
+    args.upstream_feature_selection = getattr(args, "upstream_feature_selection", None)
+    args.upstream_model_config = getattr(args, "upstream_model_config", None)
+    args.upstream_refresh = getattr(args, "upstream_refresh", False)
+    args.upstream_ckpt = getattr(args, "upstream_ckpt", None)
+    args.init_ckpt = getattr(args, "init_ckpt", None)
+    args.verbose = getattr(args, "verbose", False)
+    args.tile_factor = getattr(args, "tile_factor", 1)
+    return args
+
+
+class S3prlFrontend(nn.Module):
+    """Speech Pretrained Representation frontend structure for ASR."""
+
+    def __init__(
+        self,
+        fs: Union[int, str] = 16000,
+        frontend_conf: Optional[dict] = None,
+        download_dir: str = None,
+        multilayer_feature: bool = False,
+    ):
+        """Initialize S3prlFrontend.
+        
+            Args:
+                fs: TODO.
+                frontend_conf: Configuration dict for frontend.
+                download_dir: TODO.
+                multilayer_feature: TODO.
+            """
+        super().__init__()
+        if isinstance(fs, str):
+            if humanfriendly is not None:
+                fs = humanfriendly.parse_size(fs)
+            else:
+                fs = int(fs)
+
+        if download_dir is not None:
+            torch.hub.set_dir(download_dir)
+
+        self.multilayer_feature = multilayer_feature
+        self.upstream, self.featurizer = self._get_upstream(frontend_conf)
+        self.pretrained_params = copy.deepcopy(self.upstream.state_dict())
+        self.output_dim = self.featurizer.output_dim
+        self.frontend_type = "s3prl"
+        self.hop_length = self.upstream.get_downsample_rates("key")
+
+    def _get_upstream(self, frontend_conf):
+        """Get S3PRL upstream model."""
+        s3prl_args = base_s3prl_setup(
+            Namespace(**frontend_conf, device="cpu"),
+        )
+        self.args = s3prl_args
+
+        s3prl_path = None
+        python_path_list = os.environ.get("PYTHONPATH", "(None)").split(":")
+        for p in python_path_list:
+            if p.endswith("s3prl"):
+                s3prl_path = p
+                break
+        assert s3prl_path is not None
+
+        s3prl_upstream = torch.hub.load(
+            s3prl_path,
+            s3prl_args.upstream,
+            ckpt=s3prl_args.upstream_ckpt,
+            model_config=s3prl_args.upstream_model_config,
+            refresh=s3prl_args.upstream_refresh,
+            source="local",
+        ).to("cpu")
+
+        if getattr(
+            s3prl_upstream, "model", None
+        ) is not None and s3prl_upstream.model.__class__.__name__ in [
+            "Wav2Vec2Model",
+            "HubertModel",
+        ]:
+            s3prl_upstream.model.encoder.layerdrop = 0.0
+
+        from s3prl.upstream.interfaces import Featurizer
+
+        if self.multilayer_feature is None:
+            feature_selection = "last_hidden_state"
+        else:
+            feature_selection = "hidden_states"
+        s3prl_featurizer = Featurizer(
+            upstream=s3prl_upstream,
+            feature_selection=feature_selection,
+            upstream_device="cpu",
+        )
+
+        return s3prl_upstream, s3prl_featurizer
+
+    def _tile_representations(self, feature):
+        """Tile up the representations by `tile_factor`.
+        Input - sequence of representations
+                shape: (batch_size, seq_len, feature_dim)
+        Output - sequence of tiled representations
+                 shape: (batch_size, seq_len * factor, feature_dim)
+        """
+        assert len(feature.shape) == 3, "Input argument `feature` has invalid shape: {}".format(
+            feature.shape
+        )
+        tiled_feature = feature.repeat(1, 1, self.args.tile_factor)
+        tiled_feature = tiled_feature.reshape(
+            feature.size(0), feature.size(1) * self.args.tile_factor, feature.size(2)
+        )
+        return tiled_feature
+
+    def output_size(self) -> int:
+        """Output size."""
+        return self.output_dim
+
+    def forward(
+        self, input: torch.Tensor, input_lengths: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass for training.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+            """
+        wavs = [wav[: input_lengths[i]] for i, wav in enumerate(input)]
+        self.upstream.eval()
+        with torch.no_grad():
+            feats = self.upstream(wavs)
+        feats = self.featurizer(wavs, feats)
+
+        if self.args.tile_factor != 1:
+            feats = self._tile_representations(feats)
+
+        input_feats = pad_list(feats, 0.0)
+        feats_lens = torch.tensor([f.shape[0] for f in feats], dtype=torch.long)
+
+        # Saving CUDA Memory
+        del feats
+
+        return input_feats, feats_lens
+
+    def reload_pretrained_parameters(self):
+        """Reload pretrained parameters."""
+        self.upstream.load_state_dict(self.pretrained_params)
+        logging.info("Pretrained S3PRL frontend model parameters reloaded!")

funasr/frontends/utils/__init__.py

@@ -0,0 +1 @@
+"""Initialize sub package."""

funasr/frontends/utils/beamformer.py

@@ -0,0 +1,88 @@
+import torch
+from torch_complex import functional as FC
+from torch_complex.tensor import ComplexTensor
+
+
+def get_power_spectral_density_matrix(
+    xs: ComplexTensor, mask: torch.Tensor, normalization=True, eps: float = 1e-15
+) -> ComplexTensor:
+    """Return cross-channel power spectral density (PSD) matrix
+
+    Args:
+        xs (ComplexTensor): (..., F, C, T)
+        mask (torch.Tensor): (..., F, C, T)
+        normalization (bool):
+        eps (float):
+    Returns
+        psd (ComplexTensor): (..., F, C, C)
+
+    """
+    # outer product: (..., C_1, T) x (..., C_2, T) -> (..., T, C, C_2)
+    psd_Y = FC.einsum("...ct,...et->...tce", [xs, xs.conj()])
+
+    # Averaging mask along C: (..., C, T) -> (..., T)
+    mask = mask.mean(dim=-2)
+
+    # Normalized mask along T: (..., T)
+    if normalization:
+        # If assuming the tensor is padded with zero, the summation along
+        # the time axis is same regardless of the padding length.
+        mask = mask / (mask.sum(dim=-1, keepdim=True) + eps)
+
+    # psd: (..., T, C, C)
+    psd = psd_Y * mask[..., None, None]
+    # (..., T, C, C) -> (..., C, C)
+    psd = psd.sum(dim=-3)
+
+    return psd
+
+
+def get_mvdr_vector(
+    psd_s: ComplexTensor,
+    psd_n: ComplexTensor,
+    reference_vector: torch.Tensor,
+    eps: float = 1e-15,
+) -> ComplexTensor:
+    """Return the MVDR(Minimum Variance Distortionless Response) vector:
+
+        h = (Npsd^-1 @ Spsd) / (Tr(Npsd^-1 @ Spsd)) @ u
+
+    Reference:
+        On optimal frequency-domain multichannel linear filtering
+        for noise reduction; M. Souden et al., 2010;
+        https://ieeexplore.ieee.org/document/5089420
+
+    Args:
+        psd_s (ComplexTensor): (..., F, C, C)
+        psd_n (ComplexTensor): (..., F, C, C)
+        reference_vector (torch.Tensor): (..., C)
+        eps (float):
+    Returns:
+        beamform_vector (ComplexTensor)r: (..., F, C)
+    """
+    # Add eps
+    C = psd_n.size(-1)
+    eye = torch.eye(C, dtype=psd_n.dtype, device=psd_n.device)
+    shape = [1 for _ in range(psd_n.dim() - 2)] + [C, C]
+    eye = eye.view(*shape)
+    psd_n += eps * eye
+
+    # numerator: (..., C_1, C_2) x (..., C_2, C_3) -> (..., C_1, C_3)
+    numerator = FC.einsum("...ec,...cd->...ed", [psd_n.inverse(), psd_s])
+    # ws: (..., C, C) / (...,) -> (..., C, C)
+    ws = numerator / (FC.trace(numerator)[..., None, None] + eps)
+    # h: (..., F, C_1, C_2) x (..., C_2) -> (..., F, C_1)
+    beamform_vector = FC.einsum("...fec,...c->...fe", [ws, reference_vector])
+    return beamform_vector
+
+
+def apply_beamforming_vector(beamform_vector: ComplexTensor, mix: ComplexTensor) -> ComplexTensor:
+    # (..., C) x (..., C, T) -> (..., T)
+    """Apply beamforming vector.
+    
+        Args:
+            beamform_vector: TODO.
+            mix: TODO.
+        """
+    es = FC.einsum("...c,...ct->...t", [beamform_vector.conj(), mix])
+    return es

funasr/frontends/utils/complex_utils.py

@@ -0,0 +1,259 @@
+"""Beamformer module."""
+
+from distutils.version import LooseVersion
+from typing import Sequence
+from typing import Tuple
+from typing import Union
+
+import torch
+
+try:
+    from torch_complex import functional as FC
+    from torch_complex.tensor import ComplexTensor
+except:
+    print("Please install torch_complex firstly")
+
+
+EPS = torch.finfo(torch.double).eps
+is_torch_1_8_plus = LooseVersion(torch.__version__) >= LooseVersion("1.8.0")
+is_torch_1_9_plus = LooseVersion(torch.__version__) >= LooseVersion("1.9.0")
+
+
+def new_complex_like(
+    ref: Union[torch.Tensor, ComplexTensor],
+    real_imag: Tuple[torch.Tensor, torch.Tensor],
+):
+    """New complex like.
+    
+        Args:
+            ref: TODO.
+            real_imag: TODO.
+        """
+    if isinstance(ref, ComplexTensor):
+        return ComplexTensor(*real_imag)
+    elif is_torch_complex_tensor(ref):
+        return torch.complex(*real_imag)
+    else:
+        raise ValueError("Please update your PyTorch version to 1.9+ for complex support.")
+
+
+def is_torch_complex_tensor(c):
+    """Is torch complex tensor.
+    
+        Args:
+            c: TODO.
+        """
+    return not isinstance(c, ComplexTensor) and is_torch_1_9_plus and torch.is_complex(c)
+
+
+def is_complex(c):
+    """Is complex.
+    
+        Args:
+            c: TODO.
+        """
+    return isinstance(c, ComplexTensor) or is_torch_complex_tensor(c)
+
+
+def to_double(c):
+    """To double.
+    
+        Args:
+            c: TODO.
+        """
+    if not isinstance(c, ComplexTensor) and is_torch_1_9_plus and torch.is_complex(c):
+        return c.to(dtype=torch.complex128)
+    else:
+        return c.double()
+
+
+def to_float(c):
+    """To float.
+    
+        Args:
+            c: TODO.
+        """
+    if not isinstance(c, ComplexTensor) and is_torch_1_9_plus and torch.is_complex(c):
+        return c.to(dtype=torch.complex64)
+    else:
+        return c.float()
+
+
+def cat(seq: Sequence[Union[ComplexTensor, torch.Tensor]], *args, **kwargs):
+    """Cat.
+    
+        Args:
+            seq: TODO.
+            *args: Variable positional arguments.
+            **kwargs: Additional keyword arguments.
+        """
+    if not isinstance(seq, (list, tuple)):
+        raise TypeError(
+            "cat(): argument 'tensors' (position 1) must be tuple of Tensors, " "not Tensor"
+        )
+    if isinstance(seq[0], ComplexTensor):
+        return FC.cat(seq, *args, **kwargs)
+    else:
+        return torch.cat(seq, *args, **kwargs)
+
+
+def complex_norm(c: Union[torch.Tensor, ComplexTensor], dim=-1, keepdim=False) -> torch.Tensor:
+    """Complex norm.
+    
+        Args:
+            c: TODO.
+            dim: TODO.
+            keepdim: TODO.
+        """
+    if not is_complex(c):
+        raise TypeError("Input is not a complex tensor.")
+    if is_torch_complex_tensor(c):
+        return torch.norm(c, dim=dim, keepdim=keepdim)
+    else:
+        return torch.sqrt((c.real**2 + c.imag**2).sum(dim=dim, keepdim=keepdim) + EPS)
+
+
+def einsum(equation, *operands):
+    # NOTE: Do not mix ComplexTensor and torch.complex in the input!
+    # NOTE (wangyou): Until PyTorch 1.9.0, torch.einsum does not support
+    # mixed input with complex and real tensors.
+    """Einsum.
+    
+        Args:
+            equation: TODO.
+            *operands: Variable positional arguments.
+        """
+    if len(operands) == 1:
+        if isinstance(operands[0], (tuple, list)):
+            operands = operands[0]
+        complex_module = FC if isinstance(operands[0], ComplexTensor) else torch
+        return complex_module.einsum(equation, *operands)
+    elif len(operands) != 2:
+        op0 = operands[0]
+        same_type = all(op.dtype == op0.dtype for op in operands[1:])
+        if same_type:
+            _einsum = FC.einsum if isinstance(op0, ComplexTensor) else torch.einsum
+            return _einsum(equation, *operands)
+        else:
+            raise ValueError("0 or More than 2 operands are not supported.")
+    a, b = operands
+    if isinstance(a, ComplexTensor) or isinstance(b, ComplexTensor):
+        return FC.einsum(equation, a, b)
+    elif is_torch_1_9_plus and (torch.is_complex(a) or torch.is_complex(b)):
+        if not torch.is_complex(a):
+            o_real = torch.einsum(equation, a, b.real)
+            o_imag = torch.einsum(equation, a, b.imag)
+            return torch.complex(o_real, o_imag)
+        elif not torch.is_complex(b):
+            o_real = torch.einsum(equation, a.real, b)
+            o_imag = torch.einsum(equation, a.imag, b)
+            return torch.complex(o_real, o_imag)
+        else:
+            return torch.einsum(equation, a, b)
+    else:
+        return torch.einsum(equation, a, b)
+
+
+def inverse(c: Union[torch.Tensor, ComplexTensor]) -> Union[torch.Tensor, ComplexTensor]:
+    """Inverse.
+    
+        Args:
+            c: TODO.
+        """
+    if isinstance(c, ComplexTensor):
+        return c.inverse2()
+    else:
+        return c.inverse()
+
+
+def matmul(
+    a: Union[torch.Tensor, ComplexTensor], b: Union[torch.Tensor, ComplexTensor]
+) -> Union[torch.Tensor, ComplexTensor]:
+    # NOTE: Do not mix ComplexTensor and torch.complex in the input!
+    # NOTE (wangyou): Until PyTorch 1.9.0, torch.matmul does not support
+    # multiplication between complex and real tensors.
+    """Matmul.
+    
+        Args:
+            a: TODO.
+            b: TODO.
+        """
+    if isinstance(a, ComplexTensor) or isinstance(b, ComplexTensor):
+        return FC.matmul(a, b)
+    elif is_torch_1_9_plus and (torch.is_complex(a) or torch.is_complex(b)):
+        if not torch.is_complex(a):
+            o_real = torch.matmul(a, b.real)
+            o_imag = torch.matmul(a, b.imag)
+            return torch.complex(o_real, o_imag)
+        elif not torch.is_complex(b):
+            o_real = torch.matmul(a.real, b)
+            o_imag = torch.matmul(a.imag, b)
+            return torch.complex(o_real, o_imag)
+        else:
+            return torch.matmul(a, b)
+    else:
+        return torch.matmul(a, b)
+
+
+def trace(a: Union[torch.Tensor, ComplexTensor]):
+    # NOTE (wangyou): until PyTorch 1.9.0, torch.trace does not
+    # support bacth processing. Use FC.trace() as fallback.
+    """Trace.
+    
+        Args:
+            a: TODO.
+        """
+    return FC.trace(a)
+
+
+def reverse(a: Union[torch.Tensor, ComplexTensor], dim=0):
+    """Reverse.
+    
+        Args:
+            a: TODO.
+            dim: TODO.
+        """
+    if isinstance(a, ComplexTensor):
+        return FC.reverse(a, dim=dim)
+    else:
+        return torch.flip(a, dims=(dim,))
+
+
+def solve(b: Union[torch.Tensor, ComplexTensor], a: Union[torch.Tensor, ComplexTensor]):
+    """Solve the linear equation ax = b."""
+    # NOTE: Do not mix ComplexTensor and torch.complex in the input!
+    # NOTE (wangyou): Until PyTorch 1.9.0, torch.solve does not support
+    # mixed input with complex and real tensors.
+    if isinstance(a, ComplexTensor) or isinstance(b, ComplexTensor):
+        if isinstance(a, ComplexTensor) and isinstance(b, ComplexTensor):
+            return FC.solve(b, a, return_LU=False)
+        else:
+            return matmul(inverse(a), b)
+    elif is_torch_1_9_plus and (torch.is_complex(a) or torch.is_complex(b)):
+        if torch.is_complex(a) and torch.is_complex(b):
+            return torch.linalg.solve(a, b)
+        else:
+            return matmul(inverse(a), b)
+    else:
+        if is_torch_1_8_plus:
+            return torch.linalg.solve(a, b)
+        else:
+            return torch.solve(b, a)[0]
+
+
+def stack(seq: Sequence[Union[ComplexTensor, torch.Tensor]], *args, **kwargs):
+    """Stack.
+    
+        Args:
+            seq: TODO.
+            *args: Variable positional arguments.
+            **kwargs: Additional keyword arguments.
+        """
+    if not isinstance(seq, (list, tuple)):
+        raise TypeError(
+            "stack(): argument 'tensors' (position 1) must be tuple of Tensors, " "not Tensor"
+        )
+    if isinstance(seq[0], ComplexTensor):
+        return FC.stack(seq, *args, **kwargs)
+    else:
+        return torch.stack(seq, *args, **kwargs)

funasr/frontends/utils/dnn_beamformer.py

@@ -0,0 +1,189 @@
+"""DNN beamformer module."""
+
+from typing import Tuple
+
+import torch
+from torch.nn import functional as F
+
+from funasr.frontends.utils.beamformer import apply_beamforming_vector
+from funasr.frontends.utils.beamformer import get_mvdr_vector
+from funasr.frontends.utils.beamformer import (
+    get_power_spectral_density_matrix,  # noqa: H301
+)
+from funasr.frontends.utils.mask_estimator import MaskEstimator
+from torch_complex.tensor import ComplexTensor
+
+
+class DNN_Beamformer(torch.nn.Module):
+    """DNN mask based Beamformer
+
+    Citation:
+        Multichannel End-to-end Speech Recognition; T. Ochiai et al., 2017;
+        https://arxiv.org/abs/1703.04783
+
+    """
+
+    def __init__(
+        self,
+        bidim,
+        btype="blstmp",
+        blayers=3,
+        bunits=300,
+        bprojs=320,
+        bnmask=2,
+        dropout_rate=0.0,
+        badim=320,
+        ref_channel: int = -1,
+        beamformer_type="mvdr",
+    ):
+        """Initialize DNN_Beamformer.
+        
+            Args:
+                bidim: TODO.
+                btype: TODO.
+                blayers: TODO.
+                bunits: TODO.
+                bprojs: TODO.
+                bnmask: TODO.
+                dropout_rate: TODO.
+                badim: TODO.
+                ref_channel: TODO.
+                beamformer_type: TODO.
+            """
+        super().__init__()
+        self.mask = MaskEstimator(btype, bidim, blayers, bunits, bprojs, dropout_rate, nmask=bnmask)
+        self.ref = AttentionReference(bidim, badim)
+        self.ref_channel = ref_channel
+
+        self.nmask = bnmask
+
+        if beamformer_type != "mvdr":
+            raise ValueError("Not supporting beamformer_type={}".format(beamformer_type))
+        self.beamformer_type = beamformer_type
+
+    def forward(
+        self, data: ComplexTensor, ilens: torch.LongTensor
+    ) -> Tuple[ComplexTensor, torch.LongTensor, ComplexTensor]:
+        """The forward function
+
+        Notation:
+            B: Batch
+            C: Channel
+            T: Time or Sequence length
+            F: Freq
+
+        Args:
+            data (ComplexTensor): (B, T, C, F)
+            ilens (torch.Tensor): (B,)
+        Returns:
+            enhanced (ComplexTensor): (B, T, F)
+            ilens (torch.Tensor): (B,)
+
+        """
+
+        def apply_beamforming(data, ilens, psd_speech, psd_noise):
+            # u: (B, C)
+            """Apply beamforming.
+            
+                Args:
+                    data: TODO.
+                    ilens: TODO.
+                    psd_speech: TODO.
+                    psd_noise: TODO.
+                """
+            if self.ref_channel < 0:
+                u, _ = self.ref(psd_speech, ilens)
+            else:
+                # (optional) Create onehot vector for fixed reference microphone
+                u = torch.zeros(*(data.size()[:-3] + (data.size(-2),)), device=data.device)
+                u[..., self.ref_channel].fill_(1)
+
+            ws = get_mvdr_vector(psd_speech, psd_noise, u)
+            enhanced = apply_beamforming_vector(ws, data)
+
+            return enhanced, ws
+
+        # data (B, T, C, F) -> (B, F, C, T)
+        data = data.permute(0, 3, 2, 1)
+
+        # mask: (B, F, C, T)
+        masks, _ = self.mask(data, ilens)
+        assert self.nmask == len(masks)
+
+        if self.nmask == 2:  # (mask_speech, mask_noise)
+            mask_speech, mask_noise = masks
+
+            psd_speech = get_power_spectral_density_matrix(data, mask_speech)
+            psd_noise = get_power_spectral_density_matrix(data, mask_noise)
+
+            enhanced, ws = apply_beamforming(data, ilens, psd_speech, psd_noise)
+
+            # (..., F, T) -> (..., T, F)
+            enhanced = enhanced.transpose(-1, -2)
+            mask_speech = mask_speech.transpose(-1, -3)
+        else:  # multi-speaker case: (mask_speech1, ..., mask_noise)
+            mask_speech = list(masks[:-1])
+            mask_noise = masks[-1]
+
+            psd_speeches = [get_power_spectral_density_matrix(data, mask) for mask in mask_speech]
+            psd_noise = get_power_spectral_density_matrix(data, mask_noise)
+
+            enhanced = []
+            ws = []
+            for i in range(self.nmask - 1):
+                psd_speech = psd_speeches.pop(i)
+                # treat all other speakers' psd_speech as noises
+                enh, w = apply_beamforming(data, ilens, psd_speech, sum(psd_speeches) + psd_noise)
+                psd_speeches.insert(i, psd_speech)
+
+                # (..., F, T) -> (..., T, F)
+                enh = enh.transpose(-1, -2)
+                mask_speech[i] = mask_speech[i].transpose(-1, -3)
+
+                enhanced.append(enh)
+                ws.append(w)
+
+        return enhanced, ilens, mask_speech
+
+
+class AttentionReference(torch.nn.Module):
+    def __init__(self, bidim, att_dim):
+        """Initialize AttentionReference.
+        
+            Args:
+                bidim: TODO.
+                att_dim: Size/dimension parameter.
+            """
+        super().__init__()
+        self.mlp_psd = torch.nn.Linear(bidim, att_dim)
+        self.gvec = torch.nn.Linear(att_dim, 1)
+
+    def forward(
+        self, psd_in: ComplexTensor, ilens: torch.LongTensor, scaling: float = 2.0
+    ) -> Tuple[torch.Tensor, torch.LongTensor]:
+        """The forward function
+
+        Args:
+            psd_in (ComplexTensor): (B, F, C, C)
+            ilens (torch.Tensor): (B,)
+            scaling (float):
+        Returns:
+            u (torch.Tensor): (B, C)
+            ilens (torch.Tensor): (B,)
+        """
+        B, _, C = psd_in.size()[:3]
+        assert psd_in.size(2) == psd_in.size(3), psd_in.size()
+        # psd_in: (B, F, C, C)
+        psd = psd_in.masked_fill(torch.eye(C, dtype=torch.bool, device=psd_in.device), 0)
+        # psd: (B, F, C, C) -> (B, C, F)
+        psd = (psd.sum(dim=-1) / (C - 1)).transpose(-1, -2)
+
+        # Calculate amplitude
+        psd_feat = (psd.real**2 + psd.imag**2) ** 0.5
+
+        # (B, C, F) -> (B, C, F2)
+        mlp_psd = self.mlp_psd(psd_feat)
+        # (B, C, F2) -> (B, C, 1) -> (B, C)
+        e = self.gvec(torch.tanh(mlp_psd)).squeeze(-1)
+        u = F.softmax(scaling * e, dim=-1)
+        return u, ilens

funasr/frontends/utils/dnn_wpe.py

@@ -0,0 +1,108 @@
+from typing import Tuple
+
+from pytorch_wpe import wpe_one_iteration
+import torch
+from torch_complex.tensor import ComplexTensor
+
+from funasr.frontends.utils.mask_estimator import MaskEstimator
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+
+
+class DNN_WPE(torch.nn.Module):
+    def __init__(
+        self,
+        wtype: str = "blstmp",
+        widim: int = 257,
+        wlayers: int = 3,
+        wunits: int = 300,
+        wprojs: int = 320,
+        dropout_rate: float = 0.0,
+        taps: int = 5,
+        delay: int = 3,
+        use_dnn_mask: bool = True,
+        iterations: int = 1,
+        normalization: bool = False,
+    ):
+        """Initialize DNN_WPE.
+        
+            Args:
+                wtype: TODO.
+                widim: TODO.
+                wlayers: TODO.
+                wunits: TODO.
+                wprojs: TODO.
+                dropout_rate: TODO.
+                taps: TODO.
+                delay: TODO.
+                use_dnn_mask: TODO.
+                iterations: TODO.
+                normalization: TODO.
+            """
+        super().__init__()
+        self.iterations = iterations
+        self.taps = taps
+        self.delay = delay
+
+        self.normalization = normalization
+        self.use_dnn_mask = use_dnn_mask
+
+        self.inverse_power = True
+
+        if self.use_dnn_mask:
+            self.mask_est = MaskEstimator(
+                wtype, widim, wlayers, wunits, wprojs, dropout_rate, nmask=1
+            )
+
+    def forward(
+        self, data: ComplexTensor, ilens: torch.LongTensor
+    ) -> Tuple[ComplexTensor, torch.LongTensor, ComplexTensor]:
+        """The forward function
+
+        Notation:
+            B: Batch
+            C: Channel
+            T: Time or Sequence length
+            F: Freq or Some dimension of the feature vector
+
+        Args:
+            data: (B, C, T, F)
+            ilens: (B,)
+        Returns:
+            data: (B, C, T, F)
+            ilens: (B,)
+        """
+        # (B, T, C, F) -> (B, F, C, T)
+        enhanced = data = data.permute(0, 3, 2, 1)
+        mask = None
+
+        for i in range(self.iterations):
+            # Calculate power: (..., C, T)
+            power = enhanced.real**2 + enhanced.imag**2
+            if i == 0 and self.use_dnn_mask:
+                # mask: (B, F, C, T)
+                (mask,), _ = self.mask_est(enhanced, ilens)
+                if self.normalization:
+                    # Normalize along T
+                    mask = mask / mask.sum(dim=-1)[..., None]
+                # (..., C, T) * (..., C, T) -> (..., C, T)
+                power = power * mask
+
+            # Averaging along the channel axis: (..., C, T) -> (..., T)
+            power = power.mean(dim=-2)
+
+            # enhanced: (..., C, T) -> (..., C, T)
+            enhanced = wpe_one_iteration(
+                data.contiguous(),
+                power,
+                taps=self.taps,
+                delay=self.delay,
+                inverse_power=self.inverse_power,
+            )
+
+            enhanced.masked_fill_(make_pad_mask(ilens, enhanced.real), 0)
+
+        # (B, F, C, T) -> (B, T, C, F)
+        enhanced = enhanced.permute(0, 3, 2, 1)
+        if mask is not None:
+            mask = mask.transpose(-1, -3)
+        return enhanced, ilens, mask

funasr/frontends/utils/feature_transform.py

@@ -0,0 +1,331 @@
+from typing import List
+from typing import Tuple
+from typing import Union
+
+import librosa
+import numpy as np
+import torch
+from torch_complex.tensor import ComplexTensor
+
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+
+
+class FeatureTransform(torch.nn.Module):
+    def __init__(
+        self,
+        # Mel options,
+        fs: int = 16000,
+        n_fft: int = 512,
+        n_mels: int = 80,
+        fmin: float = 0.0,
+        fmax: float = None,
+        # Normalization
+        stats_file: str = None,
+        apply_uttmvn: bool = True,
+        uttmvn_norm_means: bool = True,
+        uttmvn_norm_vars: bool = False,
+    ):
+        """Initialize FeatureTransform.
+        
+            Args:
+                fs: TODO.
+                n_fft: TODO.
+                n_mels: TODO.
+                fmin: TODO.
+                fmax: TODO.
+                stats_file: TODO.
+                apply_uttmvn: TODO.
+                uttmvn_norm_means: TODO.
+                uttmvn_norm_vars: TODO.
+            """
+        super().__init__()
+        self.apply_uttmvn = apply_uttmvn
+
+        self.logmel = LogMel(fs=fs, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax)
+        self.stats_file = stats_file
+        if stats_file is not None:
+            self.global_mvn = GlobalMVN(stats_file)
+        else:
+            self.global_mvn = None
+
+        if self.apply_uttmvn is not None:
+            self.uttmvn = UtteranceMVN(norm_means=uttmvn_norm_means, norm_vars=uttmvn_norm_vars)
+        else:
+            self.uttmvn = None
+
+    def forward(
+        self, x: ComplexTensor, ilens: Union[torch.LongTensor, np.ndarray, List[int]]
+    ) -> Tuple[torch.Tensor, torch.LongTensor]:
+        # (B, T, F) or (B, T, C, F)
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+                ilens: TODO.
+            """
+        if x.dim() not in (3, 4):
+            raise ValueError(f"Input dim must be 3 or 4: {x.dim()}")
+        if not torch.is_tensor(ilens):
+            ilens = torch.from_numpy(np.asarray(ilens)).to(x.device)
+
+        if x.dim() == 4:
+            # h: (B, T, C, F) -> h: (B, T, F)
+            if self.training:
+                # Select 1ch randomly
+                ch = np.random.randint(x.size(2))
+                h = x[:, :, ch, :]
+            else:
+                # Use the first channel
+                h = x[:, :, 0, :]
+        else:
+            h = x
+
+        # h: ComplexTensor(B, T, F) -> torch.Tensor(B, T, F)
+        h = h.real**2 + h.imag**2
+
+        h, _ = self.logmel(h, ilens)
+        if self.stats_file is not None:
+            h, _ = self.global_mvn(h, ilens)
+        if self.apply_uttmvn:
+            h, _ = self.uttmvn(h, ilens)
+
+        return h, ilens
+
+
+class LogMel(torch.nn.Module):
+    """Convert STFT to fbank feats
+
+    The arguments is same as librosa.filters.mel
+
+    Args:
+        fs: number > 0 [scalar] sampling rate of the incoming signal
+        n_fft: int > 0 [scalar] number of FFT components
+        n_mels: int > 0 [scalar] number of Mel bands to generate
+        fmin: float >= 0 [scalar] lowest frequency (in Hz)
+        fmax: float >= 0 [scalar] highest frequency (in Hz).
+            If `None`, use `fmax = fs / 2.0`
+        htk: use HTK formula instead of Slaney
+        norm: {None, 1, np.inf} [scalar]
+            if 1, divide the triangular mel weights by the width of the mel band
+            (area normalization).  Otherwise, leave all the triangles aiming for
+            a peak value of 1.0
+
+    """
+
+    def __init__(
+        self,
+        fs: int = 16000,
+        n_fft: int = 512,
+        n_mels: int = 80,
+        fmin: float = 0.0,
+        fmax: float = None,
+        htk: bool = False,
+        norm=1,
+    ):
+        """Initialize LogMel.
+        
+            Args:
+                fs: TODO.
+                n_fft: TODO.
+                n_mels: TODO.
+                fmin: TODO.
+                fmax: TODO.
+                htk: TODO.
+                norm: TODO.
+            """
+        super().__init__()
+
+        _mel_options = dict(
+            sr=fs, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax, htk=htk, norm=norm
+        )
+        self.mel_options = _mel_options
+
+        # Note(kamo): The mel matrix of librosa is different from kaldi.
+        melmat = librosa.filters.mel(**_mel_options)
+        # melmat: (D2, D1) -> (D1, D2)
+        self.register_buffer("melmat", torch.from_numpy(melmat.T).float())
+
+    def extra_repr(self):
+        """Extra repr."""
+        return ", ".join(f"{k}={v}" for k, v in self.mel_options.items())
+
+    def forward(
+        self, feat: torch.Tensor, ilens: torch.LongTensor
+    ) -> Tuple[torch.Tensor, torch.LongTensor]:
+        # feat: (B, T, D1) x melmat: (D1, D2) -> mel_feat: (B, T, D2)
+        """Forward pass for training.
+        
+            Args:
+                feat: TODO.
+                ilens: TODO.
+            """
+        mel_feat = torch.matmul(feat, self.melmat)
+
+        logmel_feat = (mel_feat + 1e-20).log()
+        # Zero padding
+        logmel_feat = logmel_feat.masked_fill(make_pad_mask(ilens, logmel_feat, 1), 0.0)
+        return logmel_feat, ilens
+
+
+class GlobalMVN(torch.nn.Module):
+    """Apply global mean and variance normalization
+
+    Args:
+        stats_file(str): npy file of 1-dim array or text file.
+            From the _first element to
+            the {(len(array) - 1) / 2}th element are treated as
+            the sum of features,
+            and the rest excluding the last elements are
+            treated as the sum of the square value of features,
+            and the last elements eqauls to the number of samples.
+        std_floor(float):
+    """
+
+    def __init__(
+        self,
+        stats_file: str,
+        norm_means: bool = True,
+        norm_vars: bool = True,
+        eps: float = 1.0e-20,
+    ):
+        """Initialize GlobalMVN.
+        
+            Args:
+                stats_file: TODO.
+                norm_means: TODO.
+                norm_vars: TODO.
+                eps: TODO.
+            """
+        super().__init__()
+        self.norm_means = norm_means
+        self.norm_vars = norm_vars
+
+        self.stats_file = stats_file
+        stats = np.load(stats_file)
+
+        stats = stats.astype(float)
+        assert (len(stats) - 1) % 2 == 0, stats.shape
+
+        count = stats.flatten()[-1]
+        mean = stats[: (len(stats) - 1) // 2] / count
+        var = stats[(len(stats) - 1) // 2 : -1] / count - mean * mean
+        std = np.maximum(np.sqrt(var), eps)
+
+        self.register_buffer("bias", torch.from_numpy(-mean.astype(np.float32)))
+        self.register_buffer("scale", torch.from_numpy(1 / std.astype(np.float32)))
+
+    def extra_repr(self):
+        """Extra repr."""
+        return (
+            f"stats_file={self.stats_file}, "
+            f"norm_means={self.norm_means}, norm_vars={self.norm_vars}"
+        )
+
+    def forward(
+        self, x: torch.Tensor, ilens: torch.LongTensor
+    ) -> Tuple[torch.Tensor, torch.LongTensor]:
+        # feat: (B, T, D)
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+                ilens: TODO.
+            """
+        if self.norm_means:
+            x += self.bias.type_as(x)
+            x.masked_fill(make_pad_mask(ilens, x, 1), 0.0)
+
+        if self.norm_vars:
+            x *= self.scale.type_as(x)
+        return x, ilens
+
+
+class UtteranceMVN(torch.nn.Module):
+    def __init__(self, norm_means: bool = True, norm_vars: bool = False, eps: float = 1.0e-20):
+        """Initialize UtteranceMVN.
+        
+            Args:
+                norm_means: TODO.
+                norm_vars: TODO.
+                eps: TODO.
+            """
+        super().__init__()
+        self.norm_means = norm_means
+        self.norm_vars = norm_vars
+        self.eps = eps
+
+    def extra_repr(self):
+        """Extra repr."""
+        return f"norm_means={self.norm_means}, norm_vars={self.norm_vars}"
+
+    def forward(
+        self, x: torch.Tensor, ilens: torch.LongTensor
+    ) -> Tuple[torch.Tensor, torch.LongTensor]:
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+                ilens: TODO.
+            """
+        return utterance_mvn(
+            x, ilens, norm_means=self.norm_means, norm_vars=self.norm_vars, eps=self.eps
+        )
+
+
+def utterance_mvn(
+    x: torch.Tensor,
+    ilens: torch.LongTensor,
+    norm_means: bool = True,
+    norm_vars: bool = False,
+    eps: float = 1.0e-20,
+) -> Tuple[torch.Tensor, torch.LongTensor]:
+    """Apply utterance mean and variance normalization
+
+    Args:
+        x: (B, T, D), assumed zero padded
+        ilens: (B, T, D)
+        norm_means:
+        norm_vars:
+        eps:
+
+    """
+    ilens_ = ilens.type_as(x)
+    # mean: (B, D)
+    mean = x.sum(dim=1) / ilens_[:, None]
+
+    if norm_means:
+        x -= mean[:, None, :]
+        x_ = x
+    else:
+        x_ = x - mean[:, None, :]
+
+    # Zero padding
+    x_.masked_fill(make_pad_mask(ilens, x_, 1), 0.0)
+    if norm_vars:
+        var = x_.pow(2).sum(dim=1) / ilens_[:, None]
+        var = torch.clamp(var, min=eps)
+        x /= var.sqrt()[:, None, :]
+        x_ = x
+    return x_, ilens
+
+
+def feature_transform_for(args, n_fft):
+    """Feature transform for.
+    
+        Args:
+            args: TODO.
+            n_fft: TODO.
+        """
+    return FeatureTransform(
+        # Mel options,
+        fs=args.fbank_fs,
+        n_fft=n_fft,
+        n_mels=args.n_mels,
+        fmin=args.fbank_fmin,
+        fmax=args.fbank_fmax,
+        # Normalization
+        stats_file=args.stats_file,
+        apply_uttmvn=args.apply_uttmvn,
+        uttmvn_norm_means=args.uttmvn_norm_means,
+        uttmvn_norm_vars=args.uttmvn_norm_vars,
+    )

funasr/frontends/utils/frontend.py

@@ -0,0 +1,186 @@
+from typing import List
+from typing import Optional
+from typing import Tuple
+from typing import Union
+
+import numpy
+import torch
+import torch.nn as nn
+from torch_complex.tensor import ComplexTensor
+
+from funasr.frontends.utils.dnn_beamformer import DNN_Beamformer
+from funasr.frontends.utils.dnn_wpe import DNN_WPE
+
+
+class Frontend(nn.Module):
+    def __init__(
+        self,
+        idim: int,
+        # WPE options
+        use_wpe: bool = False,
+        wtype: str = "blstmp",
+        wlayers: int = 3,
+        wunits: int = 300,
+        wprojs: int = 320,
+        wdropout_rate: float = 0.0,
+        taps: int = 5,
+        delay: int = 3,
+        use_dnn_mask_for_wpe: bool = True,
+        # Beamformer options
+        use_beamformer: bool = False,
+        btype: str = "blstmp",
+        blayers: int = 3,
+        bunits: int = 300,
+        bprojs: int = 320,
+        bnmask: int = 2,
+        badim: int = 320,
+        ref_channel: int = -1,
+        bdropout_rate=0.0,
+    ):
+        """Initialize Frontend.
+        
+            Args:
+                idim: TODO.
+                use_wpe: TODO.
+                wtype: TODO.
+                wlayers: TODO.
+                wunits: TODO.
+                wprojs: TODO.
+                wdropout_rate: TODO.
+                taps: TODO.
+                delay: TODO.
+                use_dnn_mask_for_wpe: TODO.
+                use_beamformer: TODO.
+                btype: TODO.
+                blayers: TODO.
+                bunits: TODO.
+                bprojs: TODO.
+                bnmask: TODO.
+                badim: TODO.
+                ref_channel: TODO.
+                bdropout_rate: TODO.
+            """
+        super().__init__()
+
+        self.use_beamformer = use_beamformer
+        self.use_wpe = use_wpe
+        self.use_dnn_mask_for_wpe = use_dnn_mask_for_wpe
+        # use frontend for all the data,
+        # e.g. in the case of multi-speaker speech separation
+        self.use_frontend_for_all = bnmask > 2
+
+        if self.use_wpe:
+            if self.use_dnn_mask_for_wpe:
+                # Use DNN for power estimation
+                # (Not observed significant gains)
+                iterations = 1
+            else:
+                # Performing as conventional WPE, without DNN Estimator
+                iterations = 2
+
+            self.wpe = DNN_WPE(
+                wtype=wtype,
+                widim=idim,
+                wunits=wunits,
+                wprojs=wprojs,
+                wlayers=wlayers,
+                taps=taps,
+                delay=delay,
+                dropout_rate=wdropout_rate,
+                iterations=iterations,
+                use_dnn_mask=use_dnn_mask_for_wpe,
+            )
+        else:
+            self.wpe = None
+
+        if self.use_beamformer:
+            self.beamformer = DNN_Beamformer(
+                btype=btype,
+                bidim=idim,
+                bunits=bunits,
+                bprojs=bprojs,
+                blayers=blayers,
+                bnmask=bnmask,
+                dropout_rate=bdropout_rate,
+                badim=badim,
+                ref_channel=ref_channel,
+            )
+        else:
+            self.beamformer = None
+
+    def forward(
+        self, x: ComplexTensor, ilens: Union[torch.LongTensor, numpy.ndarray, List[int]]
+    ) -> Tuple[ComplexTensor, torch.LongTensor, Optional[ComplexTensor]]:
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+                ilens: TODO.
+            """
+        assert len(x) == len(ilens), (len(x), len(ilens))
+        # (B, T, F) or (B, T, C, F)
+        if x.dim() not in (3, 4):
+            raise ValueError(f"Input dim must be 3 or 4: {x.dim()}")
+        if not torch.is_tensor(ilens):
+            ilens = torch.from_numpy(numpy.asarray(ilens)).to(x.device)
+
+        mask = None
+        h = x
+        if h.dim() == 4:
+            if self.training:
+                choices = [(False, False)] if not self.use_frontend_for_all else []
+                if self.use_wpe:
+                    choices.append((True, False))
+
+                if self.use_beamformer:
+                    choices.append((False, True))
+
+                use_wpe, use_beamformer = choices[numpy.random.randint(len(choices))]
+
+            else:
+                use_wpe = self.use_wpe
+                use_beamformer = self.use_beamformer
+
+            # 1. WPE
+            if use_wpe:
+                # h: (B, T, C, F) -> h: (B, T, C, F)
+                h, ilens, mask = self.wpe(h, ilens)
+
+            # 2. Beamformer
+            if use_beamformer:
+                # h: (B, T, C, F) -> h: (B, T, F)
+                h, ilens, mask = self.beamformer(h, ilens)
+
+        return h, ilens, mask
+
+
+def frontend_for(args, idim):
+    """Frontend for.
+    
+        Args:
+            args: TODO.
+            idim: TODO.
+        """
+    return Frontend(
+        idim=idim,
+        # WPE options
+        use_wpe=args.use_wpe,
+        wtype=args.wtype,
+        wlayers=args.wlayers,
+        wunits=args.wunits,
+        wprojs=args.wprojs,
+        wdropout_rate=args.wdropout_rate,
+        taps=args.wpe_taps,
+        delay=args.wpe_delay,
+        use_dnn_mask_for_wpe=args.use_dnn_mask_for_wpe,
+        # Beamformer options
+        use_beamformer=args.use_beamformer,
+        btype=args.btype,
+        blayers=args.blayers,
+        bunits=args.bunits,
+        bprojs=args.bprojs,
+        bnmask=args.bnmask,
+        badim=args.badim,
+        ref_channel=args.ref_channel,
+        bdropout_rate=args.bdropout_rate,
+    )

funasr/frontends/utils/log_mel.py

@@ -0,0 +1,97 @@
+import librosa
+import torch
+from typing import Tuple
+
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+
+
+class LogMel(torch.nn.Module):
+    """Convert STFT to fbank feats
+
+    The arguments is same as librosa.filters.mel
+
+    Args:
+        fs: number > 0 [scalar] sampling rate of the incoming signal
+        n_fft: int > 0 [scalar] number of FFT components
+        n_mels: int > 0 [scalar] number of Mel bands to generate
+        fmin: float >= 0 [scalar] lowest frequency (in Hz)
+        fmax: float >= 0 [scalar] highest frequency (in Hz).
+            If `None`, use `fmax = fs / 2.0`
+        htk: use HTK formula instead of Slaney
+    """
+
+    def __init__(
+        self,
+        fs: int = 16000,
+        n_fft: int = 512,
+        n_mels: int = 80,
+        fmin: float = None,
+        fmax: float = None,
+        htk: bool = False,
+        log_base: float = None,
+    ):
+        """Initialize LogMel.
+        
+            Args:
+                fs: TODO.
+                n_fft: TODO.
+                n_mels: TODO.
+                fmin: TODO.
+                fmax: TODO.
+                htk: TODO.
+                log_base: TODO.
+            """
+        super().__init__()
+
+        fmin = 0 if fmin is None else fmin
+        fmax = fs / 2 if fmax is None else fmax
+        _mel_options = dict(
+            sr=fs,
+            n_fft=n_fft,
+            n_mels=n_mels,
+            fmin=fmin,
+            fmax=fmax,
+            htk=htk,
+        )
+        self.mel_options = _mel_options
+        self.log_base = log_base
+
+        # Note(kamo): The mel matrix of librosa is different from kaldi.
+        melmat = librosa.filters.mel(**_mel_options)
+        # melmat: (D2, D1) -> (D1, D2)
+        self.register_buffer("melmat", torch.from_numpy(melmat.T).float())
+
+    def extra_repr(self):
+        """Extra repr."""
+        return ", ".join(f"{k}={v}" for k, v in self.mel_options.items())
+
+    def forward(
+        self,
+        feat: torch.Tensor,
+        ilens: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # feat: (B, T, D1) x melmat: (D1, D2) -> mel_feat: (B, T, D2)
+        """Forward pass for training.
+        
+            Args:
+                feat: TODO.
+                ilens: TODO.
+            """
+        mel_feat = torch.matmul(feat, self.melmat)
+        mel_feat = torch.clamp(mel_feat, min=1e-10)
+
+        if self.log_base is None:
+            logmel_feat = mel_feat.log()
+        elif self.log_base == 2.0:
+            logmel_feat = mel_feat.log2()
+        elif self.log_base == 10.0:
+            logmel_feat = mel_feat.log10()
+        else:
+            logmel_feat = mel_feat.log() / torch.log(self.log_base)
+
+        # Zero padding
+        if ilens is not None:
+            logmel_feat = logmel_feat.masked_fill(make_pad_mask(ilens, logmel_feat, 1), 0.0)
+        else:
+            ilens = feat.new_full([feat.size(0)], fill_value=feat.size(1), dtype=torch.long)
+        return logmel_feat, ilens

funasr/frontends/utils/mask_estimator.py

@@ -0,0 +1,86 @@
+from typing import Tuple
+
+import numpy as np
+import torch
+from torch.nn import functional as F
+from torch_complex.tensor import ComplexTensor
+
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+from funasr.models.language_model.rnn.encoders import RNN
+from funasr.models.language_model.rnn.encoders import RNNP
+
+
+class MaskEstimator(torch.nn.Module):
+    def __init__(self, type, idim, layers, units, projs, dropout, nmask=1):
+        """Initialize MaskEstimator.
+        
+            Args:
+                type: TODO.
+                idim: TODO.
+                layers: TODO.
+                units: TODO.
+                projs: TODO.
+                dropout: TODO.
+                nmask: TODO.
+            """
+        super().__init__()
+        subsample = np.ones(layers + 1, dtype=np.int32)
+
+        typ = type.lstrip("vgg").rstrip("p")
+        if type[-1] == "p":
+            self.brnn = RNNP(idim, layers, units, projs, subsample, dropout, typ=typ)
+        else:
+            self.brnn = RNN(idim, layers, units, projs, dropout, typ=typ)
+
+        self.type = type
+        self.nmask = nmask
+        self.linears = torch.nn.ModuleList([torch.nn.Linear(projs, idim) for _ in range(nmask)])
+
+    def forward(
+        self, xs: ComplexTensor, ilens: torch.LongTensor
+    ) -> Tuple[Tuple[torch.Tensor, ...], torch.LongTensor]:
+        """The forward function
+
+        Args:
+            xs: (B, F, C, T)
+            ilens: (B,)
+        Returns:
+            hs (torch.Tensor): The hidden vector (B, F, C, T)
+            masks: A tuple of the masks. (B, F, C, T)
+            ilens: (B,)
+        """
+        assert xs.size(0) == ilens.size(0), (xs.size(0), ilens.size(0))
+        _, _, C, input_length = xs.size()
+        # (B, F, C, T) -> (B, C, T, F)
+        xs = xs.permute(0, 2, 3, 1)
+
+        # Calculate amplitude: (B, C, T, F) -> (B, C, T, F)
+        xs = (xs.real**2 + xs.imag**2) ** 0.5
+        # xs: (B, C, T, F) -> xs: (B * C, T, F)
+        xs = xs.contiguous().view(-1, xs.size(-2), xs.size(-1))
+        # ilens: (B,) -> ilens_: (B * C)
+        ilens_ = ilens[:, None].expand(-1, C).contiguous().view(-1)
+
+        # xs: (B * C, T, F) -> xs: (B * C, T, D)
+        xs, _, _ = self.brnn(xs, ilens_)
+        # xs: (B * C, T, D) -> xs: (B, C, T, D)
+        xs = xs.view(-1, C, xs.size(-2), xs.size(-1))
+
+        masks = []
+        for linear in self.linears:
+            # xs: (B, C, T, D) -> mask:(B, C, T, F)
+            mask = linear(xs)
+
+            mask = torch.sigmoid(mask)
+            # Zero padding
+            mask.masked_fill(make_pad_mask(ilens, mask, length_dim=2), 0)
+
+            # (B, C, T, F) -> (B, F, C, T)
+            mask = mask.permute(0, 3, 1, 2)
+
+            # Take cares of multi gpu cases: If input_length > max(ilens)
+            if mask.size(-1) < input_length:
+                mask = F.pad(mask, [0, input_length - mask.size(-1)], value=0)
+            masks.append(mask)
+
+        return tuple(masks), ilens

funasr/frontends/utils/stft.py

@@ -0,0 +1,238 @@
+from distutils.version import LooseVersion
+from typing import Optional
+from typing import Tuple
+from typing import Union
+
+import torch
+
+try:
+    from torch_complex.tensor import ComplexTensor
+except:
+    print("Please install torch_complex firstly")
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+from funasr.frontends.utils.complex_utils import is_complex
+
+import librosa
+import numpy as np
+
+is_torch_1_9_plus = LooseVersion(torch.__version__) >= LooseVersion("1.9.0")
+
+
+is_torch_1_7_plus = LooseVersion(torch.__version__) >= LooseVersion("1.7")
+
+
+class Stft(torch.nn.Module):
+    def __init__(
+        self,
+        n_fft: int = 512,
+        win_length: int = None,
+        hop_length: int = 128,
+        window: Optional[str] = "hann",
+        center: bool = True,
+        normalized: bool = False,
+        onesided: bool = True,
+    ):
+        """Initialize Stft.
+        
+            Args:
+                n_fft: TODO.
+                win_length: TODO.
+                hop_length: TODO.
+                window: TODO.
+                center: TODO.
+                normalized: TODO.
+                onesided: TODO.
+            """
+        super().__init__()
+        self.n_fft = n_fft
+        if win_length is None:
+            self.win_length = n_fft
+        else:
+            self.win_length = win_length
+        self.hop_length = hop_length
+        self.center = center
+        self.normalized = normalized
+        self.onesided = onesided
+        if window is not None and not hasattr(torch, f"{window}_window"):
+            if window.lower() != "povey":
+                raise ValueError(f"{window} window is not implemented")
+        self.window = window
+
+    def extra_repr(self):
+        """Extra repr."""
+        return (
+            f"n_fft={self.n_fft}, "
+            f"win_length={self.win_length}, "
+            f"hop_length={self.hop_length}, "
+            f"center={self.center}, "
+            f"normalized={self.normalized}, "
+            f"onesided={self.onesided}"
+        )
+
+    def forward(
+        self, input: torch.Tensor, ilens: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """STFT forward function.
+
+        Args:
+            input: (Batch, Nsamples) or (Batch, Nsample, Channels)
+            ilens: (Batch)
+        Returns:
+            output: (Batch, Frames, Freq, 2) or (Batch, Frames, Channels, Freq, 2)
+
+        """
+        bs = input.size(0)
+        if input.dim() == 3:
+            multi_channel = True
+            # input: (Batch, Nsample, Channels) -> (Batch * Channels, Nsample)
+            input = input.transpose(1, 2).reshape(-1, input.size(1))
+        else:
+            multi_channel = False
+
+        # NOTE(kamo):
+        #   The default behaviour of torch.stft is compatible with librosa.stft
+        #   about padding and scaling.
+        #   Note that it's different from scipy.signal.stft
+
+        # output: (Batch, Freq, Frames, 2=real_imag)
+        # or (Batch, Channel, Freq, Frames, 2=real_imag)
+        if self.window is not None:
+            if self.window.lower() == "povey":
+                window = torch.hann_window(
+                    self.win_length, periodic=False, device=input.device, dtype=input.dtype
+                ).pow(0.85)
+            else:
+                window_func = getattr(torch, f"{self.window}_window")
+                window = window_func(self.win_length, dtype=input.dtype, device=input.device)
+        else:
+            window = None
+
+        # For the compatibility of ARM devices, which do not support
+        # torch.stft() due to the lake of MKL.
+        if input.is_cuda or torch.backends.mkl.is_available():
+            stft_kwargs = dict(
+                n_fft=self.n_fft,
+                win_length=self.win_length,
+                hop_length=self.hop_length,
+                center=self.center,
+                window=window,
+                normalized=self.normalized,
+                onesided=self.onesided,
+            )
+            if is_torch_1_7_plus:
+                stft_kwargs["return_complex"] = False
+            output = torch.stft(input, **stft_kwargs)
+        else:
+            if self.training:
+                raise NotImplementedError(
+                    "stft is implemented with librosa on this device, which does not "
+                    "support the training mode."
+                )
+
+            # use stft_kwargs to flexibly control different PyTorch versions' kwargs
+            stft_kwargs = dict(
+                n_fft=self.n_fft,
+                win_length=self.win_length,
+                hop_length=self.hop_length,
+                center=self.center,
+                window=window,
+            )
+
+            if window is not None:
+                # pad the given window to n_fft
+                n_pad_left = (self.n_fft - window.shape[0]) // 2
+                n_pad_right = self.n_fft - window.shape[0] - n_pad_left
+                stft_kwargs["window"] = torch.cat(
+                    [torch.zeros(n_pad_left), window, torch.zeros(n_pad_right)], 0
+                ).numpy()
+            else:
+                win_length = self.win_length if self.win_length is not None else self.n_fft
+                stft_kwargs["window"] = torch.ones(win_length)
+
+            output = []
+            # iterate over istances in a batch
+            for i, instance in enumerate(input):
+                stft = librosa.stft(input[i].numpy(), **stft_kwargs)
+                output.append(torch.tensor(np.stack([stft.real, stft.imag], -1)))
+            output = torch.stack(output, 0)
+            if not self.onesided:
+                len_conj = self.n_fft - output.shape[1]
+                conj = output[:, 1 : 1 + len_conj].flip(1)
+                conj[:, :, :, -1].data *= -1
+                output = torch.cat([output, conj], 1)
+            if self.normalized:
+                output = output * (stft_kwargs["window"].shape[0] ** (-0.5))
+
+        # output: (Batch, Freq, Frames, 2=real_imag)
+        # -> (Batch, Frames, Freq, 2=real_imag)
+        output = output.transpose(1, 2)
+        if multi_channel:
+            # output: (Batch * Channel, Frames, Freq, 2=real_imag)
+            # -> (Batch, Frame, Channel, Freq, 2=real_imag)
+            output = output.view(bs, -1, output.size(1), output.size(2), 2).transpose(1, 2)
+
+        if ilens is not None:
+            if self.center:
+                pad = self.n_fft // 2
+                ilens = ilens + 2 * pad
+
+            olens = (ilens - self.n_fft) // self.hop_length + 1
+            output.masked_fill_(make_pad_mask(olens, output, 1), 0.0)
+        else:
+            olens = None
+
+        return output, olens
+
+    def inverse(
+        self, input: Union[torch.Tensor, ComplexTensor], ilens: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Inverse STFT.
+
+        Args:
+            input: Tensor(batch, T, F, 2) or ComplexTensor(batch, T, F)
+            ilens: (batch,)
+        Returns:
+            wavs: (batch, samples)
+            ilens: (batch,)
+        """
+        if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
+            istft = torch.functional.istft
+        else:
+            try:
+                import torchaudio
+            except ImportError:
+                raise ImportError("Please install torchaudio>=0.3.0 or use torch>=1.6.0")
+
+            if not hasattr(torchaudio.functional, "istft"):
+                raise ImportError("Please install torchaudio>=0.3.0 or use torch>=1.6.0")
+            istft = torchaudio.functional.istft
+
+        if self.window is not None:
+            window_func = getattr(torch, f"{self.window}_window")
+            if is_complex(input):
+                datatype = input.real.dtype
+            else:
+                datatype = input.dtype
+            window = window_func(self.win_length, dtype=datatype, device=input.device)
+        else:
+            window = None
+
+        if is_complex(input):
+            input = torch.stack([input.real, input.imag], dim=-1)
+        elif input.shape[-1] != 2:
+            raise TypeError("Invalid input type")
+        input = input.transpose(1, 2)
+
+        wavs = istft(
+            input,
+            n_fft=self.n_fft,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+            window=window,
+            center=self.center,
+            normalized=self.normalized,
+            onesided=self.onesided,
+            length=ilens.max() if ilens is not None else ilens,
+        )
+
+        return wavs, ilens

funasr/frontends/wav_frontend.py

@@ -0,0 +1,672 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+# Part of the implementation is borrowed from espnet/espnet.
+from typing import Tuple
+import copy
+import numpy as np
+import torch
+import torch.nn as nn
+import torchaudio.compliance.kaldi as kaldi
+from torch.nn.utils.rnn import pad_sequence
+
+import funasr.frontends.eend_ola_feature as eend_ola_feature
+from funasr.register import tables
+
+
+def load_cmvn(cmvn_file):
+    """Load cmvn.
+    
+        Args:
+            cmvn_file: TODO.
+        """
+    with open(cmvn_file, "r", encoding="utf-8") as f:
+        lines = f.readlines()
+    means_list = []
+    vars_list = []
+    for i in range(len(lines)):
+        line_item = lines[i].split()
+        if line_item[0] == "<AddShift>":
+            line_item = lines[i + 1].split()
+            if line_item[0] == "<LearnRateCoef>":
+                add_shift_line = line_item[3 : (len(line_item) - 1)]
+                means_list = list(add_shift_line)
+                continue
+        elif line_item[0] == "<Rescale>":
+            line_item = lines[i + 1].split()
+            if line_item[0] == "<LearnRateCoef>":
+                rescale_line = line_item[3 : (len(line_item) - 1)]
+                vars_list = list(rescale_line)
+                continue
+    means = np.array(means_list).astype(np.float32)
+    vars = np.array(vars_list).astype(np.float32)
+    cmvn = np.array([means, vars])
+    cmvn = torch.as_tensor(cmvn, dtype=torch.float32)
+    return cmvn
+
+
+def apply_cmvn(inputs, cmvn):  # noqa
+    """
+    Apply CMVN with mvn data
+    """
+
+    device = inputs.device
+    dtype = inputs.dtype
+    frame, dim = inputs.shape
+
+    means = cmvn[0:1, :dim]
+    vars = cmvn[1:2, :dim]
+    inputs += means.to(device)
+    inputs *= vars.to(device)
+
+    return inputs.type(torch.float32)
+
+
+def apply_lfr(inputs, lfr_m, lfr_n):
+    """Apply lfr.
+    
+        Args:
+            inputs: TODO.
+            lfr_m: TODO.
+            lfr_n: TODO.
+        """
+    LFR_inputs = []
+    T = inputs.shape[0]
+    T_lfr = int(np.ceil(T / lfr_n))
+    left_padding = inputs[0].repeat((lfr_m - 1) // 2, 1)
+    inputs = torch.vstack((left_padding, inputs))
+    T = T + (lfr_m - 1) // 2
+    feat_dim = inputs.shape[-1]
+    strides = (lfr_n * feat_dim, 1)
+    sizes = (T_lfr, lfr_m * feat_dim)
+    last_idx = (T - lfr_m) // lfr_n + 1
+    num_padding = lfr_m - (T - last_idx * lfr_n)
+    if num_padding > 0:
+        num_padding = (2 * lfr_m - 2 * T + (T_lfr - 1 + last_idx) * lfr_n) / 2 * (T_lfr - last_idx)
+        inputs = torch.vstack([inputs] + [inputs[-1:]] * int(num_padding))
+    LFR_outputs = inputs.as_strided(sizes, strides)
+    return LFR_outputs.clone().type(torch.float32)
+
+
+@tables.register("frontend_classes", "wav_frontend")
+@tables.register("frontend_classes", "WavFrontend")
+class WavFrontend(nn.Module):
+    """Conventional frontend structure for ASR."""
+
+    def __init__(
+        self,
+        cmvn_file: str = None,
+        fs: int = 16000,
+        window: str = "hamming",
+        n_mels: int = 80,
+        frame_length: int = 25,
+        frame_shift: int = 10,
+        filter_length_min: int = -1,
+        filter_length_max: int = -1,
+        lfr_m: int = 1,
+        lfr_n: int = 1,
+        dither: float = 1.0,
+        snip_edges: bool = True,
+        upsacle_samples: bool = True,
+        **kwargs,
+    ):
+        """Initialize WavFrontend.
+        
+            Args:
+                cmvn_file: TODO.
+                fs: TODO.
+                window: TODO.
+                n_mels: TODO.
+                frame_length: TODO.
+                frame_shift: TODO.
+                filter_length_min: TODO.
+                filter_length_max: TODO.
+                lfr_m: TODO.
+                lfr_n: TODO.
+                dither: TODO.
+                snip_edges: TODO.
+                upsacle_samples: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        self.fs = fs
+        self.window = window
+        self.n_mels = n_mels
+        self.frame_length = frame_length
+        self.frame_shift = frame_shift
+        self.filter_length_min = filter_length_min
+        self.filter_length_max = filter_length_max
+        self.lfr_m = lfr_m
+        self.lfr_n = lfr_n
+        self.cmvn_file = cmvn_file
+        self.dither = dither
+        self.snip_edges = snip_edges
+        self.upsacle_samples = upsacle_samples
+        self.cmvn = None if self.cmvn_file is None else load_cmvn(self.cmvn_file)
+
+    def output_size(self) -> int:
+        """Output size."""
+        return self.n_mels * self.lfr_m
+
+    def forward(
+        self,
+        input: torch.Tensor,
+        input_lengths,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass for training.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+                **kwargs: Additional keyword arguments.
+            """
+        batch_size = input.size(0)
+        feats = []
+        feats_lens = []
+        for i in range(batch_size):
+            waveform_length = input_lengths[i]
+            waveform = input[i][:waveform_length]
+            if self.upsacle_samples:
+                waveform = waveform * (1 << 15)
+            waveform = waveform.unsqueeze(0)
+            mat = kaldi.fbank(
+                waveform,
+                num_mel_bins=self.n_mels,
+                frame_length=min(self.frame_length,waveform_length/self.fs*1000),
+                frame_shift=self.frame_shift,
+                dither=self.dither,
+                energy_floor=0.0,
+                window_type=self.window,
+                sample_frequency=self.fs,
+                snip_edges=self.snip_edges,
+            )
+
+            if self.lfr_m != 1 or self.lfr_n != 1:
+                mat = apply_lfr(mat, self.lfr_m, self.lfr_n)
+            if self.cmvn is not None:
+                mat = apply_cmvn(mat, self.cmvn)
+            feat_length = mat.size(0)
+            feats.append(mat)
+            feats_lens.append(feat_length)
+
+        feats_lens = torch.as_tensor(feats_lens)
+        if batch_size == 1:
+            feats_pad = feats[0][None, :, :]
+        else:
+            feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0)
+        return feats_pad, feats_lens
+
+    def forward_fbank(
+        self, input: torch.Tensor, input_lengths: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward fbank.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+            """
+        batch_size = input.size(0)
+        feats = []
+        feats_lens = []
+        for i in range(batch_size):
+            waveform_length = input_lengths[i]
+            waveform = input[i][:waveform_length]
+            waveform = waveform * (1 << 15)
+            waveform = waveform.unsqueeze(0)
+            mat = kaldi.fbank(
+                waveform,
+                num_mel_bins=self.n_mels,
+                frame_length=self.frame_length,
+                frame_shift=self.frame_shift,
+                dither=self.dither,
+                energy_floor=0.0,
+                window_type=self.window,
+                sample_frequency=self.fs,
+            )
+
+            feat_length = mat.size(0)
+            feats.append(mat)
+            feats_lens.append(feat_length)
+
+        feats_lens = torch.as_tensor(feats_lens)
+        feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0)
+        return feats_pad, feats_lens
+
+    def forward_lfr_cmvn(
+        self, input: torch.Tensor, input_lengths: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward lfr cmvn.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+            """
+        batch_size = input.size(0)
+        feats = []
+        feats_lens = []
+        for i in range(batch_size):
+            mat = input[i, : input_lengths[i], :]
+            if self.lfr_m != 1 or self.lfr_n != 1:
+                mat = apply_lfr(mat, self.lfr_m, self.lfr_n)
+            if self.cmvn is not None:
+                mat = apply_cmvn(mat, self.cmvn)
+            feat_length = mat.size(0)
+            feats.append(mat)
+            feats_lens.append(feat_length)
+
+        feats_lens = torch.as_tensor(feats_lens)
+        feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0)
+        return feats_pad, feats_lens
+
+
+@tables.register("frontend_classes", "WavFrontendOnline")
+class WavFrontendOnline(nn.Module):
+    """Conventional frontend structure for streaming ASR/VAD."""
+
+    def __init__(
+        self,
+        cmvn_file: str = None,
+        fs: int = 16000,
+        window: str = "hamming",
+        n_mels: int = 80,
+        frame_length: int = 25,
+        frame_shift: int = 10,
+        filter_length_min: int = -1,
+        filter_length_max: int = -1,
+        lfr_m: int = 1,
+        lfr_n: int = 1,
+        dither: float = 1.0,
+        snip_edges: bool = True,
+        upsacle_samples: bool = True,
+        **kwargs,
+    ):
+        """Initialize WavFrontendOnline.
+        
+            Args:
+                cmvn_file: TODO.
+                fs: TODO.
+                window: TODO.
+                n_mels: TODO.
+                frame_length: TODO.
+                frame_shift: TODO.
+                filter_length_min: TODO.
+                filter_length_max: TODO.
+                lfr_m: TODO.
+                lfr_n: TODO.
+                dither: TODO.
+                snip_edges: TODO.
+                upsacle_samples: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        self.fs = fs
+        self.window = window
+        self.n_mels = n_mels
+        self.frame_length = frame_length
+        self.frame_shift = frame_shift
+        self.frame_sample_length = int(self.frame_length * self.fs / 1000)
+        self.frame_shift_sample_length = int(self.frame_shift * self.fs / 1000)
+        self.filter_length_min = filter_length_min
+        self.filter_length_max = filter_length_max
+        self.lfr_m = lfr_m
+        self.lfr_n = lfr_n
+        self.cmvn_file = cmvn_file
+        self.dither = dither
+        self.snip_edges = snip_edges
+        self.upsacle_samples = upsacle_samples
+        # self.waveforms = None
+        # self.reserve_waveforms = None
+        # self.fbanks = None
+        # self.fbanks_lens = None
+        self.cmvn = None if self.cmvn_file is None else load_cmvn(self.cmvn_file)
+        # self.input_cache = None
+        # self.lfr_splice_cache = []
+
+    def output_size(self) -> int:
+        """Output size."""
+        return self.n_mels * self.lfr_m
+
+    @staticmethod
+    def apply_cmvn(inputs: torch.Tensor, cmvn: torch.Tensor) -> torch.Tensor:
+        """
+        Apply CMVN with mvn data
+        """
+
+        device = inputs.device
+        dtype = inputs.dtype
+        frame, dim = inputs.shape
+
+        means = np.tile(cmvn[0:1, :dim], (frame, 1))
+        vars = np.tile(cmvn[1:2, :dim], (frame, 1))
+        inputs += torch.from_numpy(means).type(dtype).to(device)
+        inputs *= torch.from_numpy(vars).type(dtype).to(device)
+
+        return inputs.type(torch.float32)
+
+    @staticmethod
+    def apply_lfr(
+        inputs: torch.Tensor, lfr_m: int, lfr_n: int, is_final: bool = False
+    ) -> Tuple[torch.Tensor, torch.Tensor, int]:
+        """
+        Apply lfr with data
+        """
+
+        LFR_inputs = []
+        # inputs = torch.vstack((inputs_lfr_cache, inputs))
+        T = inputs.shape[0]  # include the right context
+        T_lfr = int(
+            np.ceil((T - (lfr_m - 1) // 2) / lfr_n)
+        )  # minus the right context: (lfr_m - 1) // 2
+        splice_idx = T_lfr
+        feat_dim = inputs.shape[-1]
+        ori_inputs = inputs
+        strides = (lfr_n * feat_dim, 1)
+        sizes = (T_lfr, lfr_m * feat_dim)
+        last_idx = (T - lfr_m) // lfr_n + 1
+        num_padding = lfr_m - (T - last_idx * lfr_n)
+        if is_final:
+            if num_padding > 0:
+                num_padding = (2 * lfr_m - 2 * T + (T_lfr - 1 + last_idx) * lfr_n) / 2 * (T_lfr - last_idx)
+                inputs = torch.vstack([inputs] + [inputs[-1:]] * int(num_padding))
+        else:
+            if num_padding > 0:
+                sizes = (last_idx, lfr_m * feat_dim)
+                splice_idx = last_idx
+        splice_idx = min(T - 1, splice_idx * lfr_n)
+        LFR_outputs = inputs[:splice_idx].as_strided(sizes, strides)
+        lfr_splice_cache = ori_inputs[splice_idx:, :]
+        return LFR_outputs.clone().type(torch.float32), lfr_splice_cache, splice_idx
+
+    @staticmethod
+    def compute_frame_num(
+        sample_length: int, frame_sample_length: int, frame_shift_sample_length: int
+    ) -> int:
+        """Compute frame num.
+        
+            Args:
+                sample_length: TODO.
+                frame_sample_length: TODO.
+                frame_shift_sample_length: TODO.
+            """
+        frame_num = int((sample_length - frame_sample_length) / frame_shift_sample_length + 1)
+        return frame_num if frame_num >= 1 and sample_length >= frame_sample_length else 0
+
+    def forward_fbank(
+        self,
+        input: torch.Tensor,
+        input_lengths: torch.Tensor,
+        cache: dict = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Forward fbank.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+                cache: State cache dict for streaming inference.
+                **kwargs: Additional keyword arguments.
+            """
+        if cache is None:
+            cache = {}
+        batch_size = input.size(0)
+
+        input = torch.cat((cache["input_cache"], input), dim=1)
+        frame_num = self.compute_frame_num(
+            input.shape[-1], self.frame_sample_length, self.frame_shift_sample_length
+        )
+        # update self.in_cache
+        cache["input_cache"] = input[
+            :, -(input.shape[-1] - frame_num * self.frame_shift_sample_length) :
+        ]
+        waveforms = torch.empty(0)
+        feats_pad = torch.empty(0)
+        feats_lens = torch.empty(0)
+        if frame_num:
+            waveforms = []
+            feats = []
+            feats_lens = []
+            for i in range(batch_size):
+                waveform = input[i]
+                # we need accurate wave samples that used for fbank extracting
+                waveforms.append(
+                    waveform[
+                        : (
+                            (frame_num - 1) * self.frame_shift_sample_length
+                            + self.frame_sample_length
+                        )
+                    ]
+                )
+                waveform = waveform * (1 << 15)
+                waveform = waveform.unsqueeze(0)
+                mat = kaldi.fbank(
+                    waveform,
+                    num_mel_bins=self.n_mels,
+                    frame_length=self.frame_length,
+                    frame_shift=self.frame_shift,
+                    dither=self.dither,
+                    energy_floor=0.0,
+                    window_type=self.window,
+                    sample_frequency=self.fs,
+                )
+
+                feat_length = mat.size(0)
+                feats.append(mat)
+                feats_lens.append(feat_length)
+
+            waveforms = torch.stack(waveforms)
+            feats_lens = torch.as_tensor(feats_lens)
+            feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0)
+        cache["fbanks"] = feats_pad
+        cache["fbanks_lens"] = copy.deepcopy(feats_lens)
+        return waveforms, feats_pad, feats_lens
+
+    def forward_lfr_cmvn(
+        self,
+        input: torch.Tensor,
+        input_lengths: torch.Tensor,
+        is_final: bool = False,
+        cache: dict = None,
+        **kwargs,
+    ):
+        """Forward lfr cmvn.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+                is_final: Whether this is the final chunk in streaming.
+                cache: State cache dict for streaming inference.
+                **kwargs: Additional keyword arguments.
+            """
+        if cache is None:
+            cache = {}
+        batch_size = input.size(0)
+        feats = []
+        feats_lens = []
+        lfr_splice_frame_idxs = []
+        for i in range(batch_size):
+            mat = input[i, : input_lengths[i], :]
+            if self.lfr_m != 1 or self.lfr_n != 1:
+                # update self.lfr_splice_cache in self.apply_lfr
+                # mat, self.lfr_splice_cache[i], lfr_splice_frame_idx = self.apply_lfr(mat, self.lfr_m, self.lfr_n, self.lfr_splice_cache[i],
+                mat, cache["lfr_splice_cache"][i], lfr_splice_frame_idx = self.apply_lfr(
+                    mat, self.lfr_m, self.lfr_n, is_final
+                )
+            if self.cmvn_file is not None:
+                mat = self.apply_cmvn(mat, self.cmvn)
+            feat_length = mat.size(0)
+            feats.append(mat)
+            feats_lens.append(feat_length)
+            lfr_splice_frame_idxs.append(lfr_splice_frame_idx)
+
+        feats_lens = torch.as_tensor(feats_lens)
+        feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0)
+        lfr_splice_frame_idxs = torch.as_tensor(lfr_splice_frame_idxs)
+        return feats_pad, feats_lens, lfr_splice_frame_idxs
+
+    def forward(self, input: torch.Tensor, input_lengths: torch.Tensor, **kwargs):
+        """Forward pass for training.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+                **kwargs: Additional keyword arguments.
+            """
+        is_final = kwargs.get("is_final", False)
+        cache = kwargs.get("cache", {})
+        if len(cache) == 0:
+            self.init_cache(cache)
+
+        batch_size = input.shape[0]
+        assert (
+            batch_size == 1
+        ), "we support to extract feature online only when the batch size is equal to 1 now"
+
+        waveforms, feats, feats_lengths = self.forward_fbank(
+            input, input_lengths, cache=cache
+        )  # input shape: B T D
+
+        if feats.shape[0]:
+
+            cache["waveforms"] = torch.cat((cache["reserve_waveforms"], waveforms), dim=1)
+
+            if not cache["lfr_splice_cache"]:  # 初始化splice_cache
+                for i in range(batch_size):
+                    cache["lfr_splice_cache"].append(
+                        feats[i][0, :].unsqueeze(dim=0).repeat((self.lfr_m - 1) // 2, 1)
+                    )
+            # need the number of the input frames + self.lfr_splice_cache[0].shape[0] is greater than self.lfr_m
+            if feats_lengths[0] + cache["lfr_splice_cache"][0].shape[0] >= self.lfr_m:
+                lfr_splice_cache_tensor = torch.stack(cache["lfr_splice_cache"])  # B T D
+                feats = torch.cat((lfr_splice_cache_tensor, feats), dim=1)
+                feats_lengths += lfr_splice_cache_tensor[0].shape[0]
+                frame_from_waveforms = int(
+                    (cache["waveforms"].shape[1] - self.frame_sample_length)
+                    / self.frame_shift_sample_length
+                    + 1
+                )
+                minus_frame = (
+                    (self.lfr_m - 1) // 2 if cache["reserve_waveforms"].numel() == 0 else 0
+                )
+                feats, feats_lengths, lfr_splice_frame_idxs = self.forward_lfr_cmvn(
+                    feats, feats_lengths, is_final, cache=cache
+                )
+                if self.lfr_m == 1:
+                    cache["reserve_waveforms"] = torch.empty(0)
+                else:
+                    reserve_frame_idx = lfr_splice_frame_idxs[0] - minus_frame
+                    # print('reserve_frame_idx:  ' + str(reserve_frame_idx))
+                    # print('frame_frame:  ' + str(frame_from_waveforms))
+                    cache["reserve_waveforms"] = cache["waveforms"][
+                        :,
+                        reserve_frame_idx
+                        * self.frame_shift_sample_length : frame_from_waveforms
+                        * self.frame_shift_sample_length,
+                    ]
+                    sample_length = (
+                        frame_from_waveforms - 1
+                    ) * self.frame_shift_sample_length + self.frame_sample_length
+                    cache["waveforms"] = cache["waveforms"][:, :sample_length]
+            else:
+                # update self.reserve_waveforms and self.lfr_splice_cache
+                cache["reserve_waveforms"] = cache["waveforms"][
+                    :, : -(self.frame_sample_length - self.frame_shift_sample_length)
+                ]
+                for i in range(batch_size):
+                    cache["lfr_splice_cache"][i] = torch.cat(
+                        (cache["lfr_splice_cache"][i], feats[i]), dim=0
+                    )
+                return torch.empty(0), feats_lengths
+        else:
+            if is_final:
+                cache["waveforms"] = (
+                    waveforms
+                    if cache["reserve_waveforms"].numel() == 0
+                    else cache["reserve_waveforms"]
+                )
+                feats = torch.stack(cache["lfr_splice_cache"])
+                feats_lengths = torch.zeros(batch_size, dtype=torch.int) + feats.shape[1]
+                feats, feats_lengths, _ = self.forward_lfr_cmvn(
+                    feats, feats_lengths, is_final, cache=cache
+                )
+        # if is_final:
+        #     self.init_cache(cache)
+        return feats, feats_lengths
+
+    def init_cache(self, cache: dict = None):
+        """Init cache.
+        
+            Args:
+                cache: State cache dict for streaming inference.
+            """
+        if cache is None:
+            cache = {}
+        cache["reserve_waveforms"] = torch.empty(0)
+        cache["input_cache"] = torch.empty(0)
+        cache["lfr_splice_cache"] = []
+        cache["waveforms"] = None
+        cache["fbanks"] = None
+        cache["fbanks_lens"] = None
+        return cache
+
+
+class WavFrontendMel23(nn.Module):
+    """Conventional frontend structure for ASR."""
+
+    def __init__(
+        self,
+        fs: int = 16000,
+        frame_length: int = 25,
+        frame_shift: int = 10,
+        lfr_m: int = 1,
+        lfr_n: int = 1,
+        **kwargs,
+    ):
+        """Initialize WavFrontendMel23.
+        
+            Args:
+                fs: TODO.
+                frame_length: TODO.
+                frame_shift: TODO.
+                lfr_m: TODO.
+                lfr_n: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        self.fs = fs
+        self.frame_length = frame_length
+        self.frame_shift = frame_shift
+        self.lfr_m = lfr_m
+        self.lfr_n = lfr_n
+        self.n_mels = 23
+
+    def output_size(self) -> int:
+        """Output size."""
+        return self.n_mels * (2 * self.lfr_m + 1)
+
+    def forward(
+        self, input: torch.Tensor, input_lengths: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass for training.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+            """
+        batch_size = input.size(0)
+        feats = []
+        feats_lens = []
+        for i in range(batch_size):
+            waveform_length = input_lengths[i]
+            waveform = input[i][:waveform_length]
+            waveform = waveform.numpy()
+            mat = eend_ola_feature.stft(waveform, self.frame_length, self.frame_shift)
+            mat = eend_ola_feature.transform(mat)
+            mat = eend_ola_feature.splice(mat, context_size=self.lfr_m)
+            mat = mat[:: self.lfr_n]
+            mat = torch.from_numpy(mat)
+            feat_length = mat.size(0)
+            feats.append(mat)
+            feats_lens.append(feat_length)
+
+        feats_lens = torch.as_tensor(feats_lens)
+        feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0)
+        return feats_pad, feats_lens

funasr/frontends/whisper_frontend.py

@@ -0,0 +1,141 @@
+from typing import Tuple
+import torch
+import torch.nn as nn
+
+
+from funasr.register import tables
+from torch.nn.utils.rnn import pad_sequence
+
+
+@tables.register("frontend_classes", "WhisperFrontend")
+class WhisperFrontend(nn.Module):
+    """Speech Representation Using Encoder Outputs from OpenAI's Whisper Model:
+
+    URL: https://github.com/openai/whisper
+    """
+
+    def __init__(
+        self,
+        fs: int = 16000,
+        whisper_model: str = None,
+        do_pad_trim: bool = True,
+        n_mels: int = 80,
+        permute: bool = False,
+        **kwargs,
+    ):
+        """Initialize WhisperFrontend.
+        
+            Args:
+                fs: TODO.
+                whisper_model: Whisper Model instance.
+                do_pad_trim: TODO.
+                n_mels: TODO.
+                permute: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        assert fs == 16000
+        self.fs = fs
+        import whisper
+        from whisper.audio import HOP_LENGTH, N_FFT, N_SAMPLES
+
+        self.n_fft = N_FFT
+        self.win_length = N_FFT
+        self.hop_length = HOP_LENGTH
+        self.pad_samples = N_SAMPLES
+        self.frame_shift = int(self.hop_length / self.fs * 1000)
+        self.lfr_n = 1
+        self.n_mels = n_mels
+        if whisper_model == "large-v3" or whisper_model == "large":
+            self.n_mels = 128
+
+        filters_path = kwargs.get("filters_path", None)
+        self.filters_path = filters_path
+        if filters_path is not None:
+            from funasr.models.sense_voice.whisper_lib.audio import mel_filters
+
+            self.mel_filters = mel_filters
+        else:
+            self.mel_filters = whisper.audio.mel_filters
+        self.do_pad_trim = do_pad_trim
+        if do_pad_trim:
+            self.pad_or_trim = whisper.pad_or_trim
+        self.permute = permute
+
+        # assert whisper_model in whisper.available_models()
+
+    def output_size(self) -> int:
+        """Output size."""
+        return self.n_mels
+
+    def log_mel_spectrogram(
+        self,
+        audio: torch.Tensor,
+        ilens: torch.Tensor = None,
+    ) -> torch.Tensor:
+        """Log mel spectrogram.
+        
+            Args:
+                audio: TODO.
+                ilens: TODO.
+            """
+        window = torch.hann_window(self.win_length).to(audio.device)
+        stft = torch.stft(audio, self.n_fft, self.hop_length, window=window, return_complex=True)
+
+        # whisper deletes the last frame by default (Shih-Lun)
+        magnitudes = stft[..., :-1].abs() ** 2
+        if self.filters_path is not None:
+            filters = self.mel_filters(audio.device, self.n_mels, self.filters_path)
+        else:
+            filters = self.mel_filters(audio.device, self.n_mels)
+        mel_spec = filters @ magnitudes
+
+        log_spec = torch.clamp(mel_spec, min=1e-10).log10()
+
+        if ilens is not None:
+            olens = ilens // self.hop_length
+        else:
+            olens = None
+
+        log_spec = torch.maximum(
+            log_spec,
+            log_spec.view(audio.size(0), -1).max(dim=-1)[0][:, None, None] - 8.0,
+        )
+        log_spec = (log_spec + 4.0) / 4.0
+
+        return log_spec, olens
+
+    def forward(
+        self,
+        input: torch.Tensor,
+        input_lengths: torch.Tensor,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass for training.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+                **kwargs: Additional keyword arguments.
+            """
+        batch_size = input.size(0)
+        feats = []
+        feats_lens = []
+        input = input.to(torch.float32)
+        for i in range(batch_size):
+            if self.do_pad_trim:
+                feat = self.pad_or_trim(input[i], self.pad_samples)
+            else:
+                feat = input[i]
+            feat, feat_len = self.log_mel_spectrogram(feat[None, :], input_lengths[0])
+            feats.append(feat[0])
+            feats_lens.append(feat_len)
+        feats_lens = torch.as_tensor(feats_lens)
+
+        if batch_size == 1:
+            feats_pad = feats[0][None, :, :]
+        else:
+            feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0)
+        if self.permute:
+            feats_pad = feats_pad.permute(0, 2, 1)
+        return feats_pad, feats_lens

funasr/frontends/windowing.py

@@ -0,0 +1,74 @@
+#!/usr/bin/env python3
+#  2020, Technische Universität München;  Ludwig Kürzinger
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Sliding Window for raw audio input data."""
+
+import torch
+import torch.nn as nn
+from typing import Tuple
+
+
+class SlidingWindow(nn.Module):
+    """Sliding Window.
+    Provides a sliding window over a batched continuous raw audio tensor.
+    Optionally, provides padding (Currently not implemented).
+    Combine this module with a pre-encoder compatible with raw audio data,
+    for example Sinc convolutions.
+    Known issues:
+    Output length is calculated incorrectly if audio shorter than win_length.
+    WARNING: trailing values are discarded - padding not implemented yet.
+    There is currently no additional window function applied to input values.
+    """
+
+    def __init__(
+        self,
+        win_length: int = 400,
+        hop_length: int = 160,
+        channels: int = 1,
+        padding: int = None,
+        fs=None,
+    ):
+        """Initialize.
+        Args:
+            win_length: Length of frame.
+            hop_length: Relative starting point of next frame.
+            channels: Number of input channels.
+            padding: Padding (placeholder, currently not implemented).
+            fs:  Sampling rate (placeholder for compatibility, not used).
+        """
+        super().__init__()
+        self.fs = fs
+        self.win_length = win_length
+        self.hop_length = hop_length
+        self.channels = channels
+        self.padding = padding
+
+    def forward(
+        self, input: torch.Tensor, input_lengths: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Apply a sliding window on the input.
+        Args:
+            input: Input (B, T, C*D) or (B, T*C*D), with D=C=1.
+            input_lengths: Input lengths within batch.
+        Returns:
+            Tensor: Output with dimensions (B, T, C, D), with D=win_length.
+            Tensor: Output lengths within batch.
+        """
+        input_size = input.size()
+        B = input_size[0]
+        T = input_size[1]
+        C = self.channels
+        D = self.win_length
+        # (B, T, C) --> (T, B, C)
+        continuous = input.view(B, T, C).permute(1, 0, 2)
+        windowed = continuous.unfold(0, D, self.hop_length)
+        # (T, B, C, D) --> (B, T, C, D)
+        output = windowed.permute(1, 0, 2, 3).contiguous()
+        # After unfold(), windowed lengths change:
+        output_lengths = (input_lengths - self.win_length) // self.hop_length + 1
+        return output, output_lengths
+
+    def output_size(self) -> int:
+        """Return output length of feature dimension D, i.e. the window length."""
+        return self.win_length