import gradio as gr
from transformers import T5Tokenizer,AutoModelForCausalLM

tokenizer = T5Tokenizer.from_pretrained("rinna/japanese-gpt2-small")
model3 = AutoModelForCausalLM.from_pretrained("./models")
model3.to("cpu")
descriptions = "文書を入力すると、その後に、サザンオールスターズ風の歌詞を生成します"

def getarate_sentences3(seed_sentence):
    x = tokenizer.encode(seed_sentence, return_tensors="pt", add_special_tokens=False)  # 入力
    x = x.cpu()  # CPU対応
    y = model3.generate(x, #　入力
                       min_length=50,  # 文章の最小長
                       max_length=100,  # 文章の最大長
                       do_sample=True,   # 次の単語を確率で選ぶ
                       top_k=50, # Top-Kサンプリング
                       top_p=0.95,  # Top-pサンプリング
                       temperature=1.2,  # 確率分布の調整
                       num_return_sequences=1,  # 生成する文章の数
                       pad_token_id=tokenizer.pad_token_id,  # パディングのトークンID
                       bos_token_id=tokenizer.bos_token_id,  # テキスト先頭のトークンID
                       eos_token_id=tokenizer.eos_token_id,  # テキスト終端のトークンID
                       bad_word_ids=[[tokenizer.unk_token_id]]  # 生成が許可されないトークンID
                       )  
    generated_sentences = tokenizer.batch_decode(y, skip_special_tokens=True)  # 特殊トークンをスキップして文章に変換
    return generated_sentences[0]

demo = gr.Interface(fn=getarate_sentences3, inputs=gr.Textbox(lines=3, placeholder="文章を入力してください"), outputs="text", title="Southern All Stars style lyrics", description=descriptions)


demo.launch()