spaCy/spacy/tests/pipeline/test_senter.py

83 lines
2.4 KiB
Python

import pytest
from spacy import util
from spacy.training import Example
from spacy.lang.en import English
from spacy.language import Language
from spacy.tests.util import make_tempdir
def test_label_types():
nlp = Language()
senter = nlp.add_pipe("senter")
with pytest.raises(NotImplementedError):
senter.add_label("A")
SENT_STARTS = [0] * 14
SENT_STARTS[0] = 1
SENT_STARTS[5] = 1
SENT_STARTS[9] = 1
TRAIN_DATA = [
(
"I like green eggs. Eat blue ham. I like purple eggs.",
{"sent_starts": SENT_STARTS},
),
(
"She likes purple eggs. They hate ham. You like yellow eggs.",
{"sent_starts": SENT_STARTS},
),
]
def test_initialize_examples():
nlp = Language()
nlp.add_pipe("senter")
train_examples = []
for t in TRAIN_DATA:
train_examples.append(Example.from_dict(nlp.make_doc(t[0]), t[1]))
# you shouldn't really call this more than once, but for testing it should be fine
nlp.initialize()
nlp.initialize(get_examples=lambda: train_examples)
with pytest.raises(TypeError):
nlp.initialize(get_examples=lambda: None)
with pytest.raises(TypeError):
nlp.initialize(get_examples=train_examples)
def test_overfitting_IO():
# Simple test to try and quickly overfit the senter - ensuring the ML models work correctly
nlp = English()
train_examples = []
for t in TRAIN_DATA:
train_examples.append(Example.from_dict(nlp.make_doc(t[0]), t[1]))
# add some cases where SENT_START == -1
train_examples[0].reference[10].is_sent_start = False
train_examples[1].reference[1].is_sent_start = False
train_examples[1].reference[11].is_sent_start = False
nlp.add_pipe("senter")
optimizer = nlp.initialize()
for i in range(200):
losses = {}
nlp.update(train_examples, sgd=optimizer, losses=losses)
assert losses["senter"] < 0.001
# test the trained model
test_text = TRAIN_DATA[0][0]
doc = nlp(test_text)
gold_sent_starts = [0] * 14
gold_sent_starts[0] = 1
gold_sent_starts[5] = 1
gold_sent_starts[9] = 1
assert [int(t.is_sent_start) for t in doc] == gold_sent_starts
# Also test the results are still the same after IO
with make_tempdir() as tmp_dir:
nlp.to_disk(tmp_dir)
nlp2 = util.load_model_from_path(tmp_dir)
doc2 = nlp2(test_text)
assert [int(t.is_sent_start) for t in doc2] == gold_sent_starts