Introducing onnx and OpenVino support to Missing Annotators

src/main/scala/com/johnsnowlabs/ml/ai/AlbertClassification.scala

@@ -17,13 +17,16 @@
 package com.johnsnowlabs.ml.ai
 
 import ai.onnxruntime.OnnxTensor
+import com.johnsnowlabs.ml.ai.util.PrepareEmbeddings
 import com.johnsnowlabs.ml.onnx.{OnnxSession, OnnxWrapper}
+import com.johnsnowlabs.ml.openvino.OpenvinoWrapper
 import com.johnsnowlabs.ml.tensorflow.sentencepiece.{SentencePieceWrapper, SentencepieceEncoder}
 import com.johnsnowlabs.ml.tensorflow.sign.{ModelSignatureConstants, ModelSignatureManager}
 import com.johnsnowlabs.ml.tensorflow.{TensorResources, TensorflowWrapper}
-import com.johnsnowlabs.ml.util.{ONNX, TensorFlow}
+import com.johnsnowlabs.ml.util.{ONNX, Openvino, TensorFlow}
 import com.johnsnowlabs.nlp.annotators.common._
 import com.johnsnowlabs.nlp.{ActivationFunction, Annotation}
+import org.intel.openvino.Tensor
 import org.tensorflow.ndarray.buffer.IntDataBuffer
 import org.slf4j.{Logger, LoggerFactory}
 
