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 = """
{}
"""
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)