summaryrefslogtreecommitdiffstats
path: root/toolkit/components/ml/content/SummarizerModel.sys.mjs
diff options
context:
space:
mode:
authorDaniel Baumann <daniel.baumann@progress-linux.org>2024-04-19 01:13:27 +0000
committerDaniel Baumann <daniel.baumann@progress-linux.org>2024-04-19 01:13:27 +0000
commit40a355a42d4a9444dc753c04c6608dade2f06a23 (patch)
tree871fc667d2de662f171103ce5ec067014ef85e61 /toolkit/components/ml/content/SummarizerModel.sys.mjs
parentAdding upstream version 124.0.1. (diff)
downloadfirefox-adbda400be353e676059e335c3c0aaf99e719475.tar.xz
firefox-adbda400be353e676059e335c3c0aaf99e719475.zip
Adding upstream version 125.0.1.upstream/125.0.1
Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
Diffstat (limited to 'toolkit/components/ml/content/SummarizerModel.sys.mjs')
-rw-r--r--toolkit/components/ml/content/SummarizerModel.sys.mjs160
1 files changed, 160 insertions, 0 deletions
diff --git a/toolkit/components/ml/content/SummarizerModel.sys.mjs b/toolkit/components/ml/content/SummarizerModel.sys.mjs
new file mode 100644
index 0000000000..7cac55d92f
--- /dev/null
+++ b/toolkit/components/ml/content/SummarizerModel.sys.mjs
@@ -0,0 +1,160 @@
+/* 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/. */
+
+/**
+ * @typedef {object} LazyImports
+ * @property {typeof import("../actors/MLEngineParent.sys.mjs").MLEngineParent} MLEngineParent
+ */
+
+/** @type {LazyImports} */
+const lazy = {};
+
+ChromeUtils.defineESModuleGetters(lazy, {
+ RemoteSettings: "resource://services-settings/remote-settings.sys.mjs",
+ TranslationsParent: "resource://gre/actors/TranslationsParent.sys.mjs",
+});
+
+ChromeUtils.defineLazyGetter(lazy, "console", () => {
+ return console.createInstance({
+ maxLogLevelPref: "browser.ml.logLevel",
+ prefix: "ML",
+ });
+});
+
+export class SummarizerModel {
+ /**
+ * The RemoteSettingsClient that downloads the summarizer model.
+ *
+ * @type {RemoteSettingsClient | null}
+ */
+ static #remoteClient = null;
+
+ /** @type {Promise<WasmRecord> | null} */
+ static #modelRecord = null;
+
+ /**
+ * The following constant controls the major version for wasm downloaded from
+ * Remote Settings. When a breaking change is introduced, Nightly will have these
+ * numbers incremented by one, but Beta and Release will still be on the previous
+ * version. Remote Settings will ship both versions of the records, and the latest
+ * asset released in that version will be used. For instance, with a major version
+ * of "1", assets can be downloaded for "1.0", "1.2", "1.3beta", but assets marked
+ * as "2.0", "2.1", etc will not be downloaded.
+ */
+ static MODEL_MAJOR_VERSION = 1;
+
+ /**
+ * Remote settings isn't available in tests, so provide mocked responses.
+ */
+ static mockRemoteSettings(remoteClient) {
+ lazy.console.log("Mocking remote client in SummarizerModel.");
+ SummarizerModel.#remoteClient = remoteClient;
+ SummarizerModel.#modelRecord = null;
+ }
+
+ /**
+ * Remove anything that could have been mocked.
+ */
+ static removeMocks() {
+ lazy.console.log("Removing mocked remote client in SummarizerModel.");
+ SummarizerModel.#remoteClient = null;
+ SummarizerModel.#modelRecord = null;
+ }
+ /**
+ * Download or load the model from remote settings.
+ *
+ * @returns {Promise<ArrayBuffer>}
+ */
+ static async getModel() {
+ const client = SummarizerModel.#getRemoteClient();
+
+ if (!SummarizerModel.#modelRecord) {
+ // Place the records into a promise to prevent any races.
+ SummarizerModel.#modelRecord = (async () => {
+ // Load the wasm binary from remote settings, if it hasn't been already.
+ lazy.console.log(`Getting the summarizer model record.`);
+
+ // TODO - The getMaxVersionRecords should eventually migrated to some kind of
+ // shared utility.
+ const { getMaxVersionRecords } = lazy.TranslationsParent;
+
+ /** @type {WasmRecord[]} */
+ const wasmRecords = await getMaxVersionRecords(client, {
+ // TODO - This record needs to be created with the engine wasm payload.
+ filters: { name: "summarizer-model" },
+ majorVersion: SummarizerModel.MODEL_MAJOR_VERSION,
+ });
+
+ if (wasmRecords.length === 0) {
+ // The remote settings client provides an empty list of records when there is
+ // an error.
+ throw new Error("Unable to get the models from Remote Settings.");
+ }
+
+ if (wasmRecords.length > 1) {
+ SummarizerModel.reportError(
+ new Error("Expected the ml engine to only have 1 record."),
+ wasmRecords
+ );
+ }
+ const [record] = wasmRecords;
+ lazy.console.log(
+ `Using ${record.name}@${record.release} release version ${record.version} first released on Fx${record.fx_release}`,
+ record
+ );
+ return record;
+ })();
+ }
+
+ try {
+ /** @type {{buffer: ArrayBuffer}} */
+ const { buffer } = await client.attachments.download(
+ await SummarizerModel.#modelRecord
+ );
+
+ return buffer;
+ } catch (error) {
+ SummarizerModel.#modelRecord = null;
+ throw error;
+ }
+ }
+
+ /**
+ * Lazily initializes the RemoteSettingsClient.
+ *
+ * @returns {RemoteSettingsClient}
+ */
+ static #getRemoteClient() {
+ if (SummarizerModel.#remoteClient) {
+ return SummarizerModel.#remoteClient;
+ }
+
+ /** @type {RemoteSettingsClient} */
+ const client = lazy.RemoteSettings("ml-model");
+
+ SummarizerModel.#remoteClient = client;
+
+ client.on("sync", async ({ data: { created, updated, deleted } }) => {
+ lazy.console.log(`"sync" event for ml-model`, {
+ created,
+ updated,
+ deleted,
+ });
+
+ // Remove all the deleted records.
+ for (const record of deleted) {
+ await client.attachments.deleteDownloaded(record);
+ }
+
+ // Remove any updated records, and download the new ones.
+ for (const { old: oldRecord } of updated) {
+ await client.attachments.deleteDownloaded(oldRecord);
+ }
+
+ // Do nothing for the created records.
+ });
+
+ return client;
+ }
+}