import gradio as gr
from website_script import load, run

print("Loading model")
tokenizer, model, gazetteers_for_matching = load()
print("Loaded model")

examples = [
    "Masarykova univerzita se nachází v Brně .",
    "Barack Obama navštívil Prahu minulý týden .",
    "Angela Merkelová se setkala s francouzským prezidentem v Paříži .",
    "Nobelova cena za fyziku byla udělena týmu vědců z MIT ."
]

def ner(text):
    result = run(tokenizer, model, gazetteers_for_matching, text)
    return {"text": text, "entities": result}
with gr.Blocks(css="footer{display:none !important}", theme=gr.themes.Default(primary_hue="blue", secondary_hue="sky")) as demo:
# with gr.Blocks(theme=gr.themes.Soft()) as demo:

    gr.Interface(ner,
        gr.Textbox(lines=5, placeholder="Enter sentence here..."), 
        gr.HighlightedText(show_legend=True, color_map={"PER": "red", "ORG": "green", "LOC": "blue"}),
        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")

    gr.Interface(ner,
        gr.File(label="Upload a JSON file"),
        None,
        allow_flagging="never",
        description="Here you can upload your own gazetteers.",        
)

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