wip: enhance model selection and display across components

src/App.tsx

@@ -1,4 +1,4 @@
-import { useState } from 'react'
+import { useEffect, useState } from 'react'
 import PipelineSelector from './components/PipelineSelector'
 import ZeroShotClassification from './components/ZeroShotClassification'
 import TextClassification from './components/TextClassification'
@@ -6,11 +6,11 @@ import Header from './Header'
 import Footer from './Footer'
 import { useModel } from './contexts/ModelContext'
 import { Bot, Heart, Download, Cpu, DatabaseIcon } from 'lucide-react'
-import { getModelSize } from './lib/huggingface'
+import { getModelsByPipeline, getModelSize } from './lib/huggingface'
+import ModelSelector from './components/ModelSelector'
 
 function App() {
-  const [pipeline, setPipeline] = useState('zero-shot-classification')
-  const { progress, status, modelInfo } = useModel()
+  const { pipeline, setPipeline, progress, status, modelInfo, setModels } = useModel()
 
   const formatNumber = (num: number) => {
     if (num >= 1000000000) {
@@ -23,6 +23,14 @@ function App() {
     return num.toString()
   }
 
+    useEffect(() => {
+      const fetchModels = async () => {
+        const fetchedModels = await getModelsByPipeline(pipeline);
+        setModels(fetchedModels);
+      };
+      fetchModels();
+    }, [setModels, pipeline]);
+
   return (
     <div className="min-h-screen bg-gradient-to-br from-blue-50 to-indigo-100">
       <Header />
@@ -32,10 +40,6 @@ function App() {
         <div className="mb-8">
           <div className="bg-white rounded-lg shadow-sm border p-6">
             <div className="flex items-center justify-between mb-4">
-              <h2 className="text-lg font-semibold text-gray-900">
-                Choose a Pipeline
-              </h2>
-
               {/* Model Info Display */}
               {modelInfo.name && (
                 <div className="bg-gradient-to-r from-blue-50 to-indigo-50 px-4 py-3 rounded-lg border border-blue-200 space-y-2">
@@ -77,9 +81,10 @@ function App() {
                       <div className="flex items-center space-x-1">
                         <DatabaseIcon className="w-3 h-3 text-purple-500" />
                         <span>
-                          {`~${getModelSize(modelInfo.parameters, 'INT8').toFixed(
-                            1
-                          )}MB`}
+                          {`~${getModelSize(
+                            modelInfo.parameters,
+                            'INT8'
+                          ).toFixed(1)}MB`}
                         </span>
                       </div>
                     )}
@@ -88,7 +93,13 @@ function App() {
               )}
             </div>
 
-            <PipelineSelector pipeline={pipeline} setPipeline={setPipeline} />
+            <div className="flex flex-row items-center space-x-4">
+              <span className="text-lg font-semibold text-gray-900">
+                Choose a Pipeline
+              </span>
+              <PipelineSelector pipeline={pipeline} setPipeline={setPipeline} />
+            </div>
+            <ModelSelector />
 
             {/* Model Loading Progress */}
             {status === 'progress' && (

src/components/ModelSelector.tsx

@@ -1,25 +1,180 @@
-import React from 'react';
+import React, { useEffect, useState } from 'react'
+import { useModel } from '../contexts/ModelContext'
+import { getModelInfo } from '../lib/huggingface'
+import { Heart, Download, ChevronDown } from 'lucide-react'
 
-interface ModelSelectorProps {
-  model: string;
-  setModel: (model: string) => void;
-  models: string[];
-}
+const ModelSelector: React.FC = () => {
+  const { models, setModelInfo, modelInfo } = useModel()
+  const [isOpen, setIsOpen] = useState(false)
+  const [modelStats, setModelStats] = useState<
+    Record<string, { likes: number; downloads: number; createdAt: string }>
+  >({})
+
+  const formatNumber = (num: number) => {
+    if (num >= 1000000000) {
+      return (num / 1000000000).toFixed(1) + 'B'
+    } else if (num >= 1000000) {
+      return (num / 1000000).toFixed(1) + 'M'
+    } else if (num >= 1000) {
+      return (num / 1000).toFixed(1) + 'K'
+    }
+    return num.toString()
+  }
+
+  // Separate function to fetch only stats without updating selected model
+  const fetchModelStats = async (modelId: string) => {
+    try {
+      const modelInfoResponse = await getModelInfo(modelId)
+
+      setModelStats((prev) => ({
+        ...prev,
+        [modelId]: {
+          likes: modelInfoResponse.likes || 0,
+          downloads: modelInfoResponse.downloads || 0,
+          createdAt: modelInfoResponse.createdAt || ''
+        }
+      }))
+    } catch (error) {
+      console.error('Error fetching model stats:', error)
+    }
+  }
+
+  // Function to fetch full model info and set as selected
+  const fetchModelAndSetInfo = async (modelId: string) => {
+    try {
+      const modelInfoResponse = await getModelInfo(modelId)
+      let parameters = 0
+      if (modelInfoResponse.safetensors) {
+        const safetensors = modelInfoResponse.safetensors
+        parameters =
+          safetensors.parameters.F16 ||
+          safetensors.parameters.F32 ||
+          safetensors.parameters.total ||
+          0
+      }
+
+      // Transform ModelInfoResponse to ModelInfo
+      const modelInfo = {
+        id: modelId,
+        name: modelInfoResponse.id || modelId,
+        architecture: modelInfoResponse.config?.architectures?.[0] || 'Unknown',
+        parameters,
+        likes: modelInfoResponse.likes || 0,
+        downloads: modelInfoResponse.downloads || 0,
+        createdAt: modelInfoResponse.createdAt || ''
+      }
+
+      // Also update stats
+      setModelStats((prev) => ({
+        ...prev,
+        [modelId]: {
+          likes: modelInfoResponse.likes || 0,
+          downloads: modelInfoResponse.downloads || 0,
+          createdAt: modelInfoResponse.createdAt || ''
+        }
+      }))
+
+      console.log(modelInfoResponse)
+
+      setModelInfo(modelInfo)
+    } catch (error) {
+      console.error('Error fetching model info:', error)
+    }
+  }
+
+  // Fetch stats for all models when component mounts (without setting as selected)
+  useEffect(() => {
+    models.forEach((model) => {
+      if (!modelStats[model.id]) {
+        fetchModelStats(model.id)
+      }
+    })
+  }, [models])
+
+  // Only fetch full info when a model is actually selected
+  useEffect(() => {
+    if (!modelInfo.id) return
+    // Only fetch if we don't already have the full info
+    if (!modelStats[modelInfo.id]) {
+      fetchModelAndSetInfo(modelInfo.id)
+    }
+  }, [modelInfo.id])
+
+  const handleModelSelect = (modelId: string) => {
+    fetchModelAndSetInfo(modelId)
+    setIsOpen(false)
+  }
 
-const ModelSelector: React.FC<ModelSelectorProps> = ({
-  model,
-  setModel,
-  models
-}) => {
   return (
-    <select value={model} onChange={(e) => setModel(e.target.value)}>
-      {models.map((m) => (
-        <option key={m} value={m}>
-          {m}
-        </option>
-      ))}
-    </select>
-  );
-};
-
-export default ModelSelector;
+    <div className="relative">
+      {/* Custom Dropdown Button */}
+      <button
+        onClick={() => setIsOpen(!isOpen)}
+        className="w-full px-3 py-2 border border-gray-300 rounded-md focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-transparent bg-white text-left flex items-center justify-between"
+      >
+        <span className="truncate">{modelInfo.id || 'Select a model'}</span>
+        <ChevronDown
+          className={`w-4 h-4 transition-transform ${
+            isOpen ? 'rotate-180' : ''
+          }`}
+        />
+      </button>
+
+      {/* Custom Dropdown Options */}
+      {isOpen && (
+        <div className="absolute z-10 w-full mt-1 bg-white border border-gray-300 rounded-md shadow-lg max-h-60 overflow-auto">
+          {models.map((model) => (
+            <div
+              key={model.id}
+              onClick={() => handleModelSelect(model.id)}
+              className="px-3 py-2 hover:bg-gray-50 cursor-pointer border-b border-gray-100 last:border-b-0"
+            >
+              <div className="flex items-center justify-between">
+                <span className="text-sm font-medium truncate flex-1 mr-2">
+                  {model.id}
+                </span>
+
+                {/* Stats Display */}
+                {modelStats[model.id] &&
+                  (modelStats[model.id].likes > 0 ||
+                    modelStats[model.id].downloads > 0) && (
+                    <div className="flex items-center space-x-3 text-xs text-gray-500 flex-shrink-0">
+                      {modelStats[model.id].likes > 0 && (
+                        <div className="flex items-center space-x-1">
+                          <Heart className="w-3 h-3 text-red-500" />
+                          <span>
+                            {formatNumber(modelStats[model.id].likes)}
+                          </span>
+                        </div>
+                      )}
+
+                      {modelStats[model.id].downloads > 0 && (
+                        <div className="flex items-center space-x-1">
+                          <Download className="w-3 h-3 text-green-500" />
+                          <span>
+                            {formatNumber(modelStats[model.id].downloads)}
+                          </span>
+                        </div>
+                      )}
+                      {modelStats[model.id].createdAt !== '' && (
+                        <span className="text-xs text-gray-400">
+                          {modelStats[model.id].createdAt.split('T')[0]}
+                        </span>
+                      )}
+                    </div>
+                  )}
+              </div>
+            </div>
+          ))}
+        </div>
+      )}
+
+      {/* Click outside to close */}
+      {isOpen && (
+        <div className="fixed inset-0 z-0" onClick={() => setIsOpen(false)} />
+      )}
+    </div>
+  )
+}
+
+export default ModelSelector

src/components/PipelineSelector.tsx

@@ -3,6 +3,10 @@ import React from 'react';
 const pipelines = [
   'zero-shot-classification',
   'text-classification',
+  'text-generation',
+  'summarization',
+  'feature-extraction',
+  'sentiment-analysis',
   'image-classification',
   'question-answering',
   'translation'

src/components/TextClassification.tsx

@@ -3,7 +3,6 @@ import {
   ClassificationOutput,
   TextClassificationWorkerInput,
   WorkerMessage,
-  ModelInfo
 } from '../types';
 import { useModel } from '../contexts/ModelContext';
 import { getModelInfo } from '../lib/huggingface';
@@ -25,13 +24,14 @@ const PLACEHOLDER_TEXTS: string[] = [
 function TextClassification() {
   const [text, setText] = useState<string>(PLACEHOLDER_TEXTS.join('\n'))
   const [results, setResults] = useState<ClassificationOutput[]>([])
-  const { setProgress, status, setStatus, modelInfo, setModelInfo} = useModel()
+  const { setProgress, status, setStatus, modelInfo, setModelInfo, models, setModels} = useModel()
+
+
   useEffect(() => {
-    const modelName = 'distilbert/distilbert-base-uncased-finetuned-sst-2-english'
+    if (!modelInfo.id) return;
     const fetchModelInfo = async () => {
       try {
-        const modelInfoResponse = await getModelInfo(modelName)
-        console.log(modelInfoResponse)
+        const modelInfoResponse = await getModelInfo(modelInfo.id)
         let parameters = 0
         if (modelInfoResponse.safetensors) {
           const safetensors = modelInfoResponse.safetensors
@@ -42,8 +42,8 @@ function TextClassification() {
               0)
         }
         setModelInfo({
-          name: modelName,
-          architecture: modelInfoResponse.config.architectures[0],
+          ...modelInfo,
+          architecture: modelInfoResponse.config?.architectures[0] ?? '',
           parameters,
           likes: modelInfoResponse.likes,
           downloads: modelInfoResponse.downloads
@@ -54,7 +54,7 @@ function TextClassification() {
     }
 
     fetchModelInfo()
-  }, [setModelInfo])
+  }, [modelInfo.id, setModelInfo])
 
   // Create a reference to the worker object.
   const worker = useRef<Worker | null>(null)
@@ -110,9 +110,9 @@ function TextClassification() {
   const classify = useCallback(() => {
     setStatus('processing')
     setResults([]) // Clear previous results
-    const message: TextClassificationWorkerInput = { text, model: modelInfo.name }
+    const message: TextClassificationWorkerInput = { text, model: modelInfo.id }
     worker.current?.postMessage(message)
-  }, [text, modelInfo.name])
+  }, [text, modelInfo.id])
 
   const busy: boolean = status !== 'idle'
 

src/components/ZeroShotClassification.tsx

@@ -68,11 +68,13 @@ function ZeroShotClassification() {
             0
         }
         setModelInfo({
+          id: modelInfoResponse.id,
           name: modelName,
-          architecture: modelInfoResponse.config.architectures[0],
+          architecture: modelInfoResponse.config?.architectures[0] ?? '',
           parameters,
           likes: modelInfoResponse.likes,
-          downloads: modelInfoResponse.downloads
+          downloads: modelInfoResponse.downloads,
+          createdAt: modelInfoResponse.createdAt,
         })
       } catch (error) {
         console.error('Error fetching model info:', error)

src/contexts/ModelContext.tsx

@@ -1,5 +1,5 @@
 import React, { createContext, useContext, useEffect, useState } from 'react'
-import { ModelInfo } from '../types'
+import { ModelInfo, ModelInfoResponse } from '../types'
 
 interface ModelContextType {
   progress: number
@@ -8,6 +8,10 @@ interface ModelContextType {
   setStatus: (status: string) => void
   modelInfo: ModelInfo
   setModelInfo: (model: ModelInfo) => void
+  pipeline: string
+  setPipeline: (pipeline: string) => void
+  models: ModelInfoResponse[]
+  setModels: (models: ModelInfoResponse[]) => void
 }
 
 const ModelContext = createContext<ModelContextType | undefined>(undefined)
@@ -16,6 +20,8 @@ export function ModelProvider({ children }: { children: React.ReactNode }) {
   const [progress, setProgress] = useState<number>(0)
   const [status, setStatus] = useState<string>('idle')
   const [modelInfo, setModelInfo] = useState<ModelInfo>({} as ModelInfo)
+  const [models, setModels] = useState<ModelInfoResponse[]>([] as ModelInfoResponse[])
+  const [pipeline, setPipeline] = useState<string>('zero-shot-classification')
 
   // set progress to 0 when model is changed
   useEffect(() => {
@@ -30,7 +36,11 @@ export function ModelProvider({ children }: { children: React.ReactNode }) {
         status,
         setStatus,
         modelInfo,
-        setModelInfo
+        setModelInfo,
+        models,
+        setModels,
+        pipeline,
+        setPipeline,
       }}
     >
       {children}

src/lib/huggingface.ts

@@ -1,27 +1,5 @@
-interface ModelInfoResponse {
-  id: string
-  config: {
-    architectures: string[]
-    model_type: string
-  }
-  lastModified: string
-  pipeline_tag: string
-  tags: string[]
-  transformersInfo: {
-    pipeline_tag: string
-    auto_model: string
-    processor: string
-  }
-  safetensors?: {
-    parameters: {
-      F16?: number
-      F32?: number
-      total?: number
-    }
-  }
-  likes: number
-  downloads: number
-}
+import { Mode } from "fs"
+import { ModelInfoResponse } from "../types"
 
 const getModelInfo = async (modelName: string): Promise<ModelInfoResponse> => {
   const token = process.env.REACT_APP_HUGGINGFACE_TOKEN
@@ -48,6 +26,37 @@ const getModelInfo = async (modelName: string): Promise<ModelInfoResponse> => {
   return response.json()
 }
 
+const getModelsByPipeline = async (
+  pipeline_tag: string
+): Promise<ModelInfoResponse[]> => {
+  const token = process.env.REACT_APP_HUGGINGFACE_TOKEN
+
+  if (!token) {
+    throw new Error(
+      'Hugging Face token not found. Please set REACT_APP_HUGGINGFACE_TOKEN in your .env file'
+    )
+  }
+
+  const response = await fetch(
+    `https://huggingface.co/api/models?filter=${pipeline_tag}&filter=transformers.js&sort=downloads`,
+    {
+      method: 'GET',
+      headers: {
+        Authorization: `Bearer ${token}`
+      }
+    }
+  )
+
+  if (!response.ok) {
+    throw new Error(`Failed to fetch models for pipeline: ${response.statusText}`)
+  }
+  const models = await response.json()
+  if (pipeline_tag === 'text-classification') {
+    return models.filter((model: ModelInfoResponse) => !model.tags.includes('reranker') && !model.id.includes('reranker')).slice(0, 10)
+  }
+  return models.slice(0, 10)
+}
+
 // Define the possible quantization types for clarity and type safety
 type QuantizationType = 'FP32' | 'FP16' | 'INT8' | 'Q4'
 function getModelSize(
@@ -81,5 +90,5 @@ function getModelSize(
 }
 
 
-export { getModelInfo, getModelSize }
+export { getModelInfo, getModelSize, getModelsByPipeline }
 

src/types.ts

@@ -28,9 +28,38 @@ export interface TextClassificationWorkerInput {
 export type AppStatus = 'idle' | 'loading' | 'processing'
 
 export interface ModelInfo {
+  id: string
   name: string
   architecture: string
   parameters: number
   likes: number
   downloads: number
+  createdAt: string
+}
+
+
+export interface ModelInfoResponse {
+  id: string
+  createdAt: string
+  config?: {
+    architectures: string[]
+    model_type: string
+  }
+  lastModified: string
+  pipeline_tag: string
+  tags: string[]
+  transformersInfo: {
+    pipeline_tag: string
+    auto_model: string
+    processor: string
+  }
+  safetensors?: {
+    parameters: {
+      F16?: number
+      F32?: number
+      total?: number
+    }
+  }
+  likes: number
+  downloads: number
 }