Update app.py

app.py

@@ -14,6 +14,7 @@ model.eval()
 prompt =("川柳：朝寝坊 言い訳考え また寝なおす 川柳：財布見る まるで冷蔵庫 空っぽだ 川柳：ダイエット 明日からって 五年目だ 川柳：マスク取る 顔が誰だか わからない 川柳：ため息が わいふぁいよりも よく飛ぶよ 川柳：")
 def haiku_generate():
     inputs = tokenizer(prompt,return_tensors="pt")
+    prompt_len = inputs["input_ids"].shape[-1]  # プロンプトのトークン数！
     if not torch.cuda.is_available():
         inputs ={k:v.to("cpu") for k,v in inputs.items()}
     with torch.no_grad():
@@ -27,15 +28,11 @@ def haiku_generate():
             temperature=1.3,
             top_p=0.9
         )
-    print(tokenizer.decode(outputs[0],skip_special_tokens=True))
-    return tokenizer.decode(outputs[0],skip_special_tokens=True)
-def output_extracter():
-    output_haiku = haiku_generate()
-    generated_haiku = output_haiku.replace(prompt,"")
-    return generated_haiku
+    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=output_extracter,outputs=output)
+    btn.click(fn=haiku_generate,outputs=output)
 
 demo.launch()