Prevent pipeline warning when using spaCy

tomaarsen · Oct 31, 2023 · de1d681 · de1d681
1 parent 4e414b1
commit de1d681
Showing 1 changed file with 8 additions and 7 deletions.
diff --git a/span_marker/__init__.py b/span_marker/__init__.py
@@ -18,13 +18,14 @@
 # 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")},
-)
+if "span-marker" not in PIPELINE_REGISTRY.supported_tasks:
+    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: