Fix the bad dimensions in BGE output [skip test]

src/main/scala/com/johnsnowlabs/ml/ai/BGE.scala

@@ -183,20 +183,17 @@ private[johnsnowlabs] class BGE(
       val info = lastHiddenState.getInfo.asInstanceOf[TensorInfo]
       val shape = info.getShape
       try {
-        val embeddings = lastHiddenState
+        val flattenEmbeddings = lastHiddenState
           .asInstanceOf[OnnxTensor]
           .getFloatBuffer
           .array()
         tokenTensors.close()
         maskTensors.close()
         segmentTensors.close()
 
-        val dim = shape.last.toInt
-        // Perfom CLS pooling (the first element of each sequence)
-        val clsPooling = embeddings.grouped(dim).map(_.head).toArray
-        val normalizedSentenceEmbeddings = LinAlg.lpNormalizeArray(clsPooling, 2)
-
-        Array(normalizedSentenceEmbeddings)
+        val embeddings = LinAlg.avgPooling(flattenEmbeddings, attentionMask, shape)
+        val normalizedEmbeddings = LinAlg.l2Normalize(embeddings)
+        LinAlg.denseMatrixToArray(normalizedEmbeddings)
       } finally if (results != null) results.close()
     }
   }

src/test/scala/com/johnsnowlabs/nlp/embeddings/BGEEmbeddingsTestSpec.scala

@@ -113,4 +113,48 @@ class BGEEmbeddingsTestSpec extends AnyFlatSpec {
     pipelineDF.select("bge.embeddings").show(false)
   }
 
+  it should "not return empty embeddings" taggedAs SlowTest in {
+    import ResourceHelper.spark.implicits._
+    val interests = Seq(
+      "I like music",
+      "I like movies",
+      "I like books",
+      "I like sports",
+      "I like travel",
+      "I like food",
+      "I like games",
+      "I like art",
+      "I like nature",
+      "I like science",
+      "I like technology",
+      "I like history",
+      "I like fashion",
+      "I like cars",
+      "I like animals",
+      "I like gardening")
+    val testDf = interests.toDF("text")
+
+    val document = new DocumentAssembler()
+      .setInputCol("text")
+      .setOutputCol("document")
+
+    val embeddings = BGEEmbeddings
+      .pretrained()
+      .setInputCols(Array("document"))
+      .setOutputCol("bge")
+
+    val pipeline = new Pipeline().setStages(Array(document, embeddings))
+
+    val pipelineDF = pipeline.fit(testDf).transform(testDf)
+
+    val embeddingsDF = pipelineDF.withColumn("embeddings", col("bge.embeddings").getItem(0))
+
+    val sizesArray: Array[Int] = embeddingsDF
+      .select(size(col("embeddings")).as("size"))
+      .collect()
+      .map(row => row.getAs[Int]("size"))
+
+    assert(sizesArray.forall(_ > 0))
+  }
+
 }