From c2406e92bcaaf6cfce300e4fec273d67075dfd8d Mon Sep 17 00:00:00 2001 From: Ines Montani Date: Thu, 12 Jan 2017 12:25:10 +0100 Subject: [PATCH] Allow setting ents in get_doc --- spacy/tests/util.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/spacy/tests/util.py b/spacy/tests/util.py index df4b7eea8..60b7c29ee 100644 --- a/spacy/tests/util.py +++ b/spacy/tests/util.py @@ -5,7 +5,7 @@ from ..tokens import Doc from ..attrs import ORTH, POS, HEAD, DEP -def get_doc(vocab, words=[], pos=None, heads=None, deps=None, tags=None): +def get_doc(vocab, words=[], pos=None, heads=None, deps=None, tags=None, ents=None): """Create Doc object from given vocab, words and annotations.""" pos = pos or [''] * len(words) heads = heads or [0] * len(words) @@ -18,6 +18,8 @@ def get_doc(vocab, words=[], pos=None, heads=None, deps=None, tags=None): attrs[i, 1] = head attrs[i, 2] = doc.vocab.strings[dep] doc.from_array([POS, HEAD, DEP], attrs) + if ents: + doc.ents = [(ent_id, doc.vocab.strings[label], start, end) for ent_id, label, start, end in ents] if tags: for token in doc: token.tag_ = tags[token.i]