# transformersとtorch、Gradioをインポート
from transformers import AutoTokenizer, AutoModel
import torch
import numpy as np
import gradio as gr

# トークナイザーとモデルをロードする
tokenizer = AutoTokenizer.from_pretrained("studio-ousia/luke-japanese-base-lite")
model = AutoModel.from_pretrained("TomokiFujihara/luke-japanese-base-lite-offensiveness-estimation")

# txt入力を受け取り、それを使用して新しいtxtを生成する関数
def generate(text):
    # 入力テキストをトークナイズ
    inputs = tokenizer.encode_plus(text, return_tensors='pt')
    # トークナイズされたテキストを使用して攻撃性を推定
    outputs = np.array(model(**inputs)[:3])
    # スコアを算出
    min = min(outputs)
    if min < 0:
        outputs = outputs - min
    score = outputs / np.sum(outputs)
    prediction = f'Not Offensive: {outputs[0]}, Gray-area: {outputs[1]}, Offensive: {outputs[1]}'
    # デコードされたテキストを関数から返す
    return prediction

# Gradio UIを作成
iface = gr.Interface(
    # 上記で定義したテキスト生成関数を使用
    generate, 
    # 入力としてテキストボックスを使用
    inputs = gr.Textbox(label = "Input a text", default = "攻撃性を評価したいコメントを入力してください."), 
    # 出力はテキスト形式
    outputs="text", 
    # UIのタイトル設定
    title = "Mr Bot🤖")

# Gradio UIを起動
iface.launch()