diff --git a/spacy/tests/doc/test_add_entities.py b/spacy/tests/doc/test_add_entities.py index cd444ba81..31d2b8420 100644 --- a/spacy/tests/doc/test_add_entities.py +++ b/spacy/tests/doc/test_add_entities.py @@ -18,7 +18,7 @@ def test_doc_add_entities_set_ents_iob(en_vocab): assert [w.ent_iob_ for w in doc] == (['O'] * len(doc)) doc.ents = [(doc.vocab.strings['ANIMAL'], 3, 4)] - assert [w.ent_iob_ for w in doc] == ['O', 'O', 'O', 'B'] + assert [w.ent_iob_ for w in doc] == ['', '', '', 'B'] doc.ents = [(doc.vocab.strings['WORD'], 0, 2)] - assert [w.ent_iob_ for w in doc] == ['B', 'I', 'O', 'O'] + assert [w.ent_iob_ for w in doc] == ['B', 'I', '', ''] diff --git a/spacy/tests/doc/test_span_merge.py b/spacy/tests/doc/test_span_merge.py index 61f8ca50d..ae1f4f4a1 100644 --- a/spacy/tests/doc/test_span_merge.py +++ b/spacy/tests/doc/test_span_merge.py @@ -2,6 +2,8 @@ from __future__ import unicode_literals from ..util import get_doc +from ...vocab import Vocab +from ...tokens import Doc import pytest @@ -95,6 +97,21 @@ def test_spans_entity_merge(en_tokenizer): assert len(doc) == 15 +def test_spans_entity_merge_iob(): + # Test entity IOB stays consistent after merging + words = ["a", "b", "c", "d", "e"] + doc = Doc(Vocab(), words=words) + doc.ents = [(doc.vocab.strings.add('ent-abc'), 0, 3), + (doc.vocab.strings.add('ent-d'), 3, 4)] + assert doc[0].ent_iob_ == "B" + assert doc[1].ent_iob_ == "I" + assert doc[2].ent_iob_ == "I" + assert doc[3].ent_iob_ == "B" + doc[0:1].merge() + assert doc[0].ent_iob_ == "B" + assert doc[1].ent_iob_ == "I" + + def test_spans_sentence_update_after_merge(en_tokenizer): text = "Stewart Lee is a stand up comedian. He lives in England and loves Joe Pasquale." heads = [1, 1, 0, 1, 2, -1, -4, -5, 1, 0, -1, -1, -3, -4, 1, -2, -7] diff --git a/spacy/tokens/doc.pyx b/spacy/tokens/doc.pyx index 8c5e04ea6..1885dc872 100644 --- a/spacy/tokens/doc.pyx +++ b/spacy/tokens/doc.pyx @@ -421,7 +421,12 @@ cdef class Doc: for i in range(self.length): token = &self.c[i] if token.ent_iob == 1: - assert start != -1 + if start == -1: + seq = ['%s|%s' % (t.text, t.ent_iob_) for t in self[i-5:i+5]] + raise ValueError( + "token.ent_iob values make invalid sequence: " + "I without B\n" + "{seq}".format(seq=' '.join(seq))) elif token.ent_iob == 2 or token.ent_iob == 0: if start != -1: output.append(Span(self, start, i, label=label)) @@ -446,10 +451,7 @@ cdef class Doc: cdef int i for i in range(self.length): self.c[i].ent_type = 0 - # At this point we don't know whether the NER has run over the - # Doc. If the ent_iob is missing, leave it missing. - if self.c[i].ent_iob != 0: - self.c[i].ent_iob = 2 # Means O. Non-O are set from ents. + self.c[i].ent_iob = 0 # Means missing. cdef attr_t ent_type cdef int start, end for ent_info in ents: @@ -947,6 +949,13 @@ cdef class Doc: self.vocab.morphology.assign_tag(token, attr_value) else: Token.set_struct_attr(token, attr_name, attr_value) + # Make sure ent_iob remains consistent + if self.c[end].ent_iob == 1 and token.ent_iob in (0, 2): + if token.ent_type == self.c[end].ent_type: + token.ent_iob = 3 + else: + # If they're not the same entity type, let them be two entities + self.c[end].ent_iob = 3 # Begin by setting all the head indices to absolute token positions # This is easier to work with for now than the offsets # Before thinking of something simpler, beware the case where a