container topic inference

This commit is contained in:
Fabian Freund
2025-07-20 00:37:51 +02:00
parent 2a3e0b6b6d
commit dc14a7d718
29 changed files with 1117 additions and 12 deletions
@@ -0,0 +1,151 @@
/* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
* file, You can obtain one at http://mozilla.org/MPL/2.0/. */
"use strict";
const { createEngine } = ChromeUtils.importESModule("chrome://global/content/ml/EngineProcess.sys.mjs");
const ML_TASK_FEATURE_EXTRACTION = "feature-extraction";
const ML_TASK_TEXT2TEXT = "text2text-generation";
const SMART_TAB_GROUPING_CONFIG = {
embedding: {
dtype: "q8",
timeoutMS: 2 * 60 * 1000, // 2 minutes
taskName: ML_TASK_FEATURE_EXTRACTION,
featureId: "smart-tab-embedding",
backend: "onnx",
},
topicGeneration: {
dtype: "q8",
timeoutMS: 2 * 60 * 1000, // 2 minutes
taskName: ML_TASK_TEXT2TEXT,
featureId: "smart-tab-topic",
backend: "onnx",
},
// dataConfig: {
// titleKey: "label",
// descriptionKey: "description",
// },
// clustering: {
// dimReductionMethod: null, // Not completed.
// clusterImplementation: CLUSTER_METHODS.KMEANS,
// clusteringTriesPerK: 3,
// anchorMethod: ANCHOR_METHODS.FIXED,
// pregroupedHandlingMethod: PREGROUPED_HANDLING_METHODS.EXCLUDE,
// pregroupedSilhouetteBoost: 2, // Relative weight of the cluster's score and all other cluster's combined
// suggestOtherTabsMethod: SUGGEST_OTHER_TABS_METHODS.NEAREST_NEIGHBOR,
// },
};
/**
* Generate model input from keywords and documents
* @param {string []} keywords
* @param {string []} documents
*/
function createModelInput(keywords, documents) {
if (!keywords || keywords.length === 0) {
return `Topic from keywords: titles: \n${documents.join(" \n")}`;
}
return `Topic from keywords: ${keywords.join(", ")}. titles: \n${documents.join(" \n")}`;
}
/**
* One artifact of the LLM output is that sometimes words are duplicated
* This function cuts the phrase when it sees the first duplicate word.
* Handles simple singluar / plural duplicates (-s only).
* @param {string} phrase Input phrase
* @returns {string} phrase cut before any duplicate word
*/
function cutAtDuplicateWords(phrase) {
if (!phrase.length) {
return phrase;
}
const wordsSet = new Set();
const wordList = phrase.split(" ");
for (let i = 0; i < wordList.length; i++) {
let baseWord = wordList[i].toLowerCase();
if (baseWord.length > 3) {
if (baseWord.slice(-1) === "s") {
baseWord = baseWord.slice(0, -1);
}
}
if (wordsSet.has(baseWord)) {
// We are seeing a baseWord word. Exit with just the words so far and don't
// add any new words
return wordList.slice(0, i).join(" ");
}
wordsSet.add(baseWord);
}
return phrase; // return original phrase
}
/**
*
* @param {MLEngine} engine the engine to check
* @return {boolean} true if the engine has not been initialized or closed
*/
function isEngineClosed(engine) {
return !engine || engine?.engineStatus === "closed";
}
this.ml = class extends ExtensionAPI {
getAPI(context) {
return {
experiments: {
ml: {
async containerTopic(keywords, documents) {
if (isEngineClosed(this.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,
};
this.topicEngine = await createEngine(initData);
}
const inputArgs = createModelInput(
keywords,
documents
);
const requestInfo = {
inputArgs,
runOptions: {
max_length: 6,
},
};
const request = {
args: [requestInfo.inputArgs],
options: requestInfo.runOptions,
};
const res = await this.topicEngine.run(request);
const generated = cutAtDuplicateWords((res[0]["generated_text"] || "").trim());
return generated;
}
}
}
};
}
};
@@ -0,0 +1,29 @@
/* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
* file, You can obtain one at http://mozilla.org/MPL/2.0/. */
"use strict";
const { KeywordExtractor } = ChromeUtils.importESModule(
"chrome://global/content/ml/NLPUtils.sys.mjs"
);
this.nlp = class extends ExtensionAPI {
getAPI(context) {
return {
experiments: {
nlp: {
async extractKeywords(corpus, maxKeywords = 3) {
try {
const keywordExtractor = new KeywordExtractor();
const keywords = keywordExtractor.fitTransform(corpus, maxKeywords);
return keywords;
} catch (error) {
throw new ExtensionError(`Keyword extraction failed: ${error.message}`);
}
},
}
}
};
}
};
@@ -0,0 +1,38 @@
'use strict';
const port = browser.runtime.connectNative("mlEngine");
function sendJsonResultForRequest(id) {
return function (result) {
port.postMessage({
"id": id,
"status": "success",
"result": result
})
}
}
function sendErrorForRequest(id) {
return function (error) {
console.error(error);
port.postMessage({
"id": id,
"status": "error",
"error": error
});
}
}
port.onMessage.addListener(async (message) => {
let requestId = message["id"]
switch (message["action"]) {
case "getContainerTopic":
const documents = message["args"];
const keywords = await browser.experiments.nlp.extractKeywords([documents.slice(0, 3).join(" ")]);
browser.experiments.ml.containerTopic(keywords[0], documents)
.then(sendJsonResultForRequest(requestId))
.catch(sendErrorForRequest(requestId))
break
}
});
@@ -0,0 +1,61 @@
{
"manifest_version": 2,
"name": "ml-engine",
"version": "1.0",
"description": "WebLibre ML Engine",
"browser_specific_settings": {
"gecko": {
"id": "ml-engine@weblibre.eu"
}
},
"experiment_apis": {
"nlp": {
"schema": "schema.json",
"parent": {
"scopes": [
"addon_parent"
],
"script": "api/nlp.js",
"paths": [
[
"experiments",
"nlp"
]
]
}
},
"ml": {
"schema": "schema.json",
"parent": {
"scopes": [
"addon_parent"
],
"script": "api/ml.js",
"paths": [
[
"experiments",
"ml"
]
]
}
}
},
"background": {
"scripts": [
"background.js"
]
},
"optional_permissions": [
"trialML"
],
"permissions": [
"nativeMessaging",
"nativeMessagingFromContent",
"geckoViewAddons",
"cookies",
"menus",
"scripting",
"storage",
"<all_urls>"
]
}
@@ -0,0 +1,61 @@
[
{
"namespace": "experiments.nlp",
"description": "Natural Language Processing utilities",
"functions": [
{
"name": "extractKeywords",
"type": "function",
"description": "Extract keywords from a corpus of text documents",
"async": true,
"parameters": [
{
"name": "corpus",
"type": "array",
"items": {
"type": "string"
},
"description": "Array of text documents to extract keywords from"
},
{
"name": "maxKeywords",
"type": "integer",
"optional": true,
"default": 3,
"description": "Maximum number of keywords to extract per document"
}
]
}
]
},
{
"namespace": "experiments.ml",
"description": "Machine Learning utilities",
"functions": [
{
"name": "containerTopic",
"type": "function",
"description": "Generate topic from keywords and documents using ML engine",
"async": true,
"parameters": [
{
"name": "keywords",
"type": "array",
"items": {
"type": "string"
},
"description": "Array of keywords to generate topic from"
},
{
"name": "documents",
"type": "array",
"items": {
"type": "string"
},
"description": "Array of document titles/content"
}
]
}
]
}
]
@@ -9,6 +9,7 @@ import eu.weblibre.flutter_mozilla_components.feature.ContainerProxyFeature
import eu.weblibre.flutter_mozilla_components.feature.CookieManagerFeature
import eu.weblibre.flutter_mozilla_components.feature.PrefManagerFeature
import eu.weblibre.flutter_mozilla_components.feature.BrowserExtensionFeature
import eu.weblibre.flutter_mozilla_components.feature.MLEngineFeature
import eu.weblibre.flutter_mozilla_components.pigeons.BrowserExtensionEvents
import mozilla.components.browser.engine.gecko.GeckoEngine
import mozilla.components.browser.engine.gecko.fetch.GeckoViewFetchClient
@@ -39,7 +40,7 @@ object EngineProvider {
builder.extensionsWebAPIEnabled(true)
// Disable output for now to improve performance
//builder.consoleOutput(true)
builder.consoleOutput(true)
runtime = GeckoRuntime.create(context, builder.build())
}
@@ -57,6 +58,7 @@ object EngineProvider {
PrefManagerFeature.install(it)
ContainerProxyFeature.install(it)
BrowserExtensionFeature.install(it, extensionEvents)
MLEngineFeature.install(it)
}
}
@@ -27,6 +27,7 @@ import eu.weblibre.flutter_mozilla_components.pigeons.GeckoDownloadsApi
import eu.weblibre.flutter_mozilla_components.pigeons.GeckoEngineSettingsApi
import eu.weblibre.flutter_mozilla_components.pigeons.GeckoFindApi
import eu.weblibre.flutter_mozilla_components.pigeons.GeckoIconsApi
import eu.weblibre.flutter_mozilla_components.pigeons.GeckoMlApi
import eu.weblibre.flutter_mozilla_components.pigeons.GeckoPrefApi
import eu.weblibre.flutter_mozilla_components.pigeons.GeckoSelectionActionController
import eu.weblibre.flutter_mozilla_components.pigeons.GeckoSelectionActionEvents
@@ -152,6 +153,7 @@ class GeckoBrowserApiImpl : GeckoBrowserApi {
GeckoTabsApi.setUp(_flutterPluginBinding.binaryMessenger, GeckoTabsApiImpl())
GeckoIconsApi.setUp(_flutterPluginBinding.binaryMessenger, GeckoIconsApiImpl())
GeckoCookieApi.setUp(_flutterPluginBinding.binaryMessenger, GeckoCookieApiImpl())
GeckoMlApi.setUp(_flutterPluginBinding.binaryMessenger, GeckoMlApiImpl())
GeckoPrefApi.setUp(_flutterPluginBinding.binaryMessenger, GeckoPrefApiImpl())
GeckoContainerProxyApi.setUp(_flutterPluginBinding.binaryMessenger, GeckoContainerProxyApiImpl())
GeckoFindApi.setUp(_flutterPluginBinding.binaryMessenger, GeckoFindApiImpl())
@@ -0,0 +1,27 @@
package eu.weblibre.flutter_mozilla_components.api
import eu.weblibre.flutter_mozilla_components.feature.MLEngineFeature
import eu.weblibre.flutter_mozilla_components.feature.ResultConsumer
import eu.weblibre.flutter_mozilla_components.pigeons.GeckoMlApi
import org.json.JSONArray
import org.json.JSONObject
class GeckoMlApiImpl : GeckoMlApi {
private fun List<String>?.toJson(): JSONArray {
return JSONArray().apply {
this@toJson?.forEach { put(it) }
}
}
override fun getContainerTopic(titles: List<String>, callback: (Result<String>) -> Unit) {
MLEngineFeature.scheduleRequest("getContainerTopic", titles.toJson(), object : ResultConsumer<JSONObject> {
override fun success(result: JSONObject) {
callback(Result.success(result.getString("result")))
}
override fun error(errorCode: String, errorMessage: String?, errorDetails: Any?) {
callback(Result.failure(Exception("$errorCode $errorMessage $errorDetails")))
}
})
}
}
@@ -0,0 +1,126 @@
/*
* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
* file, You can obtain one at https://mozilla.org/MPL/2.0/.
*/
package eu.weblibre.flutter_mozilla_components.feature
import androidx.annotation.VisibleForTesting
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlinx.coroutines.withContext
import mozilla.components.concept.engine.webextension.MessageHandler
import mozilla.components.concept.engine.webextension.Port
import mozilla.components.concept.engine.webextension.WebExtension
import mozilla.components.concept.engine.webextension.WebExtensionRuntime
import mozilla.components.support.base.log.logger.Logger
import mozilla.components.support.webextensions.BuiltInWebExtensionController
import org.json.JSONObject
object MLEngineFeature {
private val logger = Logger("ml-engine")
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 var nextRequestId: Int = 0
private val requestHandlers = HashMap<Int, ResultConsumer<JSONObject>>()
private val mutex = Mutex()
@VisibleForTesting
// This is an internal var to make it mutable for unit testing purposes only
internal var extensionController = BuiltInWebExtensionController(
ML_ENGINE_REPORTER_EXTENSION_ID,
ML_ENGINE_REPORTER_EXTENSION_URL,
ML_ENGINE_REPORTER_MESSAGING_ID,
)
fun scheduleRequest(command: String, args: Any, callback: ResultConsumer<JSONObject>) {
val message = JSONObject()
message.put("action", command);
message.put("args", args)
runBlocking {
withContext(Dispatchers.Default) {
mutex.withLock {
message.put("id", nextRequestId)
requestHandlers[nextRequestId] = callback
nextRequestId += 1
extensionController.sendBackgroundMessage(message)
}
}
}
}
private class PrefManagerReporterBackgroundMessageHandler() : MessageHandler {
override fun onPortMessage(message: Any, port: Port) {
runBlocking {
withContext(Dispatchers.Default) {
mutex.withLock {
val messageJSON = message as JSONObject;
val requestId = messageJSON.getInt("id")
val status = messageJSON.getString("status")
if (status == "success") {
requestHandlers[requestId]?.success(message)
} else {
requestHandlers[requestId]?.error(
"ML Engine",
"Failed to perform operation",
message.getString("error")
)
}
}
}
}
}
}
/**
* Installs the web extension in the runtime through the WebExtensionRuntime install method
*
* @param runtime a WebExtensionRuntime.
* @param productName a custom product name used to automatically label reports. Defaults to
* "android-components".
*/
fun install(runtime: WebExtensionRuntime) {
extensionController.registerBackgroundMessageHandler(
PrefManagerReporterBackgroundMessageHandler(),
)
extensionController.install(
runtime,
onSuccess = {
logger.debug("Installed ml-engine webextension: ${it.id}")
grantPermissions(runtime, it)
},
onError = { throwable ->
logger.error("Failed to install ml-engine webextension: ", throwable)
},
)
}
private fun grantPermissions(runtime: WebExtensionRuntime, extension: WebExtension) {
val permissions = listOf("trialML")
val origins = emptyList<String>() // Add any host permissions if needed
runtime.addOptionalPermissions(
ML_ENGINE_REPORTER_EXTENSION_ID,
permissions,
origins,
onSuccess = { grantedExtension ->
logger.debug("Successfully granted permissions to extension: ${grantedExtension.id}")
},
onError = { throwable ->
logger.error("Failed to grant permissions to extension: ", throwable)
}
)
}
}
@@ -1,4 +1,4 @@
// Autogenerated from Pigeon (v25.3.2), do not edit directly.
// Autogenerated from Pigeon (v25.5.0), do not edit directly.
// See also: https://pub.dev/packages/pigeon
@file:Suppress("UNCHECKED_CAST", "ArrayInDataClass")
@@ -3424,6 +3424,42 @@ interface GeckoPrefApi {
}
}
/** Generated interface from Pigeon that represents a handler of messages from Flutter. */
interface GeckoMlApi {
fun getContainerTopic(titles: List<String>, callback: (Result<String>) -> Unit)
companion object {
/** The codec used by GeckoMlApi. */
val codec: MessageCodec<Any?> by lazy {
GeckoPigeonCodec()
}
/** Sets up an instance of `GeckoMlApi` to handle messages through the `binaryMessenger`. */
@JvmOverloads
fun setUp(binaryMessenger: BinaryMessenger, api: GeckoMlApi?, messageChannelSuffix: String = "") {
val separatedMessageChannelSuffix = if (messageChannelSuffix.isNotEmpty()) ".$messageChannelSuffix" else ""
run {
val channel = BasicMessageChannel<Any?>(binaryMessenger, "dev.flutter.pigeon.flutter_mozilla_components.GeckoMlApi.getContainerTopic$separatedMessageChannelSuffix", codec)
if (api != null) {
channel.setMessageHandler { message, reply ->
val args = message as List<Any?>
val titlesArg = args[0] as List<String>
api.getContainerTopic(titlesArg) { result: Result<String> ->
val error = result.exceptionOrNull()
if (error != null) {
reply.reply(GeckoPigeonUtils.wrapError(error))
} else {
val data = result.getOrNull()
reply.reply(GeckoPigeonUtils.wrapResult(data))
}
}
}
} else {
channel.setMessageHandler(null)
}
}
}
}
}
/** Generated interface from Pigeon that represents a handler of messages from Flutter. */
interface GeckoBrowserExtensionApi {
fun getMarkdown(htmlList: List<String>, callback: (Result<List<Any>>) -> Unit)
@@ -18,6 +18,7 @@ export 'src/domain/services/gecko_engine_settings.dart';
export 'src/domain/services/gecko_event.dart';
export 'src/domain/services/gecko_find_in_page.dart';
export 'src/domain/services/gecko_icon.dart';
export 'src/domain/services/gecko_ml.dart';
export 'src/domain/services/gecko_pref.dart';
export 'src/domain/services/gecko_readerable.dart';
export 'src/domain/services/gecko_selection_action.dart';
@@ -0,0 +1,15 @@
/*
* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
* file, You can obtain one at https://mozilla.org/MPL/2.0/.
*/
import 'package:flutter_mozilla_components/src/pigeons/gecko.g.dart';
final _apiInstance = GeckoMlApi();
class GeckoMlService {
Future<String> getContainerTopic(Set<String> titles, {int maxCount = 10}) {
return _apiInstance.getContainerTopic(titles.take(maxCount).toList());
}
}
@@ -1,4 +1,4 @@
// Autogenerated from Pigeon (v25.3.2), do not edit directly.
// Autogenerated from Pigeon (v25.5.0), do not edit directly.
// See also: https://pub.dev/packages/pigeon
// ignore_for_file: public_member_api_docs, non_constant_identifier_names, avoid_as, unused_import, unnecessary_parenthesis, prefer_null_aware_operators, omit_local_variable_types, unused_shown_name, unnecessary_import, no_leading_underscores_for_local_identifiers
@@ -3955,6 +3955,48 @@ class GeckoPrefApi {
}
}
class GeckoMlApi {
/// Constructor for [GeckoMlApi]. The [binaryMessenger] named argument is
/// available for dependency injection. If it is left null, the default
/// BinaryMessenger will be used which routes to the host platform.
GeckoMlApi({BinaryMessenger? binaryMessenger, String messageChannelSuffix = ''})
: pigeonVar_binaryMessenger = binaryMessenger,
pigeonVar_messageChannelSuffix = messageChannelSuffix.isNotEmpty ? '.$messageChannelSuffix' : '';
final BinaryMessenger? pigeonVar_binaryMessenger;
static const MessageCodec<Object?> pigeonChannelCodec = _PigeonCodec();
final String pigeonVar_messageChannelSuffix;
Future<String> getContainerTopic(List<String> titles) async {
final String pigeonVar_channelName = 'dev.flutter.pigeon.flutter_mozilla_components.GeckoMlApi.getContainerTopic$pigeonVar_messageChannelSuffix';
final BasicMessageChannel<Object?> pigeonVar_channel = BasicMessageChannel<Object?>(
pigeonVar_channelName,
pigeonChannelCodec,
binaryMessenger: pigeonVar_binaryMessenger,
);
final Future<Object?> pigeonVar_sendFuture = pigeonVar_channel.send(<Object?>[titles]);
final List<Object?>? pigeonVar_replyList =
await pigeonVar_sendFuture as List<Object?>?;
if (pigeonVar_replyList == null) {
throw _createConnectionError(pigeonVar_channelName);
} else if (pigeonVar_replyList.length > 1) {
throw PlatformException(
code: pigeonVar_replyList[0]! as String,
message: pigeonVar_replyList[1] as String?,
details: pigeonVar_replyList[2],
);
} else if (pigeonVar_replyList[0] == null) {
throw PlatformException(
code: 'null-error',
message: 'Host platform returned null value for non-null return value.',
);
} else {
return (pigeonVar_replyList[0] as String?)!;
}
}
}
class GeckoBrowserExtensionApi {
/// Constructor for [GeckoBrowserExtensionApi]. The [binaryMessenger] named argument is
/// available for dependency injection. If it is left null, the default
@@ -953,6 +953,12 @@ abstract class GeckoPrefApi {
void resetPrefs(List<String>? preferenceNames);
}
@HostApi()
abstract class GeckoMlApi {
@async
String getContainerTopic(List<String> titles);
}
@HostApi()
abstract class GeckoBrowserExtensionApi {
@async