Proper embedding tokenization

agg23 · agg23 · commit 4d71a34f6cbc · 2025-10-17T20:28:38.000-07:00
diff --git a/.github/workflows/release_on_push.yml b/.github/workflows/release_on_push.yml
@@ -34,6 +34,7 @@ jobs:
         run: |
           mkdir -p mabl/src/main/assets
           curl -L -f -o mabl/src/main/assets/minilm-l6-v2-qint8-arm64.onnx "https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/onnx/model_qint8_arm64.onnx?download=true"
+          curl -L -f -o mabl/src/main/assets/minilm-l6-v2-tokenizer.json "https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/tokenizer.json?download=true"
 
       - name: Generate version for build
         run: |
diff --git a/build.sh b/build.sh
@@ -1,6 +1,9 @@
 if ! [ -f mabl/src/main/assets/minilm-l6-v2-qint8-arm64.onnx ]; then
   curl -L -o mabl/src/main/assets/minilm-l6-v2-qint8-arm64.onnx https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/onnx/model_qint8_arm64.onnx?download=true
 fi
+if ! [ -f mabl/src/main/assets/minilm-l6-v2-tokenizer.json ]; then
+  curl -L -o mabl/src/main/assets/minilm-l6-v2-tokenizer.json https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/tokenizer.json?download=true
+fi
 ./gradlew :plugins:demo:installDebug :plugins:aipinsystem:installDebug :plugins:system:installDebug :plugins:openai:installDebug :plugins:googlesearch:installDebug :mabl:installAipinDebug
 adb shell pm grant com.penumbraos.mabl.pin android.permission.CAMERA
 adb shell appops set com.penumbraos.mabl.pin MANAGE_EXTERNAL_STORAGE allow
diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
@@ -21,6 +21,7 @@ ktor-client = "3.0.0"
 kotlinx-serialization = "1.7.1"
 kotlinx-coroutines = "1.8.1"
 onnx-runtime = "1.20.0"
+sentence-embeddings = "v6"
 room = "2.7.2"
 jsoup = "1.17.2"
 # The first number needs to match the Kotlin version
