import gradio as gr
from transformers import pipeline

ner_pipeline = pipeline(task='ner', model='cwchang/ner_model', device=-1)

def ner(text):
    output = ner_pipeline(text, aggregation_strategy="simple")
    return {"text": text, "entities": output}

demo = gr.Interface(
    fn=ner,
    inputs=gr.Textbox(placeholder="Please enter a sentence here..."),
    outputs=gr.HighlightedText(),
    examples=[["Emily Johnson will attend a United Nations meeting in New York on December 5, 2023."],
        ["Microsoft plans to launch its new tech product in Tokyo at the beginning of 2024."],
        ["Professor Michael Brown from the University of Toronto will present his latest research on climate change."],
        ["The famous Christ the Redeemer statue in Rio de Janeiro, Brazil, is scheduled for renovation in July 2023."],
        ["Maria Fernandez will open a new art gallery in Paris in June 2023."],
        ["The 2023 World Economic Forum is set to take place in Davos, Switzerland, expected to attract numerous international leaders."],
        ["Japanese scientist Ichiro Tanaka discovered a new renewable energy technology at the University of Tokyo."],
        ["The International Red Cross will initiate a humanitarian aid project in Nairobi, Kenya, in May 2023."],
        ["The Sydney Opera House in Australia will host a special concert during Christmas 2023."],
        ["The history museum in Berlin, Germany, will organize a World War II exhibition in January 2024."],]
    )

demo.launch(debug=True)