ml fixes and improvements

2026-05-28 21:05:01 +02:00
parent 2d23b9b33e
commit bdca6d78a4
4 changed files with 167 additions and 72 deletions
@@ -202,6 +202,11 @@ class GeckoInferenceRepository extends _$GeckoInferenceRepository {
                  .toList(),
        );

+        final idsByTitle = <String, List<String>>{};
+        for (final MapEntry(:key, :value) in unassignedDocumentsInput.entries) {
+          (idsByTitle[value] ??= []).add(key);
+        }
+
        final clusterResult = await clusters.mapNotNull(
          (cluster) => Future.wait(
            cluster.map((clusterTitles) async {
@@ -224,11 +229,7 @@ class GeckoInferenceRepository extends _$GeckoInferenceRepository {
                  },
                ),
                tabIds: originalTitles
-                    .map(
-                      (title) => unassignedDocumentsInput.entries
-                          .firstWhere((entry) => entry.value == title)
-                          .key,
-                    )
+                    .expand((title) => idsByTitle[title] ?? const <String>[])
                    .toList(),
              );
            }),
@@ -266,7 +267,8 @@ class GeckoInferenceRepository extends _$GeckoInferenceRepository {
          await _initialLoadComplete.future;

          final embeddings = await Result.fromAsync(
-            () async => await _service.generateDocumentEmbeddings(documents),
+            () async =>
+                await _service.generateDocumentEmbeddings(embeddingsToGenerate),
          );

          return embeddings;
@@ -276,8 +278,25 @@ class GeckoInferenceRepository extends _$GeckoInferenceRepository {
          return Result.failure(generatedEmbeddings.error!);
        }

+        final generated = generatedEmbeddings.value;
+        if (generated.length != embeddingsToGenerate.length) {
+          return Result.failure(
+            ErrorMessage(
+              source: 'Document Embeddings',
+              message: 'Unexpected embedding count',
+              details: {
+                'expected': embeddingsToGenerate.length,
+                'actual': generated.length,
+              },
+            ),
+          );
+        }
+
        for (var i = 0; i < embeddingsToGenerate.length; i++) {
-          embeddings[embeddingsToGenerate[i]] = generatedEmbeddings.value[i];
+          final embedding = generated[i];
+          final document = embeddingsToGenerate[i];
+          embeddings[document] = embedding;
+          _embeddingCache.set(document, embedding);
        }
      }

@@ -289,6 +308,15 @@ class GeckoInferenceRepository extends _$GeckoInferenceRepository {
      return Result.failure(
        const ErrorMessage(source: 'Document Embeddings', message: 'Timeout'),
      );
+    } catch (e, s) {
+      return Result.failure(
+        ErrorMessage(
+          source: 'Document Embeddings',
+          message: e.toString(),
+          details: e,
+          stackTrace: s,
+        ),
+      );
    }
  }

@@ -456,15 +484,18 @@ Future<List<String>?> containerTabSuggestions(
                .toList(),
          );

-      return suggestedTitles.mapNotNull(
-        (titles) => titles
-            .map(
-              (title) => unassignedTitles.value
-                  .firstWhere((tab) => tab.$2 == title)
-                  .$1,
-            )
-            .toList(),
-      );
+      final idsByTitle = <String, List<String>>{};
+      for (final (id, title) in unassignedTitles.value) {
+        (idsByTitle[title] ??= []).add(id);
+      }
+
+      return suggestedTitles.mapNotNull((titles) {
+        final tabIds = titles
+            .expand((title) => idsByTitle[title] ?? const <String>[])
+            .toList();
+
+        return tabIds.isEmpty ? null : tabIds;
+      });
    }
  }

@@ -4,7 +4,10 @@

 "use strict";

-const { createEngine } = ChromeUtils.importESModule("chrome://global/content/ml/EngineProcess.sys.mjs");
+const {
+    createEngine,
+    FEATURES,
+} = ChromeUtils.importESModule("chrome://global/content/ml/EngineProcess.sys.mjs");

 const ML_TASK_FEATURE_EXTRACTION = "feature-extraction";
 const ML_TASK_TEXT2TEXT = "text2text-generation";
