Bug 1940906 - Select the speech recognition model per locale. r=media-playback-reviewers,ai-platform-reviewers,alwu,valentinp

Add a models.yaml manifest (gen_speech_models.py) so
LanguagesToSpeechModel (SpeechRecognitionModelMapping.{h,cpp}) picks a model
per locale prefix, overridable via media.webspeech.recognition.model.<prefix>
prefs, instead of always using a single hardcoded parakeet model.

Model size will be used in later patches.

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

Author: padenot

dom/media/webspeech/recognition/SpeechRecognitionModelMapping.cpp

@@ -0,0 +1,111 @@
+/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* vim: set ts=8  et sw=2 tw=80: */
+/* 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/. */
+
+#include "SpeechRecognitionModelMapping.h"
+
+#include "SpeechRecognitionModels.h"
+#include "mozilla/Assertions.h"
+#include "mozilla/Preferences.h"
+#include "nsFmtString.h"
+#include "nsReadableUtils.h"
+
+namespace mozilla::dom {
+
+nsCString SpeechModelIdentifier::ToString() const {
+  return nsFmtCString("{}/{}/{}", mModelName.get(), mFileName.get(),
+                      mRevision.get());
+}
+
+nsCString LanguagesToSpeechModelId(const nsTArray<nsCString>& aLanguages) {
+  // Determine the primary-subtag locale prefix (e.g. "en" from "en-US"). This
+  // discards script/region subtags, so e.g. zh-Hans and zh-Hant both collapse
+  // to zh; proper BCP-47 canonicalization is tracked in bug 2060247.
+  nsCString prefix;
+  if (!aLanguages.IsEmpty()) {
+    prefix = aLanguages[0];
+    int32_t dash = prefix.FindChar('-');
+    if (dash != kNotFound) {
+      prefix.Truncate(dash);
+    }
+  }
+
+  // A pref may override the default model for a locale prefix:
+  // media.webspeech.recognition.model.<prefix> (or .multilingual for the
+  // fallback). Empty prefix uses the multilingual fallback.
+  nsAutoCString prefKey("media.webspeech.recognition.model.");
+  prefKey.Append(prefix.IsEmpty() ? "multilingual"_ns : prefix);
+  nsAutoCString prefModelId;
+  Preferences::GetCString(prefKey.get(), prefModelId);
+
+  if (!prefModelId.IsEmpty()) {
+    for (const auto& m : kSpeechRecognitionModels) {
+      if (m.id && prefModelId.Equals(m.id)) {
+        return nsCString(m.id);
+      }
+    }
+  }
+
+  // No usable pref: pick the default model whose locale list matches the
+  // prefix, falling back to the default fallback model (empty locale list).
+  const SpeechRecognitionModelInfo* fallback = nullptr;
+  for (const auto& m : kSpeechRecognitionModels) {
+    if (!m.id) {
+      break;
+    }
+    if (!m.locales[0]) {
+      if (m.is_default && !fallback) {
+        fallback = &m;
+      }
+      continue;
+    }
+    for (const char* const* l = m.locales; *l; ++l) {
+      if (!prefix.IsEmpty() &&
+          StringBeginsWith(prefix, nsDependentCString(*l))) {
+        if (m.is_default) {
+          return nsCString(m.id);
+        }
+      }
+    }
+  }
+
+  if (fallback) {
+    return nsCString(fallback->id);
+  }
+
+  MOZ_ASSERT_UNREACHABLE("No default model found in kSpeechRecognitionModels");
+  return {};
+}
+
+bool ResolveSpeechModelId(const nsACString& aId, SpeechModelIdentifier& aOut) {
+  for (const auto& m : kSpeechRecognitionModels) {
+    if (!m.id) {
+      break;
+    }
+    if (aId.Equals(m.id)) {
+      aOut = {nsCString(m.repo), nsCString(m.filename), nsCString(m.revision),
+              m.size_mb};
+      return true;
+    }
+  }
+  return false;
+}
+
+uint32_t SpeechModelSizeMB(const nsACString& aModel,
+                           const nsACString& aRevision,
+                           const nsACString& aFilename) {
+  for (const auto& m : kSpeechRecognitionModels) {
+    if (!m.id) {
+      break;
+    }
+    if (aModel.Equals(m.repo) && aRevision.Equals(m.revision) &&
+        aFilename.Equals(m.filename)) {
+      return m.size_mb;
+    }
+  }
+  return 0;
+}
+
+}  // namespace mozilla::dom

