Untitled
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