Spaces:
Sleeping
Sleeping
| # -*- 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) | |