Add pipeline support for SpanMarker

CHANGELOG.md

@@ -15,6 +15,12 @@ Types of changes
 * "Security" in case of vulnerabilities.
 -->
 
+## [Unreleased]
+
+### Added
+
+- Added `transformers` `pipeline` support, e.g. `pipeline(task="span-marker", model="tomaarsen/span-marker-mbert-base-multinerd")`.
+
 ## [1.3.0]
 
 ### Added

docs/api/span_marker.pipeline_component.rst

@@ -0,0 +1,14 @@
+
+:autogenerated:
+
+..
+    This file is autogenerated by `sphinx-api`.
+
+span_marker.pipeline_component module
+=====================================
+
+.. currentmodule:: span_marker.pipeline_component
+
+.. automodule:: span_marker.pipeline_component
+    :members:
+    :show-inheritance:

docs/api/span_marker.rst

@@ -19,6 +19,7 @@ span_marker package
        span_marker.modeling
        span_marker.trainer
        span_marker.configuration
+       span_marker.pipeline_component
        span_marker.data_collator
        span_marker.tokenizer
        span_marker.evaluation

span_marker/__init__.py

@@ -5,14 +5,23 @@
 
 import torch
 from transformers import AutoConfig, AutoModel, TrainingArguments
+from transformers.pipelines import PIPELINE_REGISTRY, pipeline
 
 from span_marker.configuration import SpanMarkerConfig
 from span_marker.modeling import SpanMarkerModel
+from span_marker.pipeline_component import SpanMarkerPipeline
 from span_marker.trainer import Trainer
 
 # Set up for Transformers
 AutoConfig.register("span-marker", SpanMarkerConfig)
 AutoModel.register(SpanMarkerConfig, SpanMarkerModel)
+PIPELINE_REGISTRY.register_pipeline(
+    "span-marker",
+    pipeline_class=SpanMarkerPipeline,
+    pt_model=SpanMarkerModel,
+    type="text",
+    default={"pt": ("tomaarsen/span-marker-bert-base-fewnerd-fine-super", "main")},
+)
 
 # Set up for spaCy
 try:

span_marker/pipeline_component.py

@@ -0,0 +1,38 @@
+from typing import Any, Dict, List, Tuple, Union
+
+from transformers import Pipeline
+
+INPUT_TYPES = Union[str, List[str], List[List[str]]]
+OUTPUT_TYPES = Union[List[Dict[str, Union[str, int, float]]], List[List[Dict[str, Union[str, int, float]]]]]
+
+
+class SpanMarkerPipeline(Pipeline):
+    """A Pipeline component for SpanMarker.
+
+    The `pipeline` function is :func:`~transformers.pipeline`, which you can also import with
+    ``from transformers import pipeline``, but you must also import ``span_marker`` to register the
+    ``"span-marker"`` pipeline task.
+
+    Example::
+
+        >>> from span_marker import pipeline
+        >>> pipe = pipeline(task="span-marker", model="tomaarsen/span-marker-mbert-base-multinerd", device_map="auto")
+        >>> pipe("Amelia Earhart flew her single engine Lockheed Vega 5B across the Atlantic to Paris.")
+        [{'span': 'Amelia Earhart', 'label': 'PER', 'score': 0.9999709129333496, 'char_start_index': 0, 'char_end_index': 14},
+         {'span': 'Lockheed Vega 5B', 'label': 'VEHI', 'score': 0.9050095677375793, 'char_start_index': 38, 'char_end_index': 54},
+         {'span': 'Atlantic', 'label': 'LOC', 'score': 0.9991973042488098, 'char_start_index': 66, 'char_end_index': 74},
+         {'span': 'Paris', 'label': 'LOC', 'score': 0.9999232292175293, 'char_start_index': 78, 'char_end_index': 83}]
+
+    """
+
+    def _sanitize_parameters(self, **kwargs) -> Tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any]]:
+        return {}, {}, {}
+
+    def preprocess(self, inputs: INPUT_TYPES) -> INPUT_TYPES:
+        return inputs
+
+    def _forward(self, inputs: INPUT_TYPES) -> OUTPUT_TYPES:
+        return self.model.predict(inputs)
+
+    def postprocess(self, outputs: OUTPUT_TYPES) -> OUTPUT_TYPES:
+        return outputs

tests/test_pipeline.py

@@ -0,0 +1,11 @@
+from transformers import pipeline
+
+import span_marker
+
+
+def test_pipeline() -> None:
+    pipe = pipeline(task="span-marker", model="tomaarsen/span-marker-bert-tiny-fewnerd-coarse-super")
+    outputs = pipe("Tom lives in the Netherlands.")
+    assert len(outputs) == 2
+    assert outputs[0]["span"] == "Tom"
+    assert outputs[1]["span"] == "Netherlands"