Initial commit: FunASR Speech Recognition Toolkit

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

funasr/models/transducer/beam_search_transducer.py

@@ -0,0 +1,694 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import torch
+import numpy as np
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from funasr.models.transducer.joint_network import JointNetwork
+
+
+@dataclass
+class Hypothesis:
+    """Default hypothesis definition for Transducer search algorithms.
+
+    Args:
+        score: Total log-probability.
+        yseq: Label sequence as integer ID sequence.
+        dec_state: RNNDecoder or StatelessDecoder state.
+                     ((N, 1, D_dec), (N, 1, D_dec) or None) or None
+        lm_state: RNNLM state. ((N, D_lm), (N, D_lm)) or None
+
+    """
+
+    score: float
+    yseq: List[int]
+    dec_state: Optional[Tuple[torch.Tensor, Optional[torch.Tensor]]] = None
+    lm_state: Optional[Union[Dict[str, Any], List[Any]]] = None
+
+
+@dataclass
+class ExtendedHypothesis(Hypothesis):
+    """Extended hypothesis definition for NSC beam search and mAES.
+
+    Args:
+        : Hypothesis dataclass arguments.
+        dec_out: Decoder output sequence. (B, D_dec)
+        lm_score: Log-probabilities of the LM for given label. (vocab_size)
+
+    """
+
+    dec_out: torch.Tensor = None
+    lm_score: torch.Tensor = None
+
+
+class BeamSearchTransducer:
+    """Beam search implementation for Transducer.
+
+    Args:
+        decoder: Decoder module.
+        joint_network: Joint network module.
+        beam_size: Size of the beam.
+        lm: LM class.
+        lm_weight: LM weight for soft fusion.
+        search_type: Search algorithm to use during inference.
+        max_sym_exp: Number of maximum symbol expansions at each time step. (TSD)
+        u_max: Maximum expected target sequence length. (ALSD)
+        nstep: Number of maximum expansion steps at each time step. (mAES)
+        expansion_gamma: Allowed logp difference for prune-by-value method. (mAES)
+        expansion_beta:
+             Number of additional candidates for expanded hypotheses selection. (mAES)
+        score_norm: Normalize final scores by length.
+        nbest: Number of final hypothesis.
+        streaming: Whether to perform chunk-by-chunk beam search.
+
+    """
+
+    def __init__(
+        self,
+        decoder,
+        joint_network: JointNetwork,
+        beam_size: int,
+        lm: Optional[torch.nn.Module] = None,
+        lm_weight: float = 0.1,
+        search_type: str = "default",
+        max_sym_exp: int = 3,
+        u_max: int = 50,
+        nstep: int = 2,
+        expansion_gamma: float = 2.3,
+        expansion_beta: int = 2,
+        score_norm: bool = False,
+        nbest: int = 1,
+        streaming: bool = False,
+    ) -> None:
+        """Construct a BeamSearchTransducer object."""
+        super().__init__()
+
+        self.decoder = decoder
+        self.joint_network = joint_network
+
+        self.vocab_size = decoder.vocab_size
+
+        assert (
+            beam_size <= self.vocab_size
+        ), "beam_size (%d) should be smaller than or equal to vocabulary size (%d)." % (
+            beam_size,
+            self.vocab_size,
+        )
+        self.beam_size = beam_size
+
+        if search_type == "default":
+            self.search_algorithm = self.default_beam_search
+        elif search_type == "tsd":
+            assert max_sym_exp > 1, "max_sym_exp (%d) should be greater than one." % (max_sym_exp)
+            self.max_sym_exp = max_sym_exp
+
+            self.search_algorithm = self.time_sync_decoding
+        elif search_type == "alsd":
+            assert not streaming, "ALSD is not available in streaming mode."
+
+            assert u_max >= 0, "u_max should be a positive integer, a portion of max_T."
+            self.u_max = u_max
+
+            self.search_algorithm = self.align_length_sync_decoding
+        elif search_type == "maes":
+            assert self.vocab_size >= beam_size + expansion_beta, (
+                "beam_size (%d) + expansion_beta (%d) "
+                " should be smaller than or equal to vocab size (%d)."
+                % (beam_size, expansion_beta, self.vocab_size)
+            )
+            self.max_candidates = beam_size + expansion_beta
+
+            self.nstep = nstep
+            self.expansion_gamma = expansion_gamma
+
+            self.search_algorithm = self.modified_adaptive_expansion_search
+        else:
+            raise NotImplementedError("Specified search type (%s) is not supported." % search_type)
+
+        self.use_lm = lm is not None
+
+        if self.use_lm:
+            assert hasattr(lm, "rnn_type"), "Transformer LM is currently not supported."
+
+            self.sos = self.vocab_size - 1
+
+            self.lm = lm
+            self.lm_weight = lm_weight
+
+        self.score_norm = score_norm
+        self.nbest = nbest
+
+        self.reset_inference_cache()
+
+    def __call__(
+        self,
+        enc_out: torch.Tensor,
+        is_final: bool = True,
+    ) -> List[Hypothesis]:
+        """Perform beam search.
+
+        Args:
+            enc_out: Encoder output sequence. (T, D_enc)
+            is_final: Whether enc_out is the final chunk of data.
+
+        Returns:
+            nbest_hyps: N-best decoding results
+
+        """
+        self.decoder.set_device(enc_out.device)
+
+        hyps = self.search_algorithm(enc_out)
+
+        if is_final:
+            self.reset_inference_cache()
+
+            return self.sort_nbest(hyps)
+
+        self.search_cache = hyps
+
+        return hyps
+
+    def reset_inference_cache(self) -> None:
+        """Reset cache for decoder scoring and streaming."""
+        self.decoder.score_cache = {}
+        self.search_cache = None
+
+    def sort_nbest(self, hyps: List[Hypothesis]) -> List[Hypothesis]:
+        """Sort in-place hypotheses by score or score given sequence length.
+
+        Args:
+            hyps: Hypothesis.
+
+        Return:
+            hyps: Sorted hypothesis.
+
+        """
+        if self.score_norm:
+            hyps.sort(key=lambda x: x.score / len(x.yseq), reverse=True)
+        else:
+            hyps.sort(key=lambda x: x.score, reverse=True)
+
+        return hyps[: self.nbest]
+
+    def recombine_hyps(self, hyps: List[Hypothesis]) -> List[Hypothesis]:
+        """Recombine hypotheses with same label ID sequence.
+
+        Args:
+            hyps: Hypotheses.
+
+        Returns:
+            final: Recombined hypotheses.
+
+        """
+        final = {}
+
+        for hyp in hyps:
+            str_yseq = "_".join(map(str, hyp.yseq))
+
+            if str_yseq in final:
+                final[str_yseq].score = np.logaddexp(final[str_yseq].score, hyp.score)
+            else:
+                final[str_yseq] = hyp
+
+        return [*final.values()]
+
+    def select_k_expansions(
+        self,
+        hyps: List[ExtendedHypothesis],
+        topk_idx: torch.Tensor,
+        topk_logp: torch.Tensor,
+    ) -> List[ExtendedHypothesis]:
+        """Return K hypotheses candidates for expansion from a list of hypothesis.
+
+        K candidates are selected according to the extended hypotheses probabilities
+        and a prune-by-value method. Where K is equal to beam_size + beta.
+
+        Args:
+            hyps: Hypotheses.
+            topk_idx: Indices of candidates hypothesis.
+            topk_logp: Log-probabilities of candidates hypothesis.
+
+        Returns:
+            k_expansions: Best K expansion hypotheses candidates.
+
+        """
+        k_expansions = []
+
+        for i, hyp in enumerate(hyps):
+            hyp_i = [(int(k), hyp.score + float(v)) for k, v in zip(topk_idx[i], topk_logp[i])]
+            k_best_exp = max(hyp_i, key=lambda x: x[1])[1]
+
+            k_expansions.append(
+                sorted(
+                    filter(lambda x: (k_best_exp - self.expansion_gamma) <= x[1], hyp_i),
+                    key=lambda x: x[1],
+                    reverse=True,
+                )
+            )
+
+        return k_expansions
+
+    def create_lm_batch_inputs(self, hyps_seq: List[List[int]]) -> torch.Tensor:
+        """Make batch of inputs with left padding for LM scoring.
+
+        Args:
+            hyps_seq: Hypothesis sequences.
+
+        Returns:
+            : Padded batch of sequences.
+
+        """
+        max_len = max([len(h) for h in hyps_seq])
+
+        return torch.LongTensor(
+            [[self.sos] + ([0] * (max_len - len(h))) + h[1:] for h in hyps_seq],
+            device=self.decoder.device,
+        )
+
+    def default_beam_search(self, enc_out: torch.Tensor) -> List[Hypothesis]:
+        """Beam search implementation without prefix search.
+
+        Modified from https://arxiv.org/pdf/1211.3711.pdf
+
+        Args:
+            enc_out: Encoder output sequence. (T, D)
+
+        Returns:
+            nbest_hyps: N-best hypothesis.
+
+        """
+        beam_k = min(self.beam_size, (self.vocab_size - 1))
+        max_t = len(enc_out)
+
+        if self.search_cache is not None:
+            kept_hyps = self.search_cache
+        else:
+            kept_hyps = [
+                Hypothesis(
+                    score=0.0,
+                    yseq=[0],
+                    dec_state=self.decoder.init_state(1),
+                )
+            ]
+
+        for t in range(max_t):
+            hyps = kept_hyps
+            kept_hyps = []
+
+            while True:
+                max_hyp = max(hyps, key=lambda x: x.score)
+                hyps.remove(max_hyp)
+
+                label = torch.full(
+                    (1, 1),
+                    max_hyp.yseq[-1],
+                    dtype=torch.long,
+                    device=self.decoder.device,
+                )
+                dec_out, state = self.decoder.score(
+                    label,
+                    max_hyp.yseq,
+                    max_hyp.dec_state,
+                )
+
+                logp = torch.log_softmax(
+                    self.joint_network(enc_out[t : t + 1, :], dec_out),
+                    dim=-1,
+                ).squeeze(0)
+                top_k = logp[1:].topk(beam_k, dim=-1)
+
+                kept_hyps.append(
+                    Hypothesis(
+                        score=(max_hyp.score + float(logp[0:1])),
+                        yseq=max_hyp.yseq,
+                        dec_state=max_hyp.dec_state,
+                        lm_state=max_hyp.lm_state,
+                    )
+                )
+
+                if self.use_lm:
+                    lm_scores, lm_state = self.lm.score(
+                        torch.LongTensor([self.sos] + max_hyp.yseq[1:], device=self.decoder.device),
+                        max_hyp.lm_state,
+                        None,
+                    )
+                else:
+                    lm_state = max_hyp.lm_state
+
+                for logp, k in zip(*top_k):
+                    score = max_hyp.score + float(logp)
+
+                    if self.use_lm:
+                        score += self.lm_weight * lm_scores[k + 1]
+
+                    hyps.append(
+                        Hypothesis(
+                            score=score,
+                            yseq=max_hyp.yseq + [int(k + 1)],
+                            dec_state=state,
+                            lm_state=lm_state,
+                        )
+                    )
+
+                hyps_max = float(max(hyps, key=lambda x: x.score).score)
+                kept_most_prob = sorted(
+                    [hyp for hyp in kept_hyps if hyp.score > hyps_max],
+                    key=lambda x: x.score,
+                )
+                if len(kept_most_prob) >= self.beam_size:
+                    kept_hyps = kept_most_prob
+                    break
+
+        return kept_hyps
+
+    def align_length_sync_decoding(
+        self,
+        enc_out: torch.Tensor,
+    ) -> List[Hypothesis]:
+        """Alignment-length synchronous beam search implementation.
+
+        Based on https://ieeexplore.ieee.org/document/9053040
+
+        Args:
+            h: Encoder output sequences. (T, D)
+
+        Returns:
+            nbest_hyps: N-best hypothesis.
+
+        """
+        t_max = int(enc_out.size(0))
+        u_max = min(self.u_max, (t_max - 1))
+
+        B = [Hypothesis(yseq=[0], score=0.0, dec_state=self.decoder.init_state(1))]
+        final = []
+
+        if self.use_lm:
+            B[0].lm_state = self.lm.zero_state()
+
+        for i in range(t_max + u_max):
+            A = []
+
+            B_ = []
+            B_enc_out = []
+            for hyp in B:
+                u = len(hyp.yseq) - 1
+                t = i - u
+
+                if t > (t_max - 1):
+                    continue
+
+                B_.append(hyp)
+                B_enc_out.append((t, enc_out[t]))
+
+            if B_:
+                beam_enc_out = torch.stack([b[1] for b in B_enc_out])
+                beam_dec_out, beam_state = self.decoder.batch_score(B_)
+
+                beam_logp = torch.log_softmax(
+                    self.joint_network(beam_enc_out, beam_dec_out),
+                    dim=-1,
+                )
+                beam_topk = beam_logp[:, 1:].topk(self.beam_size, dim=-1)
+
+                if self.use_lm:
+                    beam_lm_scores, beam_lm_states = self.lm.batch_score(
+                        self.create_lm_batch_inputs([b.yseq for b in B_]),
+                        [b.lm_state for b in B_],
+                        None,
+                    )
+
+                for i, hyp in enumerate(B_):
+                    new_hyp = Hypothesis(
+                        score=(hyp.score + float(beam_logp[i, 0])),
+                        yseq=hyp.yseq[:],
+                        dec_state=hyp.dec_state,
+                        lm_state=hyp.lm_state,
+                    )
+
+                    A.append(new_hyp)
+
+                    if B_enc_out[i][0] == (t_max - 1):
+                        final.append(new_hyp)
+
+                    for logp, k in zip(beam_topk[0][i], beam_topk[1][i] + 1):
+                        new_hyp = Hypothesis(
+                            score=(hyp.score + float(logp)),
+                            yseq=(hyp.yseq[:] + [int(k)]),
+                            dec_state=self.decoder.select_state(beam_state, i),
+                            lm_state=hyp.lm_state,
+                        )
+
+                        if self.use_lm:
+                            new_hyp.score += self.lm_weight * beam_lm_scores[i, k]
+                            new_hyp.lm_state = beam_lm_states[i]
+
+                        A.append(new_hyp)
+
+                B = sorted(A, key=lambda x: x.score, reverse=True)[: self.beam_size]
+                B = self.recombine_hyps(B)
+
+        if final:
+            return final
+
+        return B
+
+    def time_sync_decoding(self, enc_out: torch.Tensor) -> List[Hypothesis]:
+        """Time synchronous beam search implementation.
+
+        Based on https://ieeexplore.ieee.org/document/9053040
+
+        Args:
+            enc_out: Encoder output sequence. (T, D)
+
+        Returns:
+            nbest_hyps: N-best hypothesis.
+
+        """
+        if self.search_cache is not None:
+            B = self.search_cache
+        else:
+            B = [
+                Hypothesis(
+                    yseq=[0],
+                    score=0.0,
+                    dec_state=self.decoder.init_state(1),
+                )
+            ]
+
+            if self.use_lm:
+                B[0].lm_state = self.lm.zero_state()
+
+        for enc_out_t in enc_out:
+            A = []
+            C = B
+
+            enc_out_t = enc_out_t.unsqueeze(0)
+
+            for v in range(self.max_sym_exp):
+                D = []
+
+                beam_dec_out, beam_state = self.decoder.batch_score(C)
+
+                beam_logp = torch.log_softmax(
+                    self.joint_network(enc_out_t, beam_dec_out),
+                    dim=-1,
+                )
+                beam_topk = beam_logp[:, 1:].topk(self.beam_size, dim=-1)
+
+                seq_A = [h.yseq for h in A]
+
+                for i, hyp in enumerate(C):
+                    if hyp.yseq not in seq_A:
+                        A.append(
+                            Hypothesis(
+                                score=(hyp.score + float(beam_logp[i, 0])),
+                                yseq=hyp.yseq[:],
+                                dec_state=hyp.dec_state,
+                                lm_state=hyp.lm_state,
+                            )
+                        )
+                    else:
+                        dict_pos = seq_A.index(hyp.yseq)
+
+                        A[dict_pos].score = np.logaddexp(
+                            A[dict_pos].score, (hyp.score + float(beam_logp[i, 0]))
+                        )
+
+                if v < (self.max_sym_exp - 1):
+                    if self.use_lm:
+                        beam_lm_scores, beam_lm_states = self.lm.batch_score(
+                            self.create_lm_batch_inputs([c.yseq for c in C]),
+                            [c.lm_state for c in C],
+                            None,
+                        )
+
+                    for i, hyp in enumerate(C):
+                        for logp, k in zip(beam_topk[0][i], beam_topk[1][i] + 1):
+                            new_hyp = Hypothesis(
+                                score=(hyp.score + float(logp)),
+                                yseq=(hyp.yseq + [int(k)]),
+                                dec_state=self.decoder.select_state(beam_state, i),
+                                lm_state=hyp.lm_state,
+                            )
+
+                            if self.use_lm:
+                                new_hyp.score += self.lm_weight * beam_lm_scores[i, k]
+                                new_hyp.lm_state = beam_lm_states[i]
+
+                            D.append(new_hyp)
+
+                C = sorted(D, key=lambda x: x.score, reverse=True)[: self.beam_size]
+
+            B = sorted(A, key=lambda x: x.score, reverse=True)[: self.beam_size]
+
+        return B
+
+    def modified_adaptive_expansion_search(
+        self,
+        enc_out: torch.Tensor,
+    ) -> List[ExtendedHypothesis]:
+        """Modified version of Adaptive Expansion Search (mAES).
+
+        Based on AES (https://ieeexplore.ieee.org/document/9250505) and
+                 NSC (https://arxiv.org/abs/2201.05420).
+
+        Args:
+            enc_out: Encoder output sequence. (T, D_enc)
+
+        Returns:
+            nbest_hyps: N-best hypothesis.
+
+        """
+        if self.search_cache is not None:
+            kept_hyps = self.search_cache
+        else:
+            init_tokens = [
+                ExtendedHypothesis(
+                    yseq=[0],
+                    score=0.0,
+                    dec_state=self.decoder.init_state(1),
+                )
+            ]
+
+            beam_dec_out, beam_state = self.decoder.batch_score(
+                init_tokens,
+            )
+
+            if self.use_lm:
+                beam_lm_scores, beam_lm_states = self.lm.batch_score(
+                    self.create_lm_batch_inputs([h.yseq for h in init_tokens]),
+                    [h.lm_state for h in init_tokens],
+                    None,
+                )
+
+                lm_state = beam_lm_states[0]
+                lm_score = beam_lm_scores[0]
+            else:
+                lm_state = None
+                lm_score = None
+
+            kept_hyps = [
+                ExtendedHypothesis(
+                    yseq=[0],
+                    score=0.0,
+                    dec_state=self.decoder.select_state(beam_state, 0),
+                    dec_out=beam_dec_out[0],
+                    lm_state=lm_state,
+                    lm_score=lm_score,
+                )
+            ]
+
+        for enc_out_t in enc_out:
+            hyps = kept_hyps
+            kept_hyps = []
+
+            beam_enc_out = enc_out_t.unsqueeze(0)
+
+            list_b = []
+            for n in range(self.nstep):
+                beam_dec_out = torch.stack([h.dec_out for h in hyps])
+
+                beam_logp, beam_idx = torch.log_softmax(
+                    self.joint_network(beam_enc_out, beam_dec_out),
+                    dim=-1,
+                ).topk(self.max_candidates, dim=-1)
+
+                k_expansions = self.select_k_expansions(hyps, beam_idx, beam_logp)
+
+                list_exp = []
+                for i, hyp in enumerate(hyps):
+                    for k, new_score in k_expansions[i]:
+                        new_hyp = ExtendedHypothesis(
+                            yseq=hyp.yseq[:],
+                            score=new_score,
+                            dec_out=hyp.dec_out,
+                            dec_state=hyp.dec_state,
+                            lm_state=hyp.lm_state,
+                            lm_score=hyp.lm_score,
+                        )
+
+                        if k == 0:
+                            list_b.append(new_hyp)
+                        else:
+                            new_hyp.yseq.append(int(k))
+
+                            if self.use_lm:
+                                new_hyp.score += self.lm_weight * float(hyp.lm_score[k])
+
+                            list_exp.append(new_hyp)
+
+                if not list_exp:
+                    kept_hyps = sorted(
+                        self.recombine_hyps(list_b), key=lambda x: x.score, reverse=True
+                    )[: self.beam_size]
+
+                    break
+                else:
+                    beam_dec_out, beam_state = self.decoder.batch_score(
+                        list_exp,
+                    )
+
+                    if self.use_lm:
+                        beam_lm_scores, beam_lm_states = self.lm.batch_score(
+                            self.create_lm_batch_inputs([h.yseq for h in list_exp]),
+                            [h.lm_state for h in list_exp],
+                            None,
+                        )
+
+                    if n < (self.nstep - 1):
+                        for i, hyp in enumerate(list_exp):
+                            hyp.dec_out = beam_dec_out[i]
+                            hyp.dec_state = self.decoder.select_state(beam_state, i)
+
+                            if self.use_lm:
+                                hyp.lm_state = beam_lm_states[i]
+                                hyp.lm_score = beam_lm_scores[i]
+
+                        hyps = list_exp[:]
+                    else:
+                        beam_logp = torch.log_softmax(
+                            self.joint_network(beam_enc_out, beam_dec_out),
+                            dim=-1,
+                        )
+
+                        for i, hyp in enumerate(list_exp):
+                            hyp.score += float(beam_logp[i, 0])
+
+                            hyp.dec_out = beam_dec_out[i]
+                            hyp.dec_state = self.decoder.select_state(beam_state, i)
+
+                            if self.use_lm:
+                                hyp.lm_state = beam_lm_states[i]
+                                hyp.lm_score = beam_lm_scores[i]
+
+                        kept_hyps = sorted(
+                            self.recombine_hyps(list_b + list_exp),
+                            key=lambda x: x.score,
+                            reverse=True,
+                        )[: self.beam_size]
+
+        return kept_hyps

