merge main

wenet-e2e · robin1001 · Oct 30, 2023 · Oct 16, 2023 · Oct 20, 2023 · Oct 23, 2023
commit 0b0eea79fe578cecce76169a9f74bc87fc58d9be
diff --git a/wenet/cli/transcribe.py b/wenet/cli/transcribe.py
@@ -14,8 +14,7 @@
 
 import argparse
 
-from wenet.cli.model import Model
-from wenet.cli.paraformer_model import Paraformer
+from wenet.cli.model import load_model
 
 
 def get_args():
@@ -30,19 +29,30 @@ def get_args():
                         ],
                         default='chinese',
                         help='language type')
-    parser.add_argument('--model_dir', default='', help='wenet jit model dirs')
-
+    parser.add_argument('-m',
+                        '--model_dir',
+                        default=None,
+                        help='specify your own model dir')
+    parser.add_argument('-t',
+                        '--show_tokens_info',
+                        action='store_true',
+                        help='whether to output token(word) level information'
+                        ', such times/confidence')
+    parser.add_argument('--align',
+                        action='store_true',
+                        help='force align the input audio and transcript')
+    parser.add_argument('--label', type=str, help='the input label to align')
     args = parser.parse_args()
     return args
 
 
 def main():
     args = get_args()
-    if args.language == 'chinese-paraformer':
-        model = Paraformer(args.model_dir)
+    model = load_model(args.language, args.model_dir)
+    if args.align:
+        result = model.align(args.audio_file, args.label)
     else:
-        model = Model(args.language)
-    result = model.transcribe(args.audio_file)
+        result = model.transcribe(args.audio_file, args.show_tokens_info)
     print(result)