Improve browser extension sample/template (huggingface#196)

* Update extension to be module

* Update example extension

* Allow user to specify a custom cache system

* Implement custom cache system

Emulates the Web Cache API using chrome's local storage API

* Use custom cache system in extension

* Fix serialization

* Remove old folders

* Update extension readme

* Add note about JSON requirement for local storage

Author: xenova

examples/extension/README.md

@@ -14,28 +14,11 @@ An example project to show how to run 🤗 Transformers in a browser extension.
     npm install 
     ```
 
-1. Add your model files to `./public/models/`. For this demo, we use [distilbert-base-uncased-finetuned-sst-2-english](https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english/tree/main) from the Hugging Face Hub. It should look something like this:
-    ```
-    distilbert-base-uncased-finetuned-sst-2-english/
-    ├── config.json
-    ├── tokenizer.json
-    ├── tokenizer_config.json
-    └── onnx/
-        ├── model.onnx
-        └── model_quantized.onnx
-    ```
-
-1. Add the WASM files to `./public/wasm/`. You can download them from the jsDelivr CDN [here](https://www.jsdelivr.com/package/npm/@xenova/transformers?tab=files&path=dist):
-    ```
-    ort-wasm.wasm
-    ort-wasm-simd.wasm
-    ort-wasm-simd-threaded.wasm
-    ort-wasm-threaded.wasm
-    ```
 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!
@@ -45,5 +28,8 @@ An example project to show how to run 🤗 Transformers in a browser extension.
 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` - contains the service worker code which runs in the background. It handles all the requests from the UI, does processing on a separate thread, 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.
-- `popup.html`, `popup.css`, `popup.js` - 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.
+- `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/package.json

@@ -6,6 +6,7 @@
     "build": "webpack",
     "dev": "webpack --watch"
   },
+  "type": "module",
   "author": "Xenova",
   "license": "MIT",
   "devDependencies": {

examples/extension/public/manifest.json

@@ -3,17 +3,43 @@
   "name": "extension",
   "description": "Transformers.js | Sample browser extension",
   "version": "0.0.1",
-  "permissions": [],
+  "permissions": [
+    "activeTab",
+    "scripting",
+    "contextMenus",
+    "storage",
+    "unlimitedStorage"
+  ],
   "background": {
-    "service_worker": "background.js"
+    "service_worker": "background.js",
+    "type": "module"
   },
+  "content_scripts": [
+    {
+      "matches": [
+        "<all_urls>"
+      ],
+      "js": [
+        "content.js"
+      ]
+    }
+  ],
   "minimum_chrome_version": "92",
   "action": {
-    "default_icon": "icons/icon.png",
-    "default_title": "extension",
+    "default_icon": {
+      "16": "icons/icon.png",
+      "24": "icons/icon.png",
+      "32": "icons/icon.png"
+    },
+    "default_title": "Transformers.js",
     "default_popup": "popup.html"
   },
   "content_security_policy": {
     "extension_pages": "script-src 'self' 'wasm-unsafe-eval'"
+  },
+  "icons": {
+    "16": "icons/icon.png",
+    "48": "icons/icon.png",
+    "128": "icons/icon.png"
   }
 }

examples/extension/public/models/.gitignore

@@ -1,44 +1,102 @@
-// background.js - Handles requests from the frontend, runs the model, then sends back a response
-// TODO - make persistent (i.e., do not close after inactivity)
-
-if (typeof ServiceWorkerGlobalScope !== 'undefined' && self instanceof ServiceWorkerGlobalScope) {
-    // Load the library
-    const { pipeline, env } = require('@xenova/transformers');
-
-    // Set environment variables to only use local models.
-    env.useBrowserCache = false;
-    env.remoteModels = false;
-    env.localModelPath = chrome.runtime.getURL('models/')
-    env.backends.onnx.wasm.wasmPaths = chrome.runtime.getURL('wasm/')
-    env.backends.onnx.wasm.numThreads = 1;
-
-    // TODO: Replace this with your own task and model
-    const task = 'text-classification';
-    const model = 'distilbert-base-uncased-finetuned-sst-2-english';
-
-    // Load model, storing the promise that is returned from the pipeline function.
-    // Doing it this way will load the model in the background as soon as the worker is created.
-    // To actually use the model, you must call `await modelPromise` to get the actual classifier.
-    const modelPromise = pipeline(task, model, {
-        progress_callback: (data) => {
-            // If you would like to add a progress bar for model loading,
-            // you can send `data` back to the UI.
+// background.js - Handles requests from the UI, runs the model, then sends back a response
+
+import { pipeline, env } from '@xenova/transformers';
+import { CustomCache } from "./cache.js";
+
+// Define caching parameters
+env.useBrowserCache = false;
+env.useCustomCache = true;
+env.customCache = new CustomCache('transformers-cache');
+
+// Skip initial check for local models, since we are not loading any local models.
+env.allowLocalModels = false;
+
+// Due to a bug in onnxruntime-web, we must disable multithreading for now.
+// See https://github.com/microsoft/onnxruntime/issues/14445 for more information.
+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;
+    }
+}
+
+// Create generic classify function, which will be reused for the different types of events.
+const classify = async (text) => {
+    // Get the pipeline instance. This will load and build the model when run for the first time.
+    let model = await PipelineSingleton.getInstance((data) => {
+        // You can track the progress of the pipeline creation here.
+        // e.g., you can send `data` back to the UI to indicate a progress bar
+        // console.log('progress', data)
     });
 
+    // Actually run the model on the input text
+    let result = await model(text);
+    return result;
+};
+
+////////////////////// 1. Context Menus //////////////////////
+//
+// Add a listener to create the initial context menu items,
+// context menu items only need to be created at runtime.onInstalled
+chrome.runtime.onInstalled.addListener(function () {
+    // Register a context menu item that will only show up for selection text.
+    chrome.contextMenus.create({
+        id: 'classify-selection',
+        title: 'Classify "%s"',
+        contexts: ['selection'],
+    });
+});
 
-    // Listen for messages from the UI, process it, and send the result back.
-    chrome.runtime.onMessage.addListener((message, sender, sendResponse) => {
+// 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;
 
-        // Run model prediction asynchronously
-        (async function () {
-            let model = await modelPromise;     // 1. Load model if not already loaded
-            let result = await model(message);  // 2. Run model prediction
-            sendResponse(result);               // 3. Send response back to UI
-        })();
+    // Perform classification on the selected text
+    let result = await classify(info.selectionText);
 
-        // return true to indicate we will send a response asynchronously
-        // see https://stackoverflow.com/a/46628145 for more information
-        return true;
+    // Do something with the result
+    chrome.scripting.executeScript({
+        target: { tabId: tab.id },    // Run in the tab that the user clicked in
+        args: [result],               // The arguments to pass to the function
+        function: (result) => {       // The function to run
+            // NOTE: This function is run in the context of the web page, meaning that `document` is available.
+            console.log('result', result)
+            console.log('document', document)
+        },
     });
-}
+});
+//////////////////////////////////////////////////////////////
+
+////////////////////// 2. Message Events /////////////////////
+// 
+// Listen for messages from the UI, process it, and send the result back.
+chrome.runtime.onMessage.addListener((message, sender, sendResponse) => {
+    console.log('sender', sender)
+    if (message.action !== 'classify') return; // Ignore messages that are not meant for classification.
+
+    // Run model prediction asynchronously
+    (async function () {
+        // Perform classification
+        let result = await classify(message.text);
+
+        // Send response back to UI
+        sendResponse(result);
+    })();
+
+    // return true to indicate we will send a response asynchronously
+    // see https://stackoverflow.com/a/46628145 for more information
+    return true;
+});
+//////////////////////////////////////////////////////////////
+

examples/extension/public/wasm/.gitignore

@@ -0,0 +1,80 @@
+// Design a caching API to be used by the extension which implements the same interface as
+// the browser's native Cache API (https://developer.mozilla.org/en-US/docs/Web/API/Cache)
+// but uses the browser's local storage API (https://developer.chrome.com/docs/extensions/reference/storage/).
+// 
+// Since the local storage API requires all data to be stored as JSON (which doesn't allow some ASCII chars),
+// one of the better approaches is to store the response body as a base64-encoded string. This is not ideal,
+// as it increases the size of the response body by ~33%, but it's the best we can do with the local storage API.
+// See https://stackoverflow.com/a/1443240/13989043 for more information about this.
+// 
+// For serialization (arraybuffer -> string) and unserialization (string -> arraybuffer),
+// use the `FileReader` and `Blob` APIs. Although other options are also possible, this approach
+// is considered to be better for larger files (like models).
+// 
+// Other references:
+//  - https://developer.chrome.com/docs/extensions/reference/storage/#property-local
+//  - https://stackoverflow.com/questions/6965107/converting-between-strings-and-arraybuffers
+
+export class CustomCache {
+    /**
+     * Instantiate a `CustomCache` object.
+     * @param {string} path 
+     */
+    constructor(cacheName) {
+        this.cacheName = cacheName;
+    }
+
+    /**
+     * Checks whether the given request is in the cache.
+     * @param {Request|string} request 
+     * @returns {Promise<Response | undefined>}
+     */
+    async match(request) {
+        const url = request instanceof Request ? request.url : request;
+        const cached = await chrome.storage.local.get([url]);
+
+        if (cached[url]) {
+            return await fetch(cached[url]._body);
+        } else {
+            return undefined;
+        }
+    }
+
+    /**
+     * Adds the given response to the cache.
+     * @param {Request|string} request 
+     * @param {Response} response 
+     * @returns {Promise<void>}
+     */
+    async put(request, response) {
+        const url = request instanceof Request ? request.url : request;
+        const buffer = await response.arrayBuffer();
+
+        const body = await new Promise((resolve, reject) => {
+            const reader = new FileReader();
+            reader.onload = e => resolve(e.target.result);
+            reader.onerror = e => reject(e.target.error);
+            reader.readAsDataURL(new Blob([buffer], { type: 'application/octet-stream' }));
+        });
+
+        try {
+            await chrome.storage.local.set({
+                [url]: {
+                    _body: body,
+
+                    // Save original response in case
+                    status: response.status,
+                    statusText: response.statusText,
+                    headers: Object.fromEntries(response.headers.entries()),
+                    url: response.url,
+                    redirected: response.redirected,
+                    type: response.type,
+                    ok: response.ok,
+                }
+            });
+
+        } catch (err) {
+            console.warn('An error occurred while writing the file to cache:', err)
+        }
+    }
+}

examples/extension/src/background.js

@@ -0,0 +1,11 @@
+// content.js - the content scripts which is run in the context of web pages, and has access
+// to the DOM and other web APIs.
+
+// Example usage:
+// const message = {
+//     action: 'classify',
+//     text: 'text to classify',
+// }
+// chrome.runtime.sendMessage(message, (response) => {
+//     console.log('received user data', response)
+// });

examples/extension/src/cache.js

@@ -1,13 +1,21 @@
-// This script handles interaction with the user interface, as well as handling
-// the communication between the main thread (UI) and the background thread (processing).
+// popup.js - handles interaction with the extension's popup, sends requests to the
+// service worker (background.js), and updates the popup's UI (popup.html) on completion.
 
 const inputElement = document.getElementById('text');
 const outputElement = document.getElementById('output');
 
-// 1. Send input data to the worker thread when it changes.
+// Listen for changes made to the textbox.
 inputElement.addEventListener('input', (event) => {
-    chrome.runtime.sendMessage(event.target.value, (response) => {
-        // 2. Handle results returned by the service worker (`background.js`) and update the UI.
+
+    // Bundle the input data into a message.
+    const message = {
+        action: 'classify',
+        text: event.target.value,
+    }
+
+    // Send this message to the service worker.
+    chrome.runtime.sendMessage(message, (response) => {
+        // Handle results returned by the service worker (`background.js`) and update the popup's UI.
         outputElement.innerText = JSON.stringify(response, null, 2);
     });
 });