remove pooling max since is not supported by transformers.js, and fix

src/components/ModelLoader.tsx

@@ -1,13 +1,13 @@
 import { useEffect, useCallback, useState } from 'react'
-import { ChevronDown, Loader, X } from 'lucide-react'
+import { ChevronDown, Loader2, X } from 'lucide-react'
 import { QuantizationType, WorkerMessage } from '../types'
 import { useModel } from '../contexts/ModelContext'
 import { getWorker } from '../lib/workerManager'
 import { Alert, AlertDescription } from './ui/alert'
 
 const ModelLoader = () => {
-  const [isError, setIsError] = useState(false)
-  const [errorMessage, setErrorMessage] = useState('')
+  const [showAlert, setShowAlert] = useState(false)
+  const [alertMessage, setAlertMessage] = useState<React.ReactNode>('')
   const {
     modelInfo,
     selectedQuantization,
@@ -75,16 +75,23 @@ const ModelLoader = () => {
           output.file.startsWith('onnx')
         ) {
           setProgress(output.progress)
+          setShowAlert(true)
+          setAlertMessage(
+            <div className="flex items-center">
+              <Loader2 className="animate-spin h-4 w-4 mr-2" />
+              Loading Model
+            </div>
+          )
         }
       } else if (status === 'error') {
         setStatus('error')
         const error = e.data.output
         console.error(error)
-        setErrorMessage(error.split('.')[0] + '. See console for details.')
-        setIsError(true)
+        setAlertMessage(error.split('.')[0] + '. See console for details.')
+        setShowAlert(true)
         setTimeout(() => {
-          setIsError(false)
-          setErrorMessage('')
+          setShowAlert(false)
+          setAlertMessage('')
         }, 3000)
       }
     }
@@ -106,6 +113,15 @@ const ModelLoader = () => {
     setHasBeenLoaded
   ])
 
