model.py

import time
import torch
from torch import nn
import torch.nn.functional as F
from typing import Iterable, Optional, Sequence, Tuple, Union

from funasr.register import tables
from funasr.models.ctc.ctc import CTC
from funasr.utils.datadir_writer import DatadirWriter
from funasr.models.paraformer.search import Hypothesis
from funasr.models.scama.chunk_utilis import overlap_chunk
from funasr.models.transformer.embedding import StreamSinusoidalPositionEncoder
from funasr.train_utils.device_funcs import force_gatherable
from funasr.losses.label_smoothing_loss import LabelSmoothingLoss
from funasr.metrics.compute_acc import compute_accuracy, th_accuracy
from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
from utils.ctc_alignment import ctc_forced_align

class SinusoidalPositionEncoder(torch.nn.Module):
    """ """

    def __init__(self, d_model=80, dropout_rate=0.1):
        super().__init__()

    def encode(
        self, positions: torch.Tensor = None, depth: int = None, dtype: torch.dtype = torch.float32
    ):
        batch_size = positions.size(0)
        positions = positions.type(dtype)
        device = positions.device
        log_timescale_increment = torch.log(torch.tensor([10000], dtype=dtype, device=device)) / (
            depth / 2 - 1
        )
        inv_timescales = torch.exp(
            torch.arange(depth / 2, device=device).type(dtype) * (-log_timescale_increment)
        )
        inv_timescales = torch.reshape(inv_timescales, [batch_size, -1])
        scaled_time = torch.reshape(positions, [1, -1, 1]) * torch.reshape(
            inv_timescales, [1, 1, -1]
        )
        encoding = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=2)
        return encoding.type(dtype)

    def forward(self, x):
        batch_size, timesteps, input_dim = x.size()
        positions = torch.arange(1, timesteps + 1, device=x.device)[None, :]
        position_encoding = self.encode(positions, input_dim, x.dtype).to(x.device)

        return x + position_encoding


class PositionwiseFeedForward(torch.nn.Module):
    """Positionwise feed forward layer.

    Args:
        idim (int): Input dimenstion.
        hidden_units (int): The number of hidden units.
        dropout_rate (float): Dropout rate.

    """

    def __init__(self, idim, hidden_units, dropout_rate, activation=torch.nn.ReLU()):
        """Construct an PositionwiseFeedForward object."""
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = torch.nn.Linear(idim, hidden_units)
        self.w_2 = torch.nn.Linear(hidden_units, idim)
        self.dropout = torch.nn.Dropout(dropout_rate)
        self.activation = activation

    def forward(self, x):
        """Forward function."""
        return self.w_2(self.dropout(self.activation(self.w_1(x))))


