refactor call args

wenet/cli/punc_model.py

@@ -1,6 +1,7 @@
 import os
 from typing import List
 
+import jieba
 import torch
 from wenet.cli.hub import Hub
 from wenet.paraformer.search import _isAllAlpha
@@ -26,8 +27,6 @@ def split_words(self, text: str):
             self.use_jieba = True
             import logging
 
-            import jieba
-
             # Disable jieba's logger
             logging.getLogger('jieba').disabled = True
             jieba.load_userdict(os.path.join(self.model_dir, 'jieba_usr_dict'))
@@ -95,12 +94,11 @@ def add_punc_batch(self, texts: List[str]):
             result.append(sentence.replace('▁', ' '))
         return result
 
-    def __call__(self, result):
-        text = result['text']
+    def __call__(self, text: str):
         if text != '':
             r = self.add_punc_batch([text])[0]
-            result['text_with_punc'] = r
-        return result
+            return r
+        return ''
 
 
 def load_model(model_dir: str = None,

wenet/cli/transcribe.py

@@ -94,7 +94,7 @@ def main():
         result = model.transcribe(args.audio_file, args.show_tokens_info)
         if args.punc:
             assert punc_model is not None
-            result = punc_model(result)
+            result['text_with_punc'] = punc_model(result['text'])
     print(result)