wip: refactor model loading and classification components

public/workers/text-classification.js

@@ -29,7 +29,7 @@ self.addEventListener('message', async (event) => {
   const classifier = await MyTextClassificationPipeline.getInstance(
     model,
     (x) => {
-      self.postMessage({ status: 'progress', output: x })
+      self.postMessage({ status: 'loading', output: x })
     }
   )
 
@@ -40,7 +40,7 @@ self.addEventListener('message', async (event) => {
 
   if (type === 'classify') {
     if (!text) {
-      self.postMessage({ status: 'complete' }) // Nothing to process
+      self.postMessage({ status: 'ready' }) // Nothing to process
       return
     }
     const split = text.split('\n')
@@ -57,6 +57,6 @@ self.addEventListener('message', async (event) => {
         })
       }
     }
-    self.postMessage({ status: 'complete' })
+    self.postMessage({ status: 'ready' })
   }
 })

src/App.tsx

@@ -1,4 +1,4 @@
-import { useEffect, useState } from 'react'
+import { useEffect } from 'react'
 import PipelineSelector from './components/PipelineSelector'
 import ZeroShotClassification from './components/ZeroShotClassification'
 import TextClassification from './components/TextClassification'
@@ -10,14 +10,12 @@ import ModelSelector from './components/ModelSelector'
 import ModelInfo from './components/ModelInfo'
 
 function App() {
-  const { pipeline, setPipeline, progress, status, modelInfo, setModels } =
-    useModel()
+  const { pipeline, setPipeline, setModels } = useModel()
 
   useEffect(() => {
     const fetchModels = async () => {
       const fetchedModels = await getModelsByPipeline(pipeline)
       setModels(fetchedModels)
-      console.log(fetchedModels)
     }
     fetchModels()
   }, [setModels, pipeline])
@@ -26,76 +24,34 @@ function App() {
     <div className="min-h-screen bg-gradient-to-br from-blue-50 to-indigo-100">
       <Header />
 
-      <main className="max-w-8xl mx-auto px-4 sm:px-6 lg:px-8 py-8">
-        {/* Pipeline Selection */}
+      <main className="max-w-6xl mx-auto px-4 sm:px-6 lg:px-8 py-8">
         <div className="mb-8">
-          <div className="bg-white rounded-lg shadow-sm border p-6">
-            <div className="flex items-center justify-between mb-4">
-              <ModelInfo />
-            </div>
+          <div className="bg-white rounded-lg border p-6">
+            <div className="flex items-start justify-between max-w-6xl mx-auto">
+              <div className="space-y-2 flex-1">
+                <div className="space-y-2">
+                  <span className="text-lg font-semibold text-gray-900 block">
+                    Choose a Pipeline
+                  </span>
+                  <PipelineSelector
+                    pipeline={pipeline}
+                    setPipeline={setPipeline}
+                  />
+                </div>
 
-            <div className="grid grid-cols-1 xl:grid-cols-2 gap-4 items-start">
-              <div className="space-y-2">
-                <span className="text-lg font-semibold text-gray-900 block">
-                  Choose a Pipeline
-                </span>
-                <PipelineSelector
-                  pipeline={pipeline}
-                  setPipeline={setPipeline}
-                />
+                <div className="space-y-2">
+                  <span className="text-lg font-semibold text-gray-900 block">
+                    Select Model
+                  </span>
+                  <ModelSelector />
+                </div>
               </div>
 
-              <div className="space-y-2">
-                <span className="text-lg font-semibold text-gray-900 block">
-                  Select Model
-                </span>
-                <ModelSelector />
+              <div className="ml-6">
+                <ModelInfo />
               </div>
             </div>
 
-            {/* Model Loading Progress */}
-            {status === 'progress' && (
-              <div className="mt-4 p-4 bg-blue-50 rounded-lg">
-                <div className="flex items-center space-x-3">
-                  <div className="flex-shrink-0">
-                    <svg
-                      className="animate-spin h-5 w-5 text-blue-500"
-                      fill="none"
-                      viewBox="0 0 24 24"
-                    >
-                      <circle
-                        className="opacity-25"
-                        cx="12"
-                        cy="12"
-                        r="10"
-                        stroke="currentColor"
-                        strokeWidth="4"
-                      ></circle>
-                      <path
-                        className="opacity-75"
-                        fill="currentColor"
-                        d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z"
-                      ></path>
-                    </svg>
-                  </div>
-                  <div className="flex-1">
-                    <p className="text-sm font-medium text-blue-900">
-                      Loading Model...
-                    </p>
-                    <div className="mt-2 bg-blue-200 rounded-full h-2">
-                      <div
-                        className="bg-blue-500 h-2 rounded-full transition-all duration-300"
-                        style={{ width: `${progress.toFixed(2)}%` }}
-                      ></div>
-                    </div>
-                    <p className="text-xs text-blue-700 mt-1">
-                      {progress.toFixed(2)}%
-                    </p>
-                  </div>
-                </div>
-              </div>
-            )}
-
             {/* Pipeline Description */}
             <div className="mt-4 p-4 bg-gray-50 rounded-lg">
               <div className="flex items-start space-x-3">
@@ -129,7 +85,6 @@ function App() {
           </div>
         </div>
 
-        {/* Pipeline Component */}
         <div className="bg-white rounded-lg shadow-sm border overflow-hidden">
           {pipeline === 'zero-shot-classification' && (
             <ZeroShotClassification />

src/components/ModelInfo.tsx

@@ -6,14 +6,11 @@ import {
   DatabaseIcon,
   CheckCircle,
   XCircle,
-  ExternalLink,
-  ChevronDown
+  ExternalLink
 } from 'lucide-react'
 import { getModelSize } from '../lib/huggingface'
 import { useModel } from '../contexts/ModelContext'
-import { useEffect, useCallback } from 'react'
-import { QuantizationType, WorkerMessage } from '../types'
-import { getWorker } from '../lib/workerManager'
+import ModelLoader from './ModelLoader'
 
 const ModelInfo = () => {
   const formatNumber = (num: number) => {
@@ -29,96 +26,50 @@ const ModelInfo = () => {
 
   const {
     modelInfo,
-    selectedQuantization,
-    setSelectedQuantization,
-    status,
-    setStatus,
-    setProgress,
-    activeWorker,
-    setActiveWorker,
-    pipeline,
-    workerLoaded,
-    setWorkerLoaded
+    selectedQuantization
   } = useModel()
 
-  useEffect(() => {
-    if (modelInfo.isCompatible && modelInfo.supportedQuantizations.length > 0) {
-      const quantizations = modelInfo.supportedQuantizations
-      let defaultQuant: QuantizationType = 'fp32'
-
-      if (quantizations.includes('int8')) {
-        defaultQuant = 'int8'
-      } else if (quantizations.includes('q8')) {
-        defaultQuant = 'q8'
-      } else if (quantizations.includes('q4')) {
-        defaultQuant = 'q4'
-      }
-
-      setSelectedQuantization(defaultQuant)
-    }
-  }, [
-    modelInfo.supportedQuantizations,
-    modelInfo.isCompatible,
-    setSelectedQuantization
-  ])
-
-  useEffect(() => {
-    const newWorker = getWorker(pipeline)
-    if (!newWorker) {
-      return
-    }
-
-    setStatus('idle')
-    setWorkerLoaded(false)
-    setActiveWorker(newWorker)
-
-    const onMessageReceived = (e: MessageEvent<WorkerMessage>) => {
-      const { status, output } = e.data
-      if (status === 'initiate') {
-        setStatus('loading')
-      } else if (status === 'ready') {
-        setStatus('ready')
-        setWorkerLoaded(true)
-      } else if (status === 'progress' && output) {
-        setStatus('progress')
-        if (
-          output.progress &&
-          typeof output.file === 'string' &&
-          output.file.startsWith('onnx')
-        ) {
-          setProgress(output.progress)
-        }
-      }
-    }
-
-    newWorker.addEventListener('message', onMessageReceived)
-
-    return () => {
-      newWorker.removeEventListener('message', onMessageReceived)
-      // terminateWorker(pipeline);
-    }
-  }, [pipeline, selectedQuantization, setActiveWorker, setStatus, setProgress, setWorkerLoaded]) 
-
-  const loadModel = useCallback(() => {
-    if (!modelInfo.name || !selectedQuantization) return
+  const ModelInfoSkeleton = () => (
+    <div className="mt-5 bg-gradient-to-r from-blue-50 to-indigo-50 px-4 py-3 rounded-lg border border-blue-200 space-y-4 h-full min-h-[160px] animate-pulse w-[400px]">
+      <div className="flex items-center space-x-2">
+        <Bot className="w-4 h-4 text-blue-300" />
+        <div className="h-4 bg-gray-300 rounded w-48"></div>
+        <div className="w-4 h-4 bg-gray-300 rounded-full"></div>
+      </div>
 
-    setStatus('loading')
-    const message = {
-      type: 'load',
-      model: modelInfo.name,
-      quantization: selectedQuantization
-    }
-    activeWorker?.postMessage(message)
-  }, [modelInfo.name, selectedQuantization, setStatus, activeWorker])
+      <div className="flex items-center space-x-2 ml-6">
+        <div className="h-3 bg-gray-200 rounded w-32"></div>
+      </div>
 
-  const busy: boolean = status !== 'idle'
+      <div className="flex items-center justify-self-end space-x-4">
+        <div className="flex items-center space-x-1">
+          <Heart className="w-3 h-3 text-red-300" />
+          <div className="h-3 bg-gray-200 rounded w-8"></div>
+        </div>
+        <div className="flex items-center space-x-1">
+          <Download className="w-3 h-3 text-green-300" />
+          <div className="h-3 bg-gray-200 rounded w-8"></div>
+        </div>
+        <div className="flex items-center space-x-1">
+          <Cpu className="w-3 h-3 text-purple-300" />
+          <div className="h-3 bg-gray-200 rounded w-8"></div>
+        </div>
+        <div className="flex items-center space-x-1">
+          <DatabaseIcon className="w-3 h-3 text-purple-300" />
+          <div className="h-3 bg-gray-200 rounded w-12"></div>
+        </div>
+      </div>
+      <hr className="border-gray-200" />
+      <div className="h-8 bg-gray-200 rounded w-full"></div>
+    </div>
+  )
 
   if (!modelInfo.name) {
-    return null
+    return <ModelInfoSkeleton />
   }
 
   return (
-    <div className="bg-gradient-to-r from-blue-50 to-indigo-50 px-4 py-3 rounded-lg border border-blue-200 space-y-3">
+    <div className="mt-5 bg-gradient-to-r from-blue-50 to-indigo-50 px-4 py-3 rounded-lg border border-blue-200 space-y-3 h-full min-h-[150px]">
       {/* Model Name Row */}
       <div className="flex items-center space-x-2">
         <Bot className="w-4 h-4 text-blue-600" />
@@ -150,16 +101,16 @@ const ModelInfo = () => {
 
       {/* Base Model Link */}
       {modelInfo.baseId && (
-        <div className="flex items-center space-x-2 ml-6">
+        <div className="flex items-center space-x-2 ml-6 text-xs text-gray-600 truncate max-w-100">
           <a
             href={`https://huggingface.co/${modelInfo.baseId}`}
             target="_blank"
             rel="noopener noreferrer"
-            className="text-xs text-gray-600 truncate max-w-100 hover:underline"
+            className=" hover:underline"
             title={`Base model: ${modelInfo.baseId}`}
           >
             <ExternalLink className="w-3 h-3 inline-block mr-1" />
-            {modelInfo.baseId}
+            ({modelInfo.baseId})
           </a>
         </div>
       )}
@@ -180,75 +131,31 @@ const ModelInfo = () => {
           </div>
         )}
 
-        {modelInfo.parameters > 0 && (
-          <div className="flex items-center space-x-1">
-            <Cpu className="w-3 h-3 text-purple-500" />
+        <div className="flex items-center space-x-1">
+          <Cpu className="w-3 h-3 text-purple-500" />
+          {modelInfo.parameters ? (
             <span>{formatNumber(modelInfo.parameters)}</span>
-          </div>
-        )}
+          ) : (
+            <span>?</span>
+          )}
+        </div>
 
-        {modelInfo.parameters > 0 && (
-          <div className="flex items-center space-x-1">
-            <DatabaseIcon className="w-3 h-3 text-purple-500" />
+        <div className="flex items-center space-x-1">
+          <DatabaseIcon className="w-3 h-3 text-purple-500" />
+          {modelInfo.parameters ? (
             <span>
               {`~${getModelSize(
                 modelInfo.parameters,
                 selectedQuantization
               ).toFixed(1)}MB`}
             </span>
-          </div>
-        )}
+          ) : (
+            <span>?</span>
+          )}
+        </div>
       </div>
 
-      {/* Separator */}
-      {modelInfo.isCompatible &&
-        modelInfo.supportedQuantizations.length > 0 && (
-          <hr className="border-gray-200" />
-        )}
-
-      {/* Quantization Dropdown */}
-      {modelInfo.isCompatible &&
-        modelInfo.supportedQuantizations.length > 0 && (
-          <div className="flex items-center space-x-2">
-            <span className="text-xs text-gray-600 font-medium">
-              Quantization:
-            </span>
-            <div className="relative">
-              <select
-                value={selectedQuantization || ''}
-                onChange={(e) =>
-                  setSelectedQuantization(e.target.value as QuantizationType)
-                }
-                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}
-                  </option>
-                ))}
-              </select>
-              <ChevronDown className="absolute right-2 top-1/2 transform -translate-y-1/2 w-3 h-3 text-gray-400 pointer-events-none" />
-            </div>
-          </div>
-        )}
-
-      {/* Load Model Button */}
-      {modelInfo.isCompatible && selectedQuantization && (
-        <div className="flex justify-center">
-          <button
-            className="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"
-            disabled={busy || !selectedQuantization || workerLoaded}
-            onClick={loadModel}
-          >
-            {status === 'loading'
-              ? 'Loading Model...'
-              : workerLoaded
-              ? 'Model Ready'
-              : 'Load Model'}
-          </button>
-        </div>
-      )}
+      <ModelLoader />
 
       {/* Incompatibility Message */}
       {modelInfo.isCompatible === false && modelInfo.incompatibilityReason && (

src/components/ModelLoader.tsx

@@ -0,0 +1,144 @@
+import { useEffect, useCallback } from 'react'
+import { ChevronDown, Loader } from 'lucide-react'
+import { QuantizationType, WorkerMessage } from '../types'
+import { useModel } from '../contexts/ModelContext'
+import { getWorker } from '../lib/workerManager'
+
+const ModelLoader = () => {
+  const {
+    modelInfo,
+    selectedQuantization,
+    setSelectedQuantization,
+    status,
+    progress,
+    setStatus,
+    setProgress,
+    activeWorker,
+    setActiveWorker,
+    pipeline
+  } = useModel()
+
+  useEffect(() => {
+    if (modelInfo.isCompatible && modelInfo.supportedQuantizations.length > 0) {
+      const quantizations = modelInfo.supportedQuantizations
+      let defaultQuant: QuantizationType = 'fp32'
+
+      if (quantizations.includes('int8')) {
+        defaultQuant = 'int8'
+      } else if (quantizations.includes('q8')) {
+        defaultQuant = 'q8'
+      } else if (quantizations.includes('q4')) {
+        defaultQuant = 'q4'
+      }
+
+      setSelectedQuantization(defaultQuant)
+    }
+  }, [
+    modelInfo.supportedQuantizations,
+    modelInfo.isCompatible,
+    setSelectedQuantization
+  ])
+
+  useEffect(() => {
+    const newWorker = getWorker(pipeline)
+    if (!newWorker) {
+      return
+    }
+
+    setStatus('initiate')
+    setActiveWorker(newWorker)
+
+    const onMessageReceived = (e: MessageEvent<WorkerMessage>) => {
+      const { status, output } = e.data
+      if (status === 'ready') {
+        setStatus('ready')
+      } else if (status === 'loading' && output) {
+        setStatus('loading')
+        if (
+          output.progress &&
+          typeof output.file === 'string' &&
+          output.file.startsWith('onnx')
+        ) {
+          setProgress(output.progress)
+        }
+      }
+    }
+
+    newWorker.addEventListener('message', onMessageReceived)
+
+    return () => {
+      newWorker.removeEventListener('message', onMessageReceived)
+      // terminateWorker(pipeline);
+    }
+  }, [pipeline, modelInfo.name, selectedQuantization, setActiveWorker, setStatus, setProgress])
+
+  const loadModel = useCallback(() => {
+    if (!modelInfo.name || !selectedQuantization) return
+
+    setStatus('loading')
+    const message = {
+      type: 'load',
+      model: modelInfo.name,
+      quantization: selectedQuantization
+    }
+    activeWorker?.postMessage(message)
+  }, [modelInfo.name, selectedQuantization, setStatus, activeWorker])
+
+   const ready: boolean = status === 'ready'
+   const busy: boolean = status === 'loading'
+
+  if (!modelInfo.isCompatible || modelInfo.supportedQuantizations.length === 0) {
+    return null
+  }
+
+  return (
+    <div className="space-y-3">
+      <hr className="border-gray-200" />
+
+      <div className="flex items-center justify-between space-x-4">
+        <div className="flex items-center space-x-2">
+          <span className="text-xs text-gray-600 font-medium">
+            Quantization:
+          </span>
+          <div className="relative">
+            <select
+              value={selectedQuantization || ''}
+              onChange={(e) =>
+                setSelectedQuantization(e.target.value as QuantizationType)
+              }
+              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}
+                </option>
+              ))}
+            </select>
+            <ChevronDown className="absolute right-2 top-1/2 transform -translate-y-1/2 w-3 h-3 text-gray-400 pointer-events-none" />
+          </div>
+        </div>
+
+        {selectedQuantization && (
+          <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={(busy && !ready) || !selectedQuantization || ready}
+              onClick={loadModel}
+            >
+              {status === 'loading' && (
+                <>
+                  <Loader className="animate-spin h-4 w-4" />
+                  <span>{progress.toFixed(0)}%</span>
+                </>
+              )}
+              {!ready && !busy ? <span>Load Model</span> : !ready ? null : <span>Model Ready</span>}
+            </button>
+          </div>
+        )}
+      </div>
+    </div>
+  )
+}
+
+export default ModelLoader

src/components/ModelSelector.tsx

@@ -1,5 +1,11 @@
-import React, { useEffect, useState } from 'react'
-import { Listbox, ListboxButton, ListboxOption, ListboxOptions, Transition } from '@headlessui/react'
+import React, { useCallback, useEffect, useState } from 'react'
+import {
+  Listbox,
+  ListboxButton,
+  ListboxOption,
+  ListboxOptions,
+  Transition
+} from '@headlessui/react'
 import { useModel } from '../contexts/ModelContext'
 import { getModelInfo } from '../lib/huggingface'
 import { Heart, Download, ChevronDown, Check, ArrowUpDown } from 'lucide-react'
@@ -50,42 +56,46 @@ const ModelSelector: React.FC = () => {
   }, [models, sortBy, sortOrder])
 
   // Function to fetch detailed model info and set as selected
-  const fetchAndSetModelInfo = async (modelId: string) => {
-    try {
-      const modelInfoResponse = await getModelInfo(modelId)
-      
-      let parameters = 0
-      if (modelInfoResponse.safetensors) {
-        const safetensors = modelInfoResponse.safetensors
-        parameters =
-          safetensors.parameters.BF16 ||
-          safetensors.parameters.F16 ||
-          safetensors.parameters.F32 ||
-          safetensors.parameters.total ||
-          0
-      }
+  const fetchAndSetModelInfo = useCallback(
+    async (modelId: string) => {
+      try {
+        const modelInfoResponse = await getModelInfo(modelId)
 
-      const modelInfo = {
-        id: modelId,
-        name: modelInfoResponse.id || modelId,
-        architecture: modelInfoResponse.config?.architectures?.[0] || 'Unknown',
-        parameters,
-        likes: modelInfoResponse.likes || 0,
-        downloads: modelInfoResponse.downloads || 0,
-        createdAt: modelInfoResponse.createdAt || '',
-        isCompatible: modelInfoResponse.isCompatible,
-        incompatibilityReason: modelInfoResponse.incompatibilityReason,
-        supportedQuantizations: modelInfoResponse.supportedQuantizations,
-        baseId: modelInfoResponse.baseId
-      }
+        let parameters = 0
+        if (modelInfoResponse.safetensors) {
+          const safetensors = modelInfoResponse.safetensors
+          parameters =
+            safetensors.parameters.BF16 ||
+            safetensors.parameters.F16 ||
+            safetensors.parameters.F32 ||
+            safetensors.parameters.total ||
+            0
+        }
 
-      console.log('Fetched model info:', modelInfoResponse)
+        const modelInfo = {
+          id: modelId,
+          name: modelInfoResponse.id || modelId,
+          architecture:
+            modelInfoResponse.config?.architectures?.[0] || 'Unknown',
+          parameters,
+          likes: modelInfoResponse.likes || 0,
+          downloads: modelInfoResponse.downloads || 0,
+          createdAt: modelInfoResponse.createdAt || '',
+          isCompatible: modelInfoResponse.isCompatible,
+          incompatibilityReason: modelInfoResponse.incompatibilityReason,
+          supportedQuantizations: modelInfoResponse.supportedQuantizations,
+          baseId: modelInfoResponse.baseId
+        }
 
-      setModelInfo(modelInfo)
-    } catch (error) {
-      console.error('Error fetching model info:', error)
-    }
-  }
+        console.log('Fetched model info:', modelInfoResponse)
+
+        setModelInfo(modelInfo)
+      } catch (error) {
+        console.error('Error fetching model info:', error)
+      }
+    },
+    [setModelInfo]
+  )
 
   // Update modelInfo to first model when pipeline changes
   useEffect(() => {
@@ -93,7 +103,7 @@ const ModelSelector: React.FC = () => {
       const firstModel = models[0]
       fetchAndSetModelInfo(firstModel.id)
     }
-  }, [pipeline, models])
+  }, [pipeline, models, fetchAndSetModelInfo])
 
   const handleModelSelect = (modelId: string) => {
     fetchAndSetModelInfo(modelId)
@@ -108,35 +118,42 @@ const ModelSelector: React.FC = () => {
     }
   }
 
