File size: 3,840 Bytes
3baa9da |
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 118 119 120 121 122 123 124 125 126 127 128 129 130 131 |
import { pipeline, env } from "@huggingface/transformers";
// Skip local model check
env.allowLocalModels = false;
async function supportsWebGPU() {
try {
if (!navigator.gpu) return false;
await navigator.gpu.requestAdapter();
return true;
} catch (e) {
return false;
}
}
const device = (await supportsWebGPU()) ? "webgpu" : "wasm";
class PipelineManager {
static defaultConfigs = {
"text-classification": {
model: "onnx-community/rubert-tiny-sentiment-balanced-ONNX",
},
"image-classification": {
model: "onnx-community/mobilenet_v2_1.0_224",
},
};
static instances = {}; // key: `${task}:${modelName}` -> pipeline instance
static currentTask = "text-classification";
static currentModel = PipelineManager.defaultConfigs["text-classification"].model;
static queue = [];
static isProcessing = false;
static async getInstance(task, modelName, progress_callback = null) {
const key = `${task}:${modelName}`;
if (!this.instances[key]) {
self.postMessage({ status: "initiate", file: modelName, task });
this.instances[key] = await pipeline(task, modelName, { progress_callback, device: device});
self.postMessage({ status: "ready", file: modelName, task });
}
return this.instances[key];
}
static async processQueue() {
if (this.isProcessing || this.queue.length === 0) return;
this.isProcessing = true;
const { input, task, modelName } = this.queue[this.queue.length - 1];
this.queue = [];
try {
const classifier = await this.getInstance(task, modelName, (x) => {
self.postMessage({
...x,
status: x.status || "progress",
file: x.file || modelName,
name: modelName,
task,
loaded: x.loaded,
total: x.total,
progress: x.loaded && x.total ? (x.loaded / x.total) * 100 : 0,
});
});
let output;
if (task === "image-classification") {
// input is a data URL or Blob
output = await classifier(input, { top_k: 5 });
} else if (task === "automatic-speech-recognition") {
output = await classifier(input);
} else {
output = await classifier(input, { top_k: 5 });
}
self.postMessage({
status: "complete",
output,
file: modelName,
task,
});
} catch (error) {
self.postMessage({
status: "error",
error: error.message,
file: modelName,
task,
});
}
this.isProcessing = false;
if (this.queue.length > 0) {
this.processQueue();
}
}
}
// Listen for messages from the main thread
self.addEventListener("message", async (event) => {
const { input, modelName, task, action } = event.data;
// console.log("Worker received message:", event.data); // Add this line to log the received message t
if (action === "load-model") {
PipelineManager.currentTask = task || "text-classification";
PipelineManager.currentModel =
modelName ||
PipelineManager.defaultConfigs[PipelineManager.currentTask].model;
await PipelineManager.getInstance(
PipelineManager.currentTask,
PipelineManager.currentModel,
(x) => {
self.postMessage({
...x,
file: PipelineManager.currentModel,
status: x.status || "progress",
loaded: x.loaded,
total: x.total,
task: PipelineManager.currentTask,
});
}
);
return;
}
PipelineManager.queue.push({
input,
task: task || PipelineManager.currentTask,
modelName: modelName || PipelineManager.currentModel,
});
PipelineManager.processQueue();
}); |