Update app.py

app.py

@@ -26,23 +26,23 @@ def classify_review(text):
     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
     with torch.no_grad():
         outputs = model(**inputs)
-    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
-    label_idx = torch.argmax(probs).item()
-    predicted_label = labels[label_idx]
+    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)[0]
+    prob_dict = {label: float(prob) for label, prob in zip(labels, probs)}
+    predicted_label = labels[torch.argmax(probs).item()]
 
     # Logging
     with open(log_path, mode='a', newline='') as file:
         writer = csv.writer(file)
         writer.writerow([datetime.now().isoformat(), text, predicted_label])
     
-    return {label: float(prob) for label, prob in zip(labels, probs[0])}, text, predicted_label
+    return prob_dict, text, predicted_label
 
 def save_correction(text, predicted_label, user_correction):
     if user_correction != predicted_label:
         with open(corrections_path, mode='a', newline='') as file:
             writer = csv.writer(file)
             writer.writerow([datetime.now().isoformat(), text, predicted_label, user_correction])
-    return f"Thanks! Correction recorded: {user_correction}"
+    return f"✅ Thanks! Correction recorded: {user_correction}"
 
 # Gradio UI
 with gr.Blocks() as demo:
@@ -50,7 +50,7 @@ with gr.Blocks() as demo:
     gr.Markdown("Enter a course review and get the sentiment prediction. You can correct the result if needed.")
 
     input_text = gr.Textbox(lines=4, placeholder="Enter course review here...")
-    output_label = gr.Label(num_top_classes=3)
+    output_label = gr.Label(num_top_classes=3, label="Predicted Sentiment")
     predict_btn = gr.Button("Classify")
 
     with gr.Row(visible=False) as correction_row:
@@ -62,13 +62,12 @@ with gr.Blocks() as demo:
     hidden_text = gr.Textbox(visible=False)
     hidden_pred = gr.Textbox(visible=False)
 
-    def show_correction_ui(result, text, pred):
-        return result, gr.update(visible=True), text, pred
-
-    predict_btn.click(classify_review, inputs=input_text, outputs=[output_label, hidden_text, hidden_pred])\
-           .then(show_correction_ui, inputs=[output_label, hidden_text, hidden_pred],
-                 outputs=[output_label, correction_row, hidden_text, hidden_pred])
+    def show_correction_ui(_, text, pred):
+        return gr.update(visible=True), text, pred
 
+    predict_btn.click(classify_review, inputs=input_text, outputs=[output_label, hidden_text, hidden_pred]) \
+               .then(show_correction_ui, inputs=[output_label, hidden_text, hidden_pred],
+                     outputs=[correction_row, hidden_text, hidden_pred])
 
     submit_btn.click(save_correction, inputs=[hidden_text, hidden_pred, correction_dropdown], outputs=correction_status)