[SPARKNLP-1011] Adding unique suffix to avoid duplication in spark files

danilojsl · danilojsl · commit d34af15b9d5f · 2024-04-24T12:54:48.000-05:00
diff --git a/src/main/scala/com/johnsnowlabs/ml/onnx/OnnxWrapper.scala b/src/main/scala/com/johnsnowlabs/ml/onnx/OnnxWrapper.scala
@@ -99,7 +99,6 @@ object OnnxWrapper {
       }
     }
 
-  // TODO: the parts related to onnx_data should be refactored once we support addFile()
   def read(
       sparkSession: SparkSession,
       modelPath: String,
@@ -108,7 +107,6 @@ object OnnxWrapper {
       modelName: String = "model",
       dataFileSuffix: Option[String] = Some("_data"),
       onnxFileSuffix: Option[String] = None): OnnxWrapper = {
-    // TODO: onnxFileSuffix default none??
     // 1. Create tmp folder
     val tmpFolder = Files
       .createTempDirectory(UUID.randomUUID().toString.takeRight(12) + "_onnx")
@@ -136,7 +134,6 @@ object OnnxWrapper {
       if (onnxFileSuffix.isDefined && dataFileSuffix.isDefined) {
         val onnxDataFilePath = s"${onnxFileSuffix.get}_$modelName${dataFileSuffix.get}"
         onnxDataFile = Paths.get(parentDir, onnxDataFilePath).toFile
-        // TODO: Here is the error it does not exist
         onnxDataFile.exists()
       } else false
     }
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/AlbertForQuestionAnswering.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/AlbertForQuestionAnswering.scala
@@ -328,12 +328,7 @@ trait ReadAlbertForQuestionAnsweringDLModel
         instance.setModelIfNotSet(spark, Some(tf), None, spp)
       case ONNX.name =>
         val onnxWrapper =
-          readOnnxModel(
-            path,
-            spark,
-            "_albert_classification_onnx",
-            zipped = true,
-            useBundle = false)
+          readOnnxModel(path, spark, "albert_qa_classification_onnx")
         instance.setModelIfNotSet(spark, None, Some(onnxWrapper), spp)
       case _ =>
         throw new Exception(notSupportedEngineError)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/AlbertForSequenceClassification.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/AlbertForSequenceClassification.scala
@@ -382,12 +382,7 @@ trait ReadAlbertForSequenceDLModel
         instance.setModelIfNotSet(spark, Some(tf), None, spp)
       case ONNX.name =>
         val onnxWrapper =
-          readOnnxModel(
-            path,
-            spark,
-            "_albert_classification_onnx",
-            zipped = true,
-            useBundle = false)
+          readOnnxModel(path, spark, "albert_sequence_classification_onnx")
         instance.setModelIfNotSet(spark, None, Some(onnxWrapper), spp)
       case _ =>
         throw new Exception(notSupportedEngineError)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/AlbertForTokenClassification.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/AlbertForTokenClassification.scala
@@ -353,12 +353,7 @@ trait ReadAlbertForTokenDLModel
         instance.setModelIfNotSet(spark, Some(tfWrapper), None, spp)
       case ONNX.name =>
         val onnxWrapper =
-          readOnnxModel(
-            path,
-            spark,
-            "_albert_classification_onnx",
-            zipped = true,
-            useBundle = false)
+          readOnnxModel(path, spark, "albert_token_classification_onnx")
         instance.setModelIfNotSet(spark, None, Some(onnxWrapper), spp)
       case _ =>
         throw new Exception(notSupportedEngineError)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/BertForQuestionAnswering.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/BertForQuestionAnswering.scala
@@ -330,7 +330,7 @@ trait ReadBertForQuestionAnsweringDLModel extends ReadTensorflowModel with ReadO
         instance.setModelIfNotSet(spark, Some(tensorFlow), None)
       case ONNX.name =>
         val onnxWrapper =
-          readOnnxModel(path, spark, "_bert_classification_onnx")
+          readOnnxModel(path, spark, "bert_qa_classification_onnx")
         instance.setModelIfNotSet(spark, None, Some(onnxWrapper))
       case _ =>
         throw new Exception(notSupportedEngineError)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/BertForSequenceClassification.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/BertForSequenceClassification.scala
@@ -384,7 +384,7 @@ trait ReadBertForSequenceDLModel extends ReadTensorflowModel with ReadOnnxModel
         instance.setModelIfNotSet(spark, Some(tensorFlow), None)
       case ONNX.name =>
         val onnxWrapper =
-          readOnnxModel(path, spark, "_bert_classification_onnx")
+          readOnnxModel(path, spark, "bert_sequence_classification_onnx")
         instance.setModelIfNotSet(spark, None, Some(onnxWrapper))
       case _ =>
         throw new Exception(notSupportedEngineError)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/BertForTokenClassification.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/BertForTokenClassification.scala
@@ -350,7 +350,7 @@ trait ReadBertForTokenDLModel extends ReadTensorflowModel with ReadOnnxModel {
         instance.setModelIfNotSet(spark, Some(tensorFlow), None)
       case ONNX.name =>
         val onnxWrapper =
-          readOnnxModel(path, spark, "_bert_classification_onnx")
+          readOnnxModel(path, spark, "bert_token_classification_onnx")
         instance.setModelIfNotSet(spark, None, Some(onnxWrapper))
       case _ =>
         throw new Exception(notSupportedEngineError)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/BertForZeroShotClassification.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/BertForZeroShotClassification.scala
@@ -400,7 +400,7 @@ trait ReadBertForZeroShotDLModel extends ReadTensorflowModel with ReadOnnxModel
         instance.setModelIfNotSet(spark, Some(tensorFlow), None)
       case ONNX.name =>
         val onnxWrapper =
-          readOnnxModel(path, spark, "_bert_classification_onnx")
+          readOnnxModel(path, spark, "bert_zs_classification_onnx")
         instance.setModelIfNotSet(spark, None, Some(onnxWrapper))
       case _ =>
         throw new Exception(notSupportedEngineError)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/CamemBertForQuestionAnswering.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/CamemBertForQuestionAnswering.scala
@@ -336,7 +336,7 @@ trait ReadCamemBertForQADLModel
           readOnnxModel(
             path,
             spark,
-            "_camembert_classification_onnx",
+            "camembert_qa_classification_onnx",
             zipped = true,
             useBundle = false,
             None)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/CamemBertForSequenceClassification.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/CamemBertForSequenceClassification.scala
@@ -388,7 +388,7 @@ trait ReadCamemBertForSequenceDLModel
           readOnnxModel(
             path,
             spark,
-            "_camembert_classification_onnx",
+            "camembert_sequence_classification_onnx",
             zipped = true,
             useBundle = false,
             None)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/CamemBertForTokenClassification.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/CamemBertForTokenClassification.scala
@@ -357,7 +357,7 @@ trait ReadCamemBertForTokenDLModel
           readOnnxModel(
             path,
             spark,
-            "_camembert_classification_onnx",
+            "camembert_token_classification_onnx",
             zipped = true,
             useBundle = false,
             None)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/DeBertaForQuestionAnswering.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/DeBertaForQuestionAnswering.scala
@@ -333,7 +333,7 @@ trait ReadDeBertaForQuestionAnsweringDLModel
           readOnnxModel(
             path,
             spark,
-            "_deberta_classification_onnx",
+            "deberta_qa_classification_onnx",
             zipped = true,
             useBundle = false,
             None)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/DeBertaForSequenceClassification.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/DeBertaForSequenceClassification.scala
@@ -385,7 +385,7 @@ trait ReadDeBertaForSequenceDLModel
           readOnnxModel(
             path,
             spark,
-            "_deberta_classification_onnx",
+            "deberta_sequence_classification_onnx",
             zipped = true,
             useBundle = false,
             None)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/DeBertaForTokenClassification.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/DeBertaForTokenClassification.scala
@@ -356,7 +356,7 @@ trait ReadDeBertaForTokenDLModel
           readOnnxModel(
             path,
             spark,
-            "_deberta_classification_onnx",
+            "deberta_token_classification_onnx",
             zipped = true,
             useBundle = false,
             None)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/DistilBertForQuestionAnswering.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/DistilBertForQuestionAnswering.scala
@@ -336,7 +336,7 @@ trait ReadDistilBertForQuestionAnsweringDLModel extends ReadTensorflowModel with
           readOnnxModel(
             path,
             spark,
-            "_distilbert_classification_onnx",
+            "distilbert_qa_classification_onnx",
             zipped = true,
             useBundle = false,
             None)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/DistilBertForSequenceClassification.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/DistilBertForSequenceClassification.scala
@@ -388,7 +388,7 @@ trait ReadDistilBertForSequenceDLModel extends ReadTensorflowModel with ReadOnnx
           readOnnxModel(
             path,
             spark,
-            "_albert_classification_onnx",
+            "distilbert_sequence_classification_onnx",
             zipped = true,
             useBundle = false,
             None)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/DistilBertForTokenClassification.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/DistilBertForTokenClassification.scala
@@ -358,7 +358,7 @@ trait ReadDistilBertForTokenDLModel extends ReadTensorflowModel with ReadOnnxMod
           readOnnxModel(
             path,
             spark,
-            "_distilbert_classification_onnx",
+            "distilbert_token_classification_onnx",
             zipped = true,
             useBundle = false,
             None)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/MPNetForQuestionAnswering.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/MPNetForQuestionAnswering.scala
@@ -296,13 +296,7 @@ trait ReadMPNetForQuestionAnsweringDLModel extends ReadOnnxModel {
     instance.getEngine match {
       case ONNX.name =>
         val onnxWrapper =
-          readOnnxModel(
-            path,
-            spark,
-            "_mpnet_question_answering_onnx",
-            zipped = true,
-            useBundle = false,
-            None)
+          readOnnxModel(path, spark, "mpnet_qa_onnx", zipped = true, useBundle = false, None)
         instance.setModelIfNotSet(spark, Some(onnxWrapper))
       case _ =>
         throw new NotImplementedError("Tensorflow models are not supported.")
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/MPNetForSequenceClassification.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/MPNetForSequenceClassification.scala
@@ -358,7 +358,7 @@ trait ReadMPNetForSequenceDLModel extends ReadOnnxModel {
           readOnnxModel(
             path,
             spark,
-            "_mpnet_classification_onnx",
+            "mpnet_sequence_classification_onnx",
             zipped = true,
             useBundle = false,
             None)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/RoBertaForQuestionAnswering.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/RoBertaForQuestionAnswering.scala
@@ -348,7 +348,7 @@ trait ReadRoBertaForQuestionAnsweringDLModel extends ReadTensorflowModel with Re
           readOnnxModel(
             path,
             spark,
-            "roberta_classification_onnx",
+            "roberta_qa_classification_onnx",
             zipped = true,
             useBundle = false,
             None)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/RoBertaForSequenceClassification.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/RoBertaForSequenceClassification.scala
@@ -400,7 +400,7 @@ trait ReadRoBertaForSequenceDLModel extends ReadTensorflowModel with ReadOnnxMod
           readOnnxModel(
             path,
             spark,
-            "roberta_classification_onnx",
+            "roberta_sequence_classification_onnx",
             zipped = true,
             useBundle = false,
             None)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/RoBertaForTokenClassification.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/RoBertaForTokenClassification.scala
@@ -370,7 +370,7 @@ trait ReadRoBertaForTokenDLModel extends ReadTensorflowModel with ReadOnnxModel
           readOnnxModel(
             path,
             spark,
-            "roberta_classification_onnx",
+            "roberta_token_classification_onnx",
             zipped = true,
             useBundle = false,
             None)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/XlmRoBertaForQuestionAnswering.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/XlmRoBertaForQuestionAnswering.scala
@@ -326,7 +326,7 @@ trait ReadXlmRoBertaForQuestionAnsweringDLModel
           readOnnxModel(
             path,
             spark,
-            "xlm_roberta_classification_onnx",
+            "xlm_roberta_qa_classification_onnx",
             zipped = true,
             useBundle = false,
             None)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/XlmRoBertaForSequenceClassification.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/XlmRoBertaForSequenceClassification.scala
@@ -377,7 +377,7 @@ trait ReadXlmRoBertaForSequenceDLModel
           readOnnxModel(
             path,
             spark,
-            "xlm_roberta_classification_onnx",
+            "xlm_roberta_sequence_classification_onnx",
             zipped = true,
             useBundle = false,
             None)
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/XlmRoBertaForTokenClassification.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/XlmRoBertaForTokenClassification.scala
@@ -349,7 +349,7 @@ trait ReadXlmRoBertaForTokenDLModel
           readOnnxModel(
             path,
             spark,
-            "xlm_roberta_classification_onnx",
+            "xlm_roberta_token_classification_onnx",
             zipped = true,
             useBundle = false,
             None)
