update repo

嘉渊

2023-06-15 04528cb2921b236c95588ff611b3e16efa0dba9c

 funasr/build_utils/build_streaming_iterator.py

@@ -5,7 +5,7 @@
from funasr.datasets.iterable_dataset import IterableESPnetDataset
from funasr.datasets.small_datasets.collate_fn import CommonCollateFn
from funasr.datasets.small_datasets.preprocessor import build_preprocess

from funasr.build_utils.build_model_from_file import build_model_from_file

def build_streaming_iterator(
        task_name,
@@ -17,6 +17,8 @@
        mc: bool = False,
        dtype: str = np.float32,
        num_workers: int = 1,
        use_collate_fn: bool = True,
        preprocess_fn=None,
        ngpu: int = 0,
        train: bool=False,
) -> DataLoader:
@@ -24,10 +26,18 @@
    assert check_argument_types()

    # preprocess
    preprocess_fn = build_preprocess(preprocess_args, train)
    if preprocess_fn is not None:
        preprocess_fn = preprocess_fn
    elif 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)