# transformersとtorch、Gradioをインポート
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import numpy as np
import gradio as gr
import os
SECRET_TOKEN = os.getenv("TOKEN_OFFENSIVENESS_ESTIMATION")

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

# txt入力を受け取り、それを使用して新しいtxtを生成する関数
def generate(text):
    # 入力テキストをトークナイズ
    inputs = tokenizer.encode_plus(text, return_tensors='pt')
    # トークナイズされたテキストを使用して攻撃性を推定
    outputs = model(inputs['input_ids'], inputs['attention_mask']).detach().numpy()[0][:3]
    # スコアを算出
    minimum = np.min(outputs)
    if minimum < 0:
        outputs = outputs - minimum
    score = outputs / np.sum(outputs)
    prediction = f'攻撃的でない発言: {score[0]:.1%},\nグレーゾーンの発言: {score[1]:.1%},\n攻撃的な発言: {score[2]:.1%}'
    # デコードされたテキストを関数から返す
    return prediction

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

# Gradio UIを起動
iface.launch(share=True)