# Bibliotecas
import streamlit as st
import spacy
from spacy import displacy

# Título:
st.title("Web APP LeNER-Br")

# Descrição:
st.write("Este Web App foi criado para proporcionar a comparação dos modelos spaCy [sm](https://huggingface.co/flaviaggp/pt_pipeline) e [lg](https://huggingface.co/flaviaggp/pt_lg_pipeline), voltados para a tarefa de NER com o dataset LeNER-Br.")

# Legenda dos modelos:
st.write("*sm: pipeline mais rápido e menor, mas menos preciso.*")
st.write("*lg: pipeline maior e mais lento, mas mais preciso.*")

#Texto:
input_text = st.text_input('Insira o texto a ser analisado:')

# Função que carrega os modelos:
def load_models():
    sm_model = spacy.load("modelo_lener_sm")
    lg_model = spacy.load("modelo_lener_lg")
    models = {"sm": sm_model, "lg": lg_model}
    return models

models = load_models()
selected_type = st.sidebar.selectbox("Selecione o tipo do modelo", options=["sm", "lg"]) # caixa de seleção do modelo
selected_model = models[selected_type]
doc= selected_model(input_text) # função doc que processa o texto de acordo com a opção escolhida acima

# Cabeçalho
st.header("Visualizador de entidades")

# Cores:
colors = {"LEGISLACAO": "linear-gradient(90deg, #aa9cfc, #fc9ce7)", 'JURISPRUDENCIA': "#ccfbf1", 'LOCAL': "#ffedd5", 'ORGANIZACAO': "#fae8ff", 'PESSOA': "#e0f2fe", 'TEMPO': "#fefde0", }
options = {"ents": ["LEGISLACAO", "JURISPRUDENCIA", "LOCAL", "ORGANIZACAO", "PESSOA", "TEMPO",], "colors": colors}

# Html:
ent_html = displacy.render(doc, style="ent", options=options, jupyter=False) # https://spacy.io/usage/visualizers

st.markdown(ent_html, unsafe_allow_html=True)

# streamlit run render.py