refactor search

wenet-e2e · robin1001 · Oct 30, 2023 · Oct 16, 2023 · Oct 20, 2023 · Oct 23, 2023
commit d264297937a4c9830cc4747ca26a3890c066df77
diff --git a/wenet/paraformer/ali_paraformer/model.py b/wenet/paraformer/ali_paraformer/model.py
@@ -9,8 +9,12 @@
                                                        MultiHeadAttentionCross,
                                                        MultiHeadedAttentionSANM
                                                        )
+from wenet.paraformer.paraformer import Paraformer
+from wenet.paraformer.search import paraformer_beam_search, paraformer_greedy_search
+from wenet.transducer.predictor import PredictorBase
+from wenet.transformer.ctc import CTC
 from wenet.transformer.search import DecodeResult
-from wenet.transformer.encoder import BaseEncoder
+from wenet.transformer.encoder import BaseEncoder, TransformerEncoder
 from wenet.transformer.decoder import TransformerDecoder
 from wenet.transformer.decoder_layer import DecoderLayer
 from wenet.transformer.encoder_layer import TransformerEncoderLayer
@@ -435,7 +439,9 @@ def forward(
         encoder_out_mask: torch.Tensor,
         sematic_embeds: torch.Tensor,
         ys_pad_lens: torch.Tensor,
-    ) -> torch.Tensor:
+        r_ys_in_pad: torch.Tensor = torch.empty(0),
+        reverse_weight: float = 0.0,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """ only for inference now
         """
         ys_pad_mask = make_non_pad_mask(ys_pad_lens).unsqueeze(1)
