mentalcare-chatbot / modelGuards /emotionModel.py
mudabbirbhat's picture
Upload folder using huggingface_hub
e2eccc0 verified
from transformers import pipeline
classifier = pipeline(task="text-classification", model="vibhorag101/roberta-base-emotion-prediction-phr", top_k=None,device_map="auto")
# Can classify 28 emotions
def predict_emotion(text):
emotions = classifier(text)[0]
print(emotions)
res = []
for emotion in emotions:
if(emotion["score"] > 0.5):
res.append(emotion["label"])
if(len(res) == 0):
res.append("neutral")
return(res)
if __name__ == "__main__":
print(predict_emotion("My son is so caring."))