fix: fix item issue for mobile cases

unit-mesh · Nov 9, 2023 · 401f55c · 401f55c
1 parent 72e60e0
commit 401f55c
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/llm-modules/sentence-transformers/src/main/kotlin/cc/unitmesh/cf/STSemantic.kt b/llm-modules/sentence-transformers/src/main/kotlin/cc/unitmesh/cf/STSemantic.kt
@@ -44,11 +44,13 @@ class STSemantic(
 
  val floatArray = outputTensor.floatBuffer.array()
  // floatArray is an inputIds.size * 384 array, we need to mean it to 384 * 1
- val meanArray = FloatArray(384)
- for (i in 0 until 384) {
+ // 1, shape, shape.length
+ val shapeSize = outputTensor.info.shape[2].toInt()
+ val meanArray = FloatArray(shapeSzie)
+ for (i in 0 until shapeSzie) {
  var sum = 0f
  for (j in inputIds.indices) {
- sum += floatArray[j * 384 + i]
+ sum += floatArray[j * shapeSzie + i]
  }
 
  meanArray[i] = sum / inputIds.size