feat: enhance model information display and improve context management

- Added model information display in the pipeline selection area, including model name, likes, downloads, and parameters.
- Updated the context to manage model information instead of just the model name.
- Integrated model information fetching from Hugging Face API in both TextClassification and ZeroShotClassification components.
- Modified worker scripts to accept dynamic model names for text classification and zero-shot classification tasks.
- Improved formatting for large numbers in the UI.
- Added environment variable support for Hugging Face API token.
- Updated types to include ModelInfo interface for better type safety.

.env.example

@@ -0,0 +1 @@
+REACT_APP_HUGGINGFACE_TOKEN="hf_..."

.gitignore

@@ -21,3 +21,5 @@
 npm-debug.log*
 yarn-debug.log*
 yarn-error.log*
+
+.env

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "CodeGPT.apiKey": "CodeGPT Plus Beta"
+}

package.json

@@ -13,6 +13,9 @@
     "@types/react": "^19.1.8",
     "@types/react-dom": "^19.1.6",
     "build": "^0.1.4",
+    "dotenv": "^17.0.1",
+    "lucide-react": "^0.525.0",
+    "path-browserify": "^1.0.1",
     "react": "^19.1.0",
     "react-dom": "^19.1.0",
     "react-scripts": "5.0.1",

src/App.tsx

@@ -5,10 +5,22 @@ import TextClassification from './components/TextClassification';
 import Header from './Header';
 import Footer from './Footer';
 import { useModel } from './contexts/ModelContext';
+import { Bot, Heart, Download, Cpu } from 'lucide-react';
 
 function App() {
   const [pipeline, setPipeline] = useState('zero-shot-classification');
-  const { progress, status, model } = useModel();
+  const { progress, status, modelInfo } = useModel();
+
+  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();
+  };
 
   return (
     <div className="min-h-screen bg-gradient-to-br from-blue-50 to-indigo-100">
@@ -18,9 +30,47 @@ function App() {
         {/* Pipeline Selection */}
         <div className="mb-8">
           <div className="bg-white rounded-lg shadow-sm border p-6">
-            <h2 className="text-lg font-semibold text-gray-900 mb-4">
-              Choose a Pipeline
-            </h2>
+            <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="flex items-center space-x-4 bg-gradient-to-r from-blue-50 to-indigo-50 px-4 py-2 rounded-lg border border-blue-200">
+                  <div className="flex items-center space-x-2">
+                    <Bot className="w-4 h-4 text-blue-600" />
+                    <span className="text-sm font-medium text-gray-700 truncate max-w-80" title={modelInfo.name}>
+                      {modelInfo.name.split('/').pop()}
+                    </span>
+                  </div>
+                  
+                  <div className="flex items-center space-x-4 text-xs text-gray-600">
+                    {modelInfo.likes > 0 && (
+                      <div className="flex items-center space-x-1">
+                        <Heart className="w-3 h-3 text-red-500" />
+                        <span>{formatNumber(modelInfo.likes)}</span>
+                      </div>
+                    )}
+                    
+                    {modelInfo.downloads > 0 && (
+                      <div className="flex items-center space-x-1">
+                        <Download className="w-3 h-3 text-green-500" />
+                        <span>{formatNumber(modelInfo.downloads)}</span>
+                      </div>
+                    )}
+                    
+                    {modelInfo.parameters > 0 && (
+                      <div className="flex items-center space-x-1">
+                        <Cpu className="w-3 h-3 text-purple-500" />
+                        <span>{formatNumber(modelInfo.parameters)}</span>
+                      </div>
+                    )}
+                  </div>
+                </div>
+              )}
+            </div>
+            
             <PipelineSelector pipeline={pipeline} setPipeline={setPipeline} />
 
             {/* Model Loading Progress */}
