Update app.py

app.py

@@ -80,16 +80,49 @@ templates = {
 }
 
 # Prediction logic
+# Emotion to Emoji mapping
+emotion_emojis = {
+    "Angry": "😠",
+    "Disgust": "🤢",
+    "Fear": "😨",
+    "Happy": "😄",
+    "Sad": "😢",
+    "Surprise": "😲",
+    "Neutral": "😐"
+}
+
 def predict_emotion_and_icebreaker(image, tone):
     image = Image.fromarray(image).convert("RGB")
     image = transform(image).unsqueeze(0).to(device)
-    
+
     with torch.no_grad():
         output = model(image)
         pred = output.argmax(dim=1).item()
-    
+
     emotion = classes[pred]
-    prompt = templates[tone][emotion]
+    emoji = emotion_emojis.get(emotion, "🙂")
+
+    # Refined prompt for short, relevant output with emoji
+    prompt = templates[tone][emotion] + " Respond in one short sentence with an emoji."
+
+    result = gen(
+        prompt,
+        max_length=30,
+        num_return_sequences=1,
+        do_sample=True,
+        temperature=0.8,
+        top_p=0.9,
+        pad_token_id=50256
+    )[0]['generated_text']
+
+    # Clean generated text by stripping prompt
+    response = result[len(prompt):].strip()
+
+    # Ensure minimal valid output
+    if not response or len(response) < 5:
+        response = "Let's get the conversation going! 🎉"
+
+    return f"🧠 Emotion Detected: {emotion} {emoji}\n💬 Icebreaker ({tone}):\n{response}"
 
     # Generate icebreaker with better sampling
     generated = gen(