Sparknlp 967 add onnx support to xlm roberta classifiers (#14130)

* fixing typo + adding support for ONNX to XLM-Roberta * adding conversion notebooks
JohnSnowLabs · Jan 18, 2024 · 818de8d · 818de8d
1 parent 8585b7e
commit 818de8d
Show file tree

Hide file tree

Showing 14 changed files with 7,116 additions and 219 deletions.
diff --git a/...thon/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_XlmRoBertaForQuestionAnswering.ipynb b/...thon/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_XlmRoBertaForQuestionAnswering.ipynb
diff --git a/...transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_XlmRoBertaForSequenceClassification.ipynb b/...transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_XlmRoBertaForSequenceClassification.ipynb
diff --git a/...on/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_XlmRoBertaForTokenClassification.ipynb b/...on/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_XlmRoBertaForTokenClassification.ipynb
diff --git a/src/main/scala/com/johnsnowlabs/ml/ai/AlbertClassification.scala b/src/main/scala/com/johnsnowlabs/ml/ai/AlbertClassification.scala
@@ -108,7 +108,7 @@ private[johnsnowlabs] class AlbertClassification(
     val maxSentenceLength = batch.map(encodedSentence => encodedSentence.length).max
 
     val rawScores = detectedEngine match {
-      case ONNX.name => getRowScoresWithOnnx(batch, maxSentenceLength, sequence = true)
+      case ONNX.name => getRawScoresWithOnnx(batch, maxSentenceLength, sequence = true)
       case _ => getRawScoresWithTF(batch, maxSentenceLength)
     }
 
@@ -128,7 +128,7 @@ private[johnsnowlabs] class AlbertClassification(
     val maxSentenceLength = batch.map(encodedSentence => encodedSentence.length).max
 
     val rawScores = detectedEngine match {
-      case ONNX.name => getRowScoresWithOnnx(batch, maxSentenceLength, sequence = true)
+      case ONNX.name => getRawScoresWithOnnx(batch, maxSentenceLength, sequence = true)
       case _ => getRawScoresWithTF(batch, maxSentenceLength)
     }
 
@@ -203,7 +203,7 @@ private[johnsnowlabs] class AlbertClassification(
     rawScores
   }
 
-  private def getRowScoresWithOnnx(
+  private def getRawScoresWithOnnx(
       batch: Seq[Array[Int]],
       maxSentenceLength: Int,
       sequence: Boolean): Array[Float] = {

diff --git a/src/main/scala/com/johnsnowlabs/ml/ai/BertClassification.scala b/src/main/scala/com/johnsnowlabs/ml/ai/BertClassification.scala
@@ -149,7 +149,7 @@ private[johnsnowlabs] class BertClassification(
 
     val rawScores = detectedEngine match {
       case ONNX.name =>
-        getRowScoresWithOnnx(batch, maxSentenceLength)
+        getRawScoresWithOnnx(batch, maxSentenceLength)
       case _ => getRawScoresWithTF(batch, maxSentenceLength)
     }
 
@@ -218,7 +218,7 @@ private[johnsnowlabs] class BertClassification(
     rawScores
   }
 
-  private def getRowScoresWithOnnx(
+  private def getRawScoresWithOnnx(
       batch: Seq[Array[Int]],
       maxSentenceLength: Int): Array[Float] = {
 
@@ -265,7 +265,7 @@ private[johnsnowlabs] class BertClassification(
     val maxSentenceLength = batch.map(encodedSentence => encodedSentence.length).max
     val rawScores = detectedEngine match {
       case ONNX.name =>
-        getRowScoresWithOnnx(batch, maxSentenceLength)
+        getRawScoresWithOnnx(batch, maxSentenceLength)
       case _ => getRawScoresWithTF(batch, maxSentenceLength)
     }
 

diff --git a/src/main/scala/com/johnsnowlabs/ml/ai/CamemBertClassification.scala b/src/main/scala/com/johnsnowlabs/ml/ai/CamemBertClassification.scala
@@ -123,7 +123,7 @@ private[johnsnowlabs] class CamemBertClassification(
     val maxSentenceLength = batch.map(encodedSentence => encodedSentence.length).max
 
     val rawScores = detectedEngine match {
-      case ONNX.name => getRowScoresWithOnnx(batch)
+      case ONNX.name => getRawScoresWithOnnx(batch)
       case _ => getRawScoresWithTF(batch, maxSentenceLength)
     }
 
@@ -189,7 +189,7 @@ private[johnsnowlabs] class CamemBertClassification(
     rawScores
   }
 
-  private def getRowScoresWithOnnx(batch: Seq[Array[Int]]): Array[Float] = {
+  private def getRawScoresWithOnnx(batch: Seq[Array[Int]]): Array[Float] = {
 
     // [nb of encoded sentences , maxSentenceLength]
     val (runner, env) = onnxWrapper.get.getSession(onnxSessionOptions)
@@ -227,7 +227,7 @@ private[johnsnowlabs] class CamemBertClassification(
     val batchLength = batch.length
 
     val rawScores = detectedEngine match {
-      case ONNX.name => getRowScoresWithOnnx(batch)
+      case ONNX.name => getRawScoresWithOnnx(batch)
       case _ => getRawScoresWithTF(batch, maxSentenceLength)
     }
 

diff --git a/src/main/scala/com/johnsnowlabs/ml/ai/DeBertaClassification.scala b/src/main/scala/com/johnsnowlabs/ml/ai/DeBertaClassification.scala
@@ -109,7 +109,7 @@ private[johnsnowlabs] class DeBertaClassification(
     val batchLength = batch.length
 
     val rawScores = detectedEngine match {
-      case ONNX.name => getRowScoresWithOnnx(batch)
+      case ONNX.name => getRawScoresWithOnnx(batch)
       case _ => getRawScoresWithTF(batch)
     }
 
@@ -182,7 +182,7 @@ private[johnsnowlabs] class DeBertaClassification(
     rawScores
   }
 
-  private def getRowScoresWithOnnx(batch: Seq[Array[Int]]): Array[Float] = {
+  private def getRawScoresWithOnnx(batch: Seq[Array[Int]]): Array[Float] = {
 
     // [nb of encoded sentences , maxSentenceLength]
     val (runner, env) = onnxWrapper.get.getSession(onnxSessionOptions)
@@ -219,7 +219,7 @@ private[johnsnowlabs] class DeBertaClassification(
     val batchLength = batch.length
 
     val rawScores = detectedEngine match {
-      case ONNX.name => getRowScoresWithOnnx(batch)
+      case ONNX.name => getRawScoresWithOnnx(batch)
       case _ => getRawScoresWithTF(batch)
     }
 

diff --git a/src/main/scala/com/johnsnowlabs/ml/ai/DistilBertClassification.scala b/src/main/scala/com/johnsnowlabs/ml/ai/DistilBertClassification.scala
@@ -148,7 +148,7 @@ private[johnsnowlabs] class DistilBertClassification(
     val maxSentenceLength = batch.map(encodedSentence => encodedSentence.length).max
 
     val rawScores = detectedEngine match {
-      case ONNX.name => getRowScoresWithOnnx(batch)
+      case ONNX.name => getRawScoresWithOnnx(batch)
       case _ => getRawScoresWithTF(batch, maxSentenceLength)
     }
 
@@ -211,7 +211,7 @@ private[johnsnowlabs] class DistilBertClassification(
     rawScores
   }
 
-  private def getRowScoresWithOnnx(batch: Seq[Array[Int]]): Array[Float] = {
+  private def getRawScoresWithOnnx(batch: Seq[Array[Int]]): Array[Float] = {
 
     val (runner, env) = onnxWrapper.get.getSession(onnxSessionOptions)
 
@@ -247,7 +247,7 @@ private[johnsnowlabs] class DistilBertClassification(
     val maxSentenceLength = batch.map(encodedSentence => encodedSentence.length).max
 
     val rawScores = detectedEngine match {
-      case ONNX.name => getRowScoresWithOnnx(batch)
+      case ONNX.name => getRawScoresWithOnnx(batch)
       case _ => getRawScoresWithTF(batch, maxSentenceLength)
     }
 

diff --git a/src/main/scala/com/johnsnowlabs/ml/ai/RoBertaClassification.scala b/src/main/scala/com/johnsnowlabs/ml/ai/RoBertaClassification.scala
@@ -141,7 +141,7 @@ private[johnsnowlabs] class RoBertaClassification(
     val maxSentenceLength = batch.map(encodedSentence => encodedSentence.length).max
 
     val rawScores = detectedEngine match {
-      case ONNX.name => getRowScoresWithOnnx(batch)
+      case ONNX.name => getRawScoresWithOnnx(batch)
       case _ => getRawScoresWithTF(batch, maxSentenceLength)
     }
 
@@ -207,7 +207,7 @@ private[johnsnowlabs] class RoBertaClassification(
     rawScores
   }
 
-  private def getRowScoresWithOnnx(batch: Seq[Array[Int]]): Array[Float] = {
+  private def getRawScoresWithOnnx(batch: Seq[Array[Int]]): Array[Float] = {
 
     // [nb of encoded sentences , maxSentenceLength]
     val (runner, env) = onnxWrapper.get.getSession(onnxSessionOptions)
@@ -244,7 +244,7 @@ private[johnsnowlabs] class RoBertaClassification(
     val maxSentenceLength = batch.map(encodedSentence => encodedSentence.length).max
 
     val rawScores = detectedEngine match {
-      case ONNX.name => getRowScoresWithOnnx(batch)
+      case ONNX.name => getRawScoresWithOnnx(batch)
       case _ => getRawScoresWithTF(batch, maxSentenceLength)
     }