Source code
Revision control
Copy as Markdown
Other Tools
Test Info:
/* Any copyright is dedicated to the Public Domain.
"use strict";
const PREFIX = "inference";
const METRICS = [
`${PREFIX}-${PIPELINE_READY_LATENCY}`,
`${PREFIX}-${INITIALIZATION_LATENCY}`,
`${PREFIX}-${MODEL_RUN_LATENCY}`,
`${PREFIX}-${TOTAL_MEMORY_USAGE}`,
];
const journal = {};
for (let metric of METRICS) {
journal[metric] = [1];
}
const perfMetadata = {
owner: "GenAI Team",
name: "browser_ml_suggest_inference.js",
description: "Template test for ML suggest Inference Model",
options: {
default: {
perfherder: true,
perfherder_metrics: [
{
name: "inference-pipeline-ready-latency",
unit: "ms",
shouldAlert: false,
},
{
name: "inference-initialization-latency",
unit: "ms",
shouldAlert: false,
},
{
name: "inference-model-run-latency",
unit: "ms",
shouldAlert: false,
},
{
name: "inference-total-memory-usage",
unit: "ms",
shouldAlert: false,
},
],
verbose: true,
manifest: "perftest.toml",
manifest_flavor: "browser-chrome",
try_platform: ["linux", "mac", "win"],
},
},
};
requestLongerTimeout(120);
const CUSTOM_INTENT_OPTIONS = {
taskName: "text-classification",
featureId: "suggest-intent-classification",
engineId: "ml-suggest-intent",
modelId: "Mozilla/mobilebert-uncased-finetuned-LoRA-intent-classifier",
modelHubUrlTemplate: "{model}/{revision}",
dtype: "q8",
modelRevision: "main",
numThreads: 2,
timeoutMS: -1,
};
const CUSTOM_NER_OPTIONS = {
taskName: "token-classification",
featureId: "suggest-NER",
engineId: "ml-suggest-ner",
modelId: "Mozilla/distilbert-uncased-NER-LoRA",
modelHubUrlTemplate: "{model}/{revision}",
dtype: "q8",
modelRevision: "main",
numThreads: 2,
timeoutMS: -1,
};
const ROOT_URL =
const YELP_KEYWORDS_DATA = `${ROOT_URL}/yelp_val_keywords_data.json`;
const YELP_VAL_DATA = `${ROOT_URL}/yelp_val_generated_data.json`;
const NER_VAL_DATA = `${ROOT_URL}/named_entity_val_generated_data.json`;
async function get_val_data(inputDataPath) {
const response = await fetch(inputDataPath);
if (!response.ok) {
throw new Error(
`Failed to fetch data: ${response.statusText} from ${inputDataPath}`
);
}
return response.json();
}
async function read_data_by_type(type) {
let data;
if (type === "YELP_KEYWORDS_DATA") {
data = await get_val_data(YELP_KEYWORDS_DATA);
return data[0].subjects;
} else if (type === "YELP_VAL_DATA") {
data = await get_val_data(YELP_VAL_DATA);
return data.queries;
} else if (type === "NER_VAL_DATA") {
data = await get_val_data(NER_VAL_DATA);
return data.queries;
}
return [];
}
// Utility to write results to a local JSON file using IOUtils
async function writeResultsToFile(results, type) {
try {
const json = JSON.stringify(results, null, 2);
const OUTPUT_FILE_PATH = normalizePathForOS(
`${
Services.dirsvc.get("DfltDwnld", Ci.nsIFile).path
}/ML_output_${type}.json`
);
await IOUtils.writeUTF8(OUTPUT_FILE_PATH, json);
console.log("Results successfully written to:", OUTPUT_FILE_PATH);
} catch (error) {
console.error("Failed to write results to file:", error);
Assert.ok(false, "Failed to write results to file");
}
}
async function perform_inference(queries, type) {
// Ensure MLSuggest is initialized
await MLSuggest.initialize();
const results = [];
for (const query of queries) {
const suggestion = await MLSuggest.makeSuggestions(query);
const res = {
query,
intent: suggestion.intent,
city: suggestion.location.city,
state: suggestion.location.state,
};
results.push(res);
}
Assert.ok(
results.length === queries.length,
"results size should be equal to queries size."
);
// Write results to a file
await writeResultsToFile(results, type);
}
const runInference2 = async () => {
ChromeUtils.defineESModuleGetters(this, {
MLSuggest: "resource:///modules/urlbar/private/MLSuggest.sys.mjs",
});
// Override INTENT and NER options within MLSuggest
MLSuggest.INTENT_OPTIONS = CUSTOM_INTENT_OPTIONS;
MLSuggest.NER_OPTIONS = CUSTOM_NER_OPTIONS;
const modelHubRootUrl = Services.env.get("MOZ_MODELS_HUB");
if (!modelHubRootUrl) {
throw new Error(
"MOZ_MODELS_HUB is not set, you need to run with --hooks toolkit/components/ml/tests/tools/hook_local_hub.py"
);
}
info(`ModelHubRootUrl: ${modelHubRootUrl}`);
const { cleanup } = await perfSetup({
prefs: [
["browser.ml.modelHubRootUrl", modelHubRootUrl],
["javascript.options.wasm_lazy_tiering", true],
],
});
const TYPES = ["YELP_KEYWORDS_DATA", "YELP_VAL_DATA", "NER_VAL_DATA"];
for (const type of TYPES) {
info(`processing ${type} now`);
// Read data for the current type
const queries = await read_data_by_type(type);
if (!queries) {
info(`No queries found for type: ${type}`);
continue;
}
// Run inference for each query
await perform_inference(queries, type);
}
await MLSuggest.shutdown();
await EngineProcess.destroyMLEngine();
await cleanup();
};
/**
* Tests remote ml model
*/
add_task(async function test_ml_generic_pipeline() {
await runInference2();
reportMetrics(journal);
});