雾聪
2024-03-29 9ba0dbd98bf69c830dfcfde8f109a400cb65e4e5
funasr/auto/auto_model.py
@@ -157,13 +157,13 @@
        kwargs["device"] = device
        torch.set_num_threads(kwargs.get("ncpu", 4))
        
        # build tokenizer
        tokenizer = kwargs.get("tokenizer", None)
        if tokenizer is not None:
            tokenizer_class = tables.tokenizer_classes.get(tokenizer)
            tokenizer = tokenizer_class(**kwargs["tokenizer_conf"])
            tokenizer_conf = kwargs.get("tokenizer_conf", {})
            tokenizer = tokenizer_class(**tokenizer_conf)
            kwargs["tokenizer"] = tokenizer
            kwargs["token_list"] = tokenizer.token_list if hasattr(tokenizer, "token_list") else None
@@ -495,11 +495,19 @@
                export_dir = export_utils.export_onnx(
                                        model=model,
                                        data_in=data_list,
                                        quantize=quantize,
                                        fallback_num=fallback_num,
                                        calib_num=calib_num,
                                        opset_version=opset_version,
                                        **kwargs)
            else:
                export_dir = export_utils.export_torchscripts(
                                        model=model,
                                        data_in=data_list,
                                        quantize=quantize,
                                        fallback_num=fallback_num,
                                        calib_num=calib_num,
                                        opset_version=opset_version,
                                        **kwargs)
        return export_dir