import gradio as gr
from website_script import load, run

tokenizer, model, gazetteers_for_matching = load()

examples = [
    ["Masarykova univerzita se nachází v Brně.", None],
    ["Barack Obama navštívil Prahu minulý týden.", None],
    ["Angela Merkelová se setkala s francouzským prezidentem v Paříži.", None],
    ["Nobelova cena za fyziku byla udělena týmu vědců z MIT.", None],
    ["Eiffelova věž je ikonickou památkou v Paříži.", None],
    ["Bill Gates, spoluzakladatel společnosti Microsoft, oznámil nový grant pro výzkum umělé inteligence.", None],
    ["Britská královna Alžběta II. navštívila Kanadu v rámci svého posledního zahraničního turné, během kterého zdůraznila důležitost spolupráce a přátelství mezi oběma národy.", None],
    ["Francouzský prezident Emmanuel Macron oznámil nový plán na podporu start-upů a inovací ve Francii, který zahrnuje investice ve výši několika miliard eur.", None],
    ["Světová zdravotnická organizace spustila nový program na boj proti malárii v subsaharské Africe, který zahrnuje rozdělování sítí proti komárům a očkování milionů lidí.", None]
 ]


def ner(text, file_names):
    text = text.replace(".", " .")
    result = run(tokenizer, model, gazetteers_for_matching, text, file_names)
    return {"text": text, "entities": result}

with gr.Blocks(css="./style.css", theme=gr.themes.Default(primary_hue="blue", secondary_hue="sky")) as demo:
    gr.Interface(ner,
        gr.Textbox(lines=5, placeholder="Enter sentence here..."), 
        gr.HighlightedText(show_legend=True, color_map={"PER": "#f7a7a3", "ORG": "#77fc6a", "LOC": "#87CEFF"}),
        examples=examples,
        title="NerROB-czech",
        description="This is an implementation of a Named Entity Recognition model for the Czech language using gazetteers.",
        allow_flagging="never",
        additional_inputs=gr.File(label="Upload a JSON file containing gazetteers", file_count="multiple", file_types=[".json"]),
        )

if __name__ == "__main__":
    demo.launch()