@@ -85,11 +135,6 @@ function App() {
                     {pipeline === 'zero-shot-classification'
                       ? 'Zero-Shot Classification'
                       : 'Text-Classification'}
-                    {model && (
-                      <span className="ml-2 text-xs text-gray-500 font-normal">
-                        ({model})
-                      </span>
-                    )}
                   </h3>
                   <p className="text-sm text-gray-600 mt-1">
                     {pipeline === 'zero-shot-classification'

src/components/TextClassification.tsx

@@ -2,9 +2,11 @@ import { useState, useRef, useEffect, useCallback } from 'react';
 import {
   ClassificationOutput,
   TextClassificationWorkerInput,
-  WorkerMessage
+  WorkerMessage,
+  ModelInfo
 } from '../types';
 import { useModel } from '../contexts/ModelContext';
+import { getModelInfo } from '../lib/huggingface';
 
 
 const PLACEHOLDER_TEXTS: string[] = [
@@ -21,13 +23,41 @@ const PLACEHOLDER_TEXTS: string[] = [
 ].sort(() => Math.random() - 0.5);
 
 function TextClassification() {
-  const [text, setText] = useState<string>(PLACEHOLDER_TEXTS.join('\n'));
-  const [results, setResults] = useState<ClassificationOutput[]>([]);
-  const { setProgress, status, setStatus, setModel } = useModel();
-  setModel('Xenova/bert-base-multilingual-uncased-sentiment')
+  const [text, setText] = useState<string>(PLACEHOLDER_TEXTS.join('\n'))
+  const [results, setResults] = useState<ClassificationOutput[]>([])
+  const { setProgress, status, setStatus, modelInfo, setModelInfo} = useModel()
+  useEffect(() => {
+    const modelName = 'Xenova/distilbert-base-uncased-finetuned-sst-2-english'
+    const fetchModelInfo = async () => {
+      try {
+        const modelInfoResponse = await getModelInfo(modelName)
+        console.log(modelInfoResponse)
+        let parameters = 0
+        if (modelInfoResponse.safetensors) {
+          const safetensors = modelInfoResponse.safetensors
+          parameters =
+            (safetensors.parameters.F16 ||
+              safetensors.parameters.F32 ||
+              safetensors.parameters.total ||
+              0)
+        }
+        setModelInfo({
+          name: modelName,
+          architecture: modelInfoResponse.config.architectures[0],
+          parameters,
+          likes: modelInfoResponse.likes,
+          downloads: modelInfoResponse.downloads
+        })
+      } catch (error) {
+        console.error('Error fetching model info:', error)
+      }
+    }
+
+    fetchModelInfo()
+  }, [setModelInfo])
 
   // Create a reference to the worker object.
-  const worker = useRef<Worker | null>(null);
+  const worker = useRef<Worker | null>(null)
 
   // We use the `useEffect` hook to setup the worker as soon as the component is mounted.
   useEffect(() => {
@@ -38,54 +68,57 @@ function TextClassification() {
         {
           type: 'module'
         }
-      );
+      )
     }
 
     // Create a callback function for messages from the worker thread.
     const onMessageReceived = (e: MessageEvent<WorkerMessage>) => {
-      const status = e.data.status;
+      const status = e.data.status
       if (status === 'initiate') {
-        setStatus('loading');
+        setStatus('loading')
       } else if (status === 'ready') {
-        setStatus('ready');
+        setStatus('ready')
       } else if (status === 'progress') {
-        setStatus('progress');
+        setStatus('progress')
         if (
           e.data.output.progress &&
           (e.data.output.file as string).startsWith('onnx')
         )
-          setProgress(e.data.output.progress);
+          setProgress(e.data.output.progress)
       } else if (status === 'output') {
-        setStatus('output');
-        const result = e.data.output!;
-        setResults((prevResults) => [...prevResults, result]);
-        console.log(result);
+        setStatus('output')
+        const result = e.data.output!
+        setResults((prevResults) => [...prevResults, result])
+        console.log(result)
       } else if (status === 'complete') {
-        setStatus('idle');
-        setProgress(100);
+        setStatus('idle')
+        setProgress(100)
+      } else if (status === 'error') {
+        setStatus('error')
+        console.error(e.data.output)
       }
-    };
+    }
 
     // Attach the callback function as an event listener.
-    worker.current.addEventListener('message', onMessageReceived);
+    worker.current.addEventListener('message', onMessageReceived)
 
     // Define a cleanup function for when the component is unmounted.
     return () =>
-      worker.current?.removeEventListener('message', onMessageReceived);
-  }, []);
+      worker.current?.removeEventListener('message', onMessageReceived)
+  }, [])
 
   const classify = useCallback(() => {
-    setStatus('processing');
-    setResults([]); // Clear previous results
-    const message: TextClassificationWorkerInput = { text };
-    worker.current?.postMessage(message);
-  }, [text]);
+    setStatus('processing')
+    setResults([]) // Clear previous results
+    const message: TextClassificationWorkerInput = { text, model: modelInfo.name }
+    worker.current?.postMessage(message)
+  }, [text, modelInfo.name])
 
