[paraformer] add decoding args (#2359)

wenet-e2e · Feb 21, 2024 · 87831da · 87831da
1 parent 7987a20
commit 87831da
Showing 1 changed file with 17 additions and 13 deletions.
diff --git a/wenet/paraformer/paraformer.py b/wenet/paraformer/paraformer.py
@@ -355,19 +355,23 @@ def _forward_paraformer(
             "decoder_out_lens": token_num
         }
 
-    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,
-               context_graph=None,
-               blank_id: int = 0,
-               blank_penalty: float = 0.0) -> Dict[str, List[DecodeResult]]:
+    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,
+        context_graph=None,
+        blank_id: int = 0,
+        blank_penalty: float = 0.0,
+        length_penalty: float = 0.0,
+        infos: Dict[str, List[str]] = None,
+    ) -> Dict[str, List[DecodeResult]]:
         res = self._forward_paraformer(speech, speech_lengths,
                                        decoding_chunk_size,
                                        num_decoding_left_chunks)