al

app.py

@@ -2,7 +2,6 @@ import gradio as gr
 from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer
 import pandas as pd
 from sklearn.metrics import accuracy_score, precision_recall_fscore_support
-import torch
 
 # Define the available models and tasks
 TASKS = ["sentiment-analysis", "ner", "text-classification"]
@@ -14,23 +13,20 @@ MODELS = {
     # Add other models here
 }
 
-def load_pipeline(task, model_name):
-    model = MODELS.get(model_name)
-    if not model:
-        raise ValueError(f"Model {model_name} is not available.")
+def load_pipeline(task, model):
     return pipeline(task, model=model)
 
-def predict(task, model_name, text):
-    selected_pipeline = load_pipeline(task, model_name)
+def predict(task, model, text):
+    selected_pipeline = load_pipeline(task, model)
     results = selected_pipeline(text)
     return results
 
-def benchmark(task, model_name, file):
+def benchmark(task, model, file):
     data = pd.read_csv(file.name)
     texts = data['text'].tolist()
     true_labels = data['label'].tolist()
     
-    selected_pipeline = load_pipeline(task, model_name)
+    selected_pipeline = load_pipeline(task, model)
     predictions = [selected_pipeline(text)[0]['label'] for text in texts]
     
     accuracy = accuracy_score(true_labels, predictions)
@@ -62,5 +58,5 @@ with gr.Blocks() as demo:
             benchmark_button = gr.Button("Benchmark")
         benchmark_output = gr.JSON(label="Benchmark Output")
         benchmark_button.click(benchmark, inputs=[task_input, model_input, file_input], outputs=benchmark_output)
-    
-    demo.launch()
+
+demo.launch()