def test_preprocess(nlp: BertPipeline):
# We split a document into spans (in this case, into sentences).
nlp.text_splitter = lambda text: text.split('\n')
raw_document = ("This is the test sentence 1.\n"
"This is the test sentence 2.\n"
"This is the test sentence 3.")
task = nlp.preprocess(text=raw_document, aspects=['aspect_1', 'aspect_2'])
assert len(task.subtasks) == 2
assert list(task.subtasks) == ['aspect_1', 'aspect_2']
assert len(task.batch) == 6
assert task.indices == [(0, 3), (3, 6)]
subtask_1, subtask_2 = task
assert subtask_1.text == subtask_2.text == raw_document
assert subtask_1.aspect == 'aspect_1'
assert len(subtask_1.examples) == 3
def test_get_completed_task(nlp: BertPipeline):
text = ("We are great fans of Slack.\n"
"The Slack often has bugs.\n"
"best of all is the warm vibe")
# Make sure we have defined a text_splitter, even naive.
nlp.text_splitter = lambda text: text.split('\n')
task = nlp.preprocess(text, aspects=['slack', 'price'])
tokenized_examples = task.batch
input_batch = nlp.encode(tokenized_examples)
output_batch = nlp.predict(input_batch)
aspect_span_labeled = nlp.label(tokenized_examples, output_batch)
completed_task = nlp.get_completed_task(task, aspect_span_labeled)
assert len(completed_task.batch) == 6
assert completed_task.indices == [(0, 3), (3, 6)]
slack, price = completed_task
assert slack.text == price.text == text
# The sentiment among fragments are different. We normalize scores.
assert np.allclose(slack.scores, [0.06, 0.46, 0.48], atol=0.01)
# Please note once gain that there is a problem
# with the neutral sentiment, model is over-fitted.
assert np.allclose(price.scores, [0.06, 0.42, 0.52], atol=0.01)