mr-bot / app.py
yasunai's picture
fix context
90ab8fb
# 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()