-  const busy: boolean = status !== 'idle';
+  const busy: boolean = status !== 'idle'
 
   const handleClear = (): void => {
-    setResults([]);
-  };
+    setResults([])
+  }
 
   return (
     <div className="flex flex-col h-[40vh] max-h-[80vh] w-full p-4">
@@ -157,7 +190,7 @@ function TextClassification() {
         </div>
       </div>
     </div>
-  );
+  )
 }
 
 export default TextClassification;

src/components/ZeroShotClassification.tsx

@@ -1,7 +1,13 @@
 // src/App.tsx
-import { useState, useRef, useEffect, useCallback } from 'react';
-import { Section, WorkerMessage, ZeroShotWorkerInput } from '../types';
-import { useModel } from '../contexts/ModelContext';
+import { useState, useRef, useEffect, useCallback } from 'react'
+import {
+  Section,
+  WorkerMessage,
+  ZeroShotWorkerInput,
+  ModelInfo
+} from '../types'
+import { useModel } from '../contexts/ModelContext'
+import { getModelInfo } from '../lib/huggingface'
 
 const PLACEHOLDER_REVIEWS: string[] = [
   // battery/charging problems
@@ -28,7 +34,7 @@ const PLACEHOLDER_REVIEWS: string[] = [
   "I'm not sure what to make of this phone. It's not bad, but it's not great either. I'm on the fence about it.",
   "I hate the color of this phone. It's so ugly!",
   "This phone sucks! I'm returning it."
-].sort(() => Math.random() - 0.5);
+].sort(() => Math.random() - 0.5)
 
 const PLACEHOLDER_SECTIONS: string[] = [
   'Battery and charging problems',
@@ -36,20 +42,48 @@ const PLACEHOLDER_SECTIONS: string[] = [
   'Poor build quality',
   'Software issues',
   'Other'
-];
+]
 
 function ZeroShotClassification() {
-  const [text, setText] = useState<string>(PLACEHOLDER_REVIEWS.join('\n'));
+  const [text, setText] = useState<string>(PLACEHOLDER_REVIEWS.join('\n'))
 
   const [sections, setSections] = useState<Section[]>(
     PLACEHOLDER_SECTIONS.map((title) => ({ title, items: [] }))
-  );
+  )
 
-  const { setProgress, status, setStatus, setModel } = useModel();
-  setModel('MoritzLaurer/deberta-v3-xsmall-zeroshot-v1.1-all-33')
+  const { setProgress, status, setStatus, modelInfo, setModelInfo } = useModel()
+  useEffect(() => {
+    const modelName = 'MoritzLaurer/deberta-v3-xsmall-zeroshot-v1.1-all-33'
+    const fetchModelInfo = async () => {
+      try {
+        const modelInfoResponse = await getModelInfo(modelName)
+        console.log(modelInfoResponse)
+        let parameters = 0
+        if (modelInfoResponse.safetensors) {
+          const safetensors = modelInfoResponse.safetensors
+          parameters =
+            safetensors.parameters.F16 ||
+            safetensors.parameters.F32 ||
+            safetensors.parameters.total ||
+            0
+        }
+        setModelInfo({
+          name: modelName,
+          architecture: modelInfoResponse.config.architectures[0],
+          parameters,
+          likes: modelInfoResponse.likes,
+          downloads: modelInfoResponse.downloads
+        })
+      } catch (error) {
+        console.error('Error fetching model info:', error)
+      }
+    }
+
+    fetchModelInfo()
+  }, [setModelInfo])
 
   // Create a reference to the worker object.
