app.py

import gradio as gr
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel
import numpy as np

# Configuration
MODEL_NAME = "roberta-base"
MAX_LEN = 200
EMOTIONS = ["anger", "fear", "joy", "sadness", "surprise"]
EMOTION_EMOJIS = ["😠", "😨", "😊", "😢", "😲"]
EMOTION_COLORS = ["#ef4444", "#f59e0b", "#10b981", "#3b82f6", "#8b5cf6"]

# Model Architecture (MUST MATCH TRAINING)
class RobertaEmotion(nn.Module):
    def __init__(self, model_name=MODEL_NAME, dropout=0.35, num_labels=5):
        super().__init__()
        self.backbone = AutoModel.from_pretrained(model_name)
        hidden_size = self.backbone.config.hidden_size
        self.dropout = nn.Dropout(dropout)
        self.head = nn.Linear(hidden_size, num_labels)

    def forward(self, input_ids, attention_mask):
        out = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
        if hasattr(out, "pooler_output") and out.pooler_output is not None:
            pooled = out.pooler_output
        else:
            pooled = out.last_hidden_state[:, 0]
        x = self.dropout(pooled)
        logits = self.head(x)
        return logits

# Load model and tokenizer
print("🔄 Loading EmotiScan model...")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"📱 Device: {device}")

try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = RobertaEmotion(num_labels=len(EMOTIONS))
    
    # Load trained weights
    state_dict = torch.load('roberta.pth', map_location=device)
    model.load_state_dict(state_dict)
    model = model.to(device)
    model.eval()
    
    print("✅ EmotiScan ready!")
except Exception as e:
    print(f"⚠️ Error loading model: {e}")
    raise e

# Optimized thresholds from training
BEST_THRESHOLDS = np.array([0.5, 0.5, 0.5, 0.5, 0.5])

