import spacy
import streamlit as st

# from transformers import AutoModelForTokenClassification, AutoTokenizer, pipeline


def render_entities(entities):
    colors = {"LOCATION": "#5cff84"}
    options = {"ents": ["LOCATION"], "colors": colors}
    html = spacy.displacy.render(entities, style="ent", options=options, manual=True)
    html = html.replace("\n", " ")

    return html


HTML_WRAPPER = """<div style="overflow-x: auto; border: 1px solid #e6e9ef; border-radius: 0.25rem; padding: 1rem; margin-bottom: 2.5rem">{}</div>"""

st.header("Location Entity Recognition Demo 🔎🌆🌍")
threshold = st.sidebar.slider("Threshold", value=0.5, min_value=0.0, max_value=1.0)
display_probabilities = st.sidebar.checkbox("Display probabilities")


text = st.text_area("Text input", value="This text is about Malaria", height=400)

nlp = spacy.load("en_core_web_trf")

doc = nlp(text)

ents = [
    {"start": ent.start_char, "end": ent.end_char, "label": "LOCATION"}
    for ent in doc.ents
]
foo = {"text": text, "ents": ents}


print(ents)
print(doc.ents)

html = render_entities(foo)
st.write(HTML_WRAPPER.format(html), unsafe_allow_html=True)