@@ -43,6 +46,7 @@ import scala.collection.JavaConverters._
 private[johnsnowlabs] class AlbertClassification(
     val tensorflowWrapper: Option[TensorflowWrapper],
     val onnxWrapper: Option[OnnxWrapper],
+    val openvinoWrapper: Option[OpenvinoWrapper],
     val spp: SentencePieceWrapper,
     configProtoBytes: Option[Array[Byte]] = None,
     tags: Map[String, Int],
@@ -56,6 +60,7 @@ private[johnsnowlabs] class AlbertClassification(
   val detectedEngine: String =
     if (tensorflowWrapper.isDefined) TensorFlow.name
     else if (onnxWrapper.isDefined) ONNX.name
+    else if (openvinoWrapper.isDefined) Openvino.name
     else TensorFlow.name
   private val onnxSessionOptions: Map[String, String] = new OnnxSession().getSessionOptions
 
@@ -69,6 +74,8 @@ private[johnsnowlabs] class AlbertClassification(
 
   protected val logger: Logger = LoggerFactory.getLogger("AlbertClassification")
 
+
+
   def tokenizeWithAlignment(
       sentences: Seq[TokenizedSentence],
       maxSeqLength: Int,
@@ -112,7 +119,8 @@ private[johnsnowlabs] class AlbertClassification(
 
     val rawScores = detectedEngine match {
       case ONNX.name => getRawScoresWithOnnx(batch, maxSentenceLength, sequence = true)
-      case _ => getRawScoresWithTF(batch, maxSentenceLength)
+      case Openvino.name => getRawScoresWithOv(batch, maxSentenceLength)
+      case TensorFlow.name => getRawScoresWithTF(batch, maxSentenceLength)
     }
 
     val dim = rawScores.length / (batchLength * maxSentenceLength)
@@ -132,7 +140,8 @@ private[johnsnowlabs] class AlbertClassification(
 
     val rawScores = detectedEngine match {
       case ONNX.name => getRawScoresWithOnnx(batch, maxSentenceLength, sequence = true)
-      case _ => getRawScoresWithTF(batch, maxSentenceLength)
+      case TensorFlow.name => getRawScoresWithTF(batch, maxSentenceLength)
+      case Openvino.name => getRawScoresWithOv(batch, maxSentenceLength)
     }
 
     val dim = rawScores.length / batchLength
@@ -206,6 +215,42 @@ private[johnsnowlabs] class AlbertClassification(
     rawScores
   }
 
+  private def getRawScoresWithOv(
+                                    batch: Seq[Array[Int]],
+                                    maxSentenceLength: Int
+                                    ): Array[Float] = {
+
+
+
+    val batchLength = batch.length
+    val shape = Array(batchLength, maxSentenceLength)
+    val (tokenTensors, maskTensors) =
+      PrepareEmbeddings.prepareOvLongBatchTensors(batch, maxSentenceLength, batchLength)
+    val segmentTensors = new Tensor(shape, Array.fill(batchLength * maxSentenceLength)(0L))
+
+    val inferRequest = openvinoWrapper.get.getCompiledModel().create_infer_request()
+    inferRequest.set_tensor("input_ids", tokenTensors)
+    inferRequest.set_tensor("attention_mask", maskTensors)
+    inferRequest.set_tensor("token_type_ids", segmentTensors)
+
+    inferRequest.infer()
+
+    try {
+      try {
+        inferRequest
+          .get_tensor("logits")
+          .data()
+      }
+    } catch {
+      case e: Exception =>
+        // Log the exception as a warning
+        logger.warn("Exception in getRawScoresWithOnnx", e)
+        // Rethrow the exception to propagate it further
+        throw e
+    }
+
+  }
+
   private def getRawScoresWithOnnx(
       batch: Seq[Array[Int]],
       maxSentenceLength: Int,
@@ -258,6 +303,9 @@ private[johnsnowlabs] class AlbertClassification(
     }
   }
 
+
+
+
   def tagZeroShotSequence(
       batch: Seq[Array[Int]],
       entailmentId: Int,
@@ -269,7 +317,8 @@ private[johnsnowlabs] class AlbertClassification(
     val maxSentenceLength = batch.map(encodedSentence => encodedSentence.length).max
     val (startLogits, endLogits) = detectedEngine match {
       case ONNX.name => computeLogitsWithOnnx(batch, maxSentenceLength)
-      case _ => computeLogitsWithTF(batch, maxSentenceLength)
+      case TensorFlow.name => computeLogitsWithTF(batch, maxSentenceLength)
+      case Openvino.name => computeLogitsWithOv(batch, maxSentenceLength)
     }
 
     val endDim = endLogits.length / batchLength
@@ -357,8 +406,8 @@ private[johnsnowlabs] class AlbertClassification(
   }
 
   private def computeLogitsWithOnnx(
-      batch: Seq[Array[Int]],
-      maxSentenceLength: Int): (Array[Float], Array[Float]) = {
+                                     batch: Seq[Array[Int]],
+                                     maxSentenceLength: Int): (Array[Float], Array[Float]) = {
     // [nb of encoded sentences , maxSentenceLength]
     val (runner, env) = onnxWrapper.get.getSession(onnxSessionOptions)
 
@@ -410,6 +459,45 @@ private[johnsnowlabs] class AlbertClassification(
     }
   }
 
+
+
+  private def computeLogitsWithOv(
+                                     batch: Seq[Array[Int]],
+                                     maxSentenceLength: Int): (Array[Float], Array[Float]) = {
+    // [nb of encoded sentences , maxSentenceLength]
+
+    val batchLength = batch.length
+    val shape = Array(batchLength, maxSentenceLength)
+    val (tokenTensors, maskTensors) =
+      PrepareEmbeddings.prepareOvLongBatchTensors(batch, maxSentenceLength, batchLength)
+    val segmentTensors = new Tensor(shape, Array.fill(batchLength * maxSentenceLength)(0L))
+
+    val inferRequest = openvinoWrapper.get.getCompiledModel().create_infer_request()
+    inferRequest.set_tensor("input_ids", tokenTensors)
+    inferRequest.set_tensor("attention_mask", maskTensors)
+    inferRequest.set_tensor("token_type_ids", segmentTensors)
+
+    inferRequest.infer()
+
+    try {
+      try {
+        val startLogits =  inferRequest
+          .get_tensor("start_logits")
+          .data()
+        val endLogits = inferRequest
+          .get_tensor("end_logits")
+          .data()
+
+        (startLogits.slice(1, startLogits.length), endLogits.slice(1, endLogits.length))
+      }
+    } catch {
+      case e: Exception =>
+        // Log the exception as a warning
+        logger.warn("Exception in getRawScoresWithOnnx", e)
+        // Rethrow the exception to propagate it further
+        throw e
+    }
+  }
   def findIndexedToken(
       tokenizedSentences: Seq[TokenizedSentence],
       sentence: (WordpieceTokenizedSentence, Int),