| 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() | |