wip: public access to workers

public/workers/text-classification.js

@@ -0,0 +1,62 @@
+/* eslint-disable no-restricted-globals */
+import { pipeline } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.6.3';
+
+class MyTextClassificationPipeline {
+  static task = 'text-classification'
+  static instance = null
+
+  static async getInstance(model, progress_callback = null) {
+    this.instance = pipeline(this.task, model, {
+      progress_callback
+    })
+    return this.instance
+  }
+}
+
+// Listen for messages from the main thread
+self.addEventListener('message', async (event) => {
+  const { type, model, text } = event.data // Destructure 'type'
+
+  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,
+    (x) => {
+      self.postMessage({ status: 'progress', output: x })
+    }
+  )
+
+  if (type === 'load') {
+    self.postMessage({ status: 'ready' })
+    return
+  }
+
+  if (type === 'classify') {
+    if (!text) {
+      self.postMessage({ status: 'complete' }) // 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: 'complete' })
+  }
+})

src/components/ModelInfo.tsx

@@ -1,8 +1,19 @@
-import { Bot, Heart, Download, Cpu, DatabaseIcon, CheckCircle, XCircle, ExternalLink, ChevronDown } from 'lucide-react'
+import {
+  Bot,
+  Heart,
+  Download,
+  Cpu,
+  DatabaseIcon,
+  CheckCircle,
+  XCircle,
+  ExternalLink,
+  ChevronDown
+} from 'lucide-react'
 import { getModelSize } from '../lib/huggingface'
 import { useModel } from '../contexts/ModelContext'
-import { useEffect } from 'react'
-import { QuantizationType } from '../types'
+import { useEffect, useCallback } from 'react'
+import { QuantizationType, WorkerMessage } from '../types'
+import { getWorker } from '../lib/workerManager'
 
 const ModelInfo = () => {
   const formatNumber = (num: number) => {
@@ -16,14 +27,25 @@ const ModelInfo = () => {
     return num.toString()
   }
 
-  const { modelInfo, selectedQuantization, setSelectedQuantization } = useModel()
+  const {
+    modelInfo,
+    selectedQuantization,
+    setSelectedQuantization,
+    status,
+    setStatus,
+    setProgress,
+    activeWorker,
+    setActiveWorker,
+    pipeline,
+    workerLoaded,
+    setWorkerLoaded
+  } = useModel()
 
-  // Set default quantization when model changes
   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')) {
@@ -31,17 +53,72 @@ const ModelInfo = () => {
       } else if (quantizations.includes('q4')) {
         defaultQuant = 'q4'
       }
-      
+
       setSelectedQuantization(defaultQuant)
     }
-  }, [modelInfo.supportedQuantizations, modelInfo.isCompatible, setSelectedQuantization])
+  }, [
+    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
+
+    setStatus('loading')
+    const message = {
+      type: 'load',
+      model: modelInfo.name,
+      quantization: selectedQuantization
+    }
+    activeWorker?.postMessage(message)
+  }, [modelInfo.name, selectedQuantization, setStatus, activeWorker])
+
+  const busy: boolean = status !== 'idle'
 
   if (!modelInfo.name) {
     return null
   }
 
   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="bg-gradient-to-r from-blue-50 to-indigo-50 px-4 py-3 rounded-lg border border-blue-200 space-y-3">
       {/* Model Name Row */}
       <div className="flex items-center space-x-2">
         <Bot className="w-4 h-4 text-blue-600" />
@@ -70,7 +147,7 @@ const ModelInfo = () => {
           </div>
         )}
       </div>
