spaCy/spacy/tests/doc/test_add_entities.py

36 lines
1019 B
Python
Raw Normal View History

2017-01-11 20:54:24 +03:00
# coding: utf-8
from __future__ import unicode_literals
2017-01-11 20:54:24 +03:00
from ...pipeline import EntityRecognizer
from ..util import get_doc
from ...tokens import Span
2017-01-11 20:54:24 +03:00
import pytest
2017-01-11 21:00:52 +03:00
def test_doc_add_entities_set_ents_iob(en_vocab):
2017-01-11 20:54:24 +03:00
text = ["This", "is", "a", "lion"]
doc = get_doc(en_vocab, text)
ner = EntityRecognizer(en_vocab)
ner.begin_training([])
2017-01-11 20:54:24 +03:00
ner(doc)
assert len(list(doc.ents)) == 0
2017-01-11 20:54:24 +03:00
assert [w.ent_iob_ for w in doc] == (['O'] * len(doc))
doc.ents = [(doc.vocab.strings['ANIMAL'], 3, 4)]
2018-03-26 08:14:35 +03:00
assert [w.ent_iob_ for w in doc] == ['', '', '', 'B']
2017-01-11 20:54:24 +03:00
doc.ents = [(doc.vocab.strings['WORD'], 0, 2)]
2018-03-26 08:14:35 +03:00
assert [w.ent_iob_ for w in doc] == ['B', 'I', '', '']
def test_add_overlapping_entities(en_vocab):
text = ["Louisiana", "Office", "of", "Conservation"]
doc = get_doc(en_vocab, text)
entity = Span(doc, 0, 4, label=391)
doc.ents = [entity]
new_entity = Span(doc, 0, 1, label=392)
with pytest.raises(ValueError):
doc.ents = list(doc.ents) + [new_entity]