Views
No views yet
distilbert-base-uncased[openness, conscientiousness, extraversion, agreeableness, neuroticism]1from transformers import AutoTokenizer, AutoConfig, AutoModelForSequenceClassification
2import torch, numpy as np
3
4repo_id = "{REPO_ID}"
5tok = AutoTokenizer.from_pretrained(repo_id)
6cfg = AutoConfig.from_pretrained(repo_id)
7model = AutoModelForSequenceClassification.from_pretrained(repo_id, config=cfg).eval()
8
9# single text
10text = "I enjoy planning ahead and keeping things organized."
11inp = tok(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
12with torch.no_grad():
13 out = model(**inp).logits.squeeze(0).tolist()
14OCEAN = ["openness","conscientiousness","extraversion","agreeableness","neuroticism"]
15print(dict(zip(OCEAN, out)))
16# OCEAN: {'openness': 0.6432028412818909, 'conscientiousness': 0.7445886135101318, 'extraversion': 0.20433923602104187, 'agreeableness': 0.4930797815322876, 'neuroticism': 0.33562132716178894}
17
18# mini-batch
19def predict_personality(texts, batch_size=32, max_len=512):
20 preds = []
21 for i in range(0, len(texts), batch_size):
22 b = texts[i:i+batch_size]
23 enc = tok(b, padding=True, truncation=True, max_length=max_len, return_tensors="pt")
24 with torch.no_grad():
25 o = model(**enc).logits
26 preds.append(o.cpu().numpy())
27 return np.vstack(preds)