avatar/sentiment_transformer.py

"""
Sentiment Transformer - Fast neural sentiment/emotion detection

Uses DistilBERT + emotion classification head for real-time emotion detection
No keywords required — pure transformer-based analysis
"""

import torch
from transformers import pipeline
from typing import Dict, Any
import warnings

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


class SentimentAnalyzer:
    """
    Transformer-based sentiment analyzer using DistilBERT
    
    Fast, accurate emotion detection without keyword hardcoding
    Supports 6+ emotion classes: joy, sadness, anger, fear, surprise, neutral
    """
    
    # Emotion label mappings from HF model outputs
    EMOTION_LABEL_MAP = {
        "POSITIVE": "happiness",
        "NEGATIVE": "sadness",
        "NEUTRAL": "neutral",
        # Extended emotion labels if using multi-class model
        "joy": "joy",
        "sadness": "sadness",
        "anger": "anger",
        "fear": "fear",
        "surprise": "surprise",
        "disgust": "disgust",
    }
    
    # Greeting patterns that should be neutral (not sad/negative)
    GREETING_PATTERNS = [
        "how are you", "how're you", "how r u", "how r you",
        "how is it going", "how's it going", "hows it going",
        "what's up", "whats up", "wassup", "sup",
        "how do you do", "how ya doing", "how you doing",
        "what is up", "what is going on", "what's going on",
        "how have you been", "how've you been",
        "are you okay", "are you ok", "you okay", "you ok",
        "how was your day", "how's your day",
        "how do you feel", "how are things",
    ]
    
    def __init__(self, model_name: str = "distilbert-base-uncased-finetuned-sst-2-english"):
        """
        Initialize transformer-based sentiment classifier
        
        Args:
            model_name: HuggingFace model identifier
                - distilbert-base-uncased-finetuned-sst-2-english (3-class: positive/negative/neutral)
                - Use local cache to avoid repeated downloads
        """
        self.model_name = model_name
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        
        print(f"📊 Loading sentiment model: {model_name}")
        print(f"   Device: {self.device}")
        
        # Initialize the pipeline
        self.pipeline = pipeline(
            "sentiment-analysis",
            model=model_name,
            device=0 if self.device == "cuda" else -1,
        )
        
        print(f"✓ Model loaded and ready")
    
    def _get_last_sentence(self, text: str) -> str:
        """Extract last sentence for real-time accuracy"""
        import re
        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 _is_greeting_or_question(self, text: str) -> bool:
        """Check if text is a common greeting/question that should be neutral"""
        text_lower = text.lower().strip()
        # Remove punctuation for matching
        text_clean = text_lower.rstrip('?!.')
        
        for pattern in self.GREETING_PATTERNS:
            if pattern in text_clean:
                return True
        return False
    
    def analyze(self, text: str) -> Dict[str, Any]:
        """
        Analyze text and return detected emotion
        
        Focuses on LAST SENTENCE for real-time updates
        
        Args:
            text: Input text to analyze
            
        Returns:
            Dict with 'label' (emotion), 'score', and 'details'
        """
        if not text or not text.strip():
            return {
                "label": "neutral",
                "score": 0.0,
                "details": {"segment": "empty"}
            }
        
        # Use last sentence for real-time accuracy
        last_sentence = self._get_last_sentence(text)
        
        # Check for greetings/common questions first - should be neutral
        if self._is_greeting_or_question(last_sentence):
            return {
                "label": "neutral",
                "score": 0.85,
                "details": {
                    "segment": "greeting",
                    "text": last_sentence[:50],
                    "model": "greeting_override",
                }
            }
        
        # Truncate to max 512 tokens (BERT limit) for efficiency
        truncated = last_sentence[:512]
        
        try:
            # Get prediction from transformer
            result = self.pipeline(truncated, truncation=True)
            
            if isinstance(result, list):
                result = result[0]
            
            # Extract label and score
            raw_label = result.get("label", "NEUTRAL").upper()
            raw_score = result.get("score", 0.0)
            
            # Map to emotion label
            emotion_label = self.EMOTION_LABEL_MAP.get(raw_label, "neutral").lower()
            
            return {
                "label": emotion_label,
                "score": min(raw_score, 1.0),
                "details": {
                    "segment": "last_sentence",
                    "text": truncated[:50],
                    "model": self.model_name,
                }
            }
        except Exception as e:
            print(f"⚠️  Transformer error: {e}")
            return {
                "label": "neutral",
                "score": 0.0,
                "details": {"error": str(e)}
            }


if __name__ == "__main__":
    print("=" * 70)
    print("Testing Sentiment Transformer Analyzer")
    print("=" * 70)
    
    analyzer = SentimentAnalyzer()
    
    test_cases = [
        "I am so happy today!",
        "I am good",
        "I'm okay",
        "I love this!",
        "This is exciting!",
        "I am really sad",
        "This makes me angry",
        "I am scared",
        "I am confused",
        "I miss you",
        "The weather is nice",
        "This is terrible",
        "I don't know what to think",
        "Absolutely amazing experience!",
        "Completely disappointed with this product",
    ]
    
    print("\nSentiment Predictions:")
    print("-" * 70)
    for text in test_cases:
        result = analyzer.analyze(text)
        emoji = "😊" if "happiness" in result["label"] or result["label"] == "positive" else \
                "😢" if "sadness" in result["label"] else \
                "😠" if result["label"] == "anger" else \
                "😨" if result["label"] == "fear" else \
                "😐"
        
        print(f"{emoji} '{text[:50]:50}' → {result['label']:15} ({result['score']:.2f})")
    
    print("\n" + "=" * 70)
    print("✅ Sentiment Transformer ready!")
    print("=" * 70)