def predict_emotions(text):
    """Predict emotions from text with enhanced visualization"""
    if not text or not text.strip():
        return """
        <div style="text-align: center; padding: 40px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 16px; color: white;">
            <div style="font-size: 48px; margin-bottom: 16px;">🤔</div>
            <div style="font-size: 20px; font-weight: 600;">Waiting for your text...</div>
            <div style="font-size: 14px; opacity: 0.9; margin-top: 8px;">Enter some text above to analyze emotions</div>
        </div>
        """
    
    try:
        # Tokenize
        encoding = tokenizer(
            text,
            truncation=True,
            padding="max_length",
            max_length=MAX_LEN,
            return_tensors="pt"
        )
        
        input_ids = encoding["input_ids"].to(device)
        attention_mask = encoding["attention_mask"].to(device)
        
        # Predict
        with torch.no_grad():
            logits = model(input_ids, attention_mask)
            probs = torch.sigmoid(logits).cpu().numpy()[0]
        
        # Apply thresholds
        predictions = (probs > BEST_THRESHOLDS).astype(int)
        
        # Build beautiful HTML output
        html = """
        <style>
            @keyframes fadeIn {
                from { opacity: 0; transform: translateY(10px); }
                to { opacity: 1; transform: translateY(0); }
            }
            @keyframes pulse {
                0%, 100% { transform: scale(1); }
                50% { transform: scale(1.05); }
            }
            .emotion-card {
                animation: fadeIn 0.5s ease-out;
                transition: all 0.3s ease;
            }
            .emotion-card:hover {
                transform: translateY(-4px);
                box-shadow: 0 8px 24px rgba(0,0,0,0.15);
            }
            .detected-badge {
                animation: pulse 2s infinite;
            }
            .progress-bar {
                transition: width 0.8s ease-out;
            }
        </style>
        """
        
        # Detected emotions section
        detected = [(emotion, emoji, prob, color) for emotion, emoji, prob, pred, color 
                   in zip(EMOTIONS, EMOTION_EMOJIS, probs, predictions, EMOTION_COLORS) if pred == 1]
        
        if detected:
            html += """
            <div style="background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); 
                        padding: 24px; border-radius: 16px; margin-bottom: 24px; text-align: center;">
                <div style="color: white; font-size: 18px; font-weight: 600; margin-bottom: 16px;">
                    🎯 Detected Emotions
                </div>
                <div style="display: flex; gap: 12px; flex-wrap: wrap; justify-content: center;">
            """
            for emotion, emoji, prob, color in detected:
                html += f"""
                <div class="detected-badge" style="background: white; padding: 12px 20px; 
                     border-radius: 24px; display: flex; align-items: center; gap: 8px;
                     box-shadow: 0 4px 12px rgba(0,0,0,0.1);">
                    <span style="font-size: 24px;">{emoji}</span>
                    <span style="font-weight: 600; color: {color}; text-transform: capitalize;">
                        {emotion}
                    </span>
                    <span style="background: {color}; color: white; padding: 2px 8px; 
                          border-radius: 12px; font-size: 12px; font-weight: 600;">
                        {prob:.0%}
                    </span>
                </div>
                """
            html += "</div></div>"
        else:
            html += """
            <div style="background: linear-gradient(135deg, #6b7280 0%, #4b5563 100%); 
                        padding: 24px; border-radius: 16px; margin-bottom: 24px; text-align: center; color: white;">
                <div style="font-size: 32px; margin-bottom: 8px;">😐</div>
                <div style="font-size: 16px; font-weight: 600;">No Strong Emotions Detected</div>
                <div style="font-size: 14px; opacity: 0.8; margin-top: 4px;">All emotions below threshold</div>
            </div>
            """
        
        # All emotions with progress bars
        html += """
        <div style="background: white; padding: 24px; border-radius: 16px; box-shadow: 0 2px 8px rgba(0,0,0,0.1);">
            <div style="font-size: 18px; font-weight: 600; margin-bottom: 20px; color: #1f2937;">
                📊 Emotion Breakdown
            </div>
            <div style="display: flex; flex-direction: column; gap: 16px;">
        """
        
        for emotion, emoji, prob, color in zip(EMOTIONS, EMOTION_EMOJIS, probs, EMOTION_COLORS):
            html += f"""
            <div class="emotion-card" style="background: #f9fafb; padding: 16px; border-radius: 12px; 
                 border-left: 4px solid {color};">
                <div style="display: flex; justify-content: space-between; align-items: center; margin-bottom: 8px;">
                    <div style="display: flex; align-items: center; gap: 10px;">
                        <span style="font-size: 28px;">{emoji}</span>
                        <span style="font-weight: 600; color: #374151; text-transform: capitalize; font-size: 16px;">
                            {emotion}
                        </span>
                    </div>
                    <span style="font-weight: 700; color: {color}; font-size: 18px;">
                        {prob:.1%}
                    </span>
                </div>
                <div style="background: #e5e7eb; height: 12px; border-radius: 6px; overflow: hidden;">
                    <div class="progress-bar" style="background: linear-gradient(90deg, {color}, {color}dd); 
                         height: 100%; width: {prob*100}%; border-radius: 6px;
                         box-shadow: 0 0 8px {color}66;"></div>
                </div>
            </div>
            """
        
        html += "</div></div>"
        
        return html
        
    except Exception as e:
        return f"""
        <div style="background: #fef2f2; border: 2px solid #ef4444; padding: 20px; 
             border-radius: 12px; color: #991b1b;">
            <div style="font-size: 24px; margin-bottom: 8px;">⚠️</div>
            <div style="font-weight: 600; margin-bottom: 4px;">Analysis Error</div>
            <div style="font-size: 14px;">{str(e)}</div>
        </div>
        """

# Example texts
examples = [
    ["I just got promoted at work! I can't believe it!"],
    ["I'm so worried about the exam tomorrow. What if I fail?"],
    ["This is absolutely unacceptable! I demand to speak to the manager!"],
    ["I miss my family so much. It's been months since I've seen them."],
    ["Wow! I never expected to see you here!"],
    ["I'm excited but also nervous about starting my new job next week."],
]

