MrBlake/MrbWebLibre
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 2 Wiki Activity

ml fixes and improvements

Browse Source
This commit is contained in:
Fabian Freund
2026-05-28 21:05:01 +02:00
parent 2d23b9b33e
commit bdca6d78a4
4 changed files with 167 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/flutter_mozilla_components/android/src/main/assets/extensions/ml_engine/api/ml.js
+58 -29
View File
@@ -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 {
}
};
}
};
};
packages/flutter_mozilla_components/android/src/main/assets/extensions/ml_engine/background.js
+24 -20
View File
@@ -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);
}
});
packages/flutter_mozilla_components/android/src/main/kotlin/eu/weblibre/flutter_mozilla_components/feature/MLEngineFeature.kt
+38 -7
View File
@@ -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")
)
}
}