-  const worker = useRef<Worker | null>(null);
+  const worker = useRef<Worker | null>(null)
 
   // We use the `useEffect` hook to setup the worker as soon as the `App` component is mounted.
   useEffect(() => {
@@ -60,86 +94,90 @@ function ZeroShotClassification() {
         {
           type: 'module'
         }
-      );
+      )
     }
 
     // Create a callback function for messages from the worker thread.
     const onMessageReceived = (e: MessageEvent<WorkerMessage>) => {
-      const status = e.data.status;
+      const status = e.data.status
       if (status === 'initiate') {
-        setStatus('loading');
+        setStatus('loading')
       } else if (status === 'ready') {
-        setStatus('ready');
+        setStatus('ready')
       } else if (status === 'progress') {
-        setStatus('progress');
+        setStatus('progress')
         if (
           e.data.output.progress &&
           (e.data.output.file as string).startsWith('onnx')
         )
-          setProgress(e.data.output.progress);
+          setProgress(e.data.output.progress)
       } else if (status === 'output') {
-        setStatus('output');
-        const { sequence, labels, scores } = e.data.output!;
+        setStatus('output')
+        const { sequence, labels, scores } = e.data.output!
 
         // Threshold for classification
-        const label = scores[0] > 0.5 ? labels[0] : 'Other';
+        const label = scores[0] > 0.5 ? labels[0] : 'Other'
 
         const sectionID =
-          sections.map((x) => x.title).indexOf(label) ?? sections.length - 1;
+          sections.map((x) => x.title).indexOf(label) ?? sections.length - 1
         setSections((sections) => {
-          const newSections = [...sections];
+          const newSections = [...sections]
           newSections[sectionID] = {
             ...newSections[sectionID],
             items: [...newSections[sectionID].items, sequence]
-          };
-          return newSections;
-        });
+          }
+          return newSections
+        })
       } else if (status === 'complete') {
-        setStatus('idle');
-        setProgress(100);
+        setStatus('idle')
+        setProgress(100)
+      } else if (status === 'error') {
+        setStatus('error')
+        console.error(e.data.output)
       }
-    };
+    }
 
     // Attach the callback function as an event listener.
-    worker.current.addEventListener('message', onMessageReceived);
+    worker.current.addEventListener('message', onMessageReceived)
 
     // Define a cleanup function for when the component is unmounted.
     return () =>
-      worker.current?.removeEventListener('message', onMessageReceived);
-  }, [sections]);
+      worker.current?.removeEventListener('message', onMessageReceived)
+  }, [sections])
 
   const classify = useCallback(() => {
-    setStatus('processing');
+    setStatus('processing')
     const message: ZeroShotWorkerInput = {
       text,
       labels: sections
         .slice(0, sections.length - 1)
-        .map((section) => section.title)
-    };
-    worker.current?.postMessage(message);
-  }, [text, sections]);
+        .map((section) => section.title),
+      model: modelInfo.name
+    }
+    worker.current?.postMessage(message)
+  }, [text, sections, modelInfo.name])
 
-  const busy: boolean = status !== 'idle';
+  const busy: boolean = status !== 'idle'
 
   const handleAddCategory = (): void => {
     setSections((sections) => {
-      const newSections = [...sections];
+      const newSections = [...sections]
       // add at position 2 from the end
       newSections.splice(newSections.length - 1, 0, {
         title: 'New Category',
         items: []
-      });
-      return newSections;
-    });
-  };
+      })
+      return newSections
+    })
+  }
 
   const handleRemoveCategory = (): void => {
     setSections((sections) => {
-      const newSections = [...sections];
-      newSections.splice(newSections.length - 2, 1); // Remove second last element
-      return newSections;
-    });
-  };
+      const newSections = [...sections]
+      newSections.splice(newSections.length - 2, 1) // Remove second last element
+      return newSections
+    })
+  }
 
   const handleClear = (): void => {
     setSections((sections) =>
@@ -147,16 +185,16 @@ function ZeroShotClassification() {
         ...section,
         items: []
       }))