@@ -450,49 +456,81 @@ def forward(
             x = self.after_norm(x)
         if self.output_layer is not None:
             x = self.output_layer(x)
-        return x
+        return x, torch.tensor(0.0), ys_pad_lens
 
 
 class AliParaformer(torch.nn.Module):
 
-    def __init__(self, encoder: SanmEncoder, decoder: SanmDecoer,
-                 predictor: Predictor) -> None:
+    def __init__(self,
+                 encoder: SanmEncoder,
+                 decoder: SanmDecoer,
+                 predictor: Predictor,
+                 sos: int = -1,
+                 eos: int = -1):
         super().__init__()
-
         self.encoder = encoder
         self.decoder = decoder
         self.predictor = predictor
         self.lfr = LFR()
+        if eos != -1:
+            self.eos = eos
+        if sos != -1:
+            self.sos = sos
 
+    @torch.jit.ignore(drop=True)
     def forward(
-            self, speech: torch.Tensor,
-            speech_lens: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+            self, speech: torch.Tensor, speech_lengths: torch.Tensor,
+            text: torch.Tensor,
+            text_lengths: torch.Tensor) -> Dict[str, Optional[torch.Tensor]]:
+        raise NotImplementedError
 
-        features, features_lens = self.lfr(speech, speech_lens)
-        features_lens = features_lens.to(speech_lens.dtype)
+    @torch.jit.export
+    def forward_paraformer(
+        self,
+        speech: torch.Tensor,
+        speech_lengths: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        features, features_lens = self.lfr(speech, speech_lengths)
+        features_lens = features_lens.to(speech_lengths.dtype)
         # encoder
         encoder_out, encoder_out_mask = self.encoder(features, features_lens)
 
         # cif predictor
         acoustic_embed, token_num, _, _ = self.predictor(encoder_out,
                                                          mask=encoder_out_mask)
-        token_num = token_num.floor().to(speech_lens.dtype)
+        token_num = token_num.floor().to(speech_lengths.dtype)
 
         # decoder
-        decoder_out = self.decoder(encoder_out, encoder_out_mask,
-                                   acoustic_embed, token_num)
+        decoder_out, _, _ = self.decoder(encoder_out, encoder_out_mask,
+                                         acoustic_embed, token_num)
         # decoder_out = decoder_out.log_softmax(dim=-1)
         return decoder_out, token_num
 
-    def decode(self, methods: List[str], speech: torch.Tensor,
-               speech_lens: torch.Tensor,
-               **kwrgs) -> Dict[str, List[DecodeResult]]:
-        assert 'paraformer_greedy_search' in methods
-        results_dict = {}
-        results = []
-        out, out_lens = self.forward(speech, speech_lens)
-        for (i, value) in enumerate(out.argmax(-1).numpy()):
-            results.append(DecodeResult(value[:out_lens[i]]))
-
-        results_dict['paraformer_greedy_search'] = results
-        return results_dict
+    def decode(self,
+               methods: List[str],
+               speech: torch.Tensor,
+               speech_lengths: torch.Tensor,
+               beam_size: int,
+               decoding_chunk_size: int = -1,
+               num_decoding_left_chunks: int = -1,
+               ctc_weight: float = 0,
+               simulate_streaming: bool = False,
+               reverse_weight: float = 0) -> Dict[str, List[DecodeResult]]:
+        decoder_out, decoder_out_lens = self.forward_paraformer(
+            speech, speech_lengths)
+
+        results = {}
+        if 'paraformer_greedy_search' in methods:
+            assert decoder_out is not None
+            assert decoder_out_lens is not None
+            paraformer_greedy_result = paraformer_greedy_search(
+                decoder_out, decoder_out_lens)
+            results['paraformer_greedy_search'] = paraformer_greedy_result
+        if 'paraformer_beam_search' in methods:
+            assert decoder_out is not None
+            assert decoder_out_lens is not None
+            paraformer_beam_result = paraformer_beam_search(
+                decoder_out, decoder_out_lens, beam_size=beam_size)
+            results['paraformer_beam_search'] = paraformer_beam_result
+
+        return results
diff --git a/wenet/paraformer/ali_paraformer/test_infer_jit.py b/wenet/paraformer/ali_paraformer/test_infer_jit.py
@@ -54,8 +54,10 @@ def main():
     feats = feats.unsqueeze(0)
     feats_lens = torch.tensor([feats.size(1)], dtype=torch.int64)
 
-    decode_results = model.decode(['paraformer_greedy_search'], feats,
-                                  feats_lens)
+    decode_results = model.decode(['paraformer_greedy_search'],
+                                  feats,
+                                  feats_lens,
+                                  beam_size=10)
     print("".join([
         char_dict[id]
         for id in decode_results['paraformer_greedy_search'][0].tokens
@@ -66,7 +68,7 @@ def main():
         script_model.save(args.output_file)
 
     model = torch.jit.load(args.output_file)
-    out, token_nums = model.forward(feats, feats_lens)
+    out, token_nums = model.forward_paraformer(feats, feats_lens)
     print("".join([char_dict[id] for id in out.argmax(-1)[0].numpy()]))
     print(token_nums)
 

diff --git a/wenet/paraformer/paraformer.py b/wenet/paraformer/paraformer.py
@@ -20,15 +20,15 @@
 import torch
 
 from wenet.cif.predictor import MAELoss
+from wenet.paraformer.search import paraformer_beam_search, paraformer_greedy_search
 from wenet.transformer.asr_model import ASRModel
 from wenet.transformer.ctc import CTC
 from wenet.transformer.decoder import TransformerDecoder
 from wenet.transformer.encoder import TransformerEncoder
 from wenet.transformer.search import (DecodeResult, ctc_greedy_search,
                                       ctc_prefix_beam_search)
 from wenet.utils.common import (IGNORE_ID, add_sos_eos, th_accuracy)
-from wenet.utils.mask import (make_non_pad_mask, make_pad_mask,
-                              mask_finished_preds, mask_finished_scores)
+from wenet.utils.mask import make_pad_mask
 
 
 class Paraformer(ASRModel):
@@ -164,98 +164,8 @@ def cal_decoder_with_predictor(self, encoder_out, encoder_mask,
                                    sematic_embeds, ys_pad_lens):
         decoder_out, _, _ = self.decoder(encoder_out, encoder_mask,
                                          sematic_embeds, ys_pad_lens)
-        decoder_out = torch.log_softmax(decoder_out, dim=-1)
         return decoder_out, ys_pad_lens
 
-    def paraformer_greedy_search(
-            self, decoder_out: torch.Tensor,
-            decoder_out_lens: torch.Tensor) -> List[DecodeResult]:
-        batch_size = decoder_out.shape[0]
-        maxlen = decoder_out.size(1)
-        topk_prob, topk_index = decoder_out.topk(1, dim=2)
-        topk_index = topk_index.view(batch_size, maxlen)  # (B, maxlen)
-        results = []
-        # TODO(Mddct): scores, times etc
-        for (i, hyp) in enumerate(topk_index.tolist()):
-            r = DecodeResult(hyp[:decoder_out_lens.numpy()[i]])
-            results.append(r)
-        return results
-
-    def _batch_beam_search(
-        self,
-        logit: torch.Tensor,
-        masks: torch.Tensor,
-        beam_size: int = 10,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """ Perform batch beam search
-
-        Args:
-            logit: shape (batch_size, seq_length, vocab_size)
-            masks: shape (batch_size, seq_length)
-            beam_size: beam size
-
-        Returns:
-            indices: shape (batch_size, beam_size, seq_length)
-            log_prob: shape (batch_size, beam_size)
-
-        """
-
-        batch_size, seq_length, vocab_size = logit.shape
-        eos = self.eos
-        masks = ~masks
-        # beam search
-        with torch.no_grad():
-            # b,t,v
-            log_post = torch.nn.functional.log_softmax(logit, dim=-1)
-            # b,k
-            log_prob, indices = log_post[:, 0, :].topk(beam_size, sorted=True)
-            end_flag = torch.eq(masks[:, 0], 1).view(-1, 1)
-            # mask predictor and scores if end
-            log_prob = mask_finished_scores(log_prob, end_flag)
-            indices = mask_finished_preds(indices, end_flag, eos)
-            # b,k,1
-            indices = indices.unsqueeze(-1)
-
-            for i in range(1, seq_length):
-                # b,v
-                scores = mask_finished_scores(log_post[:, i, :], end_flag)
-                # b,v -> b,k,v
-                topk_scores = scores.unsqueeze(1).repeat(1, beam_size, 1)
-                # b,k,1 + b,k,v -> b,k,v
-                top_k_logp = log_prob.unsqueeze(-1) + topk_scores
-
-                # b,k,v -> b,k*v -> b,k
-                log_prob, top_k_index = top_k_logp.view(batch_size,
-                                                        -1).topk(beam_size,
-                                                                 sorted=True)
-
-                index = mask_finished_preds(top_k_index, end_flag, eos)
-
-                indices = torch.cat([indices, index.unsqueeze(-1)], dim=-1)
-
-                end_flag = torch.eq(masks[:, i], 1).view(-1, 1)
-
-            indices = torch.fmod(indices, vocab_size)
-
-        return indices, log_prob
-
-    def paraformer_beam_search(self,
-                               decoder_out: torch.Tensor,
-                               decoder_out_lens: torch.Tensor,
-                               beam_size: int = 10) -> List[DecodeResult]:
-        mask = make_non_pad_mask(decoder_out_lens)
-        indices, _ = self._batch_beam_search(decoder_out,
-                                             mask,
-                                             beam_size=beam_size)
-
-        best_hyps = indices[:, 0, :]
-        results = []
-        # TODO(Mddct): scores, times etc
-        for (i, hyp) in enumerate(best_hyps.tolist()):
-            r = DecodeResult(hyp[:decoder_out_lens.numpy()[i]])
-            results.append(r)
-        return results
-
     def decode(self,
                methods: List[str],
                speech: torch.Tensor,
@@ -298,15 +208,14 @@ def decode(self,
         if 'paraformer_greedy_search' in methods:
             assert decoder_out is not None
             assert decoder_out_lens is not None
-
-            paraformer_greedy_result = self.paraformer_greedy_search(
+            paraformer_greedy_result = paraformer_greedy_search(
                 decoder_out, decoder_out_lens)
             results['paraformer_greedy_search'] = paraformer_greedy_result
         if 'paraformer_beam_search' in methods:
             assert decoder_out is not None
             assert decoder_out_lens is not None
-            paraformer_beam_result = self.paraformer_beam_search(
-                decoder_out, decoder_out_lens)
+            paraformer_beam_result = paraformer_beam_search(
+                decoder_out, decoder_out_lens, beam_size=beam_size)
             results['paraformer_beam_search'] = paraformer_beam_result
 
         return results
diff --git a/wenet/paraformer/search.py b/wenet/paraformer/search.py
@@ -0,0 +1,93 @@
+from typing import List, Tuple
+import torch
+
+from wenet.transformer.search import DecodeResult
+from wenet.utils.mask import make_non_pad_mask, mask_finished_preds, mask_finished_scores
+
+
+def paraformer_greedy_search(
+        decoder_out: torch.Tensor,
+        decoder_out_lens: torch.Tensor) -> List[DecodeResult]:
+    batch_size = decoder_out.shape[0]
+    maxlen = decoder_out.size(1)
+    topk_prob, topk_index = decoder_out.topk(1, dim=2)
+    topk_index = topk_index.view(batch_size, maxlen)  # (B, maxlen)
+    results = []
+    # TODO(Mddct): scores, times etc
+    for (i, hyp) in enumerate(topk_index.tolist()):
+        r = DecodeResult(hyp[:decoder_out_lens.numpy()[i]])
+        results.append(r)
+    return results
+
+
+def paraformer_beam_search(decoder_out: torch.Tensor,
+                           decoder_out_lens: torch.Tensor,
+                           beam_size: int = 10) -> List[DecodeResult]:
+    mask = make_non_pad_mask(decoder_out_lens)
+    indices, _ = _batch_beam_search(decoder_out, mask, beam_size=beam_size)
+
+    best_hyps = indices[:, 0, :]
+    results = []
+    # TODO(Mddct): scores, times etc
+    for (i, hyp) in enumerate(best_hyps.tolist()):
+        r = DecodeResult(hyp[:decoder_out_lens.numpy()[i]])
+        results.append(r)
+    return results
+
+
+def _batch_beam_search(
+    logit: torch.Tensor,
+    masks: torch.Tensor,
+    beam_size: int = 10,
+    eos: int = -1,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """ Perform batch beam search
+
+        Args:
+            logit: shape (batch_size, seq_length, vocab_size)
+            masks: shape (batch_size, seq_length)
+            beam_size: beam size
+
+        Returns:
+            indices: shape (batch_size, beam_size, seq_length)
+            log_prob: shape (batch_size, beam_size)
+
+        """
+
+    batch_size, seq_length, vocab_size = logit.shape
+    masks = ~masks
+    # beam search
+    with torch.no_grad():
+        # b,t,v
+        log_post = torch.nn.functional.log_softmax(logit, dim=-1)
+        # b,k
+        log_prob, indices = log_post[:, 0, :].topk(beam_size, sorted=True)
+        end_flag = torch.eq(masks[:, 0], 1).view(-1, 1)
+        # mask predictor and scores if end
+        log_prob = mask_finished_scores(log_prob, end_flag)
+        indices = mask_finished_preds(indices, end_flag, eos)
+        # b,k,1
+        indices = indices.unsqueeze(-1)
+
+        for i in range(1, seq_length):
+            # b,v
+            scores = mask_finished_scores(log_post[:, i, :], end_flag)
+            # b,v -> b,k,v
+            topk_scores = scores.unsqueeze(1).repeat(1, beam_size, 1)
+            # b,k,1 + b,k,v -> b,k,v
+            top_k_logp = log_prob.unsqueeze(-1) + topk_scores
+
+            # b,k,v -> b,k*v -> b,k
+            log_prob, top_k_index = top_k_logp.view(batch_size,
+                                                    -1).topk(beam_size,
+                                                             sorted=True)
+
+            index = mask_finished_preds(top_k_index, end_flag, eos)
+
+            indices = torch.cat([indices, index.unsqueeze(-1)], dim=-1)
+
+            end_flag = torch.eq(masks[:, i], 1).view(-1, 1)
+
+        indices = torch.fmod(indices, vocab_size)
+
+    return indices, log_prob
diff --git a/wenet/utils/init_model.py b/wenet/utils/init_model.py
@@ -134,7 +134,9 @@ def init_model(configs):
         if isinstance(encoder, SanmEncoder):
             assert isinstance(decoder, SanmDecoer)
             # NOTE(Mddct): only support inference for now
-            model = AliParaformer(encoder, decoder, predictor)
+            model = AliParaformer(encoder=encoder,
+                                  decoder=decoder,
+                                  predictor=predictor)
         else:
             model = Paraformer(vocab_size=vocab_size,
                                encoder=encoder,