shixian.shi
2023-03-10 7ce644eefdee3315012fd9eab691a94cc6014e65
timestamp pipeline (alone)
1个文件已修改
18 ■■■■ 已修改文件
funasr/bin/tp_inference.py 18 ●●●● 补丁 | 查看 | 原始文档 | blame | 历史
funasr/bin/tp_inference.py
@@ -110,7 +110,7 @@
            timestamp_infer_config, timestamp_model_file, device
        )
        if 'cuda' in device:
            tp_model = tp_model.cuda()
            tp_model = tp_model.cuda()  # force model to cuda
        frontend = None
        if tp_train_args.frontend is not None:
@@ -299,9 +299,6 @@
            inference=True,
        )
        finish_count = 0
        file_count = 1
        tp_result_list = []
        for keys, batch in loader:
            assert isinstance(batch, dict), type(batch)
@@ -321,7 +318,6 @@
                ts_str, ts_list = time_stamp_lfr6_advance(us_alphas[batch_id], us_cif_peak[batch_id], token)
                logging.warning(ts_str)
                item = {'key': key, 'value': ts_str, 'timestamp':ts_list}
                # tp_result_list.append({'text':"".join([i for i in token if i != '<sil>']), 'timestamp': ts_list})
                tp_result_list.append(item)
        return tp_result_list
@@ -407,6 +403,18 @@
        default=1,
        help="The batch size for inference",
    )
    group.add_argument(
        "--seg_dict_file",
        type=str,
        default=None,
        help="The batch size for inference",
    )
    group.add_argument(
        "--split_with_space",
        type=bool,
        default=False,
        help="The batch size for inference",
    )
    return parser