yhliang
2023-05-12 248408f92fab400d4de451742d085391d9ee7dec
funasr/export/export_model.py
@@ -27,15 +27,13 @@
    ):
        assert check_argument_types()
        self.set_all_random_seed(0)
        if cache_dir is None:
            cache_dir = Path.home() / ".cache" / "export"
        self.cache_dir = Path(cache_dir)
        self.cache_dir = cache_dir
        self.export_config = dict(
            feats_dim=560,
            onnx=False,
        )
        print("output dir: {}".format(self.cache_dir))
        self.onnx = onnx
        self.device = device
        self.quant = quant
@@ -52,7 +50,7 @@
        verbose: bool = False,
    ):
        export_dir = self.cache_dir / tag_name.replace(' ', '-')
        export_dir = self.cache_dir
        os.makedirs(export_dir, exist_ok=True)
        # export encoder1
@@ -174,6 +172,7 @@
        if model_dir.startswith('damo'):
            from modelscope.hub.snapshot_download import snapshot_download
            model_dir = snapshot_download(model_dir, cache_dir=self.cache_dir)
        self.cache_dir = model_dir
        if mode is None:
            import json