SPARKNLP-985: Make Whisper compatible with onnx_data

JohnSnowLabs · Feb 6, 2024 · 6bab822 · 6bab822
1 parent a2cb06b
commit 6bab822
Showing 1 changed file with 10 additions and 7 deletions.
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/audio/WhisperForCTC.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/audio/WhisperForCTC.scala
@@ -360,9 +360,9 @@ class WhisperForCTC(override val uid: String)
           path,
           spark,
           Seq(
-            (wrappers.encoder, "encoder_model"),
-            (wrappers.decoder, "decoder_model"),
-            (wrappers.decoderWithPast, "decoder_with_past_model")),
+            (wrappers.encoder, "encoder_model.onnx"),
+            (wrappers.decoder, "decoder_model.onnx"),
+            (wrappers.decoderWithPast, "decoder_with_past_model.onnx")),
           WhisperForCTC.suffix)
     }
 
@@ -425,6 +425,9 @@ trait ReadWhisperForCTCDLModel extends ReadTensorflowModel with ReadOnnxModel {
   override val tfFile: String = "whisper_ctc_tensorflow"
   override val onnxFile: String = "whisper_ctc_onnx"
   val suffix: String = "_whisper_ctc"
+  val encoderModel: String = "encoder_model.onnx"
+  val decoderModel: String = "decoder_model.onnx"
+  val decoderWithPastModel: String = "decoder_with_past_model.onnx"
 
   private def checkVersion(spark: SparkSession): Unit = {
     val version = Version.parse(spark.version).toFloat
@@ -447,13 +450,13 @@ trait ReadWhisperForCTCDLModel extends ReadTensorflowModel with ReadOnnxModel {
           readOnnxModels(
             path,
             spark,
-            Seq("encoder_model", "decoder_model", "decoder_with_past_model"),
+            Seq(encoderModel, decoderModel, decoderWithPastModel),
             WhisperForCTC.suffix)
 
         val onnxWrappers = EncoderDecoderWrappers(
-          wrappers("encoder_model"),
-          decoder = wrappers("decoder_model"),
-          decoderWithPast = wrappers("decoder_with_past_model"))
+          wrappers(encoderModel),
+          decoder = wrappers(decoderModel),
+          decoderWithPast = wrappers(decoderWithPastModel))
 
         instance.setModelIfNotSet(spark, None, Some(onnxWrappers))
       case _ =>