From 9d2456936be826df95288d0a0798ca91379e4de6 Mon Sep 17 00:00:00 2001
From: Prabod Rathnayaka <prabod@rathnayaka.me>
Date: Fri, 8 Nov 2024 06:33:37 +0000
Subject: [PATCH] Add custom model requirements

Signed-off-by: Prabod Rathnayaka <prabod@rathnayaka.me>
---
 python/sparknlp/internal/__init__.py          |  8 +++
 .../ml/util/LoadExternalModel.scala           | 58 +++++++++++++------
 2 files changed, 47 insertions(+), 19 deletions(-)

diff --git a/python/sparknlp/internal/__init__.py b/python/sparknlp/internal/__init__.py
index eec3544dc41c6f..5953056d2f3449 100644
--- 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
index cd0761f0f9daa3..1b48fdf2c7b6d5 100644
--- 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,22 +104,39 @@ 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()
+      }
     }
   }
 
@@ -126,7 +144,8 @@ object LoadExternalModel {
       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] = {