model/inference.py

"""
inference.py — Load trained model and predict emotions for new text.
Used by both the FastAPI backend and direct Python usage.
"""

import os
import re
import html
import json
import numpy as np
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel
from huggingface_hub import hf_hub_download

EMOTIONS = [
    "admiration", "amusement", "anger", "annoyance", "approval",
    "caring", "confusion", "curiosity", "desire", "disappointment",
    "disapproval", "disgust", "embarrassment", "excitement", "fear",
    "gratitude", "grief", "joy", "love", "nervousness",
    "optimism", "pride", "realization", "relief", "remorse",
    "sadness", "surprise", "neutral"
]

POSITIVE = {"admiration","amusement","approval","caring","curiosity","desire",
            "excitement","gratitude","joy","love","optimism","pride",
            "realization","relief","surprise"}
NEGATIVE = {"anger","annoyance","disappointment","disapproval","disgust",
            "embarrassment","fear","grief","nervousness","remorse","sadness"}
AMBIGUOUS = {"confusion"}
NEUTRAL_SET = {"neutral"}

EMOTION_COLORS = {
    "admiration": "#FFD700", "amusement": "#FF69B4", "anger": "#FF4444",
    "annoyance": "#FF8C00", "approval": "#90EE90", "caring": "#FFB6C1",
    "confusion": "#DDA0DD", "curiosity": "#87CEEB", "desire": "#FF1493",
    "disappointment": "#708090", "disapproval": "#CD853F", "disgust": "#556B2F",
    "embarrassment": "#FF6347", "excitement": "#FFD700", "fear": "#8B008B",
    "gratitude": "#32CD32", "grief": "#4169E1", "joy": "#FFD700",
    "love": "#FF69B4", "nervousness": "#9370DB", "optimism": "#FFA500",
    "pride": "#DAA520", "realization": "#00CED1", "relief": "#98FB98",
    "remorse": "#808080", "sadness": "#4682B4", "surprise": "#FF8C00",
    "neutral": "#A9A9A9"
}


