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

# トークナイザーとモデルをロードする
tokenizer = AutoTokenizer.from_pretrained("yasunamiura/mr-ja")
model = AutoModelForCausalLM.from_pretrained("yasunamiura/mr-ja")

# txt入力を受け取り、それを使用して新しいtxtを生成する関数
def generate(prompt):
    # 入力テキストをトークナイズ
    inputs = tokenizer.encode(prompt, return_tensors='pt')
    # トークナイズされたテキストを使用して新しいテキストを生成
    outputs = model.generate(inputs, max_length=300, temperature=1.0, top_k=50, top_p=0.95, do_sample=True)
    # 生成されたテキストをデコード
    predicted_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # デコードされたテキストを関数から返す
    return predicted_text

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

# Gradio UIを起動
iface.launch()