From ac00b7deee093773ee2f42f2694746dfbbd8163f Mon Sep 17 00:00:00 2001
From: 嘉渊 <wangjiaming.wjm@alibaba-inc.com>
Date: 星期四, 15 六月 2023 16:53:26 +0800
Subject: [PATCH] update repo

---
 funasr/build_utils/build_streaming_iterator.py |    6 +++++-
 1 files changed, 5 insertions(+), 1 deletions(-)

diff --git a/funasr/build_utils/build_streaming_iterator.py b/funasr/build_utils/build_streaming_iterator.py
index 732fe09..ad36b4e 100644
--- a/funasr/build_utils/build_streaming_iterator.py
+++ b/funasr/build_utils/build_streaming_iterator.py
@@ -17,6 +17,7 @@
         mc: bool = False,
         dtype: str = np.float32,
         num_workers: int = 1,
+        use_collate_fn: bool = True,
         ngpu: int = 0,
         train: bool=False,
 ) -> DataLoader:
@@ -25,12 +26,15 @@
 
     # preprocess
     if preprocess_args is not None:
+        preprocess_args.task_name = task_name
         preprocess_fn = build_preprocess(preprocess_args, train)
     else:
         preprocess_fn = None
 
     # collate
-    if task_name in ["punc", "lm"]:
+    if not use_collate_fn:
+        collate_fn = None
+    elif task_name in ["punc", "lm"]:
         collate_fn = CommonCollateFn(int_pad_value=0)
     else:
         collate_fn = CommonCollateFn(float_pad_value=0.0, int_pad_value=-1)

--
Gitblit v1.9.1