Minej

/

bert-base-personality

like 36

from transformers import BertTokenizer, BertForSequenceClassification

def personality_detection(text):
tokenizer = BertTokenizer.from_pretrained(pretrained_model_name_or_path = "personality-bert-base")
model = BertForSequenceClassification.from_pretrained(pretrained_model_name_or_path = "personality-bert-base")
inputs = tokenizer(text, truncation=True, padding=True, return_tensors="pt")
outputs = model(**inputs)
predictions = outputs.logits.squeeze().detach().numpy()

label_names = ['Extroversion', 'Neuroticism', 'Agreeableness', 'Conscientiousness', 'Openness']
result = {label_names[i]: predictions[i] for i in range(len(label_names))}

return result

text = "I am feeling excited about the upcoming event"
result = personality_detection(text)
print(result)
{'Extroversion': -0.58424836, 'Neuroticism': 0.10371539, 'Agreeableness': -0.15687147, 'Conscientiousness': -0.78882587, 'Openness': -0.5310378}