configured ner with better accuracy

introduck/inference.py

@@ -1,7 +1,22 @@
 import spacy
 
 
+def _load_spacy_model() -> spacy.Language:
+    spacy_model_name: str = "en_core_web_trf"
+
+    if not hasattr(_load_spacy_model, "nlp"):
+        nlp = spacy.load(name=spacy_model_name)
+        print(f"Loaded {nlp.meta.get('name', 'unknown')} model from {nlp.path}")
+
+        _load_spacy_model.nlp = nlp
+
+    return _load_spacy_model.nlp
+
+
 def extract_contacts_from_email(payload: str) -> list[list[str]]:
+    if not payload:
+        return []
+
     return [["John Doe", "Acme Corp.", "john@example.com", "ai, nn"]]
 
 
@@ -14,18 +29,13 @@ def generate_rejection_reply(payload: str) -> str:
 
 
 def highlight_named_entities(payload: str, labels: list[str] = None) -> dict:
+    if not payload:
+        return {"text": "", "entities": []}
+
     if labels is None:
         labels = ["ORG", "PERSON"]
 
-    nlp: spacy.Language | None = None
-    if not hasattr(highlight_named_entities, "nlp"):
-        nlp = spacy.load(name="en_core_web_sm")
-        highlight_named_entities.nlp = nlp
-        print(f"Loaded {nlp.meta.get('name', 'unknown')} model from {nlp.path}")
-    else:
-        nlp = highlight_named_entities.nlp
-        print(f"Reused {nlp.meta.get('name', 'unknown')} model from {nlp.path}")
-
+    nlp: spacy.Language = _load_spacy_model()
     doc = nlp(payload)
 
     entities: list = []

introduck/routes.py

@@ -7,7 +7,7 @@ from introduck.inference import extract_contacts_from_email
 from introduck.inference import generate_acceptance_reply
 from introduck.inference import generate_rejection_reply
 from introduck.inference import highlight_named_entities
-from introduck.utils import create_email_message
+from introduck.utils import dump_email_as_string
 
 _INTRO_SUBJECT_EXAMPLE: str = "Could you make an intro?"
 _INTRO_MESSAGE_EXAMPLE: str = """\
@@ -41,14 +41,14 @@ def _analyze_message(sender: str, recipients: str, subject: str, body: str):
         "subject": subject,
         "body": body}
 
-    msg: str = create_email_message(data=msg_data)
+    msg: str = dump_email_as_string(data=msg_data)
 
     contacts: list = extract_contacts_from_email(payload=msg)
     acceptance_reply: str = generate_acceptance_reply(payload=msg)
     rejection_reply: str = generate_rejection_reply(payload=msg)
-    text_with_named_entities: dict = highlight_named_entities(payload=msg)
+    highlighted_text: dict = highlight_named_entities(payload=msg)
 
-    return contacts, acceptance_reply, rejection_reply, text_with_named_entities
+    return contacts, acceptance_reply, rejection_reply, highlighted_text
 
 
 def _use_message_template() -> (str, str):
@@ -58,7 +58,12 @@ def _use_message_template() -> (str, str):
 def create_base_route(title: str = "") -> FastAPI:
     # TODO: Fix once resolved https://github.com/gradio-app/gradio/issues/1683
     # TODO: ...and remove elem_id="htxt" from HighlightedText for RF822 output
-    css_workaround: str = "#htxt span {white-space: pre}"
+    css_workaround: str = """
+        #htxt span {
+            font-family: monospace;
+            white-space: pre;
+        }
+        """
 
     base_blocks: gr.Blocks = gr.Blocks(
         analytics_enabled=False,

introduck/utils.py

@@ -1,16 +1,13 @@
-from email.message import EmailMessage
-
-
-def create_email_message(data: dict) -> str:
-    msg: EmailMessage = EmailMessage()
-
-    body_text: str = data.get("body", "")
-    # body_html: str = ""
-
-    msg["From"] = data.get("from", "")
-    msg["To"] = data.get("to", "")
-    msg["Subject"] = data.get("subject", "")
-    msg.set_content(body_text)
-    # msg.add_alternative(body_html, subtype="html")
-
-    return msg.as_string()
+def dump_email_as_string(data: dict) -> str:
+    msg_from: str = data.get("from", "") or "*"
+    msg_to: str = data.get("to", "") or "*"
+    msg_subject: str = data.get("subject", "") or "*"
+    msg_body: str = data.get("body", "") or "***"
+
+    msg: str = ""
+    msg += f"From: {msg_from}\n"
+    msg += f"To: {msg_to}\n"
+    msg += f"Subject: {msg_subject}\n\n"
+    msg += f"{msg_body}\n"
+
+    return msg

requirements.txt

@@ -4,3 +4,4 @@ spacy==3.4.0
 uvicorn[standard]
 
 https://huggingface.co/spacy/en_core_web_sm/resolve/main/en_core_web_sm-any-py3-none-any.whl
+https://huggingface.co/spacy/en_core_web_trf/resolve/main/en_core_web_trf-any-py3-none-any.whl