from transformers import pipeline
import gradio as gr

model_repository_name = "Kosuke-Yamada/t5_ad_generation"
ad_pipeline = pipeline("text2text-generation", model_repository_name, no_repeat_ngram_size=1)

def main(text):
    return ad_pipeline(text)[0]["generated_text"]

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            gr.Markdown("# 広告生成デモ")

            input_text = gr.components.Textbox(lines=1, label="キーワードを空白区切りで入力して！")
            submit_button = gr.Button("Submit")
            output_text = gr.components.Textbox(lines=1, label="生成広告")

    submit_button.click(main, inputs=[input_text], outputs=[output_text])
demo.launch()