openvino

export.py

@@ -91,11 +91,20 @@ def convert_onnx():
     torch.onnx.export(traced_vision_model, inputs.data['pixel_values'], "vision.onnx")
     torch.onnx.export(traced_text_model, (inputs.data['input_ids'], inputs.data['input_ids']), "text.onnx")
 
-convert_onnx()
+# convert_onnx()
 
-# ov_model = ov.convert_model(traced_model, example_input=inputs.data['pixel_values'])
-# ov.save_model(ov_model, "openvino/vision.xml")
 
+def convert_openvino():
+    import openvino as ov
+    ov_vision_model = ov.convert_model(traced_vision_model, example_input=inputs.data['pixel_values'])
+    ov.save_model(ov_vision_model, "openvino/vision.xml")
+
+    ov_text_model = ov.convert_model(traced_text_model, example_input=(inputs.data['input_ids'], inputs.data['attention_mask']))
+    ov.save_model(ov_text_model, "openvino/text.xml")
+
+convert_openvino()
+
+# convert_coreml()
 
 # Create an ONNX Runtime session
 # session = ort.InferenceSession(model_path)

requirements.txt

@@ -5,4 +5,5 @@ transformers
 coremltools
 torch
 huggingface
-requests
+requests
+openvino