File size: 1,555 Bytes
e7ba29d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
/* eslint-disable no-restricted-globals */
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
  }
}

// Listen for messages from the main thread
self.addEventListener('message', async (event) => {
  const { type, model, text } = event.data // Destructure 'type'

  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 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' }) // Nothing to process
      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' })
  }
})