def clean_text(text: str) -> str:
    text = html.unescape(text)
    text = re.sub(r"^>+\s?", "", text, flags=re.MULTILINE)
    text = re.sub(r"http\S+|www\.\S+", "[URL]", text)
    text = re.sub(r"\[deleted\]|\[removed\]", "", text)
    text = re.sub(r"r/\w+|u/\w+", "", text)
    text = re.sub(r"([!?])\1{2,}", r"\1\1", text)
    text = re.sub(r"(.)\1{2,}", r"\1\1", text)
    text = re.sub(r"[^\x00-\x7F]+", " ", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text


class MeanPooling(nn.Module):
    def forward(self, last_hidden, attention_mask):
        mask   = attention_mask.unsqueeze(-1).float()
        summed = (last_hidden * mask).sum(1)
        count  = mask.sum(1).clamp(min=1e-9)
        return summed / count

class EmotionClassifier(nn.Module):
    def __init__(self, encoder, num_labels=28, dropout=0.1):
        super().__init__()
        self.encoder    = encoder
        hidden          = encoder.config.hidden_size
        self.pool       = MeanPooling()
        self.norm       = nn.LayerNorm(hidden * 2)
        self.drop       = nn.Dropout(dropout)
        self.classifier = nn.Linear(hidden * 2, num_labels)

    def forward(self, input_ids, attention_mask, token_type_ids=None):
        out  = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        cls  = out.last_hidden_state[:, 0, :]
        mean = self.pool(out.last_hidden_state, attention_mask)
        feat = torch.cat([cls, mean], dim=-1)
        return self.classifier(self.drop(self.norm(feat)))


class EmotionPredictor:
    """
    Singleton-style predictor. Call EmotionPredictor.get_instance() to reuse.
    Falls back to SamLowe/roberta-base-go_emotions directly if custom
    trained weights are not found.
    """
    _instance = None

    def __init__(self, model_dir: str = "Harsh-1611/emosense-model",
                 thresholds_path: str = "model/thresholds.npy",
                 device: str = None):
        self.device = torch.device(
            device if device else ("cuda" if torch.cuda.is_available() else "cpu")
        )
        self.model_dir = model_dir
        self.thresholds_path = thresholds_path
        self._load()

    @classmethod
    def get_instance(cls, **kwargs):
        if cls._instance is None:
            cls._instance = cls(**kwargs)
        return cls._instance

    def _load(self):
        from huggingface_hub import hf_hub_download

        try:
            print(f"Loading model from HF Hub: {self.model_dir}")

        # Load encoder + tokenizer
            encoder = AutoModel.from_pretrained(self.model_dir)
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_dir)

        # Build model
            self.model = EmotionClassifier(encoder, num_labels=28).to(self.device)

        # Download head.pt from HF
            head_path = hf_hub_download(
                repo_id=self.model_dir,
                filename="head.pt"
            )

            ckpt = torch.load(head_path, map_location=self.device)

        # Load classifier weights
            if isinstance(ckpt, dict) and "classifier" in ckpt:
                self.model.norm.load_state_dict(ckpt["norm"])
                self.model.drop.load_state_dict(ckpt["drop"])
                self.model.classifier.load_state_dict(ckpt["classifier"])
            else:
                self.model.classifier.load_state_dict(ckpt)

            self.use_pipeline = False

        except Exception as e:
            print("Custom model loading failed. Falling back to pipeline.")
            print(str(e))

            try:
                from transformers import pipeline

                self._pipeline = pipeline(
                    "text-classification",
                    model="SamLowe/roberta-base-go_emotions",
                    top_k=None,
                    device=0 if self.device.type == "cuda" else -1
                )

                self.use_pipeline = True
                self.thresholds = np.full(28, 0.1)

            except Exception as e:
                raise RuntimeError(f"Could not load any model: {e}")

        if not self.use_pipeline:
            # Load thresholds
            from huggingface_hub import hf_hub_download
            try:
                threshold_path=hf_hub_download(
                    repo_id=self.model_dir,
                    filename="thresholds.npy"
                )
                self.thresholds = np.load(threshold_path)
            except:
                self.thresholds = np.full(28, 0.5)
            self.model.eval()

    @torch.no_grad()
    def predict(self, text: str) -> dict:
        """
        Returns a dict with:
          - scores: {emotion: float}  (all 28 probabilities)
          - predictions: list of emotion names above threshold
          - top_emotion: highest scoring emotion
          - sentiment: overall sentiment group
          - sentiment_breakdown: {positive/negative/neutral/ambiguous: float}
        """
        cleaned = clean_text(text)
        if not cleaned:
            return self._empty_result()

        if self.use_pipeline:
            raw = self._pipeline(cleaned)[0]
            scores = {item["label"]: item["score"] for item in raw}
            # Normalize to all 28
            full_scores = {e: scores.get(e, 0.0) for e in EMOTIONS}
            predictions = [e for e, s in full_scores.items() if s >= self.thresholds[EMOTIONS.index(e)]]
        else:
            enc = self.tokenizer(
                cleaned, max_length=128, padding="max_length",
                truncation=True, return_tensors="pt"
            )
            input_ids      = enc["input_ids"].to(self.device)
            attention_mask = enc["attention_mask"].to(self.device)

            logits = self.model(input_ids, attention_mask)
            probs = torch.sigmoid(logits).cpu().numpy()[0]
            full_scores = {e: float(probs[i]) for i, e in enumerate(EMOTIONS)}
            predictions = [
                e for i, e in enumerate(EMOTIONS) if probs[i] >= self.thresholds[i]
            ]

        if not predictions:
            # Always return at least the top emotion
            predictions = [max(full_scores, key=full_scores.get)]

        top_emotion = max(full_scores, key=full_scores.get)

        # Sentiment breakdown (weighted average of scores in each group)
        pos_score = np.mean([full_scores[e] for e in POSITIVE])
        neg_score = np.mean([full_scores[e] for e in NEGATIVE])
        neu_score = full_scores["neutral"]
        amb_score = full_scores["confusion"]

        total = pos_score + neg_score + neu_score + amb_score + 1e-9
        sentiment_breakdown = {
            "positive":  round(pos_score / total, 4),
            "negative":  round(neg_score / total, 4),
            "neutral":   round(neu_score / total, 4),
            "ambiguous": round(amb_score / total, 4),
        }
        dominant_sentiment = max(sentiment_breakdown, key=sentiment_breakdown.get)

        return {
            "text":               cleaned,
            "scores":             {e: round(float(full_scores[e]), 4) for e in EMOTIONS},
            "predictions":        predictions,
            "top_emotion":        top_emotion,
            "sentiment":          dominant_sentiment,
            "sentiment_breakdown": sentiment_breakdown,
            "colors":             {e: EMOTION_COLORS[e] for e in EMOTIONS},
        }

    def _empty_result(self):
        return {
            "text": "",
            "scores": {e: 0.0 for e in EMOTIONS},
            "predictions": ["neutral"],
            "top_emotion": "neutral",
            "sentiment": "neutral",
            "sentiment_breakdown": {"positive": 0, "negative": 0, "neutral": 1, "ambiguous": 0},
            "colors": EMOTION_COLORS,
        }