-  const selectedModel = models.find(model => model.id === modelInfo.id) || models[0]
+  const selectedModel =
+    models.find((model) => model.id === modelInfo.id) || models[0]
 
   return (
     <div className="relative">
-      <Listbox value={selectedModel} onChange={(model) => handleModelSelect(model.id)}>
+      <Listbox
+        value={selectedModel}
+        onChange={(model) => handleModelSelect(model.id)}
+      >
         <div className="relative">
           <ListboxButton className="w-full px-3 py-2 border border-gray-300 rounded-md focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-transparent bg-white text-left flex items-center justify-between">
             <div className="flex items-center justify-between w-full">
               <div className="flex flex-col flex-1 min-w-0">
-                <span className="truncate font-medium">{modelInfo.id || 'Select a model'}</span>
+                <span className="truncate font-medium">
+                  {modelInfo.id || 'Select a model'}
+                </span>
               </div>
-              
+
               <div className="flex items-center space-x-3">
-                {selectedModel && (selectedModel.likes > 0 || selectedModel.downloads > 0) && (
-                  <div className="flex items-center space-x-3 text-xs text-gray-500">
-                    {selectedModel.likes > 0 && (
-                      <div className="flex items-center space-x-1">
-                        <Heart className="w-3 h-3 text-red-500" />
-                        <span>{formatNumber(selectedModel.likes)}</span>
-                      </div>
-                    )}
-                    {selectedModel.downloads > 0 && (
-                      <div className="flex items-center space-x-1">
-                        <Download className="w-3 h-3 text-green-500" />
-                        <span>{formatNumber(selectedModel.downloads)}</span>
-                      </div>
-                    )}
-                  </div>
-                )}
+                {selectedModel &&
+                  (selectedModel.likes > 0 || selectedModel.downloads > 0) && (
+                    <div className="flex items-center space-x-3 text-xs text-gray-500">
+                      {selectedModel.likes > 0 && (
+                        <div className="flex items-center space-x-1">
+                          <Heart className="w-3 h-3 text-red-500" />
+                          <span>{formatNumber(selectedModel.likes)}</span>
+                        </div>
+                      )}
+                      {selectedModel.downloads > 0 && (
+                        <div className="flex items-center space-x-1">
+                          <Download className="w-3 h-3 text-green-500" />
+                          <span>{formatNumber(selectedModel.downloads)}</span>
+                        </div>
+                      )}
+                    </div>
+                  )}
                 <ChevronDown className="w-4 h-4 ui-open:rotate-180 transition-transform flex-shrink-0" />
               </div>
             </div>
@@ -158,7 +175,9 @@ const ModelSelector: React.FC = () => {
                   <button
                     onClick={() => handleSortChange('name')}
                     className={`px-2 py-1 rounded flex items-center space-x-1 ${
-                      sortBy === 'name' ? 'bg-blue-100 text-blue-700' : 'text-gray-600 hover:bg-gray-100'
+                      sortBy === 'name'
+                        ? 'bg-blue-100 text-blue-700'
+                        : 'text-gray-600 hover:bg-gray-100'
                     }`}
                   >
                     <span>Name</span>
@@ -167,7 +186,9 @@ const ModelSelector: React.FC = () => {
                   <button
                     onClick={() => handleSortChange('likes')}
                     className={`px-2 py-1 rounded flex items-center space-x-1 ${
-                      sortBy === 'likes' ? 'bg-blue-100 text-blue-700' : 'text-gray-600 hover:bg-gray-100'
+                      sortBy === 'likes'
+                        ? 'bg-blue-100 text-blue-700'
+                        : 'text-gray-600 hover:bg-gray-100'
                     }`}
                   >
                     <Heart className="w-3 h-3" />
@@ -177,21 +198,29 @@ const ModelSelector: React.FC = () => {
                   <button
                     onClick={() => handleSortChange('downloads')}
                     className={`px-2 py-1 rounded flex items-center space-x-1 ${
-                      sortBy === 'downloads' ? 'bg-blue-100 text-blue-700' : 'text-gray-600 hover:bg-gray-100'
+                      sortBy === 'downloads'
+                        ? 'bg-blue-100 text-blue-700'
+                        : 'text-gray-600 hover:bg-gray-100'
                     }`}
                   >
                     <Download className="w-3 h-3" />
                     <span>Downloads</span>
-                    {sortBy === 'downloads' && <ArrowUpDown className="w-3 h-3" />}
+                    {sortBy === 'downloads' && (
+                      <ArrowUpDown className="w-3 h-3" />
+                    )}
                   </button>
                   <button
                     onClick={() => handleSortChange('createdAt')}
                     className={`px-2 py-1 rounded flex items-center space-x-1 ${
-                      sortBy === 'createdAt' ? 'bg-blue-100 text-blue-700' : 'text-gray-600 hover:bg-gray-100'
+                      sortBy === 'createdAt'
+                        ? 'bg-blue-100 text-blue-700'
+                        : 'text-gray-600 hover:bg-gray-100'
                     }`}
                   >
                     <span>Date</span>
-                    {sortBy === 'createdAt' && <ArrowUpDown className="w-3 h-3" />}
+                    {sortBy === 'createdAt' && (
+                      <ArrowUpDown className="w-3 h-3" />
+                    )}
                   </button>
                 </div>
               </div>
@@ -200,7 +229,7 @@ const ModelSelector: React.FC = () => {
               <div className="overflow-auto max-h-48">
                 {sortedModels.map((model) => {
                   const hasStats = model.likes > 0 || model.downloads > 0
-                  
+
                   return (
                     <ListboxOption
                       key={model.id}
@@ -238,7 +267,7 @@ const ModelSelector: React.FC = () => {
                                   <span>{formatNumber(model.downloads)}</span>
                                 </div>
                               )}
-                              
+
                               {model.createdAt && (
                                 <span className="text-xs text-gray-400">
                                   {model.createdAt.split('T')[0]}

src/components/PipelineSelector.tsx

@@ -42,7 +42,7 @@ const PipelineSelector: React.FC<PipelineSelectorProps> = ({
     <div className="relative">
       <Listbox value={selectedPipeline} onChange={setPipeline}>
         <div className="relative">
-          <ListboxButton className="relative w-full cursor-default rounded-lg bg-white py-2 pl-3 pr-10 text-left shadow-md focus:outline-none focus-visible:border-indigo-500 focus-visible:ring-2 focus-visible:ring-white focus-visible:ring-opacity-75 focus-visible:ring-offset-2 focus-visible:ring-offset-orange-300 sm:text-sm border border-gray-300">
+          <ListboxButton className="relative w-full cursor-default rounded-lg bg-white py-2 pl-3 pr-10 text-left focus:outline-none focus-visible:border-indigo-500 focus-visible:ring-2 focus-visible:ring-white focus-visible:ring-opacity-75 focus-visible:ring-offset-2 focus-visible:ring-offset-orange-300 sm:text-sm border border-gray-300">
             <span className="block truncate font-medium">
               {formatPipelineName(selectedPipeline)}
             </span>
@@ -62,7 +62,7 @@ const PipelineSelector: React.FC<PipelineSelectorProps> = ({
             leaveFrom="transform scale-100 opacity-100"
             leaveTo="transform scale-95 opacity-0"
           >
-            <ListboxOptions className="absolute z-10 mt-1 max-h-60 w-full overflow-auto rounded-md bg-white py-1 text-base shadow-lg ring-1 ring-black ring-opacity-5 focus:outline-none sm:text-sm">
+            <ListboxOptions className="absolute z-10 mt-1 max-h-60 w-full overflow-auto rounded-md bg-white py-1 text-base ring-1 ring-black ring-opacity-5 focus:outline-none sm:text-sm">
               {pipelines.map((p) => (
                 <ListboxOption
                   key={p}

src/components/TextClassification.tsx

@@ -1,13 +1,11 @@
-import { useState, useRef, useEffect, useCallback } from 'react';
+import { useState, useRef, useEffect, useCallback } from 'react'
 import {
   ClassificationOutput,
   TextClassificationWorkerInput,
-  WorkerMessage,
-} from '../types';
-import { useModel } from '../contexts/ModelContext';
-import { getModelInfo } from '../lib/huggingface';
-import { getWorker } from '../lib/workerManager';
-
+  WorkerMessage
+} from '../types'
+import { useModel } from '../contexts/ModelContext'
+import { getWorker } from '../lib/workerManager'
 
 const PLACEHOLDER_TEXTS: string[] = [
   'I absolutely love this product! It exceeded all my expectations.',
@@ -20,61 +18,31 @@ const PLACEHOLDER_TEXTS: string[] = [
   'The product arrived damaged and the return process was a nightmare.',
   'Pretty good overall. A few minor issues but mostly positive experience.',
   'Outstanding! This company really knows how to treat their customers.'
-].sort(() => Math.random() - 0.5);
+].sort(() => Math.random() - 0.5)
 
 function TextClassification() {
   const [text, setText] = useState<string>(PLACEHOLDER_TEXTS.join('\n'))
   const [results, setResults] = useState<ClassificationOutput[]>([])
-  const { setProgress, status, setStatus, modelInfo, setModelInfo, workerLoaded} = useModel()
+  const { status, setStatus, modelInfo } = useModel()
   const workerRef = useRef<Worker | null>(null)
 
 
-  useEffect(() => {
-    if (!modelInfo.id) return;
-    const fetchModelInfo = async () => {
-      try {
-        const modelInfoResponse = await getModelInfo(modelInfo.id)
-        let parameters = 0
-        if (modelInfoResponse.safetensors) {
-          const safetensors = modelInfoResponse.safetensors
-          parameters =
-            (safetensors.parameters.F16 ||
-              safetensors.parameters.F32 ||
-              safetensors.parameters.total ||
-              0)
-        }
-        setModelInfo({
-          ...modelInfo,
-          architecture: modelInfoResponse.config?.architectures[0] ?? '',
-          parameters,
-          likes: modelInfoResponse.likes,
-          downloads: modelInfoResponse.downloads
-        })
-      } catch (error) {
-        console.error('Error fetching model info:', error)
-      }
-    }
-
-    fetchModelInfo()
-  }, [modelInfo.id, setModelInfo])
-
   // We use the `useEffect` hook to setup the worker as soon as the component is mounted.
   useEffect(() => {
-    if(!workerRef.current) {
+    if (!workerRef.current) {
       workerRef.current = getWorker('text-classification')
     }
 
-
     // Create a callback function for messages from the worker thread.
     const onMessageReceived = (e: MessageEvent<WorkerMessage>) => {
       const status = e.data.status
-      if (status === 'output') {
+      if (status === 'ready') {
+        setStatus('ready')
+      } else if (status === 'output') {
         setStatus('output')
         const result = e.data.output!
         setResults((prevResults) => [...prevResults, result])
         console.log(result)
-      } else if (status === 'complete') {
-        setStatus('idle')
       } else if (status === 'error') {
         setStatus('error')
         console.error(e.data.output)
@@ -87,10 +55,10 @@ function TextClassification() {
     // Define a cleanup function for when the component is unmounted.
     return () =>
       workerRef.current?.removeEventListener('message', onMessageReceived)
-  }, [])
+  }, [setStatus])
 
   const classify = useCallback(() => {
-    setStatus('processing')
+    setStatus('loading')
     setResults([]) // Clear previous results
     const message: TextClassificationWorkerInput = {
       type: 'classify',
@@ -98,17 +66,16 @@ function TextClassification() {
       model: modelInfo.id
     }
     workerRef.current?.postMessage(message)
-  }, [text, modelInfo.id])
+  }, [text, modelInfo.id, setStatus])
 
   const busy: boolean = status !== 'ready'
 
-
   const handleClear = (): void => {
     setResults([])
   }
 
   return (
-    <div className="flex flex-col h-[40vh] max-h-[80vh] w-full p-4">
+    <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>
 
       <div className="flex flex-col lg:flex-row gap-4 h-full">
@@ -125,14 +92,14 @@ function TextClassification() {
           <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 || !workerLoaded}
+              disabled={busy}
               onClick={classify}
             >
-              {workerLoaded ? (!busy
-                ? 'Classify Text'
-                : status === 'loading'
-                ? 'Model loading...'
-                : 'Processing...') : 'Load model first'}
+              {status === 'ready'
+                ? !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"
@@ -180,4 +147,4 @@ function TextClassification() {
   )
 }
 
-export default TextClassification;
+export default TextClassification

src/components/ZeroShotClassification.tsx

@@ -4,10 +4,8 @@ import {
   Section,
   WorkerMessage,
   ZeroShotWorkerInput,
-  ModelInfo
 } from '../types'
 import { useModel } from '../contexts/ModelContext'
-import { getModelInfo } from '../lib/huggingface'
 
 const PLACEHOLDER_REVIEWS: string[] = [
   // battery/charging problems
@@ -51,7 +49,7 @@ function ZeroShotClassification() {
     PLACEHOLDER_SECTIONS.map((title) => ({ title, items: [] }))
   )
 
-  const { setProgress, status, setStatus, modelInfo, setModelInfo } = useModel()
+  const { status, setStatus, modelInfo } = useModel()
 
   // Create a reference to the worker object.
   const worker = useRef<Worker | null>(null)
@@ -72,17 +70,8 @@ function ZeroShotClassification() {
     // Create a callback function for messages from the worker thread.
     const onMessageReceived = (e: MessageEvent<WorkerMessage>) => {
       const status = e.data.status
-      if (status === 'initiate') {
-        setStatus('loading')
-      } else if (status === 'ready') {
+      if (status === 'ready') {
         setStatus('ready')
-      } else if (status === 'progress') {
-        setStatus('progress')
-        if (
-          e.data.output.progress &&
-          (e.data.output.file as string).startsWith('onnx')
-        )
-          setProgress(e.data.output.progress)
       } else if (status === 'output') {
         setStatus('output')
         const { sequence, labels, scores } = e.data.output!
@@ -100,9 +89,6 @@ function ZeroShotClassification() {
           }
           return newSections
         })
-      } else if (status === 'complete') {
-        setStatus('idle')
-        setProgress(100)
       } else if (status === 'error') {
         setStatus('error')
         console.error(e.data.output)
@@ -118,7 +104,7 @@ function ZeroShotClassification() {
   }, [sections])
 
   const classify = useCallback(() => {
-    setStatus('processing')
+    setStatus('loading')
     const message: ZeroShotWorkerInput = {
       text,
       labels: sections
@@ -129,7 +115,7 @@ function ZeroShotClassification() {
     worker.current?.postMessage(message)
   }, [text, sections, modelInfo.name])
 
-  const busy: boolean = status !== 'idle'
+  const busy: boolean = status !== 'ready'
 
   const handleAddCategory = (): void => {
     setSections((sections) => {

src/contexts/ModelContext.tsx

@@ -1,11 +1,21 @@
-import React, { createContext, RefObject, useContext, useEffect, useRef, useState } from 'react'
-import { ModelInfo, ModelInfoResponse, QuantizationType } from '../types'
+import React, {
+  createContext,
+  useContext,
+  useEffect,
+  useState
+} from 'react'
+import {
+  ModelInfo,
+  ModelInfoResponse,
+  QuantizationType,
+  WorkerStatus
+} from '../types'
 
 interface ModelContextType {
+  status: WorkerStatus
+  setStatus: (status: WorkerStatus) => void
   progress: number
-  status: string
   setProgress: (progress: number) => void
-  setStatus: (status: string) => void
   modelInfo: ModelInfo
   setModelInfo: (model: ModelInfo) => void
   pipeline: string
@@ -16,22 +26,21 @@ interface ModelContextType {
   setSelectedQuantization: (quantization: QuantizationType) => void
   activeWorker: Worker | null
   setActiveWorker: (worker: Worker | null) => void
-  workerLoaded: boolean
-  setWorkerLoaded: (workerLoaded: boolean) => void
 }
 
 const ModelContext = createContext<ModelContextType | undefined>(undefined)
 
 export function ModelProvider({ children }: { children: React.ReactNode }) {
   const [progress, setProgress] = useState<number>(0)
-  const [status, setStatus] = useState<string>('idle')
+  const [status, setStatus] = useState<WorkerStatus>('initiate')
   const [modelInfo, setModelInfo] = useState<ModelInfo>({} as ModelInfo)
-  const [models, setModels] = useState<ModelInfoResponse[]>([] as ModelInfoResponse[])
+  const [models, setModels] = useState<ModelInfoResponse[]>(
+    [] as ModelInfoResponse[]
+  )
   const [pipeline, setPipeline] = useState<string>('text-classification')
-  const [selectedQuantization, setSelectedQuantization] = useState<QuantizationType>('int8')
+  const [selectedQuantization, setSelectedQuantization] =
+    useState<QuantizationType>('int8')
   const [activeWorker, setActiveWorker] = useState<Worker | null>(null)
-  const [workerLoaded, setWorkerLoaded] = useState<boolean>(false)
-  
 
   // set progress to 0 when model is changed
   useEffect(() => {
@@ -54,9 +63,7 @@ export function ModelProvider({ children }: { children: React.ReactNode }) {
         selectedQuantization,
         setSelectedQuantization,
         activeWorker,
-        setActiveWorker,
-        workerLoaded,
-        setWorkerLoaded
+        setActiveWorker
       }}
     >
       {children}

src/lib/huggingface.ts

@@ -114,7 +114,14 @@ const getModelsByPipeline = async (
   }
   const models = await response.json()
   if (pipeline_tag === 'text-classification') {
-    return models.filter((model: ModelInfoResponse) => !model.tags.includes('reranker') && !model.id.includes('reranker')).slice(0, 10)
+    return models
+      .filter(
+        (model: ModelInfoResponse) =>
+          !model.tags.includes('reranker') &&
+          !model.id.includes('reranker') &&
+          !model.tags.includes('sentence-transformers')
+      )
+      .slice(0, 10)
   }
   return models.slice(0, 10)
 }

src/types.ts

@@ -9,8 +9,10 @@ export interface ClassificationOutput {
   scores: number[]
 }
 
+export type WorkerStatus = 'initiate' | 'ready' | 'output' | 'loading' | 'error'
+
 export interface WorkerMessage {
-  status: 'initiate' | 'ready' | 'output' | 'complete' | 'progress' | 'error'
+  status: WorkerStatus
   progress?: number
   error?: string
   output?: any
@@ -28,7 +30,6 @@ export interface TextClassificationWorkerInput {
   model: string
 }
 
-export type AppStatus = 'idle' | 'loading' | 'processing'
 
 type q8 = 'q8' | 'int8' | 'bnb8' | 'uint8'
 type q4 = 'q4' | 'bnb4'