import gradio as gr
from transformers import AutoTokenizer,AutoModelForCausalLM,BitsAndBytesConfig
import torch
model_id = "rinna/japanese-gpt2-small"
tokenizer = AutoTokenizer.from_pretrained(model_id,trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    low_cpu_mem_usage=True,
    trust_remote_code=True
)
model.eval()
prompt =("川柳：朝寝坊 言い訳考え また寝なおす 川柳：財布見る まるで冷蔵庫 空っぽだ 川柳：ダイエット 明日からって 五年目だ 川柳：マスク取る 顔が誰だか わからない 川柳：ため息が わいふぁいよりも よく飛ぶよ 川柳：")
inputs = tokenizer(prompt,return_tensors="pt")
prompt_len = inputs["input_ids"].shape[-1]  # プロンプトのトークン数！
device = "cuda" if torch.cuda.is_available() else "cpu"
inputs ={k:v.to(device) for k,v in inputs.items()}
def haiku_generate():
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=20,
            eos_token_id=tokenizer.eos_token_id,  # 終わりを明示！
            pad_token_id=tokenizer.pad_token_id,
            do_sample=True,
            repetition_penalty=1.2,
            temperature=0.9,
            top_p=0.9
        )
    generated_tokens = outputs[0][prompt_len:]
    return tokenizer.decode(generated_tokens,skip_special_tokens=True)
with gr.Blocks() as demo:
    output = gr.Textbox(label="俳句")
    btn = gr.Button("生成")
    btn.click(fn=haiku_generate,outputs=output)

demo.launch()