biu

app.py

@@ -1,16 +1,14 @@
 import gradio as gr
-from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments
+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"]
+TASKS = ["sentiment-analysis", "text-classification"]
 MODELS = {
-    "DistilBERT": "distilbert-base-uncased-finetuned-sst-2-english",
+    "DistilBERT": "distilbert-base-uncased",
     "BERT": "bert-base-uncased",
-    "RoBERTa": "roberta-base",
-    "LLaMA": "decapoda-research/llama-7b-hf"
+    "RoBERTa": "roberta-base"
     # Add other models here
 }
 
@@ -40,42 +38,6 @@ def benchmark(task, model, file):
         "F1 Score": f1
     }
 
-def train_model(task, model, file):
-    data = pd.read_csv(file.name)
-    train_texts = data['text'].tolist()
-    train_labels = data['label'].tolist()
-    
-    model_name = MODELS[model]
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    
-    encodings = tokenizer(train_texts, truncation=True, padding=True)
-    dataset = torch.utils.data.TensorDataset(
-        torch.tensor(encodings['input_ids']),
-        torch.tensor(encodings['attention_mask']),
-        torch.tensor(train_labels)
-    )
-    
-    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=len(set(train_labels)))
-    training_args = TrainingArguments(
-        output_dir='./results',
-        num_train_epochs=3,
-        per_device_train_batch_size=8,
-        per_device_eval_batch_size=8,
-        warmup_steps=500,
-        weight_decay=0.01,
-        logging_dir='./logs'
-    )
-    
-    trainer = Trainer(
-        model=model,
-        args=training_args,
-        train_dataset=dataset,
-        eval_dataset=dataset
-    )
-    
-    trainer.train()
-    return "Training Completed"
-
 # Define the Gradio interface
 with gr.Blocks() as demo:
     with gr.Row():
@@ -95,12 +57,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)
-    
-    with gr.Tab("Train"):
-        with gr.Row():
-            train_file_input = gr.File(label="Upload CSV for Training")
-            train_button = gr.Button("Train Model")
-        train_output = gr.JSON(label="Training Status")
-        train_button.click(train_model, inputs=[task_input, model_input, train_file_input], outputs=train_output)
 
 demo.launch()