class MultiHeadedAttentionSANM(nn.Module):
    """Multi-Head Attention layer.

    Args:
        n_head (int): The number of heads.
        n_feat (int): The number of features.
        dropout_rate (float): Dropout rate.

    """

    def __init__(
        self,
        n_head,
        in_feat,
        n_feat,
        dropout_rate,
        kernel_size,
        sanm_shfit=0,
        lora_list=None,
        lora_rank=8,
        lora_alpha=16,
        lora_dropout=0.1,
    ):
        """Construct an MultiHeadedAttention object."""
        super().__init__()
        assert n_feat % n_head == 0
        # We assume d_v always equals d_k
        self.d_k = n_feat // n_head
        self.h = n_head
        # self.linear_q = nn.Linear(n_feat, n_feat)
        # self.linear_k = nn.Linear(n_feat, n_feat)
        # self.linear_v = nn.Linear(n_feat, n_feat)

        self.linear_out = nn.Linear(n_feat, n_feat)
        self.linear_q_k_v = nn.Linear(in_feat, n_feat * 3)
        self.attn = None
        self.dropout = nn.Dropout(p=dropout_rate)

        self.fsmn_block = nn.Conv1d(
            n_feat, n_feat, kernel_size, stride=1, padding=0, groups=n_feat, bias=False
        )
        # padding
        left_padding = (kernel_size - 1) // 2
        if sanm_shfit > 0:
            left_padding = left_padding + sanm_shfit
        right_padding = kernel_size - 1 - left_padding
        self.pad_fn = nn.ConstantPad1d((left_padding, right_padding), 0.0)

    def forward_fsmn(self, inputs, mask, mask_shfit_chunk=None):
        b, t, d = inputs.size()
        if mask is not None:
            mask = torch.reshape(mask, (b, -1, 1))
            if mask_shfit_chunk is not None:
                mask = mask * mask_shfit_chunk
            inputs = inputs * mask

        x = inputs.transpose(1, 2)
        x = self.pad_fn(x)
        x = self.fsmn_block(x)
        x = x.transpose(1, 2)
        x += inputs
        x = self.dropout(x)
        if mask is not None:
            x = x * mask
        return x

    def forward_qkv(self, x):
        """Transform query, key and value.

        Args:
            query (torch.Tensor): Query tensor (#batch, time1, size).
            key (torch.Tensor): Key tensor (#batch, time2, size).
            value (torch.Tensor): Value tensor (#batch, time2, size).

        Returns:
            torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k).
            torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k).
            torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k).

        """
        b, t, d = x.size()
        q_k_v = self.linear_q_k_v(x)
        q, k, v = torch.split(q_k_v, int(self.h * self.d_k), dim=-1)
        q_h = torch.reshape(q, (b, t, self.h, self.d_k)).transpose(
            1, 2
        )  # (batch, head, time1, d_k)
        k_h = torch.reshape(k, (b, t, self.h, self.d_k)).transpose(
            1, 2
        )  # (batch, head, time2, d_k)
        v_h = torch.reshape(v, (b, t, self.h, self.d_k)).transpose(
            1, 2
        )  # (batch, head, time2, d_k)

        return q_h, k_h, v_h, v

    def forward_attention(self, value, scores, mask, mask_att_chunk_encoder=None):
        """Compute attention context vector.

        Args:
            value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k).
            scores (torch.Tensor): Attention score (#batch, n_head, time1, time2).
            mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2).

        Returns:
            torch.Tensor: Transformed value (#batch, time1, d_model)
                weighted by the attention score (#batch, time1, time2).

        """
        n_batch = value.size(0)
        if mask is not None:
            if mask_att_chunk_encoder is not None:
                mask = mask * mask_att_chunk_encoder

            mask = mask.unsqueeze(1).eq(0)  # (batch, 1, *, time2)

            min_value = -float(
                "inf"
            )  # float(numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min)
            scores = scores.masked_fill(mask, min_value)
            attn = torch.softmax(scores, dim=-1).masked_fill(
                mask, 0.0
            )  # (batch, head, time1, time2)
        else:
            attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)

        p_attn = self.dropout(attn)
        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
        x = (
            x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)
        )  # (batch, time1, d_model)

        return self.linear_out(x)  # (batch, time1, d_model)

    def forward(self, x, mask, mask_shfit_chunk=None, mask_att_chunk_encoder=None):
        """Compute scaled dot product attention.

        Args:
            query (torch.Tensor): Query tensor (#batch, time1, size).
            key (torch.Tensor): Key tensor (#batch, time2, size).
            value (torch.Tensor): Value tensor (#batch, time2, size).
            mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
                (#batch, time1, time2).

        Returns:
            torch.Tensor: Output tensor (#batch, time1, d_model).

        """
        q_h, k_h, v_h, v = self.forward_qkv(x)
        fsmn_memory = self.forward_fsmn(v, mask, mask_shfit_chunk)
        q_h = q_h * self.d_k ** (-0.5)
        scores = torch.matmul(q_h, k_h.transpose(-2, -1))
        att_outs = self.forward_attention(v_h, scores, mask, mask_att_chunk_encoder)
        return att_outs + fsmn_memory

    def forward_chunk(self, x, cache=None, chunk_size=None, look_back=0):
        """Compute scaled dot product attention.

        Args:
            query (torch.Tensor): Query tensor (#batch, time1, size).
            key (torch.Tensor): Key tensor (#batch, time2, size).
            value (torch.Tensor): Value tensor (#batch, time2, size).
            mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
                (#batch, time1, time2).

        Returns:
            torch.Tensor: Output tensor (#batch, time1, d_model).

        """
        q_h, k_h, v_h, v = self.forward_qkv(x)
        if chunk_size is not None and look_back > 0 or look_back == -1:
            if cache is not None:
                k_h_stride = k_h[:, :, : -(chunk_size[2]), :]
                v_h_stride = v_h[:, :, : -(chunk_size[2]), :]
                k_h = torch.cat((cache["k"], k_h), dim=2)
                v_h = torch.cat((cache["v"], v_h), dim=2)

                cache["k"] = torch.cat((cache["k"], k_h_stride), dim=2)
                cache["v"] = torch.cat((cache["v"], v_h_stride), dim=2)
                if look_back != -1:
                    cache["k"] = cache["k"][:, :, -(look_back * chunk_size[1]) :, :]
                    cache["v"] = cache["v"][:, :, -(look_back * chunk_size[1]) :, :]
            else:
                cache_tmp = {
                    "k": k_h[:, :, : -(chunk_size[2]), :],
                    "v": v_h[:, :, : -(chunk_size[2]), :],
                }
                cache = cache_tmp
        fsmn_memory = self.forward_fsmn(v, None)
        q_h = q_h * self.d_k ** (-0.5)
        scores = torch.matmul(q_h, k_h.transpose(-2, -1))
        att_outs = self.forward_attention(v_h, scores, None)
        return att_outs + fsmn_memory, cache


class LayerNorm(nn.LayerNorm):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def forward(self, input):
        output = F.layer_norm(
            input.float(),
            self.normalized_shape,
            self.weight.float() if self.weight is not None else None,
            self.bias.float() if self.bias is not None else None,
            self.eps,
        )
        return output.type_as(input)


