Added Openvino support for python API

JohnSnowLabs · Aug 7, 2024 · fcdbb6c · fcdbb6c
1 parent a84e2f7
commit fcdbb6c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/python/sparknlp/annotator/embeddings/nomic_embeddings.py b/python/sparknlp/annotator/embeddings/nomic_embeddings.py
@@ -139,7 +139,7 @@ def __init__(self, classname="com.johnsnowlabs.nlp.embeddings.NomicEmbeddings",
         self._setDefault(dimension=768, batchSize=8, maxSentenceLength=512, caseSensitive=False, )
 
     @staticmethod
-    def loadSavedModel(folder, spark_session):
+    def loadSavedModel(folder, spark_session, use_openvino=False):
         """Loads a locally saved model.
 
         Parameters
@@ -155,7 +155,7 @@ def loadSavedModel(folder, spark_session):
             The restored model
         """
         from sparknlp.internal import _NomicLoader
-        jModel = _NomicLoader(folder, spark_session._jsparkSession)._java_obj
+        jModel = _NomicLoader(folder, spark_session._jsparkSession, use_openvino)._java_obj
         return NomicEmbeddings(java_model=jModel)
 
     @staticmethod

diff --git a/python/sparknlp/internal/__init__.py b/python/sparknlp/internal/__init__.py
@@ -379,8 +379,8 @@ def __init__(self, path, jspark, useCache):
         )
 
 class _NomicLoader(ExtendedJavaWrapper):
-    def __init__(self, path, jspark):
-        super(_NomicLoader, self).__init__("com.johnsnowlabs.nlp.embeddings.NomicEmbeddings.loadSavedModel", path, jspark)
+    def __init__(self, path, jspark, use_openvino=False):
+        super(_NomicLoader, self).__init__("com.johnsnowlabs.nlp.embeddings.NomicEmbeddings.loadSavedModel", path, jspark, use_openvino)
 
 class _USELoader(ExtendedJavaWrapper):
     def __init__(self, path, jspark, loadsp):