1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
|
/* 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
* file, You can obtain one at http://mozilla.org/MPL/2.0/. */
const lazy = {};
ChromeUtils.defineESModuleGetters(
lazy,
{
PromiseWorker: "resource://gre/modules/workers/PromiseWorker.mjs",
Pipeline: "chrome://global/content/ml/ONNXPipeline.mjs",
PipelineOptions: "chrome://global/content/ml/EngineProcess.sys.mjs",
},
{ global: "current" }
);
/**
* The actual MLEngine lives here in a worker.
*/
class MLEngineWorker {
#pipeline;
constructor() {
// Connect the provider to the worker.
this.#connectToPromiseWorker();
}
/** Implements the `match` function from the Cache API for Transformers.js custom cache.
*
* See https://developer.mozilla.org/en-US/docs/Web/API/Cache
*
* Attempts to match and retrieve a model file based on a provided key.
* Fetches a model file by delegating the call to the worker's main thread.
* Then wraps the fetched model file into a response object compatible with Transformers.js expectations.
*
* @param {string} key The unique identifier for the model to fetch.
* @returns {Promise<Response|null>} A promise that resolves with a Response object containing the model file or null if not found.
*/
async match(key) {
let res = await this.getModelFile(key);
if (res.fail) {
return null;
}
let headers = res.ok[1];
let modelFile = res.ok[2];
// Transformers.js expects a response object, so we wrap the array buffer
const response = new Response(modelFile, {
status: 200,
headers,
});
return response;
}
async getModelFile(...args) {
let result = await self.callMainThread("getModelFile", args);
return result;
}
/**
* Placeholder for the `put` method from the Cache API for Transformers.js custom cache.
*
* @throws {Error} Always thrown to indicate the method is not implemented.
*/
put() {
throw new Error("Method not implemented.");
}
/**
* @param {ArrayBuffer} wasm
* @param {object} options received as an object, converted to a PipelineOptions instance
*/
async initializeEngine(wasm, options) {
this.#pipeline = await lazy.Pipeline.initialize(
this,
wasm,
new lazy.PipelineOptions(options)
);
}
/**
* Run the worker.
*
* @param {string} request
*/
async run(request) {
if (request === "throw") {
throw new Error(
'Received the message "throw", so intentionally throwing an error.'
);
}
return await this.#pipeline.run(request);
}
/**
* Glue code to connect the `MLEngineWorker` to the PromiseWorker interface.
*/
#connectToPromiseWorker() {
const worker = new lazy.PromiseWorker.AbstractWorker();
worker.dispatch = (method, args = []) => {
if (!this[method]) {
throw new Error("Method does not exist: " + method);
}
return this[method](...args);
};
worker.close = () => self.close();
worker.postMessage = (message, ...transfers) => {
self.postMessage(message, ...transfers);
};
self.callMainThread = worker.callMainThread.bind(worker);
self.addEventListener("message", msg => worker.handleMessage(msg));
self.addEventListener("unhandledrejection", function (error) {
throw error.reason;
});
}
}
new MLEngineWorker();
|
