from transformers import pipeline
import gradio as gr
import os

def auth(user_name, password):
    if user_name == os.getenv("USER_NAME") and password == os.getenv("PASSWORD"):
        return True
    else:
        return False

hf_token = os.getenv("HF_TOKEN")

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

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

with gr.Blocks() as demo:
    with gr.Column():
        gr.Markdown("# 広告生成デモ")
        with gr.Row():
            with gr.Column():
                input_text = gr.components.Textbox(lines=1, label="キーワードを空白区切りで入力して！")
                submit_button = gr.Button("Submit")
            with gr.Column():
                output_text = gr.components.Textbox(lines=1, label="生成広告")

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