dom/media/webspeech/recognition/SpeechRecognitionModelMapping.h

@@ -0,0 +1,56 @@
+/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* vim: set ts=8  et sw=2 tw=80: */
+/* 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/. */
+
+#ifndef DOM_MEDIA_WEBSPEECH_RECOGNITION_SPEECHRECOGNITIONMODELMAPPING_H_
+#define DOM_MEDIA_WEBSPEECH_RECOGNITION_SPEECHRECOGNITIONMODELMAPPING_H_
+
+#include "nsString.h"
+#include "nsTArray.h"
+
+namespace mozilla::dom {
+
+// ModelHub engine id and task used for every speech recognition model. Shared
+// so the download path in the parent process and the read-only availability
+// queries in the utility process agree on them.
+inline constexpr auto kSpeechRecognitionEngineId = "parakeet-gguf"_ns;
+inline constexpr auto kSpeechRecognitionTask = "speech-recognition"_ns;
+
+// The concrete ModelHub artifact a set of requested languages maps to.
+struct SpeechModelIdentifier {
+  nsCString mModelName;
+  nsCString mFileName;
+  nsCString mRevision = "main"_ns;
+  uint32_t mSizeMB = 0;
+  nsCString ToString() const;
+};
+
+// Maps a set of requested BCP-47 languages to a model id (the model table's
+// own "id" field, e.g. "english_tdt_ctc_q6_k"), honoring the
+// media.webspeech.recognition.model.<prefix> pref override. The id, not the
+// artifact itself, is what the utility process sends onward: it can only ever
+// select among the ids this function can return, never name a
+// model/revision/filename directly. The model table (generated from
+// models.yaml) lives only in this translation unit's .cpp.
+nsCString LanguagesToSpeechModelId(const nsTArray<nsCString>& aLanguages);
+
+// Expands aId (as returned by LanguagesToSpeechModelId) to the concrete
+// ModelHub artifact it names. Returns false if aId is unknown. Used trusted-
+// side (see SpeechModelResolver) to validate an id supplied by the utility
+// process before it is passed to ModelHub.
+bool ResolveSpeechModelId(const nsACString& aId, SpeechModelIdentifier& aOut);
+
+// Approximate download size in MB of the model artifact identified by
+// aModel/aRevision/aFilename, for display in the download permission prompt.
+// Returns 0 for an unknown artifact. Keyed by the artifact (not languages) so
+// the download gate, which only sees the generic model coordinates, can look
+// it up. The model table lives only in this translation unit's .cpp.
+uint32_t SpeechModelSizeMB(const nsACString& aModel,
+                           const nsACString& aRevision,
+                           const nsACString& aFilename);
+
+}  // namespace mozilla::dom
+
+#endif  // DOM_MEDIA_WEBSPEECH_RECOGNITION_SPEECHRECOGNITIONMODELMAPPING_H_

dom/media/webspeech/recognition/SpeechRecognitionParent.cpp

@@ -10,6 +10,7 @@
 #include <chrono>
 #include <thread>
 
+#include "SpeechRecognitionModelMapping.h"
 #include "mozilla/Logging.h"
 #include "mozilla/Mutex.h"
 #include "mozilla/Preferences.h"