-    );
-  };
+    )
+  }
 
   const handleSectionTitleChange = (index: number, newTitle: string): void => {
     setSections((sections) => {
-      const newSections = [...sections];
-      newSections[index].title = newTitle;
-      return newSections;
-    });
-  };
+      const newSections = [...sections]
+      newSections[index].title = newTitle
+      return newSections
+    })
+  }
 
   return (
     <div className="flex flex-col h-screen w-full p-1">
@@ -174,8 +212,8 @@ function ZeroShotClassification() {
           {!busy
             ? 'Categorize'
             : status === 'loading'
-              ? 'Model loading...'
-              : 'Processing'}
+            ? 'Model loading...'
+            : 'Processing'}
         </button>
         <div className="flex gap-1">
           <button
@@ -223,7 +261,7 @@ function ZeroShotClassification() {
         ))}
       </div>
     </div>
-  );
+  )
 }
 
-export default ZeroShotClassification;
+export default ZeroShotClassification

src/contexts/ModelContext.tsx

@@ -1,39 +1,47 @@
-import React, { createContext, useContext, useEffect, useState } from 'react';
+import React, { createContext, useContext, useEffect, useState } from 'react'
+import { ModelInfo } from '../types'
 
 interface ModelContextType {
-  progress: number;
-  status: string;
-  setProgress: (progress: number) => void;
-  setStatus: (status: string) => void;
-  model: string;
-  setModel: (model: string) => void;
+  progress: number
+  status: string
+  setProgress: (progress: number) => void
+  setStatus: (status: string) => void
+  modelInfo: ModelInfo
+  setModelInfo: (model: ModelInfo) => void
 }
 
-const ModelContext = createContext<ModelContextType | undefined>(undefined);
+const ModelContext = createContext<ModelContextType | undefined>(undefined)
 
 export function ModelProvider({ children }: { children: React.ReactNode }) {
-  const [progress, setProgress] = useState<number>(0);
-  const [status, setStatus] = useState<string>('idle');
-  const [model, setModel] = useState<string>('');
+  const [progress, setProgress] = useState<number>(0)
+  const [status, setStatus] = useState<string>('idle')
+  const [modelInfo, setModelInfo] = useState<ModelInfo>({} as ModelInfo)
 
   // set progress to 0 when model is changed
   useEffect(() => {
-    setProgress(0);
-  }, [model]);
+    setProgress(0)
+  }, [modelInfo.name])
 
   return (
     <ModelContext.Provider
-      value={{ progress, setProgress, status, setStatus, model, setModel }}
+      value={{
+        progress,
+        setProgress,
+        status,
+        setStatus,
+        modelInfo,
+        setModelInfo
+      }}
     >
       {children}
     </ModelContext.Provider>
-  );
+  )
 }
 
 export function useModel() {
-  const context = useContext(ModelContext);
+  const context = useContext(ModelContext)
   if (context === undefined) {
-    throw new Error('useModel must be used within a ModelProvider');
+    throw new Error('useModel must be used within a ModelProvider')
   }
-  return context;
+  return context
 }

src/lib/huggingface.ts

