FunASR/funasr/bin/export.py
zhifu gao a7d7a0f3a2
Dev gzf (#1467)
* qwenaudio qwenaudiochat

* qwenaudio qwenaudiochat

* whisper

* whisper

* llm

* llm

* llm

* llm

* llm

* llm

* llm

* llm

* export onnx

* export onnx

* export onnx

* dingding

* dingding

* llm

* doc

* onnx
2024-03-11 19:24:44 +08:00

42 lines
1.3 KiB
Python

import os
import hydra
import logging
from omegaconf import DictConfig, OmegaConf, ListConfig
from funasr.auto.auto_model import AutoModel
@hydra.main(config_name=None, version_base=None)
def main_hydra(cfg: DictConfig):
def to_plain_list(cfg_item):
if isinstance(cfg_item, ListConfig):
return OmegaConf.to_container(cfg_item, resolve=True)
elif isinstance(cfg_item, DictConfig):
return {k: to_plain_list(v) for k, v in cfg_item.items()}
else:
return cfg_item
kwargs = to_plain_list(cfg)
log_level = getattr(logging, kwargs.get("log_level", "INFO").upper())
logging.basicConfig(level=log_level)
if kwargs.get("debug", False):
import pdb; pdb.set_trace()
if "device" not in kwargs:
kwargs["device"] = "cpu"
model = AutoModel(**kwargs)
res = model.export(input=kwargs.get("input", None),
type=kwargs.get("type", "onnx"),
quantize=kwargs.get("quantize", False),
fallback_num=kwargs.get("fallback-num", 5),
calib_num=kwargs.get("calib_num", 100),
opset_version=kwargs.get("opset_version", 14),
)
print(res)
if __name__ == '__main__':
main_hydra()