avatar/sentiment_multi_emotion.py

# -*- coding: utf-8 -*-
"""
Sentiment Multi-Emotion Analyzer - Advanced emotion detection using multi-class models

Uses RoBERTa-based emotion classifier for 7+ emotion detection:
- anger, disgust, fear, joy, neutral, sadness, surprise

This provides granular emotion detection instead of just positive/negative
"""

import torch
from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer
from typing import Dict, Any, List, Optional
import warnings
import re

warnings.filterwarnings("ignore", category=UserWarning)


class MultiEmotionAnalyzer:
    """
    Multi-class emotion analyzer using DistilRoBERTa
    
    Detects 7 core emotions: anger, disgust, fear, joy, neutral, sadness, surprise
    Maps to extended emotion vocabulary for nuanced emoji display
    """
    
    # Core emotion to polarity mapping
    EMOTION_TO_POLARITY = {
        "joy": "positive",
        "surprise": "positive",  # Generally positive context
        "neutral": "neutral",
        "sadness": "negative",
        "anger": "negative",
        "fear": "negative",
        "disgust": "negative",
    }
    
    # Extended emotion mapping from core emotions + context
    EXTENDED_EMOTION_MAP = {
        # Joy variations based on intensity/context
        "joy": {
            "high_confidence": "ecstasy",
            "medium_confidence": "joy",
            "low_confidence": "contentment",
        },
        # Sadness variations
        "sadness": {
            "high_confidence": "grief",
            "medium_confidence": "sadness", 
            "low_confidence": "melancholy",
        },
        # Anger variations
        "anger": {
            "high_confidence": "rage",
            "medium_confidence": "anger",
            "low_confidence": "irritation",
        },
        # Fear variations
        "fear": {
            "high_confidence": "terror",
            "medium_confidence": "fear",
            "low_confidence": "anxiety",
        },
        # Disgust variations
        "disgust": {
            "high_confidence": "revulsion",
            "medium_confidence": "disgust",
            "low_confidence": "contempt",
        },
        # Surprise variations
        "surprise": {
            "high_confidence": "astonishment",
            "medium_confidence": "surprise",
            "low_confidence": "curiosity",
        },
        # Neutral variations
        "neutral": {
            "high_confidence": "neutral",
            "medium_confidence": "neutral",
            "low_confidence": "neutral",
        },
    }
    
    # Context keywords for emotion refinement
    # IMPORTANT: More specific patterns (heartbreak, hopeless) must come BEFORE 
    # general patterns (love, hope) to avoid false matches
    CONTEXT_REFINEMENTS = {
        # === OVERRIDE PATTERNS (check first) ===
        
        # Heartbreak-related words (negative - despite "heart" being in love keywords)
        "heartbreak_keywords": ["heartbroken", "heartbreak", "broke my heart", "heart is broken", "shattered heart"],
        "heartbreak_emotion": "grief",
        "heartbreak_polarity": "negative",
        
        # Hopelessness-related words (negative - despite "hope" being positive)
        "hopelessness_keywords": ["hopeless", "hopelessness", "no hope", "lost all hope", "without hope"],
        "hopelessness_emotion": "despair",
        "hopelessness_polarity": "negative",
        
        # Sarcasm-related words (negative - often masking frustration)
        "sarcasm_keywords": ["oh great,", "just what i needed", "yeah right", "oh wonderful,", "just perfect", "oh fantastic", "exactly what i wanted", "wow, fantastic", "that's exactly", "that is exactly", "oh wonderful"],
        "sarcasm_emotion": "sarcasm",
        "sarcasm_polarity": "negative",
        
        # Indifferent/neutral-related words (neutral) - short responses
        "indifferent_keywords": ["meh", "whatever", "don't care", "doesn't matter", "indifferent", "fine i guess", "i'm fine", "just fine", "i guess so", "sure whatever"],
        "indifferent_emotion": "indifferent",
        "indifferent_polarity": "neutral",
        
        # === POSITIVE PATTERNS ===
        
        # Love-related words (but NOT heartbroken)
        "love_keywords": ["love", "adore", "beloved", "darling", "sweetheart"],
        "love_emotion": "love",
        "love_polarity": "positive",
        
        # Best/compliments (positive)
        "compliment_keywords": ["you're the best", "you are the best", "best ever", "you're amazing", "you are amazing"],
        "compliment_emotion": "admiration",
        "compliment_polarity": "positive",
        
        # Gratitude-related words
        "gratitude_keywords": ["thank", "grateful", "appreciate", "thankful", "gratitude"],
        "gratitude_emotion": "gratitude",
        "gratitude_polarity": "positive",
        
        # Excitement-related words
        "excitement_keywords": ["excited", "exciting", "thrilled", "can't wait", "pumped"],
        "excitement_emotion": "excitement",
        "excitement_polarity": "positive",
        
        # Hope-related words
        "hope_keywords": ["hope", "hopeful", "optimistic", "looking forward"],
        "hope_emotion": "hope",
        "hope_polarity": "positive",
        
        # Nostalgia-related words (positive - fond memories)
        "nostalgia_keywords": ["nostalgic", "nostalgia", "remember when", "miss the old", "good old days"],
        "nostalgia_emotion": "nostalgia",
        "nostalgia_polarity": "positive",
        
        # Confusion-related words (negative - distressing)
        "confusion_keywords": ["confused", "confusing", "puzzled", "don't understand", "baffled", "perplexed", "bewildered"],
        "confusion_emotion": "confused",
        "confusion_polarity": "negative",
        
        # Longing-related words (negative - painful desire)
        "longing_keywords": ["longing", "yearning", "yearn", "i miss you", "miss you so much", "miss him", "miss her"],
        "longing_emotion": "longing",
        "longing_polarity": "negative",
        
        # Playful-related words
        "playful_keywords": ["lol", "haha", "hehe", "😂", "🤣", "playful", "silly", "joking", "kidding"],
        "playful_emotion": "playful",
        "playful_polarity": "positive",
        
        # Pride-related words
        "pride_keywords": ["proud", "pride", "accomplished", "achieved"],
        "pride_emotion": "pride",
        "pride_polarity": "positive",
        
        # Embarrassment-related words
        "embarrassment_keywords": ["embarrassed", "embarrassing", "awkward", "cringe"],
        "embarrassment_emotion": "embarrassment",
        "embarrassment_polarity": "negative",
        
        # Sympathy-related words (positive - showing care)
        "sympathy_keywords": ["sorry for", "sympathize", "sympathy", "feel for you", "my condolences", "so sorry to hear"],
        "sympathy_emotion": "sympathy",
        "sympathy_polarity": "positive",
        
        # Empathy-related words (positive - showing understanding)
        "empathy_keywords": ["empathize", "empathy", "empathetic", "understand how you feel", "feel what you're feeling", "i understand", "i know how you feel"],
        "empathy_emotion": "empathy",
        "empathy_polarity": "positive",
        
        # Compassion-related words (positive - showing care)
        "compassion_keywords": ["compassion", "compassionate", "feel compassion", "care about", "caring"],
        "compassion_emotion": "compassion",
        "compassion_polarity": "positive",
        
        # Awe/Wonder-related words (positive - not fear)
        "awe_keywords": ["awe", "in awe", "awe-inspiring", "awesome", "awestruck", "wonder", "wondrous", "marvelous"],
        "awe_emotion": "awe",
        "awe_polarity": "positive",
        
        # Fascination/Interest-related words (positive)
        "fascination_keywords": ["fascinated", "fascinating", "intrigued", "intriguing", "captivated", "captivating", "curious"],
        "fascination_emotion": "fascination",
        "fascination_polarity": "positive",
        
        # Calm/Peace-related words (positive)
        "calm_keywords": ["calm", "peaceful", "serene", "tranquil", "relaxed", "at peace", "zen"],
        "calm_emotion": "calm",
        "calm_polarity": "positive",
        
        # Tenderness-related words (positive)
        "tenderness_keywords": ["tender", "tenderness", "gently", "softly", "warmth"],
        "tenderness_emotion": "tenderness",
        "tenderness_polarity": "positive",
        
        # Affection-related words (positive)
        "affection_keywords": ["affection", "affectionate", "fond", "fondness", "warmly"],
        "affection_emotion": "affection",
        "affection_polarity": "positive",
        
        # Shock-related words (negative - distressing)
        "shock_keywords": ["shock", "shocked", "in shock", "shocking"],
        "shock_emotion": "shock",
        "shock_polarity": "negative",
        
        # Thinking-related words (neutral)
        "thinking_keywords": ["thinking", "think about", "contemplating", "pondering", "considering", "let me think"],
        "thinking_emotion": "thinking",
        "thinking_polarity": "neutral",
        
        # Silly-related words (positive - fun)
        "silly_keywords": ["silly", "goofy", "dorky", "being silly", "acting silly"],
        "silly_emotion": "silly",
        "silly_polarity": "positive",
        
        # Determination-related words (positive)
        "determination_keywords": ["determined", "determination", "won't give up", "never give up", "nothing will stop"],
        "determination_emotion": "determination",
        "determination_polarity": "positive",
        
        # Anticipation-related words (positive)
        "anticipation_keywords": ["anticipating", "anticipation", "looking forward", "eagerly awaiting"],
        "anticipation_emotion": "anticipation",
        "anticipation_polarity": "positive",
        
        # Trust-related words (positive)
        "trust_keywords": ["trust", "believe in you", "have faith", "rely on you", "count on you"],
        "trust_emotion": "trust",
        "trust_polarity": "positive",
    }
    
    def __init__(self, model_name: str = "j-hartmann/emotion-english-distilroberta-base"):
        """
        Initialize multi-emotion classifier
        
        Args:
            model_name: HuggingFace model for 7-class emotion detection
                Default: j-hartmann/emotion-english-distilroberta-base
                Outputs: anger, disgust, fear, joy, neutral, sadness, surprise
        """
        self.model_name = model_name
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        
        print(f"📊 Loading multi-emotion model: {model_name}")
        print(f"   Device: {self.device}")
        
        try:
            # Initialize the pipeline with top_k to get all emotion scores
            self.pipeline = pipeline(
                "text-classification",
                model=model_name,
                device=0 if self.device == "cuda" else -1,
                top_k=None,  # Return all emotion scores
            )
            self.model_loaded = True
            print(f"✓ Multi-emotion model loaded (7 emotions)")
        except Exception as e:
            print(f"⚠️ Failed to load multi-emotion model: {e}")
            print(f"   Falling back to binary sentiment model")
            self.model_loaded = False
            # Fallback to binary model
            self.pipeline = pipeline(
                "sentiment-analysis",
                model="distilbert-base-uncased-finetuned-sst-2-english",
                device=0 if self.device == "cuda" else -1,
            )
    
    def _get_last_sentence(self, text: str) -> str:
        """Extract last sentence for real-time accuracy"""
        parts = re.split(r'[.!?;:\n]+', text)
        parts = [p.strip() for p in parts if p.strip()]
        return parts[-1] if parts else text.strip()
    
    def _check_context_keywords(self, text: str) -> Optional[tuple]:
        """
        Check for context keywords that indicate specific emotions
        
        Returns:
            Tuple of (emotion, polarity) if context match found, None otherwise
        """
        text_lower = text.lower().strip()
        text_clean = text_lower.rstrip('?!.,')
        
        # Special handling for very short responses (exact match)
        SHORT_NEUTRAL_WORDS = {"fine", "ok", "okay", "sure", "alright", "k", "kk", 
                               "yep", "yup", "nope", "nah", "yes", "no", "maybe",
                               "i guess", "i suppose", "perhaps"}
        if text_clean in SHORT_NEUTRAL_WORDS:
            return ("neutral", "neutral")
        
        # Check each context category
        for key, value in self.CONTEXT_REFINEMENTS.items():
            if key.endswith("_keywords"):
                base_key = key.replace("_keywords", "")
                emotion_key = f"{base_key}_emotion"
                polarity_key = f"{base_key}_polarity"
                
                keywords = value
                emotion = self.CONTEXT_REFINEMENTS.get(emotion_key)
                polarity = self.CONTEXT_REFINEMENTS.get(polarity_key)
                
                if emotion and any(kw in text_lower for kw in keywords):
                    return (emotion, polarity)
        
        return None
    
    def _get_intensity_level(self, score: float) -> str:
        """Map confidence score to intensity level"""
        if score >= 0.7:
            return "high_confidence"
        elif score >= 0.4:
            return "medium_confidence"
        else:
            return "low_confidence"
    
    def _refine_emotion(self, base_emotion: str, score: float, text: str) -> tuple:
        """
        Refine base emotion using context and intensity
        
        Returns:
            Tuple of (refined_emotion, polarity)
        """
        # First check for context-specific emotions
        context_result = self._check_context_keywords(text)
        if context_result:
            return context_result  # Returns (emotion, polarity)
        
        # Otherwise use intensity-based refinement
        intensity = self._get_intensity_level(score)
        emotion_variants = self.EXTENDED_EMOTION_MAP.get(base_emotion, {})
        refined_emotion = emotion_variants.get(intensity, base_emotion)
        
        # Use base emotion polarity
        polarity = self.EMOTION_TO_POLARITY.get(base_emotion, "neutral")
        
        return (refined_emotion, polarity)
    
    def analyze(self, text: str) -> Dict[str, Any]:
        """
        Analyze text and return detected emotion with details
        
        Args:
            text: Input text to analyze
            
        Returns:
            Dict with 'label' (emotion), 'score', 'polarity', and 'details'
        """
        if not text or not text.strip():
            return {
                "label": "neutral",
                "score": 0.0,
                "polarity": "neutral",
                "details": {"segment": "empty"}
            }
        
        # Use last sentence for real-time accuracy
        last_sentence = self._get_last_sentence(text)
        truncated = last_sentence[:512]
        
        try:
            if self.model_loaded:
                # Multi-emotion model returns all scores
                results = self.pipeline(truncated)
                
                if isinstance(results, list) and len(results) > 0:
                    if isinstance(results[0], list):
                        results = results[0]
                    
                    # Sort by score to get top emotion
                    sorted_results = sorted(results, key=lambda x: x['score'], reverse=True)
                    top_result = sorted_results[0]
                    
                    base_emotion = top_result['label'].lower()
                    score = top_result['score']
                    
                    # Refine emotion based on context - returns (emotion, polarity)
                    refined_emotion, polarity = self._refine_emotion(base_emotion, score, truncated)
                    
                    # Get all emotion scores for details
                    all_scores = {r['label'].lower(): r['score'] for r in sorted_results}
                    
                    return {
                        "label": refined_emotion,
                        "base_emotion": base_emotion,
                        "score": score,
                        "polarity": polarity,
                        "details": {
                            "segment": "last_sentence",
                            "text": truncated[:50],
                            "model": self.model_name,
                            "all_scores": all_scores,
                        }
                    }
            else:
                # Fallback binary model
                result = self.pipeline(truncated, truncation=True)
                if isinstance(result, list):
                    result = result[0]
                
                raw_label = result.get("label", "NEUTRAL").upper()
                score = result.get("score", 0.0)
                
                if raw_label == "POSITIVE":
                    base_emotion = "joy"
                    polarity = "positive"
                elif raw_label == "NEGATIVE":
                    base_emotion = "sadness"
                    polarity = "negative"
                else:
                    base_emotion = "neutral"
                    polarity = "neutral"
                
                # Refine emotion based on context - returns (emotion, polarity)
                refined_emotion, ctx_polarity = self._refine_emotion(base_emotion, score, truncated)
                # Use context polarity if it differs (more specific)
                if ctx_polarity:
                    polarity = ctx_polarity
                
                return {
                    "label": refined_emotion,
                    "base_emotion": base_emotion,
                    "score": score,
                    "polarity": polarity,
                    "details": {
                        "segment": "last_sentence",
                        "text": truncated[:50],
                        "model": "binary-fallback",
                    }
                }
                
        except Exception as e:
            print(f"⚠️  Multi-emotion analyzer error: {e}")
            return {
                "label": "neutral",
                "score": 0.0,
                "polarity": "neutral",
                "details": {"error": str(e)}
            }
    
    def get_all_emotions(self, text: str) -> Dict[str, float]:
        """Get scores for all detected emotions"""
        result = self.analyze(text)
        return result.get("details", {}).get("all_scores", {})


