From 7f78be36e135e20dd46c27af32834b4c25cd718f Mon Sep 17 00:00:00 2001
From: Devin Ha <33089471+DevinTDHa@users.noreply.github.com>
Date: Thu, 7 Sep 2023 18:17:03 +0200
Subject: [PATCH] [SPARKNLP-906] Fix reading suffix (#13945)

---
 .../scala/com/johnsnowlabs/ml/onnx/OnnxSerializeModel.scala   | 2 +-
 .../johnsnowlabs/nlp/annotators/audio/WhisperForCTCTest.scala | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/main/scala/com/johnsnowlabs/ml/onnx/OnnxSerializeModel.scala b/src/main/scala/com/johnsnowlabs/ml/onnx/OnnxSerializeModel.scala
index e67bbb3a6211fb..85578509a90869 100644
--- a/src/main/scala/com/johnsnowlabs/ml/onnx/OnnxSerializeModel.scala
+++ b/src/main/scala/com/johnsnowlabs/ml/onnx/OnnxSerializeModel.scala
@@ -126,7 +126,7 @@ trait ReadOnnxModel {
 
     val wrappers = (modelNames map { modelName: String =>
       // 2. Copy to local dir
-      val localModelFile = modelName + suffix
+      val localModelFile = modelName
       fs.copyToLocalFile(new Path(path, localModelFile), new Path(tmpFolder))
 
       val localPath = new Path(tmpFolder, localModelFile).toString
diff --git a/src/test/scala/com/johnsnowlabs/nlp/annotators/audio/WhisperForCTCTest.scala b/src/test/scala/com/johnsnowlabs/nlp/annotators/audio/WhisperForCTCTest.scala
index bc1ae6e323c437..0d37be98bcb1d7 100644
--- a/src/test/scala/com/johnsnowlabs/nlp/annotators/audio/WhisperForCTCTest.scala
+++ b/src/test/scala/com/johnsnowlabs/nlp/annotators/audio/WhisperForCTCTest.scala
@@ -32,12 +32,12 @@ class WhisperForCTCTest extends AnyFlatSpec with WhisperForCTCBehaviors {
 
   // Needs to be added manually
   lazy val modelTf: WhisperForCTC = WhisperForCTC
-    .loadSavedModel("exported_tf/openai/whisper-tiny", ResourceHelper.spark)
+    .pretrained("asr_whisper_tiny")
     .setInputCols("audio_assembler")
     .setOutputCol("document")
 
   lazy val modelOnnx: WhisperForCTC = WhisperForCTC
-    .loadSavedModel("exported_onnx/openai/whisper-tiny", ResourceHelper.spark)
+    .pretrained()
     .setInputCols("audio_assembler")
     .setOutputCol("document")