-      
+
       {/* Base Model Link */}
       {modelInfo.baseId && (
         <div className="flex items-center space-x-2 ml-6">
@@ -87,7 +164,6 @@ const ModelInfo = () => {
         </div>
       )}
 
-
       {/* Stats Row */}
       <div className="flex items-center justify-self-end space-x-4 text-xs text-gray-600">
         {modelInfo.likes > 0 && (
@@ -115,36 +191,62 @@ const ModelInfo = () => {
           <div className="flex items-center space-x-1">
             <DatabaseIcon className="w-3 h-3 text-purple-500" />
             <span>
-              {`~${getModelSize(modelInfo.parameters, selectedQuantization).toFixed(1)}MB`}
+              {`~${getModelSize(
+                modelInfo.parameters,
+                selectedQuantization
+              ).toFixed(1)}MB`}
             </span>
           </div>
         )}
       </div>
 
       {/* Separator */}
-      {modelInfo.isCompatible && modelInfo.supportedQuantizations.length > 0 && (
-        <hr className="border-gray-200" />
-      )}
-      
+      {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" />
+      {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>
       )}
 

src/components/ModelSelector.tsx

@@ -8,7 +8,7 @@ type SortOption = 'likes' | 'downloads' | 'createdAt' | 'name'
 
 const ModelSelector: React.FC = () => {
   const { models, setModelInfo, modelInfo, pipeline } = useModel()
-  const [sortBy, setSortBy] = useState<SortOption>('likes')
+  const [sortBy, setSortBy] = useState<SortOption>('downloads')
   const [sortOrder, setSortOrder] = useState<'asc' | 'desc'>('desc')
 
   const formatNumber = (num: number) => {

src/components/PipelineSelector.tsx

@@ -9,8 +9,8 @@ import {
 import { ChevronDown, Check } from 'lucide-react';
 
 const pipelines = [
-  'zero-shot-classification',
   'text-classification',
+  'zero-shot-classification',
   'text-generation',
   'summarization',
   'feature-extraction',

src/components/TextClassification.tsx

@@ -6,6 +6,7 @@ import {
 } from '../types';
 import { useModel } from '../contexts/ModelContext';
 import { getModelInfo } from '../lib/huggingface';
+import { getWorker } from '../lib/workerManager';
 
 
 const PLACEHOLDER_TEXTS: string[] = [
@@ -24,7 +25,8 @@ const PLACEHOLDER_TEXTS: string[] = [
 function TextClassification() {
   const [text, setText] = useState<string>(PLACEHOLDER_TEXTS.join('\n'))
   const [results, setResults] = useState<ClassificationOutput[]>([])
-  const { setProgress, status, setStatus, modelInfo, setModelInfo, models, setModels} = useModel()
+  const { setProgress, status, setStatus, modelInfo, setModelInfo, workerLoaded} = useModel()
+  const workerRef = useRef<Worker | null>(null)
 
 
   useEffect(() => {
@@ -56,43 +58,23 @@ function TextClassification() {
     fetchModelInfo()
   }, [modelInfo.id, setModelInfo])
 
-  // 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 component is mounted.
   useEffect(() => {
-    if (!worker.current) {
-      // Create the worker if it does not yet exist.
-      worker.current = new Worker(
-        new URL('../workers/text-classification.js', import.meta.url),
-        {
-          type: 'module'
-        }
-      )
+    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 === 'initiate') {
-        setStatus('loading')
-      } else 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') {
+      if (status === 'output') {
         setStatus('output')
         const result = e.data.output!
         setResults((prevResults) => [...prevResults, result])
         console.log(result)
       } else if (status === 'complete') {
         setStatus('idle')
-        setProgress(100)
       } else if (status === 'error') {
         setStatus('error')
         console.error(e.data.output)
@@ -100,21 +82,26 @@ function TextClassification() {
     }
 
     // Attach the callback function as an event listener.
-    worker.current.addEventListener('message', onMessageReceived)
+    workerRef.current?.addEventListener('message', onMessageReceived)
 
     // Define a cleanup function for when the component is unmounted.
     return () =>
-      worker.current?.removeEventListener('message', onMessageReceived)
+      workerRef.current?.removeEventListener('message', onMessageReceived)
   }, [])
 
   const classify = useCallback(() => {
     setStatus('processing')
     setResults([]) // Clear previous results
-    const message: TextClassificationWorkerInput = { text, model: modelInfo.id }
-    worker.current?.postMessage(message)
+    const message: TextClassificationWorkerInput = {
+      type: 'classify',
+      text,
+      model: modelInfo.id
+    }
+    workerRef.current?.postMessage(message)
   }, [text, modelInfo.id])
 
-  const busy: boolean = status !== 'idle'
+  const busy: boolean = status !== 'ready'
+
 
   const handleClear = (): void => {
     setResults([])
@@ -138,14 +125,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}
+              disabled={busy || !workerLoaded}
               onClick={classify}
             >
-              {!busy
+              {workerLoaded ? (!busy
                 ? 'Classify Text'
                 : status === 'loading'
                 ? 'Model loading...'
-                : 'Processing...'}
+                : '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"

src/components/ZeroShotClassification.tsx

@@ -59,13 +59,14 @@ function ZeroShotClassification() {
   // We use the `useEffect` hook to setup the worker as soon as the `App` component is mounted.
   useEffect(() => {
     if (!worker.current) {
+      return
       // Create the worker if it does not yet exist.
-      worker.current = new Worker(
-        new URL('../workers/zero-shot.js', import.meta.url),
-        {
-          type: 'module'
-        }
-      )
+      // 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.

src/contexts/ModelContext.tsx

@@ -1,4 +1,4 @@
-import React, { createContext, useContext, useEffect, useState } from 'react'
+import React, { createContext, RefObject, useContext, useEffect, useRef, useState } from 'react'
 import { ModelInfo, ModelInfoResponse, QuantizationType } from '../types'
 
 interface ModelContextType {
@@ -14,6 +14,10 @@ interface ModelContextType {
   setModels: (models: ModelInfoResponse[]) => void
   selectedQuantization: QuantizationType
   setSelectedQuantization: (quantization: QuantizationType) => void
+  activeWorker: Worker | null
+  setActiveWorker: (worker: Worker | null) => void
+  workerLoaded: boolean
+  setWorkerLoaded: (workerLoaded: boolean) => void
 }
 
 const ModelContext = createContext<ModelContextType | undefined>(undefined)
@@ -23,8 +27,11 @@ export function ModelProvider({ children }: { children: React.ReactNode }) {
   const [status, setStatus] = useState<string>('idle')
   const [modelInfo, setModelInfo] = useState<ModelInfo>({} as ModelInfo)
   const [models, setModels] = useState<ModelInfoResponse[]>([] as ModelInfoResponse[])
-  const [pipeline, setPipeline] = useState<string>('zero-shot-classification')
+  const [pipeline, setPipeline] = useState<string>('text-classification')
   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(() => {
@@ -46,6 +53,10 @@ export function ModelProvider({ children }: { children: React.ReactNode }) {
         setPipeline,
         selectedQuantization,
         setSelectedQuantization,
+        activeWorker,
+        setActiveWorker,
+        workerLoaded,
+        setWorkerLoaded
       }}
     >
       {children}

src/lib/workerManager.ts

@@ -0,0 +1,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];
+  }
+};

src/types.ts

@@ -10,7 +10,9 @@ export interface ClassificationOutput {
 }
 
 export interface WorkerMessage {
-  status: 'initiate' | 'ready' | 'output' | 'complete' | 'progress'
+  status: 'initiate' | 'ready' | 'output' | 'complete' | 'progress' | 'error'
+  progress?: number
+  error?: string
   output?: any
 }
 
@@ -21,6 +23,7 @@ export interface ZeroShotWorkerInput {
 }
 
 export interface TextClassificationWorkerInput {
+  type: 'classify'
   text: string
   model: string
 }

src/workers/text-classification.js

@@ -1,55 +0,0 @@
-/* eslint-disable no-restricted-globals */
-import { pipeline } from '@huggingface/transformers';
-
-class MyTextClassificationPipeline {
-  static task = 'text-classification';
-  static instance = null;
-
-  static async getInstance(model, progress_callback = null) {
-    this.instance ??= pipeline(this.task, model, {
-      progress_callback
-    });
-
-    return this.instance;
-  }
-}
-
-// Listen for messages from the main thread
-self.addEventListener('message', async (event) => {
-  const { text, model } = event.data;
-  if (!model) {
-    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 MyTextClassificationPipeline.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) {
-    if (line.trim()) {
-      const output = await classifier(line);
-      // Send the output back to the main thread
-      self.postMessage({
-        status: 'output',
-        output: {
-          sequence: line,
-          labels: [output[0].label],
-          scores: [output[0].score]
-        }
-      });
-    }
-  }
-  // Send the output back to the main thread
-  self.postMessage({ status: 'complete' });
-});

tsconfig.json

@@ -1,7 +1,12 @@
 {
   "compilerOptions": {
-    "target": "es5",
-    "lib": ["dom", "dom.iterable", "esnext"],
+    "target": "es2020",
+    "lib": [
+      "dom",
+      "dom.iterable",
+      "esnext",
+      "WebWorker"
+    ],
     "allowJs": true,
     "skipLibCheck": true,
     "esModuleInterop": true,
@@ -16,5 +21,5 @@
     "noEmit": true,
     "jsx": "react-jsx"
   },
-  "include": ["src"]
-}
+  "include": ["src", "public/workers"]
+}