|
|
|
import { pipeline } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.6.3'; |
|
|
|
class MyTextClassificationPipeline { |
|
static task = 'text-classification' |
|
static instance = null |
|
|
|
static async getInstance(model, progress_callback = null) { |
|
this.instance = pipeline(this.task, model, { |
|
progress_callback |
|
}) |
|
return this.instance |
|
} |
|
} |
|
|
|
|
|
self.addEventListener('message', async (event) => { |
|
const { type, model, text } = event.data |
|
|
|
if (!model) { |
|
self.postMessage({ |
|
status: 'error', |
|
output: 'No model provided' |
|
}) |
|
return |
|
} |
|
|
|
|
|
const classifier = await MyTextClassificationPipeline.getInstance( |
|
model, |
|
(x) => { |
|
self.postMessage({ status: 'progress', output: x }) |
|
} |
|
) |
|
|
|
if (type === 'load') { |
|
self.postMessage({ status: 'ready' }) |
|
return |
|
} |
|
|
|
if (type === 'classify') { |
|
if (!text) { |
|
self.postMessage({ status: 'complete' }) |
|
return |
|
} |
|
const split = text.split('\n') |
|
for (const line of split) { |
|
if (line.trim()) { |
|
const output = await classifier(line) |
|
self.postMessage({ |
|
status: 'output', |
|
output: { |
|
sequence: line, |
|
labels: [output[0].label], |
|
scores: [output[0].score] |
|
} |
|
}) |
|
} |
|
} |
|
self.postMessage({ status: 'complete' }) |
|
} |
|
}) |
|
|