Enhance text classification pipeline: add device support, improve error handling, and refine message processing logic

public/workers/text-classification.js

@@ -9,7 +9,7 @@ class MyTextClassificationPipeline {
     this.instance = pipeline(
       this.task,
       model,
-      { dtype, progress_callback },
+      { dtype, device: "webgpu", progress_callback },
     )
     return this.instance
   }
@@ -17,49 +17,59 @@ class MyTextClassificationPipeline {
 
 // Listen for messages from the main thread
 self.addEventListener('message', async (event) => {
-  const { type, model, dtype, text } = event.data
+  try {
+    const { type, model, dtype, text } = 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 MyTextClassificationPipeline.getInstance(
-    model,
-    dtype,
-    (x) => {
-      self.postMessage({ status: 'loading', output: x })
+    if (!model) {
+      self.postMessage({
+        status: 'error',
+        output: 'No model provided'
+      })
+      return
     }
-  )
 
-  if (type === 'load') {
-    self.postMessage({ status: 'ready' })
-    return
-  }
+    // Retrieve the pipeline. This will download the model if not already cached.
+    const classifier = await MyTextClassificationPipeline.getInstance(
+      model,
+      dtype,
+      (x) => {
+        self.postMessage({ status: 'loading', output: x })
+      }
+    )
 
-  if (type === 'classify') {
-    if (!text) {
-      self.postMessage({ status: 'ready' }) // Nothing to process
+    if (type === 'load') {
+      self.postMessage({
+        status: 'ready',
+        output: `Model ${model}, dtype ${dtype} loaded`
+      })
       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]
-          }
-        })
+
+    if (type === 'classify') {
+      if (!text) {
+        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)
+          self.postMessage({
+            status: 'output',
+            output: {
+              sequence: line,
+              labels: [output[0].label],
+              scores: [output[0].score]
+            }
+          })
+        }
       }
+      self.postMessage({ status: 'ready' })
     }
-    self.postMessage({ status: 'ready' })
+  } catch (error) {
+    self.postMessage({
+      status: 'error',
+      output: error.message || 'An error occurred during processing'
+    })
   }
 })

src/components/ModelLoader.tsx

@@ -21,6 +21,9 @@ const ModelLoader = () => {
     setHasBeenLoaded
   } = useModel()
 