@@ -44,29 +45,9 @@ static LazyLogModule gSpeechRecognitionParentLog("SpeechRecognitionParent");
 #define LOGE(fmt, ...) \
   MOZ_LOG_FMT(gSpeechRecognitionParentLog, LogLevel::Error, fmt, ##__VA_ARGS__)
 
-// Task name the model hub files speech recognition models under.
-static constexpr auto kSpeechRecognitionTask = "speech-recognition"_ns;
-
 // Sample rate the Parakeet models operate at.
 static constexpr int32_t PARAKEET_SAMPLE_RATE = 16000;
 
-SpeechRecognitionParent::ModelIdentifier
-SpeechRecognitionParent::LanguagesToModelIdentifier(
-    const nsTArray<nsCString>&) {
-  // For now this ignores the requested languages and always returns a single
-  // hardcoded model. Per-language selection and generalization to more models
-  // land in a later patch in this stack.
-  // mudler/parakeet.cpp cache-aware streaming GGUF, hosted on the Mozilla
-  // model hub under asr-test/parakeet.
-  return {"asr-test/parakeet"_ns, "realtime_eou_120m-v1-q5_k.gguf"_ns,
-          "main"_ns};
-}
-
-nsCString SpeechRecognitionParent::ModelIdentifier::ToString() const {
-  return nsFmtCString("{}/{}/{}", mModelName.get(), mFileName.get(),
-                      mRevision.get());
-}
-
 void SpeechRecognitionParent::ResolveOrRejectInitOnIPCThread(
     InitResolver&& aResolver, bool aSuccess) {
   if (GetActorEventTarget()->IsOnCurrentThread()) {
@@ -133,14 +114,14 @@ mozilla::ipc::IPCResult SpeechRecognitionParent::RecvIsModelAvailable(
                     "RecvIsModelAvailable requires at least one language");
   }
 
-  nsCString modelId = LanguagesToModelIdentifier(aLanguages).ToString();
+  nsCString modelId = dom::LanguagesToSpeechModelId(aLanguages);
   LOGD("{} languages: {} mapped to id={}", __func__,
        fmt::join(aLanguages, ", "), modelId.get());
 
   return RunHWInferenceBoolQuery(
       __func__,
       [modelId](hwinference::HWInferenceChild* aChild) {
-        return aChild->SendIsModelAvailable(nsCString(kSpeechRecognitionTask),
+        return aChild->SendIsModelAvailable(dom::kSpeechRecognitionTask,
                                             modelId);
       },
       std::move(aResolver), mIsModelAvailableRequest);
@@ -152,7 +133,7 @@ mozilla::ipc::IPCResult SpeechRecognitionParent::RecvInstallModels(
     return IPC_FAIL(this, "RecvInstallModels requires at least one language");
   }
 
-  nsCString modelId = LanguagesToModelIdentifier(aLanguages).ToString();
+  nsCString modelId = dom::LanguagesToSpeechModelId(aLanguages);
   LOGD("{} languages: {} mapped to id={}", __func__,
        fmt::join(aLanguages, ", "), modelId.get());
 
@@ -162,7 +143,7 @@ mozilla::ipc::IPCResult SpeechRecognitionParent::RecvInstallModels(
   return RunHWInferenceBoolQuery(
       __func__,
       [modelId, contentId](hwinference::HWInferenceChild* aChild) {
-        return aChild->SendInstallModel(kSpeechRecognitionTask, modelId, 0,
+        return aChild->SendInstallModel(dom::kSpeechRecognitionTask, modelId, 0,
                                         contentId);
       },
       std::move(aResolver), mInstallModelRequest);
@@ -211,7 +192,7 @@ void SpeechRecognitionParent::RetrieveModel(InitResolver&& aResolver) {
   nsCString modelId;
   {
     MutexAutoLock lock(mLock);
-    modelId = LanguagesToModelIdentifier(nsTArray{mLanguage}).ToString();
+    modelId = dom::LanguagesToSpeechModelId(nsTArray{mLanguage});
   }
 
   LOGD("{} Checking model is installed: id={}", __func__, modelId.get());
@@ -220,8 +201,7 @@ void SpeechRecognitionParent::RetrieveModel(InitResolver&& aResolver) {
   // permission doorhanger. start() requires the model to already be
   // installed, so check that before FetchModelFile() rather than letting
   // GetModelFile download it on demand.
-  hwInferenceChild
-      ->SendIsModelInstalled(nsCString(kSpeechRecognitionTask), modelId)
+  hwInferenceChild->SendIsModelInstalled(dom::kSpeechRecognitionTask, modelId)
       ->Then(GetCurrentSerialEventTarget(), __func__,
              [self = RefPtr{this}, aResolver = std::move(aResolver),
               modelId](hwinference::PHWInferenceChild::IsModelInstalledPromise::
@@ -256,8 +236,7 @@ void SpeechRecognitionParent::FetchModelFile(const nsCString& aModelId,
 
   LOGD("{} Requesting model: id={}", __func__, aModelId.get());
 
-  hwInferenceChild
-      ->SendGetModelFile(nsCString(kSpeechRecognitionTask), aModelId)
+  hwInferenceChild->SendGetModelFile(dom::kSpeechRecognitionTask, aModelId)
       ->Then(
           GetCurrentSerialEventTarget(), __func__,
           [self = RefPtr{this}, aResolver = std::move(aResolver)](

dom/media/webspeech/recognition/SpeechRecognitionParent.h

@@ -60,16 +60,6 @@ class SpeechRecognitionParent final : public PSpeechRecognitionParent {
 
   void ActorDestroy(ActorDestroyReason aReason) override;
 
-  struct ModelIdentifier {
-    nsCString mModelName;
-    nsCString mFileName;
-    nsCString mRevision = "main"_ns;
-    nsCString ToString() const;
-  };
-
-  ModelIdentifier LanguagesToModelIdentifier(
-      const nsTArray<nsCString>& aLanguages);
-
   void ResolveOrRejectInitOnIPCThread(InitResolver&& aResolver, bool aSuccess)
       MOZ_EXCLUDES(mLock);
 

dom/media/webspeech/recognition/gen_speech_models.py

@@ -0,0 +1,85 @@
+# 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/.
+
+"""Generate SpeechRecognitionModels.h and SpeechRecognitionModelDisplayInfo.sys.mjs
+from models.yaml."""
+
+import json
+
+import yaml
+
+
+def gen_models_header(output, input):
+    with open(input) as f:
+        data = yaml.safe_load(f)
+
+    models = data["models"]
+
+    output.write(
+        "/* This file is generated by gen_speech_models.py. Do not edit. */\n"
+        "#pragma once\n\n"
+        "#include <stdint.h>\n\n"
+        "namespace mozilla::dom {\n\n"
+        "struct SpeechRecognitionModelInfo {\n"
+        "  const char* id;\n"
+        "  const char* const* locales; // null-terminated; empty (first==nullptr) = fallback\n"
+        "  const char* repo;\n"
+        "  const char* filename;\n"
+        "  const char* revision;\n"
+        "  uint32_t size_mb;\n"
+        "  const char* quant;\n"
+        "  uint32_t latency_ms;\n"
+        "  bool is_default;\n"
+        "  bool is_streaming; // true = streaming RNN-T, false = offline CTC batch\n"
+        "};\n\n"
+    )
+
+    # Emit per-model locale arrays.
+    for m in models:
+        locales = m["locales"]
+        entries = (
+            ", ".join(f'"{l}"' for l in locales) + (", " if locales else "") + "nullptr"
+        )
+        output.write(
+            f"static const char* const kSpeechModelLocales_{m['id']}[] = {{{entries}}};\n"
+        )
+
+    output.write("\n")
+
+    # Emit the table, sentinel-terminated.
+    output.write(
+        "static const SpeechRecognitionModelInfo kSpeechRecognitionModels[] = {\n"
+    )
+    for m in models:
+        is_default = "true" if m.get("default", False) else "false"
+        is_streaming = "true" if m.get("streaming", True) else "false"
+        output.write(
+            f'  {{"{m["id"]}", kSpeechModelLocales_{m["id"]},\n'
+            f'   "{m["repo"]}", "{m["filename"]}", "{m["revision"]}",\n'
+            f'   {m["size_mb"]}, "{m["quant"]}", {m["latency_ms"]}, {is_default}, {is_streaming}}},\n'
+        )
+    output.write(
+        "  {nullptr, nullptr, nullptr, nullptr, nullptr, 0, nullptr, 0, false}\n"
+    )
+    output.write("};\n\n")
+
+    output.write("}  // namespace mozilla::dom\n")
+
+
+def gen_models_display_info(output, input):
+    """Generate a JS module mapping each model's downloaded file name to the
+    human-readable name and upstream Hugging Face model card shown in
+    about:addons' "Manage On-Device AI Models" panel."""
+    with open(input) as f:
+        data = yaml.safe_load(f)
+
+    info = {
+        m["filename"]: {"name": m["display_name"], "hfUrl": m["hf_url"]}
+        for m in data["models"]
+    }
+
+    output.write(
+        "/* This file is generated by gen_speech_models.py. Do not edit. */\n\n"
+        "export default " + json.dumps(info, indent=2) + ";\n"
+    )