嘉渊 2 лет назад
Родитель
Сommit
ac00b7deee
1 измененных файлов с 9 добавлено и 5 удалено
  1. 9 5
      funasr/build_utils/build_diar_model.py

+ 9 - 5
funasr/build_utils/build_diar_model.py

@@ -178,14 +178,18 @@ class_choices_list = [
 
 def build_diar_model(args):
     # token_list
-    if args.token_list is not None:
-        with open(args.token_list) as f:
+    if isinstance(args.token_list, str):
+        with open(args.token_list, encoding="utf-8") as f:
             token_list = [line.rstrip() for line in f]
+
+        # Overwriting token_list to keep it as "portable".
         args.token_list = list(token_list)
-        vocab_size = len(token_list)
-        logging.info(f"Vocabulary size: {vocab_size}")
+    elif isinstance(args.token_list, (tuple, list)):
+        token_list = list(args.token_list)
     else:
-        vocab_size = None
+        raise RuntimeError("token_list must be str or list")
+    vocab_size = len(token_list)
+    logging.info(f"Vocabulary size: {vocab_size}")
 
     # frontend
     if args.input_size is None: