app/mood_text_model.py

from functools import lru_cache
from typing import Literal

from transformers import pipeline


EmotionLabel = Literal[
    "anger",
    "disgust",
    "fear",
    "joy",
    "neutral",
    "sadness",
    "surprise",
]


@lru_cache
def _get_pipeline():
    return pipeline(
        "text-classification",
        model="j-hartmann/emotion-english-distilroberta-base",
        top_k=None,
        framework="pt",
    )


def analyze_text_emotion(text: str) -> tuple[EmotionLabel, float]:
    """Return (emotion_label, score) for the given text."""
    clf = _get_pipeline()
    outputs = clf(text)[0]
    best = max(outputs, key=lambda x: x["score"])
    label = best["label"].lower()
    score = float(best["score"])
    known_labels: set[str] = {
        "anger",
        "disgust",
        "fear",
        "joy",
        "neutral",
        "sadness",
        "surprise",
    }
    norm_label: EmotionLabel = (label if label in known_labels else "neutral")  # type: ignore[assignment]
    return norm_label, score