Spaces:
Running
Running
File size: 2,940 Bytes
c4fe4b1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | from transformers import pipeline
import numpy as np
# Define the 5 target emotions in alphabetical order
EMOTIONS = ["Bored", "Confident", "Confused", "Curious", "Frustrated"]
# Map GoEmotions category labels to target emotions
ROBERTA_MAPPING = {
"confusion": "Confused",
"curiosity": "Curious",
"annoyance": "Frustrated",
"anger": "Frustrated",
"disapproval": "Frustrated",
"approval": "Confident",
"pride": "Confident",
"admiration": "Confident",
"optimism": "Confident",
"neutral": "Bored",
"disappointment": "Bored",
"sadness": "Bored"
}
class RobertaWrapper:
def __init__(self, model_name="SamLowe/roberta-base-go_emotions"):
print(f"Loading pretrained RoBERTa model: {model_name}...")
self.classifier = pipeline("text-classification", model=model_name, top_k=None)
print("Pretrained RoBERTa model loaded successfully.")
def predict(self, text):
"""
Runs inference on the input text and returns predictions mapped to the 5 target emotions.
"""
if not text.strip():
# Handle empty input
flat_score = 1.0 / len(EMOTIONS)
return {
"primary_emotion": EMOTIONS[0],
"primary_confidence": round(flat_score, 4),
"all_emotions": {e: round(flat_score, 4) for e in EMOTIONS},
"mixed_emotions": []
}
# Run model inference
results = self.classifier(text)[0]
# Initialize target emotions scores
target_scores = {e: 0.0 for e in EMOTIONS}
# Accumulate scores based on mapping
for item in results:
label = item["label"]
score = item["score"]
if label in ROBERTA_MAPPING:
target_emotion = ROBERTA_MAPPING[label]
target_scores[target_emotion] += score
# Clean/normalize scores to be between 0.0 and 1.0 (some sums might slightly exceed 1.0 due to floating point)
for emotion in target_scores:
target_scores[emotion] = min(max(round(float(target_scores[emotion]), 4), 0.0), 1.0)
# Find primary emotion
primary_emotion = max(target_scores, key=target_scores.get)
primary_confidence = target_scores[primary_emotion]
# Mixed emotions: any emotion other than the top one with confidence >= 15%
mixed_emotions = []
for emotion, prob in target_scores.items():
if emotion != primary_emotion and prob >= 0.15:
mixed_emotions.append([emotion, round(prob, 4)])
# Sort mixed emotions descending by score
mixed_emotions.sort(key=lambda x: x[1], reverse=True)
return {
"primary_emotion": primary_emotion,
"primary_confidence": primary_confidence,
"all_emotions": target_scores,
"mixed_emotions": mixed_emotions
}
|