import streamlit as st
from spacy import displacy
from config import settings
from const import COLORS
from utils import init_model, custom_predict
from transformers import AutoTokenizer, AutoModelForTokenClassification


def main():
    st.title("Entity Checker")
    st.title("👋")
    raw_text = st.text_area("Enter Text Here", "Type Here")
    if st.button("Analyze"):
        pipe = init_model(settings.TASK, settings.MODEL_NAME)
        tokenizer = AutoTokenizer.from_pretrained(settings.MODEL_NAME)
        model = AutoModelForTokenClassification.from_pretrained(
            settings.MODEL_NAME
        )
        result = custom_predict(raw_text, pipe)
        st.subheader(f"{settings.TITLE} {settings.MODEL_NAME}")
        options = {"ents": ["LOC", "ORG", "PER", "MISC"], "colors": COLORS}
        ent_html = displacy.render(result, style="ent", manual=True, options=options)
        st.markdown(ent_html, unsafe_allow_html=True)

if __name__ == '__main__':
    main()