import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

REPO = "YMmim/klue-bert-emotion-6"
labels = ["기쁨", "슬픔", "분노", "불안", "당황", "상처"]
tok = AutoTokenizer.from_pretrained(REPO)
model = AutoModelForSequenceClassification.from_pretrained(REPO)
model.eval()

x = tok("시험 결과가 걱정돼서 잠이 안 와.", return_tensors="pt", truncation=True, max_length=128)
with torch.no_grad():
    print(labels[int(model(**x).logits.argmax())])

한계

감정 라벨이 발화 단위가 아니라 대화 단위로 부여된 데이터 특성상, 한 대화 내 모든 화자 발화에 동일 감정이 적용되어 성능 상한에 제약이 있습니다.

출처

AI Hub 「공감형 대화」 데이터셋 기반. 비상업적 용도로 사용하세요.