File size: 2,548 Bytes
85a4687 415aaef 85a4687 b1c66bb 85a4687 b1c66bb ac2af95 85a4687 9283c8b b1c66bb ac2af95 b1c66bb ac2af95 b1c66bb 96812c9 b1c66bb 85a4687 b1c66bb |
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 |
/* eslint-disable no-restricted-globals */
import { pipeline } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@latest'
class MyZeroShotClassificationPipeline {
static task = 'zero-shot-classification'
static instance = null
static async getInstance(model, dtype = 'fp32', progress_callback = null) {
try {
// Try WebGPU first
this.instance = await pipeline(this.task, model, {
dtype,
device: 'webgpu',
progress_callback
})
return this.instance
} catch (webgpuError) {
// Fallback to WASM if WebGPU fails
if (progress_callback) {
progress_callback({
status: 'fallback',
message: 'WebGPU failed, falling back to WASM'
})
}
try {
this.instance = await pipeline(this.task, model, {
dtype,
device: 'wasm',
progress_callback
})
return this.instance
} catch (wasmError) {
throw new Error(
`Both WebGPU and WASM failed. WebGPU error: ${webgpuError.message}. WASM error: ${wasmError.message}`
)
}
}
}
}
// Listen for messages from the main thread
self.addEventListener('message', async (event) => {
try {
const { type, model, dtype, text, labels } = event.data
if (!model) {
self.postMessage({
status: 'error',
output: 'No model provided'
})
return
}
// Retrieve the pipeline. This will download the model if not already cached.
const classifier = await MyZeroShotClassificationPipeline.getInstance(
model,
dtype,
(x) => {
self.postMessage({ status: 'loading', output: x })
}
)
if (type === 'load') {
self.postMessage({
status: 'ready',
output: `Model ${model}, dtype ${dtype} loaded`
})
return
}
if (type === 'classify') {
if (!text || !labels) {
self.postMessage({ status: 'ready' }) // Nothing to process
return
}
const split = text.split('\n')
for (const line of split) {
if (line.trim()) {
const output = await classifier(line, labels, {
hypothesis_template: 'This text is about {}.',
multi_label: true
})
self.postMessage({ status: 'output', output })
}
}
self.postMessage({ status: 'ready' })
}
} catch (error) {
self.postMessage({
status: 'error',
output: error.message || 'An error occurred during processing'
})
}
})
|