funasr/models/transducer/joint_network.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import torch
+
+from funasr.register import tables
+from funasr.models.transformer.utils.nets_utils import get_activation
+
+
+@tables.register("joint_network_classes", "joint_network")
+class JointNetwork(torch.nn.Module):
+    """Transducer joint network module.
+
+    Args:
+        output_size: Output size.
+        encoder_size: Encoder output size.
+        decoder_size: Decoder output size..
+        joint_space_size: Joint space size.
+        joint_act_type: Type of activation for joint network.
+        **activation_parameters: Parameters for the activation function.
+
+    """
+
+    def __init__(
+        self,
+        output_size: int,
+        encoder_size: int,
+        decoder_size: int,
+        joint_space_size: int = 256,
+        joint_activation_type: str = "tanh",
+    ) -> None:
+        """Construct a JointNetwork object."""
+        super().__init__()
+
+        self.lin_enc = torch.nn.Linear(encoder_size, joint_space_size)
+        self.lin_dec = torch.nn.Linear(decoder_size, joint_space_size, bias=False)
+
+        self.lin_out = torch.nn.Linear(joint_space_size, output_size)
+
+        self.joint_activation = get_activation(joint_activation_type)
+
+    def forward(
+        self,
+        enc_out: torch.Tensor,
+        dec_out: torch.Tensor,
+        project_input: bool = True,
+    ) -> torch.Tensor:
+        """Joint computation of encoder and decoder hidden state sequences.
+
+        Args:
+            enc_out: Expanded encoder output state sequences (B, T, 1, D_enc)
+            dec_out: Expanded decoder output state sequences (B, 1, U, D_dec)
+
+        Returns:
+            joint_out: Joint output state sequences. (B, T, U, D_out)
+
+        """
+        if project_input:
+            joint_out = self.joint_activation(self.lin_enc(enc_out) + self.lin_dec(dec_out))
+        else:
+            joint_out = self.joint_activation(enc_out + dec_out)
+        return self.lin_out(joint_out)

