Supporting Multithreading in Web Extensions

examples/extension-multithreaded/.gitignore

@@ -0,0 +1,2 @@
+node_modules/
+.DS_Store

examples/extension-multithreaded/README.md

@@ -0,0 +1,35 @@
+
+# Transformers.js - Sample browser extension
+
+An example project to show how to run 🤗 Transformers in a browser extension. Although we only provide instructions for running in Chrome, it should be similar for other browsers.
+
+## Getting Started
+1. Clone the repo and enter the project directory:
+    ```bash
+    git clone https://github.com/xenova/transformers.js.git
+    cd transformers.js/examples/extension/
+    ```
+1. Install the necessary dependencies:
+    ```bash
+    npm install 
+    ```
+
+1. Build the project:
+    ```bash
+    npm run build 
+    ```
+
+1. Add the extension to your browser. To do this, go to `chrome://extensions/`, enable developer mode (top right), and click "Load unpacked". Select the `build` directory from the dialog which appears and click "Select Folder".
+
+1. That's it! You should now be able to open the extenion's popup and use the model in your browser!
+
+## Editing the template
+
+We recommend running `npm run dev` while editing the template as it will rebuild the project when changes are made. 
+
+All source code can be found in the `./src/` directory:
+- `background.js` ([service worker](https://developer.chrome.com/docs/extensions/mv3/service_workers/)) - handles all the requests from the UI, does processing in the background, then returns the result. You will need to reload the extension (by visiting `chrome://extensions/` and clicking the refresh button) after editing this file for changes to be visible in the extension.
+
+- `content.js` ([content script](https://developer.chrome.com/docs/extensions/mv3/content_scripts/)) - contains the code which is injected into every page the user visits. You can use the `sendMessage` api to make requests to the background script. Similarly, you will need to reload the extension after editing this file for changes to be visible in the extension.
+
+- `popup.html`, `popup.css`, `popup.js` ([toolbar action](https://developer.chrome.com/docs/extensions/reference/action/)) - contains the code for the popup which is visible to the user when they click the extension's icon from the extensions bar. For development, we recommend opening the `popup.html` file in its own tab by visiting `chrome-extension://<ext_id>/popup.html` (remember to replace `<ext_id>` with the extension's ID). You will need to refresh the page while you develop to see the changes you make.

examples/extension-multithreaded/build/.gitignore

@@ -0,0 +1,3 @@
+# Running `npm run build` will build the project and output the files here.
+*
+!.gitignore

examples/extension-multithreaded/package.json

@@ -0,0 +1,20 @@
+{
+  "name": "extension",
+  "version": "0.0.1",
+  "description": "Transformers.js | Sample browser extension",
+  "scripts": {
+    "build": "webpack",
+    "dev": "webpack --watch"
+  },
+  "type": "module",
+  "author": "Xenova",
+  "license": "MIT",
+  "devDependencies": {
+    "copy-webpack-plugin": "^11.0.0",
+    "html-webpack-plugin": "^5.5.1",
+    "webpack": "^5.79.0"
+  },
+  "dependencies": {
+    "@xenova/transformers": "^2.0.0"
+  }
+}

examples/extension-multithreaded/public/manifest.json

@@ -0,0 +1,47 @@
+{
+  "manifest_version": 3,
+  "name": "extension",
+  "description": "Transformers.js | Sample browser extension",
+  "version": "0.0.1",
+  "permissions": [
+    "activeTab",
+    "scripting",
+    "contextMenus",
+    "storage",
+    "unlimitedStorage",
+    "sidePanel"
+  ],
+  "background": {
+    "service_worker": "background.js"
+  },
+  "minimum_chrome_version": "92",
+  "action": {
+    "default_icon": {
+      "16": "icons/icon.png",
+      "24": "icons/icon.png",
+      "32": "icons/icon.png"
+    },
+    "default_title": "Transformers.js"
+  },
+  "content_security_policy": {
+    "extension_pages": "script-src 'self' 'wasm-unsafe-eval'"
+  },
+  "sandbox": {
+    "pages": [
+      "sandbox.html"
+    ]
+  },
+  "icons": {
+    "16": "icons/icon.png",
+    "48": "icons/icon.png",
+    "128": "icons/icon.png"
+  },
+  "side_panel": {
+    "default_path": "sidepanel.html"
+  },
+  "sandbox": {
+    "pages": [
+      "sandbox.html"
+    ]
+  }
+}

examples/extension-multithreaded/public/offscreen.html

@@ -0,0 +1,12 @@
+<!DOCTYPE html>
+<html lang="en">
+
+<head>
+    <script src="offscreen.js" type="module"></script>
+</head>
+
+<body>
+    <iframe src="sandbox.html" id="sandbox"></iframe>
+</body>
+
+</html>

examples/extension-multithreaded/public/sandbox.html

@@ -0,0 +1,2 @@
+<!DOCTYPE html>
+<script src="sandbox.js" type="module"></script>

examples/extension-multithreaded/public/sidepanel.css

@@ -0,0 +1,50 @@
+/* Styles go here */
+
+* {
+    padding: 0;
+    margin: 0;
+    box-sizing: border-box;
+    font-family: 'Roboto', sans-serif;
+}
+
+h1 {
+    font-size: 40px;
+    text-align: center;
+    font-weight: 500;
+}
+
+h2 {
+    font-size: 20px;
+    text-align: center;
+    font-weight: 400;
+    margin-bottom: 16px;
+}
+
+.container {
+    width: 360px;
+}
+
+html,
+body {
+    min-width: 400px;
+    min-height: 500px;
+}
+
+body {
+    display: flex;
+    justify-content: center;
+    align-items: center;
+}
+
+#text {
+    width: 100%;
+    padding: 8px;
+    font-size: 20px;
+    margin-bottom: 8px;
+}
+
+#output {
+    font-size: 20px;
+    font-family: 'Roboto Mono', monospace;
+    height: 100px;
+}

examples/extension-multithreaded/public/sidepanel.html

@@ -0,0 +1,28 @@
+<!DOCTYPE html>
+<html lang="en">
+
+<head>
+    <meta charset="UTF-8">
+    <meta http-equiv="X-UA-Compatible" content="IE=edge">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Transformers.js | Sample Browser Extension</title>
+
+    <link rel="stylesheet" href="sidepanel.css" />
+    <script defer src="sidepanel.js"></script>
+</head>
+
+<body>
+    <div class="container">
+        <h1>Transformers.js</h1>
+        <h2>Run 🤗 Transformers in a Browser Extension!</h2>
+        <input id="text" placeholder="Enter text here">
+        <div id="model_loading_info">
+            Loading model files:
+        </div>
+        <pre id="output"></pre>
+    </div>
+    <!-- Inference is conducted in the context of a sandboxed iframe to enable the ONNX Runtime to load multiple threads -->
+    <iframe src="sandbox.html" id="sandbox" style="display: none;"></iframe>
+</body>
+
+</html>

examples/extension-multithreaded/src/background.js

@@ -0,0 +1,43 @@
+// background.js - Handles requests from the UI, runs the model, then sends back a response
+
+// open sidepanel when extension button is clicked
+chrome.sidePanel
+    .setPanelBehavior({ openPanelOnActionClick: true })
+    .catch((error) => console.error(error));
+
+////////////////////// Send Input for Inference Via In-Page Context Menu //////////////////////
+
+// Create the initial context menu items
+chrome.runtime.onInstalled.addListener(function () {
+    chrome.contextMenus.create({
+        id: 'classify-selection',
+        title: 'Classify "%s"',
+        contexts: ['selection'],  // only show up when text selected
+    });
+});
+
+// The sidepanel may not yet be open. 
+// If not, submitting from the context menu kicks off opening the sidepanel, but because this can take a long time 
+// we cannot pass the inference input while the sidepanel is initializing. Instead we store the inference input in 
+// session storage and have the sidepanel look there upon initialization.
+// When the sidepanel opens it also establishes a message channel to accept any subsequent inference input from the background worker.
+
+// this will contain a MessageChannel connection when sidepanel opens
+var connection
+
+// Perform inference when the user clicks a context menu
+chrome.contextMenus.onClicked.addListener(async (info, tab) => {
+    // Ignore context menu clicks that are not for classifications (or when there is no input)
+    if (info.menuItemId !== 'classify-selection' || !info.selectionText) return;
+
+    try {  // sidepanel is open, post text for inference
+        connection.postMessage({ text: info.selectionText })
+    } catch { // sidepanel not yet open; open sidepanel and persist text in session storage until panel can pick it up.
+        chrome.sidePanel.open({ windowId: tab.windowId }).catch(() => { })
+        chrome.storage.session.set({ 'inference_input': info.selectionText })
+
+        chrome.runtime.onConnect.addListener(x => {
+            connection = x
+        })
+    }
+});

examples/extension-multithreaded/src/inference.mjs

@@ -0,0 +1,25 @@
+import { pipeline, env } from '@xenova/transformers';
+
+// Skip initial check for local models, since we are not loading any local models.
+env.allowLocalModels = false;
+// The extension sandbox page (which we need to use because of the way ONNX loads multiple threads)
+// is prohibited from accesing the browser cache and will throw an erorr if it tries.
+// Perhaps this can be mitigated with a custom cache instead of the browser cache?
+env.useBrowserCache = false;
+// env.backends.onnx.wasm.numThreads = 1;
+
+class PipelineSingleton {
+    static task = 'text-classification';
+    static model = 'Xenova/distilbert-base-uncased-finetuned-sst-2-english';
+    static instance = null;
+
+    static async getInstance(progress_callback = null) {
+        if (this.instance === null) {
+            this.instance = pipeline(this.task, this.model, { progress_callback });
+        }
+
+        return this.instance;
+    }
+}
+
+export { PipelineSingleton }

examples/extension-multithreaded/src/offscreen.js

@@ -0,0 +1,37 @@
+// offscreen.js - Handles requests from the background service worker, runs the model, then sends back a response
+
+// Registering this listener when the script is first executed ensures that the
+// offscreen document will be able to receive messages when the promise returned
+// by `offscreen.createDocument()` resolves.
+chrome.runtime.onMessage.addListener((message, sender, sendResponse) => {
+    console.log("sender", sender)
+    console.log('message', message)
+    if (message.target !== 'offscreen') {
+        return false;
+    }
+    switch (message.type) {
+        case 'classify':
+            sendToSandboxForInference(message.data).then(result => {
+                sendResponse({ type: 'result', result });
+            });
+            return true;
+        default:
+            console.warn(`Unexpected message type received: '${message.type}'.`);
+            return false;
+    }
+});
+
+// The sandbox is neccesary because the ONNX Runtime loads in a way that violates the typical Chrome extension Content Security Policy
+// Using a MessageChannel allows us to use an async request/response pattern: https://advancedweb.hu/how-to-use-async-await-with-postmessage/
+const sendToSandboxForInference = input => new Promise((res, rej) => {
+    const channel = new MessageChannel();
+    channel.port1.onmessage = ({ data }) => {
+        channel.port1.close();
+        if (data.error) {
+            rej(data.error);
+        } else {
+            res(data.result);
+        }
+    };
+    document.getElementById('sandbox').contentWindow.postMessage(input, '*', [channel.port2]);
+});

examples/extension-multithreaded/src/sandbox.js

@@ -0,0 +1,26 @@
+import { PipelineSingleton } from "./inference.mjs";
+
+// Create generic inference function, which will be reused for the different types of events.
+const infer = async (text, message_port) => {
+    // Get the pipeline instance. This will load and build the model when run for the first time.
+    let model = await PipelineSingleton.getInstance((data) => {
+        // Post messages back on the same channel to indicate model loading progress.
+        data.type = 'load'
+        message_port.postMessage(data)
+    });
+
+    // Actually run the model on the input text
+    let result = await model(text);
+    return result;
+};
+
+// input, output, and errors as passed to and from the sandbox via a MessageChannel
+window.addEventListener('message', function (event) {
+    try {
+        infer(event.data, event.ports[0]).then(result => {
+            event.ports[0].postMessage({ type: 'result', result });
+        })
+    } catch (e) {
+        event.ports[0].postMessage({ type: 'error', error: e });
+    }
+}, false)