Source code
Revision control
Copy as Markdown
Other Tools
# 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
"""Generate SpeechRecognitionModels.h and SpeechRecognitionModelDisplayInfo.sys.mjs
from models.yaml."""
import json
import yaml
def locale_array(name, locales):
"""A null-terminated static array of locale strings, 8 per line."""
for locale in locales:
if not isinstance(locale, str):
raise ValueError(
f"{name}: locale {locale!r} is not a string. YAML reads a bare "
"no, yes, on, off, y or n as a boolean, so a locale tag that "
"collides with one must be quoted in models.yaml."
)
items = [f'"{locale}"' for locale in locales] + ["nullptr"]
lines = [" " + ", ".join(items[i : i + 8]) for i in range(0, len(items), 8)]
return f"static const char* const {name}[] = {{\n" + ",\n".join(lines) + "};\n"
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* supported_locales; // null-terminated\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_streaming; // true = streaming RNN-T, false = offline CTC batch\n"
"};\n\n"
)
# Emit per-model locale arrays.
for m in models:
output.write(
locale_array(
f"kSpeechModelSupportedLocales_{m['id']}", m["supported_locales"]
)
)
output.write("\n")
# Emit the table, sentinel-terminated, in models.yaml (preference) order.
output.write(
"static const SpeechRecognitionModelInfo kSpeechRecognitionModels[] = {\n"
)
for m in models:
is_streaming = "true" if m.get("streaming", True) else "false"
output.write(
f' {{"{m["id"]}", kSpeechModelSupportedLocales_{m["id"]},\n'
f' "{m["repo"]}", "{m["filename"]}", "{m["revision"]}",\n'
f' {m["size_mb"]}, "{m["quant"]}", {m["latency_ms"]}, {is_streaming}}},\n'
)
output.write(" {nullptr, nullptr, nullptr, nullptr, nullptr, 0, nullptr, 0}\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"
)