Untitled

 avatar
unknown
plain_text
5 months ago
1.6 kB
1
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