diff --git a/spacy/gold/converters/iob2docs.py b/spacy/gold/converters/iob2docs.py index aba23e1b3..27876ba7a 100644 --- a/spacy/gold/converters/iob2docs.py +++ b/spacy/gold/converters/iob2docs.py @@ -1,9 +1,9 @@ from wasabi import Printer +from .conll_ner2docs import n_sents_info from ...gold import iob_to_biluo, tags_to_entities from ...tokens import Doc, Span -from .util import merge_sentences -from .conll_ner2docs import n_sents_info +from ...util import minibatch def iob2docs(input_data, vocab, n_sents=10, no_print=False, *args, **kwargs): @@ -19,31 +19,44 @@ def iob2docs(input_data, vocab, n_sents=10, no_print=False, *args, **kwargs): I|PRP|O like|VBP|O London|NNP|B-GPE and|CC|O New|NNP|B-GPE York|NNP|I-GPE City|NNP|I-GPE .|.|O """ msg = Printer(no_print=no_print) - docs = read_iob(input_data.split("\n"), vocab) if n_sents > 0: n_sents_info(msg, n_sents) - docs = merge_sentences(docs, n_sents) + docs = read_iob(input_data.split("\n"), vocab, n_sents) return docs -def read_iob(raw_sents, vocab): +def read_iob(raw_sents, vocab, n_sents): docs = [] - for line in raw_sents: - if not line.strip(): - continue - tokens = [t.split("|") for t in line.split()] - if len(tokens[0]) == 3: - words, tags, iob = zip(*tokens) - elif len(tokens[0]) == 2: - words, iob = zip(*tokens) - tags = ["-"] * len(words) - else: - raise ValueError( - "The sentence-per-line IOB/IOB2 file is not formatted correctly. Try checking whitespace and delimiters. See https://spacy.io/api/cli#convert" - ) + for group in minibatch(raw_sents, size=n_sents): + tokens = [] + words = [] + tags = [] + iob = [] + sent_starts = [] + for line in group: + if not line.strip(): + continue + sent_tokens = [t.split("|") for t in line.split()] + if len(sent_tokens[0]) == 3: + sent_words, sent_tags, sent_iob = zip(*sent_tokens) + elif len(sent_tokens[0]) == 2: + sent_words, sent_iob = zip(*sent_tokens) + sent_tags = ["-"] * len(sent_words) + else: + raise ValueError( + "The sentence-per-line IOB/IOB2 file is not formatted correctly. Try checking whitespace and delimiters. See https://spacy.io/api/cli#convert" + ) + words.extend(sent_words) + tags.extend(sent_tags) + iob.extend(sent_iob) + tokens.extend(sent_tokens) + sent_starts.append(True) + sent_starts.extend([False for _ in sent_words[1:]]) doc = Doc(vocab, words=words) for i, tag in enumerate(tags): doc[i].tag_ = tag + for i, sent_start in enumerate(sent_starts): + doc[i].is_sent_start = sent_start biluo = iob_to_biluo(iob) entities = tags_to_entities(biluo) doc.ents = [Span(doc, start=s, end=e, label=L) for (L, s, e) in entities] diff --git a/spacy/gold/converters/util.py b/spacy/gold/converters/util.py deleted file mode 100644 index 41b3e6d24..000000000 --- a/spacy/gold/converters/util.py +++ /dev/null @@ -1,8 +0,0 @@ -from spacy.util import minibatch - - -def merge_sentences(docs, n_sents): - merged = [] - for group in minibatch(docs, size=n_sents): - raise NotImplementedError - return merged diff --git a/spacy/gold/gold_io.pyx b/spacy/gold/gold_io.pyx index 10ca427ed..61ffb2cfb 100644 --- a/spacy/gold/gold_io.pyx +++ b/spacy/gold/gold_io.pyx @@ -6,37 +6,18 @@ from ..tokens import Doc from .iob_utils import biluo_tags_from_offsets -def merge_sents(sents): - m_deps = [[], [], [], [], [], []] - m_cats = {} - m_brackets = [] - i = 0 - for (ids, words, tags, heads, labels, ner), (cats, brackets) in sents: - m_deps[0].extend(id_ + i for id_ in ids) - m_deps[1].extend(words) - m_deps[2].extend(tags) - m_deps[3].extend(head + i for head in heads) - m_deps[4].extend(labels) - m_deps[5].extend(ner) - m_brackets.extend((b["first"] + i, b["last"] + i, b["label"]) - for b in brackets) - m_cats.update(cats) - i += len(ids) - return [(m_deps, (m_cats, m_brackets))] - - -def docs_to_json(docs, id=0, ner_missing_tag="O"): +def docs_to_json(docs, doc_id=0, ner_missing_tag="O"): """Convert a list of Doc objects into the JSON-serializable format used by the spacy train command. docs (iterable / Doc): The Doc object(s) to convert. - id (int): Id for the JSON. + doc_id (int): Id for the JSON. RETURNS (dict): The data in spaCy's JSON format - each input doc will be treated as a paragraph in the output doc """ if isinstance(docs, Doc): docs = [docs] - json_doc = {"id": id, "paragraphs": []} + json_doc = {"id": doc_id, "paragraphs": []} for i, doc in enumerate(docs): json_para = {'raw': doc.text, "sentences": [], "cats": [], "entities": [], "links": []} for cat, val in doc.cats.items(): diff --git a/spacy/tests/test_cli.py b/spacy/tests/test_cli.py index 164961a5b..0cf070b61 100644 --- a/spacy/tests/test_cli.py +++ b/spacy/tests/test_cli.py @@ -117,7 +117,6 @@ def test_cli_converters_conllu2json_subtokens(): assert [t["ner"] for t in tokens] == ["O", "U-PER", "O", "O"] -@pytest.mark.xfail def test_cli_converters_iob2json(en_vocab): lines = [ "I|O like|O London|I-GPE and|O New|B-GPE York|I-GPE City|I-GPE .|O", @@ -127,19 +126,19 @@ def test_cli_converters_iob2json(en_vocab): ] input_data = "\n".join(lines) converted_docs = iob2docs(input_data, en_vocab, n_sents=10) + assert len(converted_docs) == 1 converted = docs_to_json(converted_docs) - assert len(converted) == 1 - assert converted[0]["id"] == 0 - assert len(converted[0]["paragraphs"]) == 1 - assert len(converted[0]["paragraphs"][0]["sentences"]) == 4 + assert converted["id"] == 0 + assert len(converted["paragraphs"]) == 1 + assert len(converted["paragraphs"][0]["sentences"]) == 4 for i in range(0, 4): - sent = converted[0]["paragraphs"][0]["sentences"][i] + sent = converted["paragraphs"][0]["sentences"][i] assert len(sent["tokens"]) == 8 tokens = sent["tokens"] # fmt: off assert [t["orth"] for t in tokens] == ["I", "like", "London", "and", "New", "York", "City", "."] - assert [t["ner"] for t in tokens] == ["O", "O", "U-GPE", "O", "B-GPE", "I-GPE", "L-GPE", "O"] - # fmt: on + + assert converted["paragraphs"][0]["entities"] == [(18, 26, 'GPE'), (52, 60, 'GPE'), (86, 94, 'GPE'), (120, 128, 'GPE')] @pytest.mark.xfail