Move results state to TextClassification component

src/components/ModelLoader.tsx

@@ -16,7 +16,6 @@ const ModelLoader = () => {
     activeWorker,
     setActiveWorker,
     pipeline,
-    setResults,
     hasBeenLoaded,
     setHasBeenLoaded
   } = useModel()
@@ -74,12 +73,6 @@ const ModelLoader = () => {
         ) {
           setProgress(output.progress)
         }
-      } else if (status === 'output') {
-        setStatus('output')
-        const result = e.data.output!
-        setResults((prev: any[]) => [...prev, result])
-
-        // console.log(result)
       } else if (status === 'error') {
         setStatus('error')
         console.error(e.data.output)
@@ -99,7 +92,6 @@ const ModelLoader = () => {
     setActiveWorker,
     setStatus,
     setProgress,
-    setResults,
     hasBeenLoaded,
     setHasBeenLoaded
   ])

src/components/TextClassification.tsx

@@ -1,5 +1,5 @@
 import { useState, useCallback, useEffect } from 'react'
-import { TextClassificationWorkerInput } from '../types'
+import { TextClassificationWorkerInput, WorkerMessage } from '../types'
 import { useModel } from '../contexts/ModelContext'
 
 const PLACEHOLDER_TEXTS: string[] = [
@@ -18,12 +18,12 @@ const PLACEHOLDER_TEXTS: string[] = [
 function TextClassification() {
   const [text, setText] = useState<string>(PLACEHOLDER_TEXTS.join('\n'))
   const [numberExamples, setNumberExamples] = useState(PLACEHOLDER_TEXTS.length)
+  const [results, setResults] = useState<any[]>([])
   const {
     activeWorker,
     status,
+    setStatus,
     modelInfo,
-    results,
-    setResults,
     hasBeenLoaded,
     selectedQuantization
   } = useModel()
@@ -56,6 +56,23 @@ function TextClassification() {
     activeWorker.postMessage(message)
   }, [text, modelInfo, activeWorker, selectedQuantization, setResults])
 
+  // Handle worker messages
+  useEffect(() => {
+    if (!activeWorker) return
+
+    const onMessageReceived = (e: MessageEvent<WorkerMessage>) => {
+      const status = e.data.status
+      if (status === 'output') {
+        setStatus('output')
+        const result = e.data.output!
+        setResults((prev: any[]) => [...prev, result])
+      }
+    }
+
+    activeWorker.addEventListener('message', onMessageReceived)
+    return () => activeWorker.removeEventListener('message', onMessageReceived)
+  }, [activeWorker, setStatus])
+
   const busy: boolean = status !== 'ready'
 
   const handleClear = (): void => {

src/contexts/ModelContext.tsx

@@ -1,9 +1,4 @@
-import React, {
-  createContext,
-  useContext,
-  useEffect,
-  useState
-} from 'react'
+import React, { createContext, useContext, useEffect, useState } from 'react'
 import {
   ModelInfo,
   ModelInfoResponse,
@@ -28,8 +23,6 @@ interface ModelContextType {
   setActiveWorker: (worker: Worker | null) => void
   isFetching: boolean
   setIsFetching: (isFetching: boolean) => void
-  results: any[]
-  setResults: React.Dispatch<React.SetStateAction<any[]>>
   hasBeenLoaded: boolean
   setHasBeenLoaded: (hasBeenLoaded: boolean) => void
 }
@@ -48,10 +41,8 @@ export function ModelProvider({ children }: { children: React.ReactNode }) {
     useState<QuantizationType>('int8')
   const [activeWorker, setActiveWorker] = useState<Worker | null>(null)
   const [isFetching, setIsFetching] = useState(false)
-  const [results, setResults] = useState<any[]>([])
   const [hasBeenLoaded, setHasBeenLoaded] = useState(false)
 
-
   // set progress to 0 when model is changed
   useEffect(() => {
     setProgress(0)
@@ -76,8 +67,6 @@ export function ModelProvider({ children }: { children: React.ReactNode }) {
         setActiveWorker,
         isFetching,
         setIsFetching,
-        results,
-        setResults,
         hasBeenLoaded,
         setHasBeenLoaded
       }}