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",
]

def ner(text):
    result = run(tokenizer, model, gazetteers_for_matching, text)
    return result

demo = gr.Interface(ner,
             gr.Textbox(placeholder="Enter sentence here..."), 
             "textbox",
             #gr.HighlightedText(), # TODO https://www.gradio.app/guides/named-entity-recognition
             examples=examples)

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