merge main

wenet-e2e · robin1001 · Oct 30, 2023 · Oct 16, 2023 · Oct 20, 2023 · Oct 23, 2023
commit 30b5677652a119230ce87bed68a34e140ba77c14
diff --git a/wenet/paraformer/ali_paraformer/assets/config.yaml b/wenet/paraformer/ali_paraformer/assets/config.yaml
@@ -17,6 +17,7 @@ encoder_conf:
 input_dim: 80
 output_dim: 8404
 paraformer: true
+is_json_cmvn: True
 # decoder related
 decoder: SanmDecoder
 decoder_conf:
@@ -43,9 +44,3 @@ cif_predictor_conf:
     tail_threshold: 0.45
     cnn_groups: 1
     residual: false
-
-    # smooth_factor2: 0.25
-    # noise_threshold2: 0.01
-    # upsample_times: 3
-    # use_cif1_cnn: false  # TODO: support in the future, has no effect on model wer (timestamp related)
-    # upsample_type: cnn_blstm # TODO: support in the future, has no effect on model wer (timestamp related)
diff --git a/...raformer/ali_paraformer/test_infer_jit.py → ...t/paraformer/ali_paraformer/export_jit.py b/...raformer/ali_paraformer/test_infer_jit.py → ...t/paraformer/ali_paraformer/export_jit.py
@@ -21,7 +21,6 @@ def get_args():
                         required=True,
                         help='cmvn file of paraformer in wenet style')
     parser.add_argument('--dict', required=True, help='dict file')
-    parser.add_argument('--wav', required=True, help='wav file')
     parser.add_argument('--output_file', default=None, help='output file')
     args = parser.parse_args()
     return args
@@ -41,37 +40,10 @@ def main():
     load_checkpoint(model, args.ali_paraformer)
     model.eval()
 
-    waveform, sample_rate = torchaudio.load(args.wav)
-    assert sample_rate == 16000
-    waveform = waveform * (1 << 15)
-    waveform = waveform.to(torch.float)
-    feats = kaldi.fbank(waveform,
-                        num_mel_bins=80,
-                        frame_length=25,
-                        frame_shift=10,
-                        energy_floor=0.0,
-                        sample_frequency=sample_rate)
-    feats = feats.unsqueeze(0)
-    feats_lens = torch.tensor([feats.size(1)], dtype=torch.int64)
-
-    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
-    ]))
-
     if args.output_file:
         script_model = torch.jit.script(model)
         script_model.save(args.output_file)
 
-    model = torch.jit.load(args.output_file)
-    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)
-
 
 if __name__ == "__main__":
 

diff --git a/wenet/paraformer/ali_paraformer/model.py b/wenet/paraformer/ali_paraformer/model.py
@@ -461,21 +461,15 @@ def forward(
 
 class AliParaformer(torch.nn.Module):
 
-    def __init__(self,
-                 encoder: SanmEncoder,
-                 decoder: SanmDecoer,
-                 predictor: Predictor,
-                 sos: int = -1,
-                 eos: int = -1):
+    def __init__(self, encoder: SanmEncoder, decoder: SanmDecoer,
+                 predictor: Predictor):
         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
+        self.sos = 1
+        self.eos = 2
 
     @torch.jit.ignore(drop=True)
     def forward(
@@ -530,7 +524,10 @@ def decode(self,
             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)
+                decoder_out,
+                decoder_out_lens,
+                beam_size=beam_size,
+                eos=self.eos)
             results['paraformer_beam_search'] = paraformer_beam_result
 
         return results
diff --git a/wenet/paraformer/paraformer.py b/wenet/paraformer/paraformer.py
@@ -215,7 +215,10 @@ def decode(self,
             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)
+                decoder_out,
+                decoder_out_lens,
+                beam_size=beam_size,
+                eos=self.eos)
             results['paraformer_beam_search'] = paraformer_beam_result
 
         return results
diff --git a/wenet/paraformer/search.py b/wenet/paraformer/search.py
@@ -22,9 +22,13 @@ def paraformer_greedy_search(
 
 def paraformer_beam_search(decoder_out: torch.Tensor,
                            decoder_out_lens: torch.Tensor,
-                           beam_size: int = 10) -> List[DecodeResult]:
+                           beam_size: int = 10,
+                           eos: int = -1) -> List[DecodeResult]:
     mask = make_non_pad_mask(decoder_out_lens)
-    indices, _ = _batch_beam_search(decoder_out, mask, beam_size=beam_size)
+    indices, _ = _batch_beam_search(decoder_out,
+                                    mask,
+                                    beam_size=beam_size,
+                                    eos=eos)
 
     best_hyps = indices[:, 0, :]
     results = []

diff --git a/wenet/utils/init_model.py b/wenet/utils/init_model.py
@@ -134,9 +134,11 @@ def init_model(configs):
         if isinstance(encoder, SanmEncoder):
             assert isinstance(decoder, SanmDecoer)
             # NOTE(Mddct): only support inference for now
-            model = AliParaformer(encoder=encoder,
-                                  decoder=decoder,
-                                  predictor=predictor)
+            model = AliParaformer(
+                encoder=encoder,
+                decoder=decoder,
+                predictor=predictor,
+            )
         else:
             model = Paraformer(vocab_size=vocab_size,
                                encoder=encoder,