Update app.py

app.py

@@ -1,61 +1,13 @@
-from datasets import load_dataset
-from transformers import DistilBertTokenizer, DistilBertForSequenceClassification, Trainer, TrainingArguments
+from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
 import torch
-import numpy as np
-from sklearn.metrics import accuracy_score, f1_score
+import gradio as gr
 
-# Step 2: Load dataset
-dataset = load_dataset("amazon_polarity")
-train_dataset = dataset["train"].shuffle(seed=42).select(range(10000))
-test_dataset = dataset["test"].shuffle(seed=42).select(range(2000))
+# Load your saved model and tokenizer
+model_path = "fine_tuned_distilbert"
+tokenizer = DistilBertTokenizer.from_pretrained(model_path)
+model = DistilBertForSequenceClassification.from_pretrained(model_path)
 
-# Step 3: Tokenize dataset
-tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
-def tokenize_function(examples):
-    text = [title + " " + content for title, content in zip(examples["title"], examples["content"])]
-    return tokenizer(text, padding="max_length", truncation=True, max_length=512)
-tokenized_train = train_dataset.map(tokenize_function, batched=True)
-tokenized_test = test_dataset.map(tokenize_function, batched=True)
-tokenized_train = tokenized_train.remove_columns(["title", "content"])
-tokenized_test = tokenized_test.remove_columns(["title", "content"])
-tokenized_train.set_format("torch", columns=["input_ids", "attention_mask", "label"])
-tokenized_test.set_format("torch", columns=["input_ids", "attention_mask", "label"])
-
-# Step 4: Fine-tune model
-model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=2)
-def compute_metrics(pred):
-    labels = pred.label_ids
-    preds = pred.predictions.argmax(-1)
-    acc = accuracy_score(labels, preds)
-    f1 = f1_score(labels, preds, average="weighted")
-    return {"accuracy": acc, "f1": f1}
-training_args = TrainingArguments(
-    output_dir="./results",
-    num_train_epochs=3,
-    per_device_train_batch_size=16,
-    per_device_eval_batch_size=16,
-    warmup_steps=500,
-    weight_decay=0.01,
-    logging_dir="./logs",
-    logging_steps=10,
-    eval_strategy="epoch",
-    save_strategy="epoch",
-    load_best_model_at_end=True,
-)
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    train_dataset=tokenized_train,
-    eval_dataset=tokenized_test,
-    compute_metrics=compute_metrics,
-)
-trainer.train()
-
-# Step 5: Evaluate and predict
-eval_results = trainer.evaluate()
-print("Evaluation results:", eval_results)
-model.save_pretrained("./fine_tuned_distilbert")
-tokenizer.save_pretrained("./fine_tuned_distilbert")
+# Prediction function
 def predict_sentiment(text):
     inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -66,14 +18,6 @@ def predict_sentiment(text):
     logits = outputs.logits
     predicted_class = torch.argmax(logits, dim=1).item()
     return "Positive" if predicted_class == 1 else "Negative"
-example_reviews = [
-    "Great product! Fast shipping and works perfectly as described.",
-    "Terrible quality, broke after one use. Very disappointed.",
-    "The item is okay, not amazing but does the job for the price."
-]
-for review in example_reviews:
-    sentiment = predict_sentiment(review)
-    print(f"Review: {review}\nPredicted Sentiment: {sentiment}\n")
 
 # Create Gradio interface
 interface = gr.Interface(fn=predict_sentiment, inputs="text", outputs="text", title="Amazon Sentiment Analysis Demo")