eBERTの概要

eBERTはwrimeデータセットに基づく感情予想bertモデルです。
eBERTは東北大学自然言語処理研究グループ様の日本語事前学習済みbertモデルを使用しました。

使用方法

from transformers import BertForSequenceClassification, AutoTokenizer
import numpy as np

emotions = ['喜び', '悲しみ', '期待', '驚き', '怒り', '恐れ', '嫌悪', '信頼']

model = BertForSequenceClassification.from_pretrained("beezza/eBERT")
tokenizer = AutoTokenizer.from_pretrained("beezza/eBERT")

def softmax(x):
    f_x = np.exp(x) / np.sum(np.exp(x))
    return f_x

model.eval()

tokens = tokenizer("テキスト", truncation=True, return_tensors="pt")
tokens.to(model.device)
preds = model(**tokens)
prob = softmax(preds.logits.cpu().detach().numpy()[0])
out_dict = {n: p for n, p in zip(emotions, prob)}

print(out_dict)