@@ -0,0 +1,51 @@
+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
+}
+
+const getModelInfo = async (modelName: 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/${modelName}`,
+    {
+      method: 'GET',
+      headers: {
+        Authorization: `Bearer ${token}`
+      }
+    }
+  )
+
+  if (!response.ok) {
+    throw new Error(`Failed to fetch model info: ${response.statusText}`)
+  }
+  return response.json()
+}
+
+export { getModelInfo }

src/types.ts

@@ -1,26 +1,36 @@
 export interface Section {
-  title: string;
-  items: string[];
+  title: string
+  items: string[]
 }
 
 export interface ClassificationOutput {
-  sequence: string;
-  labels: string[];
-  scores: number[];
+  sequence: string
+  labels: string[]
+  scores: number[]
 }
 
 export interface WorkerMessage {
-  status: 'initiate' | 'ready' | 'output' | 'complete' | 'progress';
-  output?: any;
+  status: 'initiate' | 'ready' | 'output' | 'complete' | 'progress'
+  output?: any
 }
 
 export interface ZeroShotWorkerInput {
-  text: string;
-  labels: string[];
+  text: string
+  labels: string[]
+  model: string
 }
 
 export interface TextClassificationWorkerInput {
-  text: string;
+  text: string
+  model: string
 }
 
-export type AppStatus = 'idle' | 'loading' | 'processing';
+export type AppStatus = 'idle' | 'loading' | 'processing'
+
+export interface ModelInfo {
+  name: string
+  architecture: string
+  parameters: number
+  likes: number
+  downloads: number
+}

src/workers/text-classification.js

@@ -3,11 +3,10 @@ import { pipeline } from '@huggingface/transformers';
 
 class MyTextClassificationPipeline {
   static task = 'text-classification';
-  static model = 'Xenova/bert-base-multilingual-uncased-sentiment';
   static instance = null;
 
-  static async getInstance(progress_callback = null) {
-    this.instance ??= pipeline(this.task, this.model, {
+  static async getInstance(model, progress_callback = null) {
+    this.instance ??= pipeline(this.task, model, {
       progress_callback
     });
 
@@ -17,15 +16,24 @@ class MyTextClassificationPipeline {
 
 // Listen for messages from the main thread
 self.addEventListener('message', async (event) => {
+  const { text, model } = event.data;
+  if (!model) {
+    self.postMessage({
+      status: 'error',
+      output: 'No model provided'
+    });
+    return;
+  }
+
   // Retrieve the pipeline. When called for the first time,
   // this will load the pipeline and save it for future use.
-  const classifier = await MyTextClassificationPipeline.getInstance((x) => {
+  const classifier = await MyTextClassificationPipeline.getInstance(model, (x) => {
     // We also add a progress callback to the pipeline so that we can
     // track model loading.
     self.postMessage({ status: 'progress', output: x });
   });
 
-  const { text } = event.data;
+
 
   const split = text.split('\n');
   for (const line of split) {

src/workers/zero-shot.js

@@ -3,11 +3,10 @@ import { pipeline } from '@huggingface/transformers';
 
 class MyZeroShotClassificationPipeline {
   static task = 'zero-shot-classification';
-  static model = 'MoritzLaurer/deberta-v3-xsmall-zeroshot-v1.1-all-33';
   static instance = null;
 
-  static async getInstance(progress_callback = null) {
-    this.instance ??= pipeline(this.task, this.model, {
+  static async getInstance(model, progress_callback = null) {
+    this.instance ??= pipeline(this.task, model, {
       progress_callback
     });
 
@@ -17,16 +16,22 @@ class MyZeroShotClassificationPipeline {
 
 // Listen for messages from the main thread
 self.addEventListener('message', async (event) => {
+  const { text, labels, model } = event.data;
+  if (!model) {
+    self.postMessage({
+      status: 'error',
+      output: 'No model provided'
+    });
+    return;
+  }
+
   // Retrieve the pipeline. When called for the first time,
   // this will load the pipeline and save it for future use.
-  const classifier = await MyZeroShotClassificationPipeline.getInstance((x) => {
+  const classifier = await MyZeroShotClassificationPipeline.getInstance(model, (x) => {
     // We also add a progress callback to the pipeline so that we can
     // track model loading.
     self.postMessage({ status: 'progress', output: x });
   });
-
-  const { text, labels } = event.data;
-
   const split = text.split('\n');
   for (const line of split) {
     const output = await classifier(line, labels, {