Update app.py

app.py

@@ -1,38 +1,49 @@
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import gradio as gr
 
 # Load model and tokenizer
-model_name = "ogflash/yelp_review_classifier"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
-
-# Inference function
-def classify(text):
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+model_path = "model"  # Your local fine-tuned model directory
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = AutoModelForSequenceClassification.from_pretrained(model_path)
 
-    # Remove token_type_ids for DistilBERT compatibility
-    if "token_type_ids" in inputs:
-        del inputs["token_type_ids"]
+# Define label mapping
+id2label = model.config.id2label or {
+    0: "Negative",
+    1: "Neutral",
+    2: "Positive"
+}
 
+def classify(text):
+    inputs = tokenizer(text, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
-        logits = outputs.logits
-        predicted_class = torch.argmax(logits, dim=1).item()
-        score = torch.softmax(logits, dim=1)[0][predicted_class].item()
+    logits = outputs.logits
+    probs = torch.softmax(logits, dim=1)[0]
+
+    predicted_class = torch.argmax(probs).item()
+    label = id2label.get(predicted_class, f"LABEL_{predicted_class}")
+    confidence = round(float(probs[predicted_class]) * 100, 2)
 
-    label = model.config.id2label[predicted_class]
-    return f"{label} ({round(score * 100, 2)}%)"
+    all_probs = {
+        id2label.get(i, f"LABEL_{i}"): f"{round(float(prob)*100, 2)}%"
+        for i, prob in enumerate(probs)
+    }
+
+    return f"Prediction: {label} ({confidence}%)", all_probs
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## 📝 Yelp Review Sentiment Classifier")
+    gr.Markdown("# Yelp Review Sentiment Classifier")
+    with gr.Row():
+        input_box = gr.Textbox(lines=4, label="Enter a review")
     with gr.Row():
-        text_input = gr.Textbox(label="Enter Review", placeholder="Type your Yelp review...", lines=4)
+        output_label = gr.Textbox(label="Predicted Sentiment")
+        output_probs = gr.JSON(label="All Class Probabilities")
     with gr.Row():
-        submit_btn = gr.Button("Classify")
-    output = gr.Textbox(label="Prediction")
+        classify_btn = gr.Button("Classify")
 
-    submit_btn.click(fn=classify, inputs=text_input, outputs=output)
+    classify_btn.click(fn=classify, inputs=input_box, outputs=[output_label, output_probs])
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()