Commit 3f094fa7 authored by Tarek Ziadé's avatar Tarek Ziadé Committed by tziade@mozilla.com
Browse files

Bug 1955060 - Add onnx-native to Firefox AI Runtime r=padenot

parent 98a59c02
Loading
Loading
Loading
Loading
+9 −2
Original line number Diff line number Diff line
@@ -131,7 +131,7 @@
                <div>
                  <label for="device">Device:</label>
                  <select id="device">
                    <option value="wasm" selected>cpu</option>
                    <option value="cpu" selected>cpu</option>
                    <option value="gpu">gpu</option>
                  </select>
                </div>
@@ -143,7 +143,6 @@
                  </span>
                  <input id="numRuns" type="number" min="1" value="1" />
                </div>

                <div>
                  <label for="numThreads">Number of threads:</label>
                  <span class="tooltip"
@@ -152,6 +151,14 @@
                  </span>
                  <select id="numThreads"></select>
                </div>
                <div>
                  <label for="backend">Backend:</label>
                  <select id="backend"></select>
                </div>
                <div>
                  <label for="timeout">Timeout:</label>
                  <input id="timeout" type="number" min="30" value="30" />
                </div>
                <div>
                  <label for="inferencePad">Input data:</label>
                  <span class="tooltip"
+202 −14
Original line number Diff line number Diff line
@@ -33,7 +33,6 @@ const MODEL_HUB_ROOT_URL = Services.prefs.getStringPref(
const MODEL_HUB_URL_TEMPLATE = Services.prefs.getStringPref(
  "browser.ml.modelHubUrlTemplate"
);
const THIRTY_SECONDS = 30 * 1000;

let modelHub = null;
let modelCache = null;
@@ -84,6 +83,171 @@ const TINY_ARTICLE =
 * Presets for the pad
 */
const INFERENCE_PAD_PRESETS = {
  "summarization-large": {
    modelId: "onnx-community/Qwen2.5-0.5B-Instruct",
    task: "text-generation",
    modelRevision: "main",
    modelHub: "huggingface",
    dtype: "q8",
    device: "cpu",
    backend: "onnx",
    inputArgs: [
      "Summarize: The exploration of space has always been a subject of fascination and curiosity for humanity. \nFrom the earliest  stargazers mapping the heavens to modern astronomers utilizing cutting-edge telescopes, the quest to \nunderstand the  cosmos has driven technological and scientific advancement. In the 20th century, the space race between \nthe United States and the Soviet Union culminated in humanity's first steps on the moon, marking a new era of space \nexploration.\n\nToday, space exploration is no longer confined to government agencies. Private companies, such as SpaceX and Blue \nOrigin, are pioneering new methods of space travel, with ambitions ranging from commercial space tourism to colonizing \nMars. These ventures have sparked debates about the ethical, environmental, and economic implications of privatizing \nspace.\n\nThe potential benefits of space exploration are vast. Satellites have revolutionized communication, weather forecasting, \nand navigation. Space-based observations have provided critical data for understanding climate change and protecting \nEarth’s ecosystems. Moreover, the possibility of mining asteroids for rare minerals could address resource shortages on \nour planet.\n\nHowever, the challenges are equally significant. Space travel remains extraordinarily expensive and risky. The harsh \nenvironment of space poses threats to human health, including exposure to radiation and the psychological effects of \nprolonged isolation. Additionally, the accumulation of space debris presents a growing hazard to both current and future \nmissions.\n\nInternational collaboration has proven to be a powerful tool in overcoming some of these challenges. The International \nSpace Station (ISS) serves as a testament to what can be achieved when nations work together, pooling resources and \nexpertise. Future projects, such as the planned Lunar Gateway and Mars missions, highlight the importance of cooperative \nefforts in pushing the boundaries of human capability.\n\nAs humanity looks to the stars, questions about our role in the universe arise. What responsibilities do we have as \nstewards of our own planet and potential explorers of others? How can we ensure that space exploration benefits all of \nhumanity and not just a privileged few?\n\nThe journey into space is a reflection of our innate desire to explore and understand. It challenges us to innovate, to \ndream, and to collaborate. While the road ahead is fraught with difficulties, the potential rewards—both tangible and \nintangible—make the endeavor worthwhile. With careful planning, international cooperation, and ethical considerations, \nspace exploration can continue to inspire and benefit humanity for generations to come.",
    ],
    runOptions: {
      max_new_tokens: 195,
      min_new_tokens: 200,
      return_full_text: false,
      return_tensors: false,
      do_sample: false,
    },
    timeout: 120,
  },
  "feature-large": {
    inputArgs: [
      [
        "The quick brown fox jumps over the lazy dog.",
        "She sells seashells by the seashore.",
        "He walked slowly across the deserted street at dawn.",
        "A red rose blossomed in the middle of the garden.",
        "The sound of rain on the roof lulled me to sleep.",
        "We must never forget the lessons of history.",
        "The sun sets, painting the sky in bright shades of orange and purple.",
        "He found a hidden treasure chest in the old abandoned house.",
        "They enjoyed a picnic by the river on a sunny afternoon.",
        "Mathematics is the language of the universe.",
        "Music has the power to heal and transform hearts.",
        "A gentle breeze rustled the leaves of the towering oak tree.",
        "Many species of birds migrate to warmer climates each winter.",
        "The city lights sparkled like stars in the distance.",
        "Her bright smile lit up the room when she entered.",
        "A library is a portal to countless worlds and possibilities.",
        "The smell of fresh coffee filled the cozy café.",
        "He missed the sound of ocean waves crashing on the shore.",
        "The distant mountains were capped with snow even in summer.",
        "An empty page can represent infinite possibilities.",
        "We watched the fireworks explode in vibrant colors above us.",
        "The cat watched a butterfly dance in the garden.",
        "My grandmother taught me how to bake delicious apple pies.",
        "The conference was an opportunity to share groundbreaking research.",
        "They adopted a small puppy from the local animal shelter.",
        "Snow covered the streets, making them look like a winter wonderland.",
        "He wrote a letter to express the depth of his feelings.",
        "A curious child asked endless questions about the world.",
        "They dreamed of visiting every continent before turning forty.",
        "The aroma of freshly baked bread welcomed them home.",
        "Her voice resonated with passion as she spoke of her dreams.",
        "He overcame his fears and made a bold career change.",
        "Learning a new language can open doors to unexpected opportunities.",
        "A gentle melody drifted from the old piano in the corner.",
        "The festival brought the entire community together for a weekend of fun.",
        "He discovered an ancient map hidden beneath the floorboards.",
        "She overcame adversity with resilience and determination.",
        "The child's laughter echoed in the empty park.",
        "A well-written book can transport you to distant worlds.",
        "They found solace in the stillness of the forest.",
        "A strong partnership can often accomplish more than an individual.",
        "His grandmother had the best recipe for homemade cookies.",
        "The large oak tree provided shade on a scorching summer day.",
        "Innovation thrives when diverse minds come together.",
        "The smell of wet paint filled the newly renovated room.",
        "A random act of kindness can brighten a stranger's day.",
        "He spent hours stargazing, marveling at the vastness of the universe.",
        "She gently rocked the baby to sleep with a lullaby.",
        "Confidence grows when you face your fears one step at a time.",
        "They shared stories by the campfire under the star-filled sky.",
        "Music can express emotions words fail to capture.",
        "He meticulously planned every detail of the surprise party.",
        "Her paintings were inspired by her travels around the world.",
        "The thunderstorm raged outside while they stayed cozy indoors.",
        "The excitement of a new adventure can be exhilarating.",
        "His dedication to his craft was evident in every masterpiece.",
        "She found beauty in the simplicity of everyday moments.",
        "He whistled a cheerful tune while waiting for the bus.",
        "They discovered rare flowers blooming deep within the rainforest.",
        "His curiosity led him to explore the abandoned mansion alone.",
        "A family vacation can create lifelong memories for everyone.",
        "Her determination helped her overcome countless obstacles.",
        "He admired the intricate design of the old cathedral's stained glass windows.",
        "She found a book in the attic that changed her perspective on life.",
        "They hiked through the mountains, mesmerized by the panoramic view.",
        "A warm cup of tea soothed her nerves after a long day.",
        "He recounted tales of his youth with a nostalgic smile.",
        "The teacher praised the students for their collaborative spirit.",
        "She embraced the quiet solitude of her countryside retreat.",
        "They discovered a mysterious artifact beneath the ancient ruins.",
        "He woke early to watch the sunrise illuminate the horizon.",
        "A kind gesture can leave a lasting impression.",
        "She wore her grandmother's necklace for good luck.",
        "The baby giggled when the puppy licked her tiny toes.",
        "Every painting in the gallery told a unique story.",
        "He spent all morning crafting a heartfelt letter of apology.",
        "Her eyes lit up when she recognized an old friend across the street.",
        "They marveled at the bioluminescent plankton glowing in the midnight sea.",
        "He cherished the sound of rain tapping on his window.",
        "The library was his refuge, a place to dream and learn.",
        "She danced gracefully, captivated by the gentle piano music.",
        "A single spark can ignite a roaring flame.",
        "He spent weeks planning the perfect proposal on the beach.",
        "Her passion for art was evident in every brushstroke.",
        "They lost track of time while playing board games all evening.",
        "He gently patted the old dog that slept peacefully by his feet.",
        "The newborn foal struggled to stand on wobbly legs.",
        "She found inspiration in nature's boundless beauty.",
        "They found hope in the darkest moments, clinging to each other for support.",
        "He discovered a passion for cooking while helping at a local bistro.",
        "She carried the memory of her best friend in a locket around her neck.",
        "He spoke softly, trying not to disturb the sleeping baby.",
        "They visited the art museum to admire the classic Renaissance paintings.",
        "She overcame her stage fright by focusing on the joy of performing.",
        "He daydreamed about exploring uncharted territories.",
        "They danced in the rain, laughing at the storm's thunderous roar.",
        "He spent a quiet Sunday afternoon reading in the backyard hammock.",
        "She wrote poetry to capture her most profound emotions.",
        "They snuggled under warm blankets to watch the first snowfall of winter.",
        "He presented his research findings at the international conference.",
        "She carefully traced the family lineage back several centuries.",
        "They planted a garden together, hoping for a bountiful harvest.",
        "He woke up early to practice yoga in the morning light.",
        "She believed that kindness could change the world one step at a time.",
        "They gazed at the shooting stars streaking across the night sky.",
        "He built a small treehouse for his younger siblings to play in.",
        "She reminisced about the childhood summers spent at her grandparents' farm.",
        "They formed a band, blending various genres to create a unique sound.",
        "He spent years saving money for the trip of a lifetime.",
        "She had a knack for solving complex puzzles in record time.",
        "They cultivated a thriving vegetable garden in their backyard.",
        "He overcame his stage fright with daily practice and meditation.",
        "She found solace in writing letters to her future self.",
        "They huddled around the fireplace sharing ghost stories late into the night.",
        "He repaired the old clock that had been silent for decades.",
        "She painted vibrant sunsets inspired by her travels in the tropics.",
        "They organized a neighborhood cleanup to celebrate Earth Day.",
        "He discovered a hidden talent for sculpting while taking an art class.",
        "She sang lullabies to soothe her restless toddler at bedtime.",
        "They embarked on a cross-country road trip to visit national parks.",
        "He studied ancient myths to glean lessons about human nature.",
        "She volunteered at the local shelter, caring for abandoned animals.",
        "They joined a choir, finding joy in the harmony of voices.",
        "He journaled daily, reflecting on personal growth and gratitude.",
        "She followed a cherished family recipe to bake a special birthday cake.",
        "They marveled at the northern lights dancing across the icy tundra.",
        "He explored a remote island, documenting rare wildlife species.",
        "She wrote a heartfelt farewell message when leaving her childhood home.",
      ],
    ],
    runOptions: {
      pooling: "mean",
      normalize: true,
    },
    task: "feature-extraction",
    modelId: "Xenova/all-MiniLM-L6-v2",
    modelRevision: "main",
    modelHub: "huggingface",
    dtype: "q8",
    device: "cpu",
    backend: "onnx",
  },
  "image-to-text": {
    inputArgs: [
      "https://huggingface.co/datasets/mishig/sample_images/resolve/main/football-match.jpg",
@@ -94,7 +258,8 @@ const INFERENCE_PAD_PRESETS = {
    modelRevision: "main",
    modelHub: "huggingface",
    dtype: "q8",
    device: "wasm",
    device: "cpu",
    backend: "onnx",
  },
  ner: {
    inputArgs: ["Sarah lives in the United States of America"],
@@ -104,7 +269,8 @@ const INFERENCE_PAD_PRESETS = {
    modelRevision: "main",
    modelHub: "huggingface",
    dtype: "q8",
    device: "wasm",
    device: "cpu",
    backend: "onnx",
  },
  summary: {
    inputArgs: [TINY_ARTICLE],
@@ -116,7 +282,8 @@ const INFERENCE_PAD_PRESETS = {
    modelRevision: "main",
    modelHub: "huggingface",
    dtype: "q8",
    device: "wasm",
    device: "cpu",
    backend: "onnx",
  },
  zero: {
    inputArgs: [
@@ -129,7 +296,8 @@ const INFERENCE_PAD_PRESETS = {
    modelRevision: "main",
    modelHub: "huggingface",
    dtype: "q8",
    device: "wasm",
    device: "cpu",
    backend: "onnx",
  },
  feature: {
    inputArgs: [["This is an example sentence", "Each sentence is converted"]],
@@ -139,7 +307,8 @@ const INFERENCE_PAD_PRESETS = {
    modelRevision: "main",
    modelHub: "huggingface",
    dtype: "q8",
    device: "wasm",
    device: "cpu",
    backend: "onnx",
  },

  "link-preview": {
@@ -155,7 +324,8 @@ const INFERENCE_PAD_PRESETS = {
    modelHubRootUrl: "https://model-hub.mozilla.org",
    numThreads: 0,
    dtype: "q8",
    device: "wasm",
    device: "cpu",
    backend: "onnx",
  },
};

@@ -607,6 +777,8 @@ function loadExample(name) {
  setSelectOption("modelHub", data.modelHub);
  setSelectOption("dtype", data.dtype);
  setSelectOption("device", data.device);
  setSelectOption("backend", data.backend);
  document.getElementById("timeout").value = data.timeout || 30;
}

function findMaxMemory(metrics) {
@@ -660,6 +832,9 @@ async function runInference() {
  const numThreads = parseInt(document.getElementById("numThreads").value);
  const numRuns = parseInt(document.getElementById("numRuns").value);
  const modelHub = document.getElementById("modelHub").value;
  const backend = document.getElementById("backend").value;
  const timeout = parseInt(document.getElementById("timeout").value);

  let additionalEngineOptions = {};

  let inputData;
@@ -706,12 +881,17 @@ async function runInference() {
    modelHub,
    device,
    dtype,
    backend,
    numThreads,
    timeoutMS: THIRTY_SECONDS,
    timeoutMS: timeout * 1000,
    executionPriority: ExecutionPriority.LOW,
    ...additionalEngineOptions,
  };

  if (initData.device === "cpu" && initData.backend === "onnx") {
    initData.device = "wasm";
  }

  resultsConsole.addLine("Creating engine if needed");
  let engine;
  let startTime;
@@ -1060,7 +1240,7 @@ async function runBenchmark() {
        modelRevision: "main",
        modelHub: "huggingface",
        dtype: "q8",
        device: "wasm",
        device: "cpu",
      },
    },
    {
@@ -1074,7 +1254,7 @@ async function runBenchmark() {
        modelRevision: "main",
        modelHub: "huggingface",
        dtype: "q8",
        device: "wasm",
        device: "cpu",
      },
    },
    {
@@ -1090,7 +1270,7 @@ async function runBenchmark() {
        modelRevision: "main",
        modelHub: "huggingface",
        dtype: "q8",
        device: "wasm",
        device: "cpu",
      },
    },
    {
@@ -1143,6 +1323,13 @@ async function runBenchmark() {

        workload.pipelineOptions.backend = currentBackend;

        if (
          workload.pipelineOptions.device === "cpu" &&
          workload.pipelineOptions.backend === "onnx"
        ) {
          workload.pipelineOptions.device = "wasm";
        }

        bench.initDuration = await measure(async () => {
          const pipelineOptions = new PipelineOptions(workload.pipelineOptions);
          engine = await engineParent.getEngine(pipelineOptions);
@@ -1211,11 +1398,12 @@ window.onload = async function () {
  fillSelect("numThreads", getNumThreadsArray());
  fillSelect("predefined", PREDEFINED);
  fillSelect("benchmark.backend", ["all"].concat(lazy.BACKENDS));
  fillSelect("backend", lazy.BACKENDS);

  document.getElementById("predefined").value = "summary";
  loadExample("summary");
  document.getElementById("console").value = "";
  document.getElementById("predefined").value = "feature-large";
  loadExample("feature-large");

  document.getElementById("console").value = "";
  document
    .getElementById("inferenceButton")
    .addEventListener("click", runInference);
+9 −6
Original line number Diff line number Diff line
@@ -25,6 +25,7 @@ ChromeUtils.defineESModuleGetters(lazy, {
  PipelineOptions: "chrome://global/content/ml/EngineProcess.sys.mjs",
  DEFAULT_ENGINE_ID: "chrome://global/content/ml/EngineProcess.sys.mjs",
  DEFAULT_MODELS: "chrome://global/content/ml/EngineProcess.sys.mjs",
  WASM_BACKENDS: "chrome://global/content/ml/EngineProcess.sys.mjs",
});

ChromeUtils.defineLazyGetter(lazy, "console", () => {
@@ -332,11 +333,6 @@ class EngineDispatcher {
   * @returns {Promise<Engine>}
   */
  async initializeInferenceEngine(pipelineOptions, notificationsCallback) {
    // Create the inference engine given the wasm runtime and the options.
    const wasm = await this.mlEngineChild.getWasmArrayBuffer(
      pipelineOptions.backend
    );

    let remoteSettingsOptions = await this.mlEngineChild.getInferenceOptions(
      this.#featureId,
      this.#taskName,
@@ -361,9 +357,16 @@ class EngineDispatcher {
    }

    lazy.console.debug("Inference engine options:", mergedOptions);

    this.pipelineOptions = mergedOptions;

    // load the wasm if required.
    let wasm = null;
    if (lazy.WASM_BACKENDS.includes(pipelineOptions.backend || "onnx")) {
      wasm = await this.mlEngineChild.getWasmArrayBuffer(
        pipelineOptions.backend
      );
    }

    return InferenceEngine.create({
      wasm,
      pipelineOptions: mergedOptions,
+9 −1
Original line number Diff line number Diff line
@@ -20,7 +20,14 @@ export const DEFAULT_ENGINE_ID = "default-engine";
 * @type {Array<string>}
 * @description Supported backends.
 */
export const BACKENDS = ["onnx", "wllama"];
export const BACKENDS = ["onnx", "wllama", "onnx-native"];

/**
 * @constant
 * @type {Array<string>}
 * @description Backends using WASM.
 */
export const WASM_BACKENDS = ["onnx", "wllama"];

/**
 * @constant
@@ -290,6 +297,7 @@ export const KVCacheQuantizationLevel = {
export const InferenceDevice = {
  GPU: "gpu",
  WASM: "wasm",
  CPU: "cpu",
};

/**
+96 −50
Original line number Diff line number Diff line
@@ -2,6 +2,9 @@
 * 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/. */

// The following globals are defined in dom/webidl/ONNX.webidl
/* global Tensor, InferenceSession */

/**
 * @typedef {import("../../content/Utils.sys.mjs").ProgressAndStatusCallbackParams} ProgressAndStatusCallbackParams
 */
@@ -40,24 +43,59 @@ ChromeUtils.defineESModuleGetters(
  { global: "current" }
);

const NATIVE_BACKEND = "onnx-native";
const WASM_BACKEND = "onnx";

/**
 * Conditional import for Transformer.js
 * A global reference to the Transformers library.
 * Initially `null` until `importTransformers` sets it.
 *
 * The library will be lazily await on first usage in the Pipeline constructor.
 * If we are in Nightly, we are using the non-minified version.
 * @global
 * @type {object | null}
 */
let transformersPromise;
let transformers = null;
let transformersDev;

/**
 * Conditionally imports the Transformers library (transformers.js or transformers-dev.js)
 * on first usage, depending on the environment. If the "onnx-native" backend is used,
 * it exposes the `onnxruntime` object on the global scope under `Symbol.for("onnxruntime")`.
 *
 * @async
 * @function importTransformers
 * @param {string} backend - The backend to use (e.g. "onnx-native" or "onnx").
 * @returns {Promise<void>} A promise that resolves once the Transformers library is imported.
 */
async function importTransformers(backend) {
  if (transformers) {
    return;
  }

  lazy.console.debug(`Using backend ${backend}`);

  if (backend === NATIVE_BACKEND) {
    // check if we have the native backend.
    if (typeof InferenceSession === "undefined") {
      throw new Error("onnx-native backend not supported");
    }

    // Exposing an onnxruntime object to the Transformers lib.
    const onnxruntime = {
      InferenceSession,
      Tensor,
      supportedDevices: ["cpu"],
      defaultDevices: ["cpu"],
    };
    globalThis[Symbol.for("onnxruntime")] = onnxruntime;
  }
  if (AppConstants.NIGHTLY_BUILD) {
  transformersPromise = import(
    lazy.console.debug("Nightly detected. Using transformers-dev.js");
    transformers = await import(
      "chrome://global/content/ml/transformers-dev.js"
    );
  transformersDev = true;
  } else {
  transformersPromise = import("chrome://global/content/ml/transformers.js");
  transformersDev = false;
    lazy.console.debug("Beta or Release detected, using transformers.js");
    transformers = await import("chrome://global/content/ml/transformers.js");
  }
}

/**
@@ -331,6 +369,9 @@ export class ONNXPipeline {
    this.#errorFactory = errorFactory;
    this.#mlEngineWorker = mlEngineWorker;
    this.#metrics = [];
    let device;
    let session_options = {};

    // Setting up the Transformers.js environment
    // See https://huggingface.co/docs/transformers.js/api/env

@@ -348,6 +389,7 @@ export class ONNXPipeline {
      lazy.generateUUID()
    );
    // using `NO_LOCAL` so when the custom cache is used, we don't try to fetch it (see MLEngineWorker.match)
    if (config.backend === WASM_BACKEND) {
      transformers.env.localModelPath = "NO_LOCAL";
      transformers.env.backends.onnx.wasm.numThreads = config.numThreads;

@@ -375,14 +417,23 @@ export class ONNXPipeline {
          transformers.env.backends.onnx.debug = false;
          break;
      }
      lazy.console.debug("Transformers.js env", transformers.env);
      if (config.device === "cpu") {
        config.device = "wasm";
      }
      device = config.device || "wasm";
    } else {
      device = "cpu";
      session_options.intraOpNumThreads = config.numThreads;
      session_options.interOpNumThreads = config.numThreads;
      session_options.execution_mode = "sequential";
    }

    transformers.NoBadWordsLogitsProcessor.prototype._call =
      badWordsProcessorPatchWithBugFix; // Fix bug with bad words filter

    lazy.console.debug("Transformers.js env", transformers.env);

    // Defaults for transformers.js if they are null in PipelineOptions.
    const device = (config.device = config.device || "wasm");
    const dtype = (config.dtype = config.dtype || "q8");
    const modelRevision = (config.modelRevision =
      config.modelRevision || "main");
@@ -447,6 +498,7 @@ export class ONNXPipeline {
            device,
            dtype,
            use_external_data_format: config.useExternalDataFormat,
            session_options,
          }
        );
      } else {
@@ -503,15 +555,9 @@ export class ONNXPipeline {

    // Overriding the defaults with the options
    options.applyToConfig(config);
    config.backend = config.backend || "onnx";

    if (!transformers) {
      if (transformersDev) {
        lazy.console.debug("Nightly detected. Using transformers-dev.js");
      } else {
        lazy.console.debug("Beta or Release detected, using transformers.js");
      }
      transformers = await transformersPromise;
    }
    await importTransformers(config.backend);

    // reapply logLevel if it has changed.
    if (lazy.console.logLevel != config.logLevel) {
Loading