custom endpoint handler

.gitignore

@@ -0,0 +1,3 @@
+.idea/
+.pytest_cache/
+.venv/

handler.py

@@ -0,0 +1,49 @@
+import spacy
+
+from pathlib import Path
+from typing import Any
+
+
+def _load_spacy_model(name: str | Path) -> spacy.Language:
+    if not hasattr(_load_spacy_model, "nlp"):
+        # pipeline info https://spacy.io/models/en#en_core_web_lg
+        all_except_ner: list[str] = [
+            "tok2vec",
+            "tagger",
+            "parser",
+            "attribute_ruler",
+            "lemmatizer"]
+
+        nlp = spacy.load(name=name, exclude=all_except_ner)
+        _load_spacy_model.nlp = nlp
+        print(f"Loaded {nlp.meta.get('name', 'unknown')} model from {nlp.path}")
+
+    return _load_spacy_model.nlp
+
+
+class EndpointHandler:
+    def __init__(self, name: str | Path = "en_core_web_lg"):
+        self._nlp: spacy.Language = _load_spacy_model(name=name)
+
+    def __call__(self, data: dict[str, Any]) -> dict[str, Any]:
+        inputs: str = data.pop("inputs", "")
+        if not inputs:
+            return {}
+
+        outputs: list[dict[str, Any]] = []
+
+        doc = self._nlp(text=inputs)
+        for ent in doc.ents:
+            if ent.label_ != "PERSON":
+                continue
+
+            entity: dict = {
+                "qid": None,
+                "entity": ent.label_,
+                "text": ent.text,
+                "start": ent.start_char,
+                "end": ent.end_char}
+
+            outputs.append(entity)
+
+        return {"outputs": outputs}

handler_test.py

@@ -0,0 +1,10 @@
+from handler import EndpointHandler
+
+
+class TestEndpointHandler:
+    _handler: EndpointHandler = EndpointHandler()
+
+    def test_endpoint_handler(self):
+        result: dict = self._handler(data={"inputs": "Who is John Doe?"})
+        outputs: list = result.get("outputs", [])
+        assert len(outputs) == 1

requirements-dev.txt

@@ -0,0 +1,2 @@
+-r ./requirements.txt
+pytest

requirements.txt

@@ -0,0 +1,4 @@
+spacy==3.4.0
+
+# trained pipeline
+https://huggingface.co/spacy/en_core_web_lg/resolve/main/en_core_web_lg-any-py3-none-any.whl