Fix M2M100 not working on the second run [skip test]

JohnSnowLabs · Mar 2, 2024 · 75d398e · 75d398e
1 parent ad5a4ea
commit 75d398e
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 12 deletions.
diff --git a/src/main/scala/com/johnsnowlabs/ml/ai/M2M100.scala b/src/main/scala/com/johnsnowlabs/ml/ai/M2M100.scala
@@ -196,12 +196,12 @@ private[johnsnowlabs] class M2M100(
       applySoftmax = false)
 
     // Run the prompt through the decoder and get the past
-//    val decoderOutputs =
-//      generateGreedyOnnx(
-//        decoderInputIds,
-//        decoderEncoderStateTensors,
-//        encoderAttentionMaskTensors,
-//        onnxSession = (decoderSession, decoderEnv))
+    //    val decoderOutputs =
+    //      generateGreedyOnnx(
+    //        decoderInputIds,
+    //        decoderEncoderStateTensors,
+    //        encoderAttentionMaskTensors,
+    //        onnxSession = (decoderSession, decoderEnv))
 
     // close sessions
     decoderEncoderStateTensors.fold(
@@ -216,8 +216,6 @@ private[johnsnowlabs] class M2M100(
       },
       onnxTensor => onnxTensor.close())
 
-    encoderSession.close()
-    decoderSession.close()
     encoderEnv.close()
     decoderEnv.close()
 

diff --git a/src/test/scala/com/johnsnowlabs/nlp/annotators/seq2seq/LLAMA2TestSpec.scala b/src/test/scala/com/johnsnowlabs/nlp/annotators/seq2seq/LLAMA2TestSpec.scala
@@ -50,6 +50,10 @@ class LLAMA2TestSpec extends AnyFlatSpec {
 
     val pipelineModel = pipeline.fit(testData)
 
+    pipelineModel
+      .transform(testData)
+      .show(truncate = false)
+
     pipelineModel
       .transform(testData)
       .show(truncate = false)

diff --git a/src/test/scala/com/johnsnowlabs/nlp/annotators/seq2seq/M2M100TestSpec.scala b/src/test/scala/com/johnsnowlabs/nlp/annotators/seq2seq/M2M100TestSpec.scala
@@ -45,11 +45,15 @@ class M2M100TestSpec extends AnyFlatSpec {
       .setOutputCol("generation")
       .setBeamSize(1)
 
-    new Pipeline()
+    val pipeline = new Pipeline()
       .setStages(Array(documentAssembler, m2m100))
-      .fit(testData)
-      .transform(testData)
-      .show(truncate = false)
+
+    val pipelineModel = pipeline.fit(testData)
+
+    val result = pipelineModel.transform(testData)
+
+    result.show(truncate = false)
+    result.show(truncate = false)
 
   }