+  useEffect(() => {
+    if (progress === 100) {
+      setTimeout(() => {
+        setShowAlert(false)
+        setAlertMessage('')
+      }, 2000)
+    }
+  }, [progress])
+
   const loadModel = useCallback(() => {
     if (!modelInfo || !selectedQuantization) return
 
@@ -164,7 +180,7 @@ const ModelLoader = () => {
             >
               {status === 'loading' && !hasBeenLoaded ? (
                 <>
-                  <Loader className="animate-spin h-4 w-4" />
+                  <Loader2 className="animate-spin h-4 w-4" />
                   <span>{progress.toFixed(0)}%</span>
                 </>
               ) : (
@@ -174,10 +190,12 @@ const ModelLoader = () => {
           </div>
         )}
       </div>
-      {isError && (
+      {showAlert && (
         <div className="fixed bottom-0 right-0 m-2">
-          <Alert variant="destructive">
-            <AlertDescription>{errorMessage}</AlertDescription>
+          <Alert
+            variant={`${typeof alertMessage === 'string' ? 'destructive' : 'default'}`}
+          >
+            <AlertDescription>{alertMessage}</AlertDescription>
           </Alert>
         </div>
       )}

src/components/PipelineSelector.tsx

@@ -5,7 +5,7 @@ import {
   SelectItem,
   SelectTrigger,
   SelectValue
-} from '@/components/ui/select' // Adjust the import path as needed
+} from '@/components/ui/select'
 
 export const supportedPipelines = [
   'feature-extraction',

src/components/pipelines/FeatureExtraction.tsx

@@ -46,7 +46,13 @@ function FeatureExtraction() {
 
   const [newExampleText, setNewExampleText] = useState<string>('')
   const [isExtracting, setIsExtracting] = useState<boolean>(false)
-  const [showVisualization, setShowVisualization] = useState<boolean>(true)
+  const [showVisualization, setShowVisualization] = useState<boolean>(() => {
+    if (typeof window !== 'undefined') {
+      return window.innerWidth >= 768
+    }
+    return true
+  })
+
   const [progress, setProgress] = useState<{
     completed: number
     total: number
@@ -60,6 +66,18 @@ function FeatureExtraction() {
     selectedQuantization
   } = useModel()
 
+  useEffect(() => {
+    const handleResize = () => {
+      setShowVisualization(window.innerWidth >= 768)
+    }
+
+    window.addEventListener('resize', handleResize)
+
+    return () => {
+      window.removeEventListener('resize', handleResize)
+    }
+  }, [])
+
   const chartRef = useRef<SVGSVGElement>(null)
 
   // PCA reduction to 2D for visualization
@@ -215,7 +233,6 @@ function FeatureExtraction() {
             })
           }
         })
-        console.log({ examples })
         setIsExtracting(false)
         setProgress(null)
       } else if (status === 'error') {

src/components/pipelines/FeatureExtractionConfig.tsx

@@ -1,5 +1,12 @@
 import React from 'react'
 import { useFeatureExtraction } from '../../contexts/FeatureExtractionContext'
+import {
+  Select,
+  SelectContent,
+  SelectItem,
+  SelectTrigger,
+  SelectValue
+} from '@/components/ui/select'
 
 const FeatureExtractionConfig = () => {
   const { config, setConfig } = useFeatureExtraction()
@@ -15,20 +22,36 @@ const FeatureExtractionConfig = () => {
           <label className="block text-sm font-medium text-foreground/80 mb-1">
             Pooling Strategy
           </label>
-          <select
+          <Select
             value={config.pooling}
-            onChange={(e) =>
+            onValueChange={(value) =>
               setConfig((prev) => ({
                 ...prev,
-                pooling: e.target.value as 'mean' | 'cls' | 'max'
+                pooling: value as 'mean' | 'cls'
               }))
             }
-            className="w-full px-3 py-2 border border-input rounded-md shadow-xs focus:outline-hidden focus:ring-2 focus:ring-ring focus:border-ring text-sm"
           >
-            <option value="mean">Mean Pooling</option>
-            <option value="cls">CLS Token</option>
-            <option value="max">Max Pooling</option>
-          </select>
+            <SelectTrigger className="w-full text-sm xl:text-base">
+              <SelectValue placeholder="Select a pooling strategy" />
+            </SelectTrigger>
+            <SelectContent>
+              <SelectItem
+                key="mean"
+                value="mean"
+                className="text-sm data-[state=checked]:font-bold"
+              >
+                Mean Pooling
+              </SelectItem>
+              <SelectItem
+                key="cls"
+                value="cls"
+                className="text-sm data-[state=checked]:font-bold"
+              >
+                CLS Token
+              </SelectItem>
+            </SelectContent>
+          </Select>
+
           <p className="text-xs text-muted-foreground mt-1">
             How to aggregate token embeddings into sentence embeddings
           </p>
@@ -66,10 +89,6 @@ const FeatureExtractionConfig = () => {
             <strong>CLS Token:</strong> Use the [CLS] token embedding (if
             available)
           </p>
-          <p>
-            <strong>Max Pooling:</strong> Take element-wise maximum across
-            tokens
-          </p>
         </div>
       </div>
     </div>

src/contexts/FeatureExtractionContext.tsx

@@ -2,7 +2,7 @@ import React, { createContext, useContext, useState, useCallback } from 'react'
 import { EmbeddingExample, SimilarityResult } from '../types'
 
 interface FeatureExtractionConfig {
-  pooling: 'mean' | 'cls' | 'max'
+  pooling: 'mean' | 'cls'
   normalize: boolean
 }
 
@@ -10,7 +10,9 @@ interface FeatureExtractionContextType {
   examples: EmbeddingExample[]
   setExamples: React.Dispatch<React.SetStateAction<EmbeddingExample[]>>
   selectedExample: EmbeddingExample | null
-  setSelectedExample: React.Dispatch<React.SetStateAction<EmbeddingExample | null>>
+  setSelectedExample: React.Dispatch<
+    React.SetStateAction<EmbeddingExample | null>
+  >
   similarities: SimilarityResult[]
   setSimilarities: React.Dispatch<React.SetStateAction<SimilarityResult[]>>
   config: FeatureExtractionConfig
@@ -22,12 +24,16 @@ interface FeatureExtractionContextType {
   clearExamples: () => void
 }
 
-const FeatureExtractionContext = createContext<FeatureExtractionContextType | undefined>(undefined)
+const FeatureExtractionContext = createContext<
+  FeatureExtractionContextType | undefined
+>(undefined)
 
 export const useFeatureExtraction = () => {
   const context = useContext(FeatureExtractionContext)
   if (!context) {
-    throw new Error('useFeatureExtraction must be used within a FeatureExtractionProvider')
+    throw new Error(
+      'useFeatureExtraction must be used within a FeatureExtractionProvider'
+    )
   }
   return context
 }
@@ -58,9 +64,12 @@ const cosineSimilarity = (a: number[], b: number[]): number => {
   return dotProduct / (normA * normB)
 }
 
-export const FeatureExtractionProvider: React.FC<{ children: React.ReactNode }> = ({ children }) => {
+export const FeatureExtractionProvider: React.FC<{
+  children: React.ReactNode
+}> = ({ children }) => {
   const [examples, setExamples] = useState<EmbeddingExample[]>([])
-  const [selectedExample, setSelectedExample] = useState<EmbeddingExample | null>(null)
+  const [selectedExample, setSelectedExample] =
+    useState<EmbeddingExample | null>(null)
   const [similarities, setSimilarities] = useState<SimilarityResult[]>([])
   const [config, setConfig] = useState<FeatureExtractionConfig>({
     pooling: 'mean',
@@ -74,39 +83,55 @@ export const FeatureExtractionProvider: React.FC<{ children: React.ReactNode }>
       embedding: undefined,
       isLoading: false
     }
-    setExamples(prev => [...prev, newExample])
+    setExamples((prev) => [...prev, newExample])
   }, [])
 
-  const removeExample = useCallback((id: string) => {
-    setExamples(prev => prev.filter(example => example.id !== id))
-    if (selectedExample?.id === id) {
-      setSelectedExample(null)
-      setSimilarities([])
-    }
-  }, [selectedExample])
-
-  const updateExample = useCallback((id: string, updates: Partial<EmbeddingExample>) => {
-    setExamples(prev => prev.map(example =>
-      example.id === id ? { ...example, ...updates } : example
-    ))
-  }, [])
-
-  const calculateSimilarities = useCallback((targetExample: EmbeddingExample) => {
-    if (!targetExample.embedding) {
-      setSimilarities([])
-      return
-    }
+  const removeExample = useCallback(
+    (id: string) => {
+      setExamples((prev) => prev.filter((example) => example.id !== id))
+      if (selectedExample?.id === id) {
+        setSelectedExample(null)
+        setSimilarities([])
+      }
+    },
+    [selectedExample]
+  )
 
-    const newSimilarities: SimilarityResult[] = examples
-      .filter(example => example.id !== targetExample.id && example.embedding)
-      .map(example => ({
-        exampleId: example.id,
-        similarity: cosineSimilarity(targetExample.embedding!, example.embedding!)
-      }))
-      .sort((a, b) => b.similarity - a.similarity)
+  const updateExample = useCallback(
+    (id: string, updates: Partial<EmbeddingExample>) => {
+      setExamples((prev) =>
+        prev.map((example) =>
+          example.id === id ? { ...example, ...updates } : example
+        )
+      )
+    },
+    []
+  )
 
-    setSimilarities(newSimilarities)
-  }, [examples])
+  const calculateSimilarities = useCallback(
+    (targetExample: EmbeddingExample) => {
+      if (!targetExample.embedding) {
+        setSimilarities([])
+        return
+      }
+
+      const newSimilarities: SimilarityResult[] = examples
+        .filter(
+          (example) => example.id !== targetExample.id && example.embedding
+        )
+        .map((example) => ({
+          exampleId: example.id,
+          similarity: cosineSimilarity(
+            targetExample.embedding!,
+            example.embedding!
+          )
+        }))
+        .sort((a, b) => b.similarity - a.similarity)
+
+      setSimilarities(newSimilarities)
+    },
+    [examples]
+  )
 
   const clearExamples = useCallback(() => {
     setExamples([])

src/types.ts

@@ -75,7 +75,7 @@ export interface FeatureExtractionWorkerInput {
   model: string
   dtype: QuantizationType
   config: {
-    pooling: 'mean' | 'cls' | 'max'
+    pooling: 'mean' | 'cls'
     normalize: boolean
   }
 }