Add custom model requirements

Signed-off-by: Prabod Rathnayaka <prabod@rathnayaka.me>
JohnSnowLabs · Nov 8, 2024 · 9d24569 · 9d24569
1 parent 43f6875
commit 9d24569
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 19 deletions.
diff --git a/python/sparknlp/internal/__init__.py b/python/sparknlp/internal/__init__.py
@@ -299,6 +299,14 @@ def __init__(self, path, jspark):
             jspark,
         )
 
+class _LLAVAForMultiModalLoader(ExtendedJavaWrapper):
+    def __init__(self, path, jspark, use_openvino=False):
+        super(_LLAVAForMultiModalLoader, self).__init__(
+            "com.johnsnowlabs.nlp.annotators.cv.LLAVAForMultiModal.loadSavedModel",
+            path,
+            jspark,
+            use_openvino
+        )
 
 class _M2M100Loader(ExtendedJavaWrapper):
     def __init__(self, path, jspark, use_openvino=False):

diff --git a/src/main/scala/com/johnsnowlabs/ml/util/LoadExternalModel.scala b/src/main/scala/com/johnsnowlabs/ml/util/LoadExternalModel.scala
@@ -18,6 +18,7 @@ package com.johnsnowlabs.ml.util
 
 import com.johnsnowlabs.ml.tensorflow.sentencepiece.SentencePieceWrapper
 import com.johnsnowlabs.nlp.util.io.{ExternalResource, ReadAs, ResourceHelper}
+import org.glassfish.jersey.internal.inject.Custom
 
 import java.io.File
 import java.nio.file.Paths
@@ -103,30 +104,48 @@ object LoadExternalModel {
 
   }
 
-  def isOpenvinoModel(modelPath: String, isEncoderDecoder: Boolean): Boolean = {
-    if (isEncoderDecoder) {
-      val ovEncoderModelXml = new File(modelPath, s"${Openvino.encoderModel}.xml")
-      val ovEncoderModelBin = new File(modelPath, s"${Openvino.encoderModel}.bin")
-      val ovDecoderModelXml = new File(modelPath, s"${Openvino.decoderModel}.xml")
-      val ovDecoderModelBin = new File(modelPath, s"${Openvino.decoderModel}.bin")
-      val ovDecoderModelWithPastXml = new File(modelPath, s"${Openvino.decoderModelWithPast}.xml")
-      val ovDecoderModelWithPastBin = new File(modelPath, s"${Openvino.decoderModelWithPast}.bin")
-
-      ovEncoderModelXml.exists() && ovEncoderModelBin.exists() &&
-      ovDecoderModelXml.exists() && ovDecoderModelBin.exists() &&
-      ovDecoderModelWithPastXml.exists() && ovDecoderModelWithPastBin.exists()
+  def isOpenvinoModel(
+      modelPath: String,
+      isEncoderDecoder: Boolean,
+      custom: Option[List[String]] = None): Boolean = {
+
+    if (custom.isDefined) {
+      for (model <- custom.get) {
+        val ovModelXml = new File(modelPath, s"${model}.xml")
+        val ovModelBin = new File(modelPath, s"${model}.bin")
+        if (!ovModelXml.exists() || !ovModelBin.exists()) {
+          return false
+        }
+      }
+      true
     } else {
-      val modelXml = new File(modelPath, s"${Openvino.ovModel}.xml")
-      val modelBin = new File(modelPath, s"${Openvino.ovModel}.bin")
-      modelXml.exists() && modelBin.exists()
+      if (isEncoderDecoder) {
+        val ovEncoderModelXml = new File(modelPath, s"${Openvino.encoderModel}.xml")
+        val ovEncoderModelBin = new File(modelPath, s"${Openvino.encoderModel}.bin")
+        val ovDecoderModelXml = new File(modelPath, s"${Openvino.decoderModel}.xml")
+        val ovDecoderModelBin = new File(modelPath, s"${Openvino.decoderModel}.bin")
+        val ovDecoderModelWithPastXml =
+          new File(modelPath, s"${Openvino.decoderModelWithPast}.xml")
+        val ovDecoderModelWithPastBin =
+          new File(modelPath, s"${Openvino.decoderModelWithPast}.bin")
+
+        ovEncoderModelXml.exists() && ovEncoderModelBin.exists() &&
+        ovDecoderModelXml.exists() && ovDecoderModelBin.exists() &&
+        ovDecoderModelWithPastXml.exists() && ovDecoderModelWithPastBin.exists()
+      } else {
+        val modelXml = new File(modelPath, s"${Openvino.ovModel}.xml")
+        val modelBin = new File(modelPath, s"${Openvino.ovModel}.bin")
+        modelXml.exists() && modelBin.exists()
+      }
     }
   }
 
   def detectEngine(
       modelPath: String,
       isEncoderDecoder: Boolean = false,
       withPast: Boolean = false,
-      isDecoder: Boolean = false): String = {
+      isDecoder: Boolean = false,
+      custom: Option[List[String]] = None): String = {
 
     /** Check if the path is correct */
     val f = new File(modelPath)
@@ -146,7 +165,7 @@ object LoadExternalModel {
     val onnxModelExist = isOnnxModel(modelPath, isEncoderDecoder, withPast, isDecoder)
 
     /*Openvino required model files*/
-    val openvinoModelExist = isOpenvinoModel(modelPath, isEncoderDecoder)
+    val openvinoModelExist = isOpenvinoModel(modelPath, isEncoderDecoder, custom)
 
     if (tfSavedModelExist) {
       TensorFlow.name
@@ -176,10 +195,11 @@ object LoadExternalModel {
       path: String,
       isEncoderDecoder: Boolean = false,
       withPast: Boolean = false,
-      isDecoder: Boolean = false): (String, String) = {
+      isDecoder: Boolean = false,
+      custom: Option[List[String]] = None): (String, String) = {
     val localPath: String = ResourceHelper.copyToLocal(path)
 
-    (localPath, detectEngine(localPath, isEncoderDecoder, withPast, isDecoder))
+    (localPath, detectEngine(localPath, isEncoderDecoder, withPast, isDecoder, custom))
   }
 
   def loadTextAsset(assetPath: String, assetName: String): Array[String] = {