@@ -64,6 +65,7 @@ ktor-serialization-kotlinx-json = { group = "io.ktor", name = "ktor-serializatio
 kotlinx-serialization-json = { group = "org.jetbrains.kotlinx", name = "kotlinx-serialization-json", version.ref = "kotlinx-serialization" }
 kotlinx-coroutines-android = { group = "org.jetbrains.kotlinx", name = "kotlinx-coroutines-android", version.ref = "kotlinx-coroutines" }
 onnx-runtime-android = { group = "com.microsoft.onnxruntime", name = "onnxruntime-android", version.ref = "onnx-runtime" }
+sentence-embeddings = { group = "io.gitlab.shubham0204", name = "sentence-embeddings", version.ref = "sentence-embeddings" }
 androidx-room-runtime = { group = "androidx.room", name = "room-runtime", version.ref = "room" }
 androidx-room-ktx = { group = "androidx.room", name = "room-ktx", version.ref = "room" }
 androidx-room-compiler = { group = "androidx.room", name = "room-compiler", version.ref = "room" }
diff --git a/mabl/build.gradle.kts b/mabl/build.gradle.kts
@@ -86,7 +86,9 @@ dependencies {
     implementation(libs.androidx.camera.camera2)
 
     implementation(libs.kotlinx.serialization.json)
+
     implementation(libs.onnx.runtime.android)
+    implementation(libs.sentence.embeddings)
 
     implementation(libs.androidx.core.ktx)
     implementation(libs.androidx.lifecycle.runtime.ktx)
diff --git a/mabl/src/main/assets/.gitignore b/mabl/src/main/assets/.gitignore
@@ -1 +1,2 @@
-*.onnx
+*.onnx
+*.json
diff --git a/mabl/src/main/java/com/penumbraos/mabl/services/ToolOrchestrator.kt b/mabl/src/main/java/com/penumbraos/mabl/services/ToolOrchestrator.kt
@@ -9,7 +9,6 @@ import com.penumbraos.mabl.sdk.IToolService
 import com.penumbraos.mabl.sdk.PluginType
 import com.penumbraos.mabl.sdk.ToolCall
 import com.penumbraos.mabl.sdk.ToolDefinition
-import java.io.ByteArrayOutputStream
 import java.util.concurrent.ConcurrentHashMap
 
 private const val TAG = "ToolOrchestrator"
@@ -60,9 +59,7 @@ class ToolOrchestrator(
         allConnected.await()
 
         try {
-            val outputStream = ByteArrayOutputStream()
-            context.assets.open("minilm-l6-v2-qint8-arm64.onnx").copyTo(outputStream)
-            toolSimilarityService.initialize(outputStream.toByteArray())
+            toolSimilarityService.initialize(context)
 
             // Precalculate embeddings for all available tools
             buildToolDefinitionsMap()
@@ -73,7 +70,7 @@ class ToolOrchestrator(
                 "Tool similarity service initialized successfully with ${allTools.size} tool embeddings precalculated"
             )
         } catch (e: Exception) {
-            Log.w(TAG, "Failed to initialize similarity service: ${e.message}")
+            Log.e(TAG, "Failed to initialize similarity service: $e")
         }
     }
 
diff --git a/mabl/src/main/java/com/penumbraos/mabl/services/ToolSimilarityService.kt b/mabl/src/main/java/com/penumbraos/mabl/services/ToolSimilarityService.kt
@@ -1,14 +1,13 @@
 package com.penumbraos.mabl.services
 
-import ai.onnxruntime.OnnxTensor
-import ai.onnxruntime.OrtEnvironment
-import ai.onnxruntime.OrtSession
+import android.content.Context
 import android.util.Log
 import com.penumbraos.mabl.sdk.ToolDefinition
+import com.penumbraos.mabl.util.SentenceEmbedding
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.withContext
-import java.nio.LongBuffer
+import java.io.ByteArrayOutputStream
 import java.util.concurrent.ConcurrentHashMap
 import kotlin.math.sqrt
 
@@ -21,60 +20,75 @@ data class OfflineIntentClassificationResult(
 )
 
 class ToolSimilarityService {
-    private var ortEnvironment: OrtEnvironment? = null
-    private var ortSession: OrtSession? = null
-    private val embeddingCache = ConcurrentHashMap<String, FloatArray>()
+    private val sentenceEmbedding = SentenceEmbedding()
     private val toolEmbeddingCache = ConcurrentHashMap<String, FloatArray>()
-    private var offlineCapableTools: List<ToolDefinition> = emptyList()
     private val intentExampleEmbeddingCache = ConcurrentHashMap<String, FloatArray>()
+    private var offlineCapableTools: List<ToolDefinition> = emptyList()
     private val scope = CoroutineScope(Dispatchers.IO)
 
     companion object {
-        private const val MAX_SEQUENCE_LENGTH = 512
         private const val SIMILARITY_THRESHOLD = 0.5f
         private const val INTENT_THRESHOLD = 0.55f
+//        private const val TOOL_CONFIRMATION_MARGIN = 0.05f
     }
 
-    suspend fun initialize(modelBytes: ByteArray) {
+    suspend fun initialize(context: Context) {
         withContext(scope.coroutineContext) {
-            ortEnvironment = OrtEnvironment.getEnvironment()
-            ortSession = ortEnvironment?.createSession(modelBytes)
+            val modelBytes = ByteArrayOutputStream().use { outputStream ->
+                context.assets.open("minilm-l6-v2-qint8-arm64.onnx").copyTo(outputStream)
+                outputStream.toByteArray()
+            }
+
+            val tokenizerBytes = ByteArrayOutputStream().use { outputStream ->
+                context.assets.open("minilm-l6-v2-tokenizer.json").copyTo(outputStream)
+                outputStream.toByteArray()
+            }
+
+            try {
+                sentenceEmbedding.init(
+                    modelBytes = modelBytes,
+                    tokenizerBytes = tokenizerBytes,
+                    useTokenTypeIds = true,
+                    outputTensorName = "last_hidden_state",
+                    normalizeEmbeddings = true
+                )
+            } catch (e: Exception) {
+                Log.w(TAG, "Failed to initialize tokenizer: ${e.message}")
+                null
+            }
         }
     }
 
     suspend fun precalculateToolEmbeddings(tools: List<ToolDefinition>) {
-        if (ortSession == null) return
-
         withContext(scope.coroutineContext) {
+            val offlineCandidates = mutableListOf<ToolDefinition>()
+            toolEmbeddingCache.clear()
+            intentExampleEmbeddingCache.clear()
+
             tools.forEach { tool ->
                 val toolText = buildToolText(tool)
-                val embedding = getEmbedding(toolText)
-                toolEmbeddingCache[tool.name] = embedding
+                toolEmbeddingCache[tool.name] = sentenceEmbedding.encode(toolText)
 
                 if (!tool.examples.isNullOrEmpty()) {
-                    offlineCapableTools += tool
+                    offlineCandidates += tool
                 }
 
                 tool.examples?.forEachIndexed { index, example ->
                     if (!example.isNullOrBlank()) {
                         val key = intentExampleKey(tool.name, index)
-                        intentExampleEmbeddingCache[key] = getEmbedding(example)
+                        intentExampleEmbeddingCache[key] = sentenceEmbedding.encode(example)
                     }
                 }
             }
+            offlineCapableTools = offlineCandidates
         }
     }
 
     suspend fun classifyIntent(
         userQuery: String,
     ): OfflineIntentClassificationResult? {
-        if (ortSession == null) {
-            Log.w(TAG, "Intent classification requested before model initialization")
-            return null
-        }
-
         return withContext(scope.coroutineContext) {
-            val queryEmbedding = getEmbedding(userQuery)
+            val queryEmbedding = sentenceEmbedding.encode(userQuery)
             var bestMatch: OfflineIntentClassificationResult? = null
 
             offlineCapableTools.forEach { tool ->
@@ -90,12 +104,28 @@ class ToolSimilarityService {
 
                     val key = intentExampleKey(tool.name, index)
                     val exampleEmbedding = intentExampleEmbeddingCache[key]
-                    if (exampleEmbedding == null) {
+                        ?: return@forEachIndexed
+
+                    val score = cosineSimilarity(queryEmbedding, exampleEmbedding)
+                    if (score < INTENT_THRESHOLD) {
                         return@forEachIndexed
                     }
-                    val score = cosineSimilarity(queryEmbedding, exampleEmbedding)
 
-                    if (score >= INTENT_THRESHOLD && (bestMatch == null || score > bestMatch!!.similarity)) {
+//                    val toolEmbedding = toolEmbeddingCache[tool.name]
+//                        ?: sentenceEmbedding.encode(buildToolText(tool)).also {
+//                            toolEmbeddingCache[tool.name] = it
+//                        }
+//                    val toolScore = cosineSimilarity(queryEmbedding, toolEmbedding)
+//                    Log.e(
+//                        "ToolSimilarityService",
+//                        "Intent classification result: ${tool.name} $score $toolScore"
+//                    )
+//
+//                    if (toolScore < INTENT_THRESHOLD + TOOL_CONFIRMATION_MARGIN) {
+//                        return@forEachIndexed
+//                    }
+
+                    if (bestMatch == null || score > bestMatch!!.similarity) {
                         val parameters = extractBooleanParameters(tool, userQuery)
                         bestMatch = OfflineIntentClassificationResult(tool, score, parameters)
                     }
@@ -111,18 +141,14 @@ class ToolSimilarityService {
         userQuery: String,
         maxTools: Int
     ): List<ToolDefinition> {
-        if (ortSession == null) {
-            throw IllegalStateException("Tool similarity service not initialized")
-        }
-
         return withContext(scope.coroutineContext) {
-            val queryEmbedding = getEmbedding(userQuery)
+            val queryEmbedding = sentenceEmbedding.encode(userQuery)
 
             val toolScores = tools.map { tool ->
                 val toolEmbedding = toolEmbeddingCache[tool.name] ?: run {
                     // Fallback: calculate embedding if not cached
                     val toolText = buildToolText(tool)
-                    getEmbedding(toolText)
+                    sentenceEmbedding.encode(toolText)
                 }
                 val similarity = cosineSimilarity(queryEmbedding, toolEmbedding)
 
@@ -141,36 +167,6 @@ class ToolSimilarityService {
         }
     }
 
-    private suspend fun getEmbedding(text: String): FloatArray {
-        val cacheKey = text.hashCode().toString()
-        embeddingCache[cacheKey]?.let { return it }
-
-        return withContext(scope.coroutineContext) {
-            val tokenIds = tokenizeText(text)
-            val inputIdsTensor = createInputTensor(tokenIds)
-            val tokenTypeIdsTensor = createTokenTypeIdsTensor(tokenIds.size)
-            val attentionMaskTensor = createAttentionMaskTensor(tokenIds.size)
-
-            val inputs = mapOf(
-                "input_ids" to inputIdsTensor,
-                "token_type_ids" to tokenTypeIdsTensor,
-                "attention_mask" to attentionMaskTensor
-            )
-            val outputs = ortSession?.run(inputs)
-
-            val embedding = outputs?.get(0)?.value as Array<*>
-            val floatEmbedding = (embedding[0] as Array<FloatArray>)[0]
-
-            inputIdsTensor.close()
-            tokenTypeIdsTensor.close()
-            attentionMaskTensor.close()
-            outputs.close()
-
-            embeddingCache[cacheKey] = floatEmbedding
-            floatEmbedding
-        }
-    }
-
     private fun buildToolText(tool: ToolDefinition): String {
         val builder = StringBuilder()
         builder.append(tool.name).append(" ")
@@ -238,57 +234,6 @@ class ToolSimilarityService {
         }
     }
 
-    private fun tokenizeText(text: String): IntArray {
-        val words = text.lowercase().split(Regex("\\W+"))
-        val tokens = mutableListOf<Int>()
-
-        words.forEach { word ->
-            if (word.isNotEmpty()) {
-                tokens.add(word.hashCode() % 30000)
-            }
-        }
-
-        return tokens.take(MAX_SEQUENCE_LENGTH).toIntArray()
-    }
-
-    private fun createInputTensor(tokenIds: IntArray): OnnxTensor {
-        val shape = longArrayOf(1, tokenIds.size.toLong())
-        val buffer = LongBuffer.allocate(tokenIds.size)
-
-        tokenIds.forEach { id ->
-            buffer.put(id.toLong())
-        }
-        buffer.flip()
-
-        return OnnxTensor.createTensor(ortEnvironment, buffer, shape)
-    }
-
-    private fun createTokenTypeIdsTensor(sequenceLength: Int): OnnxTensor {
-        val shape = longArrayOf(1, sequenceLength.toLong())
-        val buffer = LongBuffer.allocate(sequenceLength)
-
-        // All tokens are type 0 (single sentence)
-        repeat(sequenceLength) {
-            buffer.put(0L)
-        }
-        buffer.flip()
-
-        return OnnxTensor.createTensor(ortEnvironment, buffer, shape)
-    }
-
-    private fun createAttentionMaskTensor(sequenceLength: Int): OnnxTensor {
-        val shape = longArrayOf(1, sequenceLength.toLong())
-        val buffer = LongBuffer.allocate(sequenceLength)
-
-        // All tokens get attention (no padding in our case)
-        repeat(sequenceLength) {
-            buffer.put(1L)
-        }
-        buffer.flip()
-
-        return OnnxTensor.createTensor(ortEnvironment, buffer, shape)
-    }
-
     private fun cosineSimilarity(a: FloatArray, b: FloatArray): Float {
         if (a.size != b.size) return 0f
 
@@ -307,7 +252,6 @@ class ToolSimilarityService {
     }
 
     fun close() {
-        ortSession?.close()
-        ortEnvironment?.close()
+        sentenceEmbedding.close()
     }
 }
diff --git a/mabl/src/main/java/com/penumbraos/mabl/util/SentenceEmbedding.kt b/mabl/src/main/java/com/penumbraos/mabl/util/SentenceEmbedding.kt
