mirror of
https://github.com/explosion/spaCy.git
synced 2024-12-26 09:56:28 +03:00
Modernise test for adding entities
This commit is contained in:
parent
6e883f4c00
commit
05447be884
|
@ -1,42 +1,23 @@
|
|||
# coding: utf-8
|
||||
from __future__ import unicode_literals
|
||||
import spacy
|
||||
from spacy.vocab import Vocab
|
||||
from spacy.matcher import Matcher
|
||||
from spacy.tokens.doc import Doc
|
||||
from spacy.attrs import *
|
||||
from spacy.pipeline import EntityRecognizer
|
||||
|
||||
from ...pipeline import EntityRecognizer
|
||||
from ..util import get_doc
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def en_vocab():
|
||||
return spacy.get_lang_class('en').Defaults.create_vocab()
|
||||
def test_add_entities_set_ents_iob(en_vocab):
|
||||
text = ["This", "is", "a", "lion"]
|
||||
doc = get_doc(en_vocab, text)
|
||||
ner = EntityRecognizer(en_vocab, features=[(2,), (3,)])
|
||||
ner(doc)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def entity_recognizer(en_vocab):
|
||||
return EntityRecognizer(en_vocab, features=[(2,), (3,)])
|
||||
|
||||
@pytest.fixture
|
||||
def animal(en_vocab):
|
||||
return nlp.vocab.strings[u"ANIMAL"]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def doc(en_vocab, entity_recognizer):
|
||||
doc = Doc(en_vocab, words=[u"this", u"is", u"a", u"lion"])
|
||||
entity_recognizer(doc)
|
||||
return doc
|
||||
|
||||
|
||||
def test_set_ents_iob(doc):
|
||||
assert len(list(doc.ents)) == 0
|
||||
tags = [w.ent_iob_ for w in doc]
|
||||
assert tags == (['O'] * len(doc))
|
||||
assert [w.ent_iob_ for w in doc] == (['O'] * len(doc))
|
||||
|
||||
doc.ents = [(doc.vocab.strings['ANIMAL'], 3, 4)]
|
||||
tags = [w.ent_iob_ for w in doc]
|
||||
assert tags == ['O', 'O', 'O', 'B']
|
||||
assert [w.ent_iob_ for w in doc] == ['O', 'O', 'O', 'B']
|
||||
|
||||
doc.ents = [(doc.vocab.strings['WORD'], 0, 2)]
|
||||
tags = [w.ent_iob_ for w in doc]
|
||||
assert tags == ['B', 'I', 'O', 'O']
|
||||
assert [w.ent_iob_ for w in doc] == ['B', 'I', 'O', 'O']
|
||||
|
|
Loading…
Reference in New Issue
Block a user