fix: enhance zero-shot classification worker and UI components

public/workers/zero-shot-classification.js

@@ -1,46 +1,73 @@
 /* eslint-disable no-restricted-globals */
-import { pipeline } from '@huggingface/transformers';
+import { pipeline } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.6.3'
 
 class MyZeroShotClassificationPipeline {
-  static task = 'zero-shot-classification';
-  static instance = null;
+  static task = 'zero-shot-classification'
+  static instance = null
 
-  static async getInstance(model, progress_callback = null) {
-    this.instance ??= pipeline(this.task, model, {
+  static async getInstance(model, dtype = 'fp32', progress_callback = null) {
+    this.instance = pipeline(this.task, model, {
+      dtype,
+      device: 'webgpu',
       progress_callback
-    });
-
-    return this.instance;
+    })
+    return this.instance
   }
 }
 
 // Listen for messages from the main thread
 self.addEventListener('message', async (event) => {
-  const { text, labels, model } = event.data;
-  if (!model) {
+  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: 'No model provided'
-    });
-    return;
-  }
-
-  // Retrieve the pipeline. When called for the first time,
-  // this will load the pipeline and save it for future use.
-  const classifier = await MyZeroShotClassificationPipeline.getInstance(model, (x) => {
-    // We also add a progress callback to the pipeline so that we can
-    // track model loading.
-    self.postMessage({ status: 'progress', output: x });
-  });
-  const split = text.split('\n');
-  for (const line of split) {
-    const output = await classifier(line, labels, {
-      hypothesis_template: 'This text is about {}.',
-      multi_label: true
-    });
-    // Send the output back to the main thread
-    self.postMessage({ status: 'output', output });
+      output: error.message || 'An error occurred during processing'
+    })
   }
-  // Send the output back to the main thread
-  self.postMessage({ status: 'complete' });
-});
+})

src/components/ModelLoader.tsx

@@ -64,6 +64,8 @@ const ModelLoader = () => {
 
     const onMessageReceived = (e: MessageEvent<WorkerMessage>) => {
       const { status, output } = e.data
+      console.log('Received output from worker', e.data)
+
       if (status === 'ready') {
         setStatus('ready')
         if (e.data.output) console.log(e.data.output)
@@ -81,6 +83,7 @@ const ModelLoader = () => {
         setStatus('output')
         const result = e.data.output!
         setResults((prev: any[]) => [...prev, result])
+        
         // console.log(result)
       } else if (status === 'error') {
         setStatus('error')
@@ -141,7 +144,6 @@ const ModelLoader = () => {
                   }
                   className="appearance-none bg-white border border-gray-300 rounded-md px-3 py-1 pr-8 text-xs text-gray-700 focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-blue-500"
                 >
-                  <option value="">Select quantization</option>
                   {modelInfo.supportedQuantizations.map((quant) => (
                     <option key={quant} value={quant}>
                       {quant}

src/components/ZeroShotClassification.tsx

@@ -1,4 +1,3 @@
-// src/App.tsx
 import { useState, useRef, useEffect, useCallback } from 'react'
 import {
   Section,
@@ -49,31 +48,40 @@ function ZeroShotClassification() {
     PLACEHOLDER_SECTIONS.map((title) => ({ title, items: [] }))
   )
 
-  const { status, setStatus, modelInfo } = useModel()
+  const { activeWorker, status, modelInfo, hasBeenLoaded } = useModel()
 
-  // Create a reference to the worker object.
-  const worker = useRef<Worker | null>(null)
-
-  // We use the `useEffect` hook to setup the worker as soon as the `App` component is mounted.
-  useEffect(() => {
-    if (!worker.current) {
+  const classify = useCallback(() => {
+    if (!modelInfo || !activeWorker) {
+      console.error('Model info or worker is not available')
       return
-      // Create the worker if it does not yet exist.
-      // worker.current = new Worker(
-      //   new URL('../workers/zero-shot-classification.js', import.meta.url),
-      //   {
-      //     type: 'module'
-      //   }
-      // )
     }
 
-    // Create a callback function for messages from the worker thread.
+    // Clear previous results
+    setSections((sections) =>
+      sections.map((section) => ({
+        ...section,
+        items: []
+      }))
+    )
+
+    const message: ZeroShotWorkerInput = {
+      type: 'classify',
+      text,
+      labels: sections
+        .slice(0, sections.length - 1)
+        .map((section) => section.title),
+      model: modelInfo.id
+    }
+    activeWorker.postMessage(message)
+  }, [text, sections, modelInfo, activeWorker])
+
+  // Handle worker messages
+  useEffect(() => {
+    if (!activeWorker) return
+
     const onMessageReceived = (e: MessageEvent<WorkerMessage>) => {
       const status = e.data.status
-      if (status === 'ready') {
-        setStatus('ready')
-      } else if (status === 'output') {
-        setStatus('output')
+      if (status === 'output') {
         const { sequence, labels, scores } = e.data.output!
 
         // Threshold for classification
@@ -89,33 +97,12 @@ function ZeroShotClassification() {
           }
           return newSections
         })
-      } else if (status === 'error') {
-        setStatus('error')
-        console.error(e.data.output)
       }
     }
 
-    // Attach the callback function as an event listener.
-    worker.current.addEventListener('message', onMessageReceived)
-
-    // Define a cleanup function for when the component is unmounted.
-    return () =>
-      worker.current?.removeEventListener('message', onMessageReceived)
-  }, [sections])
-
-  const classify = useCallback(() => {
-    if (!modelInfo) return
-
-    setStatus('loading')
-    const message: ZeroShotWorkerInput = {
-      text,
-      labels: sections
-        .slice(0, sections.length - 1)
-        .map((section) => section.title),
-      model: modelInfo.name
-    }
-    worker.current?.postMessage(message)
-  }, [text, sections, modelInfo])
+    activeWorker.addEventListener('message', onMessageReceived)
+    return () => activeWorker.removeEventListener('message', onMessageReceived)
+  }, [sections, activeWorker])
 
   const busy: boolean = status !== 'ready'
 
@@ -169,11 +156,10 @@ function ZeroShotClassification() {
           disabled={busy}
           onClick={classify}
         >
-          {!busy
-            ? 'Categorize'
-            : status === 'loading'
-            ? 'Model loading...'
-            : 'Processing'}
+          {hasBeenLoaded ? !busy
+              ? 'Categorize'
+              : 'Processing...'
+            : 'Load model first'}
         </button>
         <div className="flex gap-1">
           <button

src/types.ts

@@ -35,6 +35,7 @@ export interface WorkerMessage {
 }
 
 export interface ZeroShotWorkerInput {
+  type: 'classify'
   text: string
   labels: string[]
   model: string