ppp57420/ocean-personality-distilbert

DistilBERT for OCEAN Personality Regression

Base: distilbert-base-uncased
Head: 5-dim regression.
Order: [openness, conscientiousness, extraversion, agreeableness, neuroticism]
Targets normalized to [0,1]. Result: metrics: mae: 0.108, rmse: 0.1762, pearson: 0.7688, spearman: 0.7498

Usage

from transformers import AutoTokenizer, AutoConfig, AutoModelForSequenceClassification
import torch, numpy as np

repo_id = "{REPO_ID}"
tok = AutoTokenizer.from_pretrained(repo_id)
cfg = AutoConfig.from_pretrained(repo_id)
model = AutoModelForSequenceClassification.from_pretrained(repo_id, config=cfg).eval()

# single text
text = "I enjoy planning ahead and keeping things organized."
inp = tok(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
with torch.no_grad():
    out = model(**inp).logits.squeeze(0).tolist()
OCEAN = ["openness","conscientiousness","extraversion","agreeableness","neuroticism"]
print(dict(zip(OCEAN, out)))
# OCEAN: {'openness': 0.6432028412818909, 'conscientiousness': 0.7445886135101318, 'extraversion': 0.20433923602104187, 'agreeableness': 0.4930797815322876, 'neuroticism': 0.33562132716178894}

# mini-batch
def predict_personality(texts, batch_size=32, max_len=512):
    preds = []
    for i in range(0, len(texts), batch_size):
        b = texts[i:i+batch_size]
        enc = tok(b, padding=True, truncation=True, max_length=max_len, return_tensors="pt")
        with torch.no_grad():
            o = model(**enc).logits
        preds.append(o.cpu().numpy())
    return np.vstack(preds)