Refactor model configuration to use config object

public/workers/image-classification.js

@@ -62,7 +62,7 @@ class MyImageClassificationPipeline {
 // Listen for messages from the main thread
 self.addEventListener('message', async (event) => {
   try {
-    const { type, image, model, dtype, topK = 1 } = event.data
+    const { type, image, model, dtype, config } = event.data
 
     if (!model) {
       self.postMessage({
@@ -100,9 +100,7 @@ self.addEventListener('message', async (event) => {
 
       try {
         // Run classification
-        const output = await classifier(image, {
-          top_k: topK
-        })
+        const output = await classifier(image, config)
 
         // Format predictions
         const predictions = output.map((item) => ({

public/workers/text-generation.js

@@ -49,20 +49,8 @@ class MyTextGenerationPipeline {
 // Listen for messages from the main thread
 self.addEventListener('message', async (event) => {
   try {
-    const {
-      type,
-      model,
-      dtype,
-      messages,
-      prompt,
-      hasChatTemplate,
-      temperature,
-      max_new_tokens,
-      top_p,
-      top_k,
-      do_sample,
-      stop_words
-    } = event.data
+    const { type, model, dtype, messages, prompt, hasChatTemplate, config } =
+      event.data
 
     if (type === 'stop') {
       MyTextGenerationPipeline.stopGeneration()
@@ -108,12 +96,11 @@ self.addEventListener('message', async (event) => {
       }
 
       const options = {
-        max_new_tokens: max_new_tokens || 100,
-        temperature: temperature || 0.7,
-        do_sample: do_sample !== false,
-        ...(top_p && { top_p }),
-        ...(top_k && { top_k }),
-        ...(stop_words && stop_words.length > 0 && { stop_words })
+        max_new_tokens: config.max_new_tokens || 100,
+        temperature: config.temperature || 0.7,
+        do_sample: config.do_sample !== false,
+        ...(config.top_p && { top_p }),
+        ...(config.top_k && { top_k })
       }
 
       // Create an AbortController for this generation

src/components/pipelines/ImageClassification.tsx

@@ -45,7 +45,7 @@ function ImageClassification() {
     removeExample,
     updateExample,
     clearExamples,
-    topK
+    config
   } = useImageClassification()
 
   const [isClassifying, setIsClassifying] = useState<boolean>(false)
@@ -75,12 +75,12 @@ function ImageClassification() {
         image: example.url,
         model: modelInfo.id,
         dtype: selectedQuantization ?? 'fp32',
-        topK
+        config
       }
 
       activeWorker.postMessage(message)
     },
-    [modelInfo, activeWorker, selectedQuantization, topK, updateExample]
+    [modelInfo, activeWorker, selectedQuantization, config, updateExample]
   )
 
   const handleFileSelect = useCallback(

src/components/pipelines/ImageClassificationConfig.tsx

@@ -3,7 +3,7 @@ import { useImageClassification } from '../../contexts/ImageClassificationContex
 import { Slider } from '../ui/slider'
 
 const ImageClassificationConfig = () => {
-  const { topK, setTopK } = useImageClassification()
+  const { config, setConfig } = useImageClassification()
 
   return (
     <div className="space-y-4">
@@ -14,14 +14,14 @@ const ImageClassificationConfig = () => {
       <div className="space-y-3">
         <div>
           <label className="block text-sm font-medium text-foreground/80 mb-1">
-            Top K Predictions: {topK}
+            Top K Predictions: {config.top_k}
           </label>
           <Slider
-            defaultValue={[topK]}
+            defaultValue={[config.top_k]}
             min={1}
             max={10}
             step={1}
-            onValueChange={(value) => setTopK(value[0])}
+            onValueChange={(value) => setConfig({ top_k: value[0] })}
             className="w-full rounded-lg"
           />
           <div className="flex justify-between text-xs text-muted-foreground/60 mt-1">

src/components/pipelines/TextGeneration.tsx

@@ -58,12 +58,8 @@ function TextGeneration() {
       messages: updatedMessages,
       hasChatTemplate: modelInfo.hasChatTemplate,
       model: modelInfo.id,
-      temperature: config.temperature,
-      max_new_tokens: config.maxTokens,
-      top_p: config.topP,
-      top_k: config.topK,
-      do_sample: config.doSample,
-      dtype: selectedQuantization ?? 'fp32'
+      dtype: selectedQuantization ?? 'fp32',
+      config
     }
 
     activeWorker.postMessage(message)
@@ -87,11 +83,7 @@ function TextGeneration() {
       prompt: prompt.trim(),
       hasChatTemplate: modelInfo.hasChatTemplate,
       model: modelInfo.id,
-      temperature: config.temperature,
-      max_new_tokens: config.maxTokens,
-      top_p: config.topP,
-      top_k: config.topK,
-      do_sample: config.doSample,
+      config,
       dtype: selectedQuantization ?? 'fp32'
     }
 

src/contexts/ImageClassificationContext.tsx

@@ -1,6 +1,10 @@
 import React, { createContext, useContext, useState, useCallback } from 'react'
 import { ImageExample } from '../types'
 
+interface ImageClassificationConfig {
+  top_k: number
+}
+
 interface ImageClassificationContextType {
   examples: ImageExample[]
   selectedExample: ImageExample | null
@@ -9,8 +13,8 @@ interface ImageClassificationContextType {
   removeExample: (id: string) => void
   updateExample: (id: string, updates: Partial<ImageExample>) => void
   clearExamples: () => void
-  topK: number
-  setTopK: (k: number) => void
+  config: ImageClassificationConfig
+  setConfig: React.Dispatch<React.SetStateAction<ImageClassificationConfig>>
 }
 
 const ImageClassificationContext = createContext<
@@ -38,7 +42,9 @@ export function ImageClassificationProvider({
   const [selectedExample, setSelectedExample] = useState<ImageExample | null>(
     null
   )
-  const [topK, setTopK] = useState<number>(5)
+  const [config, setConfig] = useState<ImageClassificationConfig>({
+    top_k: 5
+  })
 
   const addExample = useCallback((file: File) => {
     const id = Math.random().toString(36).substr(2, 9)
@@ -105,8 +111,8 @@ export function ImageClassificationProvider({
     removeExample,
     updateExample,
     clearExamples,
-    topK,
-    setTopK
+    config,
+    setConfig
   }
 
   return (

src/types.ts

@@ -56,11 +56,13 @@ export interface TextGenerationWorkerInput {
   messages?: ChatMessage[]
   hasChatTemplate: boolean
   model: string
-  temperature?: number
-  max_new_tokens?: number
-  top_p?: number
-  top_k?: number
-  do_sample?: boolean
+  config?: {
+    temperature?: number
+    max_new_tokens?: number
+    top_p?: number
+    top_k?: number
+    do_sample?: boolean
+  }
   dtype: QuantizationType
 }
 
@@ -80,7 +82,9 @@ export interface ImageClassificationWorkerInput {
   image: string | ImageData | HTMLImageElement | HTMLCanvasElement
   model: string
   dtype: QuantizationType
-  topK?: number
+  config: {
+    top_k?: number
+  }
 }
 
 export interface ImageClassificationResult {