Mirror wlc952/FunASR-bmodel from ModelScope (part 2)

.gitattributes

@@ -46,3 +46,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 *.bmodel filter=lfs diff=lfs merge=lfs -textbmodel/tpu_perf-1.2.35-py3-none-manylinux2014_aarch64.whl filter=lfs diff=lfs merge=lfs -text
+system_ref.wav filter=lfs diff=lfs merge=lfs -text

funasr/models/paraformer_streaming/export_meta.py

@@ -0,0 +1,136 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import types
+import torch
+from funasr.register import tables
+
+
+def export_rebuild_model(model, **kwargs):
+    # self.device = kwargs.get("device")
+    is_onnx = kwargs.get("type", "onnx") == "onnx"
+    encoder_class = tables.encoder_classes.get(kwargs["encoder"] + "Export")
+    model.encoder = encoder_class(model.encoder, onnx=is_onnx)
+
+    predictor_class = tables.predictor_classes.get(kwargs["predictor"] + "Export")
+    model.predictor = predictor_class(model.predictor, onnx=is_onnx)
+
+    if kwargs["decoder"] == "ParaformerSANMDecoder":
+        kwargs["decoder"] = "ParaformerSANMDecoderOnline"
+    decoder_class = tables.decoder_classes.get(kwargs["decoder"] + "Export")
+    model.decoder = decoder_class(model.decoder, onnx=is_onnx)
+
+    from funasr.utils.torch_function import sequence_mask
+
+    model.make_pad_mask = sequence_mask(max_seq_len=None, flip=False)
+
+    import copy
+    import types
+
+    encoder_model = copy.copy(model)
+    decoder_model = copy.copy(model)
+
+    # encoder
+    encoder_model.forward = types.MethodType(export_encoder_forward, encoder_model)
+    encoder_model.export_dummy_inputs = types.MethodType(export_encoder_dummy_inputs, encoder_model)
+    encoder_model.export_input_names = types.MethodType(export_encoder_input_names, encoder_model)
+    encoder_model.export_output_names = types.MethodType(export_encoder_output_names, encoder_model)
+    encoder_model.export_dynamic_axes = types.MethodType(export_encoder_dynamic_axes, encoder_model)
+    encoder_model.export_name = "model"  # types.MethodType(export_encoder_name, encoder_model)
+
+    # decoder
+    decoder_model.forward = types.MethodType(export_decoder_forward, decoder_model)
+    decoder_model.export_dummy_inputs = types.MethodType(export_decoder_dummy_inputs, decoder_model)
+    decoder_model.export_input_names = types.MethodType(export_decoder_input_names, decoder_model)
+    decoder_model.export_output_names = types.MethodType(export_decoder_output_names, decoder_model)
+    decoder_model.export_dynamic_axes = types.MethodType(export_decoder_dynamic_axes, decoder_model)
+    decoder_model.export_name = "decoder"  # types.MethodType(export_decoder_name, decoder_model)
+
+    return encoder_model, decoder_model
+
+
+def export_encoder_forward(
+    self,
+    speech: torch.Tensor,
+    speech_lengths: torch.Tensor,
+):
+    # a. To device
+    batch = {"speech": speech, "speech_lengths": speech_lengths, "online": True}
+    # batch = to_device(batch, device=self.device)
+
+    enc, enc_len = self.encoder(**batch)
+    mask = self.make_pad_mask(enc_len)[:, None, :]
+    alphas, _ = self.predictor.forward_cnn(enc, mask)
+
+    return enc, enc_len, alphas
+
+
+def export_encoder_dummy_inputs(self):
+    speech = torch.randn(2, 30, 560)
+    speech_lengths = torch.tensor([6, 30], dtype=torch.int32)
+    return (speech, speech_lengths)
+
+
+def export_encoder_input_names(self):
+    return ["speech", "speech_lengths"]
+
+
+def export_encoder_output_names(self):
+    return ["enc", "enc_len", "alphas"]
+
+
+def export_encoder_dynamic_axes(self):
+    return {
+        "speech": {0: "batch_size", 1: "feats_length"},
+        "speech_lengths": {
+            0: "batch_size",
+        },
+        "enc": {0: "batch_size", 1: "feats_length"},
+        "enc_len": {
+            0: "batch_size",
+        },
+        "alphas": {0: "batch_size", 1: "feats_length"},
+    }
+
+
+def export_encoder_name(self):
+    return "model.onnx"
+
+
+def export_decoder_forward(
+    self,
+    enc: torch.Tensor,
+    enc_len: torch.Tensor,
+    acoustic_embeds: torch.Tensor,
+    acoustic_embeds_len: torch.Tensor,
+    *args,
+):
+    decoder_out, out_caches = self.decoder(
+        enc, enc_len, acoustic_embeds, acoustic_embeds_len, *args
+    )
+    sample_ids = decoder_out.argmax(dim=-1)
+
+    return decoder_out, sample_ids, out_caches
+
+
+def export_decoder_dummy_inputs(self):
+    dummy_inputs = self.decoder.get_dummy_inputs(enc_size=self.encoder._output_size)
+    return dummy_inputs
+
+
+def export_decoder_input_names(self):
+    return self.decoder.get_input_names()
+
+
+def export_decoder_output_names(self):
+    return self.decoder.get_output_names()
+
+
+def export_decoder_dynamic_axes(self):
+    return self.decoder.get_dynamic_axes()
+
+
+def export_decoder_name(self):
+    return "decoder.onnx"

funasr/models/paraformer_streaming/model.py

@@ -0,0 +1,702 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import time
+import torch
+import logging
+from typing import Dict, Tuple
+from contextlib import contextmanager
+from distutils.version import LooseVersion
+
+from funasr.register import tables
+from funasr.models.ctc.ctc import CTC
+from funasr.utils import postprocess_utils
+from funasr.metrics.compute_acc import th_accuracy
+from funasr.utils.datadir_writer import DatadirWriter
+from funasr.models.paraformer.model import Paraformer
+from funasr.models.paraformer.search import Hypothesis
+from funasr.models.paraformer.cif_predictor import mae_loss
+from funasr.train_utils.device_funcs import force_gatherable
+from funasr.losses.label_smoothing_loss import LabelSmoothingLoss
+from funasr.models.transformer.utils.add_sos_eos import add_sos_eos
+from funasr.models.transformer.utils.nets_utils import make_pad_mask, pad_list
+from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
+
+from funasr.utils.run_bmodel import EngineOV
+import numpy as np
+
+if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
+    from torch.amp import autocast
+else:
+    # Nothing to do if torch<1.6.0
+    @contextmanager
+    def autocast(enabled=True):
+        yield
+
+
+@tables.register("model_classes", "ParaformerStreaming")
+class ParaformerStreaming(Paraformer):
+    """
+    Author: Speech Lab of DAMO Academy, Alibaba Group
+    Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition
+    https://arxiv.org/abs/2206.08317
+    """
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+
+        super().__init__(*args, **kwargs)
+
+        self.sampling_ratio = kwargs.get("sampling_ratio", 0.2)
+
+        self.scama_mask = None
+        if (
+            hasattr(self.encoder, "overlap_chunk_cls")
+            and self.encoder.overlap_chunk_cls is not None
+        ):
+            from funasr.models.scama.chunk_utilis import (
+                build_scama_mask_for_cross_attention_decoder,
+            )
+
+            self.build_scama_mask_for_cross_attention_decoder_fn = (
+                build_scama_mask_for_cross_attention_decoder
+            )
+            self.decoder_attention_chunk_type = kwargs.get("decoder_attention_chunk_type", "chunk")
+
+        self.encoder_bmodel = EngineOV(kwargs["model_path"]+"/encoder_fp32.bmodel", device_id=kwargs['dev_id'])
+        self.decoder0_bmodel = EngineOV(kwargs["model_path"]+"/decoder0_fp32.bmodel", device_id=kwargs['dev_id'])
+        self.decoder1_bmodel = EngineOV(kwargs["model_path"]+"/decoder1_fp32.bmodel", device_id=kwargs['dev_id'])
+
+        self.predictor.cif_conv1d.weight = torch.load(kwargs["model_path"]+"/cif_conv1d_weight.pt",weights_only=True)
+        self.predictor.cif_conv1d.bias   = torch.load(kwargs["model_path"]+"/cif_conv1d_bias.pt",weights_only=True)
+        self.predictor.cif_output.weight = torch.load(kwargs["model_path"]+"/cif_output_weight.pt",weights_only=True)
+        self.predictor.cif_output.bias   = torch.load(kwargs["model_path"]+"/cif_output_bias.pt",weights_only=True)
+
+    def forward(
+        self,
+        speech: torch.Tensor,
+        speech_lengths: torch.Tensor,
+        text: torch.Tensor,
+        text_lengths: torch.Tensor,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:
+        """Encoder + Decoder + Calc loss
+        Args:
+                speech: (Batch, Length, ...)
+                speech_lengths: (Batch, )
+                text: (Batch, Length)
+                text_lengths: (Batch,)
+        """
+        decoding_ind = kwargs.get("decoding_ind")
+        if len(text_lengths.size()) > 1:
+            text_lengths = text_lengths[:, 0]
+        if len(speech_lengths.size()) > 1:
+            speech_lengths = speech_lengths[:, 0]
+
+        batch_size = speech.shape[0]
+
+        # Encoder
+        if hasattr(self.encoder, "overlap_chunk_cls"):
+            ind = self.encoder.overlap_chunk_cls.random_choice(self.training, decoding_ind)
+            encoder_out, encoder_out_lens = self.encode(speech, speech_lengths, ind=ind)
+        else:
+            encoder_out, encoder_out_lens = self.encode(speech, speech_lengths)
+
+        loss_ctc, cer_ctc = None, None
+        loss_pre = None
+        stats = dict()
+
+        # decoder: CTC branch
+
+        if self.ctc_weight > 0.0:
+            if hasattr(self.encoder, "overlap_chunk_cls"):
+                encoder_out_ctc, encoder_out_lens_ctc = self.encoder.overlap_chunk_cls.remove_chunk(
+                    encoder_out, encoder_out_lens, chunk_outs=None
+                )
+            else:
+                encoder_out_ctc, encoder_out_lens_ctc = encoder_out, encoder_out_lens
+
+            loss_ctc, cer_ctc = self._calc_ctc_loss(
+                encoder_out_ctc, encoder_out_lens_ctc, text, text_lengths
+            )
+            # Collect CTC branch stats
+            stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None
+            stats["cer_ctc"] = cer_ctc
+
+        # decoder: Attention decoder branch
+        loss_att, acc_att, cer_att, wer_att, loss_pre, pre_loss_att = self._calc_att_predictor_loss(
+            encoder_out, encoder_out_lens, text, text_lengths
+        )
+
+        # 3. CTC-Att loss definition
+        if self.ctc_weight == 0.0:
+            loss = loss_att + loss_pre * self.predictor_weight
+        else:
+            loss = (
+                self.ctc_weight * loss_ctc
+                + (1 - self.ctc_weight) * loss_att
+                + loss_pre * self.predictor_weight
+            )
+
+        # Collect Attn branch stats
+        stats["loss_att"] = loss_att.detach() if loss_att is not None else None
+        stats["pre_loss_att"] = pre_loss_att.detach() if pre_loss_att is not None else None
+        stats["acc"] = acc_att
+        stats["cer"] = cer_att
+        stats["wer"] = wer_att
+        stats["loss_pre"] = loss_pre.detach().cpu() if loss_pre is not None else None
+
+        stats["loss"] = torch.clone(loss.detach())
+
+        # force_gatherable: to-device and to-tensor if scalar for DataParallel
+        if self.length_normalized_loss:
+            batch_size = (text_lengths + self.predictor_bias).sum()
+        loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device)
+        return loss, stats, weight
+
+    def encode_chunk(
+        self,
+        speech: torch.Tensor,
+        speech_lengths: torch.Tensor,
+        cache: dict = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Frontend + Encoder. Note that this method is used by asr_inference.py
+        Args:
+                speech: (Batch, Length, ...)
+                speech_lengths: (Batch, )
+                ind: int
+        """
+        with autocast("cuda",enabled=False):
+
+            # Data augmentation
+            if self.specaug is not None and self.training:
+                speech, speech_lengths = self.specaug(speech, speech_lengths)
+
+            # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN
+            if self.normalize is not None:
+                speech, speech_lengths = self.normalize(speech, speech_lengths)
+
+        # Forward encoder
+        encoder_out, encoder_out_lens, _ = self.encoder.forward_chunk(
+            speech, speech_lengths, cache=cache["encoder"]
+        )
+        if isinstance(encoder_out, tuple):
+            encoder_out = encoder_out[0]
+
+        return encoder_out, torch.tensor([encoder_out.size(1)])
+
+    def _calc_att_predictor_loss(
+        self,
+        encoder_out: torch.Tensor,
+        encoder_out_lens: torch.Tensor,
+        ys_pad: torch.Tensor,
+        ys_pad_lens: torch.Tensor,
+    ):
+        encoder_out_mask = (
+            ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]
+        ).to(encoder_out.device)
+        if self.predictor_bias == 1:
+            _, ys_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id)
+            ys_pad_lens = ys_pad_lens + self.predictor_bias
+        mask_chunk_predictor = None
+        if self.encoder.overlap_chunk_cls is not None:
+            mask_chunk_predictor = self.encoder.overlap_chunk_cls.get_mask_chunk_predictor(
+                None, device=encoder_out.device, batch_size=encoder_out.size(0)
+            )
+            mask_shfit_chunk = self.encoder.overlap_chunk_cls.get_mask_shfit_chunk(
+                None, device=encoder_out.device, batch_size=encoder_out.size(0)
+            )
+            encoder_out = encoder_out * mask_shfit_chunk
+        pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor(
+            encoder_out,
+            ys_pad,
+            encoder_out_mask,
+            ignore_id=self.ignore_id,
+            mask_chunk_predictor=mask_chunk_predictor,
+            target_label_length=ys_pad_lens,
+        )
+        predictor_alignments, predictor_alignments_len = self.predictor.gen_frame_alignments(
+            pre_alphas, encoder_out_lens
+        )
+
+        scama_mask = None
+        if (
+            self.encoder.overlap_chunk_cls is not None
+            and self.decoder_attention_chunk_type == "chunk"
+        ):
+            encoder_chunk_size = self.encoder.overlap_chunk_cls.chunk_size_pad_shift_cur
+            attention_chunk_center_bias = 0
+            attention_chunk_size = encoder_chunk_size
+            decoder_att_look_back_factor = (
+                self.encoder.overlap_chunk_cls.decoder_att_look_back_factor_cur
+            )
+            mask_shift_att_chunk_decoder = (
+                self.encoder.overlap_chunk_cls.get_mask_shift_att_chunk_decoder(
+                    None, device=encoder_out.device, batch_size=encoder_out.size(0)
+                )
+            )
+            scama_mask = self.build_scama_mask_for_cross_attention_decoder_fn(
+                predictor_alignments=predictor_alignments,
+                encoder_sequence_length=encoder_out_lens,
+                chunk_size=1,
+                encoder_chunk_size=encoder_chunk_size,
+                attention_chunk_center_bias=attention_chunk_center_bias,
+                attention_chunk_size=attention_chunk_size,
+                attention_chunk_type=self.decoder_attention_chunk_type,
+                step=None,
+                predictor_mask_chunk_hopping=mask_chunk_predictor,
+                decoder_att_look_back_factor=decoder_att_look_back_factor,
+                mask_shift_att_chunk_decoder=mask_shift_att_chunk_decoder,
+                target_length=ys_pad_lens,
+                is_training=self.training,
+            )
+        elif self.encoder.overlap_chunk_cls is not None:
+            encoder_out, encoder_out_lens = self.encoder.overlap_chunk_cls.remove_chunk(
+                encoder_out, encoder_out_lens, chunk_outs=None
+            )
+        # 0. sampler
+        decoder_out_1st = None
+        pre_loss_att = None
+        if self.sampling_ratio > 0.0:
+
+            if self.use_1st_decoder_loss:
+                sematic_embeds, decoder_out_1st, pre_loss_att = self.sampler_with_grad(
+                    encoder_out,
+                    encoder_out_lens,
+                    ys_pad,
+                    ys_pad_lens,
+                    pre_acoustic_embeds,
+                    scama_mask,
+                )
+            else:
+                sematic_embeds, decoder_out_1st = self.sampler(
+                    encoder_out,
+                    encoder_out_lens,
+                    ys_pad,
+                    ys_pad_lens,
+                    pre_acoustic_embeds,
+                    scama_mask,
+                )
+        else:
+            sematic_embeds = pre_acoustic_embeds
+
+        # 1. Forward decoder
+        decoder_outs = self.decoder(
+            encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens, scama_mask
+        )
+        decoder_out, _ = decoder_outs[0], decoder_outs[1]
+
+        if decoder_out_1st is None:
+            decoder_out_1st = decoder_out
+        # 2. Compute attention loss
+        loss_att = self.criterion_att(decoder_out, ys_pad)
+        acc_att = th_accuracy(
+            decoder_out_1st.view(-1, self.vocab_size),
+            ys_pad,
+            ignore_label=self.ignore_id,
+        )
+        loss_pre = self.criterion_pre(ys_pad_lens.type_as(pre_token_length), pre_token_length)
+
+        # Compute cer/wer using attention-decoder
+        if self.training or self.error_calculator is None:
+            cer_att, wer_att = None, None
+        else:
+            ys_hat = decoder_out_1st.argmax(dim=-1)
+            cer_att, wer_att = self.error_calculator(ys_hat.cpu(), ys_pad.cpu())
+
+        return loss_att, acc_att, cer_att, wer_att, loss_pre, pre_loss_att
+
+    def sampler(
+        self,
+        encoder_out,
+        encoder_out_lens,
+        ys_pad,
+        ys_pad_lens,
+        pre_acoustic_embeds,
+        chunk_mask=None,
+    ):
+
+        tgt_mask = (~make_pad_mask(ys_pad_lens, maxlen=ys_pad_lens.max())[:, :, None]).to(
+            ys_pad.device
+        )
+        ys_pad_masked = ys_pad * tgt_mask[:, :, 0]
+        if self.share_embedding:
+            ys_pad_embed = self.decoder.output_layer.weight[ys_pad_masked]
+        else:
+            ys_pad_embed = self.decoder.embed(ys_pad_masked)
+        with torch.no_grad():
+            decoder_outs = self.decoder(
+                encoder_out, encoder_out_lens, pre_acoustic_embeds, ys_pad_lens, chunk_mask
+            )
+            decoder_out, _ = decoder_outs[0], decoder_outs[1]
+            pred_tokens = decoder_out.argmax(-1)
+            nonpad_positions = ys_pad.ne(self.ignore_id)
+            seq_lens = (nonpad_positions).sum(1)
+            same_num = ((pred_tokens == ys_pad) & nonpad_positions).sum(1)
+            input_mask = torch.ones_like(nonpad_positions)
+            bsz, seq_len = ys_pad.size()
+            for li in range(bsz):
+                target_num = (
+                    ((seq_lens[li] - same_num[li].sum()).float()) * self.sampling_ratio
+                ).long()
+                if target_num > 0:
+                    input_mask[li].scatter_(
+                        dim=0, index=torch.randperm(seq_lens[li])[:target_num].cuda(), value=0
+                    )
+            input_mask = input_mask.eq(1)
+            input_mask = input_mask.masked_fill(~nonpad_positions, False)
+            input_mask_expand_dim = input_mask.unsqueeze(2).to(pre_acoustic_embeds.device)
+
+        sematic_embeds = pre_acoustic_embeds.masked_fill(
+            ~input_mask_expand_dim, 0
+        ) + ys_pad_embed.masked_fill(input_mask_expand_dim, 0)
+        return sematic_embeds * tgt_mask, decoder_out * tgt_mask
+
+    def calc_predictor(self, encoder_out, encoder_out_lens):
+
+        encoder_out_mask = (
+            ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]
+        ).to(encoder_out.device)
+        mask_chunk_predictor = None
+        if self.encoder.overlap_chunk_cls is not None:
+            mask_chunk_predictor = self.encoder.overlap_chunk_cls.get_mask_chunk_predictor(
+                None, device=encoder_out.device, batch_size=encoder_out.size(0)
+            )
+            mask_shfit_chunk = self.encoder.overlap_chunk_cls.get_mask_shfit_chunk(
+                None, device=encoder_out.device, batch_size=encoder_out.size(0)
+            )
+            encoder_out = encoder_out * mask_shfit_chunk
+        pre_acoustic_embeds, pre_token_length, pre_alphas, pre_peak_index = self.predictor(
+            encoder_out,
+            None,
+            encoder_out_mask,
+            ignore_id=self.ignore_id,
+            mask_chunk_predictor=mask_chunk_predictor,
+            target_label_length=None,
+        )
+        predictor_alignments, predictor_alignments_len = self.predictor.gen_frame_alignments(
+            pre_alphas,
+            encoder_out_lens + 1 if self.predictor.tail_threshold > 0.0 else encoder_out_lens,
+        )
+
+        scama_mask = None
+        if (
+            self.encoder.overlap_chunk_cls is not None
+            and self.decoder_attention_chunk_type == "chunk"
+        ):
+            encoder_chunk_size = self.encoder.overlap_chunk_cls.chunk_size_pad_shift_cur
+            attention_chunk_center_bias = 0
+            attention_chunk_size = encoder_chunk_size
+            decoder_att_look_back_factor = (
+                self.encoder.overlap_chunk_cls.decoder_att_look_back_factor_cur
+            )
+            mask_shift_att_chunk_decoder = (
+                self.encoder.overlap_chunk_cls.get_mask_shift_att_chunk_decoder(
+                    None, device=encoder_out.device, batch_size=encoder_out.size(0)
+                )
+            )
+            scama_mask = self.build_scama_mask_for_cross_attention_decoder_fn(
+                predictor_alignments=predictor_alignments,
+                encoder_sequence_length=encoder_out_lens,
+                chunk_size=1,
+                encoder_chunk_size=encoder_chunk_size,
+                attention_chunk_center_bias=attention_chunk_center_bias,
+                attention_chunk_size=attention_chunk_size,
+                attention_chunk_type=self.decoder_attention_chunk_type,
+                step=None,
+                predictor_mask_chunk_hopping=mask_chunk_predictor,
+                decoder_att_look_back_factor=decoder_att_look_back_factor,
+                mask_shift_att_chunk_decoder=mask_shift_att_chunk_decoder,
+                target_length=None,
+                is_training=self.training,
+            )
+        self.scama_mask = scama_mask
+
+        return pre_acoustic_embeds, pre_token_length, pre_alphas, pre_peak_index
+
+    def calc_predictor_chunk(self, encoder_out, encoder_out_lens, cache=None, **kwargs):
+        is_final = kwargs.get("is_final", False)
+
+        return self.predictor.forward_chunk(encoder_out, cache["encoder"], is_final=is_final)
+
+    def cal_decoder_with_predictor(
+        self, encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens
+    ):
+        decoder_outs = self.decoder(
+            encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens, self.scama_mask
+        )
+        decoder_out = decoder_outs[0]
+        decoder_out = torch.log_softmax(decoder_out, dim=-1)
+        return decoder_out, ys_pad_lens
+
+    def cal_decoder_with_predictor_chunk(
+        self, encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens, cache=None
+    ):
+        decoder_outs = self.decoder.forward_chunk(encoder_out, sematic_embeds, cache["decoder"])
+        decoder_out = decoder_outs
+        decoder_out = torch.log_softmax(decoder_out, dim=-1)
+        return decoder_out, ys_pad_lens
+
+    def init_cache(self, cache: dict = {}, **kwargs):
+        chunk_size = kwargs.get("chunk_size", [0, 10, 5])
+        encoder_chunk_look_back = kwargs.get("encoder_chunk_look_back", 0)
+        decoder_chunk_look_back = kwargs.get("decoder_chunk_look_back", 0)
+        batch_size = 1
+
+        enc_output_size = kwargs["encoder_conf"]["output_size"]
+        feats_dims = kwargs["frontend_conf"]["n_mels"] * kwargs["frontend_conf"]["lfr_m"]
+        cache_encoder = {
+            "start_idx": 0,
+            "cif_hidden": torch.zeros((batch_size, 1, enc_output_size)),
+            "cif_alphas": torch.zeros((batch_size, 1)),
+            "chunk_size": chunk_size,
+            "encoder_chunk_look_back": encoder_chunk_look_back,
+            "last_chunk": False,
+            "opt": None,
+            "feats": torch.zeros((batch_size, chunk_size[0] + chunk_size[2], feats_dims)),
+            "tail_chunk": False,
+        }
+        cache["encoder"] = cache_encoder
+
+        cache_decoder = {
+            "decode_fsmn": None,
+            "decoder_chunk_look_back": decoder_chunk_look_back,
+            "opt": None,
+            "chunk_size": chunk_size,
+        }
+        cache["decoder"] = cache_decoder
+        cache["frontend"] = {}
+        cache["prev_samples"] = torch.empty(0)
+
+        return cache
+
+    def generate_chunk(
+        self,
+        speech,
+        speech_lengths=None,
+        key: list = None,
+        tokenizer=None,
+        frontend=None,
+        **kwargs,
+    ):
+        cache = kwargs.get("cache", {})
+        speech = speech.to(device=kwargs["device"])
+        speech_lengths = speech_lengths.to(device=kwargs["device"])
+
+        # Encoder
+        """
+        encoder_out, encoder_out_lens = self.encode_chunk(
+            speech, speech_lengths, cache=cache, is_final=kwargs.get("is_final", False)
+        )
+        if isinstance(encoder_out, tuple):
+            encoder_out = encoder_out[0]
+        """
+        speech = speech.detach().numpy()
+        start_idx = (torch.tensor([0], dtype=torch.int64).detach().numpy()).astype(np.int32)
+        chunk_size = (torch.tensor(cache['encoder']['chunk_size'], dtype=torch.int64).detach().numpy()).astype(np.int32)
+        feats = cache['encoder']['feats'].detach().numpy()
+        outputs = self.encoder_bmodel([speech, start_idx, chunk_size, feats])
+        cache['encoder']['start_idx'] = torch.from_numpy(outputs[1])
+        cache['encoder']['feats'] = torch.from_numpy(outputs[2])
+        encoder_out = torch.from_numpy(outputs[0])
+        encoder_out_lens = torch.tensor([encoder_out.size(1)])
+
+        # predictor
+        predictor_outs = self.calc_predictor_chunk(
+            encoder_out, encoder_out_lens, cache=cache, is_final=kwargs.get("is_final", False)
+        )
+        pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = (
+            predictor_outs[0],
+            predictor_outs[1],
+            predictor_outs[2],
+            predictor_outs[3],
+        )
+        pre_token_length = pre_token_length.round().long()
+        if torch.max(pre_token_length) < 1:
+            return []
+        """
+        decoder_outs = self.cal_decoder_with_predictor_chunk(
+            encoder_out, encoder_out_lens, pre_acoustic_embeds, pre_token_length, cache=cache
+        )
+        decoder_out, ys_pad_lens = decoder_outs[0], decoder_outs[1]
+        """
+        encoder_out = encoder_out.detach().numpy()
+        pre_acoustic_embeds = pre_acoustic_embeds.detach().numpy()
+        lens = (pre_token_length.detach().numpy()).astype(np.int32)
+        input = [encoder_out, pre_acoustic_embeds]
+        if cache['decoder']['decode_fsmn'] is None:
+            outputs = self.decoder0_bmodel(input)
+        else:
+            input.append(lens)
+            for i in range(16):
+                input.append(cache['decoder']['decode_fsmn'][i].detach().numpy())
+            outputs = self.decoder1_bmodel(input)
+        decode_fsmn = []
+        for i in range(16):
+            decode_fsmn.append(torch.from_numpy(outputs[i+1]))
+        cache['decoder']['decode_fsmn'] = decode_fsmn
+        decoder_out = torch.from_numpy(outputs[0])
+
+        results = []
+        b, n, d = decoder_out.size()
+        if isinstance(key[0], (list, tuple)):
+            key = key[0]
+        for i in range(b):
+            x = encoder_out[i, : encoder_out_lens[i], :]
+            am_scores = decoder_out[i, : pre_token_length[i], :]
+            if self.beam_search is not None:
+                nbest_hyps = self.beam_search(
+                    x=x,
+                    am_scores=am_scores,
+                    maxlenratio=kwargs.get("maxlenratio", 0.0),
+                    minlenratio=kwargs.get("minlenratio", 0.0),
+                )
+
+                nbest_hyps = nbest_hyps[: self.nbest]
+            else:
+
+                yseq = am_scores.argmax(dim=-1)
+                score = am_scores.max(dim=-1)[0]
+                score = torch.sum(score, dim=-1)
+                # pad with mask tokens to ensure compatibility with sos/eos tokens
+                yseq = torch.tensor([self.sos] + yseq.tolist() + [self.eos], device=yseq.device)
+                nbest_hyps = [Hypothesis(yseq=yseq, score=score)]
+            for nbest_idx, hyp in enumerate(nbest_hyps):
+
+                # remove sos/eos and get results
+                last_pos = -1
+                if isinstance(hyp.yseq, list):
+                    token_int = hyp.yseq[1:last_pos]
+                else:
+                    token_int = hyp.yseq[1:last_pos].tolist()
+
+                # remove blank symbol id, which is assumed to be 0
+                token_int = list(
+                    filter(
+                        lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int
+                    )
+                )
+
+                # Change integer-ids to tokens
+                token = tokenizer.ids2tokens(token_int)
+                # text = tokenizer.tokens2text(token)
+
+                result_i = token
+
+                results.extend(result_i)
+
+        return results
+
+    def inference(
+        self,
+        data_in,
+        data_lengths=None,
+        key: list = None,
+        tokenizer=None,
+        frontend=None,
+        cache: dict = {},
+        **kwargs,
+    ):
+
+        # init beamsearch
+        is_use_ctc = kwargs.get("decoding_ctc_weight", 0.0) > 0.00001 and self.ctc != None
+        is_use_lm = (
+            kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None
+        )
+        if self.beam_search is None and (is_use_lm or is_use_ctc):
+            logging.info("enable beam_search")
+            self.init_beam_search(**kwargs)
+            self.nbest = kwargs.get("nbest", 1)
+
+        if len(cache) == 0:
+            self.init_cache(cache, **kwargs)
+
+        meta_data = {}
+        chunk_size = kwargs.get("chunk_size", [0, 10, 5])
+        chunk_stride_samples = int(chunk_size[1] * 960)  # 600ms
+
+        time1 = time.perf_counter()
+        cfg = {"is_final": kwargs.get("is_final", False)}
+        audio_sample_list = load_audio_text_image_video(
+            data_in,
+            fs=frontend.fs,
+            audio_fs=kwargs.get("fs", 16000),
+            data_type=kwargs.get("data_type", "sound"),
+            tokenizer=tokenizer,
+            cache=cfg,
+        )
+        _is_final = cfg["is_final"]  # if data_in is a file or url, set is_final=True
+
+        time2 = time.perf_counter()
+        meta_data["load_data"] = f"{time2 - time1:0.3f}"
+        assert len(audio_sample_list) == 1, "batch_size must be set 1"
+
+        audio_sample = torch.cat((cache["prev_samples"], audio_sample_list[0]))
+
+        n = int(len(audio_sample) // chunk_stride_samples + int(_is_final))
+        m = int(len(audio_sample) % chunk_stride_samples * (1 - int(_is_final)))
+        tokens = []
+        for i in range(n):
+            kwargs["is_final"] = _is_final and i == n - 1
+            audio_sample_i = audio_sample[i * chunk_stride_samples : (i + 1) * chunk_stride_samples]
+            if kwargs["is_final"] and len(audio_sample_i) < 960:
+                cache["encoder"]["tail_chunk"] = True
+                speech = cache["encoder"]["feats"]
+                speech_lengths = torch.tensor([speech.shape[1]], dtype=torch.int64).to(
+                    speech.device
+                )
+            else:
+                # extract fbank feats
+                speech, speech_lengths = extract_fbank(
+                    [audio_sample_i],
+                    data_type=kwargs.get("data_type", "sound"),
+                    frontend=frontend,
+                    cache=cache["frontend"],
+                    is_final=kwargs["is_final"],
+                )
+            time3 = time.perf_counter()
+            meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
+            meta_data["batch_data_time"] = (
+                speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000
+            )
+
+            tokens_i = self.generate_chunk(
+                speech,
+                speech_lengths,
+                key=key,
+                tokenizer=tokenizer,
+                cache=cache,
+                frontend=frontend,
+                **kwargs,
+            )
+            tokens.extend(tokens_i)
+
+        text_postprocessed, _ = postprocess_utils.sentence_postprocess(tokens)
+
+        result_i = {"key": key[0], "text": text_postprocessed}
+        result = [result_i]
+
+        cache["prev_samples"] = audio_sample[:-m]
+        if _is_final:
+            self.init_cache(cache, **kwargs)
+
+        if kwargs.get("output_dir"):
+            if not hasattr(self, "writer"):
+                self.writer = DatadirWriter(kwargs.get("output_dir"))
+            ibest_writer = self.writer[f"{1}best_recog"]
+            ibest_writer["token"][key[0]] = " ".join(tokens)
+            ibest_writer["text"][key[0]] = text_postprocessed
+
+        return result, meta_data
+
+    def export(self, **kwargs):
+        from .export_meta import export_rebuild_model
+
+        models = export_rebuild_model(model=self, **kwargs)
+        return models

funasr/models/paraformer_streaming/template.yaml

@@ -0,0 +1,143 @@
+# This is an example that demonstrates how to configure a model file.
+# You can modify the configuration according to your own requirements.
+
+# to print the register_table:
+# from funasr.register import tables
+# tables.print()
+
+# network architecture
+model: ParaformerStreaming
+model_conf:
+    ctc_weight: 0.0
+    lsm_weight: 0.1
+    length_normalized_loss: true
+    predictor_weight: 1.0
+    predictor_bias: 1
+    sampling_ratio: 0.75
+
+# encoder
+encoder: SANMEncoderChunkOpt
+encoder_conf:
+    output_size: 512
+    attention_heads: 4
+    linear_units: 2048
+    num_blocks: 50
+    dropout_rate: 0.1
+    positional_dropout_rate: 0.1
+    attention_dropout_rate: 0.1
+    input_layer: pe_online
+    pos_enc_class: SinusoidalPositionEncoder
+    normalize_before: true
+    kernel_size: 11
+    sanm_shfit: 0
+    selfattention_layer_type: sanm
+    chunk_size:
+    - 12
+    - 15
+    stride:
+    - 8
+    - 10
+    pad_left:
+    - 0
+    - 0
+    encoder_att_look_back_factor:
+    - 4
+    - 4
+    decoder_att_look_back_factor:
+    - 1
+    - 1
+
+# decoder
+decoder: ParaformerSANMDecoder
+decoder_conf:
+    attention_heads: 4
+    linear_units: 2048
+    num_blocks: 16
+    dropout_rate: 0.1
+    positional_dropout_rate: 0.1
+    self_attention_dropout_rate: 0.1
+    src_attention_dropout_rate: 0.1
+    att_layer_num: 16
+    kernel_size: 11
+    sanm_shfit: 5
+
+predictor: CifPredictorV2
+predictor_conf:
+    idim: 512
+    threshold: 1.0
+    l_order: 1
+    r_order: 1
+    tail_threshold: 0.45
+
+# frontend related
+frontend: WavFrontendOnline
+frontend_conf:
+    fs: 16000
+    window: hamming
+    n_mels: 80
+    frame_length: 25
+    frame_shift: 10
+    lfr_m: 7
+    lfr_n: 6
+
+specaug: SpecAugLFR
+specaug_conf:
+    apply_time_warp: false
+    time_warp_window: 5
+    time_warp_mode: bicubic
+    apply_freq_mask: true
+    freq_mask_width_range:
+    - 0
+    - 30
+    lfr_rate: 6
+    num_freq_mask: 1
+    apply_time_mask: true
+    time_mask_width_range:
+    - 0
+    - 12
+    num_time_mask: 1
+
+train_conf:
+  accum_grad: 1
+  grad_clip: 5
+  max_epoch: 150
+  val_scheduler_criterion:
+      - valid
+      - acc
+  best_model_criterion:
+  -   - valid
+      - acc
+      - max
+  keep_nbest_models: 10
+  log_interval: 50
+
+optim: adam
+optim_conf:
+   lr: 0.0005
+scheduler: warmuplr
+scheduler_conf:
+   warmup_steps: 30000
+
+dataset: AudioDataset
+dataset_conf:
+    index_ds: IndexDSJsonl
+    batch_sampler: BatchSampler
+    batch_type: example # example or length
+    batch_size: 1 # if batch_type is example, batch_size is the numbers of samples; if length, batch_size is source_token_len+target_token_len;
+    max_token_length: 2048 # filter samples if source_token_len+target_token_len > max_token_length,
+    buffer_size: 500
+    shuffle: True
+    num_workers: 0
+
+tokenizer: CharTokenizer
+tokenizer_conf:
+  unk_symbol: <unk>
+  split_with_space: true
+
+
+ctc_conf:
+    dropout_rate: 0.0
+    ctc_type: builtin
+    reduce: true
+    ignore_nan_grad: true
+normalize: null

funasr/models/qwen_audio/audio.py

@@ -0,0 +1,442 @@
+import base64
+import gzip
+from dataclasses import dataclass
+from typing import Dict, Iterable, Optional, List
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+from subprocess import CalledProcessError, run, Popen, PIPE
+
+import os
+from functools import lru_cache
+from typing import Optional, Union
+from funasr.register import tables
+
+
+def exact_div(x, y):
+    assert x % y == 0
+    return x // y
+
+
+# hard-coded audio hyperparameters
+SAMPLE_RATE = 16000
+N_FFT = 400
+N_MELS = 80
+HOP_LENGTH = 160
+CHUNK_LENGTH = 30
+N_SAMPLES = CHUNK_LENGTH * SAMPLE_RATE  # 480000 samples in a 30-second chunk
+N_FRAMES = exact_div(N_SAMPLES, HOP_LENGTH)  # 3000 frames in a mel spectrogram input
+
+N_SAMPLES_PER_TOKEN = HOP_LENGTH * 2  # the initial convolutions has stride 2
+FRAMES_PER_SECOND = exact_div(SAMPLE_RATE, HOP_LENGTH)  # 10ms per audio frame
+TOKENS_PER_SECOND = exact_div(SAMPLE_RATE, N_SAMPLES_PER_TOKEN)  # 20ms per audio token
+
+
+def get_T_after_cnn(L_in, dilation=1):
+    for padding, kernel_size, stride in eval("[(1,3,1)] + [(1,3,2)] "):
+        L_out = L_in + 2 * padding - dilation * (kernel_size - 1) - 1
+        L_out = 1 + L_out // stride
+        L_in = L_out
+    return L_out
+
+
+def load_bytesio_audio(content, sr: int = SAMPLE_RATE):
+    cmd = [
+        "ffmpeg",
+        "-nostdin",
+        "-threads",
+        "0",
+        "-i",
+        "pipe:",
+        "-f",
+        "s16le",
+        "-ac",
+        "1",
+        "-acodec",
+        "pcm_s16le",
+        "-ar",
+        str(sr),
+        "pipe:",
+    ]
+    p = Popen(cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE, bufsize=-1)
+    out, _ = p.communicate(input=content)
+    return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
+
+
+def load_audio(file: str, sr: int = SAMPLE_RATE):
+    """
+    Open an audio file and read as mono waveform, resampling as necessary
+
+    Parameters
+    ----------
+    file: str
+        The audio file to open
+
+    sr: int
+        The sample rate to resample the audio if necessary
+
+    Returns
+    -------
+    A NumPy array containing the audio waveform, in float32 dtype.
+    """
+
+    # This launches a subprocess to decode audio while down-mixing
+    # and resampling as necessary.  Requires the ffmpeg CLI in PATH.
+    # fmt: off
+    cmd = [
+        "ffmpeg",
+        "-nostdin",
+        "-threads", "0",
+        "-i", file,
+        "-f", "s16le",
+        "-ac", "1",
+        "-acodec", "pcm_s16le",
+        "-ar", str(sr),
+        "-"
+    ]
+    # fmt: on
+    try:
+        out = run(cmd, capture_output=True, check=True).stdout
+    except CalledProcessError as e:
+        raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
+
+    return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
+
+
+def pad_or_trim(array, length: int = N_SAMPLES, *, axis: int = -1):
+    """
+    Pad or trim the audio array to N_SAMPLES, as expected by the encoder.
+    """
+    if torch.is_tensor(array):
+        if array.shape[axis] > length:
+            array = array.index_select(dim=axis, index=torch.arange(length, device=array.device))
+
+        if array.shape[axis] < length:
+            pad_widths = [(0, 0)] * array.ndim
+            pad_widths[axis] = (0, length - array.shape[axis])
+            array = F.pad(array, [pad for sizes in pad_widths[::-1] for pad in sizes])
+    else:
+        if array.shape[axis] > length:
+            array = array.take(indices=range(length), axis=axis)
+
+        if array.shape[axis] < length:
+            pad_widths = [(0, 0)] * array.ndim
+            pad_widths[axis] = (0, length - array.shape[axis])
+            array = np.pad(array, pad_widths)
+
+    return array
+
+
+def trim(array, length: int = N_SAMPLES, *, axis: int = -1):
+    """
+    Pad or trim the audio array to N_SAMPLES, as expected by the encoder.
+    """
+    if torch.is_tensor(array):
+        if array.shape[axis] > length:
+            array = array.index_select(dim=axis, index=torch.arange(length, device=array.device))
+    else:
+        if array.shape[axis] > length:
+            array = array.take(indices=range(length), axis=axis)
+    return array
+
+
+@lru_cache(maxsize=None)
+def mel_filters(device, n_mels: int = N_MELS) -> torch.Tensor:
+    """
+    load the mel filterbank matrix for projecting STFT into a Mel spectrogram.
+    Allows decoupling librosa dependency; saved using:
+
+        np.savez_compressed(
+            "mel_filters.npz",
+            mel_80=librosa.filters.mel(sr=16000, n_fft=400, n_mels=80),
+        )
+    """
+    assert n_mels == 80, f"Unsupported n_mels: {n_mels}"
+    with np.load(
+        os.path.join(os.path.dirname(__file__), "mel_filters.npz")  # todo
+        # os.path.join("assets", "mel_filters.npz")
+    ) as f:
+        return torch.from_numpy(f[f"mel_{n_mels}"]).to(device)
+
+
+def log_mel_spectrogram(
+    audio: Union[str, np.ndarray, torch.Tensor],
+    n_mels: int = N_MELS,
+    padding: int = 0,
+    device: Optional[Union[str, torch.device]] = None,
+):
+    """
+    Compute the log-Mel spectrogram of
+
+    Parameters
+    ----------
+    audio: Union[str, np.ndarray, torch.Tensor], shape = (*)
+        The path to audio or either a NumPy array or Tensor containing the audio waveform in 16 kHz
+
+    n_mels: int
+        The number of Mel-frequency filters, only 80 is supported
+
+    padding: int
+        Number of zero samples to pad to the right
+
+    device: Optional[Union[str, torch.device]]
+        If given, the audio tensor is moved to this device before STFT
+
+    Returns
+    -------
+    torch.Tensor, shape = (80, n_frames)
+        A Tensor that contains the Mel spectrogram
+    """
+    if not torch.is_tensor(audio):
+        if isinstance(audio, str):
+            audio = load_audio(audio)
+        audio = torch.from_numpy(audio)
+
+    if device is not None:
+        audio = audio.to(device)
+    if padding > 0:
+        audio = F.pad(audio, (0, padding))
+    window = torch.hann_window(N_FFT).to(audio.device)
+    stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True)
+    magnitudes = stft[..., :-1].abs() ** 2
+
+    filters = mel_filters(audio.device, n_mels)
+    mel_spec = filters @ magnitudes
+
+    log_spec = torch.clamp(mel_spec, min=1e-10).log10()
+    log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
+    log_spec = (log_spec + 4.0) / 4.0
+    return log_spec
+
+
+@dataclass
+class ModelDimensions:
+    n_mels: int
+    n_audio_ctx: int
+    n_audio_state: int
+    n_audio_head: int
+    n_audio_layer: int
+    n_vocab: int
+    n_text_ctx: int
+    n_text_state: int
+    n_text_head: int
+    n_text_layer: int
+
+
+class LayerNorm(nn.LayerNorm):
+    def forward(self, x: Tensor) -> Tensor:
+        # return super().forward(x.float()).type(x.dtype)
+        return super().forward(x).type(x.dtype)
+
+
+class Linear(nn.Linear):
+    def forward(self, x: Tensor) -> Tensor:
+        return F.linear(
+            x,
+            self.weight.to(x.dtype),
+            None if self.bias is None else self.bias.to(x.dtype),
+        )
+
+
+class Conv1d(nn.Conv1d):
+    def _conv_forward(self, x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> Tensor:
+        return super()._conv_forward(
+            x, weight.to(x.dtype), None if bias is None else bias.to(x.dtype)
+        )
+
+
+def sinusoids(length, channels, max_timescale=10000):
+    """Returns sinusoids for positional embedding"""
+    assert channels % 2 == 0
+    log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1)
+    inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2))
+    scaled_time = torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :]
+    return torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1)
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, n_state: int, n_head: int):
+        super().__init__()
+        self.n_head = n_head
+        self.query = Linear(n_state, n_state)
+        self.key = Linear(n_state, n_state, bias=False)
+        self.value = Linear(n_state, n_state)
+        self.out = Linear(n_state, n_state)
+
+    def forward(
+        self,
+        x: Tensor,
+        xa: Optional[Tensor] = None,
+        mask: Optional[Tensor] = None,
+        kv_cache: Optional[dict] = None,
+    ):
+        q = self.query(x)
+
+        if kv_cache is None or xa is None or self.key not in kv_cache:
+            # hooks, if installed (i.e. kv_cache is not None), will prepend the cached kv tensors;
+            # otherwise, perform key/value projections for self- or cross-attention as usual.
+            k = self.key(x if xa is None else xa)
+            v = self.value(x if xa is None else xa)
+        else:
+            # for cross-attention, calculate keys and values once and reuse in subsequent calls.
+            k = kv_cache[self.key]
+            v = kv_cache[self.value]
+
+        wv, qk = self.qkv_attention(q, k, v, mask)
+        return self.out(wv), qk
+
+    def qkv_attention(self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor] = None):
+        n_batch, n_ctx, n_state = q.shape
+        scale = (n_state // self.n_head) ** -0.25
+        q = q.view(*q.shape[:2], self.n_head, -1).permute(0, 2, 1, 3) * scale
+        k = k.view(*k.shape[:2], self.n_head, -1).permute(0, 2, 3, 1) * scale
+        v = v.view(*v.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
+
+        qk = q @ k
+        if mask is not None:
+            qk += mask
+
+        w = F.softmax(qk, dim=-1).to(q.dtype)
+        return (w @ v).permute(0, 2, 1, 3).flatten(start_dim=2), qk.detach()
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, n_state: int, n_head: int, cross_attention: bool = False):
+        super().__init__()
+
+        self.attn = MultiHeadAttention(n_state, n_head)
+        self.attn_ln = LayerNorm(n_state)
+
+        self.cross_attn = MultiHeadAttention(n_state, n_head) if cross_attention else None
+        self.cross_attn_ln = LayerNorm(n_state) if cross_attention else None
+
+        n_mlp = n_state * 4
+        self.mlp = nn.Sequential(Linear(n_state, n_mlp), nn.GELU(), Linear(n_mlp, n_state))
+        self.mlp_ln = LayerNorm(n_state)
+
+    def forward(
+        self,
+        x: Tensor,
+        xa: Optional[Tensor] = None,
+        mask: Optional[Tensor] = None,
+        kv_cache: Optional[dict] = None,
+    ):
+        x = x + self.attn(self.attn_ln(x), mask=mask, kv_cache=kv_cache)[0]
+        if self.cross_attn:
+            x = x + self.cross_attn(self.cross_attn_ln(x), xa, kv_cache=kv_cache)[0]
+        x = x + self.mlp(self.mlp_ln(x))
+        return x
+
+
+@tables.register("encoder_classes", "QwenAudioEncoder")
+class QwenAudioEncoder(nn.Module):
+    def __init__(
+        self,
+        n_mels: int,
+        n_ctx: int,
+        n_state: int,
+        n_head: int,
+        n_layer: int,
+        output_dim: int = 512,
+        avg_pool: bool = True,
+        add_audio_bos_eos_token: bool = True,
+        **kwargs,
+    ):
+        super().__init__()
+        self.conv1 = Conv1d(n_mels, n_state, kernel_size=3, padding=1)
+        self.conv2 = Conv1d(n_state, n_state, kernel_size=3, stride=2, padding=1)
+        self.register_buffer("positional_embedding", sinusoids(n_ctx, n_state))
+
+        self.blocks: Iterable[ResidualAttentionBlock] = nn.ModuleList(
+            [ResidualAttentionBlock(n_state, n_head) for _ in range(n_layer)]
+        )
+        self.ln_post = LayerNorm(n_state)
+
+        if avg_pool:
+            self.avg_pooler = nn.AvgPool1d(2, stride=2)
+        else:
+            self.avg_pooler = None
+        self.proj = nn.Linear(n_state, output_dim)
+        if add_audio_bos_eos_token:
+            self.audio_bos_eos_token = nn.Embedding(2, output_dim)
+        else:
+            self.audio_bos_eos_token = None
+        self.output_dim = output_dim
+        self.n_head = n_head
+
+    def forward(self, x: Tensor, padding_mask: Tensor = None, audio_lengths: Tensor = None):
+        """
+        x : torch.Tensor, shape = (batch_size, n_mels, n_ctx)
+            the mel spectrogram of the audio
+        """
+        x = x.to(dtype=self.conv1.weight.dtype, device=self.conv1.weight.device)
+        if audio_lengths is not None:
+            input_mel_len = audio_lengths[:, 0] * 2
+            max_mel_len_in_batch = input_mel_len.max()
+            x = x[:, :, :max_mel_len_in_batch]
+        x = F.gelu(self.conv1(x))
+        x = F.gelu(self.conv2(x))
+        x = x.permute(0, 2, 1)  # B, L, D
+        bsz = x.size(0)
+        src_len = x.size(1)
+
+        self.input_positional_embedding = self.positional_embedding[:src_len]
+        assert (
+            x.shape[1:] == self.input_positional_embedding.shape
+        ), f"incorrect audio shape: {x.shape[1:], self.input_positional_embedding.shape}"
+        x = (x + self.input_positional_embedding).to(x.dtype)
+        if padding_mask is not None:
+            padding_mask = padding_mask.to(
+                dtype=self.conv1.weight.dtype, device=self.conv1.weight.device
+            )
+            batch_src_len = padding_mask.size(1)
+            x = x[:, :batch_src_len, :]
+            padding_mask = padding_mask.view(bsz, -1, batch_src_len)
+            padding_mask_ = padding_mask.all(1)
+            x[padding_mask_] = 0
+            key_padding_mask = (
+                padding_mask_.view(bsz, 1, 1, batch_src_len)
+                .expand(-1, self.n_head, -1, -1)
+                .reshape(bsz, self.n_head, 1, batch_src_len)
+            )
+            new_padding_mask = torch.zeros_like(key_padding_mask, dtype=x.dtype)
+            padding_mask = new_padding_mask.masked_fill(key_padding_mask, float("-inf"))
+
+        for block in self.blocks:
+            x = block(x, mask=padding_mask)
+
+        if self.avg_pooler:
+            x = x.permute(0, 2, 1)
+            x = self.avg_pooler(x)
+            x = x.permute(0, 2, 1)
+
+        x = self.ln_post(x)
+        x = self.proj(x)
+
+        if self.audio_bos_eos_token is not None:
+            bos = self.audio_bos_eos_token.weight[0][None, :]
+            eos = self.audio_bos_eos_token.weight[1][None, :]
+        else:
+            bos, eos = None, None
+        return x, bos, eos
+
+    def encode(self, input_audios: Tensor, input_audio_lengths: Tensor, audio_span_tokens: List):
+        real_input_audio_lens = input_audio_lengths[:, 0].tolist()
+        max_len_in_batch = max(real_input_audio_lens)
+        padding_mask = torch.ones([input_audios.size(0), max_len_in_batch]).to(
+            dtype=self.conv1.weight.dtype, device=self.conv1.weight.device
+        )
+        for index in range(len(input_audios)):
+            padding_mask[index, : input_audio_lengths[index][0].item()] = 0
+        x, bos, eos = self(input_audios, padding_mask, input_audio_lengths)
+        output_audios = []
+        for i in range(len(audio_span_tokens)):
+            audio_span = audio_span_tokens[i]
+            audio = x[i][: audio_span - 2]
+            if bos is not None:
+                audio = torch.concat([bos, audio, eos])
+            assert len(audio) == audio_span
+            output_audios.append(audio)
+        return output_audios

funasr/models/qwen_audio/model.py

@@ -0,0 +1,151 @@
+from dataclasses import dataclass
+from typing import Dict
+from typing import Iterable, Optional
+import time
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch import nn
+import whisper
+from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
+
+
+from funasr.register import tables
+
+
+@tables.register("model_classes", "Qwen/Qwen-Audio")
+@tables.register("model_classes", "Qwen-Audio")
+@tables.register("model_classes", "Qwen/QwenAudio")
+@tables.register("model_classes", "QwenAudio")
+@tables.register("model_classes", "QwenAudioWarp")
+class QwenAudioWarp(nn.Module):
+    """
+    Qwen-Audio: Advancing Universal Audio Understanding via Unified Large-Scale Audio-Language Models
+    https://arxiv.org/abs/2311.07919
+    Modified from https://github.com/QwenLM/Qwen-Audio
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+        from transformers.generation import GenerationConfig
+
+        model_or_path = kwargs.get("model_path", "QwenAudio")
+        model = AutoModelForCausalLM.from_pretrained(
+            model_or_path, device_map="cpu", trust_remote_code=True
+        )
+        tokenizer = AutoTokenizer.from_pretrained(model_or_path, trust_remote_code=True)
+
+        self.model = model
+        self.tokenizer = tokenizer
+
+    def forward(
+        self,
+    ):
+        pass
+
+    def inference(
+        self,
+        data_in,
+        data_lengths=None,
+        key: list = None,
+        tokenizer=None,
+        frontend=None,
+        **kwargs,
+    ):
+        if kwargs.get("batch_size", 1) > 1:
+            raise NotImplementedError("batch decoding is not implemented")
+
+        meta_data = {}
+        # meta_data["batch_data_time"] = -1
+        prompt = kwargs.get(
+            "prompt", "<|startoftranscription|><|en|><|transcribe|><|en|><|notimestamps|><|wo_itn|>"
+        )
+        query = f"<audio>{data_in[0]}</audio>{prompt}"
+        audio_info = self.tokenizer.process_audio(query)
+        inputs = self.tokenizer(query, return_tensors="pt", audio_info=audio_info)
+        inputs = inputs.to(self.model.device)
+        pred = self.model.generate(**inputs, audio_info=audio_info)
+        response = self.tokenizer.decode(
+            pred.cpu()[0], skip_special_tokens=False, audio_info=audio_info
+        )
+
+        results = []
+        result_i = {"key": key[0], "text": response}
+
+        results.append(result_i)
+
+        return results, meta_data
+
+
+@tables.register("model_classes", "Qwen/Qwen-Audio-Chat")
+@tables.register("model_classes", "Qwen/QwenAudioChat")
+@tables.register("model_classes", "Qwen-Audio-Chat")
+@tables.register("model_classes", "QwenAudioChat")
+@tables.register("model_classes", "QwenAudioChatWarp")
+class QwenAudioChatWarp(nn.Module):
+    def __init__(self, *args, **kwargs):
+        """
+        Qwen-Audio: Advancing Universal Audio Understanding via Unified Large-Scale Audio-Language Models
+        https://arxiv.org/abs/2311.07919
+        Modified from https://github.com/QwenLM/Qwen-Audio
+        """
+        super().__init__()
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+        from transformers.generation import GenerationConfig
+
+        model_or_path = kwargs.get("model_path", "QwenAudio")
+        bf16 = kwargs.get("bf16", False)
+        fp16 = kwargs.get("fp16", False)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_or_path, device_map="cpu", bf16=bf16, fp16=fp16, trust_remote_code=True
+        )
+        tokenizer = AutoTokenizer.from_pretrained(model_or_path, trust_remote_code=True)
+
+        self.model = model
+        self.tokenizer = tokenizer
+
+    def forward(
+        self,
+    ):
+        pass
+
+    def inference(
+        self,
+        data_in,
+        data_lengths=None,
+        key: list = None,
+        tokenizer=None,
+        frontend=None,
+        **kwargs,
+    ):
+        if kwargs.get("batch_size", 1) > 1:
+            raise NotImplementedError("batch decoding is not implemented")
+
+        meta_data = {}
+
+        prompt = kwargs.get("prompt", "what does the person say?")
+        cache = kwargs.get("cache", {})
+        history = cache.get("history", None)
+        if data_in[0] is not None:
+            # 1st dialogue turn
+            query = self.tokenizer.from_list_format(
+                [
+                    {"audio": data_in[0]},  # Either a local path or an url
+                    {"text": prompt},
+                ]
+            )
+        else:
+            query = prompt
+        response, history = self.model.chat(self.tokenizer, query=query, history=history)
+        cache["history"] = history
+        # print(response)
+        # The person says: "mister quilter is the apostle of the middle classes and we are glad to welcome his gospel".
+
+        results = []
+        result_i = {"key": key[0], "text": response}
+
+        results.append(result_i)
+
+        return results, meta_data

funasr/models/qwen_audio/template.yaml

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

funasr/models/rwkv_bat/cuda_decoder/wkv_cuda.cu

@@ -0,0 +1,135 @@
+// Copied from https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v4/cuda/wkv_cuda.cu
+
+#include <stdio.h>
+#include <assert.h>
+
+#define MIN_VALUE (-1e38)
+
+template <typename F>
+__global__ void kernel_forward(const int B, const int T, const int C,
+                               const F *__restrict__ const _w, const F *__restrict__ const _u, const F *__restrict__ const _k, const F *__restrict__ const _v,
+                               F *__restrict__ const _y) {
+    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    const int _b = idx / C;
+    const int _c = idx % C;
+    const int _offset = _b * T * C + _c;
+
+    F u = _u[_c];
+    F w = _w[_c];
+    const F *__restrict__ const k = _k + _offset;
+    const F *__restrict__ const v = _v + _offset;
+    F *__restrict__ const y = _y + _offset;
+
+    // aa and bb are running sums divided by exp(pp) (to avoid overflow)
+    F aa = 0, bb = 0, pp = MIN_VALUE;
+    for (int i = 0; i < T; i++) {
+        const int ii = i * C;
+        const F kk = k[ii];
+        const F vv = v[ii];
+
+        F ww = u + kk;
+        F p = max(pp, ww);
+        F e1 = exp(pp - p);
+        F e2 = exp(ww - p);
+        y[ii] = (e1 * aa + e2 * vv) / (e1 * bb + e2);
+
+        ww = w + pp;
+        p = max(ww, kk);
+        e1 = exp(ww - p);
+        e2 = exp(kk - p);
+        aa = e1 * aa + e2 * vv;
+        bb = e1 * bb + e2;
+        pp = p;
+    }
+}
+
+template <typename F>
+__global__ void kernel_backward(const int B, const int T, const int C,
+                                const F *__restrict__ const _w, const F *__restrict__ const _u, const F *__restrict__ const _k, const F *__restrict__ const _v,
+                                const F *__restrict__ const _y, const F *__restrict__ const _gy,
+                                F *__restrict__ const _gw, F *__restrict__ const _gu, F *__restrict__ const _gk, F *__restrict__ const _gv) {
+    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    const int _b = idx / C;
+    const int _c = idx % C;
+    const int _offset = _b * T * C + _c;
+
+    F u = _u[_c];
+    F w = _w[_c];
+    const F *__restrict__ const k = _k + _offset;
+    const F *__restrict__ const v = _v + _offset;
+    const F *__restrict__ const y = _y + _offset;
+    const F *__restrict__ const gy = _gy + _offset;
+    F *__restrict__ const gk = _gk + _offset;
+    F *__restrict__ const gv = _gv + _offset;
+
+    F q[Tmax], r[Tmax];
+
+    F gw = 0, gu = 0, aa = 0, bb = 0, ga = 0, gb = 0, pp = MIN_VALUE;
+    for (int i = 0; i < T; i++) {
+        const int ii = i * C;
+        const F kk = k[ii];
+        const F vv = v[ii];
+        const F yy = y[ii];
+
+        F ww = u + kk;
+        F p = max(pp, ww);
+        F e1 = exp(pp - p);
+        F e2 = exp(ww - p);
+        const F qq = gy[ii] / (e1 * bb + e2);
+        gw += (ga - gb * yy) * e1 * qq;
+        gu += (vv - yy) * e2 * qq;
+        q[i] = qq;
+        r[i] = ww - p;
+
+        ww = w + pp;
+        p = max(ww, kk);
+        e1 = exp(ww - p);
+        e2 = exp(kk - p);
+        ga = e1 * (aa + ga);
+        gb = e1 * (bb + gb);
+        aa = e1 * aa + e2 * vv;
+        bb = e1 * bb + e2;
+        pp = p;
+    }
+    const int _offsetBC = _b * C + _c;
+    _gw[_offsetBC] = gw * _w[_c]; // multiply by w because of w -> -exp(w) in python forward()
+    _gu[_offsetBC] = gu;
+
+    aa = 0, bb = 0, pp = MIN_VALUE;
+    for (int i = T - 1; i >= 0; i--) {
+        const int ii = i * C;
+        const F kk = k[ii];
+        const F vv = v[ii];
+        const F yy = y[ii];
+        const F qq = q[i];
+        const F rr = r[i];
+
+        F e1 = qq * exp(rr);
+        F e2 = exp(kk + pp);
+        gk[ii] = e1 * (vv - yy) + e2 * (aa * vv + bb);
+        gv[ii] = e1 + e2 * aa;
+
+        const F ww = w + pp;
+        const F www = rr - u - kk;
+        const F p = max(ww, www);
+        e1 = exp(ww - p);
+        e2 = qq * exp(www - p);
+        aa = e1 * aa + e2;
+        bb = e1 * bb - e2 * yy;
+        pp = p;
+    }
+}
+
+void cuda_forward(int B, int T, int C, float *w, float *u, float *k, float *v, float *y) {
+    dim3 threadsPerBlock( min(C, 32) ); // requires --maxrregcount 60 for optimal performance
+    assert(B * C % threadsPerBlock.x == 0);
+    dim3 numBlocks(B * C / threadsPerBlock.x);
+    kernel_forward<<<numBlocks, threadsPerBlock>>>(B, T, C, w, u, k, v, y);
+}
+
+void cuda_backward(int B, int T, int C, float *w, float *u, float *k, float *v, float *y, float *gy, float *gw, float *gu, float *gk, float *gv) {
+    dim3 threadsPerBlock( min(C, 32) ); // requires --maxrregcount 60 for optimal performance
+    assert(B * C % threadsPerBlock.x == 0);
+    dim3 numBlocks(B * C / threadsPerBlock.x);
+    kernel_backward<<<numBlocks, threadsPerBlock>>>(B, T, C, w, u, k, v, y, gy, gw, gu, gk, gv);
+}

funasr/models/rwkv_bat/cuda_decoder/wkv_op.cpp

@@ -0,0 +1,37 @@
+/*
+ *    Bsed on https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v4/cuda/wkv_op.cpp
+   Function signatures were modified based on https://github.com/huggingface/transformers/blob/main/src/transformers/kernels/rwkv/wkv_op.cpp
+
+ */
+
+#include <torch/extension.h>
+
+void cuda_forward(int B, int T, int C, float *w, float *u, float *k, float *v, float *y);
+
+void cuda_backward(int B, int T, int C, float *w, float *u, float *k, float *v, float *y, float *gy, float *gw, float *gu, float *gk, float *gv);
+
+void forward(torch::Tensor &w, torch::Tensor &u, torch::Tensor &k, torch::Tensor &v, torch::Tensor &y) {
+  const int B = k.size(0);
+  const int T = k.size(1);
+  const int C = k.size(2);
+
+  cuda_forward(B, T, C, w.data_ptr<float>(), u.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(), y.data_ptr<float>());
+}
+
+void backward(torch::Tensor &w, torch::Tensor &u, torch::Tensor &k, torch::Tensor &v, torch::Tensor &y, torch::Tensor &gy, torch::Tensor &gw, torch::Tensor &gu, torch::Tensor &gk, torch::Tensor &gv) {
+  const int B = k.size(0);
+  const int T = k.size(1);
+  const int C = k.size(2);
+
+  cuda_backward(B, T, C, w.data_ptr<float>(), u.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(), y.data_ptr<float>(), gy.data_ptr<float>(), gw.data_ptr<float>(), gu.data_ptr<float>(), gk.data_ptr<float>(), gv.data_ptr<float>());
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("forward", &forward, "wkv forward");
+    m.def("backward", &backward, "wkv backward");
+}
+
+TORCH_LIBRARY(wkv_decoder, m) {
+    m.def("forward", forward);
+    m.def("backward", backward);
+}

funasr/models/rwkv_bat/cuda_encoder/wkv_cuda.cu

@@ -0,0 +1,135 @@
+// Copied from https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v4/cuda/wkv_cuda.cu
+
+#include <stdio.h>
+#include <assert.h>
+
+#define MIN_VALUE (-1e38)
+
+template <typename F>
+__global__ void kernel_forward(const int B, const int T, const int C,
+                               const F *__restrict__ const _w, const F *__restrict__ const _u, const F *__restrict__ const _k, const F *__restrict__ const _v,
+                               F *__restrict__ const _y) {
+    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    const int _b = idx / C;
+    const int _c = idx % C;
+    const int _offset = _b * T * C + _c;
+
+    F u = _u[_c];
+    F w = _w[_c];
+    const F *__restrict__ const k = _k + _offset;
+    const F *__restrict__ const v = _v + _offset;
+    F *__restrict__ const y = _y + _offset;
+
+    // aa and bb are running sums divided by exp(pp) (to avoid overflow)
+    F aa = 0, bb = 0, pp = MIN_VALUE;
+    for (int i = 0; i < T; i++) {
+        const int ii = i * C;
+        const F kk = k[ii];
+        const F vv = v[ii];
+
+        F ww = u + kk;
+        F p = max(pp, ww);
+        F e1 = exp(pp - p);
+        F e2 = exp(ww - p);
+        y[ii] = (e1 * aa + e2 * vv) / (e1 * bb + e2);
+
+        ww = w + pp;
+        p = max(ww, kk);
+        e1 = exp(ww - p);
+        e2 = exp(kk - p);
+        aa = e1 * aa + e2 * vv;
+        bb = e1 * bb + e2;
+        pp = p;
+    }
+}
+
+template <typename F>
+__global__ void kernel_backward(const int B, const int T, const int C,
+                                const F *__restrict__ const _w, const F *__restrict__ const _u, const F *__restrict__ const _k, const F *__restrict__ const _v,
+                                const F *__restrict__ const _y, const F *__restrict__ const _gy,
+                                F *__restrict__ const _gw, F *__restrict__ const _gu, F *__restrict__ const _gk, F *__restrict__ const _gv) {
+    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    const int _b = idx / C;
+    const int _c = idx % C;
+    const int _offset = _b * T * C + _c;
+
+    F u = _u[_c];
+    F w = _w[_c];
+    const F *__restrict__ const k = _k + _offset;
+    const F *__restrict__ const v = _v + _offset;
+    const F *__restrict__ const y = _y + _offset;
+    const F *__restrict__ const gy = _gy + _offset;
+    F *__restrict__ const gk = _gk + _offset;
+    F *__restrict__ const gv = _gv + _offset;
+
+    F q[Tmax], r[Tmax];
+
+    F gw = 0, gu = 0, aa = 0, bb = 0, ga = 0, gb = 0, pp = MIN_VALUE;
+    for (int i = 0; i < T; i++) {
+        const int ii = i * C;
+        const F kk = k[ii];
+        const F vv = v[ii];
+        const F yy = y[ii];
+
+        F ww = u + kk;
+        F p = max(pp, ww);
+        F e1 = exp(pp - p);
+        F e2 = exp(ww - p);
+        const F qq = gy[ii] / (e1 * bb + e2);
+        gw += (ga - gb * yy) * e1 * qq;
+        gu += (vv - yy) * e2 * qq;
+        q[i] = qq;
+        r[i] = ww - p;
+
+        ww = w + pp;
+        p = max(ww, kk);
+        e1 = exp(ww - p);
+        e2 = exp(kk - p);
+        ga = e1 * (aa + ga);
+        gb = e1 * (bb + gb);
+        aa = e1 * aa + e2 * vv;
+        bb = e1 * bb + e2;
+        pp = p;
+    }
+    const int _offsetBC = _b * C + _c;
+    _gw[_offsetBC] = gw * _w[_c]; // multiply by w because of w -> -exp(w) in python forward()
+    _gu[_offsetBC] = gu;
+
+    aa = 0, bb = 0, pp = MIN_VALUE;
+    for (int i = T - 1; i >= 0; i--) {
+        const int ii = i * C;
+        const F kk = k[ii];
+        const F vv = v[ii];
+        const F yy = y[ii];
+        const F qq = q[i];
+        const F rr = r[i];
+
+        F e1 = qq * exp(rr);
+        F e2 = exp(kk + pp);
+        gk[ii] = e1 * (vv - yy) + e2 * (aa * vv + bb);
+        gv[ii] = e1 + e2 * aa;
+
+        const F ww = w + pp;
+        const F www = rr - u - kk;
+        const F p = max(ww, www);
+        e1 = exp(ww - p);
+        e2 = qq * exp(www - p);
+        aa = e1 * aa + e2;
+        bb = e1 * bb - e2 * yy;
+        pp = p;
+    }
+}
+
+void cuda_forward(int B, int T, int C, float *w, float *u, float *k, float *v, float *y) {
+    dim3 threadsPerBlock( min(C, 32) ); // requires --maxrregcount 60 for optimal performance
+    assert(B * C % threadsPerBlock.x == 0);
+    dim3 numBlocks(B * C / threadsPerBlock.x);
+    kernel_forward<<<numBlocks, threadsPerBlock>>>(B, T, C, w, u, k, v, y);
+}
+
+void cuda_backward(int B, int T, int C, float *w, float *u, float *k, float *v, float *y, float *gy, float *gw, float *gu, float *gk, float *gv) {
+    dim3 threadsPerBlock( min(C, 32) ); // requires --maxrregcount 60 for optimal performance
+    assert(B * C % threadsPerBlock.x == 0);
+    dim3 numBlocks(B * C / threadsPerBlock.x);
+    kernel_backward<<<numBlocks, threadsPerBlock>>>(B, T, C, w, u, k, v, y, gy, gw, gu, gk, gv);
+}

funasr/models/rwkv_bat/cuda_encoder/wkv_op.cpp

@@ -0,0 +1,37 @@
+/*
+ *    Bsed on https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v4/cuda/wkv_op.cpp
+   Function signatures were modified based on https://github.com/huggingface/transformers/blob/main/src/transformers/kernels/rwkv/wkv_op.cpp
+
+ */
+
+#include <torch/extension.h>
+
+void cuda_forward(int B, int T, int C, float *w, float *u, float *k, float *v, float *y);
+
+void cuda_backward(int B, int T, int C, float *w, float *u, float *k, float *v, float *y, float *gy, float *gw, float *gu, float *gk, float *gv);
+
+void forward(torch::Tensor &w, torch::Tensor &u, torch::Tensor &k, torch::Tensor &v, torch::Tensor &y) {
+  const int B = k.size(0);
+  const int T = k.size(1);
+  const int C = k.size(2);
+
+  cuda_forward(B, T, C, w.data_ptr<float>(), u.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(), y.data_ptr<float>());
+}
+
+void backward(torch::Tensor &w, torch::Tensor &u, torch::Tensor &k, torch::Tensor &v, torch::Tensor &y, torch::Tensor &gy, torch::Tensor &gw, torch::Tensor &gu, torch::Tensor &gk, torch::Tensor &gv) {
+  const int B = k.size(0);
+  const int T = k.size(1);
+  const int C = k.size(2);
+
+  cuda_backward(B, T, C, w.data_ptr<float>(), u.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(), y.data_ptr<float>(), gy.data_ptr<float>(), gw.data_ptr<float>(), gu.data_ptr<float>(), gk.data_ptr<float>(), gv.data_ptr<float>());
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("forward", &forward, "wkv forward");
+    m.def("backward", &backward, "wkv backward");
+}
+
+TORCH_LIBRARY(wkv_encoder, m) {
+    m.def("forward", forward);
+    m.def("backward", backward);
+}

funasr/models/rwkv_bat/rwkv.py

@@ -0,0 +1,145 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import torch
+from typing import Dict, Optional, Tuple
+
+from funasr.models.transformer.layer_norm import LayerNorm
+from funasr.models.rwkv_bat.rwkv_feed_forward import FeedForward
+from funasr.models.rwkv_bat.rwkv_attention import EncoderSelfAttention, DecoderSelfAttention
+
+
+class RWKV(torch.nn.Module):
+    """RWKV module.
+
+    Args:
+        size: Input/Output size.
+        linear_size: Feed-forward hidden size.
+        attention_size: SelfAttention hidden size.
+        context_size: Context size for WKV computation.
+        block_id: Block index.
+        num_blocks: Number of blocks in the architecture.
+        normalization_class: Normalization layer class.
+        normalization_args: Normalization layer arguments.
+        att_dropout_rate: Dropout rate for the attention module.
+        ffn_dropout_rate: Dropout rate for the feed-forward module.
+
+    """
+
+    def __init__(
+        self,
+        size: int,
+        linear_size: int,
+        attention_size: int,
+        context_size: int,
+        block_id: int,
+        num_blocks: int,
+        att_dropout_rate: float = 0.0,
+        ffn_dropout_rate: float = 0.0,
+        dropout_rate: float = 0.0,
+    ) -> None:
+        """Construct a RWKV object."""
+        super().__init__()
+
+        self.layer_norm_att = LayerNorm(size)
+        self.layer_norm_ffn = LayerNorm(size)
+
+        self.att = EncoderSelfAttention(
+            size, attention_size, context_size, block_id, att_dropout_rate, num_blocks
+        )
+        self.dropout_att = torch.nn.Dropout(p=dropout_rate)
+
+        self.ffn = FeedForward(size, linear_size, block_id, ffn_dropout_rate, num_blocks)
+        self.dropout_ffn = torch.nn.Dropout(p=dropout_rate)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        state: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Compute receptance weighted key value.
+
+        Args:
+            x: RWKV input sequences. (B, L, size)
+            state: Decoder hidden states. [5 x (B, D_att/size, N)]
+
+        Returns:
+            x: RWKV output sequences. (B, L, size)
+            x: Decoder hidden states. [5 x (B, D_att/size, N)]
+
+        """
+        att, state = self.att(self.layer_norm_att(x), state=state)
+        x = x + self.dropout_att(att)
+        ffn, state = self.ffn(self.layer_norm_ffn(x), state=state)
+        x = x + self.dropout_ffn(ffn)
+        return x, state
+
+
+class RWKVDecoderLayer(torch.nn.Module):
+    """RWKV module.
+
+    Args:
+        size: Input/Output size.
+        linear_size: Feed-forward hidden size.
+        attention_size: SelfAttention hidden size.
+        context_size: Context size for WKV computation.
+        block_id: Block index.
+        num_blocks: Number of blocks in the architecture.
+        normalization_class: Normalization layer class.
+        normalization_args: Normalization layer arguments.
+        att_dropout_rate: Dropout rate for the attention module.
+        ffn_dropout_rate: Dropout rate for the feed-forward module.
+
+    """
+
+    def __init__(
+        self,
+        size: int,
+        linear_size: int,
+        attention_size: int,
+        context_size: int,
+        block_id: int,
+        num_blocks: int,
+        att_dropout_rate: float = 0.0,
+        ffn_dropout_rate: float = 0.0,
+        dropout_rate: float = 0.0,
+    ) -> None:
+        """Construct a RWKV object."""
+        super().__init__()
+
+        self.layer_norm_att = LayerNorm(size)
+        self.layer_norm_ffn = LayerNorm(size)
+
+        self.att = DecoderSelfAttention(
+            size, attention_size, context_size, block_id, att_dropout_rate, num_blocks
+        )
+        self.dropout_att = torch.nn.Dropout(p=dropout_rate)
+
+        self.ffn = FeedForward(size, linear_size, block_id, ffn_dropout_rate, num_blocks)
+        self.dropout_ffn = torch.nn.Dropout(p=dropout_rate)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        state: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Compute receptance weighted key value.
+
+        Args:
+            x: RWKV input sequences. (B, L, size)
+            state: Decoder hidden states. [5 x (B, D_att/size, N)]
+
+        Returns:
+            x: RWKV output sequences. (B, L, size)
+            x: Decoder hidden states. [5 x (B, D_att/size, N)]
+
+        """
+        att, state = self.att(self.layer_norm_att(x), state=state)
+        x = x + self.dropout_att(att)
+
+        ffn, state = self.ffn(self.layer_norm_ffn(x), state=state)
+        x = x + self.dropout_ffn(ffn)
+
+        return x, state

funasr/models/rwkv_bat/rwkv_attention.py

@@ -0,0 +1,608 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import math
+import torch
+from pathlib import Path
+from importlib.util import find_spec
+from typing import List, Optional, Tuple, Union
+
+
+wkv_kernel_encoder = None
+wkv_kernel_decoder = None
+
+
+class WKVLinearAttentionEncoder(torch.autograd.Function):
+    """WKVLinearAttention function definition."""
+
+    @staticmethod
+    def forward(
+        ctx,
+        time_decay: torch.Tensor,
+        time_first: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.tensor,
+    ) -> torch.Tensor:
+        """WKVLinearAttention function forward pass.
+
+        Args:
+            time_decay: Channel-wise time decay vector. (D_att)
+            time_first: Channel-wise time first vector. (D_att)
+            key: Key tensor. (B, U, D_att)
+            value: Value tensor. (B, U, D_att)
+
+        Returns:
+            out: Weighted Key-Value tensor. (B, U, D_att)
+
+        """
+        batch, length, dim = key.size()
+
+        assert length <= wkv_kernel_encoder.context_size, (
+            f"Cannot process key of length {length} while context_size "
+            f"is ({wkv_kernel_encoder.context_size}). Limit should be increased."
+        )
+
+        assert batch * dim % min(dim, 32) == 0, (
+            f"batch size ({batch}) by dimension ({dim}) should be a multiple of " f"{min(dim, 32)}"
+        )
+
+        ctx.input_dtype = key.dtype
+
+        time_decay = -torch.exp(time_decay.float().contiguous())
+        time_first = time_first.float().contiguous()
+
+        key = key.float().contiguous()
+        value = value.float().contiguous()
+
+        out = torch.empty_like(key, memory_format=torch.contiguous_format)
+
+        wkv_kernel_encoder.forward(time_decay, time_first, key, value, out)
+        ctx.save_for_backward(time_decay, time_first, key, value, out)
+
+        return out
+
+    @staticmethod
+    def backward(
+        ctx, grad_output: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """WKVLinearAttention function backward pass.
+
+        Args:
+            grad_output: Output gradient. (B, U, D_att)
+
+        Returns:
+            grad_time_decay: Gradient for channel-wise time decay vector. (D_att)
+            grad_time_first: Gradient for channel-wise time first vector. (D_att)
+            grad_key: Gradient for key tensor. (B, U, D_att)
+            grad_value: Gradient for value tensor. (B, U, D_att)
+
+        """
+        time_decay, time_first, key, value, output = ctx.saved_tensors
+        grad_dtype = ctx.input_dtype
+
+        batch, _, dim = key.size()
+
+        grad_time_decay = torch.empty(
+            (batch, dim),
+            memory_format=torch.contiguous_format,
+            dtype=time_decay.dtype,
+            device=time_decay.device,
+        )
+
+        grad_time_first = torch.empty(
+            (batch, dim),
+            memory_format=torch.contiguous_format,
+            dtype=time_decay.dtype,
+            device=time_decay.device,
+        )
+
+        grad_key = torch.empty_like(key, memory_format=torch.contiguous_format)
+        grad_value = torch.empty_like(value, memory_format=torch.contiguous_format)
+
+        wkv_kernel_encoder.backward(
+            time_decay,
+            time_first,
+            key,
+            value,
+            output,
+            grad_output.contiguous(),
+            grad_time_decay,
+            grad_time_first,
+            grad_key,
+            grad_value,
+        )
+
+        grad_time_decay = torch.sum(grad_time_decay, dim=0)
+        grad_time_first = torch.sum(grad_time_first, dim=0)
+
+        return (
+            grad_time_decay,
+            grad_time_first,
+            grad_key,
+            grad_value,
+        )
+
+
+class WKVLinearAttentionDecoder(torch.autograd.Function):
+    """WKVLinearAttention function definition."""
+
+    @staticmethod
+    def forward(
+        ctx,
+        time_decay: torch.Tensor,
+        time_first: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.tensor,
+    ) -> torch.Tensor:
+        """WKVLinearAttention function forward pass.
+
+        Args:
+            time_decay: Channel-wise time decay vector. (D_att)
+            time_first: Channel-wise time first vector. (D_att)
+            key: Key tensor. (B, U, D_att)
+            value: Value tensor. (B, U, D_att)
+
+        Returns:
+            out: Weighted Key-Value tensor. (B, U, D_att)
+
+        """
+        batch, length, dim = key.size()
+
+        assert length <= wkv_kernel_decoder.context_size, (
+            f"Cannot process key of length {length} while context_size "
+            f"is ({wkv_kernel.context_size}). Limit should be increased."
+        )
+
+        assert batch * dim % min(dim, 32) == 0, (
+            f"batch size ({batch}) by dimension ({dim}) should be a multiple of " f"{min(dim, 32)}"
+        )
+
+        ctx.input_dtype = key.dtype
+
+        time_decay = -torch.exp(time_decay.float().contiguous())
+        time_first = time_first.float().contiguous()
+
+        key = key.float().contiguous()
+        value = value.float().contiguous()
+
+        out = torch.empty_like(key, memory_format=torch.contiguous_format)
+
+        wkv_kernel_decoder.forward(time_decay, time_first, key, value, out)
+        ctx.save_for_backward(time_decay, time_first, key, value, out)
+
+        return out
+
+    @staticmethod
+    def backward(
+        ctx, grad_output: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """WKVLinearAttention function backward pass.
+
+        Args:
+            grad_output: Output gradient. (B, U, D_att)
+
+        Returns:
+            grad_time_decay: Gradient for channel-wise time decay vector. (D_att)
+            grad_time_first: Gradient for channel-wise time first vector. (D_att)
+            grad_key: Gradient for key tensor. (B, U, D_att)
+            grad_value: Gradient for value tensor. (B, U, D_att)
+
+        """
+        time_decay, time_first, key, value, output = ctx.saved_tensors
+        grad_dtype = ctx.input_dtype
+
+        batch, _, dim = key.size()
+
+        grad_time_decay = torch.empty(
+            (batch, dim),
+            memory_format=torch.contiguous_format,
+            dtype=time_decay.dtype,
+            device=time_decay.device,
+        )
+
+        grad_time_first = torch.empty(
+            (batch, dim),
+            memory_format=torch.contiguous_format,
+            dtype=time_decay.dtype,
+            device=time_decay.device,
+        )
+
+        grad_key = torch.empty_like(key, memory_format=torch.contiguous_format)
+        grad_value = torch.empty_like(value, memory_format=torch.contiguous_format)
+
+        wkv_kernel_decoder.backward(
+            time_decay,
+            time_first,
+            key,
+            value,
+            output,
+            grad_output.contiguous(),
+            grad_time_decay,
+            grad_time_first,
+            grad_key,
+            grad_value,
+        )
+
+        grad_time_decay = torch.sum(grad_time_decay, dim=0)
+        grad_time_first = torch.sum(grad_time_first, dim=0)
+
+        return (
+            grad_time_decay,
+            grad_time_first,
+            grad_key,
+            grad_value,
+        )
+
+
+def load_encoder_wkv_kernel(context_size: int) -> None:
+    """Load WKV CUDA kernel.
+
+    Args:
+        context_size: Context size.
+
+    """
+    from torch.utils.cpp_extension import load
+
+    global wkv_kernel_encoder
+
+    if wkv_kernel_encoder is not None and wkv_kernel_encoder.context_size == context_size:
+        return
+
+    if find_spec("ninja") is None:
+        raise ImportError(
+            "Ninja package was not found. WKV kernel module can't be loaded "
+            "for training. Please, 'pip install ninja' in your environment."
+        )
+
+    if not torch.cuda.is_available():
+        raise ImportError(
+            "CUDA is currently a requirement for WKV kernel loading. "
+            "Please set your devices properly and launch again."
+        )
+
+    kernel_folder = Path(__file__).resolve().parent / "cuda_encoder"
+    kernel_files = [kernel_folder / f for f in ["wkv_op.cpp", "wkv_cuda.cu"]]
+
+    kernel_cflags = [
+        "-res-usage",
+        "--maxrregcount 60",
+        "--use_fast_math",
+        "-O3",
+        "-Xptxas -O3",
+        f"-DTmax={context_size}",
+    ]
+    wkv_kernel_encoder = load(
+        name=f"encoder_wkv_{context_size}",
+        sources=kernel_files,
+        verbose=True,
+        extra_cuda_cflags=kernel_cflags,
+    )
+    wkv_kernel_encoder.context_size = context_size
+
+
+def load_decoder_wkv_kernel(context_size: int) -> None:
+    """Load WKV CUDA kernel.
+
+    Args:
+        context_size: Context size.
+
+    """
+    from torch.utils.cpp_extension import load
+
+    global wkv_kernel_decoder
+
+    if wkv_kernel_decoder is not None and wkv_kernel_decoder.context_size == context_size:
+        return
+
+    if find_spec("ninja") is None:
+        raise ImportError(
+            "Ninja package was not found. WKV kernel module can't be loaded "
+            "for training. Please, 'pip install ninja' in your environment."
+        )
+
+    if not torch.cuda.is_available():
+        raise ImportError(
+            "CUDA is currently a requirement for WKV kernel loading. "
+            "Please set your devices properly and launch again."
+        )
+
+    kernel_folder = Path(__file__).resolve().parent / "cuda_decoder"
+    kernel_files = [kernel_folder / f for f in ["wkv_op.cpp", "wkv_cuda.cu"]]
+
+    kernel_cflags = [
+        "-res-usage",
+        "--maxrregcount 60",
+        "--use_fast_math",
+        "-O3",
+        "-Xptxas -O3",
+        f"-DTmax={context_size}",
+    ]
+    wkv_kernel_decoder = load(
+        name=f"decoder_wkv_{context_size}",
+        sources=kernel_files,
+        verbose=True,
+        extra_cuda_cflags=kernel_cflags,
+    )
+    wkv_kernel_decoder.context_size = context_size
+
+
+class SelfAttention(torch.nn.Module):
+    """SelfAttention module definition.
+
+    Args:
+        size: Input/Output size.
+        attention_size: Attention hidden size.
+        context_size: Context size for WKV kernel.
+        block_id: Block index.
+        num_blocks: Number of blocks in the architecture.
+
+    """
+
+    def __init__(
+        self,
+        size: int,
+        attention_size: int,
+        block_id: int,
+        dropout_rate: float,
+        num_blocks: int,
+    ) -> None:
+        """Construct a SelfAttention object."""
+        super().__init__()
+        self.time_shift = torch.nn.ZeroPad2d((0, 0, 1, -1))
+
+        self.time_decay = torch.nn.Parameter(torch.empty(attention_size))
+        self.time_first = torch.nn.Parameter(torch.empty(attention_size))
+
+        self.time_mix_key = torch.nn.Parameter(torch.empty(1, 1, size))
+        self.time_mix_value = torch.nn.Parameter(torch.empty(1, 1, size))
+        self.time_mix_receptance = torch.nn.Parameter(torch.empty(1, 1, size))
+
+        self.proj_key = torch.nn.Linear(size, attention_size, bias=True)
+        self.proj_value = torch.nn.Linear(size, attention_size, bias=True)
+        self.proj_receptance = torch.nn.Linear(size, attention_size, bias=True)
+
+        self.proj_output = torch.nn.Linear(attention_size, size, bias=True)
+
+        self.block_id = block_id
+
+        self.reset_parameters(size, attention_size, block_id, num_blocks)
+        self.dropout = torch.nn.Dropout(p=dropout_rate)
+
+    def reset_parameters(
+        self, size: int, attention_size: int, block_id: int, num_blocks: int
+    ) -> None:
+        """Reset module parameters.
+
+        Args:
+            size: Block size.
+            attention_size: Attention hidden size.
+            block_id: Block index.
+            num_blocks: Number of blocks in the architecture.
+
+        """
+        ratio_0_to_1 = block_id / (num_blocks - 1)
+        ratio_1_to_almost0 = 1.0 - (block_id / num_blocks)
+
+        time_weight = torch.ones(1, 1, size)
+
+        for i in range(size):
+            time_weight[0, 0, i] = i / size
+
+        decay_speed = [
+            -5 + 8 * (h / (attention_size - 1)) ** (0.7 + 1.3 * ratio_0_to_1)
+            for h in range(attention_size)
+        ]
+        decay_speed = torch.tensor(
+            decay_speed, dtype=self.time_decay.dtype, device=self.time_decay.device
+        )
+
+        zigzag = (
+            torch.tensor(
+                [(i + 1) % 3 - 1 for i in range(attention_size)],
+                dtype=self.time_first.dtype,
+                device=self.time_first.device,
+            )
+            * 0.5
+        )
+
+        with torch.no_grad():
+            self.time_decay.data = decay_speed
+            self.time_first.data = torch.ones_like(self.time_first * math.log(0.3) + zigzag)
+
+            self.time_mix_key.data = torch.pow(time_weight, ratio_1_to_almost0)
+            self.time_mix_value.data = (
+                torch.pow(time_weight, ratio_1_to_almost0) + 0.3 * ratio_0_to_1
+            )
+            self.time_mix_receptance.data = torch.pow(time_weight, 0.5 * ratio_1_to_almost0)
+
+    @torch.no_grad()
+    def wkv_linear_attention(
+        self,
+        time_decay: torch.Tensor,
+        time_first: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        state: Tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
+        """Compute WKV with state (i.e.: for inference).
+
+        Args:
+            time_decay: Channel-wise time decay vector. (D_att)
+            time_first: Channel-wise time first vector. (D_att)
+            key: Key tensor. (B, 1, D_att)
+            value: Value tensor. (B, 1, D_att)
+            state: Decoder hidden states. [3 x (B, D_att)]
+
+        Returns:
+            output: Weighted Key-Value. (B, 1, D_att)
+            state: Decoder hidden states. [3 x (B, 1, D_att)]
+
+        """
+        num_state, den_state, max_state = state
+        time_decay = -torch.exp(time_decay)
+        max_for_output = torch.maximum(max_state, (time_first + key))
+
+        e1 = torch.exp(max_state - max_for_output)
+        e2 = torch.exp((time_first + key) - max_for_output)
+
+        numerator = e1 * num_state + e2 * value
+        denominator = e1 * den_state + e2
+
+        max_for_state = torch.maximum(key, (max_state + time_decay))
+
+        e1 = torch.exp((max_state + time_decay) - max_for_state)
+        e2 = torch.exp(key - max_for_state)
+
+        wkv = numerator / denominator
+
+        state = [e1 * num_state + e2 * value, e1 * den_state + e2, max_for_state]
+
+        return wkv, state
+
+
+class DecoderSelfAttention(SelfAttention):
+    """SelfAttention module definition.
+
+    Args:
+        size: Input/Output size.
+        attention_size: Attention hidden size.
+        context_size: Context size for WKV kernel.
+        block_id: Block index.
+        num_blocks: Number of blocks in the architecture.
+
+    """
+
+    def __init__(
+        self,
+        size: int,
+        attention_size: int,
+        context_size: int,
+        block_id: int,
+        dropout_rate: float,
+        num_blocks: int,
+    ) -> None:
+        """Construct a SelfAttention object."""
+        super().__init__(size, attention_size, block_id, dropout_rate, num_blocks)
+        # load_decoder_wkv_kernel(context_size)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        state: Optional[List[torch.Tensor]] = None,
+    ) -> Tuple[torch.Tensor, Optional[List[torch.Tensor]]]:
+        """Compute time mixing.
+
+        Args:
+            x: SelfAttention input sequences. (B, U, size)
+            state: Decoder hidden states. [5 x (B, 1, D_att, N)]
+
+        Returns:
+            x: SelfAttention output sequences. (B, U, size)
+
+        """
+        shifted_x = self.time_shift(x) if state is None else state[1][..., self.block_id]
+
+        key = x * self.time_mix_key + shifted_x * (1 - self.time_mix_key)
+        value = x * self.time_mix_value + shifted_x * (1 - self.time_mix_value)
+        receptance = x * self.time_mix_receptance + shifted_x * (1 - self.time_mix_receptance)
+
+        key = self.proj_key(key)
+        value = self.proj_value(value)
+        receptance = torch.sigmoid(self.proj_receptance(receptance))
+
+        if state is not None:
+            state[1][..., self.block_id] = x
+
+            wkv, att_state = self.wkv_linear_attention(
+                self.time_decay,
+                self.time_first,
+                key,
+                value,
+                tuple(s[..., self.block_id] for s in state[2:]),
+            )
+
+            state[2][..., self.block_id] = att_state[0]
+            state[3][..., self.block_id] = att_state[1]
+            state[4][..., self.block_id] = att_state[2]
+        else:
+            wkv = WKVLinearAttentionDecoder.apply(self.time_decay, self.time_first, key, value)
+
+        wkv = self.dropout(wkv)
+        x = self.proj_output(receptance * wkv)
+
+        return x, state
+
+
+class EncoderSelfAttention(SelfAttention):
+    """SelfAttention module definition.
+
+    Args:
+        size: Input/Output size.
+        attention_size: Attention hidden size.
+        context_size: Context size for WKV kernel.
+        block_id: Block index.
+        num_blocks: Number of blocks in the architecture.
+
+    """
+
+    def __init__(
+        self,
+        size: int,
+        attention_size: int,
+        context_size: int,
+        block_id: int,
+        dropout_rate: float,
+        num_blocks: int,
+    ) -> None:
+        """Construct a SelfAttention object."""
+        super().__init__(size, attention_size, block_id, dropout_rate, num_blocks)
+        # load_encoder_wkv_kernel(context_size)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        state: Optional[List[torch.Tensor]] = None,
+    ) -> Tuple[torch.Tensor, Optional[List[torch.Tensor]]]:
+        """Compute time mixing.
+
+        Args:
+            x: SelfAttention input sequences. (B, U, size)
+            state: Decoder hidden states. [5 x (B, 1, D_att, N)]
+
+        Returns:
+            x: SelfAttention output sequences. (B, U, size)
+
+        """
+        shifted_x = self.time_shift(x) if state is None else state[1][..., self.block_id]
+
+        key = x * self.time_mix_key + shifted_x * (1 - self.time_mix_key)
+        value = x * self.time_mix_value + shifted_x * (1 - self.time_mix_value)
+        receptance = x * self.time_mix_receptance + shifted_x * (1 - self.time_mix_receptance)
+
+        key = self.proj_key(key)
+        value = self.proj_value(value)
+        receptance = torch.sigmoid(self.proj_receptance(receptance))
+
+        if state is not None:
+            state[1][..., self.block_id] = x
+
+            wkv, att_state = self.wkv_linear_attention(
+                self.time_decay,
+                self.time_first,
+                key,
+                value,
+                tuple(s[..., self.block_id] for s in state[2:]),
+            )
+
+            state[2][..., self.block_id] = att_state[0]
+            state[3][..., self.block_id] = att_state[1]
+            state[4][..., self.block_id] = att_state[2]
+        else:
+            wkv = WKVLinearAttentionEncoder.apply(self.time_decay, self.time_first, key, value)
+
+        wkv = self.dropout(wkv)
+        x = self.proj_output(receptance * wkv)
+
+        return x, state

funasr/models/rwkv_bat/rwkv_encoder.py

@@ -0,0 +1,158 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import torch
+from typing import Dict, List, Optional, Tuple
+
+from funasr.register import tables
+from funasr.models.rwkv_bat.rwkv import RWKV
+from funasr.models.transformer.layer_norm import LayerNorm
+from funasr.models.transformer.utils.nets_utils import make_source_mask
+from funasr.models.rwkv_bat.rwkv_subsampling import RWKVConvInput
+
+
+@tables.register("encoder_classes", "RWKVEncoder")
+class RWKVEncoder(torch.nn.Module):
+    """RWKV encoder module.
+
+    Based on https://arxiv.org/pdf/2305.13048.pdf.
+
+    Args:
+        vocab_size: Vocabulary size.
+        output_size: Input/Output size.
+        context_size: Context size for WKV computation.
+        linear_size: FeedForward hidden size.
+        attention_size: SelfAttention hidden size.
+        normalization_type: Normalization layer type.
+        normalization_args: Normalization layer arguments.
+        num_blocks: Number of RWKV blocks.
+        embed_dropout_rate: Dropout rate for embedding layer.
+        att_dropout_rate: Dropout rate for the attention module.
+        ffn_dropout_rate: Dropout rate for the feed-forward module.
+    """
+
+    def __init__(
+        self,
+        input_size: int,
+        output_size: int = 512,
+        context_size: int = 1024,
+        linear_size: Optional[int] = None,
+        attention_size: Optional[int] = None,
+        num_blocks: int = 4,
+        att_dropout_rate: float = 0.0,
+        ffn_dropout_rate: float = 0.0,
+        dropout_rate: float = 0.0,
+        subsampling_factor: int = 4,
+        time_reduction_factor: int = 1,
+        kernel: int = 3,
+        **kwargs,
+    ) -> None:
+        """Construct a RWKVEncoder object."""
+        super().__init__()
+
+        self.embed = RWKVConvInput(
+            input_size,
+            [output_size // 4, output_size // 2, output_size],
+            subsampling_factor,
+            conv_kernel_size=kernel,
+            output_size=output_size,
+        )
+
+        self.subsampling_factor = subsampling_factor
+
+        linear_size = output_size * 4 if linear_size is None else linear_size
+        attention_size = output_size if attention_size is None else attention_size
+
+        self.rwkv_blocks = torch.nn.ModuleList(
+            [
+                RWKV(
+                    output_size,
+                    linear_size,
+                    attention_size,
+                    context_size,
+                    block_id,
+                    num_blocks,
+                    att_dropout_rate=att_dropout_rate,
+                    ffn_dropout_rate=ffn_dropout_rate,
+                    dropout_rate=dropout_rate,
+                )
+                for block_id in range(num_blocks)
+            ]
+        )
+
+        self.embed_norm = LayerNorm(output_size)
+        self.final_norm = LayerNorm(output_size)
+
+        self._output_size = output_size
+        self.context_size = context_size
+
+        self.num_blocks = num_blocks
+        self.time_reduction_factor = time_reduction_factor
+
+    def output_size(self) -> int:
+        return self._output_size
+
+    def forward(self, x: torch.Tensor, x_len) -> torch.Tensor:
+        """Encode source label sequences.
+
+        Args:
+            x: Encoder input sequences. (B, L)
+
+        Returns:
+            out: Encoder output sequences. (B, U, D)
+
+        """
+        _, length, _ = x.size()
+
+        assert (
+            length <= self.context_size * self.subsampling_factor
+        ), "Context size is too short for current length: %d versus %d" % (
+            length,
+            self.context_size * self.subsampling_factor,
+        )
+        mask = make_source_mask(x_len).to(x.device)
+        x, mask = self.embed(x, mask, None)
+        x = self.embed_norm(x)
+        olens = mask.eq(0).sum(1)
+
+        if self.training:
+            for block in self.rwkv_blocks:
+                x, _ = block(x)
+        else:
+            x = self.rwkv_infer(x)
+
+        x = self.final_norm(x)
+
+        if self.time_reduction_factor > 1:
+            x = x[:, :: self.time_reduction_factor, :]
+            olens = torch.floor_divide(olens - 1, self.time_reduction_factor) + 1
+
+        return x, olens, None
+
+    def rwkv_infer(self, xs_pad):
+
+        batch_size = xs_pad.shape[0]
+
+        hidden_sizes = [self._output_size for i in range(5)]
+
+        state = [
+            torch.zeros(
+                (batch_size, 1, hidden_sizes[i], self.num_blocks),
+                dtype=torch.float32,
+                device=xs_pad.device,
+            )
+            for i in range(5)
+        ]
+
+        state[4] -= 1e-30
+
+        xs_out = []
+        for t in range(xs_pad.shape[1]):
+            x_t = xs_pad[:, t, :]
+            for idx, block in enumerate(self.rwkv_blocks):
+                x_t, state = block(x_t, state=state)
+            xs_out.append(x_t)
+        xs_out = torch.cat(xs_out, dim=1)
+        return xs_out

funasr/models/rwkv_bat/rwkv_feed_forward.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import torch
+from typing import List, Optional, Tuple
+
+
+class FeedForward(torch.nn.Module):
+    """FeedForward module definition.
+
+    Args:
+        size: Input/Output size.
+        hidden_size: Hidden size.
+        block_id: Block index.
+        num_blocks: Number of blocks in the architecture.
+
+    """
+
+    def __init__(
+        self, size: int, hidden_size: int, block_id: int, dropout_rate: float, num_blocks: int
+    ) -> None:
+        """Construct a FeedForward object."""
+        super().__init__()
+
+        self.time_shift = torch.nn.ZeroPad2d((0, 0, 1, -1))
+
+        self.time_mix_key = torch.nn.Parameter(torch.empty(1, 1, size))
+        self.time_mix_receptance = torch.nn.Parameter(torch.empty(1, 1, size))
+
+        self.proj_key = torch.nn.Linear(size, hidden_size, bias=True)
+        self.proj_value = torch.nn.Linear(hidden_size, size, bias=True)
+        self.proj_receptance = torch.nn.Linear(size, size, bias=True)
+
+        self.block_id = block_id
+
+        self.reset_parameters(size, block_id, num_blocks)
+        self.dropout = torch.nn.Dropout(p=dropout_rate)
+
+    def reset_parameters(self, size: int, block_id: int, num_blocks: int) -> None:
+        """Reset module parameters.
+
+        Args:
+            size: Block size.
+            block_id: Block index.
+            num_blocks: Number of blocks in the architecture.
+
+        """
+        ratio_1_to_almost0 = 1.0 - (block_id / num_blocks)
+
+        time_weight = torch.ones(1, 1, size)
+
+        for i in range(size):
+            time_weight[0, 0, i] = i / size
+
+        with torch.no_grad():
+            self.time_mix_key.data = torch.pow(time_weight, ratio_1_to_almost0)
+            self.time_mix_receptance.data = torch.pow(time_weight, ratio_1_to_almost0)
+
+    def forward(
+        self, x: torch.Tensor, state: Optional[List[torch.Tensor]] = None
+    ) -> Tuple[torch.Tensor, Optional[List[torch.Tensor]]]:
+        """Compute channel mixing.
+
+        Args:
+            x: FeedForward input sequences. (B, U, size)
+            state: Decoder hidden state. [5 x (B, 1, size, N)]
+
+        Returns:
+            x: FeedForward output sequences. (B, U, size)
+            state: Decoder hidden state. [5 x (B, 1, size, N)]
+
+        """
+        shifted_x = self.time_shift(x) if state is None else state[0][..., self.block_id]
+
+        key = x * self.time_mix_key + shifted_x * (1 - self.time_mix_key)
+        receptance = x * self.time_mix_receptance + shifted_x * (1 - self.time_mix_receptance)
+
+        key = torch.square(torch.relu(self.proj_key(key)))
+        value = self.proj_value(self.dropout(key))
+        receptance = torch.sigmoid(self.proj_receptance(receptance))
+
+        if state is not None:
+            state[0][..., self.block_id] = x
+
+        x = receptance * value
+
+        return x, state

funasr/models/rwkv_bat/rwkv_subsampling.py

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

funasr/models/rwkv_bat/template.yaml

@@ -0,0 +1,64 @@
+# network architecture
+model: Transducer
+model_conf:
+    auxiliary_ctc_weight: 0.0
+
+# encoder
+encoder: RWKVEncoder
+encoder_conf:
+    kernel: 3
+    subsampling_factor: 4
+    output_size: 512
+    num_blocks: 18
+    time_reduction_factor: 2
+    att_dropout_rate: 0.1
+    ffn_dropout_rate: 0.1
+    dropout_rate: 0.1
+
+# decoder (prediction network)
+decoder: rnnt_decoder
+decoder_conf:
+    embed_size: 512
+    hidden_size: 512
+    embed_dropout_rate: 0.1
+    dropout_rate: 0.1
+    use_embed_mask: false
+
+# joint network
+joint_network: joint_network
+joint_network_conf:
+    joint_space_size: 512
+
+frontend: WavFrontend
+frontend_conf:
+    fs: 16000
+    window: hamming
+    n_mels: 80
+    frame_length: 25
+    frame_shift: 10
+    lfr_m: 1
+    lfr_n: 1
+    upsacle_samples: true
+
+specaug: SpecAugLFR
+specaug_conf:
+    apply_time_warp: false
+    time_warp_window: 5
+    time_warp_mode: bicubic
+    apply_freq_mask: true
+    freq_mask_width_range:
+    - 0
+    - 30
+    lfr_rate: 6
+    num_freq_mask: 1
+    apply_time_mask: true
+    time_mask_width_range:
+    - 0
+    - 12
+    num_time_mask: 1
+
+tokenizer: CharTokenizer
+tokenizer_conf:
+  unk_symbol: <unk>
+  split_with_space: true
+  

funasr/models/sa_asr/attention.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+
+"""Multi-Head Attention layer definition."""
+
+import math
+
+import numpy
+import torch
+from torch import nn
+from typing import Optional, Tuple
+
+import torch.nn.functional as F
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+import funasr.models.lora.layers as lora
+
+
+class CosineDistanceAttention(nn.Module):
+    """Compute Cosine Distance between spk decoder output and speaker profile
+    Args:
+        profile_path: speaker profile file path (.npy file)
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, spk_decoder_out, profile, profile_lens=None):
+        """
+        Args:
+            spk_decoder_out(torch.Tensor):(B, L, D)
+            spk_profiles(torch.Tensor):(B, N, D)
+        """
+        x = spk_decoder_out.unsqueeze(2)  # (B, L, 1, D)
+        if profile_lens is not None:
+
+            mask = (make_pad_mask(profile_lens)[:, None, :]).to(profile.device)
+            min_value = float(numpy.finfo(torch.tensor(0, dtype=x.dtype).numpy().dtype).min)
+            weights_not_softmax = F.cosine_similarity(x, profile.unsqueeze(1), dim=-1).masked_fill(
+                mask, min_value
+            )
+            weights = self.softmax(weights_not_softmax).masked_fill(mask, 0.0)  # (B, L, N)
+        else:
+            x = x[:, -1:, :, :]
+            weights_not_softmax = F.cosine_similarity(x, profile.unsqueeze(1).to(x.device), dim=-1)
+            weights = self.softmax(weights_not_softmax)  # (B, 1, N)
+        spk_embedding = torch.matmul(weights, profile.to(weights.device))  # (B, L, D)
+
+        return spk_embedding, weights

funasr/models/sa_asr/beam_search_sa_asr.py

@@ -0,0 +1,540 @@
+"""Beam search module."""
+
+from itertools import chain
+import logging
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import NamedTuple
+from typing import Tuple
+from typing import Union
+
+import torch
+
+from funasr.metrics import end_detect
+from funasr.models.transformer.scorers.scorer_interface import PartialScorerInterface
+from funasr.models.transformer.scorers.scorer_interface import ScorerInterface
+from funasr.models.decoder.abs_decoder import AbsDecoder
+
+
+class Hypothesis(NamedTuple):
+    """Hypothesis data type."""
+
+    yseq: torch.Tensor
+    spk_weigths: List
+    score: Union[float, torch.Tensor] = 0
+    scores: Dict[str, Union[float, torch.Tensor]] = dict()
+    states: Dict[str, Any] = dict()
+
+    def asdict(self) -> dict:
+        """Convert data to JSON-friendly dict."""
+        return self._replace(
+            yseq=self.yseq.tolist(),
+            score=float(self.score),
+            scores={k: float(v) for k, v in self.scores.items()},
+        )._asdict()
+
+
+class BeamSearch(torch.nn.Module):
+    """Beam search implementation."""
+
+    def __init__(
+        self,
+        scorers: Dict[str, ScorerInterface],
+        weights: Dict[str, float],
+        beam_size: int,
+        vocab_size: int,
+        sos: int,
+        eos: int,
+        token_list: List[str] = None,
+        pre_beam_ratio: float = 1.5,
+        pre_beam_score_key: str = None,
+    ):
+        """Initialize beam search.
+
+        Args:
+            scorers (dict[str, ScorerInterface]): Dict of decoder modules
+                e.g., Decoder, CTCPrefixScorer, LM
+                The scorer will be ignored if it is `None`
+            weights (dict[str, float]): Dict of weights for each scorers
+                The scorer will be ignored if its weight is 0
+            beam_size (int): The number of hypotheses kept during search
+            vocab_size (int): The number of vocabulary
+            sos (int): Start of sequence id
+            eos (int): End of sequence id
+            token_list (list[str]): List of tokens for debug log
+            pre_beam_score_key (str): key of scores to perform pre-beam search
+            pre_beam_ratio (float): beam size in the pre-beam search
+                will be `int(pre_beam_ratio * beam_size)`
+
+        """
+        super().__init__()
+        # set scorers
+        self.weights = weights
+        self.scorers = dict()
+        self.full_scorers = dict()
+        self.part_scorers = dict()
+        # this module dict is required for recursive cast
+        # `self.to(device, dtype)` in `recog.py`
+        self.nn_dict = torch.nn.ModuleDict()
+        for k, v in scorers.items():
+            w = weights.get(k, 0)
+            if w == 0 or v is None:
+                continue
+            assert isinstance(
+                v, ScorerInterface
+            ), f"{k} ({type(v)}) does not implement ScorerInterface"
+            self.scorers[k] = v
+            if isinstance(v, PartialScorerInterface):
+                self.part_scorers[k] = v
+            else:
+                self.full_scorers[k] = v
+            if isinstance(v, torch.nn.Module):
+                self.nn_dict[k] = v
+
+        # set configurations
+        self.sos = sos
+        self.eos = eos
+        self.token_list = token_list
+        self.pre_beam_size = int(pre_beam_ratio * beam_size)
+        self.beam_size = beam_size
+        self.n_vocab = vocab_size
+        if (
+            pre_beam_score_key is not None
+            and pre_beam_score_key != "full"
+            and pre_beam_score_key not in self.full_scorers
+        ):
+            raise KeyError(f"{pre_beam_score_key} is not found in {self.full_scorers}")
+        self.pre_beam_score_key = pre_beam_score_key
+        self.do_pre_beam = (
+            self.pre_beam_score_key is not None
+            and self.pre_beam_size < self.n_vocab
+            and len(self.part_scorers) > 0
+        )
+
+    def init_hyp(self, x: torch.Tensor) -> List[Hypothesis]:
+        """Get an initial hypothesis data.
+
+        Args:
+            x (torch.Tensor): The encoder output feature
+
+        Returns:
+            Hypothesis: The initial hypothesis.
+
+        """
+        init_states = dict()
+        init_scores = dict()
+        for k, d in self.scorers.items():
+            init_states[k] = d.init_state(x)
+            init_scores[k] = 0.0
+        return [
+            Hypothesis(
+                score=0.0,
+                scores=init_scores,
+                states=init_states,
+                yseq=torch.tensor([self.sos], device=x.device),
+                spk_weigths=[],
+            )
+        ]
+
+    @staticmethod
+    def append_token(xs: torch.Tensor, x: int) -> torch.Tensor:
+        """Append new token to prefix tokens.
+
+        Args:
+            xs (torch.Tensor): The prefix token
+            x (int): The new token to append
+
+        Returns:
+            torch.Tensor: New tensor contains: xs + [x] with xs.dtype and xs.device
+
+        """
+        x = torch.tensor([x], dtype=xs.dtype, device=xs.device)
+        return torch.cat((xs, x))
+
+    def score_full(
+        self,
+        hyp: Hypothesis,
+        asr_enc: torch.Tensor,
+        spk_enc: torch.Tensor,
+        profile: torch.Tensor,
+    ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]:
+        """Score new hypothesis by `self.full_scorers`.
+
+        Args:
+            hyp (Hypothesis): Hypothesis with prefix tokens to score
+            x (torch.Tensor): Corresponding input feature
+
+        Returns:
+            Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of
+                score dict of `hyp` that has string keys of `self.full_scorers`
+                and tensor score values of shape: `(self.n_vocab,)`,
+                and state dict that has string keys
+                and state values of `self.full_scorers`
+
+        """
+        scores = dict()
+        states = dict()
+        for k, d in self.full_scorers.items():
+            if isinstance(d, AbsDecoder):
+                scores[k], spk_weigths, states[k] = d.score(
+                    hyp.yseq, hyp.states[k], asr_enc, spk_enc, profile
+                )
+            else:
+                scores[k], states[k] = d.score(hyp.yseq, hyp.states[k], asr_enc)
+        return scores, spk_weigths, states
+
+    def score_partial(
+        self,
+        hyp: Hypothesis,
+        ids: torch.Tensor,
+        asr_enc: torch.Tensor,
+        spk_enc: torch.Tensor,
+        profile: torch.Tensor,
+    ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]:
+        """Score new hypothesis by `self.part_scorers`.
+
+        Args:
+            hyp (Hypothesis): Hypothesis with prefix tokens to score
+            ids (torch.Tensor): 1D tensor of new partial tokens to score
+            x (torch.Tensor): Corresponding input feature
+
+        Returns:
+            Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of
+                score dict of `hyp` that has string keys of `self.part_scorers`
+                and tensor score values of shape: `(len(ids),)`,
+                and state dict that has string keys
+                and state values of `self.part_scorers`
+
+        """
+        scores = dict()
+        states = dict()
+        for k, d in self.part_scorers.items():
+            if isinstance(d, AbsDecoder):
+                scores[k], states[k] = d.score_partial(
+                    hyp.yseq, ids, hyp.states[k], asr_enc, spk_enc, profile
+                )
+            else:
+                scores[k], states[k] = d.score_partial(hyp.yseq, ids, hyp.states[k], asr_enc)
+        return scores, states
+
+    def beam(
+        self, weighted_scores: torch.Tensor, ids: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Compute topk full token ids and partial token ids.
+
+        Args:
+            weighted_scores (torch.Tensor): The weighted sum scores for each tokens.
+            Its shape is `(self.n_vocab,)`.
+            ids (torch.Tensor): The partial token ids to compute topk
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]:
+                The topk full token ids and partial token ids.
+                Their shapes are `(self.beam_size,)`
+
+        """
+        # no pre beam performed
+        if weighted_scores.size(0) == ids.size(0):
+            top_ids = weighted_scores.topk(self.beam_size)[1]
+            return top_ids, top_ids
+
+        # mask pruned in pre-beam not to select in topk
+        tmp = weighted_scores[ids]
+        weighted_scores[:] = -float("inf")
+        weighted_scores[ids] = tmp
+        top_ids = weighted_scores.topk(self.beam_size)[1]
+        local_ids = weighted_scores[ids].topk(self.beam_size)[1]
+        return top_ids, local_ids
+
+    @staticmethod
+    def merge_scores(
+        prev_scores: Dict[str, float],
+        next_full_scores: Dict[str, torch.Tensor],
+        full_idx: int,
+        next_part_scores: Dict[str, torch.Tensor],
+        part_idx: int,
+    ) -> Dict[str, torch.Tensor]:
+        """Merge scores for new hypothesis.
+
+        Args:
+            prev_scores (Dict[str, float]):
+                The previous hypothesis scores by `self.scorers`
+            next_full_scores (Dict[str, torch.Tensor]): scores by `self.full_scorers`
+            full_idx (int): The next token id for `next_full_scores`
+            next_part_scores (Dict[str, torch.Tensor]):
+                scores of partial tokens by `self.part_scorers`
+            part_idx (int): The new token id for `next_part_scores`
+
+        Returns:
+            Dict[str, torch.Tensor]: The new score dict.
+                Its keys are names of `self.full_scorers` and `self.part_scorers`.
+                Its values are scalar tensors by the scorers.
+
+        """
+        new_scores = dict()
+        for k, v in next_full_scores.items():
+            new_scores[k] = prev_scores[k] + v[full_idx]
+        for k, v in next_part_scores.items():
+            new_scores[k] = prev_scores[k] + v[part_idx]
+        return new_scores
+
+    def merge_states(self, states: Any, part_states: Any, part_idx: int) -> Any:
+        """Merge states for new hypothesis.
+
+        Args:
+            states: states of `self.full_scorers`
+            part_states: states of `self.part_scorers`
+            part_idx (int): The new token id for `part_scores`
+
+        Returns:
+            Dict[str, torch.Tensor]: The new score dict.
+                Its keys are names of `self.full_scorers` and `self.part_scorers`.
+                Its values are states of the scorers.
+
+        """
+        new_states = dict()
+        for k, v in states.items():
+            new_states[k] = v
+        for k, d in self.part_scorers.items():
+            new_states[k] = d.select_state(part_states[k], part_idx)
+        return new_states
+
+    def search(
+        self,
+        running_hyps: List[Hypothesis],
+        asr_enc: torch.Tensor,
+        spk_enc: torch.Tensor,
+        profile: torch.Tensor,
+    ) -> List[Hypothesis]:
+        """Search new tokens for running hypotheses and encoded speech x.
+
+        Args:
+            running_hyps (List[Hypothesis]): Running hypotheses on beam
+            x (torch.Tensor): Encoded speech feature (T, D)
+
+        Returns:
+            List[Hypotheses]: Best sorted hypotheses
+
+        """
+        # import ipdb;ipdb.set_trace()
+        best_hyps = []
+        part_ids = torch.arange(self.n_vocab, device=asr_enc.device)  # no pre-beam
+        for hyp in running_hyps:
+            # scoring
+            weighted_scores = torch.zeros(self.n_vocab, dtype=asr_enc.dtype, device=asr_enc.device)
+            scores, spk_weigths, states = self.score_full(hyp, asr_enc, spk_enc, profile)
+            for k in self.full_scorers:
+                weighted_scores += self.weights[k] * scores[k]
+            # partial scoring
+            if self.do_pre_beam:
+                pre_beam_scores = (
+                    weighted_scores
+                    if self.pre_beam_score_key == "full"
+                    else scores[self.pre_beam_score_key]
+                )
+                part_ids = torch.topk(pre_beam_scores, self.pre_beam_size)[1]
+            part_scores, part_states = self.score_partial(hyp, part_ids, asr_enc, spk_enc, profile)
+            for k in self.part_scorers:
+                weighted_scores[part_ids] += self.weights[k] * part_scores[k]
+            # add previous hyp score
+            weighted_scores += hyp.score
+
+            # update hyps
+            for j, part_j in zip(*self.beam(weighted_scores, part_ids)):
+                # will be (2 x beam at most)
+                best_hyps.append(
+                    Hypothesis(
+                        score=weighted_scores[j],
+                        yseq=self.append_token(hyp.yseq, j),
+                        scores=self.merge_scores(hyp.scores, scores, j, part_scores, part_j),
+                        states=self.merge_states(states, part_states, part_j),
+                        spk_weigths=hyp.spk_weigths + [spk_weigths],
+                    )
+                )
+
+            # sort and prune 2 x beam -> beam
+            best_hyps = sorted(best_hyps, key=lambda x: x.score, reverse=True)[
+                : min(len(best_hyps), self.beam_size)
+            ]
+        return best_hyps
+
+    def forward(
+        self,
+        asr_enc: torch.Tensor,
+        spk_enc: torch.Tensor,
+        profile: torch.Tensor,
+        maxlenratio: float = 0.0,
+        minlenratio: float = 0.0,
+    ) -> List[Hypothesis]:
+        """Perform beam search.
+
+        Args:
+            x (torch.Tensor): Encoded speech feature (T, D)
+            maxlenratio (float): Input length ratio to obtain max output length.
+                If maxlenratio=0.0 (default), it uses a end-detect function
+                to automatically find maximum hypothesis lengths
+            minlenratio (float): Input length ratio to obtain min output length.
+
+        Returns:
+            list[Hypothesis]: N-best decoding results
+
+        """
+        # import ipdb;ipdb.set_trace()
+        # set length bounds
+        if maxlenratio == 0:
+            maxlen = asr_enc.shape[0]
+        else:
+            maxlen = max(1, int(maxlenratio * asr_enc.size(0)))
+        minlen = int(minlenratio * asr_enc.size(0))
+        logging.info("decoder input length: " + str(asr_enc.shape[0]))
+        logging.info("max output length: " + str(maxlen))
+        logging.info("min output length: " + str(minlen))
+
+        # main loop of prefix search
+        running_hyps = self.init_hyp(asr_enc)
+        ended_hyps = []
+        for i in range(maxlen):
+            logging.debug("position " + str(i))
+            best = self.search(running_hyps, asr_enc, spk_enc, profile)
+            # import pdb;pdb.set_trace()
+            # post process of one iteration
+            running_hyps = self.post_process(i, maxlen, maxlenratio, best, ended_hyps)
+            # end detection
+            if maxlenratio == 0.0 and end_detect([h.asdict() for h in ended_hyps], i):
+                logging.info(f"end detected at {i}")
+                break
+            if len(running_hyps) == 0:
+                logging.info("no hypothesis. Finish decoding.")
+                break
+            else:
+                logging.debug(f"remained hypotheses: {len(running_hyps)}")
+
+        nbest_hyps = sorted(ended_hyps, key=lambda x: x.score, reverse=True)
+        # check the number of hypotheses reaching to eos
+        if len(nbest_hyps) == 0:
+            logging.warning(
+                "there is no N-best results, perform recognition " "again with smaller minlenratio."
+            )
+            return (
+                []
+                if minlenratio < 0.1
+                else self.forward(
+                    asr_enc, spk_enc, profile, maxlenratio, max(0.0, minlenratio - 0.1)
+                )
+            )
+
+        # report the best result
+        best = nbest_hyps[0]
+        for k, v in best.scores.items():
+            logging.info(f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}")
+        logging.info(f"total log probability: {best.score:.2f}")
+        logging.info(f"normalized log probability: {best.score / len(best.yseq):.2f}")
+        logging.info(f"total number of ended hypotheses: {len(nbest_hyps)}")
+        if self.token_list is not None:
+            logging.info(
+                "best hypo: " + "".join([self.token_list[x] for x in best.yseq[1:-1]]) + "\n"
+            )
+        return nbest_hyps
+
+    def post_process(
+        self,
+        i: int,
+        maxlen: int,
+        maxlenratio: float,
+        running_hyps: List[Hypothesis],
+        ended_hyps: List[Hypothesis],
+    ) -> List[Hypothesis]:
+        """Perform post-processing of beam search iterations.
+
+        Args:
+            i (int): The length of hypothesis tokens.
+            maxlen (int): The maximum length of tokens in beam search.
+            maxlenratio (int): The maximum length ratio in beam search.
+            running_hyps (List[Hypothesis]): The running hypotheses in beam search.
+            ended_hyps (List[Hypothesis]): The ended hypotheses in beam search.
+
+        Returns:
+            List[Hypothesis]: The new running hypotheses.
+
+        """
+        logging.debug(f"the number of running hypotheses: {len(running_hyps)}")
+        if self.token_list is not None:
+            logging.debug(
+                "best hypo: " + "".join([self.token_list[x] for x in running_hyps[0].yseq[1:]])
+            )
+        # add eos in the final loop to avoid that there are no ended hyps
+        if i == maxlen - 1:
+            logging.info("adding <eos> in the last position in the loop")
+            running_hyps = [
+                h._replace(yseq=self.append_token(h.yseq, self.eos)) for h in running_hyps
+            ]
+
+        # add ended hypotheses to a final list, and removed them from current hypotheses
+        # (this will be a problem, number of hyps < beam)
+        remained_hyps = []
+        for hyp in running_hyps:
+            if hyp.yseq[-1] == self.eos:
+                # e.g., Word LM needs to add final <eos> score
+                for k, d in chain(self.full_scorers.items(), self.part_scorers.items()):
+                    s = d.final_score(hyp.states[k])
+                    hyp.scores[k] += s
+                    hyp = hyp._replace(score=hyp.score + self.weights[k] * s)
+                ended_hyps.append(hyp)
+            else:
+                remained_hyps.append(hyp)
+        return remained_hyps
+
+
+def beam_search(
+    x: torch.Tensor,
+    sos: int,
+    eos: int,
+    beam_size: int,
+    vocab_size: int,
+    scorers: Dict[str, ScorerInterface],
+    weights: Dict[str, float],
+    token_list: List[str] = None,
+    maxlenratio: float = 0.0,
+    minlenratio: float = 0.0,
+    pre_beam_ratio: float = 1.5,
+    pre_beam_score_key: str = "full",
+) -> list:
+    """Perform beam search with scorers.
+
+    Args:
+        x (torch.Tensor): Encoded speech feature (T, D)
+        sos (int): Start of sequence id
+        eos (int): End of sequence id
+        beam_size (int): The number of hypotheses kept during search
+        vocab_size (int): The number of vocabulary
+        scorers (dict[str, ScorerInterface]): Dict of decoder modules
+            e.g., Decoder, CTCPrefixScorer, LM
+            The scorer will be ignored if it is `None`
+        weights (dict[str, float]): Dict of weights for each scorers
+            The scorer will be ignored if its weight is 0
+        token_list (list[str]): List of tokens for debug log
+        maxlenratio (float): Input length ratio to obtain max output length.
+            If maxlenratio=0.0 (default), it uses a end-detect function
+            to automatically find maximum hypothesis lengths
+        minlenratio (float): Input length ratio to obtain min output length.
+        pre_beam_score_key (str): key of scores to perform pre-beam search
+        pre_beam_ratio (float): beam size in the pre-beam search
+            will be `int(pre_beam_ratio * beam_size)`
+
+    Returns:
+        list: N-best decoding results
+
+    """
+    ret = BeamSearch(
+        scorers,
+        weights,
+        beam_size=beam_size,
+        vocab_size=vocab_size,
+        pre_beam_ratio=pre_beam_ratio,
+        pre_beam_score_key=pre_beam_score_key,
+        sos=sos,
+        eos=eos,
+        token_list=token_list,
+    ).forward(x=x, maxlenratio=maxlenratio, minlenratio=minlenratio)
+    return [h.asdict() for h in ret]

funasr/models/sa_asr/e2e_sa_asr.py

@@ -0,0 +1,508 @@
+# Copyright ESPnet (https://github.com/espnet/espnet). All Rights Reserved.
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+import logging
+from contextlib import contextmanager
+from distutils.version import LooseVersion
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Tuple
+from typing import Union
+
+import torch
+import torch.nn.functional as F
+
+from funasr.layers.abs_normalize import AbsNormalize
+from funasr.losses.label_smoothing_loss import LabelSmoothingLoss, NllLoss  # noqa: H301
+from funasr.models.ctc import CTC
+from funasr.models.decoder.abs_decoder import AbsDecoder
+from funasr.models.encoder.abs_encoder import AbsEncoder
+from funasr.frontends.abs_frontend import AbsFrontend
+from funasr.models.postencoder.abs_postencoder import AbsPostEncoder
+from funasr.models.preencoder.abs_preencoder import AbsPreEncoder
+from funasr.models.specaug.abs_specaug import AbsSpecAug
+from funasr.models.transformer.utils.add_sos_eos import add_sos_eos
+from funasr.metrics import ErrorCalculator
+from funasr.metrics.compute_acc import th_accuracy
+from funasr.train_utils.device_funcs import force_gatherable
+from funasr.models.base_model import FunASRModel
+
+if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
+    from torch.cuda.amp import autocast
+else:
+    # Nothing to do if torch<1.6.0
+    @contextmanager
+    def autocast(enabled=True):
+        yield
+
+
+class SAASRModel(FunASRModel):
+    """CTC-attention hybrid Encoder-Decoder model"""
+
+    def __init__(
+        self,
+        vocab_size: int,
+        max_spk_num: int,
+        token_list: Union[Tuple[str, ...], List[str]],
+        frontend: Optional[AbsFrontend],
+        specaug: Optional[AbsSpecAug],
+        normalize: Optional[AbsNormalize],
+        asr_encoder: AbsEncoder,
+        spk_encoder: torch.nn.Module,
+        decoder: AbsDecoder,
+        ctc: CTC,
+        spk_weight: float = 0.5,
+        ctc_weight: float = 0.5,
+        interctc_weight: float = 0.0,
+        ignore_id: int = -1,
+        lsm_weight: float = 0.0,
+        length_normalized_loss: bool = False,
+        report_cer: bool = True,
+        report_wer: bool = True,
+        sym_space: str = "<space>",
+        sym_blank: str = "<blank>",
+        extract_feats_in_collect_stats: bool = True,
+    ):
+        assert 0.0 <= ctc_weight <= 1.0, ctc_weight
+        assert 0.0 <= interctc_weight < 1.0, interctc_weight
+
+        super().__init__()
+        # note that eos is the same as sos (equivalent ID)
+        self.blank_id = 0
+        self.sos = 1
+        self.eos = 2
+        self.vocab_size = vocab_size
+        self.max_spk_num = max_spk_num
+        self.ignore_id = ignore_id
+        self.spk_weight = spk_weight
+        self.ctc_weight = ctc_weight
+        self.interctc_weight = interctc_weight
+        self.token_list = token_list.copy()
+
+        self.frontend = frontend
+        self.specaug = specaug
+        self.normalize = normalize
+        self.asr_encoder = asr_encoder
+        self.spk_encoder = spk_encoder
+
+        if not hasattr(self.asr_encoder, "interctc_use_conditioning"):
+            self.asr_encoder.interctc_use_conditioning = False
+        if self.asr_encoder.interctc_use_conditioning:
+            self.asr_encoder.conditioning_layer = torch.nn.Linear(
+                vocab_size, self.asr_encoder.output_size()
+            )
+
+        self.error_calculator = None
+
+        # we set self.decoder = None in the CTC mode since
+        # self.decoder parameters were never used and PyTorch complained
+        # and threw an Exception in the multi-GPU experiment.
+        # thanks Jeff Farris for pointing out the issue.
+        if ctc_weight == 1.0:
+            self.decoder = None
+        else:
+            self.decoder = decoder
+
+        self.criterion_att = LabelSmoothingLoss(
+            size=vocab_size,
+            padding_idx=ignore_id,
+            smoothing=lsm_weight,
+            normalize_length=length_normalized_loss,
+        )
+
+        self.criterion_spk = NllLoss(
+            size=max_spk_num,
+            padding_idx=ignore_id,
+            normalize_length=length_normalized_loss,
+        )
+
+        if report_cer or report_wer:
+            self.error_calculator = ErrorCalculator(
+                token_list, sym_space, sym_blank, report_cer, report_wer
+            )
+
+        if ctc_weight == 0.0:
+            self.ctc = None
+        else:
+            self.ctc = ctc
+
+        self.extract_feats_in_collect_stats = extract_feats_in_collect_stats
+
+    def forward(
+        self,
+        speech: torch.Tensor,
+        speech_lengths: torch.Tensor,
+        text: torch.Tensor,
+        text_lengths: torch.Tensor,
+        profile: torch.Tensor,
+        profile_lengths: torch.Tensor,
+        text_id: torch.Tensor,
+        text_id_lengths: torch.Tensor,
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:
+        """Frontend + Encoder + Decoder + Calc loss
+
+        Args:
+            speech: (Batch, Length, ...)
+            speech_lengths: (Batch, )
+            text: (Batch, Length)
+            text_lengths: (Batch,)
+            profile: (Batch, Length, Dim)
+            profile_lengths: (Batch,)
+        """
+        assert text_lengths.dim() == 1, text_lengths.shape
+        # Check that batch_size is unified
+        assert (
+            speech.shape[0] == speech_lengths.shape[0] == text.shape[0] == text_lengths.shape[0]
+        ), (speech.shape, speech_lengths.shape, text.shape, text_lengths.shape)
+        batch_size = speech.shape[0]
+
+        # for data-parallel
+        text = text[:, : text_lengths.max()]
+
+        # 1. Encoder
+        asr_encoder_out, encoder_out_lens, spk_encoder_out = self.encode(speech, speech_lengths)
+        intermediate_outs = None
+        if isinstance(asr_encoder_out, tuple):
+            intermediate_outs = asr_encoder_out[1]
+            asr_encoder_out = asr_encoder_out[0]
+
+        loss_att, loss_spk, acc_att, acc_spk, cer_att, wer_att = None, None, None, None, None, None
+        loss_ctc, cer_ctc = None, None
+        stats = dict()
+
+        # 1. CTC branch
+        if self.ctc_weight != 0.0:
+            loss_ctc, cer_ctc = self._calc_ctc_loss(
+                asr_encoder_out, encoder_out_lens, text, text_lengths
+            )
+
+        # Intermediate CTC (optional)
+        loss_interctc = 0.0
+        if self.interctc_weight != 0.0 and intermediate_outs is not None:
+            for layer_idx, intermediate_out in intermediate_outs:
+                # we assume intermediate_out has the same length & padding
+                # as those of encoder_out
+                loss_ic, cer_ic = self._calc_ctc_loss(
+                    intermediate_out, encoder_out_lens, text, text_lengths
+                )
+                loss_interctc = loss_interctc + loss_ic
+
+                # Collect Intermedaite CTC stats
+                stats["loss_interctc_layer{}".format(layer_idx)] = (
+                    loss_ic.detach() if loss_ic is not None else None
+                )
+                stats["cer_interctc_layer{}".format(layer_idx)] = cer_ic
+
+            loss_interctc = loss_interctc / len(intermediate_outs)
+
+            # calculate whole encoder loss
+            loss_ctc = (1 - self.interctc_weight) * loss_ctc + self.interctc_weight * loss_interctc
+
+        # 2b. Attention decoder branch
+        if self.ctc_weight != 1.0:
+            loss_att, loss_spk, acc_att, acc_spk, cer_att, wer_att = self._calc_att_loss(
+                asr_encoder_out,
+                spk_encoder_out,
+                encoder_out_lens,
+                text,
+                text_lengths,
+                profile,
+                profile_lengths,
+                text_id,
+                text_id_lengths,
+            )
+
+        # 3. CTC-Att loss definition
+        if self.ctc_weight == 0.0:
+            loss_asr = loss_att
+        elif self.ctc_weight == 1.0:
+            loss_asr = loss_ctc
+        else:
+            loss_asr = self.ctc_weight * loss_ctc + (1 - self.ctc_weight) * loss_att
+
+        if self.spk_weight == 0.0:
+            loss = loss_asr
+        else:
+            loss = self.spk_weight * loss_spk + (1 - self.spk_weight) * loss_asr
+
+        stats = dict(
+            loss=loss.detach(),
+            loss_asr=loss_asr.detach(),
+            loss_att=loss_att.detach() if loss_att is not None else None,
+            loss_ctc=loss_ctc.detach() if loss_ctc is not None else None,
+            loss_spk=loss_spk.detach() if loss_spk is not None else None,
+            acc=acc_att,
+            acc_spk=acc_spk,
+            cer=cer_att,
+            wer=wer_att,
+            cer_ctc=cer_ctc,
+        )
+
+        # force_gatherable: to-device and to-tensor if scalar for DataParallel
+        loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device)
+        return loss, stats, weight
+
+    def collect_feats(
+        self,
+        speech: torch.Tensor,
+        speech_lengths: torch.Tensor,
+        text: torch.Tensor,
+        text_lengths: torch.Tensor,
+    ) -> Dict[str, torch.Tensor]:
+        if self.extract_feats_in_collect_stats:
+            feats, feats_lengths = self._extract_feats(speech, speech_lengths)
+        else:
+            # Generate dummy stats if extract_feats_in_collect_stats is False
+            logging.warning(
+                "Generating dummy stats for feats and feats_lengths, "
+                "because encoder_conf.extract_feats_in_collect_stats is "
+                f"{self.extract_feats_in_collect_stats}"
+            )
+            feats, feats_lengths = speech, speech_lengths
+        return {"feats": feats, "feats_lengths": feats_lengths}
+
+    def encode(
+        self, speech: torch.Tensor, speech_lengths: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Frontend + Encoder. Note that this method is used by asr_inference.py
+
+        Args:
+            speech: (Batch, Length, ...)
+            speech_lengths: (Batch, )
+        """
+        with autocast(False):
+            # 1. Extract feats
+            feats, feats_lengths = self._extract_feats(speech, speech_lengths)
+
+            # 2. Data augmentation
+            feats_raw = feats.clone()
+            if self.specaug is not None and self.training:
+                feats, feats_lengths = self.specaug(feats, feats_lengths)
+
+            # 3. Normalization for feature: e.g. Global-CMVN, Utterance-CMVN
+            if self.normalize is not None:
+                feats, feats_lengths = self.normalize(feats, feats_lengths)
+
+        # 4. Forward encoder
+        # feats: (Batch, Length, Dim)
+        # -> encoder_out: (Batch, Length2, Dim2)
+        if self.asr_encoder.interctc_use_conditioning:
+            encoder_out, encoder_out_lens, _ = self.asr_encoder(feats, feats_lengths, ctc=self.ctc)
+        else:
+            encoder_out, encoder_out_lens, _ = self.asr_encoder(feats, feats_lengths)
+        intermediate_outs = None
+        if isinstance(encoder_out, tuple):
+            intermediate_outs = encoder_out[1]
+            encoder_out = encoder_out[0]
+
+        encoder_out_spk_ori = self.spk_encoder(feats_raw, feats_lengths)[0]
+        # import ipdb;ipdb.set_trace()
+        if encoder_out_spk_ori.size(1) != encoder_out.size(1):
+            encoder_out_spk = F.interpolate(
+                encoder_out_spk_ori.transpose(-2, -1), size=(encoder_out.size(1)), mode="nearest"
+            ).transpose(-2, -1)
+        else:
+            encoder_out_spk = encoder_out_spk_ori
+
+        assert encoder_out.size(0) == speech.size(0), (
+            encoder_out.size(),
+            speech.size(0),
+        )
+        assert encoder_out.size(1) <= encoder_out_lens.max(), (
+            encoder_out.size(),
+            encoder_out_lens.max(),
+        )
+        assert encoder_out_spk.size(0) == speech.size(0), (
+            encoder_out_spk.size(),
+            speech.size(0),
+        )
+
+        if intermediate_outs is not None:
+            return (encoder_out, intermediate_outs), encoder_out_lens, encoder_out_spk
+
+        return encoder_out, encoder_out_lens, encoder_out_spk
+
+    def _extract_feats(
+        self, speech: torch.Tensor, speech_lengths: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        assert speech_lengths.dim() == 1, speech_lengths.shape
+
+        # for data-parallel
+        speech = speech[:, : speech_lengths.max()]
+
+        if self.frontend is not None:
+            # Frontend
+            #  e.g. STFT and Feature extract
+            #       data_loader may send time-domain signal in this case
+            # speech (Batch, NSamples) -> feats: (Batch, NFrames, Dim)
+            feats, feats_lengths = self.frontend(speech, speech_lengths)
+        else:
+            # No frontend and no feature extract
+            feats, feats_lengths = speech, speech_lengths
+        return feats, feats_lengths
+
+    def nll(
+        self,
+        encoder_out: torch.Tensor,
+        encoder_out_lens: torch.Tensor,
+        ys_pad: torch.Tensor,
+        ys_pad_lens: torch.Tensor,
+    ) -> torch.Tensor:
+        """Compute negative log likelihood(nll) from transformer-decoder
+
+        Normally, this function is called in batchify_nll.
+
+        Args:
+            encoder_out: (Batch, Length, Dim)
+            encoder_out_lens: (Batch,)
+            ys_pad: (Batch, Length)
+            ys_pad_lens: (Batch,)
+        """
+        ys_in_pad, ys_out_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id)
+        ys_in_lens = ys_pad_lens + 1
+
+        # 1. Forward decoder
+        decoder_out, _ = self.decoder(
+            encoder_out, encoder_out_lens, ys_in_pad, ys_in_lens
+        )  # [batch, seqlen, dim]
+        batch_size = decoder_out.size(0)
+        decoder_num_class = decoder_out.size(2)
+        # nll: negative log-likelihood
+        nll = torch.nn.functional.cross_entropy(
+            decoder_out.view(-1, decoder_num_class),
+            ys_out_pad.view(-1),
+            ignore_index=self.ignore_id,
+            reduction="none",
+        )
+        nll = nll.view(batch_size, -1)
+        nll = nll.sum(dim=1)
+        assert nll.size(0) == batch_size
+        return nll
+
+    def batchify_nll(
+        self,
+        encoder_out: torch.Tensor,
+        encoder_out_lens: torch.Tensor,
+        ys_pad: torch.Tensor,
+        ys_pad_lens: torch.Tensor,
+        batch_size: int = 100,
+    ):
+        """Compute negative log likelihood(nll) from transformer-decoder
+
+        To avoid OOM, this fuction seperate the input into batches.
+        Then call nll for each batch and combine and return results.
+        Args:
+            encoder_out: (Batch, Length, Dim)
+            encoder_out_lens: (Batch,)
+            ys_pad: (Batch, Length)
+            ys_pad_lens: (Batch,)
+            batch_size: int, samples each batch contain when computing nll,
+                        you may change this to avoid OOM or increase
+                        GPU memory usage
+        """
+        total_num = encoder_out.size(0)
+        if total_num <= batch_size:
+            nll = self.nll(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens)
+        else:
+            nll = []
+            start_idx = 0
+            while True:
+                end_idx = min(start_idx + batch_size, total_num)
+                batch_encoder_out = encoder_out[start_idx:end_idx, :, :]
+                batch_encoder_out_lens = encoder_out_lens[start_idx:end_idx]
+                batch_ys_pad = ys_pad[start_idx:end_idx, :]
+                batch_ys_pad_lens = ys_pad_lens[start_idx:end_idx]
+                batch_nll = self.nll(
+                    batch_encoder_out,
+                    batch_encoder_out_lens,
+                    batch_ys_pad,
+                    batch_ys_pad_lens,
+                )
+                nll.append(batch_nll)
+                start_idx = end_idx
+                if start_idx == total_num:
+                    break
+            nll = torch.cat(nll)
+        assert nll.size(0) == total_num
+        return nll
+
+    def _calc_att_loss(
+        self,
+        asr_encoder_out: torch.Tensor,
+        spk_encoder_out: torch.Tensor,
+        encoder_out_lens: torch.Tensor,
+        ys_pad: torch.Tensor,
+        ys_pad_lens: torch.Tensor,
+        profile: torch.Tensor,
+        profile_lens: torch.Tensor,
+        text_id: torch.Tensor,
+        text_id_lengths: torch.Tensor,
+    ):
+        ys_in_pad, ys_out_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id)
+        ys_in_lens = ys_pad_lens + 1
+
+        # 1. Forward decoder
+        decoder_out, weights_no_pad, _ = self.decoder(
+            asr_encoder_out,
+            spk_encoder_out,
+            encoder_out_lens,
+            ys_in_pad,
+            ys_in_lens,
+            profile,
+            profile_lens,
+        )
+
+        spk_num_no_pad = weights_no_pad.size(-1)
+        pad = (0, self.max_spk_num - spk_num_no_pad)
+        weights = F.pad(weights_no_pad, pad, mode="constant", value=0)
+
+        # pre_id=weights.argmax(-1)
+        # pre_text=decoder_out.argmax(-1)
+        # id_mask=(pre_id==text_id).to(dtype=text_id.dtype)
+        # pre_text_mask=pre_text*id_mask+1-id_mask #相同的地方不变，不同的地方设为1(<unk>)
+        # padding_mask= ys_out_pad != self.ignore_id
+        # numerator = torch.sum(pre_text_mask.masked_select(padding_mask) == ys_out_pad.masked_select(padding_mask))
+        # denominator = torch.sum(padding_mask)
+        # sd_acc = float(numerator) / float(denominator)
+
+        # 2. Compute attention loss
+        loss_att = self.criterion_att(decoder_out, ys_out_pad)
+        loss_spk = self.criterion_spk(torch.log(weights), text_id)
+
+        acc_spk = th_accuracy(
+            weights.view(-1, self.max_spk_num),
+            text_id,
+            ignore_label=self.ignore_id,
+        )
+        acc_att = th_accuracy(
+            decoder_out.view(-1, self.vocab_size),
+            ys_out_pad,
+            ignore_label=self.ignore_id,
+        )
+
+        # Compute cer/wer using attention-decoder
+        if self.training or self.error_calculator is None:
+            cer_att, wer_att = None, None
+        else:
+            ys_hat = decoder_out.argmax(dim=-1)
+            cer_att, wer_att = self.error_calculator(ys_hat.cpu(), ys_pad.cpu())
+
+        return loss_att, loss_spk, acc_att, acc_spk, cer_att, wer_att
+
+    def _calc_ctc_loss(
+        self,
+        encoder_out: torch.Tensor,
+        encoder_out_lens: torch.Tensor,
+        ys_pad: torch.Tensor,
+        ys_pad_lens: torch.Tensor,
+    ):
+        # Calc CTC loss
+        loss_ctc = self.ctc(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens)
+
+        # Calc CER using CTC
+        cer_ctc = None
+        if not self.training and self.error_calculator is not None:
+            ys_hat = self.ctc.argmax(encoder_out).data
+            cer_ctc = self.error_calculator(ys_hat.cpu(), ys_pad.cpu(), is_ctc=True)
+        return loss_ctc, cer_ctc

funasr/models/sa_asr/transformer_decoder.py

@@ -0,0 +1,731 @@
+# Copyright 2019 Shigeki Karita
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Decoder definition."""
+from typing import Any
+from typing import List
+from typing import Sequence
+from typing import Tuple
+
+import torch
+from torch import nn
+
+
+from funasr.models.transformer.attention import MultiHeadedAttention
+from funasr.models.sa_asr.attention import CosineDistanceAttention
+from funasr.models.transformer.utils.dynamic_conv import DynamicConvolution
+from funasr.models.transformer.utils.dynamic_conv2d import DynamicConvolution2D
+from funasr.models.transformer.embedding import PositionalEncoding
+from funasr.models.transformer.layer_norm import LayerNorm
+from funasr.models.transformer.utils.lightconv import LightweightConvolution
+from funasr.models.transformer.utils.lightconv2d import LightweightConvolution2D
+from funasr.models.transformer.utils.mask import subsequent_mask
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+from funasr.models.transformer.positionwise_feed_forward import (
+    PositionwiseFeedForward,  # noqa: H301
+)
+from funasr.models.transformer.utils.repeat import repeat
+from funasr.models.transformer.scorers.scorer_interface import BatchScorerInterface
+
+from funasr.register import tables
+
+
+class DecoderLayer(nn.Module):
+    """Single decoder layer module.
+
+    Args:
+        size (int): Input dimension.
+        self_attn (torch.nn.Module): Self-attention module instance.
+            `MultiHeadedAttention` instance can be used as the argument.
+        src_attn (torch.nn.Module): Self-attention module instance.
+            `MultiHeadedAttention` instance can be used as the argument.
+        feed_forward (torch.nn.Module): Feed-forward module instance.
+            `PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance
+            can be used as the argument.
+        dropout_rate (float): Dropout rate.
+        normalize_before (bool): Whether to use layer_norm before the first block.
+        concat_after (bool): Whether to concat attention layer's input and output.
+            if True, additional linear will be applied.
+            i.e. x -> x + linear(concat(x, att(x)))
+            if False, no additional linear will be applied. i.e. x -> x + att(x)
+
+
+    """
+
+    def __init__(
+        self,
+        size,
+        self_attn,
+        src_attn,
+        feed_forward,
+        dropout_rate,
+        normalize_before=True,
+        concat_after=False,
+    ):
+        """Construct an DecoderLayer object."""
+        super(DecoderLayer, self).__init__()
+        self.size = size
+        self.self_attn = self_attn
+        self.src_attn = src_attn
+        self.feed_forward = feed_forward
+        self.norm1 = LayerNorm(size)
+        self.norm2 = LayerNorm(size)
+        self.norm3 = LayerNorm(size)
+        self.dropout = nn.Dropout(dropout_rate)
+        self.normalize_before = normalize_before
+        self.concat_after = concat_after
+        if self.concat_after:
+            self.concat_linear1 = nn.Linear(size + size, size)
+            self.concat_linear2 = nn.Linear(size + size, size)
+
+    def forward(self, tgt, tgt_mask, memory, memory_mask, cache=None):
+        """Compute decoded features.
+
+        Args:
+            tgt (torch.Tensor): Input tensor (#batch, maxlen_out, size).
+            tgt_mask (torch.Tensor): Mask for input tensor (#batch, maxlen_out).
+            memory (torch.Tensor): Encoded memory, float32 (#batch, maxlen_in, size).
+            memory_mask (torch.Tensor): Encoded memory mask (#batch, maxlen_in).
+            cache (List[torch.Tensor]): List of cached tensors.
+                Each tensor shape should be (#batch, maxlen_out - 1, size).
+
+        Returns:
+            torch.Tensor: Output tensor(#batch, maxlen_out, size).
+            torch.Tensor: Mask for output tensor (#batch, maxlen_out).
+            torch.Tensor: Encoded memory (#batch, maxlen_in, size).
+            torch.Tensor: Encoded memory mask (#batch, maxlen_in).
+
+        """
+        residual = tgt
+        if self.normalize_before:
+            tgt = self.norm1(tgt)
+
+        if cache is None:
+            tgt_q = tgt
+            tgt_q_mask = tgt_mask
+        else:
+            # compute only the last frame query keeping dim: max_time_out -> 1
+            assert cache.shape == (
+                tgt.shape[0],
+                tgt.shape[1] - 1,
+                self.size,
+            ), f"{cache.shape} == {(tgt.shape[0], tgt.shape[1] - 1, self.size)}"
+            tgt_q = tgt[:, -1:, :]
+            residual = residual[:, -1:, :]
+            tgt_q_mask = None
+            if tgt_mask is not None:
+                tgt_q_mask = tgt_mask[:, -1:, :]
+
+        if self.concat_after:
+            tgt_concat = torch.cat((tgt_q, self.self_attn(tgt_q, tgt, tgt, tgt_q_mask)), dim=-1)
+            x = residual + self.concat_linear1(tgt_concat)
+        else:
+            x = residual + self.dropout(self.self_attn(tgt_q, tgt, tgt, tgt_q_mask))
+        if not self.normalize_before:
+            x = self.norm1(x)
+
+        residual = x
+        if self.normalize_before:
+            x = self.norm2(x)
+        if self.concat_after:
+            x_concat = torch.cat((x, self.src_attn(x, memory, memory, memory_mask)), dim=-1)
+            x = residual + self.concat_linear2(x_concat)
+        else:
+            x = residual + self.dropout(self.src_attn(x, memory, memory, memory_mask))
+        if not self.normalize_before:
+            x = self.norm2(x)
+
+        residual = x
+        if self.normalize_before:
+            x = self.norm3(x)
+        x = residual + self.dropout(self.feed_forward(x))
+        if not self.normalize_before:
+            x = self.norm3(x)
+
+        if cache is not None:
+            x = torch.cat([cache, x], dim=1)
+
+        return x, tgt_mask, memory, memory_mask
+
+
+class BaseTransformerDecoder(nn.Module, BatchScorerInterface):
+    """Base class of Transfomer decoder module.
+
+    Args:
+        vocab_size: output dim
+        encoder_output_size: dimension of attention
+        attention_heads: the number of heads of multi head attention
+        linear_units: the number of units of position-wise feed forward
+        num_blocks: the number of decoder blocks
+        dropout_rate: dropout rate
+        self_attention_dropout_rate: dropout rate for attention
+        input_layer: input layer type
+        use_output_layer: whether to use output layer
+        pos_enc_class: PositionalEncoding or ScaledPositionalEncoding
+        normalize_before: whether to use layer_norm before the first block
+        concat_after: whether to concat attention layer's input and output
+            if True, additional linear will be applied.
+            i.e. x -> x + linear(concat(x, att(x)))
+            if False, no additional linear will be applied.
+            i.e. x -> x + att(x)
+    """
+
+    def __init__(
+        self,
+        vocab_size: int,
+        encoder_output_size: int,
+        dropout_rate: float = 0.1,
+        positional_dropout_rate: float = 0.1,
+        input_layer: str = "embed",
+        use_output_layer: bool = True,
+        pos_enc_class=PositionalEncoding,
+        normalize_before: bool = True,
+    ):
+        super().__init__()
+        attention_dim = encoder_output_size
+
+        if input_layer == "embed":
+            self.embed = torch.nn.Sequential(
+                torch.nn.Embedding(vocab_size, attention_dim),
+                pos_enc_class(attention_dim, positional_dropout_rate),
+            )
+        elif input_layer == "linear":
+            self.embed = torch.nn.Sequential(
+                torch.nn.Linear(vocab_size, attention_dim),
+                torch.nn.LayerNorm(attention_dim),
+                torch.nn.Dropout(dropout_rate),
+                torch.nn.ReLU(),
+                pos_enc_class(attention_dim, positional_dropout_rate),
+            )
+        else:
+            raise ValueError(f"only 'embed' or 'linear' is supported: {input_layer}")
+
+        self.normalize_before = normalize_before
+        if self.normalize_before:
+            self.after_norm = LayerNorm(attention_dim)
+        if use_output_layer:
+            self.output_layer = torch.nn.Linear(attention_dim, vocab_size)
+        else:
+            self.output_layer = None
+
+        # Must set by the inheritance
+        self.decoders = None
+
+    def forward(
+        self,
+        hs_pad: torch.Tensor,
+        hlens: torch.Tensor,
+        ys_in_pad: torch.Tensor,
+        ys_in_lens: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward decoder.
+
+        Args:
+            hs_pad: encoded memory, float32  (batch, maxlen_in, feat)
+            hlens: (batch)
+            ys_in_pad:
+                input token ids, int64 (batch, maxlen_out)
+                if input_layer == "embed"
+                input tensor (batch, maxlen_out, #mels) in the other cases
+            ys_in_lens: (batch)
+        Returns:
+            (tuple): tuple containing:
+
+            x: decoded token score before softmax (batch, maxlen_out, token)
+                if use_output_layer is True,
+            olens: (batch, )
+        """
+        tgt = ys_in_pad
+        # tgt_mask: (B, 1, L)
+        tgt_mask = (~make_pad_mask(ys_in_lens)[:, None, :]).to(tgt.device)
+        # m: (1, L, L)
+        m = subsequent_mask(tgt_mask.size(-1), device=tgt_mask.device).unsqueeze(0)
+        # tgt_mask: (B, L, L)
+        tgt_mask = tgt_mask & m
+
+        memory = hs_pad
+        memory_mask = (~make_pad_mask(hlens, maxlen=memory.size(1)))[:, None, :].to(memory.device)
+        # Padding for Longformer
+        if memory_mask.shape[-1] != memory.shape[1]:
+            padlen = memory.shape[1] - memory_mask.shape[-1]
+            memory_mask = torch.nn.functional.pad(memory_mask, (0, padlen), "constant", False)
+
+        x = self.embed(tgt)
+        x, tgt_mask, memory, memory_mask = self.decoders(x, tgt_mask, memory, memory_mask)
+        if self.normalize_before:
+            x = self.after_norm(x)
+        if self.output_layer is not None:
+            x = self.output_layer(x)
+
+        olens = tgt_mask.sum(1)
+        return x, olens
+
+    def forward_one_step(
+        self,
+        tgt: torch.Tensor,
+        tgt_mask: torch.Tensor,
+        memory: torch.Tensor,
+        cache: List[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+        """Forward one step.
+
+        Args:
+            tgt: input token ids, int64 (batch, maxlen_out)
+            tgt_mask: input token mask,  (batch, maxlen_out)
+                      dtype=torch.uint8 in PyTorch 1.2-
+                      dtype=torch.bool in PyTorch 1.2+ (include 1.2)
+            memory: encoded memory, float32  (batch, maxlen_in, feat)
+            cache: cached output list of (batch, max_time_out-1, size)
+        Returns:
+            y, cache: NN output value and cache per `self.decoders`.
+            y.shape` is (batch, maxlen_out, token)
+        """
+        x = self.embed(tgt)
+        if cache is None:
+            cache = [None] * len(self.decoders)
+        new_cache = []
+        for c, decoder in zip(cache, self.decoders):
+            x, tgt_mask, memory, memory_mask = decoder(x, tgt_mask, memory, None, cache=c)
+            new_cache.append(x)
+
+        if self.normalize_before:
+            y = self.after_norm(x[:, -1])
+        else:
+            y = x[:, -1]
+        if self.output_layer is not None:
+            y = torch.log_softmax(self.output_layer(y), dim=-1)
+
+        return y, new_cache
+
+    def score(self, ys, state, x):
+        """Score."""
+        ys_mask = subsequent_mask(len(ys), device=x.device).unsqueeze(0)
+        logp, state = self.forward_one_step(ys.unsqueeze(0), ys_mask, x.unsqueeze(0), cache=state)
+        return logp.squeeze(0), state
+
+    def batch_score(
+        self, ys: torch.Tensor, states: List[Any], xs: torch.Tensor
+    ) -> Tuple[torch.Tensor, List[Any]]:
+        """Score new token batch.
+
+        Args:
+            ys (torch.Tensor): torch.int64 prefix tokens (n_batch, ylen).
+            states (List[Any]): Scorer states for prefix tokens.
+            xs (torch.Tensor):
+                The encoder feature that generates ys (n_batch, xlen, n_feat).
+
+        Returns:
+            tuple[torch.Tensor, List[Any]]: Tuple of
+                batchfied scores for next token with shape of `(n_batch, n_vocab)`
+                and next state list for ys.
+
+        """
+        # merge states
+        n_batch = len(ys)
+        n_layers = len(self.decoders)
+        if states[0] is None:
+            batch_state = None
+        else:
+            # transpose state of [batch, layer] into [layer, batch]
+            batch_state = [
+                torch.stack([states[b][i] for b in range(n_batch)]) for i in range(n_layers)
+            ]
+
+        # batch decoding
+        ys_mask = subsequent_mask(ys.size(-1), device=xs.device).unsqueeze(0)
+        logp, states = self.forward_one_step(ys, ys_mask, xs, cache=batch_state)
+
+        # transpose state of [layer, batch] into [batch, layer]
+        state_list = [[states[i][b] for i in range(n_layers)] for b in range(n_batch)]
+        return logp, state_list
+
+
+@tables.register("decoder_classes", "TransformerDecoder")
+class TransformerDecoder(BaseTransformerDecoder):
+    def __init__(
+        self,
+        vocab_size: int,
+        encoder_output_size: int,
+        attention_heads: int = 4,
+        linear_units: int = 2048,
+        num_blocks: int = 6,
+        dropout_rate: float = 0.1,
+        positional_dropout_rate: float = 0.1,
+        self_attention_dropout_rate: float = 0.0,
+        src_attention_dropout_rate: float = 0.0,
+        input_layer: str = "embed",
+        use_output_layer: bool = True,
+        pos_enc_class=PositionalEncoding,
+        normalize_before: bool = True,
+        concat_after: bool = False,
+    ):
+        super().__init__(
+            vocab_size=vocab_size,
+            encoder_output_size=encoder_output_size,
+            dropout_rate=dropout_rate,
+            positional_dropout_rate=positional_dropout_rate,
+            input_layer=input_layer,
+            use_output_layer=use_output_layer,
+            pos_enc_class=pos_enc_class,
+            normalize_before=normalize_before,
+        )
+
+        attention_dim = encoder_output_size
+        self.decoders = repeat(
+            num_blocks,
+            lambda lnum: DecoderLayer(
+                attention_dim,
+                MultiHeadedAttention(attention_heads, attention_dim, self_attention_dropout_rate),
+                MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate),
+                PositionwiseFeedForward(attention_dim, linear_units, dropout_rate),
+                dropout_rate,
+                normalize_before,
+                concat_after,
+            ),
+        )
+
+
+@tables.register("decoder_classes", "ParaformerDecoderSAN")
+class ParaformerDecoderSAN(BaseTransformerDecoder):
+    """
+    Author: Speech Lab of DAMO Academy, Alibaba Group
+    Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition
+    https://arxiv.org/abs/2006.01713
+    """
+
+    def __init__(
+        self,
+        vocab_size: int,
+        encoder_output_size: int,
+        attention_heads: int = 4,
+        linear_units: int = 2048,
+        num_blocks: int = 6,
+        dropout_rate: float = 0.1,
+        positional_dropout_rate: float = 0.1,
+        self_attention_dropout_rate: float = 0.0,
+        src_attention_dropout_rate: float = 0.0,
+        input_layer: str = "embed",
+        use_output_layer: bool = True,
+        pos_enc_class=PositionalEncoding,
+        normalize_before: bool = True,
+        concat_after: bool = False,
+        embeds_id: int = -1,
+    ):
+        super().__init__(
+            vocab_size=vocab_size,
+            encoder_output_size=encoder_output_size,
+            dropout_rate=dropout_rate,
+            positional_dropout_rate=positional_dropout_rate,
+            input_layer=input_layer,
+            use_output_layer=use_output_layer,
+            pos_enc_class=pos_enc_class,
+            normalize_before=normalize_before,
+        )
+
+        attention_dim = encoder_output_size
+        self.decoders = repeat(
+            num_blocks,
+            lambda lnum: DecoderLayer(
+                attention_dim,
+                MultiHeadedAttention(attention_heads, attention_dim, self_attention_dropout_rate),
+                MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate),
+                PositionwiseFeedForward(attention_dim, linear_units, dropout_rate),
+                dropout_rate,
+                normalize_before,
+                concat_after,
+            ),
+        )
+        self.embeds_id = embeds_id
+        self.attention_dim = attention_dim
+
+    def forward(
+        self,
+        hs_pad: torch.Tensor,
+        hlens: torch.Tensor,
+        ys_in_pad: torch.Tensor,
+        ys_in_lens: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward decoder.
+
+        Args:
+            hs_pad: encoded memory, float32  (batch, maxlen_in, feat)
+            hlens: (batch)
+            ys_in_pad:
+                input token ids, int64 (batch, maxlen_out)
+                if input_layer == "embed"
+                input tensor (batch, maxlen_out, #mels) in the other cases
+            ys_in_lens: (batch)
+        Returns:
+            (tuple): tuple containing:
+
+            x: decoded token score before softmax (batch, maxlen_out, token)
+                if use_output_layer is True,
+            olens: (batch, )
+        """
+        tgt = ys_in_pad
+        tgt_mask = (~make_pad_mask(ys_in_lens)[:, None, :]).to(tgt.device)
+
+        memory = hs_pad
+        memory_mask = (~make_pad_mask(hlens, maxlen=memory.size(1)))[:, None, :].to(memory.device)
+        # Padding for Longformer
+        if memory_mask.shape[-1] != memory.shape[1]:
+            padlen = memory.shape[1] - memory_mask.shape[-1]
+            memory_mask = torch.nn.functional.pad(memory_mask, (0, padlen), "constant", False)
+
+        # x = self.embed(tgt)
+        x = tgt
+        embeds_outputs = None
+        for layer_id, decoder in enumerate(self.decoders):
+            x, tgt_mask, memory, memory_mask = decoder(x, tgt_mask, memory, memory_mask)
+            if layer_id == self.embeds_id:
+                embeds_outputs = x
+        if self.normalize_before:
+            x = self.after_norm(x)
+        if self.output_layer is not None:
+            x = self.output_layer(x)
+
+        olens = tgt_mask.sum(1)
+        if embeds_outputs is not None:
+            return x, olens, embeds_outputs
+        else:
+            return x, olens
+
+
+@tables.register("decoder_classes", "LightweightConvolutionTransformerDecoder")
+class LightweightConvolutionTransformerDecoder(BaseTransformerDecoder):
+    def __init__(
+        self,
+        vocab_size: int,
+        encoder_output_size: int,
+        attention_heads: int = 4,
+        linear_units: int = 2048,
+        num_blocks: int = 6,
+        dropout_rate: float = 0.1,
+        positional_dropout_rate: float = 0.1,
+        self_attention_dropout_rate: float = 0.0,
+        src_attention_dropout_rate: float = 0.0,
+        input_layer: str = "embed",
+        use_output_layer: bool = True,
+        pos_enc_class=PositionalEncoding,
+        normalize_before: bool = True,
+        concat_after: bool = False,
+        conv_wshare: int = 4,
+        conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11),
+        conv_usebias: int = False,
+    ):
+        if len(conv_kernel_length) != num_blocks:
+            raise ValueError(
+                "conv_kernel_length must have equal number of values to num_blocks: "
+                f"{len(conv_kernel_length)} != {num_blocks}"
+            )
+        super().__init__(
+            vocab_size=vocab_size,
+            encoder_output_size=encoder_output_size,
+            dropout_rate=dropout_rate,
+            positional_dropout_rate=positional_dropout_rate,
+            input_layer=input_layer,
+            use_output_layer=use_output_layer,
+            pos_enc_class=pos_enc_class,
+            normalize_before=normalize_before,
+        )
+
+        attention_dim = encoder_output_size
+        self.decoders = repeat(
+            num_blocks,
+            lambda lnum: DecoderLayer(
+                attention_dim,
+                LightweightConvolution(
+                    wshare=conv_wshare,
+                    n_feat=attention_dim,
+                    dropout_rate=self_attention_dropout_rate,
+                    kernel_size=conv_kernel_length[lnum],
+                    use_kernel_mask=True,
+                    use_bias=conv_usebias,
+                ),
+                MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate),
+                PositionwiseFeedForward(attention_dim, linear_units, dropout_rate),
+                dropout_rate,
+                normalize_before,
+                concat_after,
+            ),
+        )
+
+
+@tables.register("decoder_classes", "LightweightConvolution2DTransformerDecoder")
+class LightweightConvolution2DTransformerDecoder(BaseTransformerDecoder):
+    def __init__(
+        self,
+        vocab_size: int,
+        encoder_output_size: int,
+        attention_heads: int = 4,
+        linear_units: int = 2048,
+        num_blocks: int = 6,
+        dropout_rate: float = 0.1,
+        positional_dropout_rate: float = 0.1,
+        self_attention_dropout_rate: float = 0.0,
+        src_attention_dropout_rate: float = 0.0,
+        input_layer: str = "embed",
+        use_output_layer: bool = True,
+        pos_enc_class=PositionalEncoding,
+        normalize_before: bool = True,
+        concat_after: bool = False,
+        conv_wshare: int = 4,
+        conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11),
+        conv_usebias: int = False,
+    ):
+        if len(conv_kernel_length) != num_blocks:
+            raise ValueError(
+                "conv_kernel_length must have equal number of values to num_blocks: "
+                f"{len(conv_kernel_length)} != {num_blocks}"
+            )
+        super().__init__(
+            vocab_size=vocab_size,
+            encoder_output_size=encoder_output_size,
+            dropout_rate=dropout_rate,
+            positional_dropout_rate=positional_dropout_rate,
+            input_layer=input_layer,
+            use_output_layer=use_output_layer,
+            pos_enc_class=pos_enc_class,
+            normalize_before=normalize_before,
+        )
+
+        attention_dim = encoder_output_size
+        self.decoders = repeat(
+            num_blocks,
+            lambda lnum: DecoderLayer(
+                attention_dim,
+                LightweightConvolution2D(
+                    wshare=conv_wshare,
+                    n_feat=attention_dim,
+                    dropout_rate=self_attention_dropout_rate,
+                    kernel_size=conv_kernel_length[lnum],
+                    use_kernel_mask=True,
+                    use_bias=conv_usebias,
+                ),
+                MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate),
+                PositionwiseFeedForward(attention_dim, linear_units, dropout_rate),
+                dropout_rate,
+                normalize_before,
+                concat_after,
+            ),
+        )
+
+
+@tables.register("decoder_classes", "DynamicConvolutionTransformerDecoder")
+class DynamicConvolutionTransformerDecoder(BaseTransformerDecoder):
+    def __init__(
+        self,
+        vocab_size: int,
+        encoder_output_size: int,
+        attention_heads: int = 4,
+        linear_units: int = 2048,
+        num_blocks: int = 6,
+        dropout_rate: float = 0.1,
+        positional_dropout_rate: float = 0.1,
+        self_attention_dropout_rate: float = 0.0,
+        src_attention_dropout_rate: float = 0.0,
+        input_layer: str = "embed",
+        use_output_layer: bool = True,
+        pos_enc_class=PositionalEncoding,
+        normalize_before: bool = True,
+        concat_after: bool = False,
+        conv_wshare: int = 4,
+        conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11),
+        conv_usebias: int = False,
+    ):
+        if len(conv_kernel_length) != num_blocks:
+            raise ValueError(
+                "conv_kernel_length must have equal number of values to num_blocks: "
+                f"{len(conv_kernel_length)} != {num_blocks}"
+            )
+        super().__init__(
+            vocab_size=vocab_size,
+            encoder_output_size=encoder_output_size,
+            dropout_rate=dropout_rate,
+            positional_dropout_rate=positional_dropout_rate,
+            input_layer=input_layer,
+            use_output_layer=use_output_layer,
+            pos_enc_class=pos_enc_class,
+            normalize_before=normalize_before,
+        )
+        attention_dim = encoder_output_size
+
+        self.decoders = repeat(
+            num_blocks,
+            lambda lnum: DecoderLayer(
+                attention_dim,
+                DynamicConvolution(
+                    wshare=conv_wshare,
+                    n_feat=attention_dim,
+                    dropout_rate=self_attention_dropout_rate,
+                    kernel_size=conv_kernel_length[lnum],
+                    use_kernel_mask=True,
+                    use_bias=conv_usebias,
+                ),
+                MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate),
+                PositionwiseFeedForward(attention_dim, linear_units, dropout_rate),
+                dropout_rate,
+                normalize_before,
+                concat_after,
+            ),
+        )
+
+
+@tables.register("decoder_classes", "DynamicConvolution2DTransformerDecoder")
+class DynamicConvolution2DTransformerDecoder(BaseTransformerDecoder):
+    def __init__(
+        self,
+        vocab_size: int,
+        encoder_output_size: int,
+        attention_heads: int = 4,
+        linear_units: int = 2048,
+        num_blocks: int = 6,
+        dropout_rate: float = 0.1,
+        positional_dropout_rate: float = 0.1,
+        self_attention_dropout_rate: float = 0.0,
+        src_attention_dropout_rate: float = 0.0,
+        input_layer: str = "embed",
+        use_output_layer: bool = True,
+        pos_enc_class=PositionalEncoding,
+        normalize_before: bool = True,
+        concat_after: bool = False,
+        conv_wshare: int = 4,
+        conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11),
+        conv_usebias: int = False,
+    ):
+        if len(conv_kernel_length) != num_blocks:
+            raise ValueError(
+                "conv_kernel_length must have equal number of values to num_blocks: "
+                f"{len(conv_kernel_length)} != {num_blocks}"
+            )
+        super().__init__(
+            vocab_size=vocab_size,
+            encoder_output_size=encoder_output_size,
+            dropout_rate=dropout_rate,
+            positional_dropout_rate=positional_dropout_rate,
+            input_layer=input_layer,
+            use_output_layer=use_output_layer,
+            pos_enc_class=pos_enc_class,
+            normalize_before=normalize_before,
+        )
+        attention_dim = encoder_output_size
+
+        self.decoders = repeat(
+            num_blocks,
+            lambda lnum: DecoderLayer(
+                attention_dim,
+                DynamicConvolution2D(
+                    wshare=conv_wshare,
+                    n_feat=attention_dim,
+                    dropout_rate=self_attention_dropout_rate,
+                    kernel_size=conv_kernel_length[lnum],
+                    use_kernel_mask=True,
+                    use_bias=conv_usebias,
+                ),
+                MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate),
+                PositionwiseFeedForward(attention_dim, linear_units, dropout_rate),
+                dropout_rate,
+                normalize_before,
+                concat_after,
+            ),
+        )

funasr/models/sanm/attention.py

@@ -0,0 +1,903 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# Copyright 2019 Shigeki Karita
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Multi-Head Attention layer definition."""
+
+import math
+
+import numpy
+import torch
+from torch import nn
+from typing import Optional, Tuple
+
+import torch.nn.functional as F
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+import funasr.models.lora.layers as lora
+
+
+def preprocess_for_attn(x, mask, cache, pad_fn, kernel_size):
+    x = x * mask
+    x = x.transpose(1, 2)
+    if cache is None:
+        x = pad_fn(x)
+    else:
+        x = torch.cat((cache, x), dim=2)
+        cache = x[:, :, -(kernel_size - 1) :]
+    return x, cache
+
+
+torch_version = tuple([int(i) for i in torch.__version__.split(".")[:2]])
+if torch_version >= (1, 8):
+    import torch.fx
+
+    torch.fx.wrap("preprocess_for_attn")
+
+
+class MultiHeadedAttention(nn.Module):
+    """Multi-Head Attention layer.
+
+    Args:
+        n_head (int): The number of heads.
+        n_feat (int): The number of features.
+        dropout_rate (float): Dropout rate.
+
+    """
+
+    def __init__(self, n_head, n_feat, dropout_rate):
+        """Construct an MultiHeadedAttention object."""
+        super(MultiHeadedAttention, self).__init__()
+        assert n_feat % n_head == 0
+        # We assume d_v always equals d_k
+        self.d_k = n_feat // n_head
+        self.h = n_head
+        self.linear_q = nn.Linear(n_feat, n_feat)
+        self.linear_k = nn.Linear(n_feat, n_feat)
+        self.linear_v = nn.Linear(n_feat, n_feat)
+        self.linear_out = nn.Linear(n_feat, n_feat)
+        self.attn = None
+        self.dropout = nn.Dropout(p=dropout_rate)
+
+    def forward_qkv(self, query, key, value):
+        """Transform query, key and value.
+
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+
+        Returns:
+            torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k).
+            torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k).
+            torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k).
+
+        """
+        n_batch = query.size(0)
+        q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k)
+        k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k)
+        v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k)
+        q = q.transpose(1, 2)  # (batch, head, time1, d_k)
+        k = k.transpose(1, 2)  # (batch, head, time2, d_k)
+        v = v.transpose(1, 2)  # (batch, head, time2, d_k)
+
+        return q, k, v
+
+    def forward_attention(self, value, scores, mask):
+        """Compute attention context vector.
+
+        Args:
+            value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k).
+            scores (torch.Tensor): Attention score (#batch, n_head, time1, time2).
+            mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2).
+
+        Returns:
+            torch.Tensor: Transformed value (#batch, time1, d_model)
+                weighted by the attention score (#batch, time1, time2).
+
+        """
+        n_batch = value.size(0)
+        if mask is not None:
+            mask = mask.unsqueeze(1).eq(0)  # (batch, 1, *, time2)
+            min_value = -float(
+                "inf"
+            )  # float(numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min)
+            scores = scores.masked_fill(mask, min_value)
+            attn = torch.softmax(scores, dim=-1).masked_fill(
+                mask, 0.0
+            )  # (batch, head, time1, time2)
+        else:
+            attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)
+
+        p_attn = self.dropout(attn)
+        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
+        x = (
+            x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)
+        )  # (batch, time1, d_model)
+
+        return self.linear_out(x)  # (batch, time1, d_model)
+
+    def forward(self, query, key, value, mask):
+        """Compute scaled dot product attention.
+
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+            mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
+                (#batch, time1, time2).
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time1, d_model).
+
+        """
+        q, k, v = self.forward_qkv(query, key, value)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
+        return self.forward_attention(v, scores, mask)
+
+
+class MultiHeadedAttentionSANM(nn.Module):
+    """Multi-Head Attention layer.
+
+    Args:
+        n_head (int): The number of heads.
+        n_feat (int): The number of features.
+        dropout_rate (float): Dropout rate.
+
+    """
+
+    def __init__(
+        self,
+        n_head,
+        in_feat,
+        n_feat,
+        dropout_rate,
+        kernel_size,
+        sanm_shfit=0,
+        lora_list=None,
+        lora_rank=8,
+        lora_alpha=16,
+        lora_dropout=0.1,
+    ):
+        """Construct an MultiHeadedAttention object."""
+        super().__init__()
+        assert n_feat % n_head == 0
+        # We assume d_v always equals d_k
+        self.d_k = n_feat // n_head
+        self.h = n_head
+        # self.linear_q = nn.Linear(n_feat, n_feat)
+        # self.linear_k = nn.Linear(n_feat, n_feat)
+        # self.linear_v = nn.Linear(n_feat, n_feat)
+        if lora_list is not None:
+            if "o" in lora_list:
+                self.linear_out = lora.Linear(
+                    n_feat, n_feat, r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout
+                )
+            else:
+                self.linear_out = nn.Linear(n_feat, n_feat)
+            lora_qkv_list = ["q" in lora_list, "k" in lora_list, "v" in lora_list]
+            if lora_qkv_list == [False, False, False]:
+                self.linear_q_k_v = nn.Linear(in_feat, n_feat * 3)
+            else:
+                self.linear_q_k_v = lora.MergedLinear(
+                    in_feat,
+                    n_feat * 3,
+                    r=lora_rank,
+                    lora_alpha=lora_alpha,
+                    lora_dropout=lora_dropout,
+                    enable_lora=lora_qkv_list,
+                )
+        else:
+            self.linear_out = nn.Linear(n_feat, n_feat)
+            self.linear_q_k_v = nn.Linear(in_feat, n_feat * 3)
+        attn = None
+        self.dropout = nn.Dropout(p=dropout_rate)
+
+        self.fsmn_block = nn.Conv1d(
+            n_feat, n_feat, kernel_size, stride=1, padding=0, groups=n_feat, bias=False
+        )
+        # padding
+        left_padding = (kernel_size - 1) // 2
+        if sanm_shfit > 0:
+            left_padding = left_padding + sanm_shfit
+        right_padding = kernel_size - 1 - left_padding
+        self.pad_fn = nn.ConstantPad1d((left_padding, right_padding), 0.0)
+
+    def forward_fsmn(self, inputs, mask, mask_shfit_chunk=None):
+        b, t, d = inputs.size()
+        if mask is not None:
+            mask = torch.reshape(mask, (b, -1, 1))
+            if mask_shfit_chunk is not None:
+                mask = mask * mask_shfit_chunk
+            inputs = inputs * mask
+
+        x = inputs.transpose(1, 2)
+        x = self.pad_fn(x)
+        x = self.fsmn_block(x)
+        x = x.transpose(1, 2)
+        x += inputs
+        x = self.dropout(x)
+        if mask is not None:
+            x = x * mask
+        return x
+
+    def forward_qkv(self, x):
+        """Transform query, key and value.
+
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+
+        Returns:
+            torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k).
+            torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k).
+            torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k).
+
+        """
+        b, t, d = x.size()
+        q_k_v = self.linear_q_k_v(x)
+        q, k, v = torch.split(q_k_v, int(self.h * self.d_k), dim=-1)
+        q_h = torch.reshape(q, (b, t, self.h, self.d_k)).transpose(
+            1, 2
+        )  # (batch, head, time1, d_k)
+        k_h = torch.reshape(k, (b, t, self.h, self.d_k)).transpose(
+            1, 2
+        )  # (batch, head, time2, d_k)
+        v_h = torch.reshape(v, (b, t, self.h, self.d_k)).transpose(
+            1, 2
+        )  # (batch, head, time2, d_k)
+
+        return q_h, k_h, v_h, v
+
+    def forward_attention(self, value, scores, mask, mask_att_chunk_encoder=None):
+        """Compute attention context vector.
+
+        Args:
+            value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k).
+            scores (torch.Tensor): Attention score (#batch, n_head, time1, time2).
+            mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2).
+
+        Returns:
+            torch.Tensor: Transformed value (#batch, time1, d_model)
+                weighted by the attention score (#batch, time1, time2).
+
+        """
+        n_batch = value.size(0)
+        if mask is not None:
+            if mask_att_chunk_encoder is not None:
+                mask = mask * mask_att_chunk_encoder
+
+            mask = mask.unsqueeze(1).eq(0)  # (batch, 1, *, time2)
+
+            min_value = -float(
+                "inf"
+            )  # float(numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min)
+            scores = scores.masked_fill(mask, min_value)
+            attn = torch.softmax(scores, dim=-1).masked_fill(
+                mask, 0.0
+            )  # (batch, head, time1, time2)
+        else:
+            attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)
+
+        p_attn = self.dropout(attn)
+        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
+        x = (
+            x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)
+        )  # (batch, time1, d_model)
+
+        return self.linear_out(x)  # (batch, time1, d_model)
+
+    def forward(self, x, mask, mask_shfit_chunk=None, mask_att_chunk_encoder=None):
+        """Compute scaled dot product attention.
+
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+            mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
+                (#batch, time1, time2).
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time1, d_model).
+
+        """
+        q_h, k_h, v_h, v = self.forward_qkv(x)
+        fsmn_memory = self.forward_fsmn(v, mask, mask_shfit_chunk)
+        q_h = q_h * self.d_k ** (-0.5)
+        scores = torch.matmul(q_h, k_h.transpose(-2, -1))
+        att_outs = self.forward_attention(v_h, scores, mask, mask_att_chunk_encoder)
+        return att_outs + fsmn_memory
+
+    def forward_chunk(self, x, cache=None, chunk_size=None, look_back=0):
+        """Compute scaled dot product attention.
+
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+            mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
+                (#batch, time1, time2).
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time1, d_model).
+
+        """
+        q_h, k_h, v_h, v = self.forward_qkv(x)
+        if chunk_size is not None and look_back > 0 or look_back == -1:
+            if cache is not None:
+                k_h_stride = k_h[:, :, : -(chunk_size[2]), :]
+                v_h_stride = v_h[:, :, : -(chunk_size[2]), :]
+                k_h = torch.cat((cache["k"], k_h), dim=2)
+                v_h = torch.cat((cache["v"], v_h), dim=2)
+
+                cache["k"] = torch.cat((cache["k"], k_h_stride), dim=2)
+                cache["v"] = torch.cat((cache["v"], v_h_stride), dim=2)
+                if look_back != -1:
+                    cache["k"] = cache["k"][:, :, -(look_back * chunk_size[1]) :, :]
+                    cache["v"] = cache["v"][:, :, -(look_back * chunk_size[1]) :, :]
+            else:
+                cache_tmp = {
+                    "k": k_h[:, :, : -(chunk_size[2]), :],
+                    "v": v_h[:, :, : -(chunk_size[2]), :],
+                }
+                cache = cache_tmp
+        fsmn_memory = self.forward_fsmn(v, None)
+        q_h = q_h * self.d_k ** (-0.5)
+        scores = torch.matmul(q_h, k_h.transpose(-2, -1))
+        att_outs = self.forward_attention(v_h, scores, None)
+        return att_outs + fsmn_memory, cache
+
+
+class MultiHeadedAttentionSANMExport(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.d_k = model.d_k
+        self.h = model.h
+        self.linear_out = model.linear_out
+        self.linear_q_k_v = model.linear_q_k_v
+        self.fsmn_block = model.fsmn_block
+        self.pad_fn = model.pad_fn
+
+        self.attn = None
+        self.all_head_size = self.h * self.d_k
+
+    def forward(self, x, mask):
+        mask_3d_btd, mask_4d_bhlt = mask
+        q_h, k_h, v_h, v = self.forward_qkv(x)
+        fsmn_memory = self.forward_fsmn(v, mask_3d_btd)
+        q_h = q_h * self.d_k ** (-0.5)
+        scores = torch.matmul(q_h, k_h.transpose(-2, -1))
+        att_outs = self.forward_attention(v_h, scores, mask_4d_bhlt)
+        return att_outs + fsmn_memory
+
+    def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
+        new_x_shape = x.size()[:-1] + (self.h, self.d_k)
+        x = x.view(new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward_qkv(self, x):
+        q_k_v = self.linear_q_k_v(x)
+        q, k, v = torch.split(q_k_v, int(self.h * self.d_k), dim=-1)
+        q_h = self.transpose_for_scores(q)
+        k_h = self.transpose_for_scores(k)
+        v_h = self.transpose_for_scores(v)
+        return q_h, k_h, v_h, v
+
+    def forward_fsmn(self, inputs, mask):
+        # b, t, d = inputs.size()
+        # mask = torch.reshape(mask, (b, -1, 1))
+        inputs = inputs * mask
+        x = inputs.transpose(1, 2)
+        x = self.pad_fn(x)
+        x = self.fsmn_block(x)
+        x = x.transpose(1, 2)
+        x = x + inputs
+        x = x * mask
+        return x
+
+    def forward_attention(self, value, scores, mask):
+        scores = scores + mask
+
+        attn = torch.softmax(scores, dim=-1)
+        context_layer = torch.matmul(attn, value)  # (batch, head, time1, d_k)
+
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(new_context_layer_shape)
+        return self.linear_out(context_layer)  # (batch, time1, d_model)
+
+
+class MultiHeadedAttentionSANMExport(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.d_k = model.d_k
+        self.h = model.h
+        self.linear_out = model.linear_out
+        self.linear_q_k_v = model.linear_q_k_v
+        self.fsmn_block = model.fsmn_block
+        self.pad_fn = model.pad_fn
+
+        self.attn = None
+        self.all_head_size = self.h * self.d_k
+
+    def forward(self, x, mask):
+        mask_3d_btd, mask_4d_bhlt = mask
+        q_h, k_h, v_h, v = self.forward_qkv(x)
+        fsmn_memory = self.forward_fsmn(v, mask_3d_btd)
+        q_h = q_h * self.d_k ** (-0.5)
+        scores = torch.matmul(q_h, k_h.transpose(-2, -1))
+        att_outs = self.forward_attention(v_h, scores, mask_4d_bhlt)
+        return att_outs + fsmn_memory
+
+    def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
+        new_x_shape = x.size()[:-1] + (self.h, self.d_k)
+        x = x.view(new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward_qkv(self, x):
+        q_k_v = self.linear_q_k_v(x)
+        q, k, v = torch.split(q_k_v, int(self.h * self.d_k), dim=-1)
+        q_h = self.transpose_for_scores(q)
+        k_h = self.transpose_for_scores(k)
+        v_h = self.transpose_for_scores(v)
+        return q_h, k_h, v_h, v
+
+    def forward_fsmn(self, inputs, mask):
+        # b, t, d = inputs.size()
+        # mask = torch.reshape(mask, (b, -1, 1))
+        inputs = inputs * mask
+        x = inputs.transpose(1, 2)
+        x = self.pad_fn(x)
+        x = self.fsmn_block(x)
+        x = x.transpose(1, 2)
+        x = x + inputs
+        x = x * mask
+        return x
+
+    def forward_attention(self, value, scores, mask):
+        scores = scores + mask
+
+        attn = torch.softmax(scores, dim=-1)
+        context_layer = torch.matmul(attn, value)  # (batch, head, time1, d_k)
+
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(new_context_layer_shape)
+        return self.linear_out(context_layer)  # (batch, time1, d_model)
+
+
+class MultiHeadedAttentionSANMDecoder(nn.Module):
+    """Multi-Head Attention layer.
+
+    Args:
+        n_head (int): The number of heads.
+        n_feat (int): The number of features.
+        dropout_rate (float): Dropout rate.
+
+    """
+
+    def __init__(self, n_feat, dropout_rate, kernel_size, sanm_shfit=0):
+        """Construct an MultiHeadedAttention object."""
+        super().__init__()
+
+        self.dropout = nn.Dropout(p=dropout_rate)
+
+        self.fsmn_block = nn.Conv1d(
+            n_feat, n_feat, kernel_size, stride=1, padding=0, groups=n_feat, bias=False
+        )
+        # padding
+        # padding
+        left_padding = (kernel_size - 1) // 2
+        if sanm_shfit > 0:
+            left_padding = left_padding + sanm_shfit
+        right_padding = kernel_size - 1 - left_padding
+        self.pad_fn = nn.ConstantPad1d((left_padding, right_padding), 0.0)
+        self.kernel_size = kernel_size
+
+    def forward(self, inputs, mask, cache=None, mask_shfit_chunk=None):
+        """
+        :param x: (#batch, time1, size).
+        :param mask: Mask tensor (#batch, 1, time)
+        :return:
+        """
+        # print("in fsmn, inputs", inputs.size())
+        b, t, d = inputs.size()
+        # logging.info(
+        #     "mask: {}".format(mask.size()))
+        if mask is not None:
+            mask = torch.reshape(mask, (b, -1, 1))
+            # logging.info("in fsmn, mask: {}, {}".format(mask.size(), mask[0:100:50, :, :]))
+            if mask_shfit_chunk is not None:
+                # logging.info("in fsmn, mask_fsmn: {}, {}".format(mask_shfit_chunk.size(), mask_shfit_chunk[0:100:50, :, :]))
+                mask = mask * mask_shfit_chunk
+            # logging.info("in fsmn, mask_after_fsmn: {}, {}".format(mask.size(), mask[0:100:50, :, :]))
+            # print("in fsmn, mask", mask.size())
+            # print("in fsmn, inputs", inputs.size())
+            inputs = inputs * mask
+
+        x = inputs.transpose(1, 2)
+        b, d, t = x.size()
+        if cache is None:
+            # print("in fsmn, cache is None, x", x.size())
+
+            x = self.pad_fn(x)
+            if not self.training:
+                cache = x
+        else:
+            # print("in fsmn, cache is not None, x", x.size())
+            # x = torch.cat((x, cache), dim=2)[:, :, :-1]
+            # if t < self.kernel_size:
+            #     x = self.pad_fn(x)
+            x = torch.cat((cache[:, :, 1:], x), dim=2)
+            x = x[:, :, -(self.kernel_size + t - 1) :]
+            # print("in fsmn, cache is not None, x_cat", x.size())
+            cache = x
+        x = self.fsmn_block(x)
+        x = x.transpose(1, 2)
+        # print("in fsmn, fsmn_out", x.size())
+        if x.size(1) != inputs.size(1):
+            inputs = inputs[:, -1, :]
+
+        x = x + inputs
+        x = self.dropout(x)
+        if mask is not None:
+            x = x * mask
+        return x, cache
+
+
+class MultiHeadedAttentionSANMDecoderExport(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.fsmn_block = model.fsmn_block
+        self.pad_fn = model.pad_fn
+        self.kernel_size = model.kernel_size
+        self.attn = None
+
+    def forward(self, inputs, mask, cache=None):
+        x, cache = preprocess_for_attn(inputs, mask, cache, self.pad_fn, self.kernel_size)
+        x = self.fsmn_block(x)
+        x = x.transpose(1, 2)
+
+        x = x + inputs
+        x = x * mask
+        return x, cache
+
+
+class MultiHeadedAttentionCrossAtt(nn.Module):
+    """Multi-Head Attention layer.
+
+    Args:
+        n_head (int): The number of heads.
+        n_feat (int): The number of features.
+        dropout_rate (float): Dropout rate.
+
+    """
+
+    def __init__(
+        self,
+        n_head,
+        n_feat,
+        dropout_rate,
+        lora_list=None,
+        lora_rank=8,
+        lora_alpha=16,
+        lora_dropout=0.1,
+        encoder_output_size=None,
+    ):
+        """Construct an MultiHeadedAttention object."""
+        super().__init__()
+        assert n_feat % n_head == 0
+        # We assume d_v always equals d_k
+        self.d_k = n_feat // n_head
+        self.h = n_head
+        if lora_list is not None:
+            if "q" in lora_list:
+                self.linear_q = lora.Linear(
+                    n_feat, n_feat, r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout
+                )
+            else:
+                self.linear_q = nn.Linear(n_feat, n_feat)
+            lora_kv_list = ["k" in lora_list, "v" in lora_list]
+            if lora_kv_list == [False, False]:
+                self.linear_k_v = nn.Linear(
+                    n_feat if encoder_output_size is None else encoder_output_size, n_feat * 2
+                )
+            else:
+                self.linear_k_v = lora.MergedLinear(
+                    n_feat if encoder_output_size is None else encoder_output_size,
+                    n_feat * 2,
+                    r=lora_rank,
+                    lora_alpha=lora_alpha,
+                    lora_dropout=lora_dropout,
+                    enable_lora=lora_kv_list,
+                )
+            if "o" in lora_list:
+                self.linear_out = lora.Linear(
+                    n_feat, n_feat, r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout
+                )
+            else:
+                self.linear_out = nn.Linear(n_feat, n_feat)
+        else:
+            self.linear_q = nn.Linear(n_feat, n_feat)
+            self.linear_k_v = nn.Linear(
+                n_feat if encoder_output_size is None else encoder_output_size, n_feat * 2
+            )
+            self.linear_out = nn.Linear(n_feat, n_feat)
+        self.attn = None
+        self.dropout = nn.Dropout(p=dropout_rate)
+
+    def forward_qkv(self, x, memory):
+        """Transform query, key and value.
+
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+
+        Returns:
+            torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k).
+            torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k).
+            torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k).
+
+        """
+
+        # print("in forward_qkv, x", x.size())
+        b = x.size(0)
+        q = self.linear_q(x)
+        q_h = torch.reshape(q, (b, -1, self.h, self.d_k)).transpose(
+            1, 2
+        )  # (batch, head, time1, d_k)
+
+        k_v = self.linear_k_v(memory)
+        k, v = torch.split(k_v, int(self.h * self.d_k), dim=-1)
+        k_h = torch.reshape(k, (b, -1, self.h, self.d_k)).transpose(
+            1, 2
+        )  # (batch, head, time2, d_k)
+        v_h = torch.reshape(v, (b, -1, self.h, self.d_k)).transpose(
+            1, 2
+        )  # (batch, head, time2, d_k)
+
+        return q_h, k_h, v_h
+
+    def forward_attention(self, value, scores, mask, ret_attn=False):
+        """Compute attention context vector.
+
+        Args:
+            value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k).
+            scores (torch.Tensor): Attention score (#batch, n_head, time1, time2).
+            mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2).
+
+        Returns:
+            torch.Tensor: Transformed value (#batch, time1, d_model)
+                weighted by the attention score (#batch, time1, time2).
+
+        """
+        n_batch = value.size(0)
+        if mask is not None:
+            mask = mask.unsqueeze(1).eq(0)  # (batch, 1, *, time2)
+            min_value = -float(
+                "inf"
+            )  # float(numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min)
+            # logging.info(
+            #     "scores: {}, mask_size: {}".format(scores.size(), mask.size()))
+            scores = scores.masked_fill(mask, min_value)
+            attn = torch.softmax(scores, dim=-1).masked_fill(
+                mask, 0.0
+            )  # (batch, head, time1, time2)
+        else:
+            attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)
+        p_attn = self.dropout(attn)
+        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
+        x = (
+            x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)
+        )  # (batch, time1, d_model)
+        if ret_attn:
+            return self.linear_out(x), attn  # (batch, time1, d_model)
+        return self.linear_out(x)  # (batch, time1, d_model)
+
+    def forward(self, x, memory, memory_mask, ret_attn=False):
+        """Compute scaled dot product attention.
+
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+            mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
+                (#batch, time1, time2).
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time1, d_model).
+
+        """
+        q_h, k_h, v_h = self.forward_qkv(x, memory)
+        q_h = q_h * self.d_k ** (-0.5)
+        scores = torch.matmul(q_h, k_h.transpose(-2, -1))
+        return self.forward_attention(v_h, scores, memory_mask, ret_attn=ret_attn)
+
+    def forward_chunk(self, x, memory, cache=None, chunk_size=None, look_back=0):
+        """Compute scaled dot product attention.
+
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+            mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
+                (#batch, time1, time2).
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time1, d_model).
+
+        """
+        q_h, k_h, v_h = self.forward_qkv(x, memory)
+        if chunk_size is not None and look_back > 0:
+            if cache is not None:
+                k_h = torch.cat((cache["k"], k_h), dim=2)
+                v_h = torch.cat((cache["v"], v_h), dim=2)
+                cache["k"] = k_h[:, :, -(look_back * chunk_size[1]) :, :]
+                cache["v"] = v_h[:, :, -(look_back * chunk_size[1]) :, :]
+            else:
+                cache_tmp = {
+                    "k": k_h[:, :, -(look_back * chunk_size[1]) :, :],
+                    "v": v_h[:, :, -(look_back * chunk_size[1]) :, :],
+                }
+                cache = cache_tmp
+        q_h = q_h * self.d_k ** (-0.5)
+        scores = torch.matmul(q_h, k_h.transpose(-2, -1))
+        return self.forward_attention(v_h, scores, None), cache
+
+
+class MultiHeadedAttentionCrossAttExport(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.d_k = model.d_k
+        self.h = model.h
+        self.linear_q = model.linear_q
+        self.linear_k_v = model.linear_k_v
+        self.linear_out = model.linear_out
+        self.attn = None
+        self.all_head_size = self.h * self.d_k
+
+    def forward(self, x, memory, memory_mask, ret_attn=False):
+        q, k, v = self.forward_qkv(x, memory)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
+        return self.forward_attention(v, scores, memory_mask, ret_attn)
+
+    def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
+        new_x_shape = x.size()[:-1] + (self.h, self.d_k)
+        x = x.view(new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward_qkv(self, x, memory):
+        q = self.linear_q(x)
+
+        k_v = self.linear_k_v(memory)
+        k, v = torch.split(k_v, int(self.h * self.d_k), dim=-1)
+        q = self.transpose_for_scores(q)
+        k = self.transpose_for_scores(k)
+        v = self.transpose_for_scores(v)
+        return q, k, v
+
+    def forward_attention(self, value, scores, mask, ret_attn):
+        scores = scores + mask.to(scores.device)
+
+        attn = torch.softmax(scores, dim=-1)
+        context_layer = torch.matmul(attn, value)  # (batch, head, time1, d_k)
+
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(new_context_layer_shape)
+        if ret_attn:
+            return self.linear_out(context_layer), attn
+        return self.linear_out(context_layer)  # (batch, time1, d_model)
+
+
+class MultiHeadSelfAttention(nn.Module):
+    """Multi-Head Attention layer.
+
+    Args:
+        n_head (int): The number of heads.
+        n_feat (int): The number of features.
+        dropout_rate (float): Dropout rate.
+
+    """
+
+    def __init__(self, n_head, in_feat, n_feat, dropout_rate):
+        """Construct an MultiHeadedAttention object."""
+        super(MultiHeadSelfAttention, self).__init__()
+        assert n_feat % n_head == 0
+        # We assume d_v always equals d_k
+        self.d_k = n_feat // n_head
+        self.h = n_head
+        self.linear_out = nn.Linear(n_feat, n_feat)
+        self.linear_q_k_v = nn.Linear(in_feat, n_feat * 3)
+        self.attn = None
+        self.dropout = nn.Dropout(p=dropout_rate)
+
+    def forward_qkv(self, x):
+        """Transform query, key and value.
+
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+
+        Returns:
+            torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k).
+            torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k).
+            torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k).
+
+        """
+        b, t, d = x.size()
+        q_k_v = self.linear_q_k_v(x)
+        q, k, v = torch.split(q_k_v, int(self.h * self.d_k), dim=-1)
+        q_h = torch.reshape(q, (b, t, self.h, self.d_k)).transpose(
+            1, 2
+        )  # (batch, head, time1, d_k)
+        k_h = torch.reshape(k, (b, t, self.h, self.d_k)).transpose(
+            1, 2
+        )  # (batch, head, time2, d_k)
+        v_h = torch.reshape(v, (b, t, self.h, self.d_k)).transpose(
+            1, 2
+        )  # (batch, head, time2, d_k)
+
+        return q_h, k_h, v_h, v
+
+    def forward_attention(self, value, scores, mask, mask_att_chunk_encoder=None):
+        """Compute attention context vector.
+
+        Args:
+            value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k).
+            scores (torch.Tensor): Attention score (#batch, n_head, time1, time2).
+            mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2).
+
+        Returns:
+            torch.Tensor: Transformed value (#batch, time1, d_model)
+                weighted by the attention score (#batch, time1, time2).
+
+        """
+        n_batch = value.size(0)
+        if mask is not None:
+            if mask_att_chunk_encoder is not None:
+                mask = mask * mask_att_chunk_encoder
+
+            mask = mask.unsqueeze(1).eq(0)  # (batch, 1, *, time2)
+
+            min_value = -float(
+                "inf"
+            )  # float(numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min)
+            scores = scores.masked_fill(mask, min_value)
+            attn = torch.softmax(scores, dim=-1).masked_fill(
+                mask, 0.0
+            )  # (batch, head, time1, time2)
+        else:
+            attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)
+
+        p_attn = self.dropout(attn)
+        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
+        x = (
+            x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)
+        )  # (batch, time1, d_model)
+
+        return self.linear_out(x)  # (batch, time1, d_model)
+
+    def forward(self, x, mask, mask_att_chunk_encoder=None):
+        """Compute scaled dot product attention.
+
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+            mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
+                (#batch, time1, time2).
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time1, d_model).
+
+        """
+        q_h, k_h, v_h, v = self.forward_qkv(x)
+        q_h = q_h * self.d_k ** (-0.5)
+        scores = torch.matmul(q_h, k_h.transpose(-2, -1))
+        att_outs = self.forward_attention(v_h, scores, mask, mask_att_chunk_encoder)
+        return att_outs

funasr/models/sanm/decoder.py

@@ -0,0 +1,498 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+from typing import List
+from typing import Tuple
+import logging
+import torch
+import torch.nn as nn
+import numpy as np
+
+from funasr.models.scama import utils as myutils
+from funasr.models.transformer.decoder import BaseTransformerDecoder
+
+from funasr.models.sanm.attention import (
+    MultiHeadedAttentionSANMDecoder,
+    MultiHeadedAttentionCrossAtt,
+)
+from funasr.models.transformer.embedding import PositionalEncoding
+from funasr.models.transformer.layer_norm import LayerNorm
+from funasr.models.sanm.positionwise_feed_forward import PositionwiseFeedForwardDecoderSANM
+from funasr.models.transformer.utils.repeat import repeat
+
+from funasr.register import tables
+
+
+class DecoderLayerSANM(nn.Module):
+    """Single decoder layer module.
+
+    Args:
+        size (int): Input dimension.
+        self_attn (torch.nn.Module): Self-attention module instance.
+            `MultiHeadedAttention` instance can be used as the argument.
+        src_attn (torch.nn.Module): Self-attention module instance.
+            `MultiHeadedAttention` instance can be used as the argument.
+        feed_forward (torch.nn.Module): Feed-forward module instance.
+            `PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance
+            can be used as the argument.
+        dropout_rate (float): Dropout rate.
+        normalize_before (bool): Whether to use layer_norm before the first block.
+        concat_after (bool): Whether to concat attention layer's input and output.
+            if True, additional linear will be applied.
+            i.e. x -> x + linear(concat(x, att(x)))
+            if False, no additional linear will be applied. i.e. x -> x + att(x)
+
+
+    """
+
+    def __init__(
+        self,
+        size,
+        self_attn,
+        src_attn,
+        feed_forward,
+        dropout_rate,
+        normalize_before=True,
+        concat_after=False,
+    ):
+        """Construct an DecoderLayer object."""
+        super(DecoderLayerSANM, self).__init__()
+        self.size = size
+        self.self_attn = self_attn
+        self.src_attn = src_attn
+        self.feed_forward = feed_forward
+        self.norm1 = LayerNorm(size)
+        if self_attn is not None:
+            self.norm2 = LayerNorm(size)
+        if src_attn is not None:
+            self.norm3 = LayerNorm(size)
+        self.dropout = nn.Dropout(dropout_rate)
+        self.normalize_before = normalize_before
+        self.concat_after = concat_after
+        if self.concat_after:
+            self.concat_linear1 = nn.Linear(size + size, size)
+            self.concat_linear2 = nn.Linear(size + size, size)
+
+    def forward(self, tgt, tgt_mask, memory, memory_mask=None, cache=None):
+        """Compute decoded features.
+
+        Args:
+            tgt (torch.Tensor): Input tensor (#batch, maxlen_out, size).
+            tgt_mask (torch.Tensor): Mask for input tensor (#batch, maxlen_out).
+            memory (torch.Tensor): Encoded memory, float32 (#batch, maxlen_in, size).
+            memory_mask (torch.Tensor): Encoded memory mask (#batch, maxlen_in).
+            cache (List[torch.Tensor]): List of cached tensors.
+                Each tensor shape should be (#batch, maxlen_out - 1, size).
+
+        Returns:
+            torch.Tensor: Output tensor(#batch, maxlen_out, size).
+            torch.Tensor: Mask for output tensor (#batch, maxlen_out).
+            torch.Tensor: Encoded memory (#batch, maxlen_in, size).
+            torch.Tensor: Encoded memory mask (#batch, maxlen_in).
+
+        """
+        # tgt = self.dropout(tgt)
+        residual = tgt
+        if self.normalize_before:
+            tgt = self.norm1(tgt)
+        tgt = self.feed_forward(tgt)
+
+        x = tgt
+        if self.self_attn:
+            if self.normalize_before:
+                tgt = self.norm2(tgt)
+            x, _ = self.self_attn(tgt, tgt_mask)
+            x = residual + self.dropout(x)
+
+        if self.src_attn is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.norm3(x)
+
+            x = residual + self.dropout(self.src_attn(x, memory, memory_mask))
+
+        return x, tgt_mask, memory, memory_mask, cache
+
+    def forward_one_step(self, tgt, tgt_mask, memory, memory_mask=None, cache=None):
+        """Compute decoded features.
+
+        Args:
+            tgt (torch.Tensor): Input tensor (#batch, maxlen_out, size).
+            tgt_mask (torch.Tensor): Mask for input tensor (#batch, maxlen_out).
+            memory (torch.Tensor): Encoded memory, float32 (#batch, maxlen_in, size).
+            memory_mask (torch.Tensor): Encoded memory mask (#batch, maxlen_in).
+            cache (List[torch.Tensor]): List of cached tensors.
+                Each tensor shape should be (#batch, maxlen_out - 1, size).
+
+        Returns:
+            torch.Tensor: Output tensor(#batch, maxlen_out, size).
+            torch.Tensor: Mask for output tensor (#batch, maxlen_out).
+            torch.Tensor: Encoded memory (#batch, maxlen_in, size).
+            torch.Tensor: Encoded memory mask (#batch, maxlen_in).
+
+        """
+        # tgt = self.dropout(tgt)
+        residual = tgt
+        if self.normalize_before:
+            tgt = self.norm1(tgt)
+        tgt = self.feed_forward(tgt)
+
+        x = tgt
+        if self.self_attn:
+            if self.normalize_before:
+                tgt = self.norm2(tgt)
+            if self.training:
+                cache = None
+            x, cache = self.self_attn(tgt, tgt_mask, cache=cache)
+            x = residual + self.dropout(x)
+
+        if self.src_attn is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.norm3(x)
+
+            x = residual + self.dropout(self.src_attn(x, memory, memory_mask))
+
+        return x, tgt_mask, memory, memory_mask, cache
+
+    def forward_chunk(
+        self, tgt, memory, fsmn_cache=None, opt_cache=None, chunk_size=None, look_back=0
+    ):
+        """Compute decoded features.
+
+        Args:
+            tgt (torch.Tensor): Input tensor (#batch, maxlen_out, size).
+            tgt_mask (torch.Tensor): Mask for input tensor (#batch, maxlen_out).
+            memory (torch.Tensor): Encoded memory, float32 (#batch, maxlen_in, size).
+            memory_mask (torch.Tensor): Encoded memory mask (#batch, maxlen_in).
+            cache (List[torch.Tensor]): List of cached tensors.
+                Each tensor shape should be (#batch, maxlen_out - 1, size).
+
+        Returns:
+            torch.Tensor: Output tensor(#batch, maxlen_out, size).
+            torch.Tensor: Mask for output tensor (#batch, maxlen_out).
+            torch.Tensor: Encoded memory (#batch, maxlen_in, size).
+            torch.Tensor: Encoded memory mask (#batch, maxlen_in).
+
+        """
+        residual = tgt
+        if self.normalize_before:
+            tgt = self.norm1(tgt)
+        tgt = self.feed_forward(tgt)
+
+        x = tgt
+        if self.self_attn:
+            if self.normalize_before:
+                tgt = self.norm2(tgt)
+            x, fsmn_cache = self.self_attn(tgt, None, fsmn_cache)
+            x = residual + self.dropout(x)
+
+        if self.src_attn is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.norm3(x)
+
+            x, opt_cache = self.src_attn.forward_chunk(x, memory, opt_cache, chunk_size, look_back)
+            x = residual + x
+
+        return x, memory, fsmn_cache, opt_cache
+
+
+@tables.register("decoder_classes", "FsmnDecoder")
+class FsmnDecoder(BaseTransformerDecoder):
+    """
+    Author: Zhifu Gao, Shiliang Zhang, Ming Lei, Ian McLoughlin
+    San-m: Memory equipped self-attention for end-to-end speech recognition
+    https://arxiv.org/abs/2006.01713
+    """
+
+    def __init__(
+        self,
+        vocab_size: int,
+        encoder_output_size: int,
+        attention_heads: int = 4,
+        linear_units: int = 2048,
+        num_blocks: int = 6,
+        dropout_rate: float = 0.1,
+        positional_dropout_rate: float = 0.1,
+        self_attention_dropout_rate: float = 0.0,
+        src_attention_dropout_rate: float = 0.0,
+        input_layer: str = "embed",
+        use_output_layer: bool = True,
+        pos_enc_class=PositionalEncoding,
+        normalize_before: bool = True,
+        concat_after: bool = False,
+        att_layer_num: int = 6,
+        kernel_size: int = 21,
+        sanm_shfit: int = None,
+        concat_embeds: bool = False,
+        attention_dim: int = None,
+        tf2torch_tensor_name_prefix_torch: str = "decoder",
+        tf2torch_tensor_name_prefix_tf: str = "seq2seq/decoder",
+        embed_tensor_name_prefix_tf: str = None,
+    ):
+        super().__init__(
+            vocab_size=vocab_size,
+            encoder_output_size=encoder_output_size,
+            dropout_rate=dropout_rate,
+            positional_dropout_rate=positional_dropout_rate,
+            input_layer=input_layer,
+            use_output_layer=use_output_layer,
+            pos_enc_class=pos_enc_class,
+            normalize_before=normalize_before,
+        )
+        if attention_dim is None:
+            attention_dim = encoder_output_size
+
+        if input_layer == "embed":
+            self.embed = torch.nn.Sequential(
+                torch.nn.Embedding(vocab_size, attention_dim),
+            )
+        elif input_layer == "linear":
+            self.embed = torch.nn.Sequential(
+                torch.nn.Linear(vocab_size, attention_dim),
+                torch.nn.LayerNorm(attention_dim),
+                torch.nn.Dropout(dropout_rate),
+                torch.nn.ReLU(),
+                pos_enc_class(attention_dim, positional_dropout_rate),
+            )
+        else:
+            raise ValueError(f"only 'embed' or 'linear' is supported: {input_layer}")
+
+        self.normalize_before = normalize_before
+        if self.normalize_before:
+            self.after_norm = LayerNorm(attention_dim)
+        if use_output_layer:
+            self.output_layer = torch.nn.Linear(attention_dim, vocab_size)
+        else:
+            self.output_layer = None
+
+        self.att_layer_num = att_layer_num
+        self.num_blocks = num_blocks
+        if sanm_shfit is None:
+            sanm_shfit = (kernel_size - 1) // 2
+        self.decoders = repeat(
+            att_layer_num,
+            lambda lnum: DecoderLayerSANM(
+                attention_dim,
+                MultiHeadedAttentionSANMDecoder(
+                    attention_dim, self_attention_dropout_rate, kernel_size, sanm_shfit=sanm_shfit
+                ),
+                MultiHeadedAttentionCrossAtt(
+                    attention_heads,
+                    attention_dim,
+                    src_attention_dropout_rate,
+                    encoder_output_size=encoder_output_size,
+                ),
+                PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate),
+                dropout_rate,
+                normalize_before,
+                concat_after,
+            ),
+        )
+        if num_blocks - att_layer_num <= 0:
+            self.decoders2 = None
+        else:
+            self.decoders2 = repeat(
+                num_blocks - att_layer_num,
+                lambda lnum: DecoderLayerSANM(
+                    attention_dim,
+                    MultiHeadedAttentionSANMDecoder(
+                        attention_dim,
+                        self_attention_dropout_rate,
+                        kernel_size,
+                        sanm_shfit=sanm_shfit,
+                    ),
+                    None,
+                    PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate),
+                    dropout_rate,
+                    normalize_before,
+                    concat_after,
+                ),
+            )
+
+        self.decoders3 = repeat(
+            1,
+            lambda lnum: DecoderLayerSANM(
+                attention_dim,
+                None,
+                None,
+                PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate),
+                dropout_rate,
+                normalize_before,
+                concat_after,
+            ),
+        )
+        if concat_embeds:
+            self.embed_concat_ffn = repeat(
+                1,
+                lambda lnum: DecoderLayerSANM(
+                    attention_dim + encoder_output_size,
+                    None,
+                    None,
+                    PositionwiseFeedForwardDecoderSANM(
+                        attention_dim + encoder_output_size,
+                        linear_units,
+                        dropout_rate,
+                        adim=attention_dim,
+                    ),
+                    dropout_rate,
+                    normalize_before,
+                    concat_after,
+                ),
+            )
+        else:
+            self.embed_concat_ffn = None
+        self.concat_embeds = concat_embeds
+        self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch
+        self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf
+        self.embed_tensor_name_prefix_tf = embed_tensor_name_prefix_tf
+
+    def forward(
+        self,
+        hs_pad: torch.Tensor,
+        hlens: torch.Tensor,
+        ys_in_pad: torch.Tensor,
+        ys_in_lens: torch.Tensor,
+        chunk_mask: torch.Tensor = None,
+        pre_acoustic_embeds: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward decoder.
+
+        Args:
+            hs_pad: encoded memory, float32  (batch, maxlen_in, feat)
+            hlens: (batch)
+            ys_in_pad:
+                input token ids, int64 (batch, maxlen_out)
+                if input_layer == "embed"
+                input tensor (batch, maxlen_out, #mels) in the other cases
+            ys_in_lens: (batch)
+        Returns:
+            (tuple): tuple containing:
+
+            x: decoded token score before softmax (batch, maxlen_out, token)
+                if use_output_layer is True,
+            olens: (batch, )
+        """
+        tgt = ys_in_pad
+        tgt_mask = myutils.sequence_mask(ys_in_lens, device=tgt.device)[:, :, None]
+
+        memory = hs_pad
+        memory_mask = myutils.sequence_mask(hlens, device=memory.device)[:, None, :]
+        if chunk_mask is not None:
+            memory_mask = memory_mask * chunk_mask
+            if tgt_mask.size(1) != memory_mask.size(1):
+                memory_mask = torch.cat((memory_mask, memory_mask[:, -2:-1, :]), dim=1)
+
+        x = self.embed(tgt)
+
+        if pre_acoustic_embeds is not None and self.concat_embeds:
+            x = torch.cat((x, pre_acoustic_embeds), dim=-1)
+            x, _, _, _, _ = self.embed_concat_ffn(x, None, None, None, None)
+
+        x, tgt_mask, memory, memory_mask, _ = self.decoders(x, tgt_mask, memory, memory_mask)
+        if self.decoders2 is not None:
+            x, tgt_mask, memory, memory_mask, _ = self.decoders2(x, tgt_mask, memory, memory_mask)
+        x, tgt_mask, memory, memory_mask, _ = self.decoders3(x, tgt_mask, memory, memory_mask)
+        if self.normalize_before:
+            x = self.after_norm(x)
+        if self.output_layer is not None:
+            x = self.output_layer(x)
+
+        olens = tgt_mask.sum(1)
+        return x, olens
+
+    def score(
+        self,
+        ys,
+        state,
+        x,
+        x_mask=None,
+        pre_acoustic_embeds: torch.Tensor = None,
+    ):
+        """Score."""
+        ys_mask = myutils.sequence_mask(
+            torch.tensor([len(ys)], dtype=torch.int32), device=x.device
+        )[:, :, None]
+        logp, state = self.forward_one_step(
+            ys.unsqueeze(0),
+            ys_mask,
+            x.unsqueeze(0),
+            memory_mask=x_mask,
+            pre_acoustic_embeds=pre_acoustic_embeds,
+            cache=state,
+        )
+        return logp.squeeze(0), state
+
+    def forward_one_step(
+        self,
+        tgt: torch.Tensor,
+        tgt_mask: torch.Tensor,
+        memory: torch.Tensor,
+        memory_mask: torch.Tensor = None,
+        pre_acoustic_embeds: torch.Tensor = None,
+        cache: List[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+        """Forward one step.
+
+        Args:
+            tgt: input token ids, int64 (batch, maxlen_out)
+            tgt_mask: input token mask,  (batch, maxlen_out)
+                      dtype=torch.uint8 in PyTorch 1.2-
+                      dtype=torch.bool in PyTorch 1.2+ (include 1.2)
+            memory: encoded memory, float32  (batch, maxlen_in, feat)
+            cache: cached output list of (batch, max_time_out-1, size)
+        Returns:
+            y, cache: NN output value and cache per `self.decoders`.
+            y.shape` is (batch, maxlen_out, token)
+        """
+
+        x = tgt[:, -1:]
+        tgt_mask = None
+        x = self.embed(x)
+
+        if pre_acoustic_embeds is not None and self.concat_embeds:
+            x = torch.cat((x, pre_acoustic_embeds), dim=-1)
+            x, _, _, _, _ = self.embed_concat_ffn(x, None, None, None, None)
+
+        if cache is None:
+            cache_layer_num = len(self.decoders)
+            if self.decoders2 is not None:
+                cache_layer_num += len(self.decoders2)
+            cache = [None] * cache_layer_num
+        new_cache = []
+        # for c, decoder in zip(cache, self.decoders):
+        for i in range(self.att_layer_num):
+            decoder = self.decoders[i]
+            c = cache[i]
+            x, tgt_mask, memory, memory_mask, c_ret = decoder.forward_one_step(
+                x, tgt_mask, memory, memory_mask, cache=c
+            )
+            new_cache.append(c_ret)
+
+        if self.num_blocks - self.att_layer_num >= 1:
+            for i in range(self.num_blocks - self.att_layer_num):
+                j = i + self.att_layer_num
+                decoder = self.decoders2[i]
+                c = cache[j]
+                x, tgt_mask, memory, memory_mask, c_ret = decoder.forward_one_step(
+                    x, tgt_mask, memory, memory_mask, cache=c
+                )
+                new_cache.append(c_ret)
+
+        for decoder in self.decoders3:
+            x, tgt_mask, memory, memory_mask, _ = decoder.forward_one_step(
+                x, tgt_mask, memory, None, cache=None
+            )
+
+        if self.normalize_before:
+            y = self.after_norm(x[:, -1])
+        else:
+            y = x[:, -1]
+        if self.output_layer is not None:
+            y = self.output_layer(y)
+            y = torch.log_softmax(y, dim=-1)
+
+        return y, new_cache

funasr/models/sanm/encoder.py

@@ -0,0 +1,612 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+from typing import List
+from typing import Optional
+from typing import Sequence
+from typing import Tuple
+from typing import Union
+import logging
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import numpy as np
+from funasr.train_utils.device_funcs import to_device
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+from funasr.models.sanm.attention import MultiHeadedAttention, MultiHeadedAttentionSANM
+from funasr.models.transformer.embedding import (
+    SinusoidalPositionEncoder,
+    StreamSinusoidalPositionEncoder,
+)
+from funasr.models.transformer.layer_norm import LayerNorm
+from funasr.models.transformer.utils.multi_layer_conv import Conv1dLinear
+from funasr.models.transformer.utils.multi_layer_conv import MultiLayeredConv1d
+from funasr.models.transformer.positionwise_feed_forward import (
+    PositionwiseFeedForward,  # noqa: H301
+)
+from funasr.models.transformer.utils.repeat import repeat
+from funasr.models.transformer.utils.subsampling import Conv2dSubsampling
+from funasr.models.transformer.utils.subsampling import Conv2dSubsampling2
+from funasr.models.transformer.utils.subsampling import Conv2dSubsampling6
+from funasr.models.transformer.utils.subsampling import Conv2dSubsampling8
+from funasr.models.transformer.utils.subsampling import TooShortUttError
+from funasr.models.transformer.utils.subsampling import check_short_utt
+
+
+from funasr.models.ctc.ctc import CTC
+
+from funasr.register import tables
+
+
+class EncoderLayerSANM(nn.Module):
+    def __init__(
+        self,
+        in_size,
+        size,
+        self_attn,
+        feed_forward,
+        dropout_rate,
+        normalize_before=True,
+        concat_after=False,
+        stochastic_depth_rate=0.0,
+    ):
+        """Construct an EncoderLayer object."""
+        super(EncoderLayerSANM, self).__init__()
+        self.self_attn = self_attn
+        self.feed_forward = feed_forward
+        self.norm1 = LayerNorm(in_size)
+        self.norm2 = LayerNorm(size)
+        self.dropout = nn.Dropout(dropout_rate)
+        self.in_size = in_size
+        self.size = size
+        self.normalize_before = normalize_before
+        self.concat_after = concat_after
+        if self.concat_after:
+            self.concat_linear = nn.Linear(size + size, size)
+        self.stochastic_depth_rate = stochastic_depth_rate
+        self.dropout_rate = dropout_rate
+
+    def forward(self, x, mask, cache=None, mask_shfit_chunk=None, mask_att_chunk_encoder=None):
+        """Compute encoded features.
+
+        Args:
+            x_input (torch.Tensor): Input tensor (#batch, time, size).
+            mask (torch.Tensor): Mask tensor for the input (#batch, time).
+            cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size).
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time, size).
+            torch.Tensor: Mask tensor (#batch, time).
+
+        """
+        skip_layer = False
+        # with stochastic depth, residual connection `x + f(x)` becomes
+        # `x <- x + 1 / (1 - p) * f(x)` at training time.
+        stoch_layer_coeff = 1.0
+        if self.training and self.stochastic_depth_rate > 0:
+            skip_layer = torch.rand(1).item() < self.stochastic_depth_rate
+            stoch_layer_coeff = 1.0 / (1 - self.stochastic_depth_rate)
+
+        if skip_layer:
+            if cache is not None:
+                x = torch.cat([cache, x], dim=1)
+            return x, mask
+
+        residual = x
+        if self.normalize_before:
+            x = self.norm1(x)
+
+        if self.concat_after:
+            x_concat = torch.cat(
+                (
+                    x,
+                    self.self_attn(
+                        x,
+                        mask,
+                        mask_shfit_chunk=mask_shfit_chunk,
+                        mask_att_chunk_encoder=mask_att_chunk_encoder,
+                    ),
+                ),
+                dim=-1,
+            )
+            if self.in_size == self.size:
+                x = residual + stoch_layer_coeff * self.concat_linear(x_concat)
+            else:
+                x = stoch_layer_coeff * self.concat_linear(x_concat)
+        else:
+            if self.in_size == self.size:
+                x = residual + stoch_layer_coeff * self.dropout(
+                    self.self_attn(
+                        x,
+                        mask,
+                        mask_shfit_chunk=mask_shfit_chunk,
+                        mask_att_chunk_encoder=mask_att_chunk_encoder,
+                    )
+                )
+            else:
+                x = stoch_layer_coeff * self.dropout(
+                    self.self_attn(
+                        x,
+                        mask,
+                        mask_shfit_chunk=mask_shfit_chunk,
+                        mask_att_chunk_encoder=mask_att_chunk_encoder,
+                    )
+                )
+        if not self.normalize_before:
+            x = self.norm1(x)
+
+        residual = x
+        if self.normalize_before:
+            x = self.norm2(x)
+        x = residual + stoch_layer_coeff * self.dropout(self.feed_forward(x))
+        if not self.normalize_before:
+            x = self.norm2(x)
+
+        return x, mask, cache, mask_shfit_chunk, mask_att_chunk_encoder
+
+    def forward_chunk(self, x, cache=None, chunk_size=None, look_back=0):
+        """Compute encoded features.
+
+        Args:
+            x_input (torch.Tensor): Input tensor (#batch, time, size).
+            mask (torch.Tensor): Mask tensor for the input (#batch, time).
+            cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size).
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time, size).
+            torch.Tensor: Mask tensor (#batch, time).
+
+        """
+
+        residual = x
+        if self.normalize_before:
+            x = self.norm1(x)
+
+        if self.in_size == self.size:
+            attn, cache = self.self_attn.forward_chunk(x, cache, chunk_size, look_back)
+            x = residual + attn
+        else:
+            x, cache = self.self_attn.forward_chunk(x, cache, chunk_size, look_back)
+
+        if not self.normalize_before:
+            x = self.norm1(x)
+
+        residual = x
+        if self.normalize_before:
+            x = self.norm2(x)
+        x = residual + self.feed_forward(x)
+        if not self.normalize_before:
+            x = self.norm2(x)
+
+        return x, cache
+
+
+@tables.register("encoder_classes", "SANMEncoder")
+class SANMEncoder(nn.Module):
+    """
+    Author: Zhifu Gao, Shiliang Zhang, Ming Lei, Ian McLoughlin
+    San-m: Memory equipped self-attention for end-to-end speech recognition
+    https://arxiv.org/abs/2006.01713
+    """
+
+    def __init__(
+        self,
+        input_size: int,
+        output_size: int = 256,
+        attention_heads: int = 4,
+        linear_units: int = 2048,
+        num_blocks: int = 6,
+        dropout_rate: float = 0.1,
+        positional_dropout_rate: float = 0.1,
+        attention_dropout_rate: float = 0.0,
+        input_layer: Optional[str] = "conv2d",
+        pos_enc_class=SinusoidalPositionEncoder,
+        normalize_before: bool = True,
+        concat_after: bool = False,
+        positionwise_layer_type: str = "linear",
+        positionwise_conv_kernel_size: int = 1,
+        padding_idx: int = -1,
+        interctc_layer_idx: List[int] = [],
+        interctc_use_conditioning: bool = False,
+        kernel_size: int = 11,
+        sanm_shfit: int = 0,
+        lora_list: List[str] = None,
+        lora_rank: int = 8,
+        lora_alpha: int = 16,
+        lora_dropout: float = 0.1,
+        selfattention_layer_type: str = "sanm",
+        tf2torch_tensor_name_prefix_torch: str = "encoder",
+        tf2torch_tensor_name_prefix_tf: str = "seq2seq/encoder",
+    ):
+        super().__init__()
+        self._output_size = output_size
+
+        if input_layer == "linear":
+            self.embed = torch.nn.Sequential(
+                torch.nn.Linear(input_size, output_size),
+                torch.nn.LayerNorm(output_size),
+                torch.nn.Dropout(dropout_rate),
+                torch.nn.ReLU(),
+                pos_enc_class(output_size, positional_dropout_rate),
+            )
+        elif input_layer == "conv2d":
+            self.embed = Conv2dSubsampling(input_size, output_size, dropout_rate)
+        elif input_layer == "conv2d2":
+            self.embed = Conv2dSubsampling2(input_size, output_size, dropout_rate)
+        elif input_layer == "conv2d6":
+            self.embed = Conv2dSubsampling6(input_size, output_size, dropout_rate)
+        elif input_layer == "conv2d8":
+            self.embed = Conv2dSubsampling8(input_size, output_size, dropout_rate)
+        elif input_layer == "embed":
+            self.embed = torch.nn.Sequential(
+                torch.nn.Embedding(input_size, output_size, padding_idx=padding_idx),
+                SinusoidalPositionEncoder(),
+            )
+        elif input_layer is None:
+            if input_size == output_size:
+                self.embed = None
+            else:
+                self.embed = torch.nn.Linear(input_size, output_size)
+        elif input_layer == "pe":
+            self.embed = SinusoidalPositionEncoder()
+        elif input_layer == "pe_online":
+            self.embed = StreamSinusoidalPositionEncoder()
+        else:
+            raise ValueError("unknown input_layer: " + input_layer)
+        self.normalize_before = normalize_before
+        if positionwise_layer_type == "linear":
+            positionwise_layer = PositionwiseFeedForward
+            positionwise_layer_args = (
+                output_size,
+                linear_units,
+                dropout_rate,
+            )
+        elif positionwise_layer_type == "conv1d":
+            positionwise_layer = MultiLayeredConv1d
+            positionwise_layer_args = (
+                output_size,
+                linear_units,
+                positionwise_conv_kernel_size,
+                dropout_rate,
+            )
+        elif positionwise_layer_type == "conv1d-linear":
+            positionwise_layer = Conv1dLinear
+            positionwise_layer_args = (
+                output_size,
+                linear_units,
+                positionwise_conv_kernel_size,
+                dropout_rate,
+            )
+        else:
+            raise NotImplementedError("Support only linear or conv1d.")
+
+        if selfattention_layer_type == "selfattn":
+            encoder_selfattn_layer = MultiHeadedAttention
+            encoder_selfattn_layer_args = (
+                attention_heads,
+                output_size,
+                attention_dropout_rate,
+            )
+
+        elif selfattention_layer_type == "sanm":
+            encoder_selfattn_layer = MultiHeadedAttentionSANM
+            encoder_selfattn_layer_args0 = (
+                attention_heads,
+                input_size,
+                output_size,
+                attention_dropout_rate,
+                kernel_size,
+                sanm_shfit,
+                lora_list,
+                lora_rank,
+                lora_alpha,
+                lora_dropout,
+            )
+
+            encoder_selfattn_layer_args = (
+                attention_heads,
+                output_size,
+                output_size,
+                attention_dropout_rate,
+                kernel_size,
+                sanm_shfit,
+                lora_list,
+                lora_rank,
+                lora_alpha,
+                lora_dropout,
+            )
+        self.encoders0 = repeat(
+            1,
+            lambda lnum: EncoderLayerSANM(
+                input_size,
+                output_size,
+                encoder_selfattn_layer(*encoder_selfattn_layer_args0),
+                positionwise_layer(*positionwise_layer_args),
+                dropout_rate,
+                normalize_before,
+                concat_after,
+            ),
+        )
+
+        self.encoders = repeat(
+            num_blocks - 1,
+            lambda lnum: EncoderLayerSANM(
+                output_size,
+                output_size,
+                encoder_selfattn_layer(*encoder_selfattn_layer_args),
+                positionwise_layer(*positionwise_layer_args),
+                dropout_rate,
+                normalize_before,
+                concat_after,
+            ),
+        )
+        if self.normalize_before:
+            self.after_norm = LayerNorm(output_size)
+
+        self.interctc_layer_idx = interctc_layer_idx
+        if len(interctc_layer_idx) > 0:
+            assert 0 < min(interctc_layer_idx) and max(interctc_layer_idx) < num_blocks
+        self.interctc_use_conditioning = interctc_use_conditioning
+        self.conditioning_layer = None
+        self.dropout = nn.Dropout(dropout_rate)
+        self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch
+        self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf
+
+    def output_size(self) -> int:
+        return self._output_size
+
+    def forward(
+        self,
+        xs_pad: torch.Tensor,
+        ilens: torch.Tensor,
+        prev_states: torch.Tensor = None,
+        ctc: CTC = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        """Embed positions in tensor.
+
+        Args:
+            xs_pad: input tensor (B, L, D)
+            ilens: input length (B)
+            prev_states: Not to be used now.
+        Returns:
+            position embedded tensor and mask
+        """
+        masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device)
+        xs_pad = xs_pad * self.output_size() ** 0.5
+        if self.embed is None:
+            xs_pad = xs_pad
+        elif (
+            isinstance(self.embed, Conv2dSubsampling)
+            or isinstance(self.embed, Conv2dSubsampling2)
+            or isinstance(self.embed, Conv2dSubsampling6)
+            or isinstance(self.embed, Conv2dSubsampling8)
+        ):
+            short_status, limit_size = check_short_utt(self.embed, xs_pad.size(1))
+            if short_status:
+                raise TooShortUttError(
+                    f"has {xs_pad.size(1)} frames and is too short for subsampling "
+                    + f"(it needs more than {limit_size} frames), return empty results",
+                    xs_pad.size(1),
+                    limit_size,
+                )
+            xs_pad, masks = self.embed(xs_pad, masks)
+        else:
+            xs_pad = self.embed(xs_pad)
+
+        # xs_pad = self.dropout(xs_pad)
+        encoder_outs = self.encoders0(xs_pad, masks)
+        xs_pad, masks = encoder_outs[0], encoder_outs[1]
+        intermediate_outs = []
+        if len(self.interctc_layer_idx) == 0:
+            encoder_outs = self.encoders(xs_pad, masks)
+            xs_pad, masks = encoder_outs[0], encoder_outs[1]
+        else:
+            for layer_idx, encoder_layer in enumerate(self.encoders):
+                encoder_outs = encoder_layer(xs_pad, masks)
+                xs_pad, masks = encoder_outs[0], encoder_outs[1]
+
+                if layer_idx + 1 in self.interctc_layer_idx:
+                    encoder_out = xs_pad
+
+                    # intermediate outputs are also normalized
+                    if self.normalize_before:
+                        encoder_out = self.after_norm(encoder_out)
+
+                    intermediate_outs.append((layer_idx + 1, encoder_out))
+
+                    if self.interctc_use_conditioning:
+                        ctc_out = ctc.softmax(encoder_out)
+                        xs_pad = xs_pad + self.conditioning_layer(ctc_out)
+
+        if self.normalize_before:
+            xs_pad = self.after_norm(xs_pad)
+
+        olens = masks.squeeze(1).sum(1)
+        if len(intermediate_outs) > 0:
+            return (xs_pad, intermediate_outs), olens, None
+        return xs_pad, olens, None
+
+    def _add_overlap_chunk(self, feats: np.ndarray, cache: dict = {}):
+        if len(cache) == 0:
+            return feats
+        cache["feats"] = to_device(cache["feats"], device=feats.device)
+        overlap_feats = torch.cat((cache["feats"], feats), dim=1)
+        cache["feats"] = overlap_feats[:, -(cache["chunk_size"][0] + cache["chunk_size"][2]) :, :]
+        return overlap_feats
+
+    def forward_chunk(
+        self,
+        xs_pad: torch.Tensor,
+        ilens: torch.Tensor,
+        cache: dict = None,
+        ctc: CTC = None,
+    ):
+        xs_pad *= self.output_size() ** 0.5
+        if self.embed is None:
+            xs_pad = xs_pad
+        else:
+            xs_pad = self.embed(xs_pad, cache)
+        if cache["tail_chunk"]:
+            xs_pad = to_device(cache["feats"], device=xs_pad.device)
+        else:
+            xs_pad = self._add_overlap_chunk(xs_pad, cache)
+        encoder_outs = self.encoders0(xs_pad, None, None, None, None)
+        xs_pad, masks = encoder_outs[0], encoder_outs[1]
+        intermediate_outs = []
+        if len(self.interctc_layer_idx) == 0:
+            encoder_outs = self.encoders(xs_pad, None, None, None, None)
+            xs_pad, masks = encoder_outs[0], encoder_outs[1]
+        else:
+            for layer_idx, encoder_layer in enumerate(self.encoders):
+                encoder_outs = encoder_layer(xs_pad, None, None, None, None)
+                xs_pad, masks = encoder_outs[0], encoder_outs[1]
+                if layer_idx + 1 in self.interctc_layer_idx:
+                    encoder_out = xs_pad
+
+                    # intermediate outputs are also normalized
+                    if self.normalize_before:
+                        encoder_out = self.after_norm(encoder_out)
+
+                    intermediate_outs.append((layer_idx + 1, encoder_out))
+
+                    if self.interctc_use_conditioning:
+                        ctc_out = ctc.softmax(encoder_out)
+                        xs_pad = xs_pad + self.conditioning_layer(ctc_out)
+
+        if self.normalize_before:
+            xs_pad = self.after_norm(xs_pad)
+
+        if len(intermediate_outs) > 0:
+            return (xs_pad, intermediate_outs), None, None
+        return xs_pad, ilens, None
+
+
+class EncoderLayerSANMExport(nn.Module):
+    def __init__(
+        self,
+        model,
+    ):
+        """Construct an EncoderLayer object."""
+        super().__init__()
+        self.self_attn = model.self_attn
+        self.feed_forward = model.feed_forward
+        self.norm1 = model.norm1
+        self.norm2 = model.norm2
+        self.in_size = model.in_size
+        self.size = model.size
+
+    def forward(self, x, mask):
+
+        residual = x
+        x = self.norm1(x)
+        x = self.self_attn(x, mask)
+        if self.in_size == self.size:
+            x = x + residual
+        residual = x
+        x = self.norm2(x)
+        x = self.feed_forward(x)
+        x = x + residual
+
+        return x, mask
+
+
+@tables.register("encoder_classes", "SANMEncoderChunkOptExport")
+@tables.register("encoder_classes", "SANMEncoderExport")
+class SANMEncoderExport(nn.Module):
+    def __init__(
+        self,
+        model,
+        max_seq_len=512,
+        feats_dim=560,
+        model_name="encoder",
+        onnx: bool = True,
+        ctc_linear: nn.Module = None,
+    ):
+        super().__init__()
+        self.embed = model.embed
+        if isinstance(self.embed, StreamSinusoidalPositionEncoder):
+            self.embed = None
+        self.model = model
+        self.feats_dim = feats_dim
+        self._output_size = model._output_size
+
+        from funasr.utils.torch_function import sequence_mask
+
+        self.make_pad_mask = sequence_mask(max_seq_len, flip=False)
+
+        from funasr.models.sanm.attention import MultiHeadedAttentionSANMExport
+
+        if hasattr(model, "encoders0"):
+            for i, d in enumerate(self.model.encoders0):
+                if isinstance(d.self_attn, MultiHeadedAttentionSANM):
+                    d.self_attn = MultiHeadedAttentionSANMExport(d.self_attn)
+                self.model.encoders0[i] = EncoderLayerSANMExport(d)
+
+        for i, d in enumerate(self.model.encoders):
+            if isinstance(d.self_attn, MultiHeadedAttentionSANM):
+                d.self_attn = MultiHeadedAttentionSANMExport(d.self_attn)
+            self.model.encoders[i] = EncoderLayerSANMExport(d)
+
+        self.model_name = model_name
+        self.num_heads = model.encoders[0].self_attn.h
+        self.hidden_size = model.encoders[0].self_attn.linear_out.out_features
+
+        self.ctc_linear = ctc_linear
+
+    def prepare_mask(self, mask):
+        mask_3d_btd = mask[:, :, None]
+        if len(mask.shape) == 2:
+            mask_4d_bhlt = 1 - mask[:, None, None, :]
+        elif len(mask.shape) == 3:
+            mask_4d_bhlt = 1 - mask[:, None, :]
+        mask_4d_bhlt = mask_4d_bhlt * -10000.0
+
+        return mask_3d_btd, mask_4d_bhlt
+
+    def forward(self, speech: torch.Tensor, speech_lengths: torch.Tensor, online: bool = False):
+        if not online:
+            speech = speech * self._output_size**0.5
+
+        mask = self.make_pad_mask(speech_lengths)
+        mask = self.prepare_mask(mask)
+        if self.embed is None:
+            xs_pad = speech
+        else:
+            xs_pad = self.embed(speech)
+
+        encoder_outs = self.model.encoders0(xs_pad, mask)
+        xs_pad, masks = encoder_outs[0], encoder_outs[1]
+
+        encoder_outs = self.model.encoders(xs_pad, mask)
+        xs_pad, masks = encoder_outs[0], encoder_outs[1]
+
+        xs_pad = self.model.after_norm(xs_pad)
+
+        if self.ctc_linear is not None:
+            xs_pad = self.ctc_linear(xs_pad)
+            xs_pad = F.softmax(xs_pad, dim=2)
+
+        return xs_pad, speech_lengths
+
+    def get_output_size(self):
+        return self.model.encoders[0].size
+
+    def get_dummy_inputs(self):
+        feats = torch.randn(1, 100, self.feats_dim)
+        return feats
+
+    def get_input_names(self):
+        return ["feats"]
+
+    def get_output_names(self):
+        return ["encoder_out", "encoder_out_lens", "predictor_weight"]
+
+    def get_dynamic_axes(self):
+        return {
+            "feats": {1: "feats_length"},
+            "encoder_out": {1: "enc_out_length"},
+            "predictor_weight": {1: "pre_out_length"},
+        }

funasr/models/sanm/model.py

@@ -0,0 +1,28 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import logging
+
+import torch
+
+from funasr.models.transformer.model import Transformer
+from funasr.register import tables
+
+
+@tables.register("model_classes", "SANM")
+class SANM(Transformer):
+    """
+    Author: Zhifu Gao, Shiliang Zhang, Ming Lei, Ian McLoughlin
+    San-m: Memory equipped self-attention for end-to-end speech recognition
+    https://arxiv.org/abs/2006.01713
+    """
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+
+        super().__init__(*args, **kwargs)

funasr/models/sanm/multihead_att.py

@@ -0,0 +1,241 @@
+import os
+import math
+
+import torch
+import torch.nn as nn
+
+
+class MultiHeadedAttentionSANMExport(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.d_k = model.d_k
+        self.h = model.h
+        self.linear_out = model.linear_out
+        self.linear_q_k_v = model.linear_q_k_v
+        self.fsmn_block = model.fsmn_block
+        self.pad_fn = model.pad_fn
+
+        self.attn = None
+        self.all_head_size = self.h * self.d_k
+
+    def forward(self, x, mask):
+        mask_3d_btd, mask_4d_bhlt = mask
+        q_h, k_h, v_h, v = self.forward_qkv(x)
+        fsmn_memory = self.forward_fsmn(v, mask_3d_btd)
+        q_h = q_h * self.d_k ** (-0.5)
+        scores = torch.matmul(q_h, k_h.transpose(-2, -1))
+        att_outs = self.forward_attention(v_h, scores, mask_4d_bhlt)
+        return att_outs + fsmn_memory
+
+    def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
+        new_x_shape = x.size()[:-1] + (self.h, self.d_k)
+        x = x.view(new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward_qkv(self, x):
+        q_k_v = self.linear_q_k_v(x)
+        q, k, v = torch.split(q_k_v, int(self.h * self.d_k), dim=-1)
+        q_h = self.transpose_for_scores(q)
+        k_h = self.transpose_for_scores(k)
+        v_h = self.transpose_for_scores(v)
+        return q_h, k_h, v_h, v
+
+    def forward_fsmn(self, inputs, mask):
+        # b, t, d = inputs.size()
+        # mask = torch.reshape(mask, (b, -1, 1))
+        inputs = inputs * mask
+        x = inputs.transpose(1, 2)
+        x = self.pad_fn(x)
+        x = self.fsmn_block(x)
+        x = x.transpose(1, 2)
+        x = x + inputs
+        x = x * mask
+        return x
+
+    def forward_attention(self, value, scores, mask):
+        scores = scores + mask
+
+        attn = torch.softmax(scores, dim=-1)
+        context_layer = torch.matmul(attn, value)  # (batch, head, time1, d_k)
+
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(new_context_layer_shape)
+        return self.linear_out(context_layer)  # (batch, time1, d_model)
+
+
+def preprocess_for_attn(x, mask, cache, pad_fn, kernel_size):
+    x = x * mask
+    x = x.transpose(1, 2)
+    if cache is None:
+        x = pad_fn(x)
+    else:
+        x = torch.cat((cache, x), dim=2)
+        cache = x[:, :, -(kernel_size - 1) :]
+    return x, cache
+
+
+torch_version = tuple([int(i) for i in torch.__version__.split(".")[:2]])
+if torch_version >= (1, 8):
+    import torch.fx
+
+    torch.fx.wrap("preprocess_for_attn")
+
+
+class MultiHeadedAttentionSANMDecoderExport(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.fsmn_block = model.fsmn_block
+        self.pad_fn = model.pad_fn
+        self.kernel_size = model.kernel_size
+        self.attn = None
+
+    def forward(self, inputs, mask, cache=None):
+        x, cache = preprocess_for_attn(inputs, mask, cache, self.pad_fn, self.kernel_size)
+        x = self.fsmn_block(x)
+        x = x.transpose(1, 2)
+
+        x = x + inputs
+        x = x * mask
+        return x, cache
+
+
+class MultiHeadedAttentionCrossAttExport(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.d_k = model.d_k
+        self.h = model.h
+        self.linear_q = model.linear_q
+        self.linear_k_v = model.linear_k_v
+        self.linear_out = model.linear_out
+        self.attn = None
+        self.all_head_size = self.h * self.d_k
+
+    def forward(self, x, memory, memory_mask):
+        q, k, v = self.forward_qkv(x, memory)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
+        return self.forward_attention(v, scores, memory_mask)
+
+    def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
+        new_x_shape = x.size()[:-1] + (self.h, self.d_k)
+        x = x.view(new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward_qkv(self, x, memory):
+        q = self.linear_q(x)
+
+        k_v = self.linear_k_v(memory)
+        k, v = torch.split(k_v, int(self.h * self.d_k), dim=-1)
+        q = self.transpose_for_scores(q)
+        k = self.transpose_for_scores(k)
+        v = self.transpose_for_scores(v)
+        return q, k, v
+
+    def forward_attention(self, value, scores, mask):
+        scores = scores + mask
+
+        attn = torch.softmax(scores, dim=-1)
+        context_layer = torch.matmul(attn, value)  # (batch, head, time1, d_k)
+
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(new_context_layer_shape)
+        return self.linear_out(context_layer)  # (batch, time1, d_model)
+
+
+class OnnxMultiHeadedAttention(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.d_k = model.d_k
+        self.h = model.h
+        self.linear_q = model.linear_q
+        self.linear_k = model.linear_k
+        self.linear_v = model.linear_v
+        self.linear_out = model.linear_out
+        self.attn = None
+        self.all_head_size = self.h * self.d_k
+
+    def forward(self, query, key, value, mask):
+        q, k, v = self.forward_qkv(query, key, value)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
+        return self.forward_attention(v, scores, mask)
+
+    def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
+        new_x_shape = x.size()[:-1] + (self.h, self.d_k)
+        x = x.view(new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward_qkv(self, query, key, value):
+        q = self.linear_q(query)
+        k = self.linear_k(key)
+        v = self.linear_v(value)
+        q = self.transpose_for_scores(q)
+        k = self.transpose_for_scores(k)
+        v = self.transpose_for_scores(v)
+        return q, k, v
+
+    def forward_attention(self, value, scores, mask):
+        scores = scores + mask
+
+        attn = torch.softmax(scores, dim=-1)
+        context_layer = torch.matmul(attn, value)  # (batch, head, time1, d_k)
+
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(new_context_layer_shape)
+        return self.linear_out(context_layer)  # (batch, time1, d_model)
+
+
+class OnnxRelPosMultiHeadedAttention(OnnxMultiHeadedAttention):
+    def __init__(self, model):
+        super().__init__(model)
+        self.linear_pos = model.linear_pos
+        self.pos_bias_u = model.pos_bias_u
+        self.pos_bias_v = model.pos_bias_v
+
+    def forward(self, query, key, value, pos_emb, mask):
+        q, k, v = self.forward_qkv(query, key, value)
+        q = q.transpose(1, 2)  # (batch, time1, head, d_k)
+
+        p = self.transpose_for_scores(self.linear_pos(pos_emb))  # (batch, head, time1, d_k)
+
+        # (batch, head, time1, d_k)
+        q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2)
+        # (batch, head, time1, d_k)
+        q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2)
+
+        # compute attention score
+        # first compute matrix a and matrix c
+        # as described in https://arxiv.org/abs/1901.02860 Section 3.3
+        # (batch, head, time1, time2)
+        matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
+
+        # compute matrix b and matrix d
+        # (batch, head, time1, time1)
+        matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
+        matrix_bd = self.rel_shift(matrix_bd)
+
+        scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k)  # (batch, head, time1, time2)
+
+        return self.forward_attention(v, scores, mask)
+
+    def rel_shift(self, x):
+        zero_pad = torch.zeros((*x.size()[:3], 1), device=x.device, dtype=x.dtype)
+        x_padded = torch.cat([zero_pad, x], dim=-1)
+
+        x_padded = x_padded.view(*x.size()[:2], x.size(3) + 1, x.size(2))
+        x = x_padded[:, :, 1:].view_as(x)[
+            :, :, :, : x.size(-1) // 2 + 1
+        ]  # only keep the positions from 0 to time2
+        return x
+
+    def forward_attention(self, value, scores, mask):
+        scores = scores + mask
+
+        attn = torch.softmax(scores, dim=-1)
+        context_layer = torch.matmul(attn, value)  # (batch, head, time1, d_k)
+
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(new_context_layer_shape)
+        return self.linear_out(context_layer)  # (batch, time1, d_model)

funasr/models/sanm/positionwise_feed_forward.py

@@ -0,0 +1,33 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+
+"""Positionwise feed forward layer definition."""
+
+import torch
+
+from funasr.models.transformer.layer_norm import LayerNorm
+
+
+class PositionwiseFeedForwardDecoderSANM(torch.nn.Module):
+    """Positionwise feed forward layer.
+
+    Args:
+        idim (int): Input dimenstion.
+        hidden_units (int): The number of hidden units.
+        dropout_rate (float): Dropout rate.
+
+    """
+
+    def __init__(self, idim, hidden_units, dropout_rate, adim=None, activation=torch.nn.ReLU()):
+        """Construct an PositionwiseFeedForward object."""
+        super(PositionwiseFeedForwardDecoderSANM, self).__init__()
+        self.w_1 = torch.nn.Linear(idim, hidden_units)
+        self.w_2 = torch.nn.Linear(hidden_units, idim if adim is None else adim, bias=False)
+        self.dropout = torch.nn.Dropout(dropout_rate)
+        self.activation = activation
+        self.norm = LayerNorm(hidden_units)
+
+    def forward(self, x):
+        """Forward function."""
+        return self.w_2(self.norm(self.dropout(self.activation(self.w_1(x)))))

funasr/models/sanm/template.yaml

@@ -0,0 +1,121 @@
+# This is an example that demonstrates how to configure a model file.
+# You can modify the configuration according to your own requirements.
+
+# to print the register_table:
+# from funasr.register import tables
+# tables.print()
+
+# network architecture
+model: SANM
+model_conf:
+    ctc_weight: 0.0
+    lsm_weight: 0.1
+    length_normalized_loss: true
+
+# encoder
+encoder: SANMEncoder
+encoder_conf:
+    output_size: 512
+    attention_heads: 4
+    linear_units: 2048
+    num_blocks: 50
+    dropout_rate: 0.1
+    positional_dropout_rate: 0.1
+    attention_dropout_rate: 0.1
+    input_layer: pe
+    pos_enc_class: SinusoidalPositionEncoder
+    normalize_before: true
+    kernel_size: 11
+    sanm_shfit: 0
+    selfattention_layer_type: sanm
+
+# decoder
+decoder: FsmnDecoder
+decoder_conf:
+    attention_heads: 4
+    linear_units: 2048
+    num_blocks: 16
+    dropout_rate: 0.1
+    positional_dropout_rate: 0.1
+    self_attention_dropout_rate: 0.1
+    src_attention_dropout_rate: 0.1
+    att_layer_num: 16
+    kernel_size: 11
+    sanm_shfit: 0
+
+
+
+# frontend related
+frontend: WavFrontend
+frontend_conf:
+    fs: 16000
+    window: hamming
+    n_mels: 80
+    frame_length: 25
+    frame_shift: 10
+    lfr_m: 7
+    lfr_n: 6
+
+specaug: SpecAugLFR
+specaug_conf:
+    apply_time_warp: false
+    time_warp_window: 5
+    time_warp_mode: bicubic
+    apply_freq_mask: true
+    freq_mask_width_range:
+    - 0
+    - 30
+    lfr_rate: 6
+    num_freq_mask: 1
+    apply_time_mask: true
+    time_mask_width_range:
+    - 0
+    - 12
+    num_time_mask: 1
+
+train_conf:
+  accum_grad: 1
+  grad_clip: 5
+  max_epoch: 150
+  val_scheduler_criterion:
+      - valid
+      - acc
+  best_model_criterion:
+  -   - valid
+      - acc
+      - max
+  keep_nbest_models: 10
+  avg_nbest_model: 10
+  log_interval: 50
+
+optim: adam
+optim_conf:
+   lr: 0.0005
+scheduler: warmuplr
+scheduler_conf:
+   warmup_steps: 30000
+
+dataset: AudioDataset
+dataset_conf:
+    index_ds: IndexDSJsonl
+    batch_sampler: BatchSampler
+    batch_type: example # example or length
+    batch_size: 1 # if batch_type is example, batch_size is the numbers of samples; if length, batch_size is source_token_len+target_token_len;
+    max_token_length: 2048 # filter samples if source_token_len+target_token_len > max_token_length,
+    buffer_size: 500
+    shuffle: True
+    num_workers: 0
+
+tokenizer: CharTokenizer
+tokenizer_conf:
+  unk_symbol: <unk>
+  split_with_space: true
+
+
+ctc_conf:
+    dropout_rate: 0.0
+    ctc_type: builtin
+    reduce: true
+    ignore_nan_grad: true
+
+normalize: null

funasr/models/sanm_kws/export_meta.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import types
+import copy
+import torch
+from funasr.register import tables
+
+
+def export_rebuild_model(model, **kwargs):
+    # self.device = kwargs.get("device")
+    is_onnx = kwargs.get("type", "onnx") == "onnx"
+    encoder_class = tables.encoder_classes.get(kwargs["encoder"] + "Export")
+
+    if hasattr(model, "ctc"):
+        model.encoder = encoder_class(
+            model.encoder,
+            onnx=is_onnx,
+            feats_dim=kwargs.get("input_size", 560),
+            ctc_linear=model.ctc.ctc_lo
+        )
+    else:
+        assert False, print(model)
+        model.encoder = encoder_class(model.encoder, onnx=is_onnx, feats_dim=kwargs.get("input_size", 560))
+
+    # from funasr.utils.torch_function import sequence_mask
+    # model.make_pad_mask = sequence_mask(max_seq_len=None, flip=False)
+
+    encoder_model = copy.copy(model)
+
+    # encoder
+    encoder_model.forward = types.MethodType(export_encoder_forward, encoder_model)
+    encoder_model.export_dummy_inputs = types.MethodType(export_encoder_dummy_inputs, encoder_model)
+    encoder_model.export_input_names = types.MethodType(export_encoder_input_names, encoder_model)
+    encoder_model.export_output_names = types.MethodType(export_encoder_output_names, encoder_model)
+    encoder_model.export_dynamic_axes = types.MethodType(export_encoder_dynamic_axes, encoder_model)
+    encoder_model.export_name = types.MethodType(export_encoder_name, encoder_model)
+
+    return encoder_model
+
+
+def export_encoder_forward(
+    self,
+    speech: torch.Tensor,
+    speech_lengths: torch.Tensor,
+):
+    # a. To device
+    batch = {
+        "speech": speech,
+        "speech_lengths": speech_lengths,
+        "online": True
+    }
+    # batch = to_device(batch, device=self.device)
+
+    encoder_out, encoder_out_len = self.encoder(**batch)
+    # mask = self.make_pad_mask(encoder_out_len)[:, None, :]
+    # alphas, _ = self.predictor.forward_cnn(enc, mask)
+
+    # return encoder_out, encoder_out_len, alphas
+    return encoder_out, encoder_out_len
+
+
+def export_encoder_dummy_inputs(self):
+    speech = torch.randn(2, 30, 280)
+    speech_lengths = torch.tensor([6, 30], dtype=torch.int32)
+    return (speech, speech_lengths)
+
+
+def export_encoder_input_names(self):
+    return ["speech", "speech_lengths"]
+
+
+def export_encoder_output_names(self):
+    # return ["encoder_out", "encoder_out_len", "alphas"]
+    return ["encoder_out", "encoder_out_len"]
+
+
+def export_encoder_dynamic_axes(self):
+    return {
+        "speech": {
+            0: "batch_size", 1: "feats_length"
+        },
+        "speech_lengths": {
+            0: "batch_size",
+        },
+        "encoder_out": {
+            0: "batch_size", 1: "feats_length"
+        },
+        "encoder_out_len": {
+            0: "batch_size",
+        },
+    }
+
+
+def export_encoder_name(self):
+    return "encoder.onnx"

funasr/models/sanm_kws/model.py

@@ -0,0 +1,266 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import time
+import torch
+import logging
+from torch.cuda.amp import autocast
+from typing import Union, Dict, List, Tuple, Optional
+
+from funasr.register import tables
+from funasr.models.ctc.ctc import CTC
+from funasr.utils import postprocess_utils
+from funasr.metrics.compute_acc import th_accuracy
+from funasr.train_utils.device_funcs import to_device
+from funasr.utils.datadir_writer import DatadirWriter
+from funasr.models.paraformer.search import Hypothesis
+from funasr.models.paraformer.cif_predictor import mae_loss
+from funasr.train_utils.device_funcs import force_gatherable
+from funasr.losses.label_smoothing_loss import LabelSmoothingLoss
+from funasr.models.transformer.utils.add_sos_eos import add_sos_eos
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
+
+
+@tables.register("model_classes", "SanmKWS")
+class SanmKWS(torch.nn.Module):
+    """
+    Author: Speech Lab of DAMO Academy, Alibaba Group
+    Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition
+    https://arxiv.org/abs/2206.08317
+    """
+
+    def __init__(
+        self,
+        specaug: Optional[str] = None,
+        specaug_conf: Optional[Dict] = None,
+        normalize: str = None,
+        normalize_conf: Optional[Dict] = None,
+        encoder: str = None,
+        encoder_conf: Optional[Dict] = None,
+        ctc: str = None,
+        ctc_conf: Optional[Dict] = None,
+        ctc_weight: float = 1.0,
+        input_size: int = 360,
+        vocab_size: int = -1,
+        ignore_id: int = -1,
+        blank_id: int = 0,
+        sos: int = 1,
+        eos: int = 2,
+        **kwargs,
+    ):
+
+        super().__init__()
+
+        if specaug is not None:
+            specaug_class = tables.specaug_classes.get(specaug)
+            specaug = specaug_class(**specaug_conf)
+
+        if normalize is not None:
+            normalize_class = tables.normalize_classes.get(normalize)
+            normalize = normalize_class(**normalize_conf)
+
+        encoder_class = tables.encoder_classes.get(encoder)
+        encoder = encoder_class(input_size=input_size, **encoder_conf)
+        encoder_output_size = encoder.output_size()
+
+        if ctc_conf is None:
+            ctc_conf = {}
+        ctc = CTC(odim=vocab_size, encoder_output_size=encoder_output_size, **ctc_conf)
+
+        # note that eos is the same as sos (equivalent ID)
+        self.blank_id = blank_id
+        self.sos = sos if sos is not None else vocab_size - 1
+        self.eos = eos if eos is not None else vocab_size - 1
+        self.vocab_size = vocab_size
+        self.ignore_id = ignore_id
+        self.ctc_weight = ctc_weight
+        # self.token_list = token_list.copy()
+        #
+        # self.frontend = frontend
+        self.specaug = specaug
+        self.normalize = normalize
+        self.encoder = encoder
+
+        self.ctc = ctc
+        self.error_calculator = None
+
+    def forward(
+        self,
+        speech: torch.Tensor,
+        speech_lengths: torch.Tensor,
+        text: torch.Tensor,
+        text_lengths: torch.Tensor,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:
+        """Encoder + Decoder + Calc loss
+        Args:
+                speech: (Batch, Length, ...)
+                speech_lengths: (Batch, )
+                text: (Batch, Length)
+                text_lengths: (Batch,)
+        """
+        if len(text_lengths.size()) > 1:
+            text_lengths = text_lengths[:, 0]
+        if len(speech_lengths.size()) > 1:
+            speech_lengths = speech_lengths[:, 0]
+        batch_size = speech.shape[0]
+
+        # Encoder
+        encoder_out, encoder_out_lens = self.encode(speech, speech_lengths)
+
+        # decoder: CTC branch
+        loss_ctc, cer_ctc = self._calc_ctc_loss(
+            encoder_out, encoder_out_lens, text, text_lengths
+        )
+
+        # Collect CTC branch stats
+        stats = dict()
+        stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None
+        stats["cer_ctc"] = cer_ctc
+        stats["cer"] = cer_ctc
+
+        loss = loss_ctc
+
+        stats["loss"] = torch.clone(loss.detach())
+        stats["batch_size"] = batch_size
+
+        # force_gatherable: to-device and to-tensor if scalar for DataParallel
+        loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device)
+        return loss, stats, weight
+
+    def encode(
+        self,
+        speech: torch.Tensor,
+        speech_lengths: torch.Tensor,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Encoder. Note that this method is used by asr_inference.py
+        Args:
+                speech: (Batch, Length, ...)
+                speech_lengths: (Batch, )
+                ind: int
+        """
+        with autocast(False):
+            # Data augmentation
+            if self.specaug is not None and self.training:
+                speech, speech_lengths = self.specaug(speech, speech_lengths)
+
+            # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN
+            if self.normalize is not None:
+                speech, speech_lengths = self.normalize(speech, speech_lengths)
+
+        # Forward encoder
+        encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths)
+        if isinstance(encoder_out, tuple):
+            encoder_out = encoder_out[0]
+
+        return encoder_out, encoder_out_lens
+
+    def _calc_ctc_loss(
+        self,
+        encoder_out: torch.Tensor,
+        encoder_out_lens: torch.Tensor,
+        ys_pad: torch.Tensor,
+        ys_pad_lens: torch.Tensor,
+    ):
+        # Calc CTC loss
+        loss_ctc = self.ctc(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens)
+
+        # Calc CER using CTC
+        cer_ctc = None
+        if not self.training and self.error_calculator is not None:
+            ys_hat = self.ctc.argmax(encoder_out).data
+            cer_ctc = self.error_calculator(ys_hat.cpu(), ys_pad.cpu(), is_ctc=True)
+        return loss_ctc, cer_ctc
+
+    def inference(
+        self,
+        data_in,
+        data_lengths=None,
+        key: list = None,
+        tokenizer=None,
+        frontend=None,
+        **kwargs,
+    ):
+        keywords = kwargs.get("keywords")
+        from funasr.utils.kws_utils import KwsCtcPrefixDecoder
+        self.kws_decoder = KwsCtcPrefixDecoder(
+            ctc=self.ctc,
+            keywords=keywords,
+            token_list=tokenizer.token_list,
+            seg_dict=tokenizer.seg_dict,
+        )
+
+        meta_data = {}
+        if (
+            isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank"
+        ):  # fbank
+            speech, speech_lengths = data_in, data_lengths
+            if len(speech.shape) < 3:
+                speech = speech[None, :, :]
+            if speech_lengths is not None:
+                speech_lengths = speech_lengths.squeeze(-1)
+            else:
+                speech_lengths = speech.shape[1]
+        else:
+            # extract fbank feats
+            time1 = time.perf_counter()
+            audio_sample_list = load_audio_text_image_video(
+                data_in,
+                fs=frontend.fs,
+                audio_fs=kwargs.get("fs", 16000),
+                data_type=kwargs.get("data_type", "sound"),
+                tokenizer=tokenizer,
+            )
+            time2 = time.perf_counter()
+            meta_data["load_data"] = f"{time2 - time1:0.3f}"
+            speech, speech_lengths = extract_fbank(
+                audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend
+            )
+            time3 = time.perf_counter()
+            meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
+            meta_data["batch_data_time"] = (
+                speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000
+            )
+
+        speech = speech.to(device=kwargs["device"])
+        speech_lengths = speech_lengths.to(device=kwargs["device"])
+        # Encoder
+        if kwargs.get("fp16", False):
+            speech = speech.half()
+        encoder_out, encoder_out_lens = self.encode(speech, speech_lengths)
+        if isinstance(encoder_out, tuple):
+            encoder_out = encoder_out[0]
+
+        results = []
+        if kwargs.get("output_dir") is not None:
+            if not hasattr(self, "writer"):
+                self.writer = DatadirWriter(kwargs.get("output_dir"))
+
+        for i in range(encoder_out.size(0)):
+            x = encoder_out[i, : encoder_out_lens[i], :]
+            detect_result = self.kws_decoder.decode(x)
+            is_deted, det_keyword, det_score = detect_result[0], detect_result[1], detect_result[2]
+
+            if is_deted:
+                self.writer["detect"][key[i]] = "detected " + det_keyword + " " + str(det_score)
+                det_info = "detected " + det_keyword + " " + str(det_score)
+            else:
+                self.writer["detect"][key[i]] = "rejected"
+                det_info = "rejected"
+
+            result_i = {"key": key[i], "text": det_info}
+            results.append(result_i)
+
+        return results, meta_data
+
+    def export(self, **kwargs):
+        from .export_meta import export_rebuild_model
+
+        if "max_seq_len" not in kwargs:
+            kwargs["max_seq_len"] = 512
+        models = export_rebuild_model(model=self, **kwargs)
+        return models

funasr/models/sanm_kws_streaming/export_meta.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import types
+import copy
+import torch
+from funasr.register import tables
+
+
+def export_rebuild_model(model, **kwargs):
+    # self.device = kwargs.get("device")
+    is_onnx = kwargs.get("type", "onnx") == "onnx"
+    encoder_class = tables.encoder_classes.get(kwargs["encoder"] + "Export")
+
+    if hasattr(model, "ctc"):
+        model.encoder = encoder_class(
+            model.encoder,
+            onnx=is_onnx,
+            feats_dim=kwargs.get("input_size", 560),
+            ctc_linear=model.ctc.ctc_lo
+        )
+    else:
+        assert False, print(model)
+        model.encoder = encoder_class(model.encoder, onnx=is_onnx, feats_dim=kwargs.get("input_size", 560))
+
+    # from funasr.utils.torch_function import sequence_mask
+    # model.make_pad_mask = sequence_mask(max_seq_len=None, flip=False)
+
+    encoder_model = copy.copy(model)
+
+    # encoder
+    encoder_model.forward = types.MethodType(export_encoder_forward, encoder_model)
+    encoder_model.export_dummy_inputs = types.MethodType(export_encoder_dummy_inputs, encoder_model)
+    encoder_model.export_input_names = types.MethodType(export_encoder_input_names, encoder_model)
+    encoder_model.export_output_names = types.MethodType(export_encoder_output_names, encoder_model)
+    encoder_model.export_dynamic_axes = types.MethodType(export_encoder_dynamic_axes, encoder_model)
+    encoder_model.export_name = types.MethodType(export_encoder_name, encoder_model)
+
+    return encoder_model
+
+
+def export_encoder_forward(
+    self,
+    speech: torch.Tensor,
+    speech_lengths: torch.Tensor,
+):
+    # a. To device
+    batch = {
+        "speech": speech,
+        "speech_lengths": speech_lengths,
+        "online": True
+    }
+    # batch = to_device(batch, device=self.device)
+
+    encoder_out, encoder_out_len = self.encoder(**batch)
+    # mask = self.make_pad_mask(encoder_out_len)[:, None, :]
+    # alphas, _ = self.predictor.forward_cnn(enc, mask)
+
+    # return encoder_out, encoder_out_len, alphas
+    return encoder_out, encoder_out_len
+
+
+def export_encoder_dummy_inputs(self):
+    speech = torch.randn(2, 30, 280)
+    speech_lengths = torch.tensor([6, 30], dtype=torch.int32)
+    return (speech, speech_lengths)
+
+
+def export_encoder_input_names(self):
+    return ["speech", "speech_lengths"]
+
+
+def export_encoder_output_names(self):
+    # return ["encoder_out", "encoder_out_len", "alphas"]
+    return ["encoder_out", "encoder_out_len"]
+
+
+def export_encoder_dynamic_axes(self):
+    return {
+        "speech": {
+            0: "batch_size", 1: "feats_length"
+        },
+        "speech_lengths": {
+            0: "batch_size",
+        },
+        "encoder_out": {
+            0: "batch_size", 1: "feats_length"
+        },
+        "encoder_out_len": {
+            0: "batch_size",
+        },
+    }
+
+
+def export_encoder_name(self):
+    return "encoder.onnx"

funasr/models/sanm_kws_streaming/model.py

@@ -0,0 +1,442 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import time
+import torch
+import logging
+from typing import Dict, Tuple
+from contextlib import contextmanager
+from distutils.version import LooseVersion
+
+from funasr.register import tables
+from funasr.models.ctc.ctc import CTC
+from funasr.utils import postprocess_utils
+from funasr.metrics.compute_acc import th_accuracy
+from funasr.utils.datadir_writer import DatadirWriter
+from funasr.models.sanm_kws.model import SanmKWS
+from funasr.models.paraformer.search import Hypothesis
+from funasr.models.paraformer.cif_predictor import mae_loss
+from funasr.train_utils.device_funcs import force_gatherable
+from funasr.losses.label_smoothing_loss import LabelSmoothingLoss
+from funasr.models.transformer.utils.add_sos_eos import add_sos_eos
+from funasr.models.transformer.utils.nets_utils import make_pad_mask, pad_list
+from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
+
+
+if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
+    from torch.cuda.amp import autocast
+else:
+    # Nothing to do if torch<1.6.0
+    @contextmanager
+    def autocast(enabled=True):
+        yield
+
+
+@tables.register("model_classes", "SanmKWSStreaming")
+class SanmKWSStreaming(SanmKWS):
+    """
+    Author: Speech Lab of DAMO Academy, Alibaba Group
+    Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition
+    https://arxiv.org/abs/2206.08317
+    """
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+
+    def forward(
+        self,
+        speech: torch.Tensor,
+        speech_lengths: torch.Tensor,
+        text: torch.Tensor,
+        text_lengths: torch.Tensor,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:
+        """Encoder + Decoder + Calc loss
+        Args:
+                speech: (Batch, Length, ...)
+                speech_lengths: (Batch, )
+                text: (Batch, Length)
+                text_lengths: (Batch,)
+        """
+        decoding_ind = kwargs.get("decoding_ind")
+        if len(text_lengths.size()) > 1:
+            text_lengths = text_lengths[:, 0]
+        if len(speech_lengths.size()) > 1:
+            speech_lengths = speech_lengths[:, 0]
+
+        batch_size = speech.shape[0]
+
+        # Encoder
+        if hasattr(self.encoder, "overlap_chunk_cls"):
+            ind = self.encoder.overlap_chunk_cls.random_choice(self.training, decoding_ind)
+            encoder_out, encoder_out_lens = self.encode(speech, speech_lengths, ind=ind)
+        else:
+            encoder_out, encoder_out_lens = self.encode(speech, speech_lengths)
+
+        # decoder: CTC branch
+        if hasattr(self.encoder, "overlap_chunk_cls"):
+            encoder_out_ctc, encoder_out_lens_ctc = self.encoder.overlap_chunk_cls.remove_chunk(
+                encoder_out, encoder_out_lens, chunk_outs=None
+            )
+        else:
+            encoder_out_ctc, encoder_out_lens_ctc = encoder_out, encoder_out_lens
+
+        loss_ctc, cer_ctc = self._calc_ctc_loss(
+            encoder_out_ctc, encoder_out_lens_ctc, text, text_lengths
+        )
+
+        # Collect CTC branch stats
+        stats = dict()
+        stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None
+        stats["cer_ctc"] = cer_ctc
+
+        loss = loss_ctc
+
+        stats["cer"] = cer_ctc
+        stats["loss"] = torch.clone(loss.detach())
+
+        # force_gatherable: to-device and to-tensor if scalar for DataParallel
+        loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device)
+        return loss, stats, weight
+
+    def encode_chunk(
+        self,
+        speech: torch.Tensor,
+        speech_lengths: torch.Tensor,
+        cache: dict = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Frontend + Encoder. Note that this method is used by asr_inference.py
+        Args:
+                speech: (Batch, Length, ...)
+                speech_lengths: (Batch, )
+                ind: int
+        """
+        with autocast(False):
+            # Data augmentation
+            if self.specaug is not None and self.training:
+                speech, speech_lengths = self.specaug(speech, speech_lengths)
+
+            # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN
+            if self.normalize is not None:
+                speech, speech_lengths = self.normalize(speech, speech_lengths)
+
+        # Forward encoder
+        encoder_out, encoder_out_lens, _ = self.encoder.forward_chunk(
+            speech, speech_lengths, cache=cache["encoder"]
+        )
+
+        if isinstance(encoder_out, tuple):
+            encoder_out = encoder_out[0]
+
+        return encoder_out, torch.tensor([encoder_out.size(1)])
+
+    def init_cache(self, cache: dict = {}, **kwargs):
+        chunk_size = kwargs.get("chunk_size", [0, 10, 5])
+        encoder_chunk_look_back = kwargs.get("encoder_chunk_look_back", 0)
+        decoder_chunk_look_back = kwargs.get("decoder_chunk_look_back", 0)
+        batch_size = 1
+
+        enc_output_size = kwargs["encoder_conf"]["output_size"]
+        feats_dims = kwargs["frontend_conf"]["n_mels"] * kwargs["frontend_conf"]["lfr_m"]
+        cache_encoder = {
+            "start_idx": 0,
+            "cif_hidden": torch.zeros((batch_size, 1, enc_output_size)),
+            "cif_alphas": torch.zeros((batch_size, 1)),
+            "encoder_out": None,
+            "encoder_out_lens": None,
+            "chunk_size": chunk_size,
+            "encoder_chunk_look_back": encoder_chunk_look_back,
+            "last_chunk": False,
+            "opt": None,
+            "feats": torch.zeros((batch_size, chunk_size[0] + chunk_size[2], feats_dims)),
+            "tail_chunk": False,
+        }
+        cache["encoder"] = cache_encoder
+
+        cache_decoder = {
+            "decode_fsmn": None,
+            "decoder_chunk_look_back": decoder_chunk_look_back,
+            "opt": None,
+            "chunk_size": chunk_size,
+        }
+        cache["decoder"] = cache_decoder
+        cache["frontend"] = {}
+        cache["prev_samples"] = torch.empty(0)
+
+        return cache
+
+    def generate_chunk(
+        self,
+        speech,
+        speech_lengths=None,
+        key: list = None,
+        tokenizer=None,
+        frontend=None,
+        **kwargs,
+    ):
+        cache = kwargs.get("cache", {})
+        speech = speech.to(device=kwargs["device"])
+        speech_lengths = speech_lengths.to(device=kwargs["device"])
+
+        # Encoder
+        is_final = kwargs.get("is_final", False)
+        encoder_out, encoder_out_lens = self.encode_chunk(
+            speech, speech_lengths, cache=cache, is_final=is_final
+        )
+        if isinstance(encoder_out, tuple):
+            encoder_out = encoder_out[0]
+
+        chunk_size = cache["encoder"]["chunk_size"]
+        real_start_pos = chunk_size[0]
+
+        if encoder_out_lens[0] > chunk_size[0] + chunk_size[1] + chunk_size[2]:
+            assert False, print("impossible case 1 !")
+        if encoder_out_lens[0] == chunk_size[0] + chunk_size[1] + chunk_size[2]:
+            real_end_pos = chunk_size[0] + chunk_size[1]
+        elif encoder_out_lens[0] > chunk_size[0] + chunk_size[1]:
+            real_end_pos = chunk_size[0] + chunk_size[1]
+        elif encoder_out_lens[0] > chunk_size[0]:
+            real_end_pos = encoder_out_lens[0]
+        else:
+            assert False, print("impossible case 2 !")
+
+        encoder_out_accum = cache["encoder"]["encoder_out"]
+        if encoder_out_accum is not None:
+            encoder_out_accum = torch.cat((encoder_out_accum, encoder_out[:, real_start_pos:real_end_pos, :]), dim=1)
+        else:
+            encoder_out_accum = encoder_out[:, real_start_pos:real_end_pos, :]
+        cache["encoder"]["encoder_out"] = encoder_out_accum
+
+        if cache["encoder"]["encoder_out_lens"] is not None:
+            cache["encoder"]["encoder_out_lens"][0] += real_end_pos - real_start_pos
+        else:
+            cache["encoder"]["encoder_out_lens"] = encoder_out_lens
+            cache["encoder"]["encoder_out_lens"][0] = real_end_pos - real_start_pos
+
+        if is_final:
+            if kwargs.get("output_dir") is not None:
+                if not hasattr(self, "writer"):
+                    self.writer = DatadirWriter(kwargs.get("output_dir"))
+
+            results = []
+            for i in range(encoder_out_accum.size(0)):
+                x = encoder_out_accum[i, : cache["encoder"]["encoder_out_lens"][i], :]
+                detect_result = self.kws_decoder.decode(x)
+                is_deted, det_keyword, det_score = detect_result[0], detect_result[1], detect_result[2]
+
+                if is_deted:
+                    self.writer["detect"][key[i]] = "detected " + det_keyword + " " + str(det_score)
+                    det_info = "detected " + det_keyword + " " + str(det_score)
+                else:
+                    self.writer["detect"][key[i]] = "rejected"
+                    det_info = "rejected"
+
+                result_i = {"key": key[i], "text": det_info}
+                results.append(result_i)
+
+            return results
+        else:
+            return None
+
+    def inference(
+        self,
+        data_in,
+        data_lengths=None,
+        key: list = None,
+        tokenizer=None,
+        frontend=None,
+        cache: dict = {},
+        **kwargs,
+    ):
+        keywords = kwargs.get("keywords")
+        from funasr.utils.kws_utils import KwsCtcPrefixDecoder
+        self.kws_decoder = KwsCtcPrefixDecoder(
+            ctc=self.ctc,
+            keywords=keywords,
+            token_list=tokenizer.token_list,
+            seg_dict=tokenizer.seg_dict,
+        )
+
+        meta_data = {}
+        chunk_size = kwargs["chunk_size"]
+        chunk_stride_samples = int(chunk_size[1] * 960)  # 600ms
+        first_chunk_padding_samples = int(chunk_size[2] * 960)  # 600ms
+
+        if len(cache) == 0:
+            self.init_cache(cache, **kwargs)
+
+        time1 = time.perf_counter()
+        cfg = {"is_final": kwargs.get("is_final", False)}
+        audio_sample_list = load_audio_text_image_video(
+            data_in,
+            fs=frontend.fs,
+            audio_fs=kwargs.get("fs", 16000),
+            data_type=kwargs.get("data_type", "sound"),
+            tokenizer=tokenizer,
+            cache=cfg,
+        )
+        _is_final = cfg["is_final"]  # if data_in is a file or url, set is_final=True
+
+        time2 = time.perf_counter()
+        meta_data["load_data"] = f"{time2 - time1:0.3f}"
+        assert len(audio_sample_list) == 1, "batch_size must be set 1"
+
+        audio_sample = torch.cat((cache["prev_samples"], audio_sample_list[0]))
+
+        if len(audio_sample) < first_chunk_padding_samples:
+            print("key: {}, audio is too short for inference {}".format(key, len(audio_sample)))
+
+        audio_sample_pre = audio_sample[0 : first_chunk_padding_samples]
+        feat_pre, feat_pre_lengths = extract_fbank(
+            [audio_sample_pre],
+            data_type=kwargs.get("data_type", "sound"),
+            frontend=frontend,
+            cache=cache["frontend"],
+            is_final=False,
+        )
+
+        audio_sample = audio_sample[first_chunk_padding_samples:]
+        audio_chunks = int(len(audio_sample) // chunk_stride_samples)
+
+        for i in range(audio_chunks):
+            if i == 0:
+                cache["encoder"]["feats"][:, chunk_size[2] :, :] = feat_pre
+
+            kwargs["is_final"] = False
+            audio_sample_i = audio_sample[i * chunk_stride_samples : (i + 1) * chunk_stride_samples]
+
+            if kwargs["is_final"] and len(audio_sample_i) < 960:
+                cache["encoder"]["tail_chunk"] = True
+                speech = cache["encoder"]["feats"]
+                speech_lengths = torch.tensor([speech.shape[1]], dtype=torch.int64).to(
+                    speech.device
+                )
+            else:
+                # extract fbank feats
+                speech, speech_lengths = extract_fbank(
+                    [audio_sample_i],
+                    data_type=kwargs.get("data_type", "sound"),
+                    frontend=frontend,
+                    cache=cache["frontend"],
+                    is_final=kwargs["is_final"],
+                )
+            time3 = time.perf_counter()
+            meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
+            meta_data["batch_data_time"] = (
+                speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000
+            )
+
+            results_chunk_i = self.generate_chunk(
+                speech,
+                speech_lengths,
+                key=key,
+                tokenizer=tokenizer,
+                cache=cache,
+                frontend=frontend,
+                **kwargs,
+            )
+
+            # results_chunk_i must be None when is_final=False
+            assert results_chunk_i is None
+
+        # process tail samples
+        tail_audio_sample = audio_sample[ audio_chunks * chunk_stride_samples: ]
+        if len(tail_audio_sample) < 960:
+            kwargs["is_final"] = _is_final
+            cache["encoder"]["tail_chunk"] = True
+            speech = cache["encoder"]["feats"]
+            speech_lengths = torch.tensor([speech.shape[1]], dtype=torch.int64).to(
+                speech.device
+            )
+            results_chunk_tail = self.generate_chunk(
+                speech,
+                speech_lengths,
+                key=key,
+                tokenizer=tokenizer,
+                cache=cache,
+                frontend=frontend,
+                **kwargs,
+            )
+        elif len(tail_audio_sample) <= first_chunk_padding_samples:
+            kwargs["is_final"] = _is_final
+            # extract fbank feats
+            # cache["encoder"]["tail_chunk"] = True  # cannot be true
+            speech, speech_lengths = extract_fbank(
+                [ tail_audio_sample ],
+                data_type=kwargs.get("data_type", "sound"),
+                frontend=frontend,
+                cache=cache["frontend"],
+                is_final=kwargs["is_final"],
+            )
+            results_chunk_tail = self.generate_chunk(
+                speech,
+                speech_lengths,
+                key=key,
+                tokenizer=tokenizer,
+                cache=cache,
+                frontend=frontend,
+                **kwargs,
+            )
+        elif len(tail_audio_sample) > first_chunk_padding_samples and \
+             len(tail_audio_sample) < chunk_stride_samples:
+            kwargs["is_final"] = False
+            # extract fbank feats
+            speech, speech_lengths = extract_fbank(
+                [ tail_audio_sample ],
+                data_type=kwargs.get("data_type", "sound"),
+                frontend=frontend,
+                cache=cache["frontend"],
+                is_final=kwargs["is_final"],
+            )
+            results_chunk = self.generate_chunk(
+                speech,
+                speech_lengths,
+                key=key,
+                tokenizer=tokenizer,
+                cache=cache,
+                frontend=frontend,
+                **kwargs,
+            )
+            # results_chunk must be None when is_final=False
+            assert results_chunk is None
+
+            # push tail chunk
+            kwargs["is_final"] = _is_final
+            cache["encoder"]["tail_chunk"] = True
+            speech = cache["encoder"]["feats"]
+            speech_lengths = torch.tensor([speech.shape[1]], dtype=torch.int64).to(
+                speech.device
+            )
+            results_chunk_tail = self.generate_chunk(
+                speech,
+                speech_lengths,
+                key=key,
+                tokenizer=tokenizer,
+                cache=cache,
+                frontend=frontend,
+                **kwargs,
+            )
+
+        result = results_chunk_tail
+
+        if _is_final:
+            self.init_cache(cache, **kwargs)
+
+        if kwargs.get("output_dir"):
+            if not hasattr(self, "writer"):
+                self.writer = DatadirWriter(kwargs.get("output_dir"))
+
+        return result, meta_data
+
+    def export(self, **kwargs):
+        from .export_meta import export_rebuild_model
+
+        models = export_rebuild_model(model=self, **kwargs)
+        return models

funasr/models/scama/beam_search.py

@@ -0,0 +1,978 @@
+"""Beam search module."""
+
+from itertools import chain
+import logging
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import NamedTuple
+from typing import Tuple
+from typing import Union
+
+import torch
+
+from funasr.metrics.common import end_detect
+from funasr.models.transformer.scorers.scorer_interface import PartialScorerInterface
+from funasr.models.transformer.scorers.scorer_interface import ScorerInterface
+
+
+class Hypothesis(NamedTuple):
+    """Hypothesis data type."""
+
+    yseq: torch.Tensor
+    score: Union[float, torch.Tensor] = 0
+    scores: Dict[str, Union[float, torch.Tensor]] = dict()
+    states: Dict[str, Any] = dict()
+
+    def asdict(self) -> dict:
+        """Convert data to JSON-friendly dict."""
+        return self._replace(
+            yseq=self.yseq.tolist(),
+            score=float(self.score),
+            scores={k: float(v) for k, v in self.scores.items()},
+        )._asdict()
+
+
+class BeamSearchScama(torch.nn.Module):
+    """Beam search implementation."""
+
+    def __init__(
+        self,
+        scorers: Dict[str, ScorerInterface],
+        weights: Dict[str, float],
+        beam_size: int,
+        vocab_size: int,
+        sos: int,
+        eos: int,
+        token_list: List[str] = None,
+        pre_beam_ratio: float = 1.5,
+        pre_beam_score_key: str = None,
+    ):
+        """Initialize beam search.
+
+        Args:
+            scorers (dict[str, ScorerInterface]): Dict of decoder modules
+                e.g., Decoder, CTCPrefixScorer, LM
+                The scorer will be ignored if it is `None`
+            weights (dict[str, float]): Dict of weights for each scorers
+                The scorer will be ignored if its weight is 0
+            beam_size (int): The number of hypotheses kept during search
+            vocab_size (int): The number of vocabulary
+            sos (int): Start of sequence id
+            eos (int): End of sequence id
+            token_list (list[str]): List of tokens for debug log
+            pre_beam_score_key (str): key of scores to perform pre-beam search
+            pre_beam_ratio (float): beam size in the pre-beam search
+                will be `int(pre_beam_ratio * beam_size)`
+
+        """
+        super().__init__()
+        # set scorers
+        self.weights = weights
+        self.scorers = dict()
+        self.full_scorers = dict()
+        self.part_scorers = dict()
+        # this module dict is required for recursive cast
+        # `self.to(device, dtype)` in `recog.py`
+        self.nn_dict = torch.nn.ModuleDict()
+        for k, v in scorers.items():
+            w = weights.get(k, 0)
+            if w == 0 or v is None:
+                continue
+            assert isinstance(
+                v, ScorerInterface
+            ), f"{k} ({type(v)}) does not implement ScorerInterface"
+            self.scorers[k] = v
+            if isinstance(v, PartialScorerInterface):
+                self.part_scorers[k] = v
+            else:
+                self.full_scorers[k] = v
+            if isinstance(v, torch.nn.Module):
+                self.nn_dict[k] = v
+
+        # set configurations
+        self.sos = sos
+        self.eos = eos
+        self.token_list = token_list
+        self.pre_beam_size = int(pre_beam_ratio * beam_size)
+        self.beam_size = beam_size
+        self.n_vocab = vocab_size
+        if (
+            pre_beam_score_key is not None
+            and pre_beam_score_key != "full"
+            and pre_beam_score_key not in self.full_scorers
+        ):
+            raise KeyError(f"{pre_beam_score_key} is not found in {self.full_scorers}")
+        self.pre_beam_score_key = pre_beam_score_key
+        self.do_pre_beam = (
+            self.pre_beam_score_key is not None
+            and self.pre_beam_size < self.n_vocab
+            and len(self.part_scorers) > 0
+        )
+
+    def init_hyp(self, x: torch.Tensor) -> List[Hypothesis]:
+        """Get an initial hypothesis data.
+
+        Args:
+            x (torch.Tensor): The encoder output feature
+
+        Returns:
+            Hypothesis: The initial hypothesis.
+
+        """
+        init_states = dict()
+        init_scores = dict()
+        for k, d in self.scorers.items():
+            init_states[k] = d.init_state(x)
+            init_scores[k] = 0.0
+        return [
+            Hypothesis(
+                score=0.0,
+                scores=init_scores,
+                states=init_states,
+                yseq=torch.tensor([self.sos], device=x.device),
+            )
+        ]
+
+    @staticmethod
+    def append_token(xs: torch.Tensor, x: int) -> torch.Tensor:
+        """Append new token to prefix tokens.
+
+        Args:
+            xs (torch.Tensor): The prefix token
+            x (int): The new token to append
+
+        Returns:
+            torch.Tensor: New tensor contains: xs + [x] with xs.dtype and xs.device
+
+        """
+        x = torch.tensor([x], dtype=xs.dtype, device=xs.device)
+        return torch.cat((xs, x))
+
+    def score_full(
+        self,
+        hyp: Hypothesis,
+        x: torch.Tensor,
+        x_mask: torch.Tensor = None,
+        pre_acoustic_embeds: torch.Tensor = None,
+    ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]:
+        """Score new hypothesis by `self.full_scorers`.
+
+        Args:
+            hyp (Hypothesis): Hypothesis with prefix tokens to score
+            x (torch.Tensor): Corresponding input feature
+
+        Returns:
+            Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of
+                score dict of `hyp` that has string keys of `self.full_scorers`
+                and tensor score values of shape: `(self.n_vocab,)`,
+                and state dict that has string keys
+                and state values of `self.full_scorers`
+
+        """
+        scores = dict()
+        states = dict()
+        for k, d in self.full_scorers.items():
+            scores[k], states[k] = d.score(
+                hyp.yseq, hyp.states[k], x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds
+            )
+        return scores, states
+
+    def score_partial(
+        self, hyp: Hypothesis, ids: torch.Tensor, x: torch.Tensor
+    ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]:
+        """Score new hypothesis by `self.part_scorers`.
+
+        Args:
+            hyp (Hypothesis): Hypothesis with prefix tokens to score
+            ids (torch.Tensor): 1D tensor of new partial tokens to score
+            x (torch.Tensor): Corresponding input feature
+
+        Returns:
+            Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of
+                score dict of `hyp` that has string keys of `self.part_scorers`
+                and tensor score values of shape: `(len(ids),)`,
+                and state dict that has string keys
+                and state values of `self.part_scorers`
+
+        """
+        scores = dict()
+        states = dict()
+        for k, d in self.part_scorers.items():
+            scores[k], states[k] = d.score_partial(hyp.yseq, ids, hyp.states[k], x)
+        return scores, states
+
+    def beam(
+        self, weighted_scores: torch.Tensor, ids: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Compute topk full token ids and partial token ids.
+
+        Args:
+            weighted_scores (torch.Tensor): The weighted sum scores for each tokens.
+            Its shape is `(self.n_vocab,)`.
+            ids (torch.Tensor): The partial token ids to compute topk
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]:
+                The topk full token ids and partial token ids.
+                Their shapes are `(self.beam_size,)`
+
+        """
+        # no pre beam performed
+        if weighted_scores.size(0) == ids.size(0):
+            top_ids = weighted_scores.topk(self.beam_size)[1]
+            return top_ids, top_ids
+
+        # mask pruned in pre-beam not to select in topk
+        tmp = weighted_scores[ids]
+        weighted_scores[:] = -float("inf")
+        weighted_scores[ids] = tmp
+        top_ids = weighted_scores.topk(self.beam_size)[1]
+        local_ids = weighted_scores[ids].topk(self.beam_size)[1]
+        return top_ids, local_ids
+
+    @staticmethod
+    def merge_scores(
+        prev_scores: Dict[str, float],
+        next_full_scores: Dict[str, torch.Tensor],
+        full_idx: int,
+        next_part_scores: Dict[str, torch.Tensor],
+        part_idx: int,
+    ) -> Dict[str, torch.Tensor]:
+        """Merge scores for new hypothesis.
+
+        Args:
+            prev_scores (Dict[str, float]):
+                The previous hypothesis scores by `self.scorers`
+            next_full_scores (Dict[str, torch.Tensor]): scores by `self.full_scorers`
+            full_idx (int): The next token id for `next_full_scores`
+            next_part_scores (Dict[str, torch.Tensor]):
+                scores of partial tokens by `self.part_scorers`
+            part_idx (int): The new token id for `next_part_scores`
+
+        Returns:
+            Dict[str, torch.Tensor]: The new score dict.
+                Its keys are names of `self.full_scorers` and `self.part_scorers`.
+                Its values are scalar tensors by the scorers.
+
+        """
+        new_scores = dict()
+        for k, v in next_full_scores.items():
+            new_scores[k] = prev_scores[k] + v[full_idx]
+        for k, v in next_part_scores.items():
+            new_scores[k] = prev_scores[k] + v[part_idx]
+        return new_scores
+
+    def merge_states(self, states: Any, part_states: Any, part_idx: int) -> Any:
+        """Merge states for new hypothesis.
+
+        Args:
+            states: states of `self.full_scorers`
+            part_states: states of `self.part_scorers`
+            part_idx (int): The new token id for `part_scores`
+
+        Returns:
+            Dict[str, torch.Tensor]: The new score dict.
+                Its keys are names of `self.full_scorers` and `self.part_scorers`.
+                Its values are states of the scorers.
+
+        """
+        new_states = dict()
+        for k, v in states.items():
+            new_states[k] = v
+        for k, d in self.part_scorers.items():
+            new_states[k] = d.select_state(part_states[k], part_idx)
+        return new_states
+
+    def search(
+        self,
+        running_hyps: List[Hypothesis],
+        x: torch.Tensor,
+        x_mask: torch.Tensor = None,
+        pre_acoustic_embeds: torch.Tensor = None,
+    ) -> List[Hypothesis]:
+        """Search new tokens for running hypotheses and encoded speech x.
+
+        Args:
+            running_hyps (List[Hypothesis]): Running hypotheses on beam
+            x (torch.Tensor): Encoded speech feature (T, D)
+
+        Returns:
+            List[Hypotheses]: Best sorted hypotheses
+
+        """
+        best_hyps = []
+        part_ids = torch.arange(self.n_vocab, device=x.device)  # no pre-beam
+        for hyp in running_hyps:
+            # scoring
+            weighted_scores = torch.zeros(self.n_vocab, dtype=x.dtype, device=x.device)
+            scores, states = self.score_full(
+                hyp, x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds
+            )
+            for k in self.full_scorers:
+                weighted_scores += self.weights[k] * scores[k]
+            # partial scoring
+            if self.do_pre_beam:
+                pre_beam_scores = (
+                    weighted_scores
+                    if self.pre_beam_score_key == "full"
+                    else scores[self.pre_beam_score_key]
+                )
+                part_ids = torch.topk(pre_beam_scores, self.pre_beam_size)[1]
+            part_scores, part_states = self.score_partial(hyp, part_ids, x)
+            for k in self.part_scorers:
+                weighted_scores[part_ids] += self.weights[k] * part_scores[k]
+            # add previous hyp score
+            weighted_scores += hyp.score
+
+            # update hyps
+            for j, part_j in zip(*self.beam(weighted_scores, part_ids)):
+                # will be (2 x beam at most)
+                best_hyps.append(
+                    Hypothesis(
+                        score=weighted_scores[j],
+                        yseq=self.append_token(hyp.yseq, j),
+                        scores=self.merge_scores(hyp.scores, scores, j, part_scores, part_j),
+                        states=self.merge_states(states, part_states, part_j),
+                    )
+                )
+
+            # sort and prune 2 x beam -> beam
+            best_hyps = sorted(best_hyps, key=lambda x: x.score, reverse=True)[
+                : min(len(best_hyps), self.beam_size)
+            ]
+        return best_hyps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        scama_mask: torch.Tensor = None,
+        pre_acoustic_embeds: torch.Tensor = None,
+        maxlenratio: float = 0.0,
+        minlenratio: float = 0.0,
+        maxlen: int = None,
+        minlen: int = 0,
+    ) -> List[Hypothesis]:
+        """Perform beam search.
+
+        Args:
+            x (torch.Tensor): Encoded speech feature (T, D)
+            maxlenratio (float): Input length ratio to obtain max output length.
+                If maxlenratio=0.0 (default), it uses a end-detect function
+                to automatically find maximum hypothesis lengths
+                If maxlenratio<0.0, its absolute value is interpreted
+                as a constant max output length.
+            minlenratio (float): Input length ratio to obtain min output length.
+
+        Returns:
+            list[Hypothesis]: N-best decoding results
+
+        """
+        if maxlen is None:
+            # set length bounds
+            if maxlenratio == 0:
+                maxlen = x.shape[0]
+            elif maxlenratio < 0:
+                maxlen = -1 * int(maxlenratio)
+            else:
+                maxlen = max(1, int(maxlenratio * x.size(0)))
+            minlen = int(minlenratio * x.size(0))
+
+        logging.info("decoder input length: " + str(x.shape[0]))
+        logging.info("max output length: " + str(maxlen))
+        logging.info("min output length: " + str(minlen))
+
+        # main loop of prefix search
+        running_hyps = self.init_hyp(x)
+        ended_hyps = []
+        for i in range(maxlen):
+            logging.debug("position " + str(i))
+            mask_enc = None
+            if scama_mask is not None:
+                token_num_predictor = scama_mask.size(1)
+                token_id_slice = min(i, token_num_predictor - 1)
+                mask_enc = scama_mask[:, token_id_slice : token_id_slice + 1, :]
+                # if mask_enc.size(1) == 0:
+                #     mask_enc = scama_mask[:, -2:-1, :]
+                #     # mask_enc = torch.zeros_like(mask_enc)
+            pre_acoustic_embeds_cur = None
+            if pre_acoustic_embeds is not None:
+                b, t, d = pre_acoustic_embeds.size()
+                pad = torch.zeros((b, 1, d), dtype=pre_acoustic_embeds.dtype).to(
+                    device=pre_acoustic_embeds.device
+                )
+                pre_acoustic_embeds = torch.cat((pre_acoustic_embeds, pad), dim=1)
+                token_id_slice = min(i, t)
+                pre_acoustic_embeds_cur = pre_acoustic_embeds[
+                    :, token_id_slice : token_id_slice + 1, :
+                ]
+
+            best = self.search(
+                running_hyps, x, x_mask=mask_enc, pre_acoustic_embeds=pre_acoustic_embeds_cur
+            )
+            # post process of one iteration
+            running_hyps = self.post_process(i, maxlen, maxlenratio, best, ended_hyps)
+            # end detection
+            if maxlenratio == 0.0 and end_detect([h.asdict() for h in ended_hyps], i):
+                logging.info(f"end detected at {i}")
+                break
+            if len(running_hyps) == 0:
+                logging.info("no hypothesis. Finish decoding.")
+                break
+            else:
+                logging.debug(f"remained hypotheses: {len(running_hyps)}")
+
+        nbest_hyps = sorted(ended_hyps, key=lambda x: x.score, reverse=True)
+        # check the number of hypotheses reaching to eos
+        if len(nbest_hyps) == 0:
+            logging.warning(
+                "there is no N-best results, perform recognition " "again with smaller minlenratio."
+            )
+            return (
+                []
+                if minlenratio < 0.1
+                else self.forward(x, maxlenratio, max(0.0, minlenratio - 0.1))
+            )
+
+        # report the best result
+        for x in nbest_hyps:
+            yseq = "".join([self.token_list[x] for x in x.yseq])
+            logging.debug("nbest: y: {}, yseq: {}, score: {}".format(x.yseq, yseq, x.score))
+        best = nbest_hyps[0]
+        for k, v in best.scores.items():
+            logging.info(f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}")
+        logging.info(f"total log probability: {best.score:.2f}")
+        logging.info(f"normalized log probability: {best.score / len(best.yseq):.2f}")
+        logging.info(f"total number of ended hypotheses: {len(nbest_hyps)}")
+        if self.token_list is not None:
+            logging.info(
+                "best hypo: " + "".join([self.token_list[x] for x in best.yseq[1:-1]]) + "\n"
+            )
+        return nbest_hyps
+
+    def post_process(
+        self,
+        i: int,
+        maxlen: int,
+        maxlenratio: float,
+        running_hyps: List[Hypothesis],
+        ended_hyps: List[Hypothesis],
+    ) -> List[Hypothesis]:
+        """Perform post-processing of beam search iterations.
+
+        Args:
+            i (int): The length of hypothesis tokens.
+            maxlen (int): The maximum length of tokens in beam search.
+            maxlenratio (int): The maximum length ratio in beam search.
+            running_hyps (List[Hypothesis]): The running hypotheses in beam search.
+            ended_hyps (List[Hypothesis]): The ended hypotheses in beam search.
+
+        Returns:
+            List[Hypothesis]: The new running hypotheses.
+
+        """
+        logging.debug(f"the number of running hypotheses: {len(running_hyps)}")
+        if self.token_list is not None:
+            logging.debug(
+                "best hypo: " + "".join([self.token_list[x] for x in running_hyps[0].yseq[1:]])
+            )
+        # add eos in the final loop to avoid that there are no ended hyps
+        if i == maxlen - 1:
+            logging.info("adding <eos> in the last position in the loop")
+            running_hyps = [
+                h._replace(yseq=self.append_token(h.yseq, self.eos)) for h in running_hyps
+            ]
+
+        # add ended hypotheses to a final list, and removed them from current hypotheses
+        # (this will be a problem, number of hyps < beam)
+        remained_hyps = []
+        for hyp in running_hyps:
+            if hyp.yseq[-1] == self.eos:
+                # e.g., Word LM needs to add final <eos> score
+                for k, d in chain(self.full_scorers.items(), self.part_scorers.items()):
+                    s = d.final_score(hyp.states[k])
+                    hyp.scores[k] += s
+                    hyp = hyp._replace(score=hyp.score + self.weights[k] * s)
+                ended_hyps.append(hyp)
+            else:
+                remained_hyps.append(hyp)
+        return remained_hyps
+
+
+class BeamSearchScamaStreaming(torch.nn.Module):
+    """Beam search implementation."""
+
+    def __init__(
+        self,
+        scorers: Dict[str, ScorerInterface],
+        weights: Dict[str, float],
+        beam_size: int,
+        vocab_size: int,
+        sos: int,
+        eos: int,
+        token_list: List[str] = None,
+        pre_beam_ratio: float = 1.5,
+        pre_beam_score_key: str = None,
+    ):
+        """Initialize beam search.
+
+        Args:
+            scorers (dict[str, ScorerInterface]): Dict of decoder modules
+                e.g., Decoder, CTCPrefixScorer, LM
+                The scorer will be ignored if it is `None`
+            weights (dict[str, float]): Dict of weights for each scorers
+                The scorer will be ignored if its weight is 0
+            beam_size (int): The number of hypotheses kept during search
+            vocab_size (int): The number of vocabulary
+            sos (int): Start of sequence id
+            eos (int): End of sequence id
+            token_list (list[str]): List of tokens for debug log
+            pre_beam_score_key (str): key of scores to perform pre-beam search
+            pre_beam_ratio (float): beam size in the pre-beam search
+                will be `int(pre_beam_ratio * beam_size)`
+
+        """
+        super().__init__()
+        # set scorers
+        self.weights = weights
+        self.scorers = dict()
+        self.full_scorers = dict()
+        self.part_scorers = dict()
+        # this module dict is required for recursive cast
+        # `self.to(device, dtype)` in `recog.py`
+        self.nn_dict = torch.nn.ModuleDict()
+        for k, v in scorers.items():
+            w = weights.get(k, 0)
+            if w == 0 or v is None:
+                continue
+            assert isinstance(
+                v, ScorerInterface
+            ), f"{k} ({type(v)}) does not implement ScorerInterface"
+            self.scorers[k] = v
+            if isinstance(v, PartialScorerInterface):
+                self.part_scorers[k] = v
+            else:
+                self.full_scorers[k] = v
+            if isinstance(v, torch.nn.Module):
+                self.nn_dict[k] = v
+
+        # set configurations
+        self.sos = sos
+        self.eos = eos
+        self.token_list = token_list
+        self.pre_beam_size = int(pre_beam_ratio * beam_size)
+        self.beam_size = beam_size
+        self.n_vocab = vocab_size
+        if (
+            pre_beam_score_key is not None
+            and pre_beam_score_key != "full"
+            and pre_beam_score_key not in self.full_scorers
+        ):
+            raise KeyError(f"{pre_beam_score_key} is not found in {self.full_scorers}")
+        self.pre_beam_score_key = pre_beam_score_key
+        self.do_pre_beam = (
+            self.pre_beam_score_key is not None
+            and self.pre_beam_size < self.n_vocab
+            and len(self.part_scorers) > 0
+        )
+
+    def init_hyp(self, x) -> List[Hypothesis]:
+        """Get an initial hypothesis data.
+
+        Args:
+            x (torch.Tensor): The encoder output feature
+
+        Returns:
+            Hypothesis: The initial hypothesis.
+
+        """
+        init_states = dict()
+        init_scores = dict()
+        for k, d in self.scorers.items():
+            init_states[k] = d.init_state(x)
+            init_scores[k] = 0.0
+        return [
+            Hypothesis(
+                score=0.0,
+                scores=init_scores,
+                states=init_states,
+                yseq=torch.tensor([self.sos], device=x.device),
+            )
+        ]
+
+    @staticmethod
+    def append_token(xs: torch.Tensor, x: int) -> torch.Tensor:
+        """Append new token to prefix tokens.
+
+        Args:
+            xs (torch.Tensor): The prefix token
+            x (int): The new token to append
+
+        Returns:
+            torch.Tensor: New tensor contains: xs + [x] with xs.dtype and xs.device
+
+        """
+        x = torch.tensor([x], dtype=xs.dtype, device=xs.device)
+        return torch.cat((xs, x))
+
+    def score_full(
+        self,
+        hyp: Hypothesis,
+        x: torch.Tensor,
+        x_mask: torch.Tensor = None,
+        pre_acoustic_embeds: torch.Tensor = None,
+        cache: dict = {},
+    ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]:
+        """Score new hypothesis by `self.full_scorers`.
+
+        Args:
+            hyp (Hypothesis): Hypothesis with prefix tokens to score
+            x (torch.Tensor): Corresponding input feature
+
+        Returns:
+            Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of
+                score dict of `hyp` that has string keys of `self.full_scorers`
+                and tensor score values of shape: `(self.n_vocab,)`,
+                and state dict that has string keys
+                and state values of `self.full_scorers`
+
+        """
+        scores = dict()
+        states = dict()
+        for k, d in self.full_scorers.items():
+            scores[k], states[k] = d.score(
+                hyp.yseq,
+                hyp.states[k],
+                x,
+                x_mask=x_mask,
+                pre_acoustic_embeds=pre_acoustic_embeds,
+                cache=cache,
+            )
+        return scores, states
+
+    def score_partial(
+        self, hyp: Hypothesis, ids: torch.Tensor, x: torch.Tensor
+    ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]:
+        """Score new hypothesis by `self.part_scorers`.
+
+        Args:
+            hyp (Hypothesis): Hypothesis with prefix tokens to score
+            ids (torch.Tensor): 1D tensor of new partial tokens to score
+            x (torch.Tensor): Corresponding input feature
+
+        Returns:
+            Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of
+                score dict of `hyp` that has string keys of `self.part_scorers`
+                and tensor score values of shape: `(len(ids),)`,
+                and state dict that has string keys
+                and state values of `self.part_scorers`
+
+        """
+        scores = dict()
+        states = dict()
+        for k, d in self.part_scorers.items():
+            scores[k], states[k] = d.score_partial(hyp.yseq, ids, hyp.states[k], x)
+        return scores, states
+
+    def beam(
+        self, weighted_scores: torch.Tensor, ids: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Compute topk full token ids and partial token ids.
+
+        Args:
+            weighted_scores (torch.Tensor): The weighted sum scores for each tokens.
+            Its shape is `(self.n_vocab,)`.
+            ids (torch.Tensor): The partial token ids to compute topk
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]:
+                The topk full token ids and partial token ids.
+                Their shapes are `(self.beam_size,)`
+
+        """
+        # no pre beam performed
+        if weighted_scores.size(0) == ids.size(0):
+            top_ids = weighted_scores.topk(self.beam_size)[1]
+            return top_ids, top_ids
+
+        # mask pruned in pre-beam not to select in topk
+        tmp = weighted_scores[ids]
+        weighted_scores[:] = -float("inf")
+        weighted_scores[ids] = tmp
+        top_ids = weighted_scores.topk(self.beam_size)[1]
+        local_ids = weighted_scores[ids].topk(self.beam_size)[1]
+        return top_ids, local_ids
+
+    @staticmethod
+    def merge_scores(
+        prev_scores: Dict[str, float],
+        next_full_scores: Dict[str, torch.Tensor],
+        full_idx: int,
+        next_part_scores: Dict[str, torch.Tensor],
+        part_idx: int,
+    ) -> Dict[str, torch.Tensor]:
+        """Merge scores for new hypothesis.
+
+        Args:
+            prev_scores (Dict[str, float]):
+                The previous hypothesis scores by `self.scorers`
+            next_full_scores (Dict[str, torch.Tensor]): scores by `self.full_scorers`
+            full_idx (int): The next token id for `next_full_scores`
+            next_part_scores (Dict[str, torch.Tensor]):
+                scores of partial tokens by `self.part_scorers`
+            part_idx (int): The new token id for `next_part_scores`
+
+        Returns:
+            Dict[str, torch.Tensor]: The new score dict.
+                Its keys are names of `self.full_scorers` and `self.part_scorers`.
+                Its values are scalar tensors by the scorers.
+
+        """
+        new_scores = dict()
+        for k, v in next_full_scores.items():
+            new_scores[k] = prev_scores[k] + v[full_idx]
+        for k, v in next_part_scores.items():
+            new_scores[k] = prev_scores[k] + v[part_idx]
+        return new_scores
+
+    def merge_states(self, states: Any, part_states: Any, part_idx: int) -> Any:
+        """Merge states for new hypothesis.
+
+        Args:
+            states: states of `self.full_scorers`
+            part_states: states of `self.part_scorers`
+            part_idx (int): The new token id for `part_scores`
+
+        Returns:
+            Dict[str, torch.Tensor]: The new score dict.
+                Its keys are names of `self.full_scorers` and `self.part_scorers`.
+                Its values are states of the scorers.
+
+        """
+        new_states = dict()
+        for k, v in states.items():
+            new_states[k] = v
+        for k, d in self.part_scorers.items():
+            new_states[k] = d.select_state(part_states[k], part_idx)
+        return new_states
+
+    def search(
+        self,
+        running_hyps: List[Hypothesis],
+        x: torch.Tensor,
+        x_mask: torch.Tensor = None,
+        pre_acoustic_embeds: torch.Tensor = None,
+        cache: dict = {},
+    ) -> List[Hypothesis]:
+        """Search new tokens for running hypotheses and encoded speech x.
+
+        Args:
+            running_hyps (List[Hypothesis]): Running hypotheses on beam
+            x (torch.Tensor): Encoded speech feature (T, D)
+
+        Returns:
+            List[Hypotheses]: Best sorted hypotheses
+
+        """
+        best_hyps = []
+        part_ids = torch.arange(self.n_vocab, device=x.device)  # no pre-beam
+        for hyp in running_hyps:
+            # scoring
+            weighted_scores = torch.zeros(self.n_vocab, dtype=x.dtype, device=x.device)
+            scores, states = self.score_full(
+                hyp, x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds, cache=cache
+            )
+            for k in self.full_scorers:
+                weighted_scores += self.weights[k] * scores[k]
+            # partial scoring
+            if self.do_pre_beam:
+                pre_beam_scores = (
+                    weighted_scores
+                    if self.pre_beam_score_key == "full"
+                    else scores[self.pre_beam_score_key]
+                )
+                part_ids = torch.topk(pre_beam_scores, self.pre_beam_size)[1]
+            part_scores, part_states = self.score_partial(hyp, part_ids, x)
+            for k in self.part_scorers:
+                weighted_scores[part_ids] += self.weights[k] * part_scores[k]
+            # add previous hyp score
+            weighted_scores += hyp.score
+
+            # update hyps
+            for j, part_j in zip(*self.beam(weighted_scores, part_ids)):
+                # will be (2 x beam at most)
+                best_hyps.append(
+                    Hypothesis(
+                        score=weighted_scores[j],
+                        yseq=self.append_token(hyp.yseq, j),
+                        scores=self.merge_scores(hyp.scores, scores, j, part_scores, part_j),
+                        states=self.merge_states(states, part_states, part_j),
+                    )
+                )
+
+            # sort and prune 2 x beam -> beam
+            best_hyps = sorted(best_hyps, key=lambda x: x.score, reverse=True)[
+                : min(len(best_hyps), self.beam_size)
+            ]
+        return best_hyps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        scama_mask: torch.Tensor = None,
+        pre_acoustic_embeds: torch.Tensor = None,
+        maxlenratio: float = 0.0,
+        minlenratio: float = 0.0,
+        maxlen: int = None,
+        minlen: int = 0,
+        cache: dict = {},
+    ) -> List[Hypothesis]:
+        """Perform beam search.
+
+        Args:
+            x (torch.Tensor): Encoded speech feature (T, D)
+            maxlenratio (float): Input length ratio to obtain max output length.
+                If maxlenratio=0.0 (default), it uses a end-detect function
+                to automatically find maximum hypothesis lengths
+                If maxlenratio<0.0, its absolute value is interpreted
+                as a constant max output length.
+            minlenratio (float): Input length ratio to obtain min output length.
+
+        Returns:
+            list[Hypothesis]: N-best decoding results
+
+        """
+        if maxlen is None:
+            # set length bounds
+            if maxlenratio == 0:
+                maxlen = x.shape[0]
+            elif maxlenratio < 0:
+                maxlen = -1 * int(maxlenratio)
+            else:
+                maxlen = max(1, int(maxlenratio * x.size(0)))
+            minlen = int(minlenratio * x.size(0))
+
+        logging.info("decoder input length: " + str(x.shape[0]))
+        logging.info("max output length: " + str(maxlen))
+        logging.info("min output length: " + str(minlen))
+
+        # main loop of prefix search
+        # running_hyps = self.init_hyp(x)
+        running_hyps = cache["running_hyps"]
+        ended_hyps = []
+        for i in range(maxlen):
+            logging.debug("position " + str(i))
+            mask_enc = None
+            # if scama_mask is not None:
+            #     token_num_predictor = scama_mask.size(1)
+            #     token_id_slice = min(i, token_num_predictor-1)
+            #     mask_enc = scama_mask[:, token_id_slice:token_id_slice+1, :]
+            #     # if mask_enc.size(1) == 0:
+            #     #     mask_enc = scama_mask[:, -2:-1, :]
+            #     #     # mask_enc = torch.zeros_like(mask_enc)
+            pre_acoustic_embeds_cur = None
+            if pre_acoustic_embeds is not None:
+                b, t, d = pre_acoustic_embeds.size()
+                pad = torch.zeros((b, 1, d), dtype=pre_acoustic_embeds.dtype).to(
+                    device=pre_acoustic_embeds.device
+                )
+                pre_acoustic_embeds = torch.cat((pre_acoustic_embeds, pad), dim=1)
+                token_id_slice = min(i, t)
+                pre_acoustic_embeds_cur = pre_acoustic_embeds[
+                    :, token_id_slice : token_id_slice + 1, :
+                ]
+
+            best = self.search(
+                running_hyps,
+                x,
+                x_mask=mask_enc,
+                pre_acoustic_embeds=pre_acoustic_embeds_cur,
+                cache=cache["decoder"],
+            )
+            # post process of one iteration
+            running_hyps = self.post_process(i, maxlen, maxlenratio, best, ended_hyps)
+            # end detection
+            if maxlenratio == 0.0 and end_detect([h.asdict() for h in ended_hyps], i):
+                logging.info(f"end detected at {i}")
+                break
+            if len(running_hyps) == 0:
+                logging.info("no hypothesis. Finish decoding.")
+                break
+            else:
+                logging.debug(f"remained hypotheses: {len(running_hyps)}")
+
+        nbest_hyps = sorted(ended_hyps, key=lambda x: x.score, reverse=True)
+        # check the number of hypotheses reaching to eos
+        if len(nbest_hyps) == 0:
+            logging.warning(
+                "there is no N-best results, perform recognition " "again with smaller minlenratio."
+            )
+            return (
+                []
+                if minlenratio < 0.1
+                else self.forward(x, maxlenratio, max(0.0, minlenratio - 0.1))
+            )
+
+        # report the best result
+        for x in nbest_hyps:
+            yseq = "".join([self.token_list[x] for x in x.yseq])
+            logging.debug("nbest: y: {}, yseq: {}, score: {}".format(x.yseq, yseq, x.score))
+        best = nbest_hyps[0]
+        for k, v in best.scores.items():
+            logging.info(f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}")
+        logging.info(f"total log probability: {best.score:.2f}")
+        logging.info(f"normalized log probability: {best.score / len(best.yseq):.2f}")
+        logging.info(f"total number of ended hypotheses: {len(nbest_hyps)}")
+        if self.token_list is not None:
+            logging.info(
+                "best hypo: " + "".join([self.token_list[x] for x in best.yseq[1:-1]]) + "\n"
+            )
+        return nbest_hyps
+
+    def post_process(
+        self,
+        i: int,
+        maxlen: int,
+        maxlenratio: float,
+        running_hyps: List[Hypothesis],
+        ended_hyps: List[Hypothesis],
+    ) -> List[Hypothesis]:
+        """Perform post-processing of beam search iterations.
+
+        Args:
+            i (int): The length of hypothesis tokens.
+            maxlen (int): The maximum length of tokens in beam search.
+            maxlenratio (int): The maximum length ratio in beam search.
+            running_hyps (List[Hypothesis]): The running hypotheses in beam search.
+            ended_hyps (List[Hypothesis]): The ended hypotheses in beam search.
+
+        Returns:
+            List[Hypothesis]: The new running hypotheses.
+
+        """
+        logging.debug(f"the number of running hypotheses: {len(running_hyps)}")
+        if self.token_list is not None:
+            logging.debug(
+                "best hypo: " + "".join([self.token_list[x] for x in running_hyps[0].yseq[1:]])
+            )
+        # add eos in the final loop to avoid that there are no ended hyps
+        if i == maxlen - 1:
+            logging.info("adding <eos> in the last position in the loop")
+            running_hyps = [
+                h._replace(yseq=self.append_token(h.yseq, self.eos)) for h in running_hyps
+            ]
+
+        # add ended hypotheses to a final list, and removed them from current hypotheses
+        # (this will be a problem, number of hyps < beam)
+        remained_hyps = []
+        for hyp in running_hyps:
+            if hyp.yseq[-1] == self.eos:
+                # e.g., Word LM needs to add final <eos> score
+                for k, d in chain(self.full_scorers.items(), self.part_scorers.items()):
+                    s = d.final_score(hyp.states[k])
+                    hyp.scores[k] += s
+                    hyp = hyp._replace(score=hyp.score + self.weights[k] * s)
+                ended_hyps.append(hyp)
+            else:
+                remained_hyps.append(hyp)
+        return remained_hyps

funasr/models/scama/chunk_utilis.py

@@ -0,0 +1,493 @@
+import math
+import torch
+import numpy as np
+import torch.nn.functional as F
+
+from funasr.models.scama.utils import sequence_mask
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+
+
+class overlap_chunk:
+    """
+    Author: Speech Lab of DAMO Academy, Alibaba Group
+    San-m: Memory equipped self-attention for end-to-end speech recognition
+    https://arxiv.org/abs/2006.01713
+
+    """
+
+    def __init__(
+        self,
+        chunk_size: tuple = (16,),
+        stride: tuple = (10,),
+        pad_left: tuple = (0,),
+        encoder_att_look_back_factor: tuple = (1,),
+        shfit_fsmn: int = 0,
+        decoder_att_look_back_factor: tuple = (1,),
+    ):
+
+        pad_left = self.check_chunk_size_args(chunk_size, pad_left)
+        encoder_att_look_back_factor = self.check_chunk_size_args(
+            chunk_size, encoder_att_look_back_factor
+        )
+        decoder_att_look_back_factor = self.check_chunk_size_args(
+            chunk_size, decoder_att_look_back_factor
+        )
+        (
+            self.chunk_size,
+            self.stride,
+            self.pad_left,
+            self.encoder_att_look_back_factor,
+            self.decoder_att_look_back_factor,
+        ) = (
+            chunk_size,
+            stride,
+            pad_left,
+            encoder_att_look_back_factor,
+            decoder_att_look_back_factor,
+        )
+        self.shfit_fsmn = shfit_fsmn
+        self.x_add_mask = None
+        self.x_rm_mask = None
+        self.x_len = None
+        self.mask_shfit_chunk = None
+        self.mask_chunk_predictor = None
+        self.mask_att_chunk_encoder = None
+        self.mask_shift_att_chunk_decoder = None
+        self.chunk_outs = None
+        (
+            self.chunk_size_cur,
+            self.stride_cur,
+            self.pad_left_cur,
+            self.encoder_att_look_back_factor_cur,
+            self.chunk_size_pad_shift_cur,
+        ) = (None, None, None, None, None)
+
+    def check_chunk_size_args(self, chunk_size, x):
+        if len(x) < len(chunk_size):
+            x = [x[0] for i in chunk_size]
+        return x
+
+    def get_chunk_size(self, ind: int = 0):
+        # with torch.no_grad:
+        chunk_size, stride, pad_left, encoder_att_look_back_factor, decoder_att_look_back_factor = (
+            self.chunk_size[ind],
+            self.stride[ind],
+            self.pad_left[ind],
+            self.encoder_att_look_back_factor[ind],
+            self.decoder_att_look_back_factor[ind],
+        )
+        (
+            self.chunk_size_cur,
+            self.stride_cur,
+            self.pad_left_cur,
+            self.encoder_att_look_back_factor_cur,
+            self.chunk_size_pad_shift_cur,
+            self.decoder_att_look_back_factor_cur,
+        ) = (
+            chunk_size,
+            stride,
+            pad_left,
+            encoder_att_look_back_factor,
+            chunk_size + self.shfit_fsmn,
+            decoder_att_look_back_factor,
+        )
+        return (
+            self.chunk_size_cur,
+            self.stride_cur,
+            self.pad_left_cur,
+            self.encoder_att_look_back_factor_cur,
+            self.chunk_size_pad_shift_cur,
+        )
+
+    def random_choice(self, training=True, decoding_ind=None):
+        chunk_num = len(self.chunk_size)
+        ind = 0
+        if training and chunk_num > 1:
+            ind = torch.randint(0, chunk_num, ()).cpu().item()
+        if not training and decoding_ind is not None:
+            ind = int(decoding_ind)
+
+        return ind
+
+    def gen_chunk_mask(self, x_len, ind=0, num_units=1, num_units_predictor=1):
+
+        with torch.no_grad():
+            x_len = x_len.cpu().numpy()
+            x_len_max = x_len.max()
+
+            chunk_size, stride, pad_left, encoder_att_look_back_factor, chunk_size_pad_shift = (
+                self.get_chunk_size(ind)
+            )
+            shfit_fsmn = self.shfit_fsmn
+            pad_right = chunk_size - stride - pad_left
+
+            chunk_num_batch = np.ceil(x_len / stride).astype(np.int32)
+            x_len_chunk = (
+                (chunk_num_batch - 1) * chunk_size_pad_shift
+                + shfit_fsmn
+                + pad_left
+                + 0
+                + x_len
+                - (chunk_num_batch - 1) * stride
+            )
+            x_len_chunk = x_len_chunk.astype(x_len.dtype)
+            x_len_chunk_max = x_len_chunk.max()
+
+            chunk_num = int(math.ceil(x_len_max / stride))
+            dtype = np.int32
+            max_len_for_x_mask_tmp = max(chunk_size, x_len_max + pad_left)
+            x_add_mask = np.zeros([0, max_len_for_x_mask_tmp], dtype=dtype)
+            x_rm_mask = np.zeros([max_len_for_x_mask_tmp, 0], dtype=dtype)
+            mask_shfit_chunk = np.zeros([0, num_units], dtype=dtype)
+            mask_chunk_predictor = np.zeros([0, num_units_predictor], dtype=dtype)
+            mask_shift_att_chunk_decoder = np.zeros([0, 1], dtype=dtype)
+            mask_att_chunk_encoder = np.zeros([0, chunk_num * chunk_size_pad_shift], dtype=dtype)
+            for chunk_ids in range(chunk_num):
+                # x_mask add
+                fsmn_padding = np.zeros((shfit_fsmn, max_len_for_x_mask_tmp), dtype=dtype)
+                x_mask_cur = np.diag(np.ones(chunk_size, dtype=np.float32))
+                x_mask_pad_left = np.zeros((chunk_size, chunk_ids * stride), dtype=dtype)
+                x_mask_pad_right = np.zeros((chunk_size, max_len_for_x_mask_tmp), dtype=dtype)
+                x_cur_pad = np.concatenate([x_mask_pad_left, x_mask_cur, x_mask_pad_right], axis=1)
+                x_cur_pad = x_cur_pad[:chunk_size, :max_len_for_x_mask_tmp]
+                x_add_mask_fsmn = np.concatenate([fsmn_padding, x_cur_pad], axis=0)
+                x_add_mask = np.concatenate([x_add_mask, x_add_mask_fsmn], axis=0)
+
+                # x_mask rm
+                fsmn_padding = np.zeros((max_len_for_x_mask_tmp, shfit_fsmn), dtype=dtype)
+                padding_mask_left = np.zeros((max_len_for_x_mask_tmp, pad_left), dtype=dtype)
+                padding_mask_right = np.zeros((max_len_for_x_mask_tmp, pad_right), dtype=dtype)
+                x_mask_cur = np.diag(np.ones(stride, dtype=dtype))
+                x_mask_cur_pad_top = np.zeros((chunk_ids * stride, stride), dtype=dtype)
+                x_mask_cur_pad_bottom = np.zeros((max_len_for_x_mask_tmp, stride), dtype=dtype)
+                x_rm_mask_cur = np.concatenate(
+                    [x_mask_cur_pad_top, x_mask_cur, x_mask_cur_pad_bottom], axis=0
+                )
+                x_rm_mask_cur = x_rm_mask_cur[:max_len_for_x_mask_tmp, :stride]
+                x_rm_mask_cur_fsmn = np.concatenate(
+                    [fsmn_padding, padding_mask_left, x_rm_mask_cur, padding_mask_right], axis=1
+                )
+                x_rm_mask = np.concatenate([x_rm_mask, x_rm_mask_cur_fsmn], axis=1)
+
+                # fsmn_padding_mask
+                pad_shfit_mask = np.zeros([shfit_fsmn, num_units], dtype=dtype)
+                ones_1 = np.ones([chunk_size, num_units], dtype=dtype)
+                mask_shfit_chunk_cur = np.concatenate([pad_shfit_mask, ones_1], axis=0)
+                mask_shfit_chunk = np.concatenate([mask_shfit_chunk, mask_shfit_chunk_cur], axis=0)
+
+                # predictor mask
+                zeros_1 = np.zeros([shfit_fsmn + pad_left, num_units_predictor], dtype=dtype)
+                ones_2 = np.ones([stride, num_units_predictor], dtype=dtype)
+                zeros_3 = np.zeros(
+                    [chunk_size - stride - pad_left, num_units_predictor], dtype=dtype
+                )
+                ones_zeros = np.concatenate([ones_2, zeros_3], axis=0)
+                mask_chunk_predictor_cur = np.concatenate([zeros_1, ones_zeros], axis=0)
+                mask_chunk_predictor = np.concatenate(
+                    [mask_chunk_predictor, mask_chunk_predictor_cur], axis=0
+                )
+
+                # encoder att mask
+                zeros_1_top = np.zeros([shfit_fsmn, chunk_num * chunk_size_pad_shift], dtype=dtype)
+
+                zeros_2_num = max(chunk_ids - encoder_att_look_back_factor, 0)
+                zeros_2 = np.zeros([chunk_size, zeros_2_num * chunk_size_pad_shift], dtype=dtype)
+
+                encoder_att_look_back_num = max(chunk_ids - zeros_2_num, 0)
+                zeros_2_left = np.zeros([chunk_size, shfit_fsmn], dtype=dtype)
+                ones_2_mid = np.ones([stride, stride], dtype=dtype)
+                zeros_2_bottom = np.zeros([chunk_size - stride, stride], dtype=dtype)
+                zeros_2_right = np.zeros([chunk_size, chunk_size - stride], dtype=dtype)
+                ones_2 = np.concatenate([ones_2_mid, zeros_2_bottom], axis=0)
+                ones_2 = np.concatenate([zeros_2_left, ones_2, zeros_2_right], axis=1)
+                ones_2 = np.tile(ones_2, [1, encoder_att_look_back_num])
+
+                zeros_3_left = np.zeros([chunk_size, shfit_fsmn], dtype=dtype)
+                ones_3_right = np.ones([chunk_size, chunk_size], dtype=dtype)
+                ones_3 = np.concatenate([zeros_3_left, ones_3_right], axis=1)
+
+                zeros_remain_num = max(chunk_num - 1 - chunk_ids, 0)
+                zeros_remain = np.zeros(
+                    [chunk_size, zeros_remain_num * chunk_size_pad_shift], dtype=dtype
+                )
+
+                ones2_bottom = np.concatenate([zeros_2, ones_2, ones_3, zeros_remain], axis=1)
+                mask_att_chunk_encoder_cur = np.concatenate([zeros_1_top, ones2_bottom], axis=0)
+                mask_att_chunk_encoder = np.concatenate(
+                    [mask_att_chunk_encoder, mask_att_chunk_encoder_cur], axis=0
+                )
+
+                # decoder fsmn_shift_att_mask
+                zeros_1 = np.zeros([shfit_fsmn, 1])
+                ones_1 = np.ones([chunk_size, 1])
+                mask_shift_att_chunk_decoder_cur = np.concatenate([zeros_1, ones_1], axis=0)
+                mask_shift_att_chunk_decoder = np.concatenate(
+                    [mask_shift_att_chunk_decoder, mask_shift_att_chunk_decoder_cur], axis=0
+                )
+
+            self.x_add_mask = x_add_mask[:x_len_chunk_max, : x_len_max + pad_left]
+            self.x_len_chunk = x_len_chunk
+            self.x_rm_mask = x_rm_mask[:x_len_max, :x_len_chunk_max]
+            self.x_len = x_len
+            self.mask_shfit_chunk = mask_shfit_chunk[:x_len_chunk_max, :]
+            self.mask_chunk_predictor = mask_chunk_predictor[:x_len_chunk_max, :]
+            self.mask_att_chunk_encoder = mask_att_chunk_encoder[:x_len_chunk_max, :x_len_chunk_max]
+            self.mask_shift_att_chunk_decoder = mask_shift_att_chunk_decoder[:x_len_chunk_max, :]
+            self.chunk_outs = (
+                self.x_add_mask,
+                self.x_len_chunk,
+                self.x_rm_mask,
+                self.x_len,
+                self.mask_shfit_chunk,
+                self.mask_chunk_predictor,
+                self.mask_att_chunk_encoder,
+                self.mask_shift_att_chunk_decoder,
+            )
+
+        return self.chunk_outs
+
+    def split_chunk(self, x, x_len, chunk_outs):
+        """
+        :param x: (b, t, d)
+        :param x_length: (b)
+        :param ind: int
+        :return:
+        """
+        x = x[:, : x_len.max(), :]
+        b, t, d = x.size()
+        x_len_mask = (~make_pad_mask(x_len, maxlen=t)).to(x.device)
+        x *= x_len_mask[:, :, None]
+
+        x_add_mask = self.get_x_add_mask(chunk_outs, x.device, dtype=x.dtype)
+        x_len_chunk = self.get_x_len_chunk(chunk_outs, x_len.device, dtype=x_len.dtype)
+        pad = (0, 0, self.pad_left_cur, 0)
+        x = F.pad(x, pad, "constant", 0.0)
+        b, t, d = x.size()
+        x = torch.transpose(x, 1, 0)
+        x = torch.reshape(x, [t, -1])
+        x_chunk = torch.mm(x_add_mask, x)
+        x_chunk = torch.reshape(x_chunk, [-1, b, d]).transpose(1, 0)
+
+        return x_chunk, x_len_chunk
+
+    def remove_chunk(self, x_chunk, x_len_chunk, chunk_outs):
+        x_chunk = x_chunk[:, : x_len_chunk.max(), :]
+        b, t, d = x_chunk.size()
+        x_len_chunk_mask = (~make_pad_mask(x_len_chunk, maxlen=t)).to(x_chunk.device)
+        x_chunk *= x_len_chunk_mask[:, :, None]
+
+        x_rm_mask = self.get_x_rm_mask(chunk_outs, x_chunk.device, dtype=x_chunk.dtype)
+        x_len = self.get_x_len(chunk_outs, x_len_chunk.device, dtype=x_len_chunk.dtype)
+        x_chunk = torch.transpose(x_chunk, 1, 0)
+        x_chunk = torch.reshape(x_chunk, [t, -1])
+        x = torch.mm(x_rm_mask, x_chunk)
+        x = torch.reshape(x, [-1, b, d]).transpose(1, 0)
+
+        return x, x_len
+
+    def get_x_add_mask(self, chunk_outs=None, device="cpu", idx=0, dtype=torch.float32):
+        with torch.no_grad():
+            x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx]
+            x = torch.from_numpy(x).type(dtype).to(device)
+        return x
+
+    def get_x_len_chunk(self, chunk_outs=None, device="cpu", idx=1, dtype=torch.float32):
+        with torch.no_grad():
+            x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx]
+            x = torch.from_numpy(x).type(dtype).to(device)
+        return x
+
+    def get_x_rm_mask(self, chunk_outs=None, device="cpu", idx=2, dtype=torch.float32):
+        with torch.no_grad():
+            x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx]
+            x = torch.from_numpy(x).type(dtype).to(device)
+        return x
+
+    def get_x_len(self, chunk_outs=None, device="cpu", idx=3, dtype=torch.float32):
+        with torch.no_grad():
+            x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx]
+            x = torch.from_numpy(x).type(dtype).to(device)
+        return x
+
+    def get_mask_shfit_chunk(
+        self, chunk_outs=None, device="cpu", batch_size=1, num_units=1, idx=4, dtype=torch.float32
+    ):
+        with torch.no_grad():
+            x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx]
+            x = np.tile(
+                x[
+                    None,
+                    :,
+                    :,
+                ],
+                [batch_size, 1, num_units],
+            )
+            x = torch.from_numpy(x).type(dtype).to(device)
+        return x
+
+    def get_mask_chunk_predictor(
+        self, chunk_outs=None, device="cpu", batch_size=1, num_units=1, idx=5, dtype=torch.float32
+    ):
+        with torch.no_grad():
+            x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx]
+            x = np.tile(
+                x[
+                    None,
+                    :,
+                    :,
+                ],
+                [batch_size, 1, num_units],
+            )
+            x = torch.from_numpy(x).type(dtype).to(device)
+        return x
+
+    def get_mask_att_chunk_encoder(
+        self, chunk_outs=None, device="cpu", batch_size=1, idx=6, dtype=torch.float32
+    ):
+        with torch.no_grad():
+            x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx]
+            x = np.tile(
+                x[
+                    None,
+                    :,
+                    :,
+                ],
+                [batch_size, 1, 1],
+            )
+            x = torch.from_numpy(x).type(dtype).to(device)
+        return x
+
+    def get_mask_shift_att_chunk_decoder(
+        self, chunk_outs=None, device="cpu", batch_size=1, idx=7, dtype=torch.float32
+    ):
+        with torch.no_grad():
+            x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx]
+            x = np.tile(x[None, None, :, 0], [batch_size, 1, 1])
+            x = torch.from_numpy(x).type(dtype).to(device)
+        return x
+
+
+def build_scama_mask_for_cross_attention_decoder(
+    predictor_alignments: torch.Tensor,
+    encoder_sequence_length: torch.Tensor,
+    chunk_size: int = 5,
+    encoder_chunk_size: int = 5,
+    attention_chunk_center_bias: int = 0,
+    attention_chunk_size: int = 1,
+    attention_chunk_type: str = "chunk",
+    step=None,
+    predictor_mask_chunk_hopping: torch.Tensor = None,
+    decoder_att_look_back_factor: int = 1,
+    mask_shift_att_chunk_decoder: torch.Tensor = None,
+    target_length: torch.Tensor = None,
+    is_training=True,
+    dtype: torch.dtype = torch.float32,
+):
+    with torch.no_grad():
+        device = predictor_alignments.device
+        batch_size, chunk_num = predictor_alignments.size()
+        maximum_encoder_length = encoder_sequence_length.max().item()
+        int_type = predictor_alignments.dtype
+        if not is_training:
+            target_length = predictor_alignments.sum(dim=-1).type(encoder_sequence_length.dtype)
+        maximum_target_length = target_length.max()
+        predictor_alignments_cumsum = torch.cumsum(predictor_alignments, dim=1)
+        predictor_alignments_cumsum = predictor_alignments_cumsum[:, None, :].repeat(
+            1, maximum_target_length, 1
+        )
+
+        index = torch.ones([batch_size, maximum_target_length], dtype=int_type).to(device)
+        index = torch.cumsum(index, dim=1)
+        index = index[:, :, None].repeat(1, 1, chunk_num)
+
+        index_div = torch.floor(torch.divide(predictor_alignments_cumsum, index)).type(int_type)
+        index_div_bool_zeros = index_div == 0
+        index_div_bool_zeros_count = torch.sum(index_div_bool_zeros.type(int_type), dim=-1) + 1
+
+        index_div_bool_zeros_count = torch.clip(index_div_bool_zeros_count, min=1, max=chunk_num)
+
+        index_div_bool_zeros_count *= chunk_size
+        index_div_bool_zeros_count += attention_chunk_center_bias
+        index_div_bool_zeros_count = torch.clip(
+            index_div_bool_zeros_count - 1, min=0, max=maximum_encoder_length
+        )
+        index_div_bool_zeros_count_ori = index_div_bool_zeros_count
+
+        index_div_bool_zeros_count = (
+            torch.floor(index_div_bool_zeros_count / encoder_chunk_size) + 1
+        ) * encoder_chunk_size
+        max_len_chunk = math.ceil(maximum_encoder_length / encoder_chunk_size) * encoder_chunk_size
+
+        mask_flip, mask_flip2 = None, None
+        if attention_chunk_size is not None:
+            index_div_bool_zeros_count_beg = index_div_bool_zeros_count - attention_chunk_size
+            index_div_bool_zeros_count_beg = torch.clip(
+                index_div_bool_zeros_count_beg, 0, max_len_chunk
+            )
+            index_div_bool_zeros_count_beg_mask = sequence_mask(
+                index_div_bool_zeros_count_beg, maxlen=max_len_chunk, dtype=int_type, device=device
+            )
+            mask_flip = 1 - index_div_bool_zeros_count_beg_mask
+            attention_chunk_size2 = attention_chunk_size * (decoder_att_look_back_factor + 1)
+            index_div_bool_zeros_count_beg = index_div_bool_zeros_count - attention_chunk_size2
+
+            index_div_bool_zeros_count_beg = torch.clip(
+                index_div_bool_zeros_count_beg, 0, max_len_chunk
+            )
+            index_div_bool_zeros_count_beg_mask = sequence_mask(
+                index_div_bool_zeros_count_beg, maxlen=max_len_chunk, dtype=int_type, device=device
+            )
+            mask_flip2 = 1 - index_div_bool_zeros_count_beg_mask
+
+        mask = sequence_mask(
+            index_div_bool_zeros_count, maxlen=max_len_chunk, dtype=dtype, device=device
+        )
+
+        if predictor_mask_chunk_hopping is not None:
+            b, k, t = mask.size()
+            predictor_mask_chunk_hopping = predictor_mask_chunk_hopping[:, None, :, 0].repeat(
+                1, k, 1
+            )
+
+            mask_mask_flip = mask
+            if mask_flip is not None:
+                mask_mask_flip = mask_flip * mask
+
+            def _fn():
+                mask_sliced = mask[:b, :k, encoder_chunk_size:t]
+                zero_pad_right = torch.zeros(
+                    [b, k, encoder_chunk_size], dtype=mask_sliced.dtype
+                ).to(device)
+                mask_sliced = torch.cat([mask_sliced, zero_pad_right], dim=2)
+                _, _, tt = predictor_mask_chunk_hopping.size()
+                pad_right_p = max_len_chunk - tt
+                predictor_mask_chunk_hopping_pad = torch.nn.functional.pad(
+                    predictor_mask_chunk_hopping, [0, pad_right_p], "constant", 0
+                )
+                masked = mask_sliced * predictor_mask_chunk_hopping_pad
+
+                mask_true = mask_mask_flip + masked
+                return mask_true
+
+            mask = _fn() if t > chunk_size else mask_mask_flip
+
+        if mask_flip2 is not None:
+            mask *= mask_flip2
+
+        mask_target = sequence_mask(
+            target_length, maxlen=maximum_target_length, dtype=mask.dtype, device=device
+        )
+        mask = mask[:, :maximum_target_length, :] * mask_target[:, :, None]
+
+        mask_len = sequence_mask(
+            encoder_sequence_length, maxlen=maximum_encoder_length, dtype=mask.dtype, device=device
+        )
+        mask = mask[:, :, :maximum_encoder_length] * mask_len[:, None, :]
+
+        if attention_chunk_type == "full":
+            mask = torch.ones_like(mask).to(device)
+        if mask_shift_att_chunk_decoder is not None:
+            mask = mask * mask_shift_att_chunk_decoder
+        mask = mask[:, :maximum_target_length, :maximum_encoder_length].type(dtype).to(device)
+
+    return mask

funasr/models/scama/decoder.py

@@ -0,0 +1,498 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+from typing import List
+from typing import Tuple
+import logging
+import torch
+import torch.nn as nn
+import numpy as np
+
+from funasr.models.scama import utils as myutils
+from funasr.models.transformer.decoder import BaseTransformerDecoder
+
+from funasr.models.sanm.attention import (
+    MultiHeadedAttentionSANMDecoder,
+    MultiHeadedAttentionCrossAtt,
+)
+from funasr.models.transformer.embedding import PositionalEncoding
+from funasr.models.transformer.layer_norm import LayerNorm
+from funasr.models.sanm.positionwise_feed_forward import PositionwiseFeedForwardDecoderSANM
+from funasr.models.transformer.utils.repeat import repeat
+
+from funasr.register import tables
+
+
+class DecoderLayerSANM(nn.Module):
+    """Single decoder layer module.
+
+    Args:
+        size (int): Input dimension.
+        self_attn (torch.nn.Module): Self-attention module instance.
+            `MultiHeadedAttention` instance can be used as the argument.
+        src_attn (torch.nn.Module): Self-attention module instance.
+            `MultiHeadedAttention` instance can be used as the argument.
+        feed_forward (torch.nn.Module): Feed-forward module instance.
+            `PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance
+            can be used as the argument.
+        dropout_rate (float): Dropout rate.
+        normalize_before (bool): Whether to use layer_norm before the first block.
+        concat_after (bool): Whether to concat attention layer's input and output.
+            if True, additional linear will be applied.
+            i.e. x -> x + linear(concat(x, att(x)))
+            if False, no additional linear will be applied. i.e. x -> x + att(x)
+
+
+    """
+
+    def __init__(
+        self,
+        size,
+        self_attn,
+        src_attn,
+        feed_forward,
+        dropout_rate,
+        normalize_before=True,
+        concat_after=False,
+    ):
+        """Construct an DecoderLayer object."""
+        super(DecoderLayerSANM, self).__init__()
+        self.size = size
+        self.self_attn = self_attn
+        self.src_attn = src_attn
+        self.feed_forward = feed_forward
+        self.norm1 = LayerNorm(size)
+        if self_attn is not None:
+            self.norm2 = LayerNorm(size)
+        if src_attn is not None:
+            self.norm3 = LayerNorm(size)
+        self.dropout = nn.Dropout(dropout_rate)
+        self.normalize_before = normalize_before
+        self.concat_after = concat_after
+        if self.concat_after:
+            self.concat_linear1 = nn.Linear(size + size, size)
+            self.concat_linear2 = nn.Linear(size + size, size)
+
+    def forward(self, tgt, tgt_mask, memory, memory_mask=None, cache=None):
+        """Compute decoded features.
+
+        Args:
+            tgt (torch.Tensor): Input tensor (#batch, maxlen_out, size).
+            tgt_mask (torch.Tensor): Mask for input tensor (#batch, maxlen_out).
+            memory (torch.Tensor): Encoded memory, float32 (#batch, maxlen_in, size).
+            memory_mask (torch.Tensor): Encoded memory mask (#batch, maxlen_in).
+            cache (List[torch.Tensor]): List of cached tensors.
+                Each tensor shape should be (#batch, maxlen_out - 1, size).
+
+        Returns:
+            torch.Tensor: Output tensor(#batch, maxlen_out, size).
+            torch.Tensor: Mask for output tensor (#batch, maxlen_out).
+            torch.Tensor: Encoded memory (#batch, maxlen_in, size).
+            torch.Tensor: Encoded memory mask (#batch, maxlen_in).
+
+        """
+        # tgt = self.dropout(tgt)
+        residual = tgt
+        if self.normalize_before:
+            tgt = self.norm1(tgt)
+        tgt = self.feed_forward(tgt)
+
+        x = tgt
+        if self.self_attn:
+            if self.normalize_before:
+                tgt = self.norm2(tgt)
+            x, _ = self.self_attn(tgt, tgt_mask)
+            x = residual + self.dropout(x)
+
+        if self.src_attn is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.norm3(x)
+
+            x = residual + self.dropout(self.src_attn(x, memory, memory_mask))
+
+        return x, tgt_mask, memory, memory_mask, cache
+
+    def forward_one_step(self, tgt, tgt_mask, memory, memory_mask=None, cache=None):
+        """Compute decoded features.
+
+        Args:
+            tgt (torch.Tensor): Input tensor (#batch, maxlen_out, size).
+            tgt_mask (torch.Tensor): Mask for input tensor (#batch, maxlen_out).
+            memory (torch.Tensor): Encoded memory, float32 (#batch, maxlen_in, size).
+            memory_mask (torch.Tensor): Encoded memory mask (#batch, maxlen_in).
+            cache (List[torch.Tensor]): List of cached tensors.
+                Each tensor shape should be (#batch, maxlen_out - 1, size).
+
+        Returns:
+            torch.Tensor: Output tensor(#batch, maxlen_out, size).
+            torch.Tensor: Mask for output tensor (#batch, maxlen_out).
+            torch.Tensor: Encoded memory (#batch, maxlen_in, size).
+            torch.Tensor: Encoded memory mask (#batch, maxlen_in).
+
+        """
+        # tgt = self.dropout(tgt)
+        residual = tgt
+        if self.normalize_before:
+            tgt = self.norm1(tgt)
+        tgt = self.feed_forward(tgt)
+
+        x = tgt
+        if self.self_attn:
+            if self.normalize_before:
+                tgt = self.norm2(tgt)
+            if self.training:
+                cache = None
+            x, cache = self.self_attn(tgt, tgt_mask, cache=cache)
+            x = residual + self.dropout(x)
+
+        if self.src_attn is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.norm3(x)
+
+            x = residual + self.dropout(self.src_attn(x, memory, memory_mask))
+
+        return x, tgt_mask, memory, memory_mask, cache
+
+    def forward_chunk(
+        self, tgt, memory, fsmn_cache=None, opt_cache=None, chunk_size=None, look_back=0
+    ):
+        """Compute decoded features.
+
+        Args:
+            tgt (torch.Tensor): Input tensor (#batch, maxlen_out, size).
+            tgt_mask (torch.Tensor): Mask for input tensor (#batch, maxlen_out).
+            memory (torch.Tensor): Encoded memory, float32 (#batch, maxlen_in, size).
+            memory_mask (torch.Tensor): Encoded memory mask (#batch, maxlen_in).
+            cache (List[torch.Tensor]): List of cached tensors.
+                Each tensor shape should be (#batch, maxlen_out - 1, size).
+
+        Returns:
+            torch.Tensor: Output tensor(#batch, maxlen_out, size).
+            torch.Tensor: Mask for output tensor (#batch, maxlen_out).
+            torch.Tensor: Encoded memory (#batch, maxlen_in, size).
+            torch.Tensor: Encoded memory mask (#batch, maxlen_in).
+
+        """
+        residual = tgt
+        if self.normalize_before:
+            tgt = self.norm1(tgt)
+        tgt = self.feed_forward(tgt)
+
+        x = tgt
+        if self.self_attn:
+            if self.normalize_before:
+                tgt = self.norm2(tgt)
+            x, fsmn_cache = self.self_attn(tgt, None, fsmn_cache)
+            x = residual + self.dropout(x)
+
+        if self.src_attn is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.norm3(x)
+
+            x, opt_cache = self.src_attn.forward_chunk(x, memory, opt_cache, chunk_size, look_back)
+            x = residual + x
+
+        return x, memory, fsmn_cache, opt_cache
+
+
+@tables.register("decoder_classes", "FsmnDecoderSCAMAOpt")
+class FsmnDecoderSCAMAOpt(BaseTransformerDecoder):
+    """
+    Author: Shiliang Zhang, Zhifu Gao, Haoneng Luo, Ming Lei, Jie Gao, Zhijie Yan, Lei Xie
+    SCAMA: Streaming chunk-aware multihead attention for online end-to-end speech recognition
+    https://arxiv.org/abs/2006.01712
+    """
+
+    def __init__(
+        self,
+        vocab_size: int,
+        encoder_output_size: int,
+        attention_heads: int = 4,
+        linear_units: int = 2048,
+        num_blocks: int = 6,
+        dropout_rate: float = 0.1,
+        positional_dropout_rate: float = 0.1,
+        self_attention_dropout_rate: float = 0.0,
+        src_attention_dropout_rate: float = 0.0,
+        input_layer: str = "embed",
+        use_output_layer: bool = True,
+        pos_enc_class=PositionalEncoding,
+        normalize_before: bool = True,
+        concat_after: bool = False,
+        att_layer_num: int = 6,
+        kernel_size: int = 21,
+        sanm_shfit: int = None,
+        concat_embeds: bool = False,
+        attention_dim: int = None,
+        tf2torch_tensor_name_prefix_torch: str = "decoder",
+        tf2torch_tensor_name_prefix_tf: str = "seq2seq/decoder",
+        embed_tensor_name_prefix_tf: str = None,
+    ):
+        super().__init__(
+            vocab_size=vocab_size,
+            encoder_output_size=encoder_output_size,
+            dropout_rate=dropout_rate,
+            positional_dropout_rate=positional_dropout_rate,
+            input_layer=input_layer,
+            use_output_layer=use_output_layer,
+            pos_enc_class=pos_enc_class,
+            normalize_before=normalize_before,
+        )
+        if attention_dim is None:
+            attention_dim = encoder_output_size
+
+        if input_layer == "embed":
+            self.embed = torch.nn.Sequential(
+                torch.nn.Embedding(vocab_size, attention_dim),
+            )
+        elif input_layer == "linear":
+            self.embed = torch.nn.Sequential(
+                torch.nn.Linear(vocab_size, attention_dim),
+                torch.nn.LayerNorm(attention_dim),
+                torch.nn.Dropout(dropout_rate),
+                torch.nn.ReLU(),
+                pos_enc_class(attention_dim, positional_dropout_rate),
+            )
+        else:
+            raise ValueError(f"only 'embed' or 'linear' is supported: {input_layer}")
+
+        self.normalize_before = normalize_before
+        if self.normalize_before:
+            self.after_norm = LayerNorm(attention_dim)
+        if use_output_layer:
+            self.output_layer = torch.nn.Linear(attention_dim, vocab_size)
+        else:
+            self.output_layer = None
+
+        self.att_layer_num = att_layer_num
+        self.num_blocks = num_blocks
+        if sanm_shfit is None:
+            sanm_shfit = (kernel_size - 1) // 2
+        self.decoders = repeat(
+            att_layer_num,
+            lambda lnum: DecoderLayerSANM(
+                attention_dim,
+                MultiHeadedAttentionSANMDecoder(
+                    attention_dim, self_attention_dropout_rate, kernel_size, sanm_shfit=sanm_shfit
+                ),
+                MultiHeadedAttentionCrossAtt(
+                    attention_heads,
+                    attention_dim,
+                    src_attention_dropout_rate,
+                    encoder_output_size=encoder_output_size,
+                ),
+                PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate),
+                dropout_rate,
+                normalize_before,
+                concat_after,
+            ),
+        )
+        if num_blocks - att_layer_num <= 0:
+            self.decoders2 = None
+        else:
+            self.decoders2 = repeat(
+                num_blocks - att_layer_num,
+                lambda lnum: DecoderLayerSANM(
+                    attention_dim,
+                    MultiHeadedAttentionSANMDecoder(
+                        attention_dim,
+                        self_attention_dropout_rate,
+                        kernel_size,
+                        sanm_shfit=sanm_shfit,
+                    ),
+                    None,
+                    PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate),
+                    dropout_rate,
+                    normalize_before,
+                    concat_after,
+                ),
+            )
+
+        self.decoders3 = repeat(
+            1,
+            lambda lnum: DecoderLayerSANM(
+                attention_dim,
+                None,
+                None,
+                PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate),
+                dropout_rate,
+                normalize_before,
+                concat_after,
+            ),
+        )
+        if concat_embeds:
+            self.embed_concat_ffn = repeat(
+                1,
+                lambda lnum: DecoderLayerSANM(
+                    attention_dim + encoder_output_size,
+                    None,
+                    None,
+                    PositionwiseFeedForwardDecoderSANM(
+                        attention_dim + encoder_output_size,
+                        linear_units,
+                        dropout_rate,
+                        adim=attention_dim,
+                    ),
+                    dropout_rate,
+                    normalize_before,
+                    concat_after,
+                ),
+            )
+        else:
+            self.embed_concat_ffn = None
+        self.concat_embeds = concat_embeds
+        self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch
+        self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf
+        self.embed_tensor_name_prefix_tf = embed_tensor_name_prefix_tf
+
+    def forward(
+        self,
+        hs_pad: torch.Tensor,
+        hlens: torch.Tensor,
+        ys_in_pad: torch.Tensor,
+        ys_in_lens: torch.Tensor,
+        chunk_mask: torch.Tensor = None,
+        pre_acoustic_embeds: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward decoder.
+
+        Args:
+            hs_pad: encoded memory, float32  (batch, maxlen_in, feat)
+            hlens: (batch)
+            ys_in_pad:
+                input token ids, int64 (batch, maxlen_out)
+                if input_layer == "embed"
+                input tensor (batch, maxlen_out, #mels) in the other cases
+            ys_in_lens: (batch)
+        Returns:
+            (tuple): tuple containing:
+
+            x: decoded token score before softmax (batch, maxlen_out, token)
+                if use_output_layer is True,
+            olens: (batch, )
+        """
+        tgt = ys_in_pad
+        tgt_mask = myutils.sequence_mask(ys_in_lens, device=tgt.device)[:, :, None]
+
+        memory = hs_pad
+        memory_mask = myutils.sequence_mask(hlens, device=memory.device)[:, None, :]
+        if chunk_mask is not None:
+            memory_mask = memory_mask * chunk_mask
+            if tgt_mask.size(1) != memory_mask.size(1):
+                memory_mask = torch.cat((memory_mask, memory_mask[:, -2:-1, :]), dim=1)
+
+        x = self.embed(tgt)
+
+        if pre_acoustic_embeds is not None and self.concat_embeds:
+            x = torch.cat((x, pre_acoustic_embeds), dim=-1)
+            x, _, _, _, _ = self.embed_concat_ffn(x, None, None, None, None)
+
+        x, tgt_mask, memory, memory_mask, _ = self.decoders(x, tgt_mask, memory, memory_mask)
+        if self.decoders2 is not None:
+            x, tgt_mask, memory, memory_mask, _ = self.decoders2(x, tgt_mask, memory, memory_mask)
+        x, tgt_mask, memory, memory_mask, _ = self.decoders3(x, tgt_mask, memory, memory_mask)
+        if self.normalize_before:
+            x = self.after_norm(x)
+        if self.output_layer is not None:
+            x = self.output_layer(x)
+
+        olens = tgt_mask.sum(1)
+        return x, olens
+
+    def score(
+        self,
+        ys,
+        state,
+        x,
+        x_mask=None,
+        pre_acoustic_embeds: torch.Tensor = None,
+    ):
+        """Score."""
+        ys_mask = myutils.sequence_mask(
+            torch.tensor([len(ys)], dtype=torch.int32), device=x.device
+        )[:, :, None]
+        logp, state = self.forward_one_step(
+            ys.unsqueeze(0),
+            ys_mask,
+            x.unsqueeze(0),
+            memory_mask=x_mask,
+            pre_acoustic_embeds=pre_acoustic_embeds,
+            cache=state,
+        )
+        return logp.squeeze(0), state
+
+    def forward_one_step(
+        self,
+        tgt: torch.Tensor,
+        tgt_mask: torch.Tensor,
+        memory: torch.Tensor,
+        memory_mask: torch.Tensor = None,
+        pre_acoustic_embeds: torch.Tensor = None,
+        cache: List[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+        """Forward one step.
+
+        Args:
+            tgt: input token ids, int64 (batch, maxlen_out)
+            tgt_mask: input token mask,  (batch, maxlen_out)
+                      dtype=torch.uint8 in PyTorch 1.2-
+                      dtype=torch.bool in PyTorch 1.2+ (include 1.2)
+            memory: encoded memory, float32  (batch, maxlen_in, feat)
+            cache: cached output list of (batch, max_time_out-1, size)
+        Returns:
+            y, cache: NN output value and cache per `self.decoders`.
+            y.shape` is (batch, maxlen_out, token)
+        """
+
+        x = tgt[:, -1:]
+        tgt_mask = None
+        x = self.embed(x)
+
+        if pre_acoustic_embeds is not None and self.concat_embeds:
+            x = torch.cat((x, pre_acoustic_embeds), dim=-1)
+            x, _, _, _, _ = self.embed_concat_ffn(x, None, None, None, None)
+
+        if cache is None:
+            cache_layer_num = len(self.decoders)
+            if self.decoders2 is not None:
+                cache_layer_num += len(self.decoders2)
+            cache = [None] * cache_layer_num
+        new_cache = []
+        # for c, decoder in zip(cache, self.decoders):
+        for i in range(self.att_layer_num):
+            decoder = self.decoders[i]
+            c = cache[i]
+            x, tgt_mask, memory, memory_mask, c_ret = decoder.forward_one_step(
+                x, tgt_mask, memory, memory_mask, cache=c
+            )
+            new_cache.append(c_ret)
+
+        if self.num_blocks - self.att_layer_num >= 1:
+            for i in range(self.num_blocks - self.att_layer_num):
+                j = i + self.att_layer_num
+                decoder = self.decoders2[i]
+                c = cache[j]
+                x, tgt_mask, memory, memory_mask, c_ret = decoder.forward_one_step(
+                    x, tgt_mask, memory, memory_mask, cache=c
+                )
+                new_cache.append(c_ret)
+
+        for decoder in self.decoders3:
+            x, tgt_mask, memory, memory_mask, _ = decoder.forward_one_step(
+                x, tgt_mask, memory, None, cache=None
+            )
+
+        if self.normalize_before:
+            y = self.after_norm(x[:, -1])
+        else:
+            y = x[:, -1]
+        if self.output_layer is not None:
+            y = self.output_layer(y)
+            y = torch.log_softmax(y, dim=-1)
+
+        return y, new_cache

funasr/models/scama/encoder.py

@@ -0,0 +1,499 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+from typing import List
+from typing import Optional
+from typing import Sequence
+from typing import Tuple
+from typing import Union
+import logging
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from funasr.models.scama.chunk_utilis import overlap_chunk
+import numpy as np
+from funasr.train_utils.device_funcs import to_device
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+from funasr.models.sanm.attention import MultiHeadedAttention, MultiHeadedAttentionSANM
+from funasr.models.transformer.embedding import (
+    SinusoidalPositionEncoder,
+    StreamSinusoidalPositionEncoder,
+)
+from funasr.models.transformer.layer_norm import LayerNorm
+from funasr.models.transformer.utils.multi_layer_conv import Conv1dLinear
+from funasr.models.transformer.utils.multi_layer_conv import MultiLayeredConv1d
+from funasr.models.transformer.positionwise_feed_forward import (
+    PositionwiseFeedForward,  # noqa: H301
+)
+from funasr.models.transformer.utils.repeat import repeat
+from funasr.models.transformer.utils.subsampling import Conv2dSubsampling
+from funasr.models.transformer.utils.subsampling import Conv2dSubsampling2
+from funasr.models.transformer.utils.subsampling import Conv2dSubsampling6
+from funasr.models.transformer.utils.subsampling import Conv2dSubsampling8
+from funasr.models.transformer.utils.subsampling import TooShortUttError
+from funasr.models.transformer.utils.subsampling import check_short_utt
+from funasr.models.transformer.utils.mask import subsequent_mask, vad_mask
+
+from funasr.models.ctc.ctc import CTC
+
+from funasr.register import tables
+
+
+class EncoderLayerSANM(nn.Module):
+    def __init__(
+        self,
+        in_size,
+        size,
+        self_attn,
+        feed_forward,
+        dropout_rate,
+        normalize_before=True,
+        concat_after=False,
+        stochastic_depth_rate=0.0,
+    ):
+        """Construct an EncoderLayer object."""
+        super(EncoderLayerSANM, self).__init__()
+        self.self_attn = self_attn
+        self.feed_forward = feed_forward
+        self.norm1 = LayerNorm(in_size)
+        self.norm2 = LayerNorm(size)
+        self.dropout = nn.Dropout(dropout_rate)
+        self.in_size = in_size
+        self.size = size
+        self.normalize_before = normalize_before
+        self.concat_after = concat_after
+        if self.concat_after:
+            self.concat_linear = nn.Linear(size + size, size)
+        self.stochastic_depth_rate = stochastic_depth_rate
+        self.dropout_rate = dropout_rate
+
+    def forward(self, x, mask, cache=None, mask_shfit_chunk=None, mask_att_chunk_encoder=None):
+        """Compute encoded features.
+
+        Args:
+            x_input (torch.Tensor): Input tensor (#batch, time, size).
+            mask (torch.Tensor): Mask tensor for the input (#batch, time).
+            cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size).
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time, size).
+            torch.Tensor: Mask tensor (#batch, time).
+
+        """
+        skip_layer = False
+        # with stochastic depth, residual connection `x + f(x)` becomes
+        # `x <- x + 1 / (1 - p) * f(x)` at training time.
+        stoch_layer_coeff = 1.0
+        if self.training and self.stochastic_depth_rate > 0:
+            skip_layer = torch.rand(1).item() < self.stochastic_depth_rate
+            stoch_layer_coeff = 1.0 / (1 - self.stochastic_depth_rate)
+
+        if skip_layer:
+            if cache is not None:
+                x = torch.cat([cache, x], dim=1)
+            return x, mask
+
+        residual = x
+        if self.normalize_before:
+            x = self.norm1(x)
+
+        if self.concat_after:
+            x_concat = torch.cat(
+                (
+                    x,
+                    self.self_attn(
+                        x,
+                        mask,
+                        mask_shfit_chunk=mask_shfit_chunk,
+                        mask_att_chunk_encoder=mask_att_chunk_encoder,
+                    ),
+                ),
+                dim=-1,
+            )
+            if self.in_size == self.size:
+                x = residual + stoch_layer_coeff * self.concat_linear(x_concat)
+            else:
+                x = stoch_layer_coeff * self.concat_linear(x_concat)
+        else:
+            if self.in_size == self.size:
+                x = residual + stoch_layer_coeff * self.dropout(
+                    self.self_attn(
+                        x,
+                        mask,
+                        mask_shfit_chunk=mask_shfit_chunk,
+                        mask_att_chunk_encoder=mask_att_chunk_encoder,
+                    )
+                )
+            else:
+                x = stoch_layer_coeff * self.dropout(
+                    self.self_attn(
+                        x,
+                        mask,
+                        mask_shfit_chunk=mask_shfit_chunk,
+                        mask_att_chunk_encoder=mask_att_chunk_encoder,
+                    )
+                )
+        if not self.normalize_before:
+            x = self.norm1(x)
+
+        residual = x
+        if self.normalize_before:
+            x = self.norm2(x)
+        x = residual + stoch_layer_coeff * self.dropout(self.feed_forward(x))
+        if not self.normalize_before:
+            x = self.norm2(x)
+
+        return x, mask, cache, mask_shfit_chunk, mask_att_chunk_encoder
+
+    def forward_chunk(self, x, cache=None, chunk_size=None, look_back=0):
+        """Compute encoded features.
+
+        Args:
+            x_input (torch.Tensor): Input tensor (#batch, time, size).
+            mask (torch.Tensor): Mask tensor for the input (#batch, time).
+            cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size).
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time, size).
+            torch.Tensor: Mask tensor (#batch, time).
+
+        """
+
+        residual = x
+        if self.normalize_before:
+            x = self.norm1(x)
+
+        if self.in_size == self.size:
+            attn, cache = self.self_attn.forward_chunk(x, cache, chunk_size, look_back)
+            x = residual + attn
+        else:
+            x, cache = self.self_attn.forward_chunk(x, cache, chunk_size, look_back)
+
+        if not self.normalize_before:
+            x = self.norm1(x)
+
+        residual = x
+        if self.normalize_before:
+            x = self.norm2(x)
+        x = residual + self.feed_forward(x)
+        if not self.normalize_before:
+            x = self.norm2(x)
+
+        return x, cache
+
+
+@tables.register("encoder_classes", "SANMEncoderChunkOpt")
+class SANMEncoderChunkOpt(nn.Module):
+    """
+    Author: Shiliang Zhang, Zhifu Gao, Haoneng Luo, Ming Lei, Jie Gao, Zhijie Yan, Lei Xie
+    SCAMA: Streaming chunk-aware multihead attention for online end-to-end speech recognition
+    https://arxiv.org/abs/2006.01712
+    """
+
+    def __init__(
+        self,
+        input_size: int,
+        output_size: int = 256,
+        attention_heads: int = 4,
+        linear_units: int = 2048,
+        num_blocks: int = 6,
+        dropout_rate: float = 0.1,
+        positional_dropout_rate: float = 0.1,
+        attention_dropout_rate: float = 0.0,
+        input_layer: Optional[str] = "conv2d",
+        pos_enc_class=SinusoidalPositionEncoder,
+        normalize_before: bool = True,
+        concat_after: bool = False,
+        positionwise_layer_type: str = "linear",
+        positionwise_conv_kernel_size: int = 1,
+        padding_idx: int = -1,
+        interctc_layer_idx: List[int] = [],
+        interctc_use_conditioning: bool = False,
+        kernel_size: int = 11,
+        sanm_shfit: int = 0,
+        selfattention_layer_type: str = "sanm",
+        chunk_size: Union[int, Sequence[int]] = (16,),
+        stride: Union[int, Sequence[int]] = (10,),
+        pad_left: Union[int, Sequence[int]] = (0,),
+        encoder_att_look_back_factor: Union[int, Sequence[int]] = (1,),
+        decoder_att_look_back_factor: Union[int, Sequence[int]] = (1,),
+        tf2torch_tensor_name_prefix_torch: str = "encoder",
+        tf2torch_tensor_name_prefix_tf: str = "seq2seq/encoder",
+    ):
+        super().__init__()
+        self._output_size = output_size
+
+        if input_layer == "linear":
+            self.embed = torch.nn.Sequential(
+                torch.nn.Linear(input_size, output_size),
+                torch.nn.LayerNorm(output_size),
+                torch.nn.Dropout(dropout_rate),
+                torch.nn.ReLU(),
+                pos_enc_class(output_size, positional_dropout_rate),
+            )
+        elif input_layer == "conv2d":
+            self.embed = Conv2dSubsampling(input_size, output_size, dropout_rate)
+        elif input_layer == "conv2d2":
+            self.embed = Conv2dSubsampling2(input_size, output_size, dropout_rate)
+        elif input_layer == "conv2d6":
+            self.embed = Conv2dSubsampling6(input_size, output_size, dropout_rate)
+        elif input_layer == "conv2d8":
+            self.embed = Conv2dSubsampling8(input_size, output_size, dropout_rate)
+        elif input_layer == "embed":
+            self.embed = torch.nn.Sequential(
+                torch.nn.Embedding(input_size, output_size, padding_idx=padding_idx),
+                pos_enc_class(output_size, positional_dropout_rate),
+            )
+        elif input_layer is None:
+            if input_size == output_size:
+                self.embed = None
+            else:
+                self.embed = torch.nn.Linear(input_size, output_size)
+        elif input_layer == "pe":
+            self.embed = SinusoidalPositionEncoder()
+        elif input_layer == "pe_online":
+            self.embed = StreamSinusoidalPositionEncoder()
+        else:
+            raise ValueError("unknown input_layer: " + input_layer)
+        self.normalize_before = normalize_before
+        if positionwise_layer_type == "linear":
+            positionwise_layer = PositionwiseFeedForward
+            positionwise_layer_args = (
+                output_size,
+                linear_units,
+                dropout_rate,
+            )
+        elif positionwise_layer_type == "conv1d":
+            positionwise_layer = MultiLayeredConv1d
+            positionwise_layer_args = (
+                output_size,
+                linear_units,
+                positionwise_conv_kernel_size,
+                dropout_rate,
+            )
+        elif positionwise_layer_type == "conv1d-linear":
+            positionwise_layer = Conv1dLinear
+            positionwise_layer_args = (
+                output_size,
+                linear_units,
+                positionwise_conv_kernel_size,
+                dropout_rate,
+            )
+        else:
+            raise NotImplementedError("Support only linear or conv1d.")
+
+        if selfattention_layer_type == "selfattn":
+            encoder_selfattn_layer = MultiHeadedAttention
+            encoder_selfattn_layer_args = (
+                attention_heads,
+                output_size,
+                attention_dropout_rate,
+            )
+        elif selfattention_layer_type == "sanm":
+            encoder_selfattn_layer = MultiHeadedAttentionSANM
+            encoder_selfattn_layer_args0 = (
+                attention_heads,
+                input_size,
+                output_size,
+                attention_dropout_rate,
+                kernel_size,
+                sanm_shfit,
+            )
+
+            encoder_selfattn_layer_args = (
+                attention_heads,
+                output_size,
+                output_size,
+                attention_dropout_rate,
+                kernel_size,
+                sanm_shfit,
+            )
+        self.encoders0 = repeat(
+            1,
+            lambda lnum: EncoderLayerSANM(
+                input_size,
+                output_size,
+                encoder_selfattn_layer(*encoder_selfattn_layer_args0),
+                positionwise_layer(*positionwise_layer_args),
+                dropout_rate,
+                normalize_before,
+                concat_after,
+            ),
+        )
+
+        self.encoders = repeat(
+            num_blocks - 1,
+            lambda lnum: EncoderLayerSANM(
+                output_size,
+                output_size,
+                encoder_selfattn_layer(*encoder_selfattn_layer_args),
+                positionwise_layer(*positionwise_layer_args),
+                dropout_rate,
+                normalize_before,
+                concat_after,
+            ),
+        )
+        if self.normalize_before:
+            self.after_norm = LayerNorm(output_size)
+
+        self.interctc_layer_idx = interctc_layer_idx
+        if len(interctc_layer_idx) > 0:
+            assert 0 < min(interctc_layer_idx) and max(interctc_layer_idx) < num_blocks
+        self.interctc_use_conditioning = interctc_use_conditioning
+        self.conditioning_layer = None
+        shfit_fsmn = (kernel_size - 1) // 2
+        self.overlap_chunk_cls = overlap_chunk(
+            chunk_size=chunk_size,
+            stride=stride,
+            pad_left=pad_left,
+            shfit_fsmn=shfit_fsmn,
+            encoder_att_look_back_factor=encoder_att_look_back_factor,
+            decoder_att_look_back_factor=decoder_att_look_back_factor,
+        )
+        self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch
+        self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf
+
+    def output_size(self) -> int:
+        return self._output_size
+
+    def forward(
+        self,
+        xs_pad: torch.Tensor,
+        ilens: torch.Tensor,
+        prev_states: torch.Tensor = None,
+        ctc: CTC = None,
+        ind: int = 0,
+    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        """Embed positions in tensor.
+
+        Args:
+            xs_pad: input tensor (B, L, D)
+            ilens: input length (B)
+            prev_states: Not to be used now.
+        Returns:
+            position embedded tensor and mask
+        """
+        masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device)
+        xs_pad *= self.output_size() ** 0.5
+        if self.embed is None:
+            xs_pad = xs_pad
+        elif (
+            isinstance(self.embed, Conv2dSubsampling)
+            or isinstance(self.embed, Conv2dSubsampling2)
+            or isinstance(self.embed, Conv2dSubsampling6)
+            or isinstance(self.embed, Conv2dSubsampling8)
+        ):
+            short_status, limit_size = check_short_utt(self.embed, xs_pad.size(1))
+            if short_status:
+                raise TooShortUttError(
+                    f"has {xs_pad.size(1)} frames and is too short for subsampling "
+                    + f"(it needs more than {limit_size} frames), return empty results",
+                    xs_pad.size(1),
+                    limit_size,
+                )
+            xs_pad, masks = self.embed(xs_pad, masks)
+        else:
+            xs_pad = self.embed(xs_pad)
+
+        mask_shfit_chunk, mask_att_chunk_encoder = None, None
+        if self.overlap_chunk_cls is not None:
+            ilens = masks.squeeze(1).sum(1)
+            chunk_outs = self.overlap_chunk_cls.gen_chunk_mask(ilens, ind)
+            xs_pad, ilens = self.overlap_chunk_cls.split_chunk(xs_pad, ilens, chunk_outs=chunk_outs)
+            masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device)
+            mask_shfit_chunk = self.overlap_chunk_cls.get_mask_shfit_chunk(
+                chunk_outs, xs_pad.device, xs_pad.size(0), dtype=xs_pad.dtype
+            )
+            mask_att_chunk_encoder = self.overlap_chunk_cls.get_mask_att_chunk_encoder(
+                chunk_outs, xs_pad.device, xs_pad.size(0), dtype=xs_pad.dtype
+            )
+
+        encoder_outs = self.encoders0(xs_pad, masks, None, mask_shfit_chunk, mask_att_chunk_encoder)
+        xs_pad, masks = encoder_outs[0], encoder_outs[1]
+        intermediate_outs = []
+        if len(self.interctc_layer_idx) == 0:
+            encoder_outs = self.encoders(
+                xs_pad, masks, None, mask_shfit_chunk, mask_att_chunk_encoder
+            )
+            xs_pad, masks = encoder_outs[0], encoder_outs[1]
+        else:
+            for layer_idx, encoder_layer in enumerate(self.encoders):
+                encoder_outs = encoder_layer(
+                    xs_pad, masks, None, mask_shfit_chunk, mask_att_chunk_encoder
+                )
+                xs_pad, masks = encoder_outs[0], encoder_outs[1]
+                if layer_idx + 1 in self.interctc_layer_idx:
+                    encoder_out = xs_pad
+
+                    # intermediate outputs are also normalized
+                    if self.normalize_before:
+                        encoder_out = self.after_norm(encoder_out)
+
+                    intermediate_outs.append((layer_idx + 1, encoder_out))
+
+                    if self.interctc_use_conditioning:
+                        ctc_out = ctc.softmax(encoder_out)
+                        xs_pad = xs_pad + self.conditioning_layer(ctc_out)
+
+        if self.normalize_before:
+            xs_pad = self.after_norm(xs_pad)
+
+        olens = masks.squeeze(1).sum(1)
+        if len(intermediate_outs) > 0:
+            return (xs_pad, intermediate_outs), olens, None
+        return xs_pad, olens, None
+
+    def _add_overlap_chunk(self, feats: np.ndarray, cache: dict = {}):
+        if len(cache) == 0:
+            return feats
+        cache["feats"] = to_device(cache["feats"], device=feats.device)
+        overlap_feats = torch.cat((cache["feats"], feats), dim=1)
+        cache["feats"] = overlap_feats[:, -(cache["chunk_size"][0] + cache["chunk_size"][2]) :, :]
+        return overlap_feats
+
+    def forward_chunk(
+        self,
+        xs_pad: torch.Tensor,
+        ilens: torch.Tensor,
+        cache: dict = None,
+        **kwargs,
+    ):
+        is_final = kwargs.get("is_final", False)
+        xs_pad *= self.output_size() ** 0.5
+        if self.embed is None:
+            xs_pad = xs_pad
+        else:
+            xs_pad = self.embed(xs_pad, cache)
+        if cache["tail_chunk"]:
+            xs_pad = to_device(cache["feats"], device=xs_pad.device)
+        else:
+            xs_pad = self._add_overlap_chunk(xs_pad, cache)
+        if cache["opt"] is None:
+            cache_layer_num = len(self.encoders0) + len(self.encoders)
+            new_cache = [None] * cache_layer_num
+        else:
+            new_cache = cache["opt"]
+
+        for layer_idx, encoder_layer in enumerate(self.encoders0):
+            encoder_outs = encoder_layer.forward_chunk(
+                xs_pad, new_cache[layer_idx], cache["chunk_size"], cache["encoder_chunk_look_back"]
+            )
+            xs_pad, new_cache[0] = encoder_outs[0], encoder_outs[1]
+
+        for layer_idx, encoder_layer in enumerate(self.encoders):
+            encoder_outs = encoder_layer.forward_chunk(
+                xs_pad,
+                new_cache[layer_idx + len(self.encoders0)],
+                cache["chunk_size"],
+                cache["encoder_chunk_look_back"],
+            )
+            xs_pad, new_cache[layer_idx + len(self.encoders0)] = encoder_outs[0], encoder_outs[1]
+
+        if self.normalize_before:
+            xs_pad = self.after_norm(xs_pad)
+        if cache["encoder_chunk_look_back"] > 0 or cache["encoder_chunk_look_back"] == -1:
+            cache["opt"] = new_cache
+
+        return xs_pad, ilens, None

funasr/models/scama/model.py

@@ -0,0 +1,739 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import time
+import torch
+import torch.nn as nn
+import torch.functional as F
+import logging
+from typing import Dict, Tuple
+from contextlib import contextmanager
+from distutils.version import LooseVersion
+
+from funasr.register import tables
+from funasr.models.ctc.ctc import CTC
+from funasr.utils import postprocess_utils
+from funasr.metrics.compute_acc import th_accuracy
+from funasr.utils.datadir_writer import DatadirWriter
+from funasr.models.paraformer.model import Paraformer
+from funasr.models.paraformer.search import Hypothesis
+from funasr.models.paraformer.cif_predictor import mae_loss
+from funasr.train_utils.device_funcs import force_gatherable
+from funasr.losses.label_smoothing_loss import LabelSmoothingLoss
+from funasr.models.transformer.utils.add_sos_eos import add_sos_eos
+from funasr.models.transformer.utils.nets_utils import make_pad_mask, pad_list
+from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
+from funasr.models.scama.utils import sequence_mask
+
+if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
+    from torch.cuda.amp import autocast
+else:
+    # Nothing to do if torch<1.6.0
+    @contextmanager
+    def autocast(enabled=True):
+        yield
+
+
+@tables.register("model_classes", "SCAMA")
+class SCAMA(nn.Module):
+    """
+    Author: Shiliang Zhang, Zhifu Gao, Haoneng Luo, Ming Lei, Jie Gao, Zhijie Yan, Lei Xie
+    SCAMA: Streaming chunk-aware multihead attention for online end-to-end speech recognition
+    https://arxiv.org/abs/2006.01712
+    """
+
+    def __init__(
+        self,
+        specaug: str = None,
+        specaug_conf: dict = None,
+        normalize: str = None,
+        normalize_conf: dict = None,
+        encoder: str = None,
+        encoder_conf: dict = None,
+        decoder: str = None,
+        decoder_conf: dict = None,
+        ctc: str = None,
+        ctc_conf: dict = None,
+        ctc_weight: float = 0.5,
+        predictor: str = None,
+        predictor_conf: dict = None,
+        predictor_bias: int = 0,
+        predictor_weight: float = 0.0,
+        input_size: int = 80,
+        vocab_size: int = -1,
+        ignore_id: int = -1,
+        blank_id: int = 0,
+        sos: int = 1,
+        eos: int = 2,
+        lsm_weight: float = 0.0,
+        length_normalized_loss: bool = False,
+        share_embedding: bool = False,
+        **kwargs,
+    ):
+
+        super().__init__()
+
+        if specaug is not None:
+            specaug_class = tables.specaug_classes.get(specaug)
+            specaug = specaug_class(**specaug_conf)
+
+        if normalize is not None:
+            normalize_class = tables.normalize_classes.get(normalize)
+            normalize = normalize_class(**normalize_conf)
+
+        encoder_class = tables.encoder_classes.get(encoder)
+        encoder = encoder_class(input_size=input_size, **encoder_conf)
+        encoder_output_size = encoder.output_size()
+
+        decoder_class = tables.decoder_classes.get(decoder)
+        decoder = decoder_class(
+            vocab_size=vocab_size,
+            encoder_output_size=encoder_output_size,
+            **decoder_conf,
+        )
+        if ctc_weight > 0.0:
+
+            if ctc_conf is None:
+                ctc_conf = {}
+
+            ctc = CTC(odim=vocab_size, encoder_output_size=encoder_output_size, **ctc_conf)
+
+        predictor_class = tables.predictor_classes.get(predictor)
+        predictor = predictor_class(**predictor_conf)
+
+        # note that eos is the same as sos (equivalent ID)
+        self.blank_id = blank_id
+        self.sos = sos if sos is not None else vocab_size - 1
+        self.eos = eos if eos is not None else vocab_size - 1
+        self.vocab_size = vocab_size
+        self.ignore_id = ignore_id
+        self.ctc_weight = ctc_weight
+
+        self.specaug = specaug
+        self.normalize = normalize
+
+        self.encoder = encoder
+
+        if ctc_weight == 1.0:
+            self.decoder = None
+        else:
+            self.decoder = decoder
+
+        self.criterion_att = LabelSmoothingLoss(
+            size=vocab_size,
+            padding_idx=ignore_id,
+            smoothing=lsm_weight,
+            normalize_length=length_normalized_loss,
+        )
+
+        if ctc_weight == 0.0:
+            self.ctc = None
+        else:
+            self.ctc = ctc
+
+        self.predictor = predictor
+        self.predictor_weight = predictor_weight
+        self.predictor_bias = predictor_bias
+
+        self.criterion_pre = mae_loss(normalize_length=length_normalized_loss)
+
+        self.share_embedding = share_embedding
+        if self.share_embedding:
+            self.decoder.embed = None
+
+        self.length_normalized_loss = length_normalized_loss
+        self.beam_search = None
+        self.error_calculator = None
+
+        if self.encoder.overlap_chunk_cls is not None:
+            from funasr.models.scama.chunk_utilis import (
+                build_scama_mask_for_cross_attention_decoder,
+            )
+
+            self.build_scama_mask_for_cross_attention_decoder_fn = (
+                build_scama_mask_for_cross_attention_decoder
+            )
+            self.decoder_attention_chunk_type = kwargs.get("decoder_attention_chunk_type", "chunk")
+
+    def forward(
+        self,
+        speech: torch.Tensor,
+        speech_lengths: torch.Tensor,
+        text: torch.Tensor,
+        text_lengths: torch.Tensor,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:
+        """Encoder + Decoder + Calc loss
+        Args:
+                speech: (Batch, Length, ...)
+                speech_lengths: (Batch, )
+                text: (Batch, Length)
+                text_lengths: (Batch,)
+        """
+
+        decoding_ind = kwargs.get("decoding_ind")
+        if len(text_lengths.size()) > 1:
+            text_lengths = text_lengths[:, 0]
+        if len(speech_lengths.size()) > 1:
+            speech_lengths = speech_lengths[:, 0]
+
+        batch_size = speech.shape[0]
+
+        # Encoder
+        ind = self.encoder.overlap_chunk_cls.random_choice(self.training, decoding_ind)
+        encoder_out, encoder_out_lens = self.encode(speech, speech_lengths, ind=ind)
+
+        loss_ctc, cer_ctc = None, None
+        loss_pre = None
+        stats = dict()
+
+        # decoder: CTC branch
+
+        if self.ctc_weight > 0.0:
+
+            encoder_out_ctc, encoder_out_lens_ctc = self.encoder.overlap_chunk_cls.remove_chunk(
+                encoder_out, encoder_out_lens, chunk_outs=None
+            )
+
+            loss_ctc, cer_ctc = self._calc_ctc_loss(
+                encoder_out_ctc, encoder_out_lens_ctc, text, text_lengths
+            )
+            # Collect CTC branch stats
+            stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None
+            stats["cer_ctc"] = cer_ctc
+
+        # decoder: Attention decoder branch
+        loss_att, acc_att, cer_att, wer_att, loss_pre = self._calc_att_predictor_loss(
+            encoder_out, encoder_out_lens, text, text_lengths
+        )
+
+        # 3. CTC-Att loss definition
+        if self.ctc_weight == 0.0:
+            loss = loss_att + loss_pre * self.predictor_weight
+        else:
+            loss = (
+                self.ctc_weight * loss_ctc
+                + (1 - self.ctc_weight) * loss_att
+                + loss_pre * self.predictor_weight
+            )
+
+        # Collect Attn branch stats
+        stats["loss_att"] = loss_att.detach() if loss_att is not None else None
+        stats["acc"] = acc_att
+        stats["cer"] = cer_att
+        stats["wer"] = wer_att
+        stats["loss_pre"] = loss_pre.detach().cpu() if loss_pre is not None else None
+
+        stats["loss"] = torch.clone(loss.detach())
+
+        # force_gatherable: to-device and to-tensor if scalar for DataParallel
+        if self.length_normalized_loss:
+            batch_size = (text_lengths + self.predictor_bias).sum()
+        loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device)
+        return loss, stats, weight
+
+    def encode(
+        self,
+        speech: torch.Tensor,
+        speech_lengths: torch.Tensor,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Encoder. Note that this method is used by asr_inference.py
+        Args:
+                speech: (Batch, Length, ...)
+                speech_lengths: (Batch, )
+                ind: int
+        """
+        with autocast(False):
+
+            # Data augmentation
+            if self.specaug is not None and self.training:
+                speech, speech_lengths = self.specaug(speech, speech_lengths)
+
+            # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN
+            if self.normalize is not None:
+                speech, speech_lengths = self.normalize(speech, speech_lengths)
+
+        # Forward encoder
+        encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths)
+        if isinstance(encoder_out, tuple):
+            encoder_out = encoder_out[0]
+
+        return encoder_out, encoder_out_lens
+
+    def encode_chunk(
+        self,
+        speech: torch.Tensor,
+        speech_lengths: torch.Tensor,
+        cache: dict = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Frontend + Encoder. Note that this method is used by asr_inference.py
+        Args:
+                speech: (Batch, Length, ...)
+                speech_lengths: (Batch, )
+                ind: int
+        """
+        with autocast(False):
+
+            # Data augmentation
+            if self.specaug is not None and self.training:
+                speech, speech_lengths = self.specaug(speech, speech_lengths)
+
+            # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN
+            if self.normalize is not None:
+                speech, speech_lengths = self.normalize(speech, speech_lengths)
+
+        # Forward encoder
+        encoder_out, encoder_out_lens, _ = self.encoder.forward_chunk(
+            speech, speech_lengths, cache=cache["encoder"]
+        )
+        if isinstance(encoder_out, tuple):
+            encoder_out = encoder_out[0]
+
+        return encoder_out, torch.tensor([encoder_out.size(1)])
+
+    def calc_predictor_chunk(self, encoder_out, encoder_out_lens, cache=None, **kwargs):
+        is_final = kwargs.get("is_final", False)
+
+        return self.predictor.forward_chunk(encoder_out, cache["encoder"], is_final=is_final)
+
+    def _calc_att_predictor_loss(
+        self,
+        encoder_out: torch.Tensor,
+        encoder_out_lens: torch.Tensor,
+        ys_pad: torch.Tensor,
+        ys_pad_lens: torch.Tensor,
+    ):
+        ys_in_pad, ys_out_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id)
+        ys_in_lens = ys_pad_lens + 1
+
+        encoder_out_mask = sequence_mask(
+            encoder_out_lens,
+            maxlen=encoder_out.size(1),
+            dtype=encoder_out.dtype,
+            device=encoder_out.device,
+        )[:, None, :]
+        mask_chunk_predictor = None
+        if self.encoder.overlap_chunk_cls is not None:
+            mask_chunk_predictor = self.encoder.overlap_chunk_cls.get_mask_chunk_predictor(
+                None, device=encoder_out.device, batch_size=encoder_out.size(0)
+            )
+            mask_shfit_chunk = self.encoder.overlap_chunk_cls.get_mask_shfit_chunk(
+                None, device=encoder_out.device, batch_size=encoder_out.size(0)
+            )
+            encoder_out = encoder_out * mask_shfit_chunk
+        pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor(
+            encoder_out,
+            ys_out_pad,
+            encoder_out_mask,
+            ignore_id=self.ignore_id,
+            mask_chunk_predictor=mask_chunk_predictor,
+            target_label_length=ys_in_lens,
+        )
+        predictor_alignments, predictor_alignments_len = self.predictor.gen_frame_alignments(
+            pre_alphas, encoder_out_lens
+        )
+
+        encoder_chunk_size = self.encoder.overlap_chunk_cls.chunk_size_pad_shift_cur
+        attention_chunk_center_bias = 0
+        attention_chunk_size = encoder_chunk_size
+        decoder_att_look_back_factor = (
+            self.encoder.overlap_chunk_cls.decoder_att_look_back_factor_cur
+        )
+        mask_shift_att_chunk_decoder = (
+            self.encoder.overlap_chunk_cls.get_mask_shift_att_chunk_decoder(
+                None, device=encoder_out.device, batch_size=encoder_out.size(0)
+            )
+        )
+        scama_mask = self.build_scama_mask_for_cross_attention_decoder_fn(
+            predictor_alignments=predictor_alignments,
+            encoder_sequence_length=encoder_out_lens,
+            chunk_size=1,
+            encoder_chunk_size=encoder_chunk_size,
+            attention_chunk_center_bias=attention_chunk_center_bias,
+            attention_chunk_size=attention_chunk_size,
+            attention_chunk_type=self.decoder_attention_chunk_type,
+            step=None,
+            predictor_mask_chunk_hopping=mask_chunk_predictor,
+            decoder_att_look_back_factor=decoder_att_look_back_factor,
+            mask_shift_att_chunk_decoder=mask_shift_att_chunk_decoder,
+            target_length=ys_in_lens,
+            is_training=self.training,
+        )
+
+        # try:
+        # 1. Forward decoder
+        decoder_out, _ = self.decoder(
+            encoder_out,
+            encoder_out_lens,
+            ys_in_pad,
+            ys_in_lens,
+            chunk_mask=scama_mask,
+            pre_acoustic_embeds=pre_acoustic_embeds,
+        )
+
+        # 2. Compute attention loss
+        loss_att = self.criterion_att(decoder_out, ys_out_pad)
+        acc_att = th_accuracy(
+            decoder_out.view(-1, self.vocab_size),
+            ys_out_pad,
+            ignore_label=self.ignore_id,
+        )
+        # predictor loss
+        loss_pre = self.criterion_pre(ys_in_lens.type_as(pre_token_length), pre_token_length)
+        # Compute cer/wer using attention-decoder
+        if self.training or self.error_calculator is None:
+            cer_att, wer_att = None, None
+        else:
+            ys_hat = decoder_out.argmax(dim=-1)
+            cer_att, wer_att = self.error_calculator(ys_hat.cpu(), ys_pad.cpu())
+
+        return loss_att, acc_att, cer_att, wer_att, loss_pre
+
+    def calc_predictor_mask(
+        self,
+        encoder_out: torch.Tensor,
+        encoder_out_lens: torch.Tensor,
+        ys_pad: torch.Tensor = None,
+        ys_pad_lens: torch.Tensor = None,
+    ):
+        # ys_in_pad, ys_out_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id)
+        # ys_in_lens = ys_pad_lens + 1
+        ys_out_pad, ys_in_lens = None, None
+
+        encoder_out_mask = sequence_mask(
+            encoder_out_lens,
+            maxlen=encoder_out.size(1),
+            dtype=encoder_out.dtype,
+            device=encoder_out.device,
+        )[:, None, :]
+        mask_chunk_predictor = None
+
+        mask_chunk_predictor = self.encoder.overlap_chunk_cls.get_mask_chunk_predictor(
+            None, device=encoder_out.device, batch_size=encoder_out.size(0)
+        )
+        mask_shfit_chunk = self.encoder.overlap_chunk_cls.get_mask_shfit_chunk(
+            None, device=encoder_out.device, batch_size=encoder_out.size(0)
+        )
+        encoder_out = encoder_out * mask_shfit_chunk
+        pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor(
+            encoder_out,
+            ys_out_pad,
+            encoder_out_mask,
+            ignore_id=self.ignore_id,
+            mask_chunk_predictor=mask_chunk_predictor,
+            target_label_length=ys_in_lens,
+        )
+        predictor_alignments, predictor_alignments_len = self.predictor.gen_frame_alignments(
+            pre_alphas, encoder_out_lens
+        )
+
+        encoder_chunk_size = self.encoder.overlap_chunk_cls.chunk_size_pad_shift_cur
+        attention_chunk_center_bias = 0
+        attention_chunk_size = encoder_chunk_size
+        decoder_att_look_back_factor = (
+            self.encoder.overlap_chunk_cls.decoder_att_look_back_factor_cur
+        )
+        mask_shift_att_chunk_decoder = (
+            self.encoder.overlap_chunk_cls.get_mask_shift_att_chunk_decoder(
+                None, device=encoder_out.device, batch_size=encoder_out.size(0)
+            )
+        )
+        scama_mask = self.build_scama_mask_for_cross_attention_decoder_fn(
+            predictor_alignments=predictor_alignments,
+            encoder_sequence_length=encoder_out_lens,
+            chunk_size=1,
+            encoder_chunk_size=encoder_chunk_size,
+            attention_chunk_center_bias=attention_chunk_center_bias,
+            attention_chunk_size=attention_chunk_size,
+            attention_chunk_type=self.decoder_attention_chunk_type,
+            step=None,
+            predictor_mask_chunk_hopping=mask_chunk_predictor,
+            decoder_att_look_back_factor=decoder_att_look_back_factor,
+            mask_shift_att_chunk_decoder=mask_shift_att_chunk_decoder,
+            target_length=ys_in_lens,
+            is_training=self.training,
+        )
+
+        return (
+            pre_acoustic_embeds,
+            pre_token_length,
+            predictor_alignments,
+            predictor_alignments_len,
+            scama_mask,
+        )
+
+    def init_beam_search(
+        self,
+        **kwargs,
+    ):
+
+        from funasr.models.scama.beam_search import BeamSearchScamaStreaming
+
+        from funasr.models.transformer.scorers.ctc import CTCPrefixScorer
+        from funasr.models.transformer.scorers.length_bonus import LengthBonus
+
+        # 1. Build ASR model
+        scorers = {}
+
+        if self.ctc != None:
+            ctc = CTCPrefixScorer(ctc=self.ctc, eos=self.eos)
+            scorers.update(ctc=ctc)
+        token_list = kwargs.get("token_list")
+        scorers.update(
+            decoder=self.decoder,
+            length_bonus=LengthBonus(len(token_list)),
+        )
+
+        # 3. Build ngram model
+        # ngram is not supported now
+        ngram = None
+        scorers["ngram"] = ngram
+
+        weights = dict(
+            decoder=1.0 - kwargs.get("decoding_ctc_weight", 0.0),
+            ctc=kwargs.get("decoding_ctc_weight", 0.0),
+            lm=kwargs.get("lm_weight", 0.0),
+            ngram=kwargs.get("ngram_weight", 0.0),
+            length_bonus=kwargs.get("penalty", 0.0),
+        )
+
+        beam_search = BeamSearchScamaStreaming(
+            beam_size=kwargs.get("beam_size", 2),
+            weights=weights,
+            scorers=scorers,
+            sos=self.sos,
+            eos=self.eos,
+            vocab_size=len(token_list),
+            token_list=token_list,
+            pre_beam_score_key=None if self.ctc_weight == 1.0 else "full",
+        )
+        # beam_search.to(device=kwargs.get("device", "cpu"), dtype=getattr(torch, kwargs.get("dtype", "float32"))).eval()
+        # for scorer in scorers.values():
+        #     if isinstance(scorer, torch.nn.Module):
+        #         scorer.to(device=kwargs.get("device", "cpu"), dtype=getattr(torch, kwargs.get("dtype", "float32"))).eval()
+        self.beam_search = beam_search
+
+    def generate_chunk(
+        self,
+        speech,
+        speech_lengths=None,
+        key: list = None,
+        tokenizer=None,
+        frontend=None,
+        **kwargs,
+    ):
+        cache = kwargs.get("cache", {})
+        speech = speech.to(device=kwargs["device"])
+        speech_lengths = speech_lengths.to(device=kwargs["device"])
+
+        # Encoder
+        encoder_out, encoder_out_lens = self.encode_chunk(
+            speech, speech_lengths, cache=cache, is_final=kwargs.get("is_final", False)
+        )
+        if isinstance(encoder_out, tuple):
+            encoder_out = encoder_out[0]
+        if "running_hyps" not in cache:
+            running_hyps = self.beam_search.init_hyp(encoder_out)
+            cache["running_hyps"] = running_hyps
+
+        # predictor
+        predictor_outs = self.calc_predictor_chunk(
+            encoder_out,
+            encoder_out_lens,
+            cache=cache,
+            is_final=kwargs.get("is_final", False),
+        )
+        pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = (
+            predictor_outs[0],
+            predictor_outs[1],
+            predictor_outs[2],
+            predictor_outs[3],
+        )
+        pre_token_length = pre_token_length.round().long()
+
+        if torch.max(pre_token_length) < 1:
+            return []
+        maxlen = minlen = pre_token_length
+        if kwargs.get("is_final", False):
+            maxlen += kwargs.get("token_num_relax", 5)
+            minlen = max(0, minlen - kwargs.get("token_num_relax", 5))
+        # c. Passed the encoder result and the beam search
+        nbest_hyps = self.beam_search(
+            x=encoder_out[0],
+            scama_mask=None,
+            pre_acoustic_embeds=pre_acoustic_embeds,
+            maxlen=int(maxlen),
+            minlen=int(minlen),
+            cache=cache,
+        )
+
+        cache["running_hyps"] = nbest_hyps
+        nbest_hyps = nbest_hyps[: self.nbest]
+
+        results = []
+        for hyp in nbest_hyps:
+            # assert isinstance(hyp, (Hypothesis)), type(hyp)
+
+            # remove sos/eos and get results
+            last_pos = -1
+            if isinstance(hyp.yseq, list):
+                token_int = hyp.yseq[1:last_pos]
+            else:
+                token_int = hyp.yseq[1:last_pos].tolist()
+
+                # remove blank symbol id, which is assumed to be 0
+                token_int = list(
+                    filter(
+                        lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int
+                    )
+                )
+
+                # Change integer-ids to tokens
+                token = tokenizer.ids2tokens(token_int)
+                # text = tokenizer.tokens2text(token)
+
+                result_i = token
+
+                results.extend(result_i)
+
+        return results
+
+    def init_cache(self, cache: dict = {}, **kwargs):
+        device = kwargs.get("device", "cuda")
+
+        chunk_size = kwargs.get("chunk_size", [0, 10, 5])
+        encoder_chunk_look_back = kwargs.get("encoder_chunk_look_back", 0)
+        decoder_chunk_look_back = kwargs.get("decoder_chunk_look_back", 0)
+        batch_size = 1
+
+        enc_output_size = kwargs["encoder_conf"]["output_size"]
+        feats_dims = kwargs["frontend_conf"]["n_mels"] * kwargs["frontend_conf"]["lfr_m"]
+
+        cache_encoder = {
+            "start_idx": 0,
+            "cif_hidden": torch.zeros((batch_size, 1, enc_output_size)).to(device=device),
+            "cif_alphas": torch.zeros((batch_size, 1)).to(device=device),
+            "chunk_size": chunk_size,
+            "encoder_chunk_look_back": encoder_chunk_look_back,
+            "last_chunk": False,
+            "opt": None,
+            "feats": torch.zeros((batch_size, chunk_size[0] + chunk_size[2], feats_dims)).to(
+                device=device
+            ),
+            "tail_chunk": False,
+        }
+        cache["encoder"] = cache_encoder
+
+        cache_decoder = {
+            "decode_fsmn": None,
+            "decoder_chunk_look_back": decoder_chunk_look_back,
+            "opt": None,
+            "chunk_size": chunk_size,
+        }
+        cache["decoder"] = cache_decoder
+        cache["frontend"] = {}
+
+        cache["prev_samples"] = torch.empty(0).to(device=device)
+
+        return cache
+
+    def inference(
+        self,
+        data_in,
+        data_lengths=None,
+        key: list = None,
+        tokenizer=None,
+        frontend=None,
+        cache: dict = {},
+        **kwargs,
+    ):
+
+        # init beamsearch
+        is_use_ctc = kwargs.get("decoding_ctc_weight", 0.0) > 0.00001 and self.ctc != None
+        is_use_lm = (
+            kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None
+        )
+
+        if self.beam_search is None:
+
+            logging.info("enable beam_search")
+            self.init_beam_search(**kwargs)
+            self.nbest = kwargs.get("nbest", 1)
+
+        if len(cache) == 0:
+            self.init_cache(cache, **kwargs)
+
+        meta_data = {}
+        chunk_size = kwargs.get("chunk_size", [0, 10, 5])
+        chunk_stride_samples = int(chunk_size[1] * 960)  # 600ms
+
+        time1 = time.perf_counter()
+        cfg = {"is_final": kwargs.get("is_final", False)}
+        audio_sample_list = load_audio_text_image_video(
+            data_in,
+            fs=frontend.fs,
+            audio_fs=kwargs.get("fs", 16000),
+            data_type=kwargs.get("data_type", "sound"),
+            tokenizer=tokenizer,
+            cache=cfg,
+        )
+        _is_final = cfg["is_final"]  # if data_in is a file or url, set is_final=True
+
+        time2 = time.perf_counter()
+        meta_data["load_data"] = f"{time2 - time1:0.3f}"
+        assert len(audio_sample_list) == 1, "batch_size must be set 1"
+
+        audio_sample = torch.cat((cache["prev_samples"], audio_sample_list[0]))
+
+        n = int(len(audio_sample) // chunk_stride_samples + int(_is_final))
+        m = int(len(audio_sample) % chunk_stride_samples * (1 - int(_is_final)))
+        tokens = []
+        for i in range(n):
+            kwargs["is_final"] = _is_final and i == n - 1
+            audio_sample_i = audio_sample[i * chunk_stride_samples : (i + 1) * chunk_stride_samples]
+
+            # extract fbank feats
+            speech, speech_lengths = extract_fbank(
+                [audio_sample_i],
+                data_type=kwargs.get("data_type", "sound"),
+                frontend=frontend,
+                cache=cache["frontend"],
+                is_final=kwargs["is_final"],
+            )
+            time3 = time.perf_counter()
+            meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
+            meta_data["batch_data_time"] = (
+                speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000
+            )
+
+            tokens_i = self.generate_chunk(
+                speech,
+                speech_lengths,
+                key=key,
+                tokenizer=tokenizer,
+                cache=cache,
+                frontend=frontend,
+                **kwargs,
+            )
+            tokens.extend(tokens_i)
+
+        text_postprocessed, _ = postprocess_utils.sentence_postprocess(tokens)
+
+        result_i = {"key": key[0], "text": text_postprocessed}
+        result = [result_i]
+
+        cache["prev_samples"] = audio_sample[:-m]
+        if _is_final:
+            self.init_cache(cache, **kwargs)
+
+        if kwargs.get("output_dir"):
+            writer = DatadirWriter(kwargs.get("output_dir"))
+            ibest_writer = writer[f"{1}best_recog"]
+            ibest_writer["token"][key[0]] = " ".join(tokens)
+            ibest_writer["text"][key[0]] = text_postprocessed
+
+        return result, meta_data

funasr/models/scama/template.yaml

@@ -0,0 +1,127 @@
+# This is an example that demonstrates how to configure a model file.
+# You can modify the configuration according to your own requirements.
+
+# to print the register_table:
+# from funasr.register import tables
+# tables.print()
+
+# network architecture
+model: SCAMA
+model_conf:
+    ctc_weight: 0.0
+    lsm_weight: 0.1
+    length_normalized_loss: true
+
+# encoder
+encoder: SANMEncoderChunkOpt
+encoder_conf:
+    output_size: 512
+    attention_heads: 4
+    linear_units: 2048
+    num_blocks: 50
+    dropout_rate: 0.1
+    positional_dropout_rate: 0.1
+    attention_dropout_rate: 0.1
+    input_layer: pe
+    pos_enc_class: SinusoidalPositionEncoder
+    normalize_before: true
+    kernel_size: 11
+    sanm_shfit: 0
+    selfattention_layer_type: sanm
+
+# decoder
+decoder: FsmnDecoderSCAMAOpt
+decoder_conf:
+    attention_heads: 4
+    linear_units: 2048
+    num_blocks: 16
+    dropout_rate: 0.1
+    positional_dropout_rate: 0.1
+    self_attention_dropout_rate: 0.1
+    src_attention_dropout_rate: 0.1
+    att_layer_num: 16
+    kernel_size: 11
+    sanm_shfit: 0
+
+predictor: CifPredictorV2
+predictor_conf:
+    idim: 512
+    threshold: 1.0
+    l_order: 1
+    r_order: 1
+    tail_threshold: 0.45
+
+# frontend related
+frontend: WavFrontend
+frontend_conf:
+    fs: 16000
+    window: hamming
+    n_mels: 80
+    frame_length: 25
+    frame_shift: 10
+    lfr_m: 7
+    lfr_n: 6
+
+specaug: SpecAugLFR
+specaug_conf:
+    apply_time_warp: false
+    time_warp_window: 5
+    time_warp_mode: bicubic
+    apply_freq_mask: true
+    freq_mask_width_range:
+    - 0
+    - 30
+    lfr_rate: 6
+    num_freq_mask: 1
+    apply_time_mask: true
+    time_mask_width_range:
+    - 0
+    - 12
+    num_time_mask: 1
+
+train_conf:
+  accum_grad: 1
+  grad_clip: 5
+  max_epoch: 150
+  val_scheduler_criterion:
+      - valid
+      - acc
+  best_model_criterion:
+  -   - valid
+      - acc
+      - max
+  keep_nbest_models: 10
+  avg_nbest_model: 10
+  log_interval: 50
+
+optim: adam
+optim_conf:
+   lr: 0.0005
+scheduler: warmuplr
+scheduler_conf:
+   warmup_steps: 30000
+
+dataset: AudioDataset
+dataset_conf:
+    index_ds: IndexDSJsonl
+    batch_sampler: BatchSampler
+    batch_type: example # example or length
+    batch_size: 1 # if batch_type is example, batch_size is the numbers of samples; if length, batch_size is source_token_len+target_token_len;
+    max_token_length: 2048 # filter samples if source_token_len+target_token_len > max_token_length,
+    buffer_size: 500
+    shuffle: True
+    num_workers: 0
+
+tokenizer: CharTokenizer
+tokenizer_conf:
+  unk_symbol: <unk>
+  split_with_space: true
+
+
+ctc_conf:
+    dropout_rate: 0.0
+    ctc_type: builtin
+    reduce: true
+    ignore_nan_grad: true
+
+normalize: null

funasr/models/scama/utils.py

@@ -0,0 +1,89 @@
+import os
+import yaml
+import torch
+import numpy as np
+from torch.nn import functional as F
+
+
+def sequence_mask(lengths, maxlen=None, dtype=torch.float32, device=None):
+    if maxlen is None:
+        maxlen = lengths.max()
+    row_vector = torch.arange(0, maxlen, 1).to(lengths.device)
+    matrix = torch.unsqueeze(lengths, dim=-1)
+    mask = row_vector < matrix
+    mask = mask.detach()
+
+    return mask.type(dtype).to(device) if device is not None else mask.type(dtype)
+
+
+def apply_cmvn(inputs, mvn):
+    device = inputs.device
+    dtype = inputs.dtype
+    frame, dim = inputs.shape
+    meams = np.tile(mvn[0:1, :dim], (frame, 1))
+    vars = np.tile(mvn[1:2, :dim], (frame, 1))
+    inputs -= torch.from_numpy(meams).type(dtype).to(device)
+    inputs *= torch.from_numpy(vars).type(dtype).to(device)
+
+    return inputs.type(torch.float32)
+
+
+def drop_and_add(
+    inputs: torch.Tensor,
+    outputs: torch.Tensor,
+    training: bool,
+    dropout_rate: float = 0.1,
+    stoch_layer_coeff: float = 1.0,
+):
+
+    outputs = F.dropout(outputs, p=dropout_rate, training=training, inplace=True)
+    outputs *= stoch_layer_coeff
+
+    input_dim = inputs.size(-1)
+    output_dim = outputs.size(-1)
+
+    if input_dim == output_dim:
+        outputs += inputs
+    return outputs
+
+
+def proc_tf_vocab(vocab_path):
+    with open(vocab_path, encoding="utf-8") as f:
+        token_list = [line.rstrip() for line in f]
+        if "<unk>" not in token_list:
+            token_list.append("<unk>")
+    return token_list
+
+
+def gen_config_for_tfmodel(config_path, vocab_path, output_dir):
+    token_list = proc_tf_vocab(vocab_path)
+    with open(config_path, encoding="utf-8") as f:
+        config = yaml.safe_load(f)
+
+    config["token_list"] = token_list
+
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
+    with open(os.path.join(output_dir, "config.yaml"), "w", encoding="utf-8") as f:
+        yaml_no_alias_safe_dump(config, f, indent=4, sort_keys=False)
+
+
+class NoAliasSafeDumper(yaml.SafeDumper):
+    # Disable anchor/alias in yaml because looks ugly
+    def ignore_aliases(self, data):
+        return True
+
+
+def yaml_no_alias_safe_dump(data, stream=None, **kwargs):
+    """Safe-dump in yaml with no anchor/alias"""
+    return yaml.dump(data, stream, allow_unicode=True, Dumper=NoAliasSafeDumper, **kwargs)
+
+
+if __name__ == "__main__":
+    import sys
+
+    config_path = sys.argv[1]
+    vocab_path = sys.argv[2]
+    output_dir = sys.argv[3]
+    gen_config_for_tfmodel(config_path, vocab_path, output_dir)

funasr/models/seaco_paraformer/export_meta.py

@@ -0,0 +1,202 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import torch
+
+from funasr.register import tables
+
+
+class ContextualEmbedderExport(torch.nn.Module):
+    def __init__(
+        self,
+        model,
+        max_seq_len=512,
+        feats_dim=560,
+        **kwargs,
+    ):
+        super().__init__()
+        self.embedding = model.decoder.embed  # model.bias_embed
+        model.bias_encoder.batch_first = False
+        self.bias_encoder = model.bias_encoder
+
+    def forward(self, hotword):
+        hotword = self.embedding(hotword).transpose(0, 1)  # batch second
+        hw_embed, (_, _) = self.bias_encoder(hotword)
+        return hw_embed
+
+    def export_dummy_inputs(self):
+        hotword = torch.tensor(
+            [
+                [10, 11, 12, 13, 14, 10, 11, 12, 13, 14],
+                [100, 101, 0, 0, 0, 0, 0, 0, 0, 0],
+                [1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                [10, 11, 12, 13, 14, 10, 11, 12, 13, 14],
+                [100, 101, 0, 0, 0, 0, 0, 0, 0, 0],
+                [1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+            ],
+            dtype=torch.int32,
+        )
+        # hotword_length = torch.tensor([10, 2, 1], dtype=torch.int32)
+        return hotword
+
+    def export_input_names(self):
+        return ["hotword"]
+
+    def export_output_names(self):
+        return ["hw_embed"]
+
+    def export_dynamic_axes(self):
+        return {
+            "hotword": {
+                0: "num_hotwords",
+            },
+            "hw_embed": {
+                1: "num_hotwords",
+            },
+        }
+
+    def export_name(self):
+        return "model_eb.onnx"
+
+
+def export_rebuild_model(model, **kwargs):
+    model.device = kwargs.get("device")
+    is_onnx = kwargs.get("type", "onnx") == "onnx"
+    encoder_class = tables.encoder_classes.get(kwargs["encoder"] + "Export")
+    model.encoder = encoder_class(model.encoder, onnx=is_onnx)
+
+    predictor_class = tables.predictor_classes.get(kwargs["predictor"] + "Export")
+    model.predictor = predictor_class(model.predictor, onnx=is_onnx)
+
+    # before decoder convert into export class
+    embedder_class = ContextualEmbedderExport
+    embedder_model = embedder_class(model, onnx=is_onnx)
+
+    decoder_class = tables.decoder_classes.get(kwargs["decoder"] + "Export")
+    model.decoder = decoder_class(model.decoder, onnx=is_onnx)
+
+    seaco_decoder_class = tables.decoder_classes.get(kwargs["seaco_decoder"] + "Export")
+    model.seaco_decoder = seaco_decoder_class(model.seaco_decoder, onnx=is_onnx)
+
+    from funasr.utils.torch_function import sequence_mask
+
+    model.make_pad_mask = sequence_mask(kwargs["max_seq_len"], flip=False)
+
+    from funasr.utils.torch_function import sequence_mask
+
+    model.make_pad_mask = sequence_mask(kwargs["max_seq_len"], flip=False)
+    model.feats_dim = 560
+    model.NOBIAS = 8377
+
+    import copy
+    import types
+
+    backbone_model = copy.copy(model)
+
+    # backbone
+    backbone_model.forward = types.MethodType(export_backbone_forward, backbone_model)
+    backbone_model.export_dummy_inputs = types.MethodType(
+        export_backbone_dummy_inputs, backbone_model
+    )
+    backbone_model.export_input_names = types.MethodType(
+        export_backbone_input_names, backbone_model
+    )
+    backbone_model.export_output_names = types.MethodType(
+        export_backbone_output_names, backbone_model
+    )
+    backbone_model.export_dynamic_axes = types.MethodType(
+        export_backbone_dynamic_axes, backbone_model
+    )
+    
+    embedder_model.export_name = "model_eb"
+    backbone_model.export_name = "model"
+
+    return backbone_model, embedder_model
+
+
+def export_backbone_forward(
+    self,
+    speech: torch.Tensor,
+    speech_lengths: torch.Tensor,
+    bias_embed: torch.Tensor,
+    # lmbd: float,
+):
+    # a. To device
+    batch = {"speech": speech, "speech_lengths": speech_lengths}
+
+    enc, enc_len = self.encoder(**batch)
+    mask = self.make_pad_mask(enc_len)[:, None, :]
+    pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = self.predictor(enc, mask)
+    pre_token_length = pre_token_length.floor().type(torch.int32)
+
+    decoder_out, decoder_hidden, _ = self.decoder(
+        enc, enc_len, pre_acoustic_embeds, pre_token_length, return_hidden=True, return_both=True
+    )
+    decoder_out = torch.log_softmax(decoder_out, dim=-1)
+    # seaco forward
+    B, N, D = bias_embed.shape
+    _contextual_length = torch.ones(B) * N
+
+    # ASF
+    hotword_scores = self.seaco_decoder.forward_asf6(
+        bias_embed, _contextual_length, decoder_hidden, pre_token_length
+    )
+    hotword_scores = hotword_scores[0].sum(0).sum(0)
+    # _ = self.decoder2(bias_embed, _contextual_length, decoder_hidden, pre_token_length)
+    # hotword_scores = self.decoder2.model.decoders[-1].attn_mat[0][0].sum(0).sum(0)
+    dec_filter = torch.sort(hotword_scores, descending=True)[1][:51]
+    contextual_info = bias_embed[:, dec_filter]
+    num_hot_word = contextual_info.shape[1]
+    _contextual_length = torch.Tensor([num_hot_word]).int().repeat(B).to(enc.device)
+
+    # again
+    cif_attended, _ = self.seaco_decoder(
+        contextual_info, _contextual_length, pre_acoustic_embeds, pre_token_length
+    )
+    dec_attended, _ = self.seaco_decoder(
+        contextual_info, _contextual_length, decoder_hidden, pre_token_length
+    )
+    merged = cif_attended + dec_attended
+    dha_output = self.hotword_output_layer(merged)
+    dha_pred = torch.log_softmax(dha_output, dim=-1)
+    # merging logits
+    dha_ids = dha_pred.max(-1)[-1]
+    dha_mask = (dha_ids == self.NOBIAS).int().unsqueeze(-1)
+    decoder_out = decoder_out * dha_mask + dha_pred * (1 - dha_mask)
+
+    # get predicted timestamps
+    us_alphas, us_cif_peak = self.predictor.get_upsample_timestmap(enc, mask, pre_token_length)
+    
+    return decoder_out, pre_token_length, us_alphas, us_cif_peak
+
+
+def export_backbone_dummy_inputs(self):
+    speech = torch.randn(2, 30, self.feats_dim)
+    speech_lengths = torch.tensor([15, 30], dtype=torch.int32)
+    bias_embed = torch.randn(2, 1, 512)
+    return (speech, speech_lengths, bias_embed)
+
+
+def export_backbone_input_names(self):
+    return ["speech", "speech_lengths", "bias_embed"]
+
+
+def export_backbone_output_names(self):
+    return ["logits", "token_num", "us_alphas", "us_cif_peak"]
+
+
+def export_backbone_dynamic_axes(self):
+    return {
+        "speech": {0: "batch_size", 1: "feats_length"},
+        "speech_lengths": {
+            0: "batch_size",
+        },
+        "bias_embed": {0: "batch_size", 1: "num_hotwords"},
+        "logits": {0: "batch_size", 1: "logits_length"},
+        "pre_acoustic_embeds": {1: "feats_length1"},
+        "us_alphas": {0: "batch_size", 1: "alphas_length"},
+        "us_cif_peak": {0: "batch_size", 1: "alphas_length"},
+    }
+

funasr/models/seaco_paraformer/model.py

@@ -0,0 +1,661 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import os
+import re
+import time
+import copy
+import torch
+import codecs
+import logging
+import tempfile
+import requests
+import numpy as np
+from typing import Dict, Tuple
+from contextlib import contextmanager
+from distutils.version import LooseVersion
+
+from funasr.register import tables
+from funasr.utils import postprocess_utils
+from funasr.models.paraformer.model import Paraformer
+from funasr.utils.datadir_writer import DatadirWriter
+from funasr.models.paraformer.search import Hypothesis
+from funasr.train_utils.device_funcs import force_gatherable
+from funasr.models.bicif_paraformer.model import BiCifParaformer
+from funasr.losses.label_smoothing_loss import LabelSmoothingLoss
+from funasr.models.transformer.utils.add_sos_eos import add_sos_eos
+from funasr.utils.timestamp_tools import ts_prediction_lfr6_standard
+from funasr.models.transformer.utils.nets_utils import make_pad_mask, pad_list
+from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
+
+from funasr.utils.run_bmodel import EngineOV
+from funasr.models.bicif_paraformer.cif_predictor import cif
+from funasr.models.bicif_paraformer.cif_predictor import cif_wo_hidden
+
+if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
+    from torch.cuda.amp import autocast
+else:
+    # Nothing to do if torch<1.6.0
+    @contextmanager
+    def autocast(enabled=True):
+        yield
+
+
+@tables.register("model_classes", "SeacoParaformer")
+class SeacoParaformer(BiCifParaformer, Paraformer):
+    """
+    Author: Speech Lab of DAMO Academy, Alibaba Group
+    SeACo-Paraformer: A Non-Autoregressive ASR System with Flexible and Effective Hotword Customization Ability
+    https://arxiv.org/abs/2308.03266
+    """
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+
+        self.inner_dim = kwargs.get("inner_dim", 256)
+        self.bias_encoder_type = kwargs.get("bias_encoder_type", "lstm")
+        bias_encoder_dropout_rate = kwargs.get("bias_encoder_dropout_rate", 0.0)
+        bias_encoder_bid = kwargs.get("bias_encoder_bid", False)
+        seaco_lsm_weight = kwargs.get("seaco_lsm_weight", 0.0)
+        seaco_length_normalized_loss = kwargs.get("seaco_length_normalized_loss", True)
+
+        # bias encoder
+        if self.bias_encoder_type == "lstm":
+            self.bias_encoder = torch.nn.LSTM(
+                self.inner_dim,
+                self.inner_dim,
+                2,
+                batch_first=True,
+                dropout=bias_encoder_dropout_rate,
+                bidirectional=bias_encoder_bid,
+            )
+            if bias_encoder_bid:
+                self.lstm_proj = torch.nn.Linear(self.inner_dim * 2, self.inner_dim)
+            else:
+                self.lstm_proj = None
+            # self.bias_embed = torch.nn.Embedding(self.vocab_size, self.inner_dim)
+        elif self.bias_encoder_type == "mean":
+            self.bias_embed = torch.nn.Embedding(self.vocab_size, self.inner_dim)
+        else:
+            logging.error("Unsupport bias encoder type: {}".format(self.bias_encoder_type))
+
+        # seaco decoder
+        seaco_decoder = kwargs.get("seaco_decoder", None)
+        if seaco_decoder is not None:
+            seaco_decoder_conf = kwargs.get("seaco_decoder_conf")
+            seaco_decoder_class = tables.decoder_classes.get(seaco_decoder)
+            self.seaco_decoder = seaco_decoder_class(
+                vocab_size=self.vocab_size,
+                encoder_output_size=self.inner_dim,
+                **seaco_decoder_conf,
+            )
+        self.hotword_output_layer = torch.nn.Linear(self.inner_dim, self.vocab_size)
+        self.criterion_seaco = LabelSmoothingLoss(
+            size=self.vocab_size,
+            padding_idx=self.ignore_id,
+            smoothing=seaco_lsm_weight,
+            normalize_length=seaco_length_normalized_loss,
+        )
+        self.train_decoder = kwargs.get("train_decoder", True)
+        self.seaco_weight = kwargs.get("seaco_weight", 0.01)
+        self.NO_BIAS = kwargs.get("NO_BIAS", 8377)
+        self.predictor_name = kwargs.get("predictor")
+
+        # self.encoder_bmodel = EngineOV(kwargs["model_path"]+"/encoder_fp32_10b.bmodel", device_id=kwargs['dev_id'])
+        self.encoder_bmodel = EngineOV(kwargs["model_path"]+"/encoder_fp16_10b.bmodel", device_id=kwargs['dev_id'])
+
+        # self.predictor_bmodel = EngineOV(kwargs["model_path"]+"/predictor_fp32_10b.bmodel", device_id=kwargs['dev_id'])
+        self.predictor_bmodel = EngineOV(kwargs["model_path"]+"/predictor_fp16_10b.bmodel", device_id=kwargs['dev_id'])
+
+        self.decoder_bmodel = EngineOV(kwargs["model_path"]+"/decoder_fp32_10b.bmodel", device_id=kwargs['dev_id'])
+
+    def forward(
+        self,
+        speech: torch.Tensor,
+        speech_lengths: torch.Tensor,
+        text: torch.Tensor,
+        text_lengths: torch.Tensor,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:
+        """Frontend + Encoder + Decoder + Calc loss
+
+        Args:
+                speech: (Batch, Length, ...)
+                speech_lengths: (Batch, )
+                text: (Batch, Length)
+                text_lengths: (Batch,)
+        """
+        if len(text_lengths.size()) > 1:
+            text_lengths = text_lengths[:, 0]
+        if len(speech_lengths.size()) > 1:
+            speech_lengths = speech_lengths[:, 0]
+        # Check that batch_size is unified
+        assert (
+            speech.shape[0] == speech_lengths.shape[0] == text.shape[0] == text_lengths.shape[0]
+        ), (speech.shape, speech_lengths.shape, text.shape, text_lengths.shape)
+
+        hotword_pad = kwargs.get("hotword_pad")
+        hotword_lengths = kwargs.get("hotword_lengths")
+        seaco_label_pad = kwargs.get("seaco_label_pad")
+        if len(hotword_lengths.size()) > 1:
+            hotword_lengths = hotword_lengths[:, 0]
+
+        batch_size = speech.shape[0]
+        # for data-parallel
+        text = text[:, : text_lengths.max()]
+        speech = speech[:, : speech_lengths.max()]
+
+        # 1. Encoder
+        encoder_out, encoder_out_lens = self.encode(speech, speech_lengths)
+        if self.predictor_bias == 1:
+            _, ys_pad = add_sos_eos(text, self.sos, self.eos, self.ignore_id)
+            ys_lengths = text_lengths + self.predictor_bias
+
+        stats = dict()
+        loss_seaco = self._calc_seaco_loss(
+            encoder_out,
+            encoder_out_lens,
+            ys_pad,
+            ys_lengths,
+            hotword_pad,
+            hotword_lengths,
+            seaco_label_pad,
+        )
+        if self.train_decoder:
+            loss_att, acc_att, _, _, _ = self._calc_att_loss(
+                encoder_out, encoder_out_lens, text, text_lengths
+            )
+            loss = loss_seaco + loss_att * self.seaco_weight
+            stats["loss_att"] = torch.clone(loss_att.detach())
+            stats["acc_att"] = acc_att
+        else:
+            loss = loss_seaco
+
+        stats["loss_seaco"] = torch.clone(loss_seaco.detach())
+        stats["loss"] = torch.clone(loss.detach())
+
+        # force_gatherable: to-device and to-tensor if scalar for DataParallel
+        if self.length_normalized_loss:
+            batch_size = (text_lengths + self.predictor_bias).sum()
+        loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device)
+        return loss, stats, weight
+
+    def _merge(self, cif_attended, dec_attended):
+        return cif_attended + dec_attended
+
+    def calc_predictor(self, encoder_out, encoder_out_lens):
+        encoder_out_mask = (
+            ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]
+        ).to(encoder_out.device)
+        predictor_outs = self.predictor(
+            encoder_out, None, encoder_out_mask, ignore_id=self.ignore_id
+        )
+        return predictor_outs[:4]
+
+    def _calc_seaco_loss(
+        self,
+        encoder_out: torch.Tensor,
+        encoder_out_lens: torch.Tensor,
+        ys_pad: torch.Tensor,
+        ys_lengths: torch.Tensor,
+        hotword_pad: torch.Tensor,
+        hotword_lengths: torch.Tensor,
+        seaco_label_pad: torch.Tensor,
+    ):
+        # predictor forward
+        encoder_out_mask = (
+            ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]
+        ).to(encoder_out.device)
+        pre_acoustic_embeds = self.predictor(
+            encoder_out, ys_pad, encoder_out_mask, ignore_id=self.ignore_id
+        )[0]
+        # decoder forward
+        decoder_out, _ = self.decoder(
+            encoder_out, encoder_out_lens, pre_acoustic_embeds, ys_lengths, return_hidden=True
+        )
+        selected = self._hotword_representation(hotword_pad, hotword_lengths)
+        contextual_info = (
+            selected.squeeze(0).repeat(encoder_out.shape[0], 1, 1).to(encoder_out.device)
+        )
+        num_hot_word = contextual_info.shape[1]
+        _contextual_length = (
+            torch.Tensor([num_hot_word]).int().repeat(encoder_out.shape[0]).to(encoder_out.device)
+        )
+        # dha core
+        cif_attended, _ = self.seaco_decoder(
+            contextual_info, _contextual_length, pre_acoustic_embeds, ys_lengths
+        )
+        dec_attended, _ = self.seaco_decoder(
+            contextual_info, _contextual_length, decoder_out, ys_lengths
+        )
+        merged = self._merge(cif_attended, dec_attended)
+        dha_output = self.hotword_output_layer(
+            merged[:, :-1]
+        )  # remove the last token in loss calculation
+        loss_att = self.criterion_seaco(dha_output, seaco_label_pad)
+        return loss_att
+
+    def _seaco_decode_with_ASF(
+        self,
+        encoder_out,
+        encoder_out_lens,
+        sematic_embeds,
+        ys_pad_lens,
+        hw_list,
+        nfilter=50,
+        seaco_weight=1.0,
+    ):
+        # decoder forward
+
+        """
+        decoder_out, decoder_hidden, _ = self.decoder(
+            encoder_out,
+            encoder_out_lens,
+            sematic_embeds,
+            ys_pad_lens,
+            return_hidden=True,
+            return_both=True,
+        )
+
+        decoder_pred = torch.log_softmax(decoder_out, dim=-1)
+        """
+        outputs = self.decoder_bmodel([encoder_out.detach().cpu().numpy(), encoder_out_lens.detach().cpu().numpy().astype(np.int32), sematic_embeds.detach().cpu().numpy(), ys_pad_lens.detach().cpu().numpy().astype(np.int32)])
+        decoder_pred, decoder_hidden = torch.from_numpy(outputs[0]), torch.from_numpy(outputs[1])
+        if hw_list is not None:
+            hw_lengths = [len(i) for i in hw_list]
+            hw_list_ = [torch.Tensor(i).long() for i in hw_list]
+            hw_list_pad = pad_list(hw_list_, 0).to(encoder_out.device)
+            selected = self._hotword_representation(
+                hw_list_pad, torch.Tensor(hw_lengths).int().to(encoder_out.device)
+            )
+
+            contextual_info = (
+                selected.squeeze(0).repeat(encoder_out.shape[0], 1, 1).to(encoder_out.device)
+            )
+            num_hot_word = contextual_info.shape[1]
+            _contextual_length = (
+                torch.Tensor([num_hot_word])
+                .int()
+                .repeat(encoder_out.shape[0])
+                .to(encoder_out.device)
+            )
+
+            # ASF Core
+            if nfilter > 0 and nfilter < num_hot_word:
+                hotword_scores = self.seaco_decoder.forward_asf6(
+                    contextual_info, _contextual_length, decoder_hidden, ys_pad_lens
+                )
+                hotword_scores = hotword_scores[0].sum(0).sum(0)
+                # hotword_scores /= torch.sqrt(torch.tensor(hw_lengths)[:-1].float()).to(hotword_scores.device)
+                dec_filter = torch.topk(hotword_scores, min(nfilter, num_hot_word - 1))[1].tolist()
+                add_filter = dec_filter
+                add_filter.append(len(hw_list_pad) - 1)
+                # filter hotword embedding
+                selected = selected[add_filter]
+                # again
+                contextual_info = (
+                    selected.squeeze(0).repeat(encoder_out.shape[0], 1, 1).to(encoder_out.device)
+                )
+                num_hot_word = contextual_info.shape[1]
+                _contextual_length = (
+                    torch.Tensor([num_hot_word])
+                    .int()
+                    .repeat(encoder_out.shape[0])
+                    .to(encoder_out.device)
+                )
+
+            # SeACo Core
+            cif_attended, _ = self.seaco_decoder(
+                contextual_info, _contextual_length, sematic_embeds, ys_pad_lens
+            )
+            dec_attended, _ = self.seaco_decoder(
+                contextual_info, _contextual_length, decoder_hidden, ys_pad_lens
+            )
+            merged = self._merge(cif_attended, dec_attended)
+
+            dha_output = self.hotword_output_layer(
+                merged
+            )  # remove the last token in loss calculation
+            dha_pred = torch.log_softmax(dha_output, dim=-1)
+
+            def _merge_res(dec_output, dha_output):
+                lmbd = torch.Tensor([seaco_weight] * dha_output.shape[0])
+                dha_ids = dha_output.max(-1)[-1]  # [0]
+                dha_mask = (dha_ids == self.NO_BIAS).int().unsqueeze(-1)
+                a = (1 - lmbd) / lmbd
+                b = 1 / lmbd
+                a, b = a.to(dec_output.device), b.to(dec_output.device)
+                dha_mask = (dha_mask + a.reshape(-1, 1, 1)) / b.reshape(-1, 1, 1)
+                # logits = dec_output * dha_mask + dha_output[:,:,:-1] * (1-dha_mask)
+                logits = dec_output * dha_mask + dha_output[:, :, :] * (1 - dha_mask)
+                return logits
+
+            merged_pred = _merge_res(decoder_pred, dha_pred)
+            return merged_pred
+        else:
+            return decoder_pred
+
+    def _hotword_representation(self, hotword_pad, hotword_lengths):
+        if self.bias_encoder_type != "lstm":
+            logging.error("Unsupported bias encoder type")
+
+        """
+        hw_embed = self.decoder.embed(hotword_pad)
+        hw_embed, (_, _) = self.bias_encoder(hw_embed)
+        if self.lstm_proj is not None:
+            hw_embed = self.lstm_proj(hw_embed)
+        _ind = np.arange(0, hw_embed.shape[0]).tolist()
+        selected = hw_embed[_ind, [i-1 for i in hotword_lengths.detach().cpu().tolist()]]
+        return selected
+        """
+
+        # hw_embed = self.sac_embedding(hotword_pad)
+        hw_embed = self.decoder.embed(hotword_pad)
+        hw_embed = torch.nn.utils.rnn.pack_padded_sequence(
+            hw_embed,
+            hotword_lengths.cpu().type(torch.int64),
+            batch_first=True,
+            enforce_sorted=False,
+        )
+        packed_rnn_output, _ = self.bias_encoder(hw_embed)
+        rnn_output = torch.nn.utils.rnn.pad_packed_sequence(packed_rnn_output, batch_first=True)[0]
+        if self.lstm_proj is not None:
+            hw_hidden = self.lstm_proj(rnn_output)
+        else:
+            hw_hidden = rnn_output
+        _ind = np.arange(0, hw_hidden.shape[0]).tolist()
+        selected = hw_hidden[_ind, [i - 1 for i in hotword_lengths.detach().cpu().tolist()]]
+        return selected
+
+    def inference(
+        self,
+        data_in,
+        data_lengths=None,
+        key: list = None,
+        tokenizer=None,
+        frontend=None,
+        **kwargs,
+    ):
+
+        # init beamsearch
+        is_use_ctc = kwargs.get("decoding_ctc_weight", 0.0) > 0.00001 and self.ctc != None
+        is_use_lm = (
+            kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None
+        )
+        if self.beam_search is None and (is_use_lm or is_use_ctc):
+            logging.info("enable beam_search")
+            self.init_beam_search(**kwargs)
+            self.nbest = kwargs.get("nbest", 1)
+        meta_data = {}
+
+        # extract fbank feats
+        time1 = time.perf_counter()
+        audio_sample_list = load_audio_text_image_video(
+            data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000)
+        )
+        time2 = time.perf_counter()
+        meta_data["load_data"] = f"{time2 - time1:0.3f}"
+        speech, speech_lengths = extract_fbank(
+            audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend
+        )
+        time3 = time.perf_counter()
+        meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
+        meta_data["batch_data_time"] = (
+            speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000
+        )
+
+        speech = speech.to(device=kwargs["device"])
+        speech_lengths = speech_lengths.to(device=kwargs["device"])
+
+        # hotword
+        """
+        self.hotword_list = self.generate_hotwords_list(
+            kwargs.get("hotword", None), tokenizer=tokenizer, frontend=frontend
+        )
+        """
+        self.hotword_list = None
+
+        # Encoder
+        #encoder_out, encoder_out_lens = self.encode(speech, speech_lengths)
+        outputs = self.encoder_bmodel([speech.detach().cpu().numpy(), speech_lengths.detach().cpu().numpy()])
+        encoder_out = torch.from_numpy(outputs[0])
+        encoder_out_lens = speech_lengths
+        if isinstance(encoder_out, tuple):
+            encoder_out = encoder_out[0]
+
+        """
+        # predictor
+        predictor_outs = self.calc_predictor(encoder_out, encoder_out_lens)
+        pre_acoustic_embeds, pre_token_length = predictor_outs[0], predictor_outs[1]
+        """
+        hidden, alphas, pre_token_length = torch.from_numpy(outputs[1]), torch.from_numpy(outputs[2]), torch.from_numpy(outputs[3])
+        pre_acoustic_embeds, pre_peak_index = cif(hidden, alphas, 1.0)
+        token_num_int = torch.max(pre_token_length).type(torch.int32).item()
+        pre_acoustic_embeds = pre_acoustic_embeds[:, :token_num_int, :]
+        pre_token_length = pre_token_length.round().long()
+        if torch.max(pre_token_length) < 1:
+            return ([],)
+
+        decoder_out = self._seaco_decode_with_ASF(
+            encoder_out,
+            encoder_out_lens,
+            pre_acoustic_embeds,
+            pre_token_length,
+            hw_list=self.hotword_list,
+        )
+
+        # decoder_out, _ = decoder_outs[0], decoder_outs[1]
+        if self.predictor_name == "CifPredictorV3":
+            outputs = self.predictor_bmodel([encoder_out.detach().cpu().numpy(), encoder_out_lens.detach().cpu().numpy().astype(np.int32)])
+            alphas2 = torch.from_numpy(outputs[0])
+            us_alphas = alphas2 * ((pre_token_length / torch.from_numpy(outputs[1]))[:, None].repeat(1, alphas2.size(1)))
+            us_peaks = cif_wo_hidden(us_alphas, 1.0 - 1e-4)
+            """
+            _, _, us_alphas, us_peaks = self.calc_predictor_timestamp(
+                encoder_out, encoder_out_lens, pre_token_length
+            )
+            """
+        else:
+            us_alphas = None
+
+        results = []
+        b, n, d = decoder_out.size()
+        for i in range(b):
+            x = encoder_out[i, : encoder_out_lens[i], :]
+            am_scores = decoder_out[i, : pre_token_length[i], :]
+            if self.beam_search is not None:
+                nbest_hyps = self.beam_search(
+                    x=x,
+                    am_scores=am_scores,
+                    maxlenratio=kwargs.get("maxlenratio", 0.0),
+                    minlenratio=kwargs.get("minlenratio", 0.0),
+                )
+
+                nbest_hyps = nbest_hyps[: self.nbest]
+            else:
+
+                yseq = am_scores.argmax(dim=-1)
+                score = am_scores.max(dim=-1)[0]
+                score = torch.sum(score, dim=-1)
+                # pad with mask tokens to ensure compatibility with sos/eos tokens
+                yseq = torch.tensor([self.sos] + yseq.tolist() + [self.eos], device=yseq.device)
+                nbest_hyps = [Hypothesis(yseq=yseq, score=score)]
+            for nbest_idx, hyp in enumerate(nbest_hyps):
+                ibest_writer = None
+                if kwargs.get("output_dir") is not None:
+                    if not hasattr(self, "writer"):
+                        self.writer = DatadirWriter(kwargs.get("output_dir"))
+                    ibest_writer = self.writer[f"{nbest_idx + 1}best_recog"]
+
+                # remove sos/eos and get results
+                last_pos = -1
+                if isinstance(hyp.yseq, list):
+                    token_int = hyp.yseq[1:last_pos]
+                else:
+                    token_int = hyp.yseq[1:last_pos].tolist()
+
+                # remove blank symbol id, which is assumed to be 0
+                token_int = list(
+                    filter(
+                        lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int
+                    )
+                )
+
+                if tokenizer is not None:
+                    # Change integer-ids to tokens
+                    token = tokenizer.ids2tokens(token_int)
+                    text = tokenizer.tokens2text(token)
+                    if us_alphas is not None:
+                        _, timestamp = ts_prediction_lfr6_standard(
+                            us_alphas[i][: encoder_out_lens[i] * 3],
+                            us_peaks[i][: encoder_out_lens[i] * 3],
+                            copy.copy(token),
+                            vad_offset=kwargs.get("begin_time", 0),
+                        )
+                        text_postprocessed, time_stamp_postprocessed, _ = (
+                            postprocess_utils.sentence_postprocess(token, timestamp)
+                        )
+                        result_i = {
+                            "key": key[i],
+                            "text": text_postprocessed,
+                            "timestamp": time_stamp_postprocessed,
+                        }
+                        if ibest_writer is not None:
+                            ibest_writer["token"][key[i]] = " ".join(token)
+                            ibest_writer["timestamp"][key[i]] = time_stamp_postprocessed
+                            ibest_writer["text"][key[i]] = text_postprocessed
+                    else:
+                        text_postprocessed, _ = postprocess_utils.sentence_postprocess(token)
+                        result_i = {"key": key[i], "text": text_postprocessed}
+                        if ibest_writer is not None:
+                            ibest_writer["token"][key[i]] = " ".join(token)
+                            ibest_writer["text"][key[i]] = text_postprocessed
+                else:
+                    result_i = {"key": key[i], "token_int": token_int}
+                results.append(result_i)
+
+        return results, meta_data
+
+    def generate_hotwords_list(self, hotword_list_or_file, tokenizer=None, frontend=None):
+        def load_seg_dict(seg_dict_file):
+            seg_dict = {}
+            assert isinstance(seg_dict_file, str)
+            with open(seg_dict_file, "r", encoding="utf8") as f:
+                lines = f.readlines()
+                for line in lines:
+                    s = line.strip().split()
+                    key = s[0]
+                    value = s[1:]
+                    seg_dict[key] = " ".join(value)
+            return seg_dict
+
+        def seg_tokenize(txt, seg_dict):
+            pattern = re.compile(r"^[\u4E00-\u9FA50-9]+$")
+            out_txt = ""
+            for word in txt:
+                word = word.lower()
+                if word in seg_dict:
+                    out_txt += seg_dict[word] + " "
+                else:
+                    if pattern.match(word):
+                        for char in word:
+                            if char in seg_dict:
+                                out_txt += seg_dict[char] + " "
+                            else:
+                                out_txt += "<unk>" + " "
+                    else:
+                        out_txt += "<unk>" + " "
+            return out_txt.strip().split()
+
+        seg_dict = None
+        if frontend.cmvn_file is not None:
+            model_dir = os.path.dirname(frontend.cmvn_file)
+            seg_dict_file = os.path.join(model_dir, "seg_dict")
+            if os.path.exists(seg_dict_file):
+                seg_dict = load_seg_dict(seg_dict_file)
+            else:
+                seg_dict = None
+        # for None
+        if hotword_list_or_file is None:
+            hotword_list = None
+        # for local txt inputs
+        elif os.path.exists(hotword_list_or_file) and hotword_list_or_file.endswith(".txt"):
+            logging.info("Attempting to parse hotwords from local txt...")
+            hotword_list = []
+            hotword_str_list = []
+            with codecs.open(hotword_list_or_file, "r") as fin:
+                for line in fin.readlines():
+                    hw = line.strip()
+                    hw_list = hw.split()
+                    if seg_dict is not None:
+                        hw_list = seg_tokenize(hw_list, seg_dict)
+                    hotword_str_list.append(hw)
+                    hotword_list.append(tokenizer.tokens2ids(hw_list))
+                hotword_list.append([self.sos])
+                hotword_str_list.append("<s>")
+            logging.info(
+                "Initialized hotword list from file: {}, hotword list: {}.".format(
+                    hotword_list_or_file, hotword_str_list
+                )
+            )
+        # for url, download and generate txt
+        elif hotword_list_or_file.startswith("http"):
+            logging.info("Attempting to parse hotwords from url...")
+            work_dir = tempfile.TemporaryDirectory().name
+            if not os.path.exists(work_dir):
+                os.makedirs(work_dir)
+            text_file_path = os.path.join(work_dir, os.path.basename(hotword_list_or_file))
+            local_file = requests.get(hotword_list_or_file)
+            open(text_file_path, "wb").write(local_file.content)
+            hotword_list_or_file = text_file_path
+            hotword_list = []
+            hotword_str_list = []
+            with codecs.open(hotword_list_or_file, "r") as fin:
+                for line in fin.readlines():
+                    hw = line.strip()
+                    hw_list = hw.split()
+                    if seg_dict is not None:
+                        hw_list = seg_tokenize(hw_list, seg_dict)
+                    hotword_str_list.append(hw)
+                    hotword_list.append(tokenizer.tokens2ids(hw_list))
+                hotword_list.append([self.sos])
+                hotword_str_list.append("<s>")
+            logging.info(
+                "Initialized hotword list from file: {}, hotword list: {}.".format(
+                    hotword_list_or_file, hotword_str_list
+                )
+            )
+        # for text str input
+        elif not hotword_list_or_file.endswith(".txt"):
+            logging.info("Attempting to parse hotwords as str...")
+            hotword_list = []
+            hotword_str_list = []
+            for hw in hotword_list_or_file.strip().split():
+                hotword_str_list.append(hw)
+                hw_list = hw.strip().split()
+                if seg_dict is not None:
+                    hw_list = seg_tokenize(hw_list, seg_dict)
+                hotword_list.append(tokenizer.tokens2ids(hw_list))
+            hotword_list.append([self.sos])
+            hotword_str_list.append("<s>")
+            logging.info("Hotword list: {}.".format(hotword_str_list))
+        else:
+            hotword_list = None
+        return hotword_list
+
+    def export(
+        self,
+        **kwargs,
+    ):
+        if "max_seq_len" not in kwargs:
+            kwargs["max_seq_len"] = 512
+        from .export_meta import export_rebuild_model
+
+        models = export_rebuild_model(model=self, **kwargs)
+        return models

funasr/models/seaco_paraformer/template.yaml

@@ -0,0 +1,156 @@
+# This is an example that demonstrates how to configure a model file.
+# You can modify the configuration according to your own requirements.
+
+# to print the register_table:
+# from funasr.register import tables
+# tables.print()
+
+# network architecture
+model: SeacoParaformer
+model_conf:
+    ctc_weight: 0.0
+    lsm_weight: 0.1
+    length_normalized_loss: true
+    predictor_weight: 1.0
+    predictor_bias: 1
+    sampling_ratio: 0.75
+    inner_dim: 512
+    bias_encoder_type: lstm
+    bias_encoder_bid: false
+    seaco_lsm_weight: 0.1
+    seaco_length_normal: true
+    train_decoder: false
+    NO_BIAS: 8377
+
+# encoder
+encoder: SANMEncoder
+encoder_conf:
+    output_size: 512
+    attention_heads: 4
+    linear_units: 2048
+    num_blocks: 50
+    dropout_rate: 0.1
+    positional_dropout_rate: 0.1
+    attention_dropout_rate: 0.1
+    input_layer: pe
+    pos_enc_class: SinusoidalPositionEncoder
+    normalize_before: true
+    kernel_size: 11
+    sanm_shfit: 0
+    selfattention_layer_type: sanm
+
+# decoder
+decoder: ParaformerSANMDecoder
+decoder_conf:
+    attention_heads: 4
+    linear_units: 2048
+    num_blocks: 16
+    dropout_rate: 0.1
+    positional_dropout_rate: 0.1
+    self_attention_dropout_rate: 0.1
+    src_attention_dropout_rate: 0.1
+    att_layer_num: 16
+    kernel_size: 11
+    sanm_shfit: 0
+
+# seaco decoder
+seaco_decoder: ParaformerSANMDecoder
+seaco_decoder_conf:
+    attention_heads: 4
+    linear_units: 1024
+    num_blocks: 4
+    dropout_rate: 0.1
+    positional_dropout_rate: 0.1
+    self_attention_dropout_rate: 0.1
+    src_attention_dropout_rate: 0.1
+    kernel_size: 21
+    sanm_shfit: 0
+    use_output_layer: false
+    wo_input_layer: true
+
+predictor: CifPredictorV3
+predictor_conf:
+    idim: 512
+    threshold: 1.0
+    l_order: 1
+    r_order: 1
+    tail_threshold: 0.45
+    smooth_factor2: 0.25
+    noise_threshold2: 0.01
+    upsample_times: 3
+    use_cif1_cnn: false
+    upsample_type: cnn_blstm
+
+# frontend related
+frontend: WavFrontend
+frontend_conf:
+    fs: 16000
+    window: hamming
+    n_mels: 80
+    frame_length: 25
+    frame_shift: 10
+    lfr_m: 7
+    lfr_n: 6
+    dither: 0.0
+
+specaug: SpecAugLFR
+specaug_conf:
+    apply_time_warp: false
+    time_warp_window: 5
+    time_warp_mode: bicubic
+    apply_freq_mask: true
+    freq_mask_width_range:
+    - 0
+    - 30
+    lfr_rate: 6
+    num_freq_mask: 1
+    apply_time_mask: true
+    time_mask_width_range:
+    - 0
+    - 12
+    num_time_mask: 1
+
+train_conf:
+  accum_grad: 1
+  grad_clip: 5
+  max_epoch: 150
+  val_scheduler_criterion:
+      - valid
+      - acc
+  best_model_criterion:
+  -   - valid
+      - acc
+      - max
+  keep_nbest_models: 10
+  log_interval: 50
+
+optim: adam
+optim_conf:
+   lr: 0.0005
+scheduler: warmuplr
+scheduler_conf:
+   warmup_steps: 30000
+
+dataset: AudioDataset
+dataset_conf:
+    index_ds: IndexDSJsonl
+    batch_sampler: BatchSampler
+    batch_type: example # example or length
+    batch_size: 1 # if batch_type is example, batch_size is the numbers of samples; if length, batch_size is source_token_len+target_token_len;
+    max_token_length: 2048 # filter samples if source_token_len+target_token_len > max_token_length,
+    buffer_size: 500
+    shuffle: True
+    num_workers: 0
+
+tokenizer: CharTokenizer
+tokenizer_conf:
+  unk_symbol: <unk>
+  split_with_space: true
+
+
+ctc_conf:
+    dropout_rate: 0.0
+    ctc_type: builtin
+    reduce: true
+    ignore_nan_grad: true
+normalize: null