import time

import gradio as gr

from litie.pipelines import EventExtractionPipeline

pipeline = EventExtractionPipeline(task_model_name="gplinker", model_name_or_path="xusenlin/duee-gplinker")


def extract(text):
    start = time.time()
    res = pipeline(text)
    running_time = time.time() - start

    return running_time, res


demo = gr.Interface(
    extract,
    [
        gr.Textbox(
            placeholder="Enter sentence here...",
            lines=5
        ),
    ],
    [gr.Number(label="Run Time"), gr.Json(label="Result")],
    examples=[
        ["油服巨头哈里伯顿裁员650人 因美国油气开采活动放缓"],
    ],
    title="Event Extraction Demo",
)

demo.launch()