Initial commit: FunASR Speech Recognition Toolkit

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

funasr/models/scama/chunk_utilis.py

@@ -0,0 +1,624 @@
+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,),
+    ):
+
+        """Initialize overlap_chunk.
+        
+            Args:
+                chunk_size: Size/dimension parameter.
+                stride: TODO.
+                pad_left: TODO.
+                encoder_att_look_back_factor: TODO.
+                shfit_fsmn: TODO.
+                decoder_att_look_back_factor: TODO.
+            """
+        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):
+        """Check chunk size args.
+        
+            Args:
+                chunk_size: Size/dimension parameter.
+                x: TODO.
+            """
+        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:
+        """Get chunk size.
+        
+            Args:
+                ind: TODO.
+            """
+        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):
+        """Random choice.
+        
+            Args:
+                training: TODO.
+                decoding_ind: TODO.
+            """
+        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):
+
+        """Gen chunk mask.
+        
+            Args:
+                x_len: TODO.
+                ind: TODO.
+                num_units: TODO.
+                num_units_predictor: TODO.
+            """
+        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):
+        """Remove chunk.
+        
+            Args:
+                x_chunk: TODO.
+                x_len_chunk: TODO.
+                chunk_outs: TODO.
+            """
+        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):
+        """Get x add mask.
+        
+            Args:
+                chunk_outs: TODO.
+                device: Target device ("cuda:0", "cpu", etc.).
+                idx: TODO.
+                dtype: TODO.
+            """
+        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):
+        """Get x len chunk.
+        
+            Args:
+                chunk_outs: TODO.
+                device: Target device ("cuda:0", "cpu", etc.).
+                idx: TODO.
+                dtype: TODO.
+            """
+        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):
+        """Get x rm mask.
+        
+            Args:
+                chunk_outs: TODO.
+                device: Target device ("cuda:0", "cpu", etc.).
+                idx: TODO.
+                dtype: TODO.
+            """
+        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):
+        """Get x len.
+        
+            Args:
+                chunk_outs: TODO.
+                device: Target device ("cuda:0", "cpu", etc.).
+                idx: TODO.
+                dtype: TODO.
+            """
+        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
+    ):
+        """Get mask shfit chunk.
+        
+            Args:
+                chunk_outs: TODO.
+                device: Target device ("cuda:0", "cpu", etc.).
+                batch_size: Number of samples per batch.
+                num_units: TODO.
+                idx: TODO.
+                dtype: TODO.
+            """
+        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
+    ):
+        """Get mask chunk predictor.
+        
+            Args:
+                chunk_outs: TODO.
+                device: Target device ("cuda:0", "cpu", etc.).
+                batch_size: Number of samples per batch.
+                num_units: TODO.
+                idx: TODO.
+                dtype: TODO.
+            """
+        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
+    ):
+        """Get mask att chunk encoder.
+        
+            Args:
+                chunk_outs: TODO.
+                device: Target device ("cuda:0", "cpu", etc.).
+                batch_size: Number of samples per batch.
+                idx: TODO.
+                dtype: TODO.
+            """
+        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
+    ):
+        """Get mask shift att chunk decoder.
+        
+            Args:
+                chunk_outs: TODO.
+                device: Target device ("cuda:0", "cpu", etc.).
+                batch_size: Number of samples per batch.
+                idx: TODO.
+                dtype: TODO.
+            """
+        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,
+):
+    """Build scama mask for cross attention decoder.
+    
+        Args:
+            predictor_alignments: TODO.
+            encoder_sequence_length: TODO.
+            chunk_size: Size/dimension parameter.
+            encoder_chunk_size: Size/dimension parameter.
+            attention_chunk_center_bias: TODO.
+            attention_chunk_size: Size/dimension parameter.
+            attention_chunk_type: TODO.
+            step: TODO.
+            predictor_mask_chunk_hopping: TODO.
+            decoder_att_look_back_factor: TODO.
+            mask_shift_att_chunk_decoder: TODO.
+            target_length: TODO.
+            is_training: Boolean flag for training.
+            dtype: TODO.
+        """
+    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():
+                """Internal: 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,524 @@
+#!/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,
+    ):
+        """Initialize FsmnDecoderSCAMAOpt.
+        
+            Args:
+                vocab_size: Size/dimension parameter.
+                encoder_output_size: Size/dimension parameter.
+                attention_heads: TODO.
+                linear_units: TODO.
+                num_blocks: TODO.
+                dropout_rate: TODO.
+                positional_dropout_rate: TODO.
+                self_attention_dropout_rate: TODO.
+                src_attention_dropout_rate: TODO.
+                input_layer: TODO.
+                use_output_layer: TODO.
+                pos_enc_class: TODO.
+                normalize_before: TODO.
+                concat_after: TODO.
+                att_layer_num: TODO.
+                kernel_size: Size/dimension parameter.
+                sanm_shfit: TODO.
+                concat_embeds: TODO.
+                attention_dim: Size/dimension parameter.
+                tf2torch_tensor_name_prefix_torch: TODO.
+                tf2torch_tensor_name_prefix_tf: TODO.
+                embed_tensor_name_prefix_tf: TODO.
+            """
+        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,549 @@
+#!/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",
+    ):
+        """Initialize SANMEncoderChunkOpt.
+        
+            Args:
+                input_size: Size/dimension parameter.
+                output_size: Size/dimension parameter.
+                attention_heads: TODO.
+                linear_units: TODO.
+                num_blocks: TODO.
+                dropout_rate: TODO.
+                positional_dropout_rate: TODO.
+                attention_dropout_rate: TODO.
+                input_layer: TODO.
+                pos_enc_class: TODO.
+                normalize_before: TODO.
+                concat_after: TODO.
+                positionwise_layer_type: TODO.
+                positionwise_conv_kernel_size: Size/dimension parameter.
+                padding_idx: TODO.
+                interctc_layer_idx: TODO.
+                interctc_use_conditioning: TODO.
+                kernel_size: Size/dimension parameter.
+                sanm_shfit: TODO.
+                selfattention_layer_type: TODO.
+                chunk_size: Size/dimension parameter.
+                stride: TODO.
+                pad_left: TODO.
+                encoder_att_look_back_factor: TODO.
+                decoder_att_look_back_factor: TODO.
+                tf2torch_tensor_name_prefix_torch: TODO.
+                tf2torch_tensor_name_prefix_tf: TODO.
+            """
+        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:
+        """Output size."""
+        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 = None):
+        """Internal: add overlap chunk.
+        
+            Args:
+                feats: Feature tensor (e.g., fbank), shape (batch, frames, dim).
+                cache: State cache dict for streaming inference.
+            """
+        if cache is None:
+            cache = {}
+        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,
+    ):
+        """Forward chunk.
+        
+            Args:
+                xs_pad: TODO.
+                ilens: TODO.
+                cache: State cache dict for streaming inference.
+                **kwargs: Additional keyword arguments.
+            """
+        if cache is None:
+            cache = {}
+        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,849 @@
+#!/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):
+        """Autocast.
+        
+            Args:
+                enabled: TODO.
+            """
+        yield
+
+
+@tables.register("model_classes", "SCAMA")
+class SCAMA(nn.Module):
+    """SCAMA: Streaming Chunk-Aware Multi-head Attention ASR.
+
+    Streaming ASR using chunk-based encoder with configurable latency.
+    Supports 2-pass decoding for accuracy refinement.
+
+    Output: {"key": str, "text": str}
+
+    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,
+    ):
+
+        """Initialize SCAMA.
+        
+            Args:
+                specaug: TODO.
+                specaug_conf: Configuration dict for specaug.
+                normalize: TODO.
+                normalize_conf: Configuration dict for normalize.
+                encoder: TODO.
+                encoder_conf: Configuration dict for encoder.
+                decoder: TODO.
+                decoder_conf: Configuration dict for decoder.
+                ctc: TODO.
+                ctc_conf: Configuration dict for ctc.
+                ctc_weight: TODO.
+                predictor: TODO.
+                predictor_conf: Configuration dict for predictor.
+                predictor_bias: TODO.
+                predictor_weight: TODO.
+                input_size: Size/dimension parameter.
+                vocab_size: Size/dimension parameter.
+                ignore_id: TODO.
+                blank_id: TODO.
+                sos: TODO.
+                eos: TODO.
+                lsm_weight: TODO.
+                length_normalized_loss: TODO.
+                share_embedding: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        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]:
+        """Encode chunk.
+        
+            Args:
+                speech: Speech audio tensor, shape (batch, time).
+                speech_lengths: Length of each speech sample.
+                cache: State cache dict for streaming inference.
+                **kwargs: Additional keyword arguments.
+            """
+        if cache is None:
+            cache = {}
+        """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):
+        """Calc predictor chunk.
+        
+            Args:
+                encoder_out: Encoder output tensor.
+                encoder_out_lens: Encoder output lengths.
+                cache: State cache dict for streaming inference.
+                **kwargs: Additional keyword arguments.
+            """
+        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,
+    ):
+        """Internal: calc att predictor loss.
+        
+            Args:
+                encoder_out: Encoder output tensor.
+                encoder_out_lens: Encoder output lengths.
+                ys_pad: TODO.
+                ys_pad_lens: Lengths of ys_pad.
+            """
+        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
+        """Calc predictor mask.
+        
+            Args:
+                encoder_out: Encoder output tensor.
+                encoder_out_lens: Encoder output lengths.
+                ys_pad: TODO.
+                ys_pad_lens: Lengths of ys_pad.
+            """
+        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,
+    ):
+
+        """Init beam search.
+        
+            Args:
+                **kwargs: Additional keyword arguments.
+            """
+        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,
+    ):
+        """Generate chunk.
+        
+            Args:
+                speech: Speech audio tensor, shape (batch, time).
+                speech_lengths: Length of each speech sample.
+                key: Sample identifiers.
+                tokenizer: Tokenizer instance for text encoding/decoding.
+                frontend: Audio frontend for feature extraction.
+                **kwargs: Additional keyword arguments.
+            """
+        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 = None, **kwargs):
+        """Init cache.
+        
+            Args:
+                cache: State cache dict for streaming inference.
+                **kwargs: Additional keyword arguments.
+            """
+        if cache is None:
+            cache = {}
+        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)
+
+        return cache
+
+    def inference(
+        self,
+        data_in,
+        data_lengths=None,
+        key: list = None,
+        tokenizer=None,
+        frontend=None,
+        cache: dict = None,
+        **kwargs,
+    ):
+        """Run inference on input data.
+        
+            Args:
+                data_in: Input data (audio samples, file paths, or text).
+                data_lengths: Lengths of each input sample in the batch.
+                key: Sample identifiers.
+                tokenizer: Tokenizer instance for text encoding/decoding.
+                frontend: Audio frontend for feature extraction.
+                cache: State cache dict for streaming inference.
+                **kwargs: Additional keyword arguments.
+            """
+        if cache is None:
+            cache = {}
+
+        # 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 m > 0 else torch.empty(0)
+        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,129 @@
+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):
+    """Sequence mask.
+    
+        Args:
+            lengths: TODO.
+            maxlen: TODO.
+            dtype: TODO.
+            device: Target device ("cuda:0", "cpu", etc.).
+        """
+    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):
+    """Apply cmvn.
+    
+        Args:
+            inputs: TODO.
+            mvn: TODO.
+        """
+    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,
+):
+
+    """Drop and add.
+    
+        Args:
+            inputs: TODO.
+            outputs: TODO.
+            training: TODO.
+            dropout_rate: TODO.
+            stoch_layer_coeff: TODO.
+        """
+    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):
+    """Proc tf vocab.
+    
+        Args:
+            vocab_path: TODO.
+        """
+    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):
+    """Gen config for tfmodel.
+    
+        Args:
+            config_path: TODO.
+            vocab_path: TODO.
+            output_dir: Directory for saving output files.
+        """
+    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):
+        """Ignore aliases.
+        
+            Args:
+                data: TODO.
+            """
+        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)