import gradio as gr

def origin_preditcion(title, context):
    return f"Title:{title}\nContext:{context}\n..."

def edit_process(title, context):
    return f"Title:{title}\nContext:{context}\n..."

def add_process(title, context, img):
    return f"Title:{title}\nContext:{context}\n...{img}"

with gr.Blocks() as demo:
    gr.Markdown("# KGE Editing")

    # 多个tab
    with gr.Tabs():

        with gr.TabItem("E-FB15k237"):
            with gr.Row():
                with gr.Column():
                    input = gr.Textbox(label="Input", lines=1, placeholder="Mask triple input")
                    origin_button = gr.Button("Origin")

                with gr.Column():
                    origin_output = gr.Textbox(label="Before Edit", lines=2, placeholder="")

            with gr.Row():
                with gr.Column():
                    alter_label = gr.Textbox(label="Alter Entity", lines=1, placeholder="Entity Name")            
                    edit_button = gr.Button("Edit", elem_id="warning")

                with gr.Column():
                    edit_output = gr.Textbox(label="After Edit", lines=2, placeholder="")
            gr.Examples(
                examples=[["[MASK] r1 t1", "h1"], ["[MASK] r2 t2", "h2"]],
                inputs=[input, alter_label],
                outputs=edit_output,
                fn=edit_process,
                cache_examples=True,
            )

        with gr.TabItem("A-FB15k237"):
            input = gr.Textbox(label="Input", lines=1, placeholder="New triple input")

            alter_label = gr.Textbox(label="Head/Tail", lines=1, placeholder="1:head / 0:tail")
            add_button = gr.Button("Add")

            add_output = gr.Textbox(label="Add Results", lines=2, placeholder="")

    origin_button.click(fn=origin_preditcion, inputs=[input, alter_label], outputs=origin_output)
    edit_button.click(fn=edit_process, inputs=[input, alter_label], outputs=edit_output)
    add_button.click(fn=add_process, inputs=[input, alter_label], outputs=add_output)

demo.launch()