From bdca6d78a41a2c52d891c0706f1b16c52dcb0208 Mon Sep 17 00:00:00 2001 From: Fabian Freund Date: Thu, 28 May 2026 21:05:01 +0200 Subject: [PATCH] ml fixes and improvements --- .../domain/repositories/gecko_inference.dart | 63 ++++++++++---- .../assets/extensions/ml_engine/api/ml.js | 87 ++++++++++++------- .../assets/extensions/ml_engine/background.js | 44 +++++----- .../feature/MLEngineFeature.kt | 45 ++++++++-- 4 files changed, 167 insertions(+), 72 deletions(-) diff --git a/apps/weblibre/lib/features/geckoview/features/tabs/domain/repositories/gecko_inference.dart b/apps/weblibre/lib/features/geckoview/features/tabs/domain/repositories/gecko_inference.dart index e657c90f..b195e155 100644 --- a/apps/weblibre/lib/features/geckoview/features/tabs/domain/repositories/gecko_inference.dart +++ b/apps/weblibre/lib/features/geckoview/features/tabs/domain/repositories/gecko_inference.dart @@ -202,6 +202,11 @@ class GeckoInferenceRepository extends _$GeckoInferenceRepository { .toList(), ); + final idsByTitle = >{}; + 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 []) .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?> containerTabSuggestions( .toList(), ); - return suggestedTitles.mapNotNull( - (titles) => titles - .map( - (title) => unassignedTitles.value - .firstWhere((tab) => tab.$2 == title) - .$1, - ) - .toList(), - ); + final idsByTitle = >{}; + for (final (id, title) in unassignedTitles.value) { + (idsByTitle[title] ??= []).add(id); + } + + return suggestedTitles.mapNotNull((titles) { + final tabIds = titles + .expand((title) => idsByTitle[title] ?? const []) + .toList(); + + return tabIds.isEmpty ? null : tabIds; + }); } } diff --git a/packages/flutter_mozilla_components/android/src/main/assets/extensions/ml_engine/api/ml.js b/packages/flutter_mozilla_components/android/src/main/assets/extensions/ml_engine/api/ml.js index da605a1a..fd341cd9 100644 --- a/packages/flutter_mozilla_components/android/src/main/assets/extensions/ml_engine/api/ml.js +++ b/packages/flutter_mozilla_components/android/src/main/assets/extensions/ml_engine/api/ml.js @@ -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 { } }; } -}; \ No newline at end of file +}; diff --git a/packages/flutter_mozilla_components/android/src/main/assets/extensions/ml_engine/background.js b/packages/flutter_mozilla_components/android/src/main/assets/extensions/ml_engine/background.js index 1fbd0788..d6fee053 100644 --- a/packages/flutter_mozilla_components/android/src/main/assets/extensions/ml_engine/background.js +++ b/packages/flutter_mozilla_components/android/src/main/assets/extensions/ml_engine/background.js @@ -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); } }); diff --git a/packages/flutter_mozilla_components/android/src/main/kotlin/eu/weblibre/flutter_mozilla_components/feature/MLEngineFeature.kt b/packages/flutter_mozilla_components/android/src/main/kotlin/eu/weblibre/flutter_mozilla_components/feature/MLEngineFeature.kt index 2beeed99..bcd7c054 100644 --- a/packages/flutter_mozilla_components/android/src/main/kotlin/eu/weblibre/flutter_mozilla_components/feature/MLEngineFeature.kt +++ b/packages/flutter_mozilla_components/android/src/main/kotlin/eu/weblibre/flutter_mozilla_components/feature/MLEngineFeature.kt @@ -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>() 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) { 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") ) } }