transformers-js-playground / src /workers /sentiment-analysis.js
Vokturz's picture
Refactor app structure to integrate text classification feature and update pipeline selection
08476ef
raw
history blame
1.41 kB
/* eslint-disable no-restricted-globals */
import { pipeline } from "@huggingface/transformers";
class MyTextClassificationPipeline {
static task = "text-classification";
static model = "Xenova/bert-base-multilingual-uncased-sentiment";
static instance = null;
static async getInstance(progress_callback = null) {
this.instance ??= pipeline(this.task, this.model, {
progress_callback,
});
return this.instance;
}
}
// Listen for messages from the main thread
self.addEventListener("message", async (event) => {
// Retrieve the pipeline. When called for the first time,
// this will load the pipeline and save it for future use.
const classifier = await MyTextClassificationPipeline.getInstance((x) => {
// We also add a progress callback to the pipeline so that we can
// track model loading.
self.postMessage({ status: "progress", output: x });
});
const { text } = event.data;
const split = text.split("\n");
for (const line of split) {
if (line.trim()) {
const output = await classifier(line);
// Send the output back to the main thread
self.postMessage({
status: "output",
output: {
sequence: line,
labels: [output[0].label],
scores: [output[0].score]
}
});
}
}
// Send the output back to the main thread
self.postMessage({ status: "complete" });
});