Update app.py for using pre-trained sentiment analysis pipeline

app.py

@@ -1,53 +1,13 @@
 import gradio as gr
-from datasets import load_dataset
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
-import torch
+from transformers import pipeline
 
-# Check if GPU is available
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Load the IMDb dataset
-dataset = load_dataset('imdb')
-
-# Initialize the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')
-model = AutoModelForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=2)
-model.to(device)
-
-# Tokenize the dataset
-def tokenize_function(examples):
-    return tokenizer(examples['text'], padding="max_length", truncation=True)
-
-tokenized_datasets = dataset.map(tokenize_function, batched=True)
-
-# Set up training arguments
-training_args = TrainingArguments(
-    output_dir="./results",
-    evaluation_strategy="epoch",
-    learning_rate=2e-5,
-    per_device_train_batch_size=16,
-    per_device_eval_batch_size=16,
-    num_train_epochs=1,  # Start with fewer epochs for quicker runs
-    weight_decay=0.01,
-)
-
-# Initialize the Trainer
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    train_dataset=tokenized_datasets["train"].shuffle(seed=42).select(range(1000)),  # Use a subset for quicker runs
-    eval_dataset=tokenized_datasets["test"].shuffle(seed=42).select(range(1000)),
-)
-
-# Train the model
-trainer.train()
+# Load the pre-trained sentiment-analysis pipeline
+classifier = pipeline('sentiment-analysis')
 
 # Function to classify sentiment
 def classify_text(text):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
-    outputs = model(**inputs)
-    prediction = torch.argmax(outputs.logits, dim=-1).item()
-    return "Positive" if prediction == 1 else "Negative"
+    result = classifier(text)[0]
+    return f"{result['label']} with score {result['score']}"
 
 # Set up the Gradio interface
 iface = gr.Interface(fn=classify_text, inputs="text", outputs="text")