feat: OpenVINO acceleration for embeddings in transformer backend (#2190

) OpenVINO acceleration for embeddings New argument type: OVModelForFeatureExtraction
mudler · Apr 30, 2024 · e38610e · e38610e
1 parent 3754f15
commit e38610e
Showing 1 changed file with 15 additions and 0 deletions.
diff --git a/backend/python/transformers/transformers_server.py b/backend/python/transformers/transformers_server.py
@@ -153,6 +153,21 @@ def LoadModel(self, request, context):
                                                                 ov_config={"PERFORMANCE_HINT": "CUMULATIVE_THROUGHPUT"}, 
                                                                 device=device_map)
                 self.OV = True
+            elif request.Type == "OVModelForFeatureExtraction":
+                from optimum.intel.openvino import OVModelForFeatureExtraction
+                from openvino.runtime import Core
+
+                if "GPU" in Core().available_devices:
+                    device_map="GPU"
+                else:
+                    device_map="CPU"
+                self.model = OVModelForFeatureExtraction.from_pretrained(model_name, 
+                                                                compile=True,
+                                                                trust_remote_code=request.TrustRemoteCode,
+                                                                ov_config={"PERFORMANCE_HINT": "CUMULATIVE_THROUGHPUT"}, 
+                                                                export=True,
+                                                                device=device_map)
+                self.OV = True
             else:
                 self.model = AutoModel.from_pretrained(model_name, 
                                                        trust_remote_code=request.TrustRemoteCode,