Update app.py

app.py

@@ -1,44 +1,42 @@
-import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
 
-# Load model & tokenizer from HF or local path
-model_name = "ogflash/yelp_review_classifier"  # Change if needed
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
+# Load the tokenizer and model from local path (or HF if internet is available)
+model = AutoModelForSequenceClassification.from_pretrained("ogflash/yelp_review_classifier")
+tokenizer = AutoTokenizer.from_pretrained("ogflash/yelp_review_classifier")
 
-# Fix for DistilBERT models that don't accept token_type_ids
+# Prediction function
 def classify(text):
     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
-    # Remove token_type_ids if not supported
-    if "token_type_ids" in inputs and "token_type_ids" not in model.forward.__code__.co_varnames:
-        del inputs["token_type_ids"]
-
-    with torch.no_grad():
-        outputs = model(**inputs)
-        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
-        top_class = torch.argmax(probs, dim=1).item()
-        confidence = probs[0][top_class].item() * 100
-
-    # Reliable label mapping
-    id2label = model.config.id2label
-    if not id2label or not isinstance(id2label, dict) or len(id2label) == 0:
-        id2label = {
-            0: "Negative",
-            1: "Neutral",
-            2: "Positive"
-        }
-
-    label_name = id2label.get(top_class, f"LABEL_{top_class}")
-    return f"{label_name} ({confidence:.2f}%)"
-
-# UI with Gradio
-iface = gr.Interface(
-    fn=classify,
-    inputs=gr.Textbox(lines=3, placeholder="Enter text to analyze..."),
-    outputs="text",
-    title="Sentiment Classifier",
-    description="Predicts sentiment using a BERT-based model.",
-)
-
-iface.launch(share=True)
+
+    # Remove token_type_ids if using DistilBERT
+    if "token_type_ids" in inputs:
+        inputs.pop("token_type_ids")
+
+    outputs = model(**inputs)
+    logits = outputs.logits
+    predicted_class_id = torch.argmax(logits, dim=1).item()
+    score = torch.softmax(logits, dim=1)[0][predicted_class_id].item()
+
+    # Map labels using if-elif-else
+    label = f"LABEL_{predicted_class_id}"
+    if label == "LABEL_0":
+        label_name = "Negative"
+    elif label == "LABEL_1":
+        label_name = "Neutral"
+    elif label == "LABEL_2":
+        label_name = "Positive"
+    else:
+        label_name = label  # fallback
+
+    return f"{label_name} ({score * 100:.2f}%)"
+
+# Gradio UI
+iface = gr.Interface(fn=classify,
+                     inputs=gr.Textbox(lines=2, placeholder="Enter your review here..."),
+                     outputs="text",
+                     title="Sentiment Classifier",
+                     description="Classifies text into Positive, Neutral, or Negative.")
+
+iface.launch()