def test_issue4313(): """ This should not crash or exit with some strange error code """ beam_width = 16 beam_density = 0.0001 nlp = English() ner = EntityRecognizer(nlp.vocab) ner.add_label("SOME_LABEL") ner.begin_training([]) nlp.add_pipe(ner) # add a new label to the doc doc = nlp("What do you think about Apple ?") assert len(ner.labels) == 1 assert "SOME_LABEL" in ner.labels apple_ent = Span(doc, 5, 6, label="MY_ORG") doc.ents = list(doc.ents) + [apple_ent] # ensure the beam_parse still works with the new label docs = [doc] beams = nlp.entity.beam_parse(docs, beam_width=beam_width, beam_density=beam_density) for doc, beam in zip(docs, beams): entity_scores = defaultdict(float) for score, ents in nlp.entity.moves.get_beam_parses(beam): for start, end, label in ents: entity_scores[(start, end, label)] += score
def test_ents_reset(en_vocab): text = ["This", "is", "a", "lion"] doc = get_doc(en_vocab, text) ner = EntityRecognizer(en_vocab) ner.begin_training([]) ner(doc) assert [t.ent_iob_ for t in doc] == (["O"] * len(doc)) doc.ents = list(doc.ents) assert [t.ent_iob_ for t in doc] == (["O"] * len(doc))
def test_add_label_deserializes_correctly(): ner1 = EntityRecognizer(Vocab()) ner1.add_label("C") ner1.add_label("B") ner1.add_label("A") ner1.begin_training([]) ner2 = EntityRecognizer(Vocab()).from_bytes(ner1.to_bytes()) assert ner1.moves.n_moves == ner2.moves.n_moves for i in range(ner1.moves.n_moves): assert ner1.moves.get_class_name(i) == ner2.moves.get_class_name(i)
def test_doc_add_entities_set_ents_iob(en_vocab): doc = Doc(en_vocab, words=["This", "is", "a", "lion"]) ner = EntityRecognizer(en_vocab) ner.begin_training([]) ner(doc) assert len(list(doc.ents)) == 0 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] == ["", "", "", "B"] doc.ents = [(doc.vocab.strings["WORD"], 0, 2)] assert [w.ent_iob_ for w in doc] == ["B", "I", "", ""]