jmwang66
2023-02-03 566242c94bbe596fa92c347d5a0cd63b44d54e35
funasr/bin/asr_inference_paraformer_vad_punc.py
@@ -3,6 +3,7 @@
import logging
import sys
import time
import json
from pathlib import Path
from typing import Optional
from typing import Sequence
@@ -478,6 +479,7 @@
    punc_infer_config: Optional[str] = None,
    punc_model_file: Optional[str] = None,
    outputs_dict: Optional[bool] = True,
    param_dict: dict = None,
    **kwargs,
):
    assert check_argument_types()
@@ -543,6 +545,7 @@
    def _forward(data_path_and_name_and_type,
                 raw_inputs: Union[np.ndarray, torch.Tensor] = None,
                 output_dir_v2: Optional[str] = None,
                 param_dict: dict = None,
                 ):
        # 3. Build data-iterator
        if data_path_and_name_and_type is None and raw_inputs is not None:
@@ -635,9 +638,16 @@
                        text_postprocessed, time_stamp_postprocessed, word_lists = postprocessed_result[0], \
                                                                                   postprocessed_result[1], \
                                                                                   postprocessed_result[2]
                        text_postprocessed_punc, punc_id_list = text2punc(word_lists, 20)
                        text_postprocessed_punc_time_stamp = "predictions: {}  time_stamp: {}".format(
                            text_postprocessed_punc, time_stamp_postprocessed)
                        if len(word_lists) > 0:
                            text_postprocessed_punc, punc_id_list = text2punc(word_lists, 20)
                            text_postprocessed_punc_time_stamp = json.dumps({"predictions": text_postprocessed_punc,
                                                                             "time_stamp": time_stamp_postprocessed},
                                                                            ensure_ascii=False)
                        else:
                            text_postprocessed_punc = ""
                            punc_id_list = []
                            text_postprocessed_punc_time_stamp = ""
                    else:
                        text_postprocessed = ""
                        time_stamp_postprocessed = ""