# Create Gradio Interface
with gr.Blocks() as demo:
    gr.HTML("""
        <style>
            @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700&display=swap');
            
            * {
                font-family: 'Inter', sans-serif !important;
            }
            
            .gradio-container {
                max-width: 1400px !important;
                margin: 0 auto !important;
            }
            
            button {
                border-radius: 12px !important;
                font-weight: 600 !important;
                transition: all 0.3s ease !important;
            }
            
            button:hover {
                transform: translateY(-2px) !important;
                box-shadow: 0 6px 20px rgba(0,0,0,0.15) !important;
            }
            
            textarea {
                border-radius: 12px !important;
                border: 2px solid #e5e7eb !important;
                transition: all 0.3s ease !important;
            }
            
            textarea:focus {
                border-color: #667eea !important;
                box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1) !important;
            }
        </style>
        <div style="background: linear-gradient(135deg, #6ee7b7 0%, #34d399 100%); padding: 40px; border-radius: 20px; margin-bottom: 30px; text-align: center; color: white; box-shadow: 0 10px 30px rgba(52, 211, 153, 0.3);">
            <div style="font-size: 56px; margin-bottom: 16px;">🎭</div>
            <h1 style="font-size: 48px; font-weight: 700; margin: 0 0 12px 0; text-shadow: 0 2px 4px rgba(0,0,0,0.1);">
                EmotiScan
            </h1>
            <p style="font-size: 20px; opacity: 0.95; margin: 0; font-weight: 500;">
                AI-Powered Multi-Emotion Detection
            </p>
            <div style="margin-top: 20px; display: flex; gap: 16px; justify-content: center; flex-wrap: wrap;">
                <span style="background: rgba(255,255,255,0.2); padding: 8px 16px; border-radius: 20px; font-size: 14px;">
                    😠 Anger
                </span>
                <span style="background: rgba(255,255,255,0.2); padding: 8px 16px; border-radius: 20px; font-size: 14px;">
                    😨 Fear
                </span>
                <span style="background: rgba(255,255,255,0.2); padding: 8px 16px; border-radius: 20px; font-size: 14px;">
                    😊 Joy
                </span>
                <span style="background: rgba(255,255,255,0.2); padding: 8px 16px; border-radius: 20px; font-size: 14px;">
                    😢 Sadness
                </span>
                <span style="background: rgba(255,255,255,0.2); padding: 8px 16px; border-radius: 20px; font-size: 14px;">
                    😲 Surprise
                </span>
            </div>
        </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            text_input = gr.Textbox(
                label="📝 Your Text",
                placeholder="Type or paste your text here to discover the emotions within...",
                lines=8,
                max_lines=12
            )
            with gr.Row():
                analyze_btn = gr.Button("🔮 Analyze Emotions", variant="primary", size="lg")
                clear_btn = gr.Button("🗑️ Clear", size="lg")
        
        with gr.Column(scale=1):
            output = gr.HTML(label="Analysis Results", value="""
                <div style="text-align: center; padding: 60px 40px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); 
                     border-radius: 16px; color: white; height: 100%;">
                    <div style="font-size: 64px; margin-bottom: 20px;">🎭</div>
                    <div style="font-size: 24px; font-weight: 700; margin-bottom: 12px;">Welcome to EmotiScan</div>
                    <div style="font-size: 16px; opacity: 0.9;">Enter text to begin emotional analysis</div>
                </div>
            """)
    
    gr.Examples(
        examples=examples,
        inputs=text_input,
        outputs=output,
        fn=predict_emotions,
        cache_examples=False,
        label="💡 Try These Examples"
    )
    
    gr.HTML("""
        <div style="background: white; padding: 32px; border-radius: 16px; margin-top: 30px; box-shadow: 0 2px 8px rgba(0,0,0,0.1);">
            <h2 style="color: #1f2937; margin-bottom: 20px; font-size: 24px; font-weight: 700;">
                🧠 About EmotiScan
            </h2>
            <p style="color: #4b5563; line-height: 1.8; margin-bottom: 24px; font-size: 15px;">
                EmotiScan uses state-of-the-art deep learning to detect multiple emotions simultaneously in text. 
                Unlike traditional single-emotion classifiers, our model recognizes that human expression is complex 
                and nuanced—one piece of text can convey multiple emotions at once.
            </p>
            
            <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(250px, 1fr)); gap: 20px; margin-top: 24px;">
                <div style="background: linear-gradient(135deg, #667eea22 0%, #764ba222 100%); padding: 20px; border-radius: 12px;">
                    <div style="font-size: 32px; margin-bottom: 8px;">🤖</div>
                    <div style="font-weight: 600; color: #1f2937; margin-bottom: 4px;">Model</div>
                    <div style="color: #6b7280; font-size: 14px;">RoBERTa-base (125M params)</div>
                </div>
                <div style="background: linear-gradient(135deg, #10b98122 0%, #059669 22 100%); padding: 20px; border-radius: 12px;">
                    <div style="font-size: 32px; margin-bottom: 8px;">🎯</div>
                    <div style="font-weight: 600; color: #1f2937; margin-bottom: 4px;">Accuracy</div>
                    <div style="color: #6b7280; font-size: 14px;">Optimized F1-Score per class</div>
                </div>
                <div style="background: linear-gradient(135deg, #f59e0b22 0%, #d9770622 100%); padding: 20px; border-radius: 12px;">
                    <div style="font-size: 32px; margin-bottom: 8px;">⚡</div>
                    <div style="font-weight: 600; color: #1f2937; margin-bottom: 4px;">Speed</div>
                    <div style="color: #6b7280; font-size: 14px;">Real-time inference</div>
                </div>
            </div>
            
            <div style="margin-top: 32px; padding: 20px; background: #f9fafb; border-radius: 12px; border-left: 4px solid #667eea;">
                <div style="font-weight: 600; color: #1f2937; margin-bottom: 12px; font-size: 16px;">
                    📚 Technical Details
                </div>
                <ul style="color: #4b5563; line-height: 2; margin: 0; padding-left: 20px; font-size: 14px;">
                    <li><strong>Architecture:</strong> Transformer encoder with classification head</li>
                    <li><strong>Training:</strong> BCE Loss with label smoothing (0.05)</li>
                    <li><strong>Max Tokens:</strong> 200 tokens per input</li>
                    <li><strong>Dropout:</strong> 0.35 for regularization</li>
                    <li><strong>Multi-Label:</strong> Each emotion is independently predicted</li>
                </ul>
            </div>
            
            <div style="margin-top: 24px; text-align: center; color: #9ca3af; font-size: 14px;">
                <p style="margin: 0;">Built with PyTorch • Transformers • Gradio</p>
                <p style="margin: 4px 0 0 0;">2025 Sep DLGenAI Course Project</p>
            </div>
        </div>
    """)
    
    # Event handlers
    analyze_btn.click(fn=predict_emotions, inputs=text_input, outputs=output)
    clear_btn.click(
        fn=lambda: ("", """
            <div style="text-align: center; padding: 60px 40px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); 
                 border-radius: 16px; color: white; height: 100%;">
                <div style="font-size: 64px; margin-bottom: 20px;">🎭</div>
                <div style="font-size: 24px; font-weight: 700; margin-bottom: 12px;">Welcome to EmotiScan</div>
                <div style="font-size: 16px; opacity: 0.9;">Enter text to begin emotional analysis</div>
            </div>
        """), 
        inputs=None, 
        outputs=[text_input, output]
    )
    text_input.submit(fn=predict_emotions, inputs=text_input, outputs=output)

if __name__ == "__main__":
    demo.launch(share=True, server_name="0.0.0.0")