import gradio as gr
import numpy as np
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

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

def np_softmax(x):
    x_exp = torch.exp(torch.tensor(x) - torch.max(torch.tensor(x)))
    f_x = x_exp / x_exp.sum()
    return f_x

def emotion_clf(text):
    model.eval()
    tokens = tokenizer(text, truncation=True, return_tensors="pt")
    tokens.to(model.device)
    preds = model(**tokens)
    prob = np_softmax(preds.logits.cpu().detach().numpy()[0])
    out_dict = {n: p.item() for n, p in zip(labels, prob)}
    return out_dict

tokenizer = AutoTokenizer.from_pretrained("cl-tohoku/bert-base-japanese-whole-word-masking")
model = AutoModelForSequenceClassification.from_pretrained("jingwora/language-emotion-classification-ja", num_labels=8)

title = "Emotion Classification"
description = "Enter Japanese text and get the emotion probabilities."
text = gr.Textbox(placeholder="Enter japanese text", label="text", lines=3)
label = gr.Label(num_top_classes=8)
examples = [
    ["大切な友人が結婚を発表したとき、心が満ち足りた幸せを感じました。"],
    ["愛する祖母が亡くなったとき、胸に深い悲しみが広がりました。"],
    ["新しい仕事を始める前夜、明日からの未知の世界に胸が高鳴りました。"],
    ["山道を歩いていたら、美しい滝が現れ、思わず息を飲む驚きを感じました。"],
    ["公平さを欠いた決定を受けたとき、憤りを抱えて怒りが湧き上がりました。"],
    ["難しい局面で仲間がしっかりと支えてくれて、心からの信頼を感じました。"],
    ["失恋したとき、心に重い悲しみが広がり、何も手につかない状況でした。"],
    ["汚れた場所で長時間過ごしたとき、不快感と嫌悪感が募りました。"],
]

demo = gr.Interface(
    fn=emotion_clf,
    inputs=text,
    outputs=label,
    title=title,
    description=description,
    examples=examples,
    interpretation="default", 
    theme="freddyaboulton/dracula_revamped",
)

demo.launch()