mirror of
https://github.com/deepset-ai/haystack.git
synced 2026-10-12 02:12:50 +00:00
test: mock all Summarizer tests and move a few into e2e (#4299)
* stub e2e folders * simplify pipeline test * mocking * unit tests fixed * clean up e2e * pipeline tests work * pylint * leftover * small fix from #2994 and additional tests * review feedback * change summaries * black * revert models and summaries
This commit is contained in:
ZanSara
@@ -1,8 +1,12 @@
|
||||
import os
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
import random
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack.schema import Document
|
||||
@@ -20,6 +24,12 @@ from haystack.document_stores import (
|
||||
SAMPLES_PATH = Path(__file__).parent.parent / "test" / "samples"
|
||||
|
||||
|
||||
# Fix all random seeds that come to mind
|
||||
torch.manual_seed(0)
|
||||
np.random.seed(0)
|
||||
random.seed(0)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def docs_all_formats():
|
||||
return [
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
from haystack.schema import Document
|
||||
from haystack.nodes import TransformersSummarizer
|
||||
|
||||
|
||||
DOCS = [
|
||||
Document(
|
||||
content="""
|
||||
PG&E stated it scheduled the blackouts in response to forecasts for high winds amid dry conditions.
|
||||
The aim is to reduce the risk of wildfires. Nearly 800 thousand customers were scheduled to be affected
|
||||
by the shutoffs which were expected to last through at least midday tomorrow.
|
||||
"""
|
||||
),
|
||||
Document(
|
||||
content="""
|
||||
The tower is 324 metres (1,063 ft) tall, about the same height as an 81-storey building, and the tallest
|
||||
structure in Paris. Its base is square, measuring 125 metres (410 ft) on each side. During its construction,
|
||||
the Eiffel Tower surpassed the Washington Monument to become the tallest man-made structure in the world, a
|
||||
title it held for 41 years until the Chrysler Building in New York City was finished in 1930. It was the first
|
||||
structure to reach a height of 300 metres. Due to the addition of a broadcasting aerial at the top of the tower
|
||||
in 1957, it is now taller than the Chrysler Building by 5.2 metres (17 ft). Excluding transmitters, the Eiffel
|
||||
Tower is the second tallest free-standing structure in France after the Millau Viaduct.
|
||||
"""
|
||||
),
|
||||
]
|
||||
|
||||
EXPECTED_SUMMARIES = [
|
||||
"Hundreds of thousands of customers in California have been left without power after a power cut.",
|
||||
"The Eiffel Tower in Paris is the world's tallest man-made structure.",
|
||||
]
|
||||
|
||||
|
||||
def test_summarization():
|
||||
summarizer = TransformersSummarizer(model_name_or_path="sshleifer/distilbart-xsum-12-6", use_gpu=False)
|
||||
summarized_docs = summarizer.predict(documents=DOCS)
|
||||
assert len(summarized_docs) == len(DOCS)
|
||||
for expected_summary, summary in zip(EXPECTED_SUMMARIES, summarized_docs):
|
||||
assert expected_summary.strip() == summary.meta["summary"].strip()
|
||||
|
||||
|
||||
def test_summarization_batch():
|
||||
summarizer = TransformersSummarizer(model_name_or_path="sshleifer/distilbart-xsum-12-6", use_gpu=False)
|
||||
summarized_docs = summarizer.predict_batch(documents=[DOCS, DOCS])
|
||||
assert len(summarized_docs) == 2 # Number of document lists
|
||||
assert len(summarized_docs[0]) == len(DOCS)
|
||||
assert len(summarized_docs[1]) == len(DOCS)
|
||||
for expected_summary, summary_1, summary_2 in zip(EXPECTED_SUMMARIES, summarized_docs[0], summarized_docs[1]):
|
||||
assert expected_summary.strip() == summary_1.meta["summary"].strip()
|
||||
assert expected_summary.strip() == summary_2.meta["summary"].strip()
|
||||
@@ -0,0 +1,72 @@
|
||||
from haystack import Document
|
||||
from haystack.document_stores import InMemoryDocumentStore
|
||||
from haystack.pipelines import SearchSummarizationPipeline
|
||||
from haystack.nodes import BM25Retriever, TransformersSummarizer
|
||||
|
||||
|
||||
def test_summarization_pipeline():
|
||||
docs = [
|
||||
Document(
|
||||
content="""
|
||||
PG&E stated it scheduled the blackouts in response to forecasts for high winds amid dry conditions.
|
||||
The aim is to reduce the risk of wildfires. Nearly 800 thousand customers were scheduled to be affected
|
||||
by the shutoffs which were expected to last through at least midday tomorrow.
|
||||
"""
|
||||
),
|
||||
Document(
|
||||
content="""
|
||||
The Eiffel Tower is 324 metres (1,063 ft) tall, about the same height as an 81-storey building, and the tallest
|
||||
structure in Paris. Its base is square, measuring 125 metres (410 ft) on each side. During its construction,
|
||||
the Eiffel Tower surpassed the Washington Monument to become the tallest man-made structure in the world, a
|
||||
title it held for 41 years until the Chrysler Building in New York City was finished in 1930. It was the first
|
||||
structure to reach a height of 300 metres. Due to the addition of a broadcasting aerial at the top of the tower
|
||||
in 1957, it is now taller than the Chrysler Building by 5.2 metres (17 ft). Excluding transmitters, the Eiffel
|
||||
Tower is the second tallest free-standing structure in France after the Millau Viaduct.
|
||||
"""
|
||||
),
|
||||
]
|
||||
summarizer = TransformersSummarizer(model_name_or_path="sshleifer/distilbart-xsum-12-6", use_gpu=False)
|
||||
|
||||
ds = InMemoryDocumentStore(use_bm25=True)
|
||||
retriever = BM25Retriever(document_store=ds)
|
||||
ds.write_documents(docs)
|
||||
|
||||
query = "Eiffel Tower"
|
||||
pipeline = SearchSummarizationPipeline(retriever=retriever, summarizer=summarizer, return_in_answer_format=True)
|
||||
output = pipeline.run(query=query, params={"Retriever": {"top_k": 1}})
|
||||
answers = output["answers"]
|
||||
assert len(answers) == 1
|
||||
assert "The Eiffel Tower is one of the world's tallest structures" == answers[0]["answer"].strip()
|
||||
|
||||
|
||||
def test_summarization_pipeline_one_summary():
|
||||
split_docs = [
|
||||
Document(
|
||||
content="""
|
||||
The tower is 324 metres (1,063 ft) tall, about the same height as an 81-storey building, and the tallest structure in Paris.
|
||||
Its base is square, measuring 125 metres (410 ft) on each side. During its construction, the Eiffel Tower surpassed the
|
||||
Washington Monument to become the tallest man-made structure in the world, a title it held for 41 years until the Chrysler
|
||||
Building in New York City was finished in 1930.
|
||||
"""
|
||||
),
|
||||
Document(
|
||||
content="""
|
||||
It was the first structure to reach a height of 300 metres. Due to the addition of a broadcasting aerial at the
|
||||
top of the tower in 1957, it is now taller than the Chrysler Building by 5.2 metres (17 ft). Excluding transmitters,
|
||||
the Eiffel Tower is the second tallest free-standing structure in France after the Millau Viaduct.
|
||||
"""
|
||||
),
|
||||
]
|
||||
ds = InMemoryDocumentStore(use_bm25=True)
|
||||
retriever = BM25Retriever(document_store=ds)
|
||||
ds.write_documents(split_docs)
|
||||
summarizer = TransformersSummarizer(model_name_or_path="sshleifer/distilbart-xsum-12-6", use_gpu=False)
|
||||
|
||||
query = "Eiffel Tower"
|
||||
pipeline = SearchSummarizationPipeline(
|
||||
retriever=retriever, summarizer=summarizer, generate_single_summary=True, return_in_answer_format=True
|
||||
)
|
||||
output = pipeline.run(query=query, params={"Retriever": {"top_k": 2}})
|
||||
answers = output["answers"]
|
||||
assert len(answers) == 1
|
||||
assert answers[0]["answer"].strip() == "The Eiffel Tower was built in 1924 in Paris, France."
|
||||
@@ -128,10 +128,10 @@ class TransformersSummarizer(BaseSummarizer):
|
||||
:return: List of Documents, where Document.meta["summary"] contains the summarization
|
||||
"""
|
||||
if self.min_length > self.max_length:
|
||||
raise AttributeError("min_length cannot be greater than max_length")
|
||||
raise ValueError("min_length cannot be greater than max_length")
|
||||
|
||||
if len(documents) == 0:
|
||||
raise AttributeError("Summarizer needs at least one document to produce a summary.")
|
||||
raise ValueError("Summarizer needs at least one document to produce a summary.")
|
||||
|
||||
contexts: List[str] = [doc.content for doc in documents]
|
||||
|
||||
@@ -178,12 +178,12 @@ class TransformersSummarizer(BaseSummarizer):
|
||||
:param batch_size: Number of Documents to process at a time.
|
||||
"""
|
||||
if self.min_length > self.max_length:
|
||||
raise AttributeError("min_length cannot be greater than max_length")
|
||||
raise ValueError("min_length cannot be greater than max_length")
|
||||
|
||||
if len(documents) == 0 or (
|
||||
isinstance(documents[0], list) and all(len(docs) == 0 for docs in documents if isinstance(docs, list))
|
||||
):
|
||||
raise AttributeError("Summarizer needs at least one document to produce a summary.")
|
||||
raise ValueError("Summarizer needs at least one document to produce a summary.")
|
||||
|
||||
if batch_size is None:
|
||||
batch_size = self.batch_size
|
||||
|
||||
@@ -549,11 +549,6 @@ def lfqa_generator(request):
|
||||
return Seq2SeqGenerator(model_name_or_path=request.param, min_length=100, max_length=200)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def summarizer():
|
||||
return TransformersSummarizer(model_name_or_path="sshleifer/distilbart-xsum-12-6", use_gpu=False)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def reader_without_normalized_scores():
|
||||
return FARMReader(
|
||||
|
||||
@@ -1,58 +1,74 @@
|
||||
from ensurepip import version
|
||||
import pytest
|
||||
|
||||
import haystack
|
||||
from haystack.utils.torch_utils import ListDataset
|
||||
from haystack.schema import Document
|
||||
from haystack.pipelines import SearchSummarizationPipeline
|
||||
from haystack.nodes import DensePassageRetriever, EmbeddingRetriever, TransformersSummarizer
|
||||
from haystack.nodes.other.document_merger import DocumentMerger
|
||||
from haystack.nodes import TransformersSummarizer
|
||||
|
||||
|
||||
pytestmark = pytest.mark.skip("Tests are too heavy for Github runners, skipping for now")
|
||||
|
||||
|
||||
DOCS = [
|
||||
Document(
|
||||
content="""PG&E stated it scheduled the blackouts in response to forecasts for high winds amid dry conditions. The aim is to reduce the risk of wildfires. Nearly 800 thousand customers were scheduled to be affected by the shutoffs which were expected to last through at least midday tomorrow."""
|
||||
),
|
||||
Document(
|
||||
content="""The tower is 324 metres (1,063 ft) tall, about the same height as an 81-storey building, and the tallest structure in Paris. Its base is square, measuring 125 metres (410 ft) on each side. During its construction, the Eiffel Tower surpassed the Washington Monument to become the tallest man-made structure in the world, a title it held for 41 years until the Chrysler Building in New York City was finished in 1930. It was the first structure to reach a height of 300 metres. Due to the addition of a broadcasting aerial at the top of the tower in 1957, it is now taller than the Chrysler Building by 5.2 metres (17 ft). Excluding transmitters, the Eiffel Tower is the second tallest free-standing structure in France after the Millau Viaduct."""
|
||||
),
|
||||
]
|
||||
|
||||
EXPECTED_SUMMARIES = [
|
||||
"California's largest electricity provider, PG&E, has shut down power supplies to thousands of customers.",
|
||||
" The Eiffel Tower in Paris has officially opened its doors to the public.",
|
||||
]
|
||||
|
||||
SPLIT_DOCS = [
|
||||
Document(
|
||||
content="""The tower is 324 metres (1,063 ft) tall, about the same height as an 81-storey building, and the tallest structure in Paris. Its base is square, measuring 125 metres (410 ft) on each side. During its construction, the Eiffel Tower surpassed the Washington Monument to become the tallest man-made structure in the world, a title it held for 41 years until the Chrysler Building in New York City was finished in 1930."""
|
||||
),
|
||||
Document(
|
||||
content="""It was the first structure to reach a height of 300 metres. Due to the addition of a broadcasting aerial at the top of the tower in 1957, it is now taller than the Chrysler Building by 5.2 metres (17 ft). Excluding transmitters, the Eiffel Tower is the second tallest free-standing structure in France after the Millau Viaduct."""
|
||||
),
|
||||
]
|
||||
|
||||
# Documents order is very important to produce summary.
|
||||
# Different order of same documents produce different summary.
|
||||
EXPECTED_ONE_SUMMARIES = [
|
||||
" The Eiffel Tower in Paris has officially opened its doors to the public.",
|
||||
" The Eiffel Tower in Paris has become the tallest man-made structure in the world.",
|
||||
DOCS = [Document(content=doc) for doc in ["First test doc", "Second test doc"]]
|
||||
EXPECTED_SUMMARIES = ["First summary", "Second summary"]
|
||||
SUMMARIZED_DOCS = [
|
||||
Document(content=doc.content, meta={"summary": summary}) for doc, summary in zip(DOCS, EXPECTED_SUMMARIES)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.summarizer
|
||||
def test_summarization(summarizer):
|
||||
class MockHFPipeline:
|
||||
def __init__(self, *a, **k):
|
||||
pass
|
||||
|
||||
def __call__(self, docs, *a, **k):
|
||||
summaries = [{"summary_text": summary} for summary in EXPECTED_SUMMARIES]
|
||||
if isinstance(docs, ListDataset):
|
||||
return [summaries for _ in docs]
|
||||
return summaries
|
||||
|
||||
def tokenizer(self, *a, **k):
|
||||
return {"input_ids": []}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_models(monkeypatch):
|
||||
monkeypatch.setattr(haystack.nodes.summarizer.transformers, "pipeline", MockHFPipeline)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def summarizer(mock_models) -> TransformersSummarizer:
|
||||
return TransformersSummarizer(model_name_or_path="irrelevant/anyway", use_gpu=False)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_summarization_no_docs(summarizer):
|
||||
with pytest.raises(ValueError, match="at least one document"):
|
||||
summarizer.predict(documents=[])
|
||||
with pytest.raises(ValueError, match="at least one document"):
|
||||
summarizer.predict_batch(documents=[])
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_summarization_no_docs(summarizer):
|
||||
summarizer.min_length = 10
|
||||
summarizer.max_length = 1
|
||||
with pytest.raises(ValueError, match="min_length cannot be greater than max_length"):
|
||||
summarizer.predict(documents=DOCS)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_summarization_one_doc(summarizer):
|
||||
summarized_docs = summarizer.predict(documents=[DOCS[0]])
|
||||
assert len(summarized_docs) == 1
|
||||
assert EXPECTED_SUMMARIES[0] == summarized_docs[0].meta["summary"]
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_summarization_more_docs(summarizer):
|
||||
summarized_docs = summarizer.predict(documents=DOCS)
|
||||
assert len(summarized_docs) == len(DOCS)
|
||||
for expected_summary, summary in zip(EXPECTED_SUMMARIES, summarized_docs):
|
||||
assert expected_summary == summary.meta["summary"]
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.summarizer
|
||||
@pytest.mark.unit
|
||||
def test_summarization_batch_single_doc_list(summarizer):
|
||||
summarized_docs = summarizer.predict_batch(documents=DOCS)
|
||||
assert len(summarized_docs) == len(DOCS)
|
||||
@@ -60,66 +76,10 @@ def test_summarization_batch_single_doc_list(summarizer):
|
||||
assert expected_summary == summary.meta["summary"]
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.summarizer
|
||||
@pytest.mark.unit
|
||||
def test_summarization_batch_multiple_doc_lists(summarizer):
|
||||
summarized_docs = summarizer.predict_batch(documents=[DOCS, DOCS])
|
||||
assert len(summarized_docs) == 2 # Number of document lists
|
||||
assert len(summarized_docs[0]) == len(DOCS)
|
||||
for expected_summary, summary in zip(EXPECTED_SUMMARIES, summarized_docs[0]):
|
||||
assert expected_summary == summary.meta["summary"]
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.summarizer
|
||||
@pytest.mark.parametrize(
|
||||
"retriever,document_store", [("embedding", "memory"), ("bm25", "elasticsearch")], indirect=True
|
||||
)
|
||||
def test_summarization_pipeline(document_store, retriever, summarizer):
|
||||
document_store.write_documents(DOCS)
|
||||
|
||||
if isinstance(retriever, EmbeddingRetriever) or isinstance(retriever, DensePassageRetriever):
|
||||
document_store.update_embeddings(retriever=retriever)
|
||||
|
||||
query = "Where is Eiffel Tower?"
|
||||
pipeline = SearchSummarizationPipeline(retriever=retriever, summarizer=summarizer, return_in_answer_format=True)
|
||||
output = pipeline.run(query=query, params={"Retriever": {"top_k": 1}})
|
||||
answers = output["answers"]
|
||||
assert len(answers) == 1
|
||||
assert " The Eiffel Tower in Paris has officially opened its doors to the public." == answers[0]["answer"]
|
||||
|
||||
|
||||
#
|
||||
# Document Merger + Summarizer tests
|
||||
#
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.summarizer
|
||||
def test_summarization_one_summary(summarizer):
|
||||
dm = DocumentMerger()
|
||||
merged_document = dm.merge(documents=SPLIT_DOCS)
|
||||
summarized_docs = summarizer.predict(documents=merged_document)
|
||||
assert len(summarized_docs) == 1
|
||||
assert EXPECTED_ONE_SUMMARIES[0] == summarized_docs[0].meta["summary"]
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.summarizer
|
||||
@pytest.mark.parametrize(
|
||||
"retriever,document_store", [("embedding", "memory"), ("bm25", "elasticsearch")], indirect=True
|
||||
)
|
||||
def test_summarization_pipeline_one_summary(document_store, retriever, summarizer):
|
||||
document_store.write_documents(SPLIT_DOCS)
|
||||
|
||||
if isinstance(retriever, EmbeddingRetriever) or isinstance(retriever, DensePassageRetriever):
|
||||
document_store.update_embeddings(retriever=retriever)
|
||||
|
||||
query = "Where is Eiffel Tower?"
|
||||
pipeline = SearchSummarizationPipeline(
|
||||
retriever=retriever, summarizer=summarizer, generate_single_summary=True, return_in_answer_format=True
|
||||
)
|
||||
output = pipeline.run(query=query, params={"Retriever": {"top_k": 2}})
|
||||
answers = output["answers"]
|
||||
assert len(answers) == 1
|
||||
assert answers[0]["answer"] in EXPECTED_ONE_SUMMARIES
|
||||
|
||||
Reference in New Issue
Block a user