File size: 1,035 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
const workers: Record<string, Worker | null> = {};

export const getWorker = (pipeline: string) => {
  if (!workers[pipeline]) {
    let workerUrl: string;

    // Construct the public URL for the worker script.
    // process.env.PUBLIC_URL ensures this works correctly even if the
    // app is hosted in a sub-directory.
    switch (pipeline) {
      case 'text-classification':
        workerUrl = `${process.env.PUBLIC_URL}/workers/text-classification.js`;
        break;
      case 'zero-shot-classification':
        workerUrl = `${process.env.PUBLIC_URL}/workers/zero-shot-classification.js`;
        break;
      // Add other pipeline types here
      default:
        // Return null or throw an error if the pipeline is unknown
        return null;
    }
    workers[pipeline] = new Worker(workerUrl, { type: 'module' });
  }
  return workers[pipeline];
};

export const terminateWorker = (pipeline: string) => {
  const worker = workers[pipeline];
  if (worker) {
    worker.terminate();
    delete workers[pipeline];
  }
};