feat: add dtype support for quantization in model inputs and remove debug logs

src/components/ModelLoader.tsx

@@ -64,8 +64,6 @@ const ModelLoader = () => {
 
     const onMessageReceived = (e: MessageEvent<WorkerMessage>) => {
       const { status, output } = e.data
-      console.log('Received output from worker', e.data)
-
       if (status === 'ready') {
         setStatus('ready')
         if (e.data.output) console.log(e.data.output)

src/components/TextClassification.tsx

@@ -3,7 +3,7 @@ import {
   TextClassificationWorkerInput,
 } from '../types'
 import { useModel } from '../contexts/ModelContext'
-import { set } from 'lodash'
+
 const PLACEHOLDER_TEXTS: string[] = [
   'I absolutely love this product! It exceeded all my expectations.',
   "This is the worst purchase I've ever made. Complete waste of money.",
@@ -19,7 +19,7 @@ const PLACEHOLDER_TEXTS: string[] = [
 
 function TextClassification() {
   const [text, setText] = useState<string>(PLACEHOLDER_TEXTS.join('\n'))
-  const { activeWorker, status, modelInfo, results, setResults, hasBeenLoaded} = useModel()
+  const { activeWorker, status, modelInfo, results, setResults, hasBeenLoaded, selectedQuantization} = useModel()
 
   const classify = useCallback(() => {
     if (!modelInfo || !activeWorker) {
@@ -30,10 +30,11 @@ function TextClassification() {
     const message: TextClassificationWorkerInput = {
       type: 'classify',
       text,
-      model: modelInfo.id
+      model: modelInfo.id,
+      dtype: selectedQuantization ?? 'fp32'
     }
     activeWorker.postMessage(message)
-  }, [text, modelInfo, activeWorker, set])
+  }, [text, modelInfo, activeWorker, selectedQuantization, setResults])
 
   const busy: boolean = status !== 'ready'
 

src/components/TextGeneration.tsx

@@ -30,7 +30,7 @@ function TextGeneration() {
   // Generation state
   const [isGenerating, setIsGenerating] = useState<boolean>(false)
 
-  const { activeWorker, status, modelInfo, hasBeenLoaded } = useModel()
+  const { activeWorker, status, modelInfo, hasBeenLoaded, selectedQuantization } = useModel()
   const messagesEndRef = useRef<HTMLDivElement>(null)
 
   const scrollToBottom = () => {
@@ -73,10 +73,11 @@ function TextGeneration() {
       top_p: topP,
       top_k: topK,
       do_sample: doSample,
+      dtype: selectedQuantization ?? 'fp32'
     }
 
     activeWorker.postMessage(message)
-  }, [currentMessage, messages, modelInfo, activeWorker, temperature, maxTokens, topP, topK, doSample, isGenerating])
+  }, [currentMessage, messages, modelInfo, activeWorker, temperature, maxTokens, topP, topK, doSample, isGenerating, selectedQuantization])
 
   const handleGenerateText = useCallback(() => {
     if (!prompt.trim() || !modelInfo || !activeWorker || isGenerating) {
@@ -94,11 +95,12 @@ function TextGeneration() {
       max_new_tokens: maxTokens,
       top_p: topP,
       top_k: topK,
-      do_sample: doSample
+      do_sample: doSample,
+      dtype: selectedQuantization ?? 'fp32'
     }
 
     activeWorker.postMessage(message)
-  }, [prompt, modelInfo, activeWorker, temperature, maxTokens, topP, topK, doSample, isGenerating])
+  }, [prompt, modelInfo, activeWorker, temperature, maxTokens, topP, topK, doSample, isGenerating, selectedQuantization])
 
   useEffect(() => {
     if (!activeWorker) return

src/components/ZeroShotClassification.tsx

@@ -48,7 +48,7 @@ function ZeroShotClassification() {
     PLACEHOLDER_SECTIONS.map((title) => ({ title, items: [] }))
   )
 
-  const { activeWorker, status, modelInfo, hasBeenLoaded } = useModel()
+  const { activeWorker, status, modelInfo, hasBeenLoaded, selectedQuantization } = useModel()
 
   const classify = useCallback(() => {
     if (!modelInfo || !activeWorker) {
@@ -70,10 +70,11 @@ function ZeroShotClassification() {
       labels: sections
         .slice(0, sections.length - 1)
         .map((section) => section.title),
-      model: modelInfo.id
+      model: modelInfo.id,
+      dtype: selectedQuantization ?? 'fp32'
     }
     activeWorker.postMessage(message)
-  }, [text, sections, modelInfo, activeWorker])
+  }, [text, sections, modelInfo, activeWorker, selectedQuantization])
 
   // Handle worker messages
   useEffect(() => {

src/types.ts

@@ -39,12 +39,14 @@ export interface ZeroShotWorkerInput {
   text: string
   labels: string[]
   model: string
+  dtype: QuantizationType
 }
 
 export interface TextClassificationWorkerInput {
   type: 'classify'
   text: string
   model: string
+  dtype: QuantizationType
 }
 
 export interface TextGenerationWorkerInput {
@@ -58,6 +60,7 @@ export interface TextGenerationWorkerInput {
   top_p?: number
   top_k?: number
   do_sample?: boolean
+  dtype: QuantizationType
 }
 
 const q8Types = ['q8', 'int8', 'bnb8', 'uint8'] as const