@@ -15,7 +18,8 @@ const SMART_TAB_GROUPING_CONFIG = {
        timeoutMS: 2 * 60 * 1000, // 2 minutes
        taskName: ML_TASK_FEATURE_EXTRACTION,
        featureId: "smart-tab-embedding",
-        backend: "onnx",
+        engineId: FEATURES["smart-tab-embedding"].engineId,
+        backend: "onnx-native",
        fallbackBackend: "onnx",
    },
    topicGeneration: {
@@ -23,7 +27,8 @@ const SMART_TAB_GROUPING_CONFIG = {
        timeoutMS: 2 * 60 * 1000, // 2 minutes
        taskName: ML_TASK_TEXT2TEXT,
        featureId: "smart-tab-topic",
-        backend: "onnx",
+        engineId: FEATURES["smart-tab-topic"].engineId,
+        backend: "onnx-native",
        fallbackBackend: "onnx",
    },
    // dataConfig: {
@@ -117,6 +122,52 @@ function createProgressCallback(modelType, progressEmitter) {
    };
 }

+async function createMlEngine(engineConfig, progressCallback) {
+    const {
+        featureId,
+        engineId,
+        dtype,
+        taskName,
+        timeoutMS,
+        modelId,
+        modelRevision,
+        backend,
+        fallbackBackend,
+    } = engineConfig;
+    const initData = {
+        featureId,
+        engineId,
+        dtype,
+        taskName,
+        timeoutMS,
+        modelId,
+        modelRevision,
+        backend,
+    };
+
+    try {
+        return await createEngine(initData, progressCallback);
+    } catch (error) {
+        if (!fallbackBackend || fallbackBackend === backend) {
+            throw error;
+        }
+
+        try {
+            return await createEngine(
+                {
+                    ...initData,
+                    backend: fallbackBackend,
+                },
+                progressCallback
+            );
+        } catch (fallbackError) {
+            throw new Error(
+                `Failed to create ML engine with ${backend} (${error?.message || error}) or ${fallbackBackend} (${fallbackError?.message || fallbackError})`
+            );
+        }
+    }
+}
+
 this.ml = class extends ExtensionAPI {
    constructor(extension) {
        super(extension);
@@ -151,7 +202,7 @@ this.ml = class extends ExtensionAPI {
                        };

                        if (isEngineClosed(self.embeddingEngine)) {
-                            self.embeddingEngine = await createEngine(
+                            self.embeddingEngine = await createMlEngine(
                                SMART_TAB_GROUPING_CONFIG.embedding,
                                createProgressCallback("Embedding Model", self.progressEmitter)
                            );
@@ -168,30 +219,8 @@ this.ml = class extends ExtensionAPI {
                    },
                    async predictTopic(keywords, documents) {
                        if (isEngineClosed(self.topicEngine)) {
-                            const {
-                                featureId,
-                                engineId,
-                                dtype,
-                                taskName,
-                                timeoutMS,
-                                modelId,
-                                modelRevision,
-                                backend,
-                            } = SMART_TAB_GROUPING_CONFIG.topicGeneration;
-
-                            let initData = {
-                                featureId,
-                                engineId,
-                                dtype,
-                                taskName,
-                                timeoutMS,
-                                modelId,
-                                modelRevision,
-                                backend,
-                            };
-
-                            self.topicEngine = await createEngine(
-                                initData,
+                            self.topicEngine = await createMlEngine(
+                                SMART_TAB_GROUPING_CONFIG.topicGeneration,
                                createProgressCallback("Topic Generation Model", self.progressEmitter)
                            );
                        }
@@ -221,4 +250,4 @@ this.ml = class extends ExtensionAPI {
            }
        };
    }
-};
+};
@@ -18,7 +18,7 @@ function sendErrorForRequest(id) {
        port.postMessage({
            "id": id,
            "status": "error",
-            "error": error
+            "error": error?.message || String(error)
        });
    }
 }
@@ -31,27 +31,31 @@ browser.experiments.ml.onProgress.addListener((progressData) => {
 });

 port.onMessage.addListener(async (message) => {
-    let requestId = message["id"]
-    switch (message["action"]) {
-        case "predictDocumentTopic": {
-            const documents = message["args"];
-            const keywords = (documents.length > 1)
-                ? await browser.experiments.nlp.extractKeywords([documents.slice(0, 3).join(" ")])
-                : [[]];
+    const requestId = message["id"];

-            browser.experiments.ml.predictTopic(keywords[0], documents)
-                .then(sendJsonResultForRequest(requestId))
-                .catch(sendErrorForRequest(requestId));
+    try {
+        switch (message["action"]) {
+            case "predictDocumentTopic": {
+                const documents = message["args"];
+                const keywords = (documents.length > 1)
+                    ? await browser.experiments.nlp.extractKeywords([documents.slice(0, 3).join(" ")])
+                    : [[]];
+                const result = await browser.experiments.ml.predictTopic(keywords[0], documents);

-            break;
-        }
-        case "generateDocumentEmbeddings": {
-            const documents = message["args"];
-            await browser.experiments.ml.generateEmbeddings(documents)
-                .then(sendJsonResultForRequest(requestId))
-                .catch(sendErrorForRequest(requestId));
-
-            break;
+                sendJsonResultForRequest(requestId)(result);
+                break;
+            }
+            case "generateDocumentEmbeddings": {
+                const documents = message["args"];
+                const result = await browser.experiments.ml.generateEmbeddings(documents);
+
+                sendJsonResultForRequest(requestId)(result);
+                break;
+            }
+            default:
+                throw new Error(`Unsupported ML action: ${message["action"]}`);
        }
+    } catch (error) {
+        sendErrorForRequest(requestId)(error);
    }
 });
@@ -7,7 +7,11 @@
 package eu.weblibre.flutter_mozilla_components.feature

 import androidx.annotation.VisibleForTesting
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
@@ -26,10 +30,12 @@ object MLEngineFeature {
    private const val ML_ENGINE_REPORTER_EXTENSION_ID = "ml-engine@weblibre.eu"
    private const val ML_ENGINE_REPORTER_EXTENSION_URL = "resource://android/assets/extensions/ml_engine/"
    private const val ML_ENGINE_REPORTER_MESSAGING_ID = "mlEngine"
+    private const val REQUEST_TIMEOUT_MS = 130_000L

    private var nextRequestId: Int = 0
    private val requestHandlers = HashMap<Int, ResultConsumer<JSONObject>>()
    private val mutex = Mutex()
+    private val scope = CoroutineScope(SupervisorJob() + Dispatchers.Default)

    // Progress callback for ML operations
    var progressCallback: ((JSONObject) -> Unit)? = null
@@ -44,29 +50,54 @@ object MLEngineFeature {

    fun scheduleRequest(command: String, args: Any, callback: ResultConsumer<JSONObject>) {
        val message = JSONObject()
-        message.put("action", command);
+        message.put("action", command)
        message.put("args", args)

-        runBlocking {
+        val requestId = runBlocking {
            withContext(Dispatchers.Default) {
                mutex.withLock {
-                    message.put("id", nextRequestId)
+                    val requestId = nextRequestId
+                    message.put("id", requestId)

-                    requestHandlers[nextRequestId] = callback
+                    requestHandlers[requestId] = callback

                    nextRequestId += 1

-                    extensionController.sendBackgroundMessage(message)
+                    try {
+                        extensionController.sendBackgroundMessage(message)
+                    } catch (throwable: Throwable) {
+                        requestHandlers.remove(requestId)
+                        callback.error(
+                            "ML Engine",
+                            "Failed to schedule request",
+                            throwable.message,
+                        )
+                    }
+
+                    requestId
                }
            }
        }
+
+        scope.launch {
+            delay(REQUEST_TIMEOUT_MS)
+            val handler = mutex.withLock {
+                requestHandlers.remove(requestId)
+            }
+
+            handler?.error(
+                "ML Engine",
+                "Request timed out",
+                "No response received for $command",
+            )
+        }
    }

    private class PrefManagerReporterBackgroundMessageHandler() : MessageHandler {
        override fun onPortMessage(message: Any, port: Port) {
            runBlocking {
                withContext(Dispatchers.Default) {
-                    val messageJSON = message as JSONObject;
+                    val messageJSON = message as JSONObject

                    // Check if this is a progress message
                    if (messageJSON.has("type") && messageJSON.getString("type") == "mlProgress") {
@@ -86,7 +117,7 @@ object MLEngineFeature {
                            handler?.error(
                                "ML Engine",
                                "Failed to perform operation",
-                                message.getString("error")
+                                messageJSON.optString("error", "Unknown ML engine error")
                            )
                        }
                    }