chore: avoid overwriting pre-existing entities tomaarsen#17

davidberenstein1957 · Jul 12, 2023 · f34f813 · f34f813
1 parent 91b8f03
commit f34f813
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/span_marker/spacy_integration.py b/span_marker/spacy_integration.py
@@ -4,7 +4,7 @@
 import torch
 from datasets import Dataset
 from spacy.tokens import Doc
-from spacy.util import minibatch
+from spacy.util import minibatch, filter_spans
 import types
 
 from span_marker.modeling import SpanMarkerModel
@@ -104,7 +104,7 @@ def __call__(self, doc: Doc) -> Doc:
                 span.label_ = entity["label"]
                 outputs.append(span)
 
-        doc.set_ents(outputs)
+        doc.set_ents(filter_spans(list(doc.ents) + outputs))
         return doc
 
     def pipe(self, stream, batch_size=128, include_sent=None):
@@ -131,5 +131,5 @@ def pipe(self, stream, batch_size=128, include_sent=None):
                     span = doc[start:end]
                     span.label_ = entity["label"]
                     outputs.append(span)
-                doc.set_ents(outputs)
+                doc.set_ents(filter_spans(list(doc.ents) + outputs))
                 yield doc