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"):
            input = gr.Textbox(label="Input", lines=1, placeholder="Mask triple input")
            origin_button = gr.Button("Origin")
            origin_output = gr.Textbox(label="Before Edit", lines=2, placeholder="")

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

            edit_output = gr.Textbox(label="After Edit", lines=2, placeholder="")
            gr.Examples(
                examples=["[MASK] r t", "h"],
                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="")


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

demo.launch()