remote inference api

introduck/inference.py

@@ -1,47 +1,42 @@
-import pandas as pd
-import spacy
+import json
+import os
+import requests
 
+from typing import Any
 
-def _load_spacy_model() -> spacy.Language:
-    spacy_model_name: str = "en_core_web_sm"
 
-    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_text(payload: str) -> dict:
+def extract_contacts_from_text(payload: str) -> list[dict[str, Any]]:
     if not payload:
-        return {}
+        return []
 
-    nlp: spacy.Language = _load_spacy_model()
-    doc = nlp(payload)
+    api_token: str = os.environ.get("HUGGINGFACE_INFERENCE_ENDPOINT_TOKEN", "")
+    if not api_token:
+        return []
 
-    contacts_cols: list[str] = ["name", "email"]
-    contacts: pd.DataFrame = pd.DataFrame(columns=contacts_cols)
+    api_url: str = os.environ.get("HUGGINGFACE_INFERENCE_ENDPOINT_URL", "")
+    if not api_url:
+        return []
 
-    entities: list[dict] = []
-    for ent in doc.ents:
-        if ent.label_ != "PERSON":
-            continue
+    response: requests.Response = requests.post(
+        url=api_url,
+        headers={
+            "Authorization": f"Bearer {api_token}",
+            "Content-Type": "application/json"},
+        data=json.dumps({"inputs": payload}))
 
+    outputs: list[dict[str, Any]] = response.json()
+    if not outputs:
+        return []
+
+    entities: list[dict[str, Any]] = []
+    for output in outputs:
         entity: dict = {
-            "entity": ent.label_,
-            "start": ent.start_char,
-            "end": ent.end_char}
+            "text": output.get("word", ""),
+            "score": output.get("score", -1),
+            "entity": output.get("entity_group", "UNKNOWN"),
+            "start": output.get("start", -1),
+            "end": output.get("end", -1)}
 
         entities.append(entity)
 
-        contact: dict = {
-            "name": ent.text,
-            "email": "unknown"}
-
-        contacts = pd.concat(
-            [contacts, pd.DataFrame([contact])],
-            ignore_index=True)
-
-    return {"contacts": contacts, "entities": entities, "payload": payload}
+    return entities

introduck/routes.py

@@ -4,9 +4,9 @@ import pandas as pd
 from fastapi import FastAPI
 from gradio.routes import App as GradioApp
 from introduck.inference import extract_contacts_from_text
-from introduck.utils import dump_email_to_string
 from introduck.utils import validate_email
 from introduck.utils import validate_multiple_emails
+from typing import Any
 
 _INTRO_SUBJECT_EXAMPLE: str = "Could you make an intro?"
 _INTRO_MESSAGE_EXAMPLE: str = """\
@@ -59,22 +59,15 @@ def _analyze_message(
     if not body:
         return default_outputs
 
-    msg_data: dict = {
-        "from": sender,
-        "to": recipients,
-        "subject": subject,
-        "body": body}
+    payload: str = ""
+    payload += f"From: {sender or '*'}\n"
+    payload += f"To: {recipients or '*'}\n"
+    payload += f"Subject: {subject or '*'}\n\n"
+    payload += f"{body or '***'}\n"
 
-    msg: str = dump_email_to_string(data=msg_data)
+    outputs: list[dict[str, Any]] = extract_contacts_from_text(payload=payload)
 
-    output: dict = extract_contacts_from_text(payload=msg)
-
-    contacts: pd.DataFrame = output.get("contacts", pd.DataFrame())
-    highlighted_text: dict = {
-        "text": output.get("payload", str),
-        "entities": output.get("entities", [])}
-
-    return contacts, highlighted_text
+    return pd.DataFrame(outputs), {"text": payload, "entities": outputs}
 
 
 def _use_message_template() -> (str, str):

introduck/utils.py

@@ -2,21 +2,6 @@ from email_validator import EmailNotValidError
 from email_validator import validate_email as validate_email_address
 
 
-def dump_email_to_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
-
-
 def validate_email(email: str) -> str:
     if not email:
         return ""

requirements.txt

@@ -2,8 +2,5 @@ email-validator==1.3.0
 fastapi[all]
 gradio==3.4.0  # keep in sync with version from README.md metadata
 pandas==1.5.0
-spacy==3.4.0
+requests==2.28.1
 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