Bug 1940906 - Allow looking up per-file display info for on-device models in about:addons. r=extension-reviewers,ai-platform-reviewers,rpl,valentinp

Currently, about:addons relies on a specific way to organize ModelHub
and also how Hugging Face works: ModelHub slugs mirror HuggingFace slugs
(same `org/model`, a bit like on github).

For SpeechRecognition, it's a little bit different: we take multiple
models that come from different Hugging Face pages from one Hugging Face
repo (https://huggingface.co/mudler/parakeet-cpp-gguf/). They are
conversions of models from Nvidia. We'd like to link to the original
repo (that has much more informative info), and we also want better
names, so we'll be in a position to say "Speech Recognition model
(multilingual)" instead of the inscrutable model name (with a link to go
further for the nerds / interested people).

This is a simple idirection layer that any developer of a feature that
uses AI models can use to display things better. For Speech Recognition,
the information comes from a small js module generated from a yaml. This
yaml also generates some C++ for other things, so we can't risk being
out of sync here.

Differential Revision: https://phabricator.services.mozilla.com/D316337

Author: padenot

toolkit/components/ml/content/Utils.sys.mjs

@@ -1086,6 +1086,38 @@ export function featureEngineIdToFluentId(engineId) {
   return null;
 }
 
+/**
+ * Looks up display info (a human-readable name and a homepage URL) for one
+ * of a downloaded model's files, from the feature registered for one of
+ * engineIds. Features register their file display info via
+ * FEATURES[...].fileDisplayInfoModule, a module whose default export maps
+ * file name to {name, hfUrl}.
+ *
+ * @param {Array<string>} [engineIds] A model with no recorded engineIds has no
+ *   feature to look up, and gets no display info.
+ * @param {Array<string>} files
+ * @returns {Promise<{name: string, hfUrl: string}|null>}
+ */
+export async function fileDisplayInfoForEngineIds(engineIds, files) {
+  for (const config of Object.values(lazy.FEATURES)) {
+    if (
+      !config.fileDisplayInfoModule ||
+      !engineIds?.includes(config.engineId)
+    ) {
+      continue;
+    }
+    const { default: info } = await ChromeUtils.importESModule(
+      config.fileDisplayInfoModule
+    );
+    for (const file of files) {
+      if (info[file]) {
+        return info[file];
+      }
+    }
+  }
+  return null;
+}
+
 /**
  * Generates a random uuid to use where Services.uuid is not available,
  * for instance pipelines

toolkit/mozapps/extensions/internal/ModelHubProvider.sys.mjs

@@ -13,6 +13,7 @@ ChromeUtils.defineESModuleGetters(lazy, {
   ModelHub: "chrome://global/content/ml/ModelHub.sys.mjs",
   isAddonEngineId: "chrome://global/content/ml/Utils.sys.mjs",
   engineIdToAddonId: "chrome://global/content/ml/Utils.sys.mjs",
+  fileDisplayInfoForEngineIds: "chrome://global/content/ml/Utils.sys.mjs",
 });
 
 XPCOMUtils.defineLazyPreferenceGetter(
@@ -36,6 +37,7 @@ export class ModelHubAddonWrapper {
   updateDate;
   modelIconURL;
   engineIds;
+  displayInfo;
 
   constructor(params) {
     this.#provider = params.provider;
@@ -47,6 +49,7 @@ export class ModelHubAddonWrapper {
     this.updateDate = params.updateDate;
     this.modelIconURL = params.modelIconURL;
     this.engineIds = params.engineIds ?? [];
+    this.displayInfo = params.displayInfo ?? null;
   }
 
   async uninstall() {
@@ -70,11 +73,17 @@ export class ModelHubAddonWrapper {
   }
 
   get name() {
+    if (this.displayInfo) {
+      return this.displayInfo.name;
+    }
     const parts = this.model.split("/");
     return parts.slice(2).join("/");
   }
 
   get modelHomepageURL() {
+    if (this.displayInfo) {
+      return this.displayInfo.hfUrl;
+    }
     // Model card URL for models downloaded from "model-hub.mozilla.org" should point to the corresponding "https://huggingface.co" url.
     return `https://${this.model}/`.replace(
       "https://model-hub.mozilla.org/",
@@ -193,13 +202,17 @@ export const ModelHubProvider = {
     const models = await this.modelHub.listModels();
 
     for (const model of models) {
-      const { metadata } = await this.modelHub.listFiles({
+      const { files, metadata } = await this.modelHub.listFiles({
         model: model.name,
         revision: model.revision,
       });
 
       const modelIconURL = await this.modelHub.getOwnerIcon(model.name);
       const id = this.getWrapperIdForModel(model);
+      const displayInfo = await lazy.fileDisplayInfoForEngineIds(
+        metadata.engineIds,
+        files.map(file => file.path)
+      );
 
       const wrapper = new ModelHubAddonWrapper({
         provider: this,
@@ -211,6 +224,7 @@ export const ModelHubProvider = {
         updateDate: new Date(metadata.updateDate),
         modelIconURL,
         engineIds: metadata.engineIds,
+        displayInfo,
       });
       this.cache.set(wrapper.id, wrapper);
     }

toolkit/mozapps/extensions/test/xpcshell/test_ModelHubProvider.js

@@ -6,6 +6,7 @@ const { ModelHubProvider } = ChromeUtils.importESModule(
   "resource://gre/modules/addons/ModelHubProvider.sys.mjs"
 );
 ChromeUtils.defineESModuleGetters(this, {
+  FEATURES: "chrome://global/content/ml/EngineProcess.sys.mjs",
   addonIdToEngineId: "chrome://global/content/ml/Utils.sys.mjs",
   engineIdToAddonId: "chrome://global/content/ml/Utils.sys.mjs",
   isAddonEngineId: "chrome://global/content/ml/Utils.sys.mjs",
@@ -88,6 +89,9 @@ add_task(
       },
     ];
     const mockListFilesResult = {
+      // No display info is registered for these, so the wrappers fall back to
+      // deriving their name and homepage from the model string.
+      files: [{ path: "model.onnx" }, { path: "config.json" }],
       metadata: {
         totalSize: 2048,
         lastUsed: new Date("2023-10-01T12:00:00Z"),
@@ -120,6 +124,7 @@ add_task(
       .resolves(mockListFilesResult)
       .onSecondCall()
       .resolves({
+        files: mockListFilesResult.files,
         metadata: {
           ...mockListFilesResult.metadata,
           // Setting engineIds to undefined to confirm that it is
@@ -375,6 +380,7 @@ add_task(
     ];
 
     const mockListFilesResult = {
+      files: [{ path: "model.onnx" }],
       metadata: {
         totalSize: 2048,
         lastUsed: new Date("2023-10-01T12:00:00Z"),
@@ -396,6 +402,7 @@ add_task(
       .resolves(mockListFilesResult)
       .onSecondCall()
       .resolves({
+        files: mockListFilesResult.files,
         metadata: {
           ...mockListFilesResult.metadata,
           // Setting engineIds to undefined to confirm that it is
@@ -432,3 +439,55 @@ add_task(
     sandbox.restore();
   }
 );
+
+add_task(
+  {
+    pref_set: [[LOCAL_MODEL_MANAGEMENT_ENABLED_PREF, true]],
+  },
+  async function test_modelhub_provider_file_display_info() {
+    let sandbox = sinon.createSandbox();
+    ModelHubProvider.clearAddonCache();
+
+    // Speech recognition registers display info for its model files, keyed by
+    // file name. Drive the lookup off the registry itself rather than repeating
+    // its contents here.
+    const { default: displayInfoByFile } = ChromeUtils.importESModule(
+      "resource://gre/modules/SpeechRecognitionModelDisplayInfo.sys.mjs"
+    );
+    const [fileName, displayInfo] = Object.entries(displayInfoByFile)[0];
+
+    const mockModel = {
+      name: "model-hub.mozilla.org/mudler/parakeet-cpp-gguf",
+      revision: "mockRevision",
+      engineIds: [],
+    };
+
+    sandbox.stub(ModelHubProvider.modelHub, "listModels").resolves([mockModel]);
+    sandbox.stub(ModelHubProvider.modelHub, "listFiles").resolves({
+      files: [{ path: fileName }],
+      metadata: {
+        totalSize: 2048,
+        lastUsed: new Date("2023-10-01T12:00:00Z"),
+        updateDate: 0,
+        engineIds: [FEATURES["speech-recognition"].engineId],
+      },
+    });
+    sandbox.stub(ModelHubProvider.modelHub, "getOwnerIcon").resolves(null);
+
+    const [modelWrapper] = await AddonManager.getAddonsByTypes(["mlmodel"]);
+    Assert.ok(modelWrapper, "Got a model AddonWrapper instance");
+
+    Assert.equal(
+      modelWrapper.name,
+      displayInfo.name,
+      "Expect the name to come from the registered display info"
+    );
+    Assert.equal(
+      modelWrapper.modelHomepageURL,
+      displayInfo.hfUrl,
+      "Expect the homepage URL to come from the registered display info"
+    );
+
+    sandbox.restore();
+  }
+);