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}

demo = gr.Interface(ner,
             gr.Textbox(placeholder="Enter sentence here..."), 
             gr.HighlightedText(),
             examples=examples)

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