Untitled
unknown
plain_text
a year ago
1.6 kB
3
Indexable
# emotion_detection.py
import torch
from transformers import BertTokenizer, BertForSequenceClassification
class EmotionDetection:
def __init__(self):
"""
Initializes the EmotionDetection class by loading the pre-trained BERT model
and its corresponding tokenizer for emotion detection.
"""
self.tokenizer = BertTokenizer.from_pretrained('bhadresh-savani/bert-base-uncased-emotion')
self.model = BertForSequenceClassification.from_pretrained('bhadresh-savani/bert-base-uncased-emotion')
print("Emotion Detection Model Loaded Successfully!")
def detect_emotion(self, text):
"""
Detects the emotion from the provided text input.
Args:
text (str): The input text from the user.
Returns:
str: The detected emotion label.
"""
# Tokenizing the input text
inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
# Running the model to get logits
with torch.no_grad():
logits = self.model(**inputs).logits
# Getting the predicted class index
predicted_class = logits.argmax().item()
# Returning the corresponding emotion label
return self.model.config.id2label[predicted_class]
# Example Usage
if __name__ == "__main__":
emotion_detector = EmotionDetection()
sample_text = "I'm feeling very happy today!"
emotion = emotion_detector.detect_emotion(sample_text)
print(f"Detected Emotion: {emotion}")
Editor is loading...
Leave a Comment