Update app.py

app.py

@@ -41,8 +41,8 @@ emotions_target = pd.Categorical(df['emotion']).codes
 emotion_classes = pd.Categorical(df['emotion']).categories
 
 # Load pre-trained BERT model for emotion prediction
-emotion_prediction_model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased")
-emotion_prediction_tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
+emotion_prediction_model = AutoModelForSequenceClassification.from_pretrained("bhadresh-savani/distilbert-base-uncased-emotion")
+emotion_prediction_tokenizer = AutoTokenizer.from_pretrained("bhadresh-savani/distilbert-base-uncased-emotion")
 
 # Lazy loading for the fine-tuned language model (DialoGPT)
 _finetuned_lm_tokenizer = None
@@ -154,19 +154,19 @@ def predict_emotion(context):
     emotion_prediction_pipeline = pipeline('text-classification', model=emotion_prediction_model, tokenizer=emotion_prediction_tokenizer, top_k=None)
     predictions = emotion_prediction_pipeline(context)
     emotion_scores = {prediction['label']: prediction['score'] for prediction in predictions[0]}
-    predicted_label = max(emotion_scores, key=emotion_scores.get)
+    predicted_emotion = max(emotion_scores, key=emotion_scores.get)
     
-    # Map the predicted label to our emotion categories
+    # Map the predicted emotion to our emotion categories
     emotion_mapping = {
-        'LABEL_0': 'joy',
-        'LABEL_1': 'sadness',
-        'LABEL_2': 'anger',
-        'LABEL_3': 'fear',
-        'LABEL_4': 'surprise',
-        'LABEL_5': 'disgust'
+        'sadness': 'sadness',
+        'joy': 'joy',
+        'love': 'pleasure',
+        'anger': 'anger',
+        'fear': 'fear',
+        'surprise': 'surprise'
     }
     
-    return emotion_mapping.get(predicted_label, 'neutral')
+    return emotion_mapping.get(predicted_emotion, 'neutral')
 
 def generate_text(prompt, emotion=None, max_length=100):
     finetuned_lm_tokenizer, finetuned_lm_model = get_finetuned_lm_model()
@@ -231,6 +231,7 @@ def reset_emotions():
             emotions[emotion]['percentage'] = 10
             emotions[emotion]['intensity'] = 0
     emotions['ideal_state']['percentage'] = 100
+    return get_emotion_summary()
 
 def respond_to_user(user_input, chat_history):
     # Predict the emotion from the user input
@@ -261,36 +262,25 @@ def respond_to_user(user_input, chat_history):
     
     return response, chat_history, get_emotion_summary()
 
-def chat_interface(user_input, history):
-    response, updated_history, emotion_summary = respond_to_user(user_input, history)
-    return response, updated_history, emotion_summary
-
 # Gradio interface
-iface = gr.Interface(
-    fn=chat_interface,
-    inputs=[
-        gr.Textbox(lines=2, placeholder="Type your message here..."),
-        gr.State([])
-    ],
-    outputs=[
-        gr.Textbox(label="AI Response"),
-        gr.State(),
-        gr.Textbox(label="Current Emotional State", lines=10)
-    ],
-    title="Emotion-Aware AI Chatbot",
-    description="Chat with an AI that understands and responds to emotions.",
-    allow_flagging="never",
-    theme="default"
-)
-
-# Add a button to reset emotions
-reset_button = gr.Button("Reset Emotions")
-reset_button.click(
-    fn=reset_emotions,
-    inputs=None,
-    outputs=gr.Textbox(label="Current Emotional State", lines=10),
-    api_name="reset_emotions"
-)
+with gr.Blocks() as demo:
+    gr.Markdown("# Emotion-Aware AI Chatbot")
+    gr.Markdown("Chat with an AI that understands and responds to emotions.")
+    
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox(label="Type your message here...")
+    clear = gr.Button("Clear")
+    
+    emotion_state = gr.Textbox(label="Current Emotional State", lines=10)
+    reset_button = gr.Button("Reset Emotions")
+    
+    def user(user_message, history):
+        response, updated_history, emotion_summary = respond_to_user(user_message, history)
+        return "", updated_history, emotion_summary
+    
+    msg.submit(user, [msg, chatbot], [msg, chatbot, emotion_state])
+    clear.click(lambda: None, None, chatbot, queue=False)
+    reset_button.click(reset_emotions, None, emotion_state, queue=False)
 
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()