Spaces:
Paused
Paused
File size: 1,766 Bytes
0fdb130 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
import os
import openvino.runtime as ov
from setfit import SetFitModel
from setfit.exporters.onnx import export_onnx
def export_to_openvino(
model: SetFitModel,
output_path: str = "model.xml",
) -> None:
"""Export a PyTorch backed SetFit model to OpenVINO Intermediate Representation.
Args:
model_body (`SentenceTransformer`): The model_body from a SetFit model body. This should be a
SentenceTransformer.
model_head (`torch.nn.Module` or `LogisticRegression`): The SetFit model head. This can be either a
dense layer SetFitHead or a Sklearn estimator.
output_path (`str`): The path where will be stored the generated OpenVINO model. At a minimum it needs to contain
the name of the final file.
ignore_ir_version (`bool`): Whether to ignore the IR version used in sklearn. The version is often missmatched
with the transformer models. Setting this to true coerces the versions to be the same. This might
cause errors but in practice works. If this is set to False you need to ensure that the IR versions
align between the transformer and the sklearn onnx representation.
"""
# Load the model and get all of the parts.
OPENVINO_SUPPORTED_OPSET = 13
model.model_body.cpu()
onnx_path = output_path.replace(".xml", ".onnx")
export_onnx(
model.model_body,
model.model_head,
opset=OPENVINO_SUPPORTED_OPSET,
output_path=onnx_path,
ignore_ir_version=True,
use_hummingbird=True,
)
# Save the final model.
ov_model = ov.Core().read_model(onnx_path)
ov.serialize(ov_model, output_path)
os.remove(onnx_path)
|