Update app.py

app.py

@@ -1,11 +1,3 @@
-# app.py
-
-import subprocess
-
-# Install dependencies
-subprocess.run(["pip", "install", "-r", "requirements.txt"])
-
-# Rest of your code
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
@@ -15,11 +7,11 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
 # Define emotion labels used by the model
-emotion_labels = ["admiration", "amusement", "anger", "annoyance", "approval", 
-                  "caring", "confusion", "curiosity", "desire", "disappointment", 
-                  "disapproval", "disgust", "embarrassment", "excitement", 
-                  "fear", "gratitude", "grief", "joy", "love", "nervousness", 
-                  "optimism", "pride", "realization", "relief", "remorse", 
+emotion_labels = ["admiration", "amusement", "anger", "annoyance", "approval",
+                  "caring", "confusion", "curiosity", "desire", "disappointment",
+                  "disapproval", "disgust", "embarrassment", "excitement",
+                  "fear", "gratitude", "grief", "joy", "love", "nervousness",
+                  "optimism", "pride", "realization", "relief", "remorse",
                   "sadness", "surprise", "neutral"]
 
 def predict_emotion(text):
@@ -29,15 +21,34 @@ def predict_emotion(text):
     predicted_class = logits.argmax().item()
 
     predicted_emotion = emotion_labels[predicted_class]
-    return predicted_emotion  # Return the predicted emotion directly
+    confidence = logits.softmax(dim=1).squeeze()[predicted_class].item()
+    return predicted_emotion, confidence  # Return the predicted emotion and confidence
+
+def get_confidence_color(confidence):
+    # Define a color scale based on confidence
+    if confidence >= 0.8:
+        return "green"
+    elif confidence >= 0.5:
+        return "orange"
+    else:
+        return "red"
 
 iface = gr.Interface(
     fn=predict_emotion,
     inputs=gr.Textbox(),
-    outputs="text",
+    outputs=[
+        gr.Textbox(),
+        "text",
+        gr.Textbox("Confidence Score:", default=""),
+    ],
     live=True,
     title="Emotion Prediction",
     description="Enter a sentence for emotion prediction.",
 )
 
+# Customize the interface appearance
+iface.style(
+    confidence_score=get_confidence_color
+)
+
 iface.launch()