feature: add ner app

config.py

@@ -0,0 +1,9 @@
+from pydantic_settings import BaseSettings
+
+class Settings(BaseSettings):
+    TASK: str = "ner"
+    MODEL_NAME: str = "dslim/bert-base-NER"
+    TITLE: str = 'Named Entity Recog with'
+
+
+settings = Settings()

const.py

@@ -0,0 +1,6 @@
+COLORS = {
+    "LOC": "#F67DE3",   # Light pink
+    "ORG": "#7DF6D9",   # Light teal
+    "PER": "#F6E37D",   # Light yellow
+    "MISC": "#7D9BF6"   # Light blue
+}

main.py

@@ -0,0 +1,20 @@
+import streamlit as st
+from spacy import displacy
+from config import settings
+from const import COLORS
+from utils import init_model, custom_predict
+
+
+def main():
+    st.title("Entity Checker")
+    raw_text = st.text_area("Enter Text Here", "Type Here")
+    if st.button("Analyze"):
+        pipe = init_model(settings.TASK, 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()

requirements.txt

@@ -0,0 +1,4 @@
+streamlit>=1.40.2
+spacy>=spacy-3.8.2
+pydantic_settings>=2.6.1
+transformers>=4.46.3

utils.py

@@ -0,0 +1,23 @@
+from functools import lru_cache
+
+from transformers import pipeline, Pipeline
+
+
+@lru_cache
+def init_model( task: str, model: str = None,  aggregation_strategy: str = None) -> Pipeline:
+    ner_pipeline = pipeline(
+        task, model=model, aggregation_strategy=aggregation_strategy
+    )
+    return ner_pipeline
+
+
+def custom_predict(text: str, pipe: str):
+    result = pipe(text, aggregation_strategy="simple")
+    ents = [
+        {"start": dic['start'],
+         "end": dic['end'],
+         "label": dic['entity_group']}
+        for dic in result]
+    return {"text": text,
+            "ents": ents,
+            "title": None}