funasr/models/transducer/model.py

@@ -0,0 +1,595 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import time
+import torch
+import logging
+from contextlib import contextmanager
+from typing import Dict, Optional, Tuple
+from distutils.version import LooseVersion
+
+from funasr.register import tables
+from funasr.utils import postprocess_utils
+from funasr.utils.datadir_writer import DatadirWriter
+from funasr.train_utils.device_funcs import force_gatherable
+from funasr.models.transformer.scorers.ctc import CTCPrefixScorer
+from funasr.losses.label_smoothing_loss import LabelSmoothingLoss
+from funasr.models.transformer.scorers.length_bonus import LengthBonus
+from funasr.models.transformer.utils.nets_utils import get_transducer_task_io
+from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
+from funasr.models.transducer.beam_search_transducer import BeamSearchTransducer
+
+
+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", "Transducer")
+class Transducer(torch.nn.Module):
+    """Transducer (RNN-T): Streaming ASR using encoder-predictor-joint architecture.
+
+    Combines encoder (audio frames), predictor (text history), and joint network.
+    Supports beam search decoding. Suitable for low-latency streaming applications.
+
+    Output: {"key": str, "text": str}
+    """
+
+    def __init__(
+        self,
+        frontend: Optional[str] = None,
+        frontend_conf: Optional[Dict] = None,
+        specaug: Optional[str] = None,
+        specaug_conf: Optional[Dict] = None,
+        normalize: str = None,
+        normalize_conf: Optional[Dict] = None,
+        encoder: str = None,
+        encoder_conf: Optional[Dict] = None,
+        decoder: str = None,
+        decoder_conf: Optional[Dict] = None,
+        joint_network: str = None,
+        joint_network_conf: Optional[Dict] = None,
+        transducer_weight: float = 1.0,
+        fastemit_lambda: float = 0.0,
+        auxiliary_ctc_weight: float = 0.0,
+        auxiliary_ctc_dropout_rate: float = 0.0,
+        auxiliary_lm_loss_weight: float = 0.0,
+        auxiliary_lm_loss_smoothing: 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,
+        # report_cer: bool = True,
+        # report_wer: bool = True,
+        # sym_space: str = "<space>",
+        # sym_blank: str = "<blank>",
+        # extract_feats_in_collect_stats: bool = True,
+        share_embedding: bool = False,
+        # preencoder: Optional[AbsPreEncoder] = None,
+        # postencoder: Optional[AbsPostEncoder] = None,
+        **kwargs,
+    ):
+
+        """Initialize Transducer.
+        
+            Args:
+                frontend: Audio frontend for feature extraction.
+                frontend_conf: Configuration dict for frontend.
+                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.
+                joint_network: TODO.
+                joint_network_conf: Configuration dict for joint_network.
+                transducer_weight: TODO.
+                fastemit_lambda: TODO.
+                auxiliary_ctc_weight: TODO.
+                auxiliary_ctc_dropout_rate: TODO.
+                auxiliary_lm_loss_weight: TODO.
+                auxiliary_lm_loss_smoothing: 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,
+            **decoder_conf,
+        )
+        decoder_output_size = decoder.output_size
+
+        joint_network_class = tables.joint_network_classes.get(joint_network)
+        joint_network = joint_network_class(
+            vocab_size,
+            encoder_output_size,
+            decoder_output_size,
+            **joint_network_conf,
+        )
+
+        self.criterion_transducer = None
+        self.error_calculator = None
+
+        self.use_auxiliary_ctc = auxiliary_ctc_weight > 0
+        self.use_auxiliary_lm_loss = auxiliary_lm_loss_weight > 0
+
+        if self.use_auxiliary_ctc:
+            self.ctc_lin = torch.nn.Linear(encoder.output_size(), vocab_size)
+            self.ctc_dropout_rate = auxiliary_ctc_dropout_rate
+
+        if self.use_auxiliary_lm_loss:
+            self.lm_lin = torch.nn.Linear(decoder.output_size, vocab_size)
+            self.lm_loss_smoothing = auxiliary_lm_loss_smoothing
+
+        self.transducer_weight = transducer_weight
+        self.fastemit_lambda = fastemit_lambda
+
+        self.auxiliary_ctc_weight = auxiliary_ctc_weight
+        self.auxiliary_lm_loss_weight = auxiliary_lm_loss_weight
+        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.frontend = frontend
+        self.specaug = specaug
+        self.normalize = normalize
+        self.encoder = encoder
+        self.decoder = decoder
+        self.joint_network = joint_network
+
+        self.criterion_att = LabelSmoothingLoss(
+            size=vocab_size,
+            padding_idx=ignore_id,
+            smoothing=lsm_weight,
+            normalize_length=length_normalized_loss,
+        )
+
+        self.length_normalized_loss = length_normalized_loss
+        self.beam_search = None
+        self.ctc = None
+        self.ctc_weight = 0.0
+
+    def forward(
+        self,
+        speech: torch.Tensor,
+        speech_lengths: torch.Tensor,
+        text: torch.Tensor,
+        text_lengths: torch.Tensor,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:
+        """Encoder + Decoder + Calc loss
+        Args:
+                speech: (Batch, Length, ...)
+                speech_lengths: (Batch, )
+                text: (Batch, Length)
+                text_lengths: (Batch,)
+        """
+        if len(text_lengths.size()) > 1:
+            text_lengths = text_lengths[:, 0]
+        if len(speech_lengths.size()) > 1:
+            speech_lengths = speech_lengths[:, 0]
+
+        batch_size = speech.shape[0]
+        # 1. Encoder
+        encoder_out, encoder_out_lens = self.encode(speech, speech_lengths)
+        if (
+            hasattr(self.encoder, "overlap_chunk_cls")
+            and self.encoder.overlap_chunk_cls is not None
+        ):
+            encoder_out, encoder_out_lens = self.encoder.overlap_chunk_cls.remove_chunk(
+                encoder_out, encoder_out_lens, chunk_outs=None
+            )
+        # 2. Transducer-related I/O preparation
+        decoder_in, target, t_len, u_len = get_transducer_task_io(
+            text,
+            encoder_out_lens,
+            ignore_id=self.ignore_id,
+        )
+
+        # 3. Decoder
+        self.decoder.set_device(encoder_out.device)
+        decoder_out = self.decoder(decoder_in, u_len)
+
+        # 4. Joint Network
+        joint_out = self.joint_network(encoder_out.unsqueeze(2), decoder_out.unsqueeze(1))
+
+        # 5. Losses
+        loss_trans, cer_trans, wer_trans = self._calc_transducer_loss(
+            encoder_out,
+            joint_out,
+            target,
+            t_len,
+            u_len,
+        )
+
+        loss_ctc, loss_lm = 0.0, 0.0
+
+        if self.use_auxiliary_ctc:
+            loss_ctc = self._calc_ctc_loss(
+                encoder_out,
+                target,
+                t_len,
+                u_len,
+            )
+
+        if self.use_auxiliary_lm_loss:
+            loss_lm = self._calc_lm_loss(decoder_out, target)
+
+        loss = (
+            self.transducer_weight * loss_trans
+            + self.auxiliary_ctc_weight * loss_ctc
+            + self.auxiliary_lm_loss_weight * loss_lm
+        )
+
+        stats = dict(
+            loss=loss.detach(),
+            loss_transducer=loss_trans.detach(),
+            aux_ctc_loss=loss_ctc.detach() if loss_ctc > 0.0 else None,
+            aux_lm_loss=loss_lm.detach() if loss_lm > 0.0 else None,
+            cer_transducer=cer_trans,
+            wer_transducer=wer_trans,
+        )
+
+        # force_gatherable: to-device and to-tensor if scalar for DataParallel
+        loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device)
+
+        return loss, stats, weight
+
+    def encode(
+        self,
+        speech: torch.Tensor,
+        speech_lengths: torch.Tensor,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """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
+        # feats: (Batch, Length, Dim)
+        # -> encoder_out: (Batch, Length2, Dim2)
+        encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths)
+        intermediate_outs = None
+        if isinstance(encoder_out, tuple):
+            intermediate_outs = encoder_out[1]
+            encoder_out = encoder_out[0]
+
+        if intermediate_outs is not None:
+            return (encoder_out, intermediate_outs), encoder_out_lens
+
+        return encoder_out, encoder_out_lens
+
+    def _calc_transducer_loss(
+        self,
+        encoder_out: torch.Tensor,
+        joint_out: torch.Tensor,
+        target: torch.Tensor,
+        t_len: torch.Tensor,
+        u_len: torch.Tensor,
+    ) -> Tuple[torch.Tensor, Optional[float], Optional[float]]:
+        """Compute Transducer loss.
+
+        Args:
+            encoder_out: Encoder output sequences. (B, T, D_enc)
+            joint_out: Joint Network output sequences (B, T, U, D_joint)
+            target: Target label ID sequences. (B, L)
+            t_len: Encoder output sequences lengths. (B,)
+            u_len: Target label ID sequences lengths. (B,)
+
+        Return:
+            loss_transducer: Transducer loss value.
+            cer_transducer: Character error rate for Transducer.
+            wer_transducer: Word Error Rate for Transducer.
+
+        """
+        if self.criterion_transducer is None:
+            try:
+                from warp_rnnt import rnnt_loss as RNNTLoss
+
+                self.criterion_transducer = RNNTLoss
+
+            except ImportError:
+                logging.error(
+                    "warp-rnnt was not installed." "Please consult the installation documentation."
+                )
+                exit(1)
+
+        log_probs = torch.log_softmax(joint_out, dim=-1)
+
+        loss_transducer = self.criterion_transducer(
+            log_probs,
+            target,
+            t_len,
+            u_len,
+            reduction="mean",
+            blank=self.blank_id,
+            fastemit_lambda=self.fastemit_lambda,
+            gather=True,
+        )
+
+        if not self.training and (self.report_cer or self.report_wer):
+            if self.error_calculator is None:
+                from funasr.metrics import ErrorCalculatorTransducer as ErrorCalculator
+
+                self.error_calculator = ErrorCalculator(
+                    self.decoder,
+                    self.joint_network,
+                    self.token_list,
+                    self.sym_space,
+                    self.sym_blank,
+                    report_cer=self.report_cer,
+                    report_wer=self.report_wer,
+                )
+
+            cer_transducer, wer_transducer = self.error_calculator(encoder_out, target, t_len)
+
+            return loss_transducer, cer_transducer, wer_transducer
+
+        return loss_transducer, None, None
+
+    def _calc_ctc_loss(
+        self,
+        encoder_out: torch.Tensor,
+        target: torch.Tensor,
+        t_len: torch.Tensor,
+        u_len: torch.Tensor,
+    ) -> torch.Tensor:
+        """Compute CTC loss.
+
+        Args:
+            encoder_out: Encoder output sequences. (B, T, D_enc)
+            target: Target label ID sequences. (B, L)
+            t_len: Encoder output sequences lengths. (B,)
+            u_len: Target label ID sequences lengths. (B,)
+
+        Return:
+            loss_ctc: CTC loss value.
+
+        """
+        ctc_in = self.ctc_lin(torch.nn.functional.dropout(encoder_out, p=self.ctc_dropout_rate))
+        ctc_in = torch.log_softmax(ctc_in.transpose(0, 1), dim=-1)
+
+        target_mask = target != 0
+        ctc_target = target[target_mask].cpu()
+
+        with torch.backends.cudnn.flags(deterministic=True):
+            loss_ctc = torch.nn.functional.ctc_loss(
+                ctc_in,
+                ctc_target,
+                t_len,
+                u_len,
+                zero_infinity=True,
+                reduction="sum",
+            )
+        loss_ctc /= target.size(0)
+
+        return loss_ctc
+
+    def _calc_lm_loss(
+        self,
+        decoder_out: torch.Tensor,
+        target: torch.Tensor,
+    ) -> torch.Tensor:
+        """Compute LM loss.
+
+        Args:
+            decoder_out: Decoder output sequences. (B, U, D_dec)
+            target: Target label ID sequences. (B, L)
+
+        Return:
+            loss_lm: LM loss value.
+
+        """
+        lm_loss_in = self.lm_lin(decoder_out[:, :-1, :]).view(-1, self.vocab_size)
+        lm_target = target.view(-1).type(torch.int64)
+
+        with torch.no_grad():
+            true_dist = lm_loss_in.clone()
+            true_dist.fill_(self.lm_loss_smoothing / (self.vocab_size - 1))
+
+            # Ignore blank ID (0)
+            ignore = lm_target == 0
+            lm_target = lm_target.masked_fill(ignore, 0)
+
+            true_dist.scatter_(1, lm_target.unsqueeze(1), (1 - self.lm_loss_smoothing))
+
+        loss_lm = torch.nn.functional.kl_div(
+            torch.log_softmax(lm_loss_in, dim=1),
+            true_dist,
+            reduction="none",
+        )
+        loss_lm = loss_lm.masked_fill(ignore.unsqueeze(1), 0).sum() / decoder_out.size(0)
+
+        return loss_lm
+
+    def init_beam_search(
+        self,
+        **kwargs,
+    ):
+
+        # 1. Build ASR model
+        """Init beam search.
+        
+            Args:
+                **kwargs: Additional keyword arguments.
+            """
+        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(
+            length_bonus=LengthBonus(len(token_list)),
+        )
+
+        # 3. Build ngram model
+        # ngram is not supported now
+        ngram = None
+        scorers["ngram"] = ngram
+
+        beam_search = BeamSearchTransducer(
+            self.decoder,
+            self.joint_network,
+            kwargs.get("beam_size", 2),
+            nbest=1,
+        )
+        # 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 inference(
+        self,
+        data_in: list,
+        data_lengths: list = None,
+        key: list = None,
+        tokenizer=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.
+                **kwargs: Additional keyword arguments.
+            """
+        if kwargs.get("batch_size", 1) > 1:
+            raise NotImplementedError("batch decoding is not implemented")
+
+        # init beamsearch
+        is_use_ctc = kwargs.get("decoding_ctc_weight", 0.0) > 0.00001 and self.ctc != None
+        is_use_lm = (
+            kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None
+        )
+        # if self.beam_search is None and (is_use_lm or is_use_ctc):
+        logging.info("enable beam_search")
+        self.init_beam_search(**kwargs)
+        self.nbest = kwargs.get("nbest", 1)
+
+        meta_data = {}
+        # extract fbank feats
+        time1 = time.perf_counter()
+        audio_sample_list = load_audio_text_image_video(
+            data_in, fs=self.frontend.fs, audio_fs=kwargs.get("fs", 16000)
+        )
+        time2 = time.perf_counter()
+        meta_data["load_data"] = f"{time2 - time1:0.3f}"
+        speech, speech_lengths = extract_fbank(
+            audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=self.frontend
+        )
+        time3 = time.perf_counter()
+        meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
+        meta_data["batch_data_time"] = (
+            speech_lengths.sum().item() * self.frontend.frame_shift * self.frontend.lfr_n / 1000
+        )
+
+        speech = speech.to(device=kwargs["device"])
+        speech_lengths = speech_lengths.to(device=kwargs["device"])
+
+        # Encoder
+        encoder_out, encoder_out_lens = self.encode(speech, speech_lengths)
+        if isinstance(encoder_out, tuple):
+            encoder_out = encoder_out[0]
+
+        # c. Passed the encoder result and the beam search
+        nbest_hyps = self.beam_search(encoder_out[0], is_final=True)
+        nbest_hyps = nbest_hyps[: self.nbest]
+
+        results = []
+        b, n, d = encoder_out.size()
+        for i in range(b):
+
+            for nbest_idx, hyp in enumerate(nbest_hyps):
+                ibest_writer = None
+                if kwargs.get("output_dir") is not None:
+                    if not hasattr(self, "writer"):
+                        self.writer = DatadirWriter(kwargs.get("output_dir"))
+                    ibest_writer = self.writer[f"{nbest_idx + 1}best_recog"]
+                # remove sos/eos and get results
+                last_pos = -1
+                if isinstance(hyp.yseq, list):
+                    token_int = hyp.yseq  # [1:last_pos]
+                else:
+                    token_int = hyp.yseq  # [1:last_pos].tolist()
+
+                # remove blank symbol id, which is assumed to be 0
+                token_int = list(
+                    filter(
+                        lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int
+                    )
+                )
+
+                # Change integer-ids to tokens
+                token = tokenizer.ids2tokens(token_int)
+                text = tokenizer.tokens2text(token)
+
+                text_postprocessed, _ = postprocess_utils.sentence_postprocess(token)
+                result_i = {
+                    "key": key[i],
+                    "token": token,
+                    "text": text,
+                    "text_postprocessed": text_postprocessed,
+                }
+                results.append(result_i)
+
+                if ibest_writer is not None:
+                    ibest_writer["token"][key[i]] = " ".join(token)
+                    ibest_writer["text"][key[i]] = text
+                    ibest_writer["text_postprocessed"][key[i]] = text_postprocessed
+
+        return results, meta_data

funasr/models/transducer/rnn_decoder.py

@@ -0,0 +1,402 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import torch
+import random
+import numpy as np
+import torch.nn as nn
+import torch.nn.functional as F
+
+from funasr.register import tables
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+from funasr.models.transformer.utils.nets_utils import to_device
+from funasr.models.language_model.rnn.attentions import initial_att
+
+
+def build_attention_list(
+    eprojs: int,
+    dunits: int,
+    atype: str = "location",
+    num_att: int = 1,
+    num_encs: int = 1,
+    aheads: int = 4,
+    adim: int = 320,
+    awin: int = 5,
+    aconv_chans: int = 10,
+    aconv_filts: int = 100,
+    han_mode: bool = False,
+    han_type=None,
+    han_heads: int = 4,
+    han_dim: int = 320,
+    han_conv_chans: int = -1,
+    han_conv_filts: int = 100,
+    han_win: int = 5,
+):
+
+    """Build attention list.
+    
+        Args:
+            eprojs: TODO.
+            dunits: TODO.
+            atype: TODO.
+            num_att: TODO.
+            num_encs: TODO.
+            aheads: TODO.
+            adim: TODO.
+            awin: TODO.
+            aconv_chans: TODO.
+            aconv_filts: TODO.
+            han_mode: TODO.
+            han_type: TODO.
+            han_heads: TODO.
+            han_dim: Size/dimension parameter.
+            han_conv_chans: TODO.
+            han_conv_filts: TODO.
+            han_win: TODO.
+        """
+    att_list = torch.nn.ModuleList()
+    if num_encs == 1:
+        for i in range(num_att):
+            att = initial_att(
+                atype,
+                eprojs,
+                dunits,
+                aheads,
+                adim,
+                awin,
+                aconv_chans,
+                aconv_filts,
+            )
+            att_list.append(att)
+    elif num_encs > 1:  # no multi-speaker mode
+        if han_mode:
+            att = initial_att(
+                han_type,
+                eprojs,
+                dunits,
+                han_heads,
+                han_dim,
+                han_win,
+                han_conv_chans,
+                han_conv_filts,
+                han_mode=True,
+            )
+            return att
+        else:
+            att_list = torch.nn.ModuleList()
+            for idx in range(num_encs):
+                att = initial_att(
+                    atype[idx],
+                    eprojs,
+                    dunits,
+                    aheads[idx],
+                    adim[idx],
+                    awin[idx],
+                    aconv_chans[idx],
+                    aconv_filts[idx],
+                )
+                att_list.append(att)
+    else:
+        raise ValueError("Number of encoders needs to be more than one. {}".format(num_encs))
+    return att_list
+
+
+@tables.register("decoder_classes", "rnn_decoder")
+class RNNDecoder(nn.Module):
+    def __init__(
+        self,
+        vocab_size: int,
+        encoder_output_size: int,
+        rnn_type: str = "lstm",
+        num_layers: int = 1,
+        hidden_size: int = 320,
+        sampling_probability: float = 0.0,
+        dropout: float = 0.0,
+        context_residual: bool = False,
+        replace_sos: bool = False,
+        num_encs: int = 1,
+        att_conf: dict = None,
+    ):
+        # FIXME(kamo): The parts of num_spk should be refactored more more more
+        """Initialize RNNDecoder.
+        
+            Args:
+                vocab_size: Size/dimension parameter.
+                encoder_output_size: Size/dimension parameter.
+                rnn_type: TODO.
+                num_layers: TODO.
+                hidden_size: Size/dimension parameter.
+                sampling_probability: TODO.
+                dropout: TODO.
+                context_residual: TODO.
+                replace_sos: TODO.
+                num_encs: TODO.
+                att_conf: Configuration dict for att.
+            """
+        if rnn_type not in {"lstm", "gru"}:
+            raise ValueError(f"Not supported: rnn_type={rnn_type}")
+
+        super().__init__()
+        eprojs = encoder_output_size
+        self.dtype = rnn_type
+        self.dunits = hidden_size
+        self.dlayers = num_layers
+        self.context_residual = context_residual
+        self.sos = vocab_size - 1
+        self.eos = vocab_size - 1
+        self.odim = vocab_size
+        self.sampling_probability = sampling_probability
+        self.dropout = dropout
+        self.num_encs = num_encs
+
+        # for multilingual translation
+        self.replace_sos = replace_sos
+
+        self.embed = torch.nn.Embedding(vocab_size, hidden_size)
+        self.dropout_emb = torch.nn.Dropout(p=dropout)
+
+        self.decoder = torch.nn.ModuleList()
+        self.dropout_dec = torch.nn.ModuleList()
+        self.decoder += [
+            (
+                torch.nn.LSTMCell(hidden_size + eprojs, hidden_size)
+                if self.dtype == "lstm"
+                else torch.nn.GRUCell(hidden_size + eprojs, hidden_size)
+            )
+        ]
+        self.dropout_dec += [torch.nn.Dropout(p=dropout)]
+        for _ in range(1, self.dlayers):
+            self.decoder += [
+                (
+                    torch.nn.LSTMCell(hidden_size, hidden_size)
+                    if self.dtype == "lstm"
+                    else torch.nn.GRUCell(hidden_size, hidden_size)
+                )
+            ]
+            self.dropout_dec += [torch.nn.Dropout(p=dropout)]
+            # NOTE: dropout is applied only for the vertical connections
+            # see https://arxiv.org/pdf/1409.2329.pdf
+
+        if context_residual:
+            self.output = torch.nn.Linear(hidden_size + eprojs, vocab_size)
+        else:
+            self.output = torch.nn.Linear(hidden_size, vocab_size)
+
+        self.att_list = build_attention_list(eprojs=eprojs, dunits=hidden_size, **att_conf)
+
+    def zero_state(self, hs_pad):
+        """Zero state.
+        
+            Args:
+                hs_pad: TODO.
+            """
+        return hs_pad.new_zeros(hs_pad.size(0), self.dunits)
+
+    def rnn_forward(self, ey, z_list, c_list, z_prev, c_prev):
+        """Rnn forward.
+        
+            Args:
+                ey: TODO.
+                z_list: TODO.
+                c_list: TODO.
+                z_prev: TODO.
+                c_prev: TODO.
+            """
+        if self.dtype == "lstm":
+            z_list[0], c_list[0] = self.decoder[0](ey, (z_prev[0], c_prev[0]))
+            for i in range(1, self.dlayers):
+                z_list[i], c_list[i] = self.decoder[i](
+                    self.dropout_dec[i - 1](z_list[i - 1]),
+                    (z_prev[i], c_prev[i]),
+                )
+        else:
+            z_list[0] = self.decoder[0](ey, z_prev[0])
+            for i in range(1, self.dlayers):
+                z_list[i] = self.decoder[i](self.dropout_dec[i - 1](z_list[i - 1]), z_prev[i])
+        return z_list, c_list
+
+    def forward(self, hs_pad, hlens, ys_in_pad, ys_in_lens, strm_idx=0):
+        # to support mutiple encoder asr mode, in single encoder mode,
+        # convert torch.Tensor to List of torch.Tensor
+        """Forward pass for training.
+        
+            Args:
+                hs_pad: TODO.
+                hlens: TODO.
+                ys_in_pad: TODO.
+                ys_in_lens: Lengths of ys_in.
+                strm_idx: TODO.
+            """
+        if self.num_encs == 1:
+            hs_pad = [hs_pad]
+            hlens = [hlens]
+
+        # attention index for the attention module
+        # in SPA (speaker parallel attention),
+        # att_idx is used to select attention module. In other cases, it is 0.
+        att_idx = min(strm_idx, len(self.att_list) - 1)
+
+        # hlens should be list of list of integer
+        hlens = [list(map(int, hlens[idx])) for idx in range(self.num_encs)]
+
+        # get dim, length info
+        olength = ys_in_pad.size(1)
+
+        # initialization
+        c_list = [self.zero_state(hs_pad[0])]
+        z_list = [self.zero_state(hs_pad[0])]
+        for _ in range(1, self.dlayers):
+            c_list.append(self.zero_state(hs_pad[0]))
+            z_list.append(self.zero_state(hs_pad[0]))
+        z_all = []
+        if self.num_encs == 1:
+            att_w = None
+            self.att_list[att_idx].reset()  # reset pre-computation of h
+        else:
+            att_w_list = [None] * (self.num_encs + 1)  # atts + han
+            att_c_list = [None] * self.num_encs  # atts
+            for idx in range(self.num_encs + 1):
+                # reset pre-computation of h in atts and han
+                self.att_list[idx].reset()
+
+        # pre-computation of embedding
+        eys = self.dropout_emb(self.embed(ys_in_pad))  # utt x olen x zdim
+
+        # loop for an output sequence
+        for i in range(olength):
+            if self.num_encs == 1:
+                att_c, att_w = self.att_list[att_idx](
+                    hs_pad[0], hlens[0], self.dropout_dec[0](z_list[0]), att_w
+                )
+            else:
+                for idx in range(self.num_encs):
+                    att_c_list[idx], att_w_list[idx] = self.att_list[idx](
+                        hs_pad[idx],
+                        hlens[idx],
+                        self.dropout_dec[0](z_list[0]),
+                        att_w_list[idx],
+                    )
+                hs_pad_han = torch.stack(att_c_list, dim=1)
+                hlens_han = [self.num_encs] * len(ys_in_pad)
+                att_c, att_w_list[self.num_encs] = self.att_list[self.num_encs](
+                    hs_pad_han,
+                    hlens_han,
+                    self.dropout_dec[0](z_list[0]),
+                    att_w_list[self.num_encs],
+                )
+            if i > 0 and random.random() < self.sampling_probability:
+                z_out = self.output(z_all[-1])
+                z_out = np.argmax(z_out.detach().cpu(), axis=1)
+                z_out = self.dropout_emb(self.embed(to_device(self, z_out)))
+                ey = torch.cat((z_out, att_c), dim=1)  # utt x (zdim + hdim)
+            else:
+                # utt x (zdim + hdim)
+                ey = torch.cat((eys[:, i, :], att_c), dim=1)
+            z_list, c_list = self.rnn_forward(ey, z_list, c_list, z_list, c_list)
+            if self.context_residual:
+                z_all.append(
+                    torch.cat((self.dropout_dec[-1](z_list[-1]), att_c), dim=-1)
+                )  # utt x (zdim + hdim)
+            else:
+                z_all.append(self.dropout_dec[-1](z_list[-1]))  # utt x (zdim)
+
+        z_all = torch.stack(z_all, dim=1)
+        z_all = self.output(z_all)
+        z_all.masked_fill_(
+            make_pad_mask(ys_in_lens, z_all, 1),
+            0,
+        )
+        return z_all, ys_in_lens
+
+    def init_state(self, x):
+        # to support mutiple encoder asr mode, in single encoder mode,
+        # convert torch.Tensor to List of torch.Tensor
+        """Init state.
+        
+            Args:
+                x: TODO.
+            """
+        if self.num_encs == 1:
+            x = [x]
+
+        c_list = [self.zero_state(x[0].unsqueeze(0))]
+        z_list = [self.zero_state(x[0].unsqueeze(0))]
+        for _ in range(1, self.dlayers):
+            c_list.append(self.zero_state(x[0].unsqueeze(0)))
+            z_list.append(self.zero_state(x[0].unsqueeze(0)))
+        # TODO(karita): support strm_index for `asr_mix`
+        strm_index = 0
+        att_idx = min(strm_index, len(self.att_list) - 1)
+        if self.num_encs == 1:
+            a = None
+            self.att_list[att_idx].reset()  # reset pre-computation of h
+        else:
+            a = [None] * (self.num_encs + 1)  # atts + han
+            for idx in range(self.num_encs + 1):
+                # reset pre-computation of h in atts and han
+                self.att_list[idx].reset()
+        return dict(
+            c_prev=c_list[:],
+            z_prev=z_list[:],
+            a_prev=a,
+            workspace=(att_idx, z_list, c_list),
+        )
+
+    def score(self, yseq, state, x):
+        # to support mutiple encoder asr mode, in single encoder mode,
+        # convert torch.Tensor to List of torch.Tensor
+        """Score.
+        
+            Args:
+                yseq: TODO.
+                state: TODO.
+                x: TODO.
+            """
+        if self.num_encs == 1:
+            x = [x]
+
+        att_idx, z_list, c_list = state["workspace"]
+        vy = yseq[-1].unsqueeze(0)
+        ey = self.dropout_emb(self.embed(vy))  # utt list (1) x zdim
+        if self.num_encs == 1:
+            att_c, att_w = self.att_list[att_idx](
+                x[0].unsqueeze(0),
+                [x[0].size(0)],
+                self.dropout_dec[0](state["z_prev"][0]),
+                state["a_prev"],
+            )
+        else:
+            att_w = [None] * (self.num_encs + 1)  # atts + han
+            att_c_list = [None] * self.num_encs  # atts
+            for idx in range(self.num_encs):
+                att_c_list[idx], att_w[idx] = self.att_list[idx](
+                    x[idx].unsqueeze(0),
+                    [x[idx].size(0)],
+                    self.dropout_dec[0](state["z_prev"][0]),
+                    state["a_prev"][idx],
+                )
+            h_han = torch.stack(att_c_list, dim=1)
+            att_c, att_w[self.num_encs] = self.att_list[self.num_encs](
+                h_han,
+                [self.num_encs],
+                self.dropout_dec[0](state["z_prev"][0]),
+                state["a_prev"][self.num_encs],
+            )
+        ey = torch.cat((ey, att_c), dim=1)  # utt(1) x (zdim + hdim)
+        z_list, c_list = self.rnn_forward(ey, z_list, c_list, state["z_prev"], state["c_prev"])
+        if self.context_residual:
+            logits = self.output(torch.cat((self.dropout_dec[-1](z_list[-1]), att_c), dim=-1))
+        else:
+            logits = self.output(self.dropout_dec[-1](z_list[-1]))
+        logp = F.log_softmax(logits, dim=1).squeeze(0)
+        return (
+            logp,
+            dict(
+                c_prev=c_list[:],
+                z_prev=z_list[:],
+                a_prev=att_w,
+                workspace=(att_idx, z_list, c_list),
+            ),
+        )

funasr/models/transducer/rnnt_decoder.py

@@ -0,0 +1,264 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import torch
+from typing import List, Optional, Tuple
+
+from funasr.register import tables
+from funasr.models.specaug.specaug import SpecAug
+from funasr.models.transducer.beam_search_transducer import Hypothesis
+
+
+@tables.register("decoder_classes", "rnnt_decoder")
+class RNNTDecoder(torch.nn.Module):
+    """RNN decoder module.
+
+    Args:
+        vocab_size: Vocabulary size.
+        embed_size: Embedding size.
+        hidden_size: Hidden size..
+        rnn_type: Decoder layers type.
+        num_layers: Number of decoder layers.
+        dropout_rate: Dropout rate for decoder layers.
+        embed_dropout_rate: Dropout rate for embedding layer.
+        embed_pad: Embedding padding symbol ID.
+
+    """
+
+    def __init__(
+        self,
+        vocab_size: int,
+        embed_size: int = 256,
+        hidden_size: int = 256,
+        rnn_type: str = "lstm",
+        num_layers: int = 1,
+        dropout_rate: float = 0.0,
+        embed_dropout_rate: float = 0.0,
+        embed_pad: int = 0,
+        use_embed_mask: bool = False,
+    ) -> None:
+        """Construct a RNNDecoder object."""
+        super().__init__()
+
+        if rnn_type not in ("lstm", "gru"):
+            raise ValueError(f"Not supported: rnn_type={rnn_type}")
+
+        self.embed = torch.nn.Embedding(vocab_size, embed_size, padding_idx=embed_pad)
+        self.dropout_embed = torch.nn.Dropout(p=embed_dropout_rate)
+
+        rnn_class = torch.nn.LSTM if rnn_type == "lstm" else torch.nn.GRU
+
+        self.rnn = torch.nn.ModuleList([rnn_class(embed_size, hidden_size, 1, batch_first=True)])
+
+        for _ in range(1, num_layers):
+            self.rnn += [rnn_class(hidden_size, hidden_size, 1, batch_first=True)]
+
+        self.dropout_rnn = torch.nn.ModuleList(
+            [torch.nn.Dropout(p=dropout_rate) for _ in range(num_layers)]
+        )
+
+        self.dlayers = num_layers
+        self.dtype = rnn_type
+
+        self.output_size = hidden_size
+        self.vocab_size = vocab_size
+
+        self.device = next(self.parameters()).device
+        self.score_cache = {}
+
+        self.use_embed_mask = use_embed_mask
+        if self.use_embed_mask:
+            self._embed_mask = SpecAug(
+                time_mask_width_range=3,
+                num_time_mask=4,
+                apply_freq_mask=False,
+                apply_time_warp=False,
+            )
+
+    def forward(
+        self,
+        labels: torch.Tensor,
+        label_lens: torch.Tensor,
+        states: Optional[Tuple[torch.Tensor, Optional[torch.Tensor]]] = None,
+    ) -> torch.Tensor:
+        """Encode source label sequences.
+
+        Args:
+            labels: Label ID sequences. (B, L)
+            states: Decoder hidden states.
+                      ((N, B, D_dec), (N, B, D_dec) or None) or None
+
+        Returns:
+            dec_out: Decoder output sequences. (B, U, D_dec)
+
+        """
+        if states is None:
+            states = self.init_state(labels.size(0))
+
+        dec_embed = self.dropout_embed(self.embed(labels))
+        if self.use_embed_mask and self.training:
+            dec_embed = self._embed_mask(dec_embed, label_lens)[0]
+        dec_out, states = self.rnn_forward(dec_embed, states)
+        return dec_out
+
+    def rnn_forward(
+        self,
+        x: torch.Tensor,
+        state: Tuple[torch.Tensor, Optional[torch.Tensor]],
+    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, Optional[torch.Tensor]]]:
+        """Encode source label sequences.
+
+        Args:
+            x: RNN input sequences. (B, D_emb)
+            state: Decoder hidden states. ((N, B, D_dec), (N, B, D_dec) or None)
+
+        Returns:
+            x: RNN output sequences. (B, D_dec)
+            (h_next, c_next): Decoder hidden states.
+                                (N, B, D_dec), (N, B, D_dec) or None)
+
+        """
+        h_prev, c_prev = state
+        h_next, c_next = self.init_state(x.size(0))
+
+        for layer in range(self.dlayers):
+            if self.dtype == "lstm":
+                x, (h_next[layer : layer + 1], c_next[layer : layer + 1]) = self.rnn[layer](
+                    x, hx=(h_prev[layer : layer + 1], c_prev[layer : layer + 1])
+                )
+            else:
+                x, h_next[layer : layer + 1] = self.rnn[layer](x, hx=h_prev[layer : layer + 1])
+
+            x = self.dropout_rnn[layer](x)
+
+        return x, (h_next, c_next)
+
+    def score(
+        self,
+        label: torch.Tensor,
+        label_sequence: List[int],
+        dec_state: Tuple[torch.Tensor, Optional[torch.Tensor]],
+    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, Optional[torch.Tensor]]]:
+        """One-step forward hypothesis.
+
+        Args:
+            label: Previous label. (1, 1)
+            label_sequence: Current label sequence.
+            dec_state: Previous decoder hidden states.
+                         ((N, 1, D_dec), (N, 1, D_dec) or None)
+
+        Returns:
+            dec_out: Decoder output sequence. (1, D_dec)
+            dec_state: Decoder hidden states.
+                         ((N, 1, D_dec), (N, 1, D_dec) or None)
+
+        """
+        str_labels = "_".join(map(str, label_sequence))
+
+        if str_labels in self.score_cache:
+            dec_out, dec_state = self.score_cache[str_labels]
+        else:
+            dec_embed = self.embed(label)
+            dec_out, dec_state = self.rnn_forward(dec_embed, dec_state)
+
+            self.score_cache[str_labels] = (dec_out, dec_state)
+
+        return dec_out[0], dec_state
+
+    def batch_score(
+        self,
+        hyps: List[Hypothesis],
+    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, Optional[torch.Tensor]]]:
+        """One-step forward hypotheses.
+
+        Args:
+            hyps: Hypotheses.
+
+        Returns:
+            dec_out: Decoder output sequences. (B, D_dec)
+            states: Decoder hidden states. ((N, B, D_dec), (N, B, D_dec) or None)
+
+        """
+        labels = torch.LongTensor([[h.yseq[-1]] for h in hyps], device=self.device)
+        dec_embed = self.embed(labels)
+
+        states = self.create_batch_states([h.dec_state for h in hyps])
+        dec_out, states = self.rnn_forward(dec_embed, states)
+
+        return dec_out.squeeze(1), states
+
+    def set_device(self, device: torch.device) -> None:
+        """Set GPU device to use.
+
+        Args:
+            device: Device ID.
+
+        """
+        self.device = device
+
+    def init_state(self, batch_size: int) -> Tuple[torch.Tensor, Optional[torch.tensor]]:
+        """Initialize decoder states.
+
+        Args:
+            batch_size: Batch size.
+
+        Returns:
+            : Initial decoder hidden states. ((N, B, D_dec), (N, B, D_dec) or None)
+
+        """
+        h_n = torch.zeros(
+            self.dlayers,
+            batch_size,
+            self.output_size,
+            device=self.device,
+        )
+
+        if self.dtype == "lstm":
+            c_n = torch.zeros(
+                self.dlayers,
+                batch_size,
+                self.output_size,
+                device=self.device,
+            )
+
+            return (h_n, c_n)
+
+        return (h_n, None)
+
+    def select_state(
+        self, states: Tuple[torch.Tensor, Optional[torch.Tensor]], idx: int
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Get specified ID state from decoder hidden states.
+
+        Args:
+            states: Decoder hidden states. ((N, B, D_dec), (N, B, D_dec) or None)
+            idx: State ID to extract.
+
+        Returns:
+            : Decoder hidden state for given ID. ((N, 1, D_dec), (N, 1, D_dec) or None)
+
+        """
+        return (
+            states[0][:, idx : idx + 1, :],
+            states[1][:, idx : idx + 1, :] if self.dtype == "lstm" else None,
+        )
+
+    def create_batch_states(
+        self,
+        new_states: List[Tuple[torch.Tensor, Optional[torch.Tensor]]],
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Create decoder hidden states.
+
+        Args:
+            new_states: Decoder hidden states. [N x ((1, D_dec), (1, D_dec) or None)]
+
+        Returns:
+            states: Decoder hidden states. ((N, B, D_dec), (N, B, D_dec) or None)
+
+        """
+        return (
+            torch.cat([s[0] for s in new_states], dim=1),
+            torch.cat([s[1] for s in new_states], dim=1) if self.dtype == "lstm" else None,
+        )