beefing up the emotion server adding distilBERT

requirements.txt

@@ -5,3 +5,4 @@ mcp>=1.0.0
 anthropic>=0.37.0
 python-dotenv>=1.0.0
 plotly>=5.18.0
+transformers>=4.44.0

servers/emotion_server.py

@@ -116,24 +116,55 @@ app = FastMCP("emotion-server")
 # Lexicons & heuristics
 # ---------------------------
 EMO_LEX = {
-    "happy": r"\b(happy|grateful|excited|joy(?:ful)?|delighted|content|optimistic|glad|thrilled|yay|better|good|great|fine)\b",
-    "sad": r"\b(sad|down|depress(?:ed|ing)|cry(?:ing)?|lonely|alone|isolated|upset|miserable|heartbroken|devastat(?:ed|ing)|rejected|abandoned|invisible)\b",
-    "angry": r"\b(angry|mad|furious|irritated|pissed|pissy|annoyed|resentful|rage|hate|infuriat(?:ed|ing)|frustrat(?:ed|ing)|boiling|trapped)\b",
-    "anxious": r"\b(worried|anxious|nervous|stressed|overwhelmed|scared|uneasy|tense|on edge|freaking out|uncertain|afraid)\b",
-    "tired": r"\b(tired|exhaust(?:ed|ing)|drained|burnt(?:\s*out)?|sleepy|fatigued|worn out|depleted|no energy)\b",
-    "love": r"\b(love|affection|caring|fond|admire|cherish|adore)\b",
-    "fear": r"\b(afraid|fear|terrified|panic(?:ky|ked)?|panicked|shaken|petrified)\b",
+    "happy": r"\b(happy|grateful|excited|joy(?:ful)?|delighted|content|optimistic|glad|thrilled|yay|better|good|great|fine|okay-ish|doing ok)\b",
+    "sad": r"\b(sad|down|depress(?:ed|ing)|cry(?:ing)?|lonely|alone|isolated|upset|miserable|heartbroken|devastat(?:ed|ing)|rejected|abandoned|invisible|empty|numb|blah|hopeless|pointless|worthless)\b",
+    "angry": r"\b(angry|mad|furious|irritated|pissed|pissy|annoyed|resentful|rage|hate|infuriat(?:ed|ing)|frustrat(?:ed|ing)|boiling|trapped|fed up|sick of)\b",
+    "anxious": r"\b(worried|anxious|nervous|stressed|overwhelmed|scared|uneasy|tense|on edge|freaking out|uncertain|afraid|dread|panicky|jitters?|heart racing|shaky|knots? in (?:my )?stomach)\b",
+    "tired": r"\b(tired|exhaust(?:ed|ing)|drained|burnt(?:\s*out)?|sleepy|fatigued|worn out|depleted|no energy|burned out|burnt out|overstimulated|overloaded)\b",
+    "love": r"\b(love|affection|caring|fond|admire|cherish|adore|grateful for you)\b",
+    "fear": r"\b(afraid|fear|terrified|panic(?:ky|ked)?|panicked|shaken|petrified|scared stiff)\b",
 }
 
 # Emojis contribute signals even without words
 EMOJI_SIGNAL = {
     "happy": ["😀", "😄", "😊", "🙂", "😁", "🥳", "✨"],
-    "sad": ["😢", "😭", "😞", "😔", "☹️"],
-    "angry": ["😠", "😡", "🤬", "💢"],
-    "anxious": ["😰", "😱", "😬", "😟", "😧"],
-    "tired": ["🥱", "😪", "😴"],
+    "sad": ["😢", "😭", "😞", "😔", "☹️", "😿"],
+    "angry": ["😠", "😡", "🤬", "💢", "😤"],
+    "anxious": ["😰", "😱", "😬", "😟", "😧", "😵‍💫", "🤯"],
+    "tired": ["🥱", "😪", "😴", "😵"],
     "love": ["❤️", "💖", "💕", "😍", "🤍", "💗", "💓", "😘"],
-    "fear": ["🫣", "😨", "😱", "👀"],
+    "fear": ["🫣", "😨", "😱", "👀", "😳"],
+}
+
+# Heuristic valence/arousal for blending ML fallback
+VALENCE_HINT = {
+    "happy": 0.8,
+    "love": 0.7,
+    "sad": -0.75,
+    "angry": -0.6,
+    "anxious": -0.45,
+    "fear": -0.5,
+    "tired": -0.3,
+}
+AROUSAL_HINT = {
+    "happy": 0.55,
+    "love": 0.5,
+    "sad": 0.35,
+    "angry": 0.65,
+    "anxious": 0.65,
+    "fear": 0.7,
+    "tired": 0.2,
+}
+
+_ML_MODEL_ID = "bhadresh-savani/distilbert-base-uncased-emotion"
+_ML_PIPE = None
+_ML_LABEL_MAP = {
+    "joy": ["happy"],
+    "sadness": ["sad"],
+    "anger": ["angry"],
+    "fear": ["anxious", "fear"],
+    "love": ["love"],
+    "surprise": ["happy", "anxious"],
 }
 
 NEGATORS = r"\b(no|not|never|hardly|barely|scarcely|isn['’]t|aren['’]t|can['’]t|don['’]t|doesn['’]t|won['’]t|without)\b"
