import gradio as gr

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

def add_process(title, context, img):
    return f"Title:{title}\nContext:{context}\n...{img}", 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=3, 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=3, placeholder="")
            gr.Examples(
                examples=[["[MASK] r1 t1", "h1"], ["[MASK] r2 t2", "h2"]],
                inputs=[input, alter_label],
                outputs=[origin_output, edit_output],
                fn=edit_process,
                cache_examples=True,
            )

        with gr.TabItem("A-FB15k237"):
            with gr.Row():
                with gr.Column():
                    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")

                with gr.Column():
                    add_origin_output = gr.Textbox(label="Origin Results", lines=3, placeholder="")
                    add_output = gr.Textbox(label="Add Results", lines=3, placeholder="")

            gr.Examples(
                examples=[["h1 r1 t1", "1"], ["h2 r2 t2", "1"]],
                inputs=[input, alter_label],
                outputs=[add_origin_output, add_output],
                fn=add_process,
                cache_examples=True,
            )

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

demo.launch()