import json
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]
]


def ner(text, file_names):
    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=10, placeholder="Enter sentence here..."), 
        # gr.HighlightedText(show_legend=True, color_map={"PER": "red", "ORG": "green", "LOC": "blue"}),
        gr.HighlightedText(show_legend=True, color_map={"PER": "#c40e0e", "ORG": "#008a20", "LOC": "#0432d9"}, elem_id="highlighted_text"),
        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()