From 401f55c76e572edb7e858e7adc0b0dfcfef1399f Mon Sep 17 00:00:00 2001 From: Phodal Huang Date: Thu, 9 Nov 2023 15:45:43 +0800 Subject: [PATCH] fix: fix item issue for mobile cases --- .../src/main/kotlin/cc/unitmesh/cf/STSemantic.kt | 8 +++++--- 1 file changed, 5 insertions(+), 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 index 69e61a34..0abc2849 100644 --- 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