| import os | |
| import json | |
| import shutil | |
| from optimum.exporters.onnx import main_export | |
| import onnx | |
| from onnxconverter_common import float16 | |
| import onnxruntime as rt | |
| from onnxruntime.tools.onnx_model_utils import * | |
| from onnxruntime.quantization import quantize_dynamic, QuantType | |
| with open('conversion_config.json') as json_file: | |
| conversion_config = json.load(json_file) | |
| model_id = conversion_config["model_id"] | |
| number_of_generated_embeddings = conversion_config["number_of_generated_embeddings"] | |
| precision_to_filename_map = conversion_config["precision_to_filename_map"] | |
| opset = conversion_config["opset"] | |
| IR = conversion_config["IR"] | |
| op = onnx.OperatorSetIdProto() | |
| op.version = opset | |
| if not os.path.exists("onnx"): | |
| os.makedirs("onnx") | |
| print("Exporting the main model version") | |
| main_export(model_name_or_path=model_id, output="./", opset=opset, trust_remote_code=True, task="feature-extraction", dtype="fp32") | |
| if "fp32" in precision_to_filename_map: | |
| print("Exporting the fp32 onnx file...") | |
| shutil.copyfile('model.onnx', precision_to_filename_map["fp32"]) | |
| print("Done\n\n") | |
| if "int8" in precision_to_filename_map: | |
| print("Quantizing fp32 model to int8...") | |
| quantize_dynamic("model.onnx", precision_to_filename_map["int8"], weight_type=QuantType.QInt8) | |
| print("Done\n\n") | |
| if "uint8" in precision_to_filename_map: | |
| print("Quantizing fp32 model to uint8...") | |
| quantize_dynamic("model.onnx", precision_to_filename_map["uint8"], weight_type=QuantType.QUInt8) | |
| print("Done\n\n") | |
| os.remove("model.onnx") | |