""" 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, }