improve text-classification components

public/workers/text-classification.js

@@ -41,7 +41,7 @@ class MyTextClassificationPipeline {
 // Listen for messages from the main thread
 self.addEventListener('message', async (event) => {
   try {
-    const { type, model, dtype, text } = event.data
+    const { type, model, dtype, text, config } = event.data
 
     if (!model) {
       self.postMessage({
@@ -76,13 +76,13 @@ self.addEventListener('message', async (event) => {
       const split = text.split('\n')
       for (const line of split) {
         if (line.trim()) {
-          const output = await classifier(line)
+          const output = await classifier(line, config)
           self.postMessage({
             status: 'output',
             output: {
               sequence: line,
-              labels: [output[0].label],
-              scores: [output[0].score]
+              labels: output.map((item) => item.label),
+              scores: output.map((item) => item.score)
             }
           })
         }

src/components/ModelCode.tsx

@@ -38,7 +38,9 @@ const ModelCode = ({ isCodeModalOpen, setIsCodeModalOpen }: ModelCodeProps) => {
     case 'text-classification':
       classType = 'classifier'
       exampleData = 'I love this product!'
-      config = {}
+      config = {
+        top_k: 1
+      }
       break
     case 'text-generation':
       classType = 'generator'

src/components/PipelineLayout.tsx

@@ -3,6 +3,7 @@ import { TextGenerationProvider } from '../contexts/TextGenerationContext'
 import { FeatureExtractionProvider } from '../contexts/FeatureExtractionContext'
 import { ZeroShotClassificationProvider } from '../contexts/ZeroShotClassificationContext'
 import { ImageClassificationProvider } from '../contexts/ImageClassificationContext'
+import { TextClassificationProvider } from '../contexts/TextClassificationContext'
 
 export const PipelineLayout = ({ children }: { children: React.ReactNode }) => {
   const { pipeline } = useModel()
@@ -26,6 +27,9 @@ export const PipelineLayout = ({ children }: { children: React.ReactNode }) => {
         <ImageClassificationProvider>{children}</ImageClassificationProvider>
       )
 
+    case 'text-classification':
+      return <TextClassificationProvider>{children}</TextClassificationProvider>
+
     default:
       return <>{children}</>
   }

src/components/PipelineSelector.tsx

@@ -12,9 +12,9 @@ export const supportedPipelines = [
   'image-classification',
   'text-generation',
   'zero-shot-classification',
-  'text-classification',
-  'summarization',
-  'translation'
+  'text-classification'
+  // 'summarization',
+  // 'translation'
 ]
 
 interface PipelineSelectorProps {

src/components/Sidebar.tsx

@@ -7,6 +7,7 @@ import TextGenerationConfig from './pipelines/TextGenerationConfig'
 import FeatureExtractionConfig from './pipelines/FeatureExtractionConfig'
 import ZeroShotClassificationConfig from './pipelines/ZeroShotClassificationConfig'
 import ImageClassificationConfig from './pipelines/ImageClassificationConfig'
+import TextClassificationConfig from './pipelines/TextClassificationConfig'
 import { Button } from '@/components/ui/button'
 
 interface SidebarProps {
@@ -102,6 +103,7 @@ const Sidebar = ({
             {pipeline === 'image-classification' && (
               <ImageClassificationConfig />
             )}
+            {pipeline === 'text-classification' && <TextClassificationConfig />}
           </div>
         </div>
       </div>

src/components/pipelines/TextClassification.tsx

@@ -1,6 +1,11 @@
 import { useState, useCallback, useEffect } from 'react'
-import { TextClassificationWorkerInput, WorkerMessage } from '../../types'
+import {
+  ClassificationOutput,
+  TextClassificationWorkerInput,
+  WorkerMessage
+} from '../../types'
 import { useModel } from '../../contexts/ModelContext'
+import { useTextClassification } from '../../contexts/TextClassificationContext'
 
 const PLACEHOLDER_TEXTS: string[] = [
   'I absolutely love this product! It exceeded all my expectations.',
@@ -18,7 +23,7 @@ const PLACEHOLDER_TEXTS: string[] = [
 function TextClassification() {
   const [text, setText] = useState<string>(PLACEHOLDER_TEXTS.join('\n'))
   const [numberExamples, setNumberExamples] = useState(PLACEHOLDER_TEXTS.length)
-  const [results, setResults] = useState<any[]>([])
+  const [results, setResults] = useState<ClassificationOutput[]>([])
   const {
     activeWorker,
     status,
@@ -27,6 +32,7 @@ function TextClassification() {
     hasBeenLoaded,
     selectedQuantization
   } = useModel()
+  const { config } = useTextClassification()
 
   useEffect(() => {
     if (modelInfo?.widgetData) {
@@ -51,10 +57,11 @@ function TextClassification() {
       type: 'classify',
       text,
       model: modelInfo.id,
-      dtype: selectedQuantization ?? 'fp32'
+      dtype: selectedQuantization ?? 'fp32',
+      config
     }
     activeWorker.postMessage(message)
-  }, [text, modelInfo, activeWorker, selectedQuantization, setResults])
+  }, [text, modelInfo, activeWorker, selectedQuantization, config, setResults])
 
   // Handle worker messages
   useEffect(() => {
@@ -65,7 +72,7 @@ function TextClassification() {
       if (status === 'output') {
         setStatus('output')
         const result = e.data.output!
-        setResults((prev: any[]) => [...prev, result])
+        setResults((prev: ClassificationOutput[]) => [...prev, result])
       }
     }
 
@@ -135,17 +142,47 @@ function TextClassification() {
               <div className="space-y-3">
                 {results.map((result, index) => (
                   <div key={index} className="p-3 rounded-sm border-2">
-                    <div className="flex justify-between items-start mb-2">
-                      <span className="font-semibold text-sm">
-                        {result.labels[0]}
-                      </span>
-                      <span className="text-sm font-mono">
-                        {(result.scores[0] * 100).toFixed(1)}%
-                      </span>
-                    </div>
-                    <div className="text-sm text-gray-700">
+                    <div className="text-sm text-gray-700 mb-3">
                       {result.sequence}
                     </div>
+                    <div className="space-y-2">
+                      {result.labels.map(
+                        (label: string, labelIndex: number) => {
+                          const score = result.scores[labelIndex]
+                          const isTopPrediction = labelIndex === 0
+
+                          return (
+                            <div
+                              key={labelIndex}
+                              className={`flex justify-between items-center p-2 rounded ${
+                                isTopPrediction
+                                  ? 'bg-blue-50 border-l-4 border-blue-500'
+                                  : 'bg-gray-50'
+                              }`}
+                            >
+                              <span
+                                className={`font-medium text-sm ${
+                                  isTopPrediction
+                                    ? 'text-blue-700'
+                                    : 'text-gray-700'
+                                }`}
+                              >
+                                {label}
+                              </span>
+                              <span
+                                className={`text-sm font-mono ${
+                                  isTopPrediction
+                                    ? 'text-blue-600'
+                                    : 'text-gray-600'
+                                }`}
+                              >
+                                {(score * 100).toFixed(1)}%
+                              </span>
+                            </div>
+                          )
+                        }
+                      )}
+                    </div>
                   </div>
                 ))}
               </div>

src/components/pipelines/TextClassificationConfig.tsx

@@ -0,0 +1,51 @@
+import React from 'react'
+import { useTextClassification } from '../../contexts/TextClassificationContext'
+import { Slider } from '../ui/slider'
+
+const TextClassificationConfig = () => {
+  const { config, setConfig } = useTextClassification()
+
+  return (
+    <div className="space-y-4">
+      <h3 className="text-lg font-semibold text-foreground">
+        Text Classification Settings
+      </h3>
+
+      <div className="space-y-3">
+        <div>
+          <label className="block text-sm font-medium text-foreground/80 mb-1">
+            Top K Predictions: {config.top_k}
+          </label>
+          <Slider
+            defaultValue={[config.top_k]}
+            min={1}
+            max={10}
+            step={1}
+            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">
+            <span>1</span>
+            <span>4</span>
+            <span>7</span>
+            <span>10</span>
+          </div>
+          <p className="text-xs text-muted-foreground mt-1">
+            Number of top predictions to return for each text
+          </p>
+        </div>
+
+        <div className="p-3 bg-chart-4/10 border border-chart-4/20 rounded-lg">
+          <h4 className="text-sm font-medium text-chart-4 mb-2">💡 Tips</h4>
+          <div className="text-xs text-chart-4 space-y-1">
+            <p>• Use Top K = 1-3 for most cases</p>
+            <p>• Higher values show more detailed rankings</p>
+            <p>• Try quantized models for faster processing</p>
+          </div>
+        </div>
+      </div>
+    </div>
+  )
+}
+
+export default TextClassificationConfig

src/contexts/TextClassificationContext.tsx

@@ -0,0 +1,47 @@
+import React, { createContext, useContext, useState } from 'react'
+
+interface TextClassificationConfig {
+  top_k: number
+}
+
+interface TextClassificationContextType {
+  config: TextClassificationConfig
+  setConfig: React.Dispatch<React.SetStateAction<TextClassificationConfig>>
+}
+
+const TextClassificationContext = createContext<
+  TextClassificationContextType | undefined
+>(undefined)
+
+export function useTextClassification() {
+  const context = useContext(TextClassificationContext)
+  if (context === undefined) {
+    throw new Error(
+      'useTextClassification must be used within a TextClassificationProvider'
+    )
+  }
+  return context
+}
+
+interface TextClassificationProviderProps {
+  children: React.ReactNode
+}
+
+export function TextClassificationProvider({
+  children
+}: TextClassificationProviderProps) {
+  const [config, setConfig] = useState<TextClassificationConfig>({
+    top_k: 1
+  })
+
+  const value: TextClassificationContextType = {
+    config,
+    setConfig
+  }
+
+  return (
+    <TextClassificationContext.Provider value={value}>
+      {children}
+    </TextClassificationContext.Provider>
+  )
+}

src/types.ts

@@ -48,6 +48,9 @@ export interface TextClassificationWorkerInput {
   text: string
   model: string
   dtype: QuantizationType
+  config?: {
+    top_k?: number
+  }
 }
 
 export interface TextGenerationWorkerInput {