Enhance model handling and loading: add dtype support, improve fetching logic, and refine component interactions

public/workers/text-classification.js

@@ -1,21 +1,23 @@
 /* eslint-disable no-restricted-globals */
-import { pipeline } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.6.3';
+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
-    })
+  static async getInstance(model, dtype = 'fp32', progress_callback = null) {
+    this.instance = pipeline(
+      this.task,
+      model,
+      { dtype, 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'
+  const { type, model, dtype, text } = event.data
 
   if (!model) {
     self.postMessage({
@@ -28,6 +30,7 @@ self.addEventListener('message', async (event) => {
   // 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 })
     }

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,19 +10,24 @@ import ModelInfo from './components/ModelInfo'
 import ModelReadme from './components/ModelReadme'
 
 function App() {
-  const { pipeline, setPipeline, setModels, setModelInfo, modelInfo } = useModel()
-  const [isFetching, setIsFetching] = useState(false)
+  const { pipeline, setPipeline, setModels, setModelInfo, modelInfo, setIsFetching} = useModel()
 
   useEffect(() => {
     setModelInfo(null)
+    setModels([])
+    setIsFetching(true)
+    
     const fetchModels = async () => {
-      setIsFetching(true)
-      const fetchedModels = await getModelsByPipeline(pipeline)
-      setModels(fetchedModels)
-      setIsFetching(false)
+      try {
+        const fetchedModels = await getModelsByPipeline(pipeline)
+        setModels(fetchedModels)
+      } catch (error) {
+        console.error('Error fetching models:', error)
+        setIsFetching(false)
+      }
     }
     fetchModels()
-  }, [setModels, setModelInfo, pipeline])
+  }, [setModels, setModelInfo, setIsFetching, pipeline])
 
   return (
     <div className="min-h-screen bg-gradient-to-br from-blue-50 to-indigo-100">
@@ -47,12 +52,12 @@ function App() {
                   <span className="text-lg font-semibold text-gray-900 block">
                     Select Model
                   </span>
-                  <ModelSelector isFetching={isFetching} />
+                  <ModelSelector />
                 </div>
               </div>
 
               <div className="ml-6">
-                <ModelInfo isFetching={isFetching} />
+                <ModelInfo />
               </div>
             </div>
 

src/components/ModelInfo.tsx

@@ -12,7 +12,7 @@ import { getModelSize } from '../lib/huggingface'
 import { useModel } from '../contexts/ModelContext'
 import ModelLoader from './ModelLoader'
 
-const ModelInfo = ({ isFetching }: { isFetching: boolean }) => {
+const ModelInfo = () => {
   const formatNumber = (num: number) => {
     if (num >= 1000000000) {
       return (num / 1000000000).toFixed(1) + 'B'
@@ -25,8 +25,10 @@ const ModelInfo = ({ isFetching }: { isFetching: boolean }) => {
   }
 
   const {
+    models,
     modelInfo,
-    selectedQuantization
+    selectedQuantization,
+    isFetching
   } = useModel()
 
   const ModelInfoSkeleton = () => (
@@ -64,7 +66,7 @@ const ModelInfo = ({ isFetching }: { isFetching: boolean }) => {
     </div>
   )
 
-  if (!modelInfo || isFetching) {
+  if (!modelInfo || isFetching || models.length === 0) {
     return <ModelInfoSkeleton />
   }
 

src/components/ModelLoader.tsx

@@ -15,13 +15,17 @@ const ModelLoader = () => {
     setProgress,
     activeWorker,
     setActiveWorker,
-    pipeline
+    pipeline,
+    setResults,
+    hasBeenLoaded,
+    setHasBeenLoaded
   } = useModel()
 
+
   useEffect(() => {
     if (!modelInfo) return
 
-    if (modelInfo.isCompatible && modelInfo.supportedQuantizations.length > 0) {
+    if (modelInfo.isCompatible) {
       const quantizations = modelInfo.supportedQuantizations
       let defaultQuant: QuantizationType = 'fp32'
 
@@ -35,10 +39,9 @@ const ModelLoader = () => {
 
       setSelectedQuantization(defaultQuant)
     }
-  }, [
-    modelInfo,
-    setSelectedQuantization
-  ])
+
+    setHasBeenLoaded(false)
+  }, [modelInfo, setSelectedQuantization, setHasBeenLoaded])
 
   useEffect(() => {
     if (!modelInfo) return
@@ -48,14 +51,18 @@ const ModelLoader = () => {
       return
     }
 
-    setStatus('initiate')
-    setActiveWorker(newWorker)
+    if (!hasBeenLoaded) {
+      setStatus('initiate')
+      setActiveWorker(newWorker)
+    }
+
 
     const onMessageReceived = (e: MessageEvent<WorkerMessage>) => {
       const { status, output } = e.data
       if (status === 'ready') {
         setStatus('ready')
-      } else if (status === 'loading' && output) {
+        setHasBeenLoaded(true)
+      } else if (status === 'loading' && output && !hasBeenLoaded) {
         setStatus('loading')
         if (
           output.progress &&
@@ -64,6 +71,14 @@ 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)
       }
     }
 
@@ -73,24 +88,30 @@ const ModelLoader = () => {
       newWorker.removeEventListener('message', onMessageReceived)
       // terminateWorker(pipeline);
     }
-  }, [pipeline, modelInfo, selectedQuantization, setActiveWorker, setStatus, setProgress])
+  }, [
+    pipeline,
+    modelInfo,
+    selectedQuantization,
+    setActiveWorker,
+    setStatus,
+    setProgress,
+    setResults,
+    hasBeenLoaded,
+    setHasBeenLoaded
+  ])
 
   const loadModel = useCallback(() => {
     if (!modelInfo || !selectedQuantization) return
 
-    setStatus('loading')
     const message = {
       type: 'load',
       model: modelInfo.name,
-      quantization: selectedQuantization
+      dtype: selectedQuantization ?? 'fp32'
     }
     activeWorker?.postMessage(message)
-  }, [modelInfo, selectedQuantization, setStatus, activeWorker])
-
-   const ready: boolean = status === 'ready'
-   const busy: boolean = status === 'loading'
+  }, [modelInfo, selectedQuantization, activeWorker])
 
-  if (!modelInfo?.isCompatible || modelInfo.supportedQuantizations.length === 0) {
+  if (!modelInfo?.isCompatible) {
     return null
   }
 
@@ -100,42 +121,52 @@ const ModelLoader = () => {
 
       <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>
+          {modelInfo.supportedQuantizations.length > 1 ? (
+            <>
+              <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>
+            </>
+          ) : (
+            <span className="text-xs text-gray-600 font-medium white-space-break-spaces">
+              No quantization available. Using fp32
+            </span>
+          )}
         </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}
+              disabled={hasBeenLoaded}
               onClick={loadModel}
             >
-              {status === 'loading' && (
+              {status === 'loading' && !hasBeenLoaded ? (
                 <>
                   <Loader className="animate-spin h-4 w-4" />
                   <span>{progress.toFixed(0)}%</span>
                 </>
+              ) : (
+                <span>{!hasBeenLoaded ? 'Load Model' : 'Model Ready'}</span>
               )}
-              {!ready && !busy ? <span>Load Model</span> : !ready ? null : <span>Model Ready</span>}
             </button>
           </div>
         )}

src/components/ModelSelector.tsx

@@ -22,8 +22,15 @@ import {
 
 type SortOption = 'likes' | 'downloads' | 'createdAt' | 'name'
 
-function ModelSelector({ isFetching }: { isFetching: boolean }) {
-  const { models, setModelInfo, modelInfo, pipeline } = useModel()
+function ModelSelector() {
+  const {
+    models,
+    setModelInfo,
+    modelInfo,
+    pipeline,
+    isFetching,
+    setIsFetching
+  } = useModel()
   const [sortBy, setSortBy] = useState<SortOption>('downloads')
   const [sortOrder, setSortOrder] = useState<'asc' | 'desc'>('desc')
   const [showCustomInput, setShowCustomInput] = useState(false)
@@ -102,25 +109,36 @@ function ModelSelector({ isFetching }: { isFetching: boolean }) {
           baseId: modelInfoResponse.baseId,
           readme: modelInfoResponse.readme
         }
+
+        console.log('Fetched model info:', modelInfoResponse)
+
         setModelInfo(modelInfo)
         setIsCustomModel(isCustom)
+        setIsFetching(false)
       } catch (error) {
         console.error('Error fetching model info:', error)
+        setIsFetching(false)
         throw error
       }
     },
-    [setModelInfo, pipeline]
+    [setModelInfo, pipeline, setIsFetching]
   )
 
-  // Update modelInfo to first model when pipeline changes
+  // Reset custom model state when pipeline changes
   useEffect(() => {
-    if (isFetching) return
+    setIsCustomModel(false)
+    setShowCustomInput(false)
+    setCustomModelName('')
+    setCustomModelError('')
+  }, [pipeline])
 
-    if (models.length > 0 && !isCustomModel) {
-      const firstModel = models[0]
+  // Update modelInfo to first model when models are loaded and no custom model is selected
+  useEffect(() => {
+    if (models.length > 0 && !isCustomModel && !modelInfo) {
+      const firstModel = sortedModels[0]
       fetchAndSetModelInfo(firstModel.id, false)
     }
-  }, [pipeline, models, fetchAndSetModelInfo, isCustomModel, isFetching])
+  }, [models, sortedModels, fetchAndSetModelInfo, isCustomModel, modelInfo])
 
   const handleModelSelect = (modelId: string) => {
     fetchAndSetModelInfo(modelId, false)
@@ -160,8 +178,8 @@ function ModelSelector({ isFetching }: { isFetching: boolean }) {
   const handleRemoveCustomModel = () => {
     setIsCustomModel(false)
     // Load the first model from the list
-    if (models.length > 0) {
-      fetchAndSetModelInfo(models[0].id, false)
+    if (sortedModels.length > 0) {
+      fetchAndSetModelInfo(sortedModels[0].id, false)
     }
   }
 
@@ -226,7 +244,7 @@ function ModelSelector({ isFetching }: { isFetching: boolean }) {
     )
   }
 
-  if (isFetching) {
+  if (isFetching || models.length === 0) {
     return (
       <div className="relative">
         <div className="w-full px-3 py-2 border border-gray-300 rounded-md bg-white flex items-center justify-between animate-pulse h-10">

src/components/TextClassification.tsx

@@ -22,52 +22,23 @@ const PLACEHOLDER_TEXTS: string[] = [
 
 function TextClassification() {
   const [text, setText] = useState<string>(PLACEHOLDER_TEXTS.join('\n'))
-  const [results, setResults] = useState<ClassificationOutput[]>([])
-  const { status, setStatus, modelInfo } = useModel()
-  const workerRef = useRef<Worker | null>(null)
+  const { activeWorker, status, setStatus, modelInfo, results, setResults, hasBeenLoaded} = useModel()
 
 
-  // We use the `useEffect` hook to setup the worker as soon as the component is mounted.
-  useEffect(() => {
-    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 === 'ready') {
-        setStatus('ready')
-      } else if (status === 'output') {
-        setStatus('output')
-        const result = e.data.output!
-        setResults((prevResults) => [...prevResults, result])
-        console.log(result)
-      } else if (status === 'error') {
-        setStatus('error')
-        console.error(e.data.output)
-      }
-    }
-
-    // Attach the callback function as an event listener.
-    workerRef.current?.addEventListener('message', onMessageReceived)
-
-    // Define a cleanup function for when the component is unmounted.
-    return () =>
-      workerRef.current?.removeEventListener('message', onMessageReceived)
-  }, [setStatus])
 
   const classify = useCallback(() => {
-    if (!modelInfo) return
-    setStatus('loading')
+    if (!modelInfo || !activeWorker) {
+      console.error('Model info or worker is not available')
+      return
+    }
     setResults([]) // Clear previous results
     const message: TextClassificationWorkerInput = {
       type: 'classify',
       text,
       model: modelInfo.id
     }
-    workerRef.current?.postMessage(message)
-  }, [text, modelInfo, setStatus])
+    activeWorker.postMessage(message)
+  }, [text, modelInfo, setStatus, activeWorker])
 
   const busy: boolean = status !== 'ready'
 
@@ -96,8 +67,7 @@ function TextClassification() {
               disabled={busy}
               onClick={classify}
             >
-              {status === 'ready'
-                ? !busy
+              {hasBeenLoaded ? !busy
                   ? 'Classify Text'
                   : 'Processing...'
                 : 'Load model first'}

src/contexts/ModelContext.tsx

@@ -26,6 +26,12 @@ interface ModelContextType {
   setSelectedQuantization: (quantization: QuantizationType) => void
   activeWorker: Worker | null
   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
 }
 
 const ModelContext = createContext<ModelContextType | undefined>(undefined)
@@ -41,6 +47,10 @@ export function ModelProvider({ children }: { children: React.ReactNode }) {
   const [selectedQuantization, setSelectedQuantization] =
     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(() => {
@@ -63,7 +73,13 @@ export function ModelProvider({ children }: { children: React.ReactNode }) {
         selectedQuantization,
         setSelectedQuantization,
         activeWorker,
-        setActiveWorker
+        setActiveWorker,
+        isFetching,
+        setIsFetching,
+        results,
+        setResults,
+        hasBeenLoaded,
+        setHasBeenLoaded
       }}
     >
       {children}

src/lib/huggingface.ts

@@ -1,7 +1,10 @@
-import { supportedPipelines } from "../components/PipelineSelector"
-import { ModelInfoResponse, QuantizationType } from "../types"
+import { supportedPipelines } from '../components/PipelineSelector'
+import { ModelInfoResponse, QuantizationType } from '../types'
 
-const getModelInfo = async (modelName: string, pipeline: string): Promise<ModelInfoResponse> => {
+const getModelInfo = async (
+  modelName: string,
+  pipeline: string
+): Promise<ModelInfoResponse> => {
   const token = process.env.REACT_APP_HUGGINGFACE_TOKEN
 
   if (!token) {
@@ -23,36 +26,53 @@ const getModelInfo = async (modelName: string, pipeline: string): Promise<ModelI
   if (!response.ok) {
     throw new Error(`Failed to fetch model info: ${response.statusText}`)
   }
-  
+
   const modelData: ModelInfoResponse = await response.json()
-  
+
   const requiredFiles = [
     'config.json',
     'tokenizer.json',
-    'tokenizer_config.json',
+    'tokenizer_config.json'
   ]
-  
-  const siblingFiles = modelData.siblings?.map(s => s.rfilename) || []
-  const missingFiles = requiredFiles.filter(file => !siblingFiles.includes(file))
-  const hasOnnxFolder = siblingFiles.some((file) => file.endsWith('.onnx') && file.startsWith('onnx/'))
 
-  const isCompatible = missingFiles.length === 0 && hasOnnxFolder && modelData.tags.includes(pipeline)
+  const siblingFiles = modelData.siblings?.map((s) => s.rfilename) || []
+  const missingFiles = requiredFiles.filter(
+    (file) => !siblingFiles.includes(file)
+  )
+  const hasOnnxFolder = siblingFiles.some(
+    (file) => file.endsWith('.onnx') && file.startsWith('onnx/')
+  )
+
+  const isCompatible =
+    missingFiles.length === 0 &&
+    hasOnnxFolder &&
+    modelData.tags.includes(pipeline)
 
-  
   let incompatibilityReason = ''
   if (!modelData.tags.includes(pipeline)) {
-    const expectedPipelines = modelData.tags.filter(tag => supportedPipelines.includes(tag)).join(', ')
-    incompatibilityReason = expectedPipelines ? `- Model can be used with ${expectedPipelines} pipelines only\n` : `- Pipeline ${pipeline} not supported by the model\n` 
-  } if (missingFiles.length > 0) {
-    incompatibilityReason += `- Missing required files: ${missingFiles.join(', ')}\n`
-  } else if (!hasOnnxFolder)  {
+    const expectedPipelines = modelData.tags
+      .filter((tag) => supportedPipelines.includes(tag))
+      .join(', ')
+    incompatibilityReason = expectedPipelines
+      ? `- Model can be used with ${expectedPipelines} pipelines only\n`
+      : `- Pipeline ${pipeline} not supported by the model\n`
+  }
+  if (missingFiles.length > 0) {
+    incompatibilityReason += `- Missing required files: ${missingFiles.join(
+      ', '
+    )}\n`
+  } else if (!hasOnnxFolder) {
     incompatibilityReason += '- Folder onnx/ is missing\n'
   }
-  const supportedQuantizations = siblingFiles
-      .filter((file) => file.endsWith('.onnx') && file.includes('_'))
-      .map((file) => file.split('/')[1].split('_')[1].split('.')[0])
-      .filter((q) => q !== 'quantized')
-  const uniqueSupportedQuantizations = Array.from(new Set(supportedQuantizations))
+  const supportedQuantizations = hasOnnxFolder
+    ? siblingFiles
+        .filter((file) => file.endsWith('.onnx') && file.includes('_'))
+        .map((file) => file.split('/')[1].split('_')[1].split('.')[0])
+        .filter((q) => q !== 'quantized')
+    : []
+  const uniqueSupportedQuantizations = Array.from(
+    new Set(supportedQuantizations)
+  )
   uniqueSupportedQuantizations.sort((a, b) => {
     const getNumericValue = (str: string) => {
       const match = str.match(/(\d+)/)
@@ -64,7 +84,9 @@ const getModelInfo = async (modelName: string, pipeline: string): Promise<ModelI
   // Fetch README content
   const fetchReadme = async (modelId: string): Promise<string> => {
     try {
-      const readmeResponse = await fetch(`https://huggingface.co/${modelId}/raw/main/README.md`)
+      const readmeResponse = await fetch(
+        `https://huggingface.co/${modelId}/raw/main/README.md`
+      )
       if (readmeResponse.ok) {
         return await readmeResponse.text()
       }
@@ -74,7 +96,7 @@ const getModelInfo = async (modelName: string, pipeline: string): Promise<ModelI
     return ''
   }
 
-  const baseModel = modelData.cardData?.base_model ?? modelData.modelId 
+  const baseModel = modelData.cardData?.base_model ?? modelData.modelId
   if (baseModel && !modelData.safetensors) {
     const baseModelResponse = await fetch(
       `https://huggingface.co/api/models/${baseModel}`,
@@ -89,21 +111,22 @@ const getModelInfo = async (modelName: string, pipeline: string): Promise<ModelI
     if (baseModelResponse.ok) {
       const baseModelData: ModelInfoResponse = await baseModelResponse.json()
       const readme = await fetchReadme(baseModel)
-      
+
       return {
         ...baseModelData,
         id: modelData.id,
         baseId: baseModel,
         isCompatible,
         incompatibilityReason,
-        supportedQuantizations: uniqueSupportedQuantizations as QuantizationType[],
+        supportedQuantizations:
+          uniqueSupportedQuantizations as QuantizationType[],
         readme
       }
     }
   }
-  
+
   const readme = await fetchReadme(modelData.id)
-  
+
   return {
     ...modelData,
     isCompatible,
@@ -135,7 +158,9 @@ const getModelsByPipeline = async (
     }
   )
   if (!response1.ok) {
-    throw new Error(`Failed to fetch models for pipeline: ${response1.statusText}`)
+    throw new Error(
+      `Failed to fetch models for pipeline: ${response1.statusText}`
+    )
   }
   const models1 = await response1.json()
 
@@ -150,14 +175,18 @@ const getModelsByPipeline = async (
     }
   )
   if (!response2.ok) {
-    throw new Error(`Failed to fetch models for pipeline: ${response2.statusText}`)
+    throw new Error(
+      `Failed to fetch models for pipeline: ${response2.statusText}`
+    )
   }
   const models2 = await response2.json()
 
   // Combine and deduplicate models based on id
-  const combinedModels = [...models1, ...models2].filter((m: ModelInfoResponse) => m.createdAt > '2022/02/03')
-  const uniqueModels = combinedModels.filter((model, index, self) => 
-    index === self.findIndex(m => m.id === model.id)
+  const combinedModels = [...models1, ...models2].filter(
+    (m: ModelInfoResponse) => m.createdAt > '2022/02/03'
+  )
+  const uniqueModels = combinedModels.filter(
+    (model, index, self) => index === self.findIndex((m) => m.id === model.id)
   )
 
   if (pipelineTag === 'text-classification') {
@@ -171,11 +200,10 @@ const getModelsByPipeline = async (
       )
       .slice(0, 20)
   }
-  
+
   return uniqueModels.slice(0, 20)
 }
 
-
 const getModelsByPipelineCustom = async (
   searchString: string,
   pipelineTag: string
@@ -197,12 +225,16 @@ const getModelsByPipelineCustom = async (
     }
   )
 
-    if (!response.ok) {
-    throw new Error(`Failed to fetch models for pipeline: ${response.statusText}`)
+  if (!response.ok) {
+    throw new Error(
+      `Failed to fetch models for pipeline: ${response.statusText}`
+    )
   }
   const models = await response.json()
 
-  const uniqueModels = models.filter((m: ModelInfoResponse) => m.createdAt > '2022/02/03')
+  const uniqueModels = models.filter(
+    (m: ModelInfoResponse) => m.createdAt > '2022/02/03'
+  )
   if (pipelineTag === 'text-classification') {
     return uniqueModels
       .filter(
@@ -214,7 +246,7 @@ const getModelsByPipelineCustom = async (
       )
       .slice(0, 20)
   }
-  
+
   return uniqueModels.slice(0, 20)
 }
 
@@ -239,9 +271,10 @@ function getModelSize(
       bytesPerParameter = 1
       break
     case 'bnb4':
-    case 'q4': 
+    case 'q4':
+    case 'q4f16':
       bytesPerParameter = 0.5
-    break
+      break
   }
 
   const sizeInBytes = parameters * bytesPerParameter
@@ -250,5 +283,9 @@ function getModelSize(
   return sizeInMB
 }
 
-
-export { getModelInfo, getModelSize, getModelsByPipeline, getModelsByPipelineCustom }
+export {
+  getModelInfo,
+  getModelSize,
+  getModelsByPipeline,
+  getModelsByPipelineCustom
+}

src/types.ts

@@ -32,7 +32,7 @@ export interface TextClassificationWorkerInput {
 
 
 type q8 = 'q8' | 'int8' | 'bnb8' | 'uint8'
-type q4 = 'q4' | 'bnb4'
+type q4 = 'q4' | 'bnb4' | 'q4f16'
 type fp16 = 'fp16'
 type fp32 = 'fp32'