# Backward compatible SentimentAnalyzer alias
class SentimentAnalyzer(MultiEmotionAnalyzer):
    """Alias for backward compatibility with existing code"""
    pass


if __name__ == "__main__":
    print("=" * 70)
    print("Testing Multi-Emotion Analyzer")
    print("=" * 70)
    
    analyzer = MultiEmotionAnalyzer()
    
    test_cases = [
        # Core emotions
        "I am so happy today!",
        "I feel terrible and sad",
        "I'm really angry about this!",
        "I'm scared of what might happen",
        "That's disgusting!",
        "Wow, I didn't expect that!",
        "The weather is okay",
        
        # Extended emotions
        "I love you so much!",
        "Thank you, I'm so grateful!",
        "I can't wait for tomorrow!",
        "I miss you",
        "LOL that's hilarious",
        "I'm so proud of myself",
        "This is so embarrassing",
        "I'm sorry for your loss",
        
        # Edge cases
        "I'm thinking about it",
        "I feel silly today",
        "I am good",
    ]
    
    print("\nEmotion Predictions:")
    print("-" * 70)
    for text in test_cases:
        result = analyzer.analyze(text)
        emoji = "😊" if result["polarity"] == "positive" else \
                "😢" if result["polarity"] == "negative" else "😐"
        
        print(f"{emoji} '{text[:45]:45}' → {result['label']:15} ({result['base_emotion']}) [{result['score']:.2f}]")
    
    print("\n" + "=" * 70)
    print("✅ Multi-Emotion Analyzer ready!")
    print("=" * 70)