def sequence_mask(lengths, maxlen=None, dtype=torch.float32, device=None):
    if maxlen is None:
        maxlen = lengths.max()
    row_vector = torch.arange(0, maxlen, 1).to(lengths.device)
    matrix = torch.unsqueeze(lengths, dim=-1)
    mask = row_vector < matrix
    mask = mask.detach()

    return mask.type(dtype).to(device) if device is not None else mask.type(dtype)


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", "SenseVoiceEncoderSmall")
class SenseVoiceEncoderSmall(nn.Module):
    """
    Author: Speech Lab of DAMO Academy, Alibaba Group
    SCAMA: Streaming chunk-aware multihead attention for online end-to-end speech recognition
    https://arxiv.org/abs/2006.01713
    """

    def __init__(
        self,
        input_size: int,
        output_size: int = 256,
        attention_heads: int = 4,
        linear_units: int = 2048,
        num_blocks: int = 6,
        tp_blocks: int = 0,
        dropout_rate: float = 0.1,
        positional_dropout_rate: float = 0.1,
        attention_dropout_rate: float = 0.0,
        stochastic_depth_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,
        kernel_size: int = 11,
        sanm_shfit: int = 0,
        selfattention_layer_type: str = "sanm",
        chunk_size: Optional[Union[int, Sequence[int]]] = None,
        stride: Optional[Union[int, Sequence[int]]] = None,
        pad_left: Optional[Union[int, Sequence[int]]] = None,
        encoder_att_look_back_factor: Optional[Union[int, Sequence[int]]] = None,
        decoder_att_look_back_factor: Optional[Union[int, Sequence[int]]] = None,
        **kwargs,
    ):
        """Build the encoder.

        Args:
            chunk_size: Total chunk width(s) in frames used by the chunk-mask training
                path, i.e. ``pad_left + stride + pad_right``. ``None`` (the default)
                disables the chunk path completely and ``forward`` runs the unmodified
                full-attention computation. A tuple with several entries turns on
                dynamic chunk training: one entry is drawn uniformly at random per
                training step. An entry ``<= 0`` is the *full attention sentinel* -
                when that entry is drawn the step runs the plain full-attention path,
                which is how full-attention batches are mixed into chunk training.
                ``-1`` is the canonical spelling of the sentinel.
            stride: Committed frames per chunk, one entry per ``chunk_size`` entry (a
                single entry is broadcast). Defaults to ``chunk_size`` (no right
                context). Must satisfy ``0 < stride <= chunk_size``.
            pad_left: Left context frames per chunk, broadcast like ``stride``.
                Defaults to ``0``. ``pad_right`` is implied as
                ``chunk_size - stride - pad_left`` and must be ``>= 0``.
            encoder_att_look_back_factor: How many previous chunks the encoder
                self-attention may attend to. Defaults to ``(1,)``.
            decoder_att_look_back_factor: Unused by SenseVoice (there is no decoder);
                accepted and forwarded because ``overlap_chunk`` requires it.

        Note:
            The 3-tuple passed to :meth:`forward_chunk` via the cache has a *different*
            meaning: ``[pad_left, stride, pad_right]``. Use :meth:`init_chunk_cache` to
            derive it from the training configuration instead of writing it by hand.
        """
        super().__init__()
        self._output_size = output_size
        self._input_size = input_size

        self.embed = SinusoidalPositionEncoder()

        self.normalize_before = normalize_before

        positionwise_layer = PositionwiseFeedForward
        positionwise_layer_args = (
            output_size,
            linear_units,
            dropout_rate,
        )

        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 = nn.ModuleList(
            [
                EncoderLayerSANM(
                    input_size,
                    output_size,
                    encoder_selfattn_layer(*encoder_selfattn_layer_args0),
                    positionwise_layer(*positionwise_layer_args),
                    dropout_rate,
                )
                for i in range(1)
            ]
        )
        self.encoders = nn.ModuleList(
            [
                EncoderLayerSANM(
                    output_size,
                    output_size,
                    encoder_selfattn_layer(*encoder_selfattn_layer_args),
                    positionwise_layer(*positionwise_layer_args),
                    dropout_rate,
                )
                for i in range(num_blocks - 1)
            ]
        )

        self.tp_encoders = nn.ModuleList(
            [
                EncoderLayerSANM(
                    output_size,
                    output_size,
                    encoder_selfattn_layer(*encoder_selfattn_layer_args),
                    positionwise_layer(*positionwise_layer_args),
                    dropout_rate,
                )
                for i in range(tp_blocks)
            ]
        )

        self.after_norm = LayerNorm(output_size)

        self.tp_norm = LayerNorm(output_size)

        # Streaming position encoder used only by ``forward_chunk``. It holds no
        # parameters and no buffers, so ``state_dict()`` is byte-for-byte unchanged
        # and published checkpoints still load with ``strict=True``.
        self.chunk_embed = StreamSinusoidalPositionEncoder()

        self.chunk_size, self.stride, self.pad_left = None, None, None
        self.overlap_chunk_cls = None
        if chunk_size is not None:
            self.chunk_size, self.stride, self.pad_left = self._normalize_chunk_args(
                chunk_size, stride, pad_left
            )
            self.overlap_chunk_cls = overlap_chunk(
                chunk_size=self.chunk_size,
                stride=self.stride,
                pad_left=self.pad_left,
                shfit_fsmn=(kernel_size - 1) // 2,
                encoder_att_look_back_factor=self._as_int_tuple(
                    encoder_att_look_back_factor, (1,)
                ),
                decoder_att_look_back_factor=self._as_int_tuple(
                    decoder_att_look_back_factor, (1,)
                ),
            )

    @staticmethod
    def _as_int_tuple(
        value: Optional[Union[int, Sequence[int]]],
        default: Tuple[int, ...],
    ) -> Tuple[int, ...]:
        """Coerce a scalar / sequence / omitted chunk argument to a tuple of ints.

        Args:
            value: An int, any sequence of ints (including a Hydra ``ListConfig``),
                or ``None``.
            default: Returned when ``value`` is ``None``.

        Returns:
            The coerced tuple.
        """
        if value is None:
            return default
        if isinstance(value, int):
            return (value,)
        return tuple(int(v) for v in value)

    @classmethod
    def _normalize_chunk_args(
        cls,
        chunk_size: Union[int, Sequence[int]],
        stride: Optional[Union[int, Sequence[int]]],
        pad_left: Optional[Union[int, Sequence[int]]],
    ) -> Tuple[Tuple[int, ...], Tuple[int, ...], Tuple[int, ...]]:
        """Validate and broadcast the chunk-mask training arguments.

        ``overlap_chunk`` broadcasts ``pad_left`` and the look-back factors itself but
        *not* ``stride``, so every tuple is expanded to ``len(chunk_size)`` here and the
        per-entry geometry is checked up front rather than failing deep inside mask
        generation.

        Args:
            chunk_size: Total chunk width(s); entries ``<= 0`` are full-attention
                sentinels and skip geometry validation.
            stride: Committed frames per chunk, or ``None`` to default to ``chunk_size``.
            pad_left: Left context frames, or ``None`` to default to ``0``.

        Returns:
            ``(chunk_size, stride, pad_left)`` as equal-length tuples of ints.

        Raises:
            ValueError: If a tuple length cannot be broadcast, or if any non-sentinel
                entry violates ``0 < stride <= chunk_size`` or
                ``0 <= pad_left <= chunk_size - stride``.
        """
        chunk_size = cls._as_int_tuple(chunk_size, ())
        if not chunk_size:
            raise ValueError("chunk_size must contain at least one entry")

        def _broadcast(value: Tuple[int, ...], name: str) -> Tuple[int, ...]:
            if len(value) == 1:
                return value * len(chunk_size)
            if len(value) != len(chunk_size):
                raise ValueError(
                    f"{name} has {len(value)} entries but chunk_size has "
                    f"{len(chunk_size)}; pass one entry per chunk_size entry or a "
                    "single entry to broadcast"
                )
            return value

        stride = _broadcast(cls._as_int_tuple(stride, chunk_size), "stride")
        pad_left = _broadcast(cls._as_int_tuple(pad_left, (0,)), "pad_left")

        for i, total in enumerate(chunk_size):
            if total <= 0:
                continue  # full-attention sentinel: geometry is not used
            if not 0 < stride[i] <= total:
                raise ValueError(
                    f"stride[{i}]={stride[i]} must satisfy 0 < stride <= "
                    f"chunk_size[{i}]={total}"
                )
            if not 0 <= pad_left[i] <= total - stride[i]:
                raise ValueError(
                    f"pad_left[{i}]={pad_left[i]} must satisfy 0 <= pad_left <= "
                    f"chunk_size[{i}] - stride[{i}] = {total - stride[i]}"
                )
        return chunk_size, stride, pad_left

    def _check_chunk_ind(self, ind: Optional[int], name: str) -> None:
        """Range-check an index into the chunk configuration tuples.

        ``overlap_chunk.random_choice`` and ``overlap_chunk.get_chunk_size`` do no
        bounds checking, so an out-of-range index surfaces as a bare ``IndexError``
        deep inside mask generation.

        Args:
            ind: Index to check. ``None`` is accepted and means "not supplied".
            name: Argument name to quote in the error message.

        Raises:
            ValueError: If the chunk path is disabled or ``ind`` is out of range.
        """
        if ind is None:
            return
        if self.overlap_chunk_cls is None:
            raise ValueError(
                f"{name}={ind} was given but this encoder was built without chunk_size"
            )
        if not 0 <= ind < len(self.chunk_size):
            raise ValueError(
                f"{name}={ind} is out of range for {len(self.chunk_size)} configured "
                "chunk size(s)"
            )

    def output_size(self) -> int:
        return self._output_size

    def forward(
        self,
        xs_pad: torch.Tensor,
        ilens: torch.Tensor,
        decoding_ind: Optional[int] = None,
        **kwargs,
    ):
        """Embed positions in tensor.

        Args:
            xs_pad: Input features (batch, time, input_size). Scaled in place, as
                before.
            ilens: Input lengths (batch,).
            decoding_ind: Chunk configuration index to use when not training. Ignored
                unless the chunk path is enabled; ``None`` uses index 0. During
                training the index is drawn at random from the configured tuple
                (dynamic chunk training).
            **kwargs: Ignored; accepted so callers can pass extra keywords.

        Returns:
            ``(xs_pad, olens)`` with the same shape and length semantics as the
            full-attention path. When the chunk path is active the chunk-expanded
            sequence is folded back to the original time axis before returning, so the
            CTC head and loss need no changes.

        Note:
            When the chunk path is active the returned time axis is truncated to
            ``ilens.max()``; the full-attention path keeps ``xs_pad.size(1)``. These
            agree for the usual case of a batch padded exactly to its longest item.
        """
        masks = sequence_mask(ilens, device=ilens.device)[:, None, :]

        xs_pad *= self.output_size() ** 0.5

        xs_pad = self.embed(xs_pad)

        self._check_chunk_ind(decoding_ind, "decoding_ind")
        chunk_outs, olens_chunk = None, None
        mask_shfit_chunk, mask_att_chunk_encoder = None, None
        if self.overlap_chunk_cls is not None:
            ind = self.overlap_chunk_cls.random_choice(self.training, decoding_ind)
            # A non-positive entry is the full-attention sentinel: leave every chunk
            # tensor as None so this step runs the plain full-attention computation.
            if self.chunk_size[ind] > 0:
                chunk_ilens = masks.squeeze(1).sum(1).to(torch.int64)
                chunk_outs = self.overlap_chunk_cls.gen_chunk_mask(chunk_ilens, ind)
                xs_pad, olens_chunk = self.overlap_chunk_cls.split_chunk(
                    xs_pad, chunk_ilens, chunk_outs=chunk_outs
                )
                masks = sequence_mask(
                    olens_chunk, maxlen=xs_pad.size(1), device=xs_pad.device
                )[:, None, :]
                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
                )

        # forward encoder1
        for layer_idx, encoder_layer in enumerate(self.encoders0):
            encoder_outs = encoder_layer(
                xs_pad, masks, None, mask_shfit_chunk, mask_att_chunk_encoder
            )
            xs_pad, masks = encoder_outs[0], encoder_outs[1]

        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]

        xs_pad = self.after_norm(xs_pad)

        # forward encoder2
        olens = masks.squeeze(1).sum(1).int()

        for layer_idx, encoder_layer in enumerate(self.tp_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]

        xs_pad = self.tp_norm(xs_pad)

        if chunk_outs is not None:
            # Fold the chunk-expanded sequence back onto the original time axis so the
            # CTC loss is computed at full length, unlike funasr's SANMEncoderChunkOpt
            # which returns the expanded sequence. remove_chunk masks its input in
            # place, so hand it a clone and keep tp_norm's output untouched for
            # autograd.
            xs_pad, olens = self.overlap_chunk_cls.remove_chunk(
                xs_pad.clone(), olens_chunk, chunk_outs
            )
            olens = olens.int()
        return xs_pad, olens

    def init_chunk_cache(
        self,
        pad_left: Optional[int] = None,
        stride: Optional[int] = None,
        pad_right: Optional[int] = None,
        encoder_chunk_look_back: int = 0,
        batch_size: int = 1,
        device: Optional[Union[str, torch.device]] = None,
        dtype: torch.dtype = torch.float32,
        ind: int = 0,
    ) -> dict:
        """Build a fresh streaming cache for :meth:`forward_chunk`.

        Args:
            pad_left: Left context frames carried over from the previous chunk.
            stride: Committed frames per :meth:`forward_chunk` call. This is how many
                new frames the caller feeds each step.
            pad_right: Lookahead frames withheld from the output until the next call.
            encoder_chunk_look_back: Number of previous chunks the encoder
                self-attention may attend to through the per-layer key/value cache.
                ``0`` (the default) keeps attention inside the current window, ``-1``
                keeps every past chunk.
            batch_size: Batch size of the stream. Streaming decode is normally 1.
            device: Device for the cached tensors. Defaults to CPU.
            dtype: Dtype for the cached tensors.
            ind: Index into the encoder's training chunk configuration, used only when
                ``stride`` is omitted.

        Returns:
            The cache dict. It is mutated in place by :meth:`forward_chunk`, so hold on
            to it and pass the same object every call. Keys:

            - ``start_idx`` (int): absolute frame offset consumed so far, used to keep
              the sinusoidal position encoding continuous across calls.
            - ``chunk_size`` (list[int]): ``[pad_left, stride, pad_right]``.
            - ``encoder_chunk_look_back`` (int): as above.
            - ``opt`` (list | None): per-layer attention key/value caches, one entry per
              encoder layer in ``encoders0 + encoders + tp_encoders``. Each entry is
              ``None`` or ``{"k": (batch, head, time, d_k), "v": (batch, head, time,
              d_k)}``. Stays ``None`` when ``encoder_chunk_look_back == 0``.
            - ``feats`` (Tensor): ``(batch, kept, input_size)`` overlap of already
              scaled and position-encoded frames prepended to the next chunk. Starts
              with ``pad_left`` zero frames and holds ``pad_left + pad_right`` frames
              afterwards. That initial length deviates from upstream funasr; see the
              note below.
            - ``tail_chunk`` (bool): set to ``True`` by the caller before the final
              call to flush the withheld lookahead frames.

        Raises:
            ValueError: If the geometry is invalid, or if ``stride`` is omitted while
                the encoder was built without ``chunk_size``.

        Note:
            When ``stride`` is omitted the geometry is derived from the training
            configuration at ``ind``: ``pad_left = self.pad_left[ind]``,
            ``stride = self.stride[ind]``,
            ``pad_right = self.chunk_size[ind] - stride - pad_left``.

        Note:
            ``encoder_chunk_look_back != 0`` requires ``pad_right >= 1`` because
            :meth:`MultiHeadedAttentionSANM.forward_chunk` trims the lookahead with
            ``k_h[:, :, :-pad_right, :]``, which would empty the cache at
            ``pad_right == 0``. Combining it with ``pad_left > 0`` also re-appends the
            left-context frames to that cache (upstream funasr behaviour), so prefer
            ``pad_left = 0`` whenever look-back is enabled.

        Note:
            ``feats`` is seeded with ``pad_left`` zero frames, where upstream funasr
            (``funasr/models/scama/model.py``, ``init_cache``) seeds
            ``pad_left + pad_right``. On the very first call nothing has been withheld
            yet, so upstream's extra ``pad_right`` frames stand for a lookahead that
            does not exist. They are not merely wasted compute: they enter
            self-attention and the FSMN convolution, so the first call emits ``stride``
            frames instead of ``stride - pad_right`` and the stream carries
            ``pad_right`` phantom output frames for the whole utterance, breaking the
            one-to-one input/output frame alignment CTC decoding depends on.
            ``tests/test_chunk_streaming_equivalence.py`` pins this. The raw first
            layer-0 windows cannot be compared directly, since upstream's is exactly
            ``pad_right`` frames longer; those surplus leading frames are exactly
            ``0.0``, and after realigning (dropping them) the two windows are
            bit-identical for every geometry tested. The resulting output corruption is
            confined to the first emitted chunk -- it peaks at ``1.514661`` for
            ``(pad_left, stride, pad_right) = (0, 10, 5)``, and every frame from index
            ``stride`` onward is bit-identical.
        """
        if stride is None:
            if self.overlap_chunk_cls is None:
                raise ValueError(
                    "stride must be given because this encoder was built without "
                    "chunk_size, so there is no training configuration to derive from"
                )
            self._check_chunk_ind(ind, "ind")
            if self.chunk_size[ind] <= 0:
                raise ValueError(
                    f"chunk_size[{ind}]={self.chunk_size[ind]} is the full-attention "
                    "sentinel and has no streaming geometry; pick another ind or pass "
                    "stride explicitly"
                )
            pad_left = self.pad_left[ind]
            stride = self.stride[ind]
            pad_right = self.chunk_size[ind] - stride - pad_left
        else:
            pad_left = 0 if pad_left is None else int(pad_left)
            stride = int(stride)
            pad_right = 0 if pad_right is None else int(pad_right)

        if stride < 1:
            raise ValueError(f"stride must be >= 1, got {stride}")
        if pad_left < 0 or pad_right < 0:
            raise ValueError(
                f"pad_left and pad_right must be >= 0, got {pad_left} and {pad_right}"
            )
        if encoder_chunk_look_back != 0 and pad_right < 1:
            raise ValueError(
                "encoder_chunk_look_back != 0 requires pad_right >= 1; the attention "
                "cache is built by dropping the last pad_right frames and would be "
                "empty otherwise"
            )

        return {
            "start_idx": 0,
            "chunk_size": [pad_left, stride, pad_right],
            "encoder_chunk_look_back": encoder_chunk_look_back,
            "opt": None,
            "feats": torch.zeros(
                (batch_size, pad_left, self._input_size), dtype=dtype, device=device
            ),
            "tail_chunk": False,
        }

    def _add_overlap_chunk(self, feats: torch.Tensor, cache: dict) -> torch.Tensor:
        """Prepend the cached overlap to the incoming frames and refresh the cache.

        The SANM memory block has no streaming cache of its own; its context is
        supplied by overlapping the chunks instead, which is what this does.

        Args:
            feats: Scaled and position-encoded new frames (batch, time, input_size).
            cache: Cache dict from :meth:`init_chunk_cache`.

        Returns:
            The concatenated window (batch, kept + time, input_size).
        """
        cached = cache.get("feats")
        if cached is None:
            return feats
        overlap_feats = torch.cat(
            (cached.to(device=feats.device, dtype=feats.dtype), feats), dim=1
        )
        keep = cache["chunk_size"][0] + cache["chunk_size"][2]
        cache["feats"] = overlap_feats[:, max(0, overlap_feats.size(1) - keep) :, :]
        return overlap_feats

    def forward_chunk(
        self,
        xs_pad: torch.Tensor,
        cache: Optional[dict] = None,
        **kwargs,
    ) -> Tuple[torch.Tensor, dict]:
        """Encode one streaming chunk.

        Args:
            xs_pad: New frames for this step, ``(batch, stride, input_size)``. The last
                real chunk may be shorter. Not modified in place, unlike
                :meth:`forward`. Ignored when ``cache["tail_chunk"]`` is set.
            cache: Cache dict from :meth:`init_chunk_cache`, mutated in place. When
                ``None`` a default cache is built from the training configuration.
            **kwargs: ``trim`` (bool, default ``True``) returns only the frames that
                are final at this step. Pass ``trim=False`` for the raw
                ``pad_left + time + pad_right`` window, matching funasr's
                ``SANMEncoderChunkOpt.forward_chunk``.

        Returns:
            ``(xs_pad, cache)``. With ``trim=True`` the output frames line up one to
            one with the input frames, delayed by ``pad_right``: the first call returns
            ``time - pad_right`` frames, later calls return ``time`` frames, and the
            ``tail_chunk`` call flushes the final ``pad_right`` frames. Concatenating
            every call's output reproduces the whole utterance in order.

        Note:
            Run this under ``torch.no_grad()`` with the module in ``eval()`` mode; the
            per-layer caches assume a single monotonic pass over one stream.

        Note:
            The caller owns the SenseVoice prompt frames. Prepend the four
            ``[language, event, emo, textnorm]`` frames to the first chunk so they take
            absolute positions 1-4, exactly as ``SenseVoiceSmall.encode`` does.
        """
        if cache is None:
            cache = self.init_chunk_cache(
                batch_size=xs_pad.size(0), device=xs_pad.device, dtype=xs_pad.dtype
            )
        pad_left, _, pad_right = cache["chunk_size"]
        is_tail = cache.get("tail_chunk", False)

        if is_tail:
            # The cached overlap is already scaled and position-encoded; re-running it
            # is what flushes the frames whose lookahead never arrived.
            xs_pad = cache["feats"].to(device=xs_pad.device, dtype=xs_pad.dtype)
        else:
            xs_pad = xs_pad * self.output_size() ** 0.5
            xs_pad = self.chunk_embed(xs_pad, cache)
            xs_pad = self._add_overlap_chunk(xs_pad, cache)

        if cache["opt"] is None:
            cache_layer_num = (
                len(self.encoders0) + len(self.encoders) + len(self.tp_encoders)
            )
            new_cache = [None] * cache_layer_num
        else:
            new_cache = cache["opt"]

        chunk_size, look_back = cache["chunk_size"], cache["encoder_chunk_look_back"]
        offset = 0
        for layer_idx, encoder_layer in enumerate(self.encoders0):
            idx = offset + layer_idx
            xs_pad, new_cache[idx] = encoder_layer.forward_chunk(
                xs_pad, new_cache[idx], chunk_size, look_back
            )

        offset += len(self.encoders0)
        for layer_idx, encoder_layer in enumerate(self.encoders):
            idx = offset + layer_idx
            xs_pad, new_cache[idx] = encoder_layer.forward_chunk(
                xs_pad, new_cache[idx], chunk_size, look_back
            )

        xs_pad = self.after_norm(xs_pad)

        offset += len(self.encoders)
        for layer_idx, encoder_layer in enumerate(self.tp_encoders):
            idx = offset + layer_idx
            xs_pad, new_cache[idx] = encoder_layer.forward_chunk(
                xs_pad, new_cache[idx], chunk_size, look_back
            )

        xs_pad = self.tp_norm(xs_pad)

        if look_back > 0 or look_back == -1:
            cache["opt"] = new_cache

        if kwargs.get("trim", True):
            end = xs_pad.size(1) if is_tail else xs_pad.size(1) - pad_right
            xs_pad = xs_pad[:, pad_left : max(0, end), :]
        return xs_pad, cache


