Introducing onnx and OpenVino support to Missing Annotators (#14359)

* adding onnx API * Adding example notebook * Introducing ONNX support to DistilBertForZeroShotClassification * adding DistilBertForZeroShotClaissifcation onnx notebook * adding onnx support to RobertaForZeroShotClassificaiton * adding onnx support to XlmRobertaForZeroShotClassification * Create HuggingFace_ONNX_in_Spark_NLP_XlmRoBertaForZeroShotClassification.ipynb * adding ONNX support to BartTransformer * Create HuggingFace_ONNX_in_Spark_NLP_Bart.ipynb * adding onnx support to GPT2 * Refactoring GPT2 code * adding gpt 2 notebook * adding onnx support to Wav2Vec2 * adding onnx support to hubert * adding ONNX support to BartForZeroShotClassificaiton * adding ONNX support to DebertaForZeroShotClassificaiton * adding openVino support to Whisper * Create HuggingFace_OpenVINO_in_Spark_NLP_Whisper.ipynb * adding OpenVino Support to AlbertForXXX * changing unit tests back to origanl model paths * adding BartForZeroShotClassiifcation * adding openvino suppor tto BertForXXX classificaition
JohnSnowLabs · Sep 1, 2024 · f47ee50 · f47ee50
1 parent 276a84b
commit f47ee50
Show file tree

Hide file tree

Showing 51 changed files with 36,644 additions and 1,369 deletions.
diff --git a/...thon/transformers/HuggingFace_ONNX_in_Spark_NLP_XlmRoBertaForZeroShotClassification.ipynb b/...thon/transformers/HuggingFace_ONNX_in_Spark_NLP_XlmRoBertaForZeroShotClassification.ipynb
diff --git a/examples/python/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_Bart.ipynb b/examples/python/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_Bart.ipynb
diff --git a/examples/python/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_BartForZeroShot.ipynb b/examples/python/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_BartForZeroShot.ipynb
diff --git a/...on/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_DeBertaForZeroShotClassification.ipynb b/...on/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_DeBertaForZeroShotClassification.ipynb
diff --git a/...transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_DistilBertForZeroShotClassification.ipynb b/...transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_DistilBertForZeroShotClassification.ipynb
diff --git a/examples/python/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_GPT2.ipynb b/examples/python/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_GPT2.ipynb
diff --git a/examples/python/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_Hubert.ipynb b/examples/python/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_Hubert.ipynb
diff --git a/examples/python/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_InstructorEmbeddings.ipynb b/examples/python/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_InstructorEmbeddings.ipynb
diff --git a/...on/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_RoBertaForZeroShotClassification.ipynb b/...on/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_RoBertaForZeroShotClassification.ipynb
diff --git a/examples/python/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_Wav2Vec2.ipynb b/examples/python/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_Wav2Vec2.ipynb
diff --git a/.../transformers/openvino/HuggingFace_OpenVINO_in_Spark_NLP_AlbertForQuestionAnswering.ipynb b/.../transformers/openvino/HuggingFace_OpenVINO_in_Spark_NLP_AlbertForQuestionAnswering.ipynb
diff --git a/...sformers/openvino/HuggingFace_OpenVINO_in_Spark_NLP_AlbertForSequenceClassification.ipynb b/...sformers/openvino/HuggingFace_OpenVINO_in_Spark_NLP_AlbertForSequenceClassification.ipynb
diff --git a/...ransformers/openvino/HuggingFace_OpenVINO_in_Spark_NLP_AlbertForTokenClassification.ipynb b/...ransformers/openvino/HuggingFace_OpenVINO_in_Spark_NLP_AlbertForTokenClassification.ipynb
diff --git a/...ansformers/openvino/HuggingFace_OpenVINO_in_Spark_NLP_BartForZeroShotClassification.ipynb b/...ansformers/openvino/HuggingFace_OpenVINO_in_Spark_NLP_BartForZeroShotClassification.ipynb
diff --git a/...ansformers/openvino/HuggingFace_OpenVINO_in_Spark_NLP_BertForSequenceClassification.ipynb b/...ansformers/openvino/HuggingFace_OpenVINO_in_Spark_NLP_BertForSequenceClassification.ipynb
diff --git a/...ansformers/openvino/HuggingFace_OpenVINO_in_Spark_NLP_BertForZeroShotClassification.ipynb b/...ansformers/openvino/HuggingFace_OpenVINO_in_Spark_NLP_BertForZeroShotClassification.ipynb
diff --git a/examples/python/transformers/openvino/HuggingFace_OpenVINO_in_Spark_NLP_Whisper.ipynb b/examples/python/transformers/openvino/HuggingFace_OpenVINO_in_Spark_NLP_Whisper.ipynb
diff --git a/src/main/scala/com/johnsnowlabs/ml/ai/AlbertClassification.scala b/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),