[SPARKNLP-906] Fix reading suffix (#13945)

JohnSnowLabs · Sep 7, 2023 · 7f78be3 · 7f78be3
1 parent 543faaf
commit 7f78be3
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git 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
@@ -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")