import json

import gradio as gr
from paddlenlp import Taskflow

model_name = "uie-base-en"

task_instance = Taskflow("information_extraction", model=model_name)


def analyze(text, schema_name):


    schema = schemas.get(schema_name.split()[0])
    if schema is None:
        return {"error": "Bad schema"}

    task_instance.set_schema(schema)
    # task_instance.set_argument(argument)

    try:
        result = task_instance(text)
    except Exception as e:
        return {"error": e}

    return result



schema_choices = [
    "Entities (people, organizations, locations, date)",
    "Relation1 (person, work for)",
    "Event1 (start job, [date, company, employee])",
    
    # "Event2 ("
]

examples = [
    ["Nicholas is an engineer working at a company called Hugging Face.", schema_choices[0]],
    ["Nicholas is an engineer working at a company called Hugging Face.", schema_choices[1]],
    ["Chris Wigley was appointed CEO of Genomics England, starting in October 2019.", schema_choices[2]],
    ["Chris Wigley was appointed CEO of Genomics England, starting in October 2019.", schema_choices[1]],
    ["Wigley is a former McKinsey executive known for applying machine learning and artificial intelligence technology", schema_choices[0]],
    ]

schemas = {
   # Entity Extraction
        "Entities": ['Person', 'Organization', 'Location'],
        "Relation1": [{"Person":["work for"]}],
        "Event1": [{'start-position': ['date', "company", "employee"]}],
}

with gr.Blocks() as demo:
    gr.Markdown("# Universal Information Extraction (UIE) \n\n \
    For Genomics England Workshop \
    ")

    with gr.Row():
        with gr.Column():
            schema = gr.Dropdown(choices=schema_choices, value=schema_choices[0])
            # schema_setup = gr.Textbox(label="Schema setup", placeholder=json.dumps(schemas[schema]))
            text_input = gr.Textbox(label="Text Input", placeholder=examples[0])
            example_textbox1 = gr.Textbox(visible=False)
            example_textbox2 = gr.Textbox(visible=False)
            gr.Examples(
                examples=examples,
                inputs=[example_textbox1, example_textbox2]
            )
        with gr.Column():
            btn = gr.Button("Run model")
            json_output = gr.JSON(label="JSON Output")


    example_textbox1.change(fn=lambda x, y: (x, y), inputs=[example_textbox1, example_textbox2], outputs=[text_input, schema])
    btn.click(fn=analyze, inputs=[text_input, schema], outputs=json_output)

demo.launch(debug=True, )