summaryrefslogtreecommitdiffstats
path: root/toolkit/components/ml/content/SummarizerModel.sys.mjs
blob: 7cac55d92f5a8fdf2161aca0cf2b2c0f4d246061 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
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;
  }
}