+  useEffect(() => {
+    setHasBeenLoaded(false)
+  }, [selectedQuantization])
 
   useEffect(() => {
     if (!modelInfo) return
@@ -43,6 +46,8 @@ const ModelLoader = () => {
     setHasBeenLoaded(false)
   }, [modelInfo, setSelectedQuantization, setHasBeenLoaded])
 
+
+
   useEffect(() => {
     if (!modelInfo) return
 
@@ -61,6 +66,7 @@ const ModelLoader = () => {
       const { status, output } = e.data
       if (status === 'ready') {
         setStatus('ready')
+        if (e.data.output) console.log(e.data.output)
         setHasBeenLoaded(true)
       } else if (status === 'loading' && output && !hasBeenLoaded) {
         setStatus('loading')
@@ -156,7 +162,7 @@ const ModelLoader = () => {
           <div className="flex justify-center">
             <button
               className="w-32 py-2 px-4 bg-green-500 hover:bg-green-600 rounded text-white font-medium disabled:opacity-50 disabled:cursor-not-allowed transition-colors text-sm inline-flex items-center text-center justify-center space-x-2"
-              disabled={hasBeenLoaded}
+              disabled={hasBeenLoaded || status === 'loading'}
               onClick={loadModel}
             >
               {status === 'loading' && !hasBeenLoaded ? (

src/components/TextClassification.tsx

@@ -24,8 +24,6 @@ function TextClassification() {
   const [text, setText] = useState<string>(PLACEHOLDER_TEXTS.join('\n'))
   const { activeWorker, status, setStatus, modelInfo, results, setResults, hasBeenLoaded} = useModel()
 
-
-
   const classify = useCallback(() => {
     if (!modelInfo || !activeWorker) {
       console.error('Model info or worker is not available')
@@ -48,46 +46,49 @@ function TextClassification() {
 
   return (
     <div className="flex flex-col h-[60vh] max-h-[100vh] w-full p-4">
-      <h1 className="text-2xl font-bold mb-4">Text Classification</h1>
+      <h1 className="text-2xl font-bold mb-4 flex-shrink-0">Text Classification</h1>
 
-      <div className="flex flex-col lg:flex-row gap-4 h-full">
+      <div className="flex flex-col lg:flex-row gap-4 flex-1 min-h-0">
         {/* Input Section */}
-        <div className="flex flex-col w-full lg:w-1/2">
-          <label className="text-lg font-medium mb-2">Input Text:</label>
-          <textarea
-            className="border border-gray-300 rounded p-3 flex-grow resize-none"
-            value={text}
-            onChange={(e) => setText(e.target.value)}
-            placeholder="Enter text to classify (one per line)..."
-          />
+        <div className="flex flex-col w-full lg:w-1/2 min-h-0">
+          <label className="text-lg font-medium mb-2 flex-shrink-0">Input Text:</label>
+          
+          <div className="flex flex-col flex-1 min-h-0">
+            <textarea
+              className="border border-gray-300 rounded p-3 flex-1 resize-none min-h-[200px]"
+              value={text}
+              onChange={(e) => setText(e.target.value)}
+              placeholder="Enter text to classify (one per line)..."
+            />
 
-          <div className="flex gap-2 mt-4">
-            <button
-              className="flex-1 py-2 px-4 bg-blue-500 hover:bg-blue-600 rounded text-white font-medium disabled:opacity-50 disabled:cursor-not-allowed transition-colors"
-              disabled={busy}
-              onClick={classify}
-            >
-              {hasBeenLoaded ? !busy
-                  ? 'Classify Text'
-                  : 'Processing...'
-                : 'Load model first'}
-            </button>
-            <button
-              className="py-2 px-4 bg-gray-500 hover:bg-gray-600 rounded text-white font-medium transition-colors"
-              onClick={handleClear}
-            >
-              Clear Results
-            </button>
+            <div className="flex gap-2 mt-4 flex-shrink-0">
+              <button
+                className="flex-1 py-2 px-4 bg-blue-500 hover:bg-blue-600 rounded text-white font-medium disabled:opacity-50 disabled:cursor-not-allowed transition-colors"
+                disabled={busy}
+                onClick={classify}
+              >
+                {hasBeenLoaded ? !busy
+                    ? 'Classify Text'
+                    : 'Processing...'
+                  : 'Load model first'}
+              </button>
+              <button
+                className="py-2 px-4 bg-gray-500 hover:bg-gray-600 rounded text-white font-medium transition-colors"
+                onClick={handleClear}
+              >
+                Clear Results
+              </button>
+            </div>
           </div>
         </div>
 
         {/* Results Section */}
-        <div className="flex flex-col w-full lg:w-1/2">
-          <label className="text-lg font-medium mb-2">
+        <div className="flex flex-col w-full lg:w-1/2 min-h-0">
+          <label className="text-lg font-medium mb-2 flex-shrink-0">
             Classification Results ({results.length}):
           </label>
 
-          <div className="border border-gray-300 rounded p-3 flex-grow overflow-y-auto">
+          <div className="border border-gray-300 rounded p-3 flex-1 overflow-y-auto min-h-[200px]">
             {results.length === 0 ? (
               <div className="text-gray-500 text-center py-8">
                 No results yet. Click "Classify Text" to analyze your input.

src/types.ts

@@ -9,7 +9,13 @@ export interface ClassificationOutput {
   scores: number[]
 }
 
-export type WorkerStatus = 'initiate' | 'ready' | 'output' | 'loading' | 'error'
+export type WorkerStatus =
+  | 'initiate'
+  | 'ready'
+  | 'output'
+  | 'loading'
+  | 'error'
+  | 'disposed'
 
 export interface WorkerMessage {
   status: WorkerStatus