@tables.register("model_classes", "SenseVoiceSmall")
class SenseVoiceSmall(nn.Module):
    """CTC-attention hybrid Encoder-Decoder model"""

    def __init__(
        self,
        specaug: str = None,
        specaug_conf: dict = None,
        normalize: str = None,
        normalize_conf: dict = None,
        encoder: str = None,
        encoder_conf: dict = None,
        ctc_conf: dict = None,
        input_size: int = 80,
        vocab_size: int = -1,
        ignore_id: int = -1,
        blank_id: int = 0,
        sos: int = 1,
        eos: int = 2,
        length_normalized_loss: bool = False,
        **kwargs,
    ):

        super().__init__()

        if specaug is not None:
            specaug_class = tables.specaug_classes.get(specaug)
            specaug = specaug_class(**specaug_conf)
        if normalize is not None:
            normalize_class = tables.normalize_classes.get(normalize)
            normalize = normalize_class(**normalize_conf)
        encoder_class = tables.encoder_classes.get(encoder)
        encoder = encoder_class(input_size=input_size, **encoder_conf)
        encoder_output_size = encoder.output_size()

        if ctc_conf is None:
            ctc_conf = {}
        ctc = CTC(odim=vocab_size, encoder_output_size=encoder_output_size, **ctc_conf)

        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.specaug = specaug
        self.normalize = normalize
        self.encoder = encoder
        self.error_calculator = None

        self.ctc = ctc

        self.length_normalized_loss = length_normalized_loss
        self.encoder_output_size = encoder_output_size

        self.lid_dict = {"auto": 0, "zh": 3, "en": 4, "yue": 7, "ja": 11, "ko": 12, "nospeech": 13}
        self.lid_int_dict = {24884: 3, 24885: 4, 24888: 7, 24892: 11, 24896: 12, 24992: 13}
        self.textnorm_dict = {"withitn": 14, "woitn": 15}
        self.textnorm_int_dict = {25016: 14, 25017: 15}
        self.embed = torch.nn.Embedding(7 + len(self.lid_dict) + len(self.textnorm_dict), input_size)
        self.emo_dict = {"unk": 25009, "happy": 25001, "sad": 25002, "angry": 25003, "neutral": 25004}
        
        self.criterion_att = LabelSmoothingLoss(
            size=self.vocab_size,
            padding_idx=self.ignore_id,
            smoothing=kwargs.get("lsm_weight", 0.0),
            normalize_length=self.length_normalized_loss,
        )
    
    @staticmethod
    def from_pretrained(model:str=None, **kwargs):
        from funasr import AutoModel
        model, kwargs = AutoModel.build_model(model=model, trust_remote_code=True, **kwargs)
        
        return model, kwargs

    def forward(
        self,
        speech: torch.Tensor,
        speech_lengths: torch.Tensor,
        text: torch.Tensor,
        text_lengths: torch.Tensor,
        **kwargs,
    ):
        """Encoder + Decoder + Calc loss
        Args:
                speech: (Batch, Length, ...)
                speech_lengths: (Batch, )
                text: (Batch, Length)
                text_lengths: (Batch,)
        """
        # import pdb;
        # pdb.set_trace()
        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]

        # 1. Encoder
        encoder_out, encoder_out_lens = self.encode(speech, speech_lengths, text)

        loss_ctc, cer_ctc = None, None
        loss_rich, acc_rich = None, None
        stats = dict()

        loss_ctc, cer_ctc = self._calc_ctc_loss(
            encoder_out[:, 4:, :], encoder_out_lens - 4, text[:, 4:], text_lengths - 4
        )

        loss_rich, acc_rich = self._calc_rich_ce_loss(
            encoder_out[:, :4, :], text[:, :4]
        )

        loss = loss_ctc + loss_rich
        # Collect total loss stats
        stats["loss_ctc"] = torch.clone(loss_ctc.detach()) if loss_ctc is not None else None
        stats["loss_rich"] = torch.clone(loss_rich.detach()) if loss_rich is not None else None
        stats["loss"] = torch.clone(loss.detach()) if loss is not None else None
        stats["acc_rich"] = acc_rich

        # force_gatherable: to-device and to-tensor if scalar for DataParallel
        if self.length_normalized_loss:
            batch_size = int((text_lengths + 1).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,
        text: torch.Tensor,
        **kwargs,
    ):
        """Frontend + Encoder. Note that this method is used by asr_inference.py
        Args:
                speech: (Batch, Length, ...)
                speech_lengths: (Batch, )
                ind: int
        """

        # 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)


        lids = torch.LongTensor([[self.lid_int_dict[int(lid)] if torch.rand(1) > 0.2 and int(lid) in self.lid_int_dict else 0 ] for lid in text[:, 0]]).to(speech.device)
        language_query = self.embed(lids)
        
        styles = torch.LongTensor([[self.textnorm_int_dict[int(style)]] for style in text[:, 3]]).to(speech.device)
        style_query = self.embed(styles)
        speech = torch.cat((style_query, speech), dim=1)
        speech_lengths += 1

        event_emo_query = self.embed(torch.LongTensor([[1, 2]]).to(speech.device)).repeat(speech.size(0), 1, 1)
        input_query = torch.cat((language_query, event_emo_query), dim=1)
        speech = torch.cat((input_query, speech), dim=1)
        speech_lengths += 3

        encoder_out, encoder_out_lens = self.encoder(speech, speech_lengths)

        return encoder_out, encoder_out_lens

    def _calc_ctc_loss(
        self,
        encoder_out: torch.Tensor,
        encoder_out_lens: torch.Tensor,
        ys_pad: torch.Tensor,
        ys_pad_lens: torch.Tensor,
    ):
        # Calc CTC loss
        loss_ctc = self.ctc(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens)

        # Calc CER using CTC
        cer_ctc = None
        if not self.training and self.error_calculator is not None:
            ys_hat = self.ctc.argmax(encoder_out).data
            cer_ctc = self.error_calculator(ys_hat.cpu(), ys_pad.cpu(), is_ctc=True)
        return loss_ctc, cer_ctc

    def _calc_rich_ce_loss(
        self,
        encoder_out: torch.Tensor,
        ys_pad: torch.Tensor,
    ):
        decoder_out = self.ctc.ctc_lo(encoder_out)
        # 2. Compute attention loss
        loss_rich = self.criterion_att(decoder_out, ys_pad.contiguous())
        acc_rich = th_accuracy(
            decoder_out.view(-1, self.vocab_size),
            ys_pad.contiguous(),
            ignore_label=self.ignore_id,
        )

        return loss_rich, acc_rich


    def inference(
        self,
        data_in,
        data_lengths=None,
        key: list = ["wav_file_tmp_name"],
        tokenizer=None,
        frontend=None,
        **kwargs,
    ):


        meta_data = {}
        if (
            isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank"
        ):  # fbank
            speech, speech_lengths = data_in, data_lengths
            if len(speech.shape) < 3:
                speech = speech[None, :, :]
            if speech_lengths is None:
                speech_lengths = speech.shape[1]
        else:
            # extract fbank feats
            time1 = time.perf_counter()
            audio_sample_list = load_audio_text_image_video(
                data_in,
                fs=frontend.fs,
                audio_fs=kwargs.get("fs", 16000),
                data_type=kwargs.get("data_type", "sound"),
                tokenizer=tokenizer,
            )
            time2 = time.perf_counter()
            meta_data["load_data"] = f"{time2 - time1:0.3f}"
            speech, speech_lengths = extract_fbank(
                audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend
            )
            time3 = time.perf_counter()
            meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
            meta_data["batch_data_time"] = (
                speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000
            )

        speech = speech.to(device=kwargs["device"])
        speech_lengths = speech_lengths.to(device=kwargs["device"])

        language = kwargs.get("language", "auto")
        language_query = self.embed(
            torch.LongTensor(
                [[self.lid_dict[language] if language in self.lid_dict else 0]]
            ).to(speech.device)
        ).repeat(speech.size(0), 1, 1)
        
        use_itn = kwargs.get("use_itn", False)
        output_timestamp = kwargs.get("output_timestamp", False)

        textnorm = kwargs.get("text_norm", None)
        if textnorm is None:
            textnorm = "withitn" if use_itn else "woitn"
        textnorm_query = self.embed(
            torch.LongTensor([[self.textnorm_dict[textnorm]]]).to(speech.device)
        ).repeat(speech.size(0), 1, 1)
        speech = torch.cat((textnorm_query, speech), dim=1)
        speech_lengths += 1

        event_emo_query = self.embed(torch.LongTensor([[1, 2]]).to(speech.device)).repeat(
            speech.size(0), 1, 1
        )
        input_query = torch.cat((language_query, event_emo_query), dim=1)
        speech = torch.cat((input_query, speech), dim=1)
        speech_lengths += 3

        # Encoder
        encoder_out, encoder_out_lens = self.encoder(speech, speech_lengths)
        if isinstance(encoder_out, tuple):
            encoder_out = encoder_out[0]

        # c. Passed the encoder result and the beam search
        ctc_logits = self.ctc.log_softmax(encoder_out)
        if kwargs.get("ban_emo_unk", False):
            ctc_logits[:, :, self.emo_dict["unk"]] = -float("inf")

        results = []
        b, n, d = encoder_out.size()
        if isinstance(key[0], (list, tuple)):
            key = key[0]
        if len(key) < b:
            key = key * b
        for i in range(b):
            x = ctc_logits[i, : encoder_out_lens[i].item(), :]
            yseq = x.argmax(dim=-1)
            yseq = torch.unique_consecutive(yseq, dim=-1)

            ibest_writer = None
            if kwargs.get("output_dir") is not None:
                if not hasattr(self, "writer"):
                    self.writer = DatadirWriter(kwargs.get("output_dir"))
                ibest_writer = self.writer[f"1best_recog"]

            mask = yseq != self.blank_id
            token_int = yseq[mask].tolist()

            # Change integer-ids to tokens
            text = tokenizer.decode(token_int)
            if ibest_writer is not None:
                ibest_writer["text"][key[i]] = text

            if output_timestamp:
                from itertools import groupby
                timestamp = []
                tokens = tokenizer.text2tokens(text)[4:]

                logits_speech = self.ctc.softmax(encoder_out)[i, 4:encoder_out_lens[i].item(), :]

                pred = logits_speech.argmax(-1).cpu()
                logits_speech[pred==self.blank_id, self.blank_id] = 0

                align = ctc_forced_align(
                    logits_speech.unsqueeze(0).float(),
                    torch.Tensor(token_int[4:]).unsqueeze(0).long().to(logits_speech.device),
                    (encoder_out_lens-4).long()[i],
                    torch.tensor(len(token_int)-4).unsqueeze(0).long().to(logits_speech.device),
                    ignore_id=self.ignore_id,
                )

                pred = groupby(align[0, :encoder_out_lens[0]])
                _start = 0
                token_id = 0
                ts_max = encoder_out_lens[i] - 4
                for pred_token, pred_frame in pred:
                    _end = _start + len(list(pred_frame))
                    if pred_token != 0 and token_id < len(tokens):
                        ts_left = max((_start*60-30)/1000, 0)
                        ts_right = min((_end*60-30)/1000, (ts_max*60-30)/1000)
                        timestamp.append([tokens[token_id], ts_left, ts_right])
                        token_id += 1
                    _start = _end

                result_i = {"key": key[i], "text": text, "timestamp": timestamp}
                results.append(result_i)
            else:
                result_i = {"key": key[i], "text": text}
                results.append(result_i)
        return results, meta_data

    def export(self, **kwargs):
        from export_meta import export_rebuild_model

        if "max_seq_len" not in kwargs:
            kwargs["max_seq_len"] = 512
        models = export_rebuild_model(model=self, **kwargs)
        return models