Update app.py

app.py

@@ -2,92 +2,37 @@ import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import gradio as gr
 
-# ============================================================
-# CONFIG
-# ============================================================
 MODEL_NAME = "bert-base-uncased"
-WEIGHTS_PATH = "bert_sentiment_model.pt"  # Upload this file to your Space
+WEIGHTS_PATH = "bert_sentiment_model.pt"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Define label mappings (same as your training code)
-id2label = {
-    0: "Positive",
-    1: "Negative",
-    2: "Neutral"
-}
-label2id = {v: k for k, v in id2label.items()}
-
-# ============================================================
-# LOAD MODEL AND TOKENIZER
-# ============================================================
+# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
-model = AutoModelForSequenceClassification.from_pretrained(
-    MODEL_NAME,
-    num_labels=len(id2label),
-    id2label=id2label,
-    label2id=label2id
-)
-
-# Load fine-tuned weights
+# ⚠️ Make sure num_labels matches your training setup (2 if you trained with Positive/Negative)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=2)
 model.load_state_dict(torch.load(WEIGHTS_PATH, map_location=DEVICE))
 model.to(DEVICE)
 model.eval()
 
-# ============================================================
-# INFERENCE FUNCTION
-# ============================================================
-def predict_sentiment(text):
-    if not text.strip():
-        return {"Sentiment": "⚠️ Please enter some text.", "Confidence": 0.0}
-
-    encoding = tokenizer(
-        text,
-        add_special_tokens=True,
-        max_length=256,
-        padding="max_length",
-        truncation=True,
-        return_attention_mask=True,
-        return_tensors="pt"
-    )
-
-    input_ids = encoding["input_ids"].to(DEVICE)
-    attention_mask = encoding["attention_mask"].to(DEVICE)
+id2label = {0: "Positive", 1: "Negative"}
 
+def predict_sentiment(text):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=256).to(DEVICE)
     with torch.no_grad():
-        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
-        logits = outputs.logits
-        probabilities = torch.softmax(logits, dim=1)
-        predicted_class = torch.argmax(probabilities, dim=1).item()
-        confidence = probabilities[0][predicted_class].item()
-
-    sentiment = id2label[predicted_class]
-    return {"Sentiment": sentiment, "Confidence": round(confidence, 4)}
+        outputs = model(**inputs)
+        probs = torch.softmax(outputs.logits, dim=1)
+        pred = torch.argmax(probs, dim=1).item()
+        confidence = probs[0][pred].item()
+    return f"{id2label[pred]} ({confidence:.2f} confidence)"
 
-# ============================================================
-# GRADIO INTERFACE
-# ============================================================
 demo = gr.Interface(
     fn=predict_sentiment,
-    inputs=gr.Textbox(lines=4, placeholder="Type your review or feedback here...", label="💬 Enter Text"),
-    outputs=[
-        gr.Label(label="Predicted Sentiment"),
-        gr.Number(label="Confidence", precision=4)
-    ],
-    title="BERT Sentiment Analyzer",
-    description="Fine-tuned BERT model for classifying text into Positive, Negative, or Neutral sentiments.",
-    examples=[
-        ["I love this product! It's fantastic!"],
-        ["This is terrible, worst experience ever."],
-        ["It's okay, nothing special."],
-        ["Amazing quality and great service!"],
-        ["Very disappointed with this product."]
-    ],
-    theme="gradio/soft"
+    inputs=gr.Textbox(lines=3, placeholder="Type your review here..."),
+    outputs="text",
+    title="Sentiment Analyzer (BERT)",
+    description="Predicts whether text sentiment is Positive or Negative."
 )
 
-# ============================================================
-# LAUNCH
-# ============================================================
 if __name__ == "__main__":
     demo.launch()