@@ -233,6 +264,55 @@ def _softmax(d: Dict[str, float]) -> Dict[str, float]:
     return {k: exps[i] / s for i, k in enumerate(d.keys())}
 
 
+def _ensure_ml_pipe():
+    """Lazy-load small fallback classifier for low-confidence texts."""
+    global _ML_PIPE
+    if _ML_PIPE is not None:
+        return _ML_PIPE
+    try:
+        from transformers import pipeline
+
+        _ML_PIPE = pipeline(
+            "text-classification", model=_ML_MODEL_ID, top_k=None, truncation=True
+        )
+    except Exception as e:
+        print(f"[emotion-server] ML fallback unavailable: {type(e).__name__}: {e}")
+        _ML_PIPE = False  # cache failure
+    return _ML_PIPE
+
+
+def _ml_fallback_probs(text: str) -> Optional[Dict[str, float]]:
+    pipe = _ensure_ml_pipe()
+    if not pipe:
+        return None
+    try:
+        out = pipe(text)
+        records = out[0] if out and isinstance(out[0], list) else out
+        scores = {emo: 0.0 for emo in EMO_LEX}
+        max_prob = 0.0
+        for rec in records:
+            lbl = rec.get("label", "").lower()
+            prob = float(rec.get("score", 0.0))
+            max_prob = max(max_prob, prob)
+            for mapped in _ML_LABEL_MAP.get(lbl, []):
+                scores[mapped] += prob
+        total = sum(scores.values())
+        if total <= 0:
+            return None
+        scores = {k: v / total for k, v in scores.items()}
+        scores["__max_prob"] = max_prob
+        return scores
+    except Exception as e:
+        print(f"[emotion-server] ML fallback error: {type(e).__name__}: {e}")
+        return None
+
+
+def _valence_arousal_from_probs(probs: Dict[str, float]) -> Tuple[float, float]:
+    val = sum(VALENCE_HINT.get(k, 0.0) * v for k, v in probs.items())
+    aro = sum(AROUSAL_HINT.get(k, 0.5) * v for k, v in probs.items())
+    return max(-1.0, min(1.0, round(val, 3))), max(0.0, min(1.0, round(aro, 3)))
+
+
 # ---------------------------
 # Per-user calibration (in-memory)
 # ---------------------------
@@ -368,6 +448,24 @@ def _analyze(text: str, user_id: Optional[str] = None) -> dict:
     labels_sorted = sorted(probs.items(), key=lambda kv: kv[1], reverse=True)
     top_labels = [k for k, v in labels_sorted[:3] if v > 0.05] or ["neutral"]
 
+    # Always blend in ML classifier when available (gives broader coverage)
+    ml_scores = _ml_fallback_probs(t)
+    if ml_scores:
+        ml_max = ml_scores.pop("__max_prob", 0.0)
+        combined = {}
+        for emo in emo_scores.keys():
+            # bias toward ML (0.65) while keeping lexicon guardrails (0.35)
+            combined[emo] = 0.35 * probs.get(emo, 0.0) + 0.65 * ml_scores.get(
+                emo, 0.0
+            )
+        probs = _softmax(combined)
+        labels_sorted = sorted(probs.items(), key=lambda kv: kv[1], reverse=True)
+        top_labels = [k for k, v in labels_sorted[:3] if v > 0.05] or ["neutral"]
+        valence, arousal = _valence_arousal_from_probs(probs)
+        tone = quad_tone(valence, arousal)
+        conf = max(conf, round(0.45 + 0.4 * ml_max, 3))
+        reasons.append("ml blend (distilbert)")
+
     return {
         "labels": top_labels,
         "scores": {k: round(v, 3) for k, v in probs.items()},