[SPARKNLP-1037] Adding changes to transfer ONNX files on executors vi…

…a addFile
JohnSnowLabs · Apr 17, 2024 · c9380c6 · c9380c6
1 parent 4353eec
commit c9380c6
Show file tree

Hide file tree

Showing 50 changed files with 227 additions and 541 deletions.
diff --git a/src/main/scala/com/johnsnowlabs/ml/ai/LLAMA2.scala b/src/main/scala/com/johnsnowlabs/ml/ai/LLAMA2.scala
@@ -20,7 +20,6 @@ import ai.onnxruntime.{OnnxTensor, OrtEnvironment, OrtSession}
 import com.johnsnowlabs.ml.ai.util.Generation.{Generate, GenerationConfig}
 import com.johnsnowlabs.ml.onnx.OnnxSession
 import com.johnsnowlabs.ml.onnx.OnnxWrapper.DecoderWrappers
-import com.johnsnowlabs.ml.onnx.OnnxLlmWrapper.DecoderWrappersLlm
 import com.johnsnowlabs.ml.onnx.TensorResources.implicits._
 import com.johnsnowlabs.ml.tensorflow.sentencepiece.SentencePieceWrapper
 import com.johnsnowlabs.nlp.Annotation
@@ -30,7 +29,7 @@ import com.johnsnowlabs.nlp.AnnotatorType.DOCUMENT
 import org.tensorflow.{Session, Tensor}
 
 private[johnsnowlabs] class LLAMA2(
-    val onnxWrappers: DecoderWrappersLlm,
+    val onnxWrappers: DecoderWrappers,
     val spp: SentencePieceWrapper,
     generationConfig: GenerationConfig)
     extends Serializable

diff --git a/src/main/scala/com/johnsnowlabs/ml/ai/M2M100.scala b/src/main/scala/com/johnsnowlabs/ml/ai/M2M100.scala
@@ -18,8 +18,8 @@ package com.johnsnowlabs.ml.ai
 
 import ai.onnxruntime.{OnnxTensor, OrtEnvironment, OrtSession}
 import com.johnsnowlabs.ml.ai.util.Generation.{Generate, GenerationConfig}
-import com.johnsnowlabs.ml.onnx.OnnxLlmWrapper.EncoderDecoderWithoutPastWrappersLlm
 import com.johnsnowlabs.ml.onnx.OnnxSession
+import com.johnsnowlabs.ml.onnx.OnnxWrapper.EncoderDecoderWithoutPastWrappers
 import com.johnsnowlabs.ml.onnx.TensorResources.implicits._
 import com.johnsnowlabs.ml.tensorflow.sentencepiece.SentencePieceWrapper
 import com.johnsnowlabs.nlp.Annotation
@@ -29,7 +29,7 @@ import org.tensorflow.{Session, Tensor}
 import scala.collection.JavaConverters._
 
 private[johnsnowlabs] class M2M100(
-    val onnxWrappers: EncoderDecoderWithoutPastWrappersLlm,
+    val onnxWrappers: EncoderDecoderWithoutPastWrappers,
     val spp: SentencePieceWrapper,
     generationConfig: GenerationConfig,
     vocab: Map[String, Int])

diff --git a/src/main/scala/com/johnsnowlabs/ml/ai/Whisper.scala b/src/main/scala/com/johnsnowlabs/ml/ai/Whisper.scala
@@ -24,7 +24,6 @@ import com.johnsnowlabs.ml.ai.util.Generation.Logit.LogitProcess.{
   SuppressLogitProcessor
 }
 import com.johnsnowlabs.ml.ai.util.Generation.Logit.LogitProcessorList
-import com.johnsnowlabs.ml.onnx.OnnxLlmWrapper.EncoderDecoderWrappersLlm
 import com.johnsnowlabs.ml.onnx.OnnxSession
 import com.johnsnowlabs.ml.onnx.OnnxWrapper.EncoderDecoderWrappers
 import com.johnsnowlabs.ml.onnx.TensorResources.implicits._
@@ -58,7 +57,7 @@ import scala.collection.JavaConverters._
   */
 private[johnsnowlabs] class Whisper(
     val tensorflowWrapper: Option[TensorflowWrapper],
-    val onnxWrappers: Option[EncoderDecoderWrappersLlm],
+    val onnxWrappers: Option[EncoderDecoderWrappers],
     configProtoBytes: Option[Array[Byte]] = None,
     signatures: Option[Map[String, String]] = None,
     preprocessor: WhisperPreprocessor,

diff --git a/src/main/scala/com/johnsnowlabs/ml/onnx/OnnxLlmWrapper.scala b/src/main/scala/com/johnsnowlabs/ml/onnx/OnnxLlmWrapper.scala