Add entity extraction pipeline, closes #203

src/python/txtai/api/routers/__init__.py

@@ -4,6 +4,7 @@ Router imports
 
 from . import caption
 from . import embeddings
+from . import entity
 from . import extractor
 from . import labels
 from . import objects

src/python/txtai/api/routers/entity.py

@@ -0,0 +1,41 @@
+"""
+Defines API paths for entity endpoints.
+"""
+
+from typing import List
+
+from fastapi import APIRouter, Body
+
+from .. import application
+
+router = APIRouter()
+
+
+@router.get("/entity")
+def entity(text: str):
+    """
+    Applies a token classifier to text.
+
+    Args:
+        text: input text
+
+    Returns:
+        list of (entity, entity type, score) per text element
+    """
+
+    return application.get().pipeline("entity", (text,))
+
+
+@router.post("/batchentity")
+def batchentity(texts: List[str] = Body(...)):
+    """
+    Applies a token classifier to text.
+
+    Args:
+        texts: list of text
+
+    Returns:
+        list of (entity, entity type, score) per text element
+    """
+
+    return application.get().pipeline("entity", (texts,))

src/python/txtai/pipeline/text/__init__.py

@@ -2,6 +2,7 @@
 Text imports
 """
 
+from .entity import Entity
 from .extractor import Extractor
 from .labels import Labels
 from .questions import Questions

src/python/txtai/pipeline/text/entity.py

@@ -0,0 +1,41 @@
+"""
+Entity module
+"""
+
+from ..hfpipeline import HFPipeline
+
+
+class Entity(HFPipeline):
+    """
+    Applies a token classifier to text. Extracts known entity/label combinations.
+    """
+
+    def __init__(self, path=None, quantize=False, gpu=True, model=None):
+        super().__init__("token-classification", path, quantize, gpu, model)
+
+    def __call__(self, text, aggregate="simple", workers=0):
+        """
+        Applies a token classifier to text.
+
+        Args:
+            text: text|list
+            aggregate: method to combine multi token entities - options are "simple" (default), "first", "average" or "max"
+            workers: number of concurrent workers to use for processing data, defaults to None
+
+        Returns:
+            list of (entity, entity type, score)
+        """
+
+        # Run token classification pipeline
+        results = self.pipeline(text, aggregation_strategy=aggregate, num_workers=workers)
+
+        # Convert results to a list if necessary
+        if isinstance(text, str):
+            results = [results]
+
+        # Extract (entity, entity type, score) tuples
+        outputs = []
+        for result in results:
+            outputs.append([(r["word"], r["entity_group"], float(r["score"])) for r in result])
+
+        return outputs[0] if isinstance(text, str) else outputs

test/python/testapi/testpipelines.py

@@ -21,6 +21,10 @@ PIPELINES = """
 # Image captions
 caption:
 
+# Entity extraction
+entity:
+    path: dslim/bert-base-NER
+
 # Label settings
 labels:
     path: prajjwal1/bert-medium-mnli
@@ -121,6 +125,22 @@ class TestPipelines(unittest.TestCase):
         captions = self.client.post("batchcaption", json=[path, path]).json()
         self.assertEqual(captions, ["a book shelf filled with many books"] * 2)
 
+    def testEntity(self):
+        """
+        Test entity extraction via API
+        """
+
+        entities = self.client.get(f"entity?text={self.data[1]}").json()
+        self.assertEqual([e[0] for e in entities], ["Canada", "Manhattan"])
+
+    def testEntityBatch(self):
+        """
+        Test batch entity via API
+        """
+
+        entities = self.client.post("batchentity", json=[self.data[1]]).json()
+        self.assertEqual([e[0] for e in entities[0]], ["Canada", "Manhattan"])
+
     def testEmpty(self):
         """
         Test empty API configuration

test/python/testpipeline/testentity.py

@@ -0,0 +1,24 @@
+"""
+Entity module tests
+"""
+
+import unittest
+
+from txtai.pipeline import Entity
+
+
+class TestEntity(unittest.TestCase):
+    """
+    Entity tests.
+    """
+
+    def testEntity(self):
+        """
+        Test entity
+        """
+
+        # Run entity extraction
+        entity = Entity("dslim/bert-base-NER")
+        entities = entity("Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg")
+
+        self.assertEqual([e[0] for e in entities], ["Canada", "Manhattan"])