mirror of
https://github.com/deepset-ai/haystack.git
synced 2025-07-14 20:40:33 +00:00
91f57015c0
* wip: adding _split_overlapp * fixing join issue for _split_overlap * adding tests * adding release notes * cleaning and fixing tests * making mypy happy * Update haystack/components/preprocessors/document_splitter.py Co-authored-by: Silvano Cerza <3314350+silvanocerza@users.noreply.github.com> * adding docstrings --------- Co-authored-by: Silvano Cerza <3314350+silvanocerza@users.noreply.github.com>
290 lines
13 KiB
Python
290 lines
13 KiB
Python
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
import pytest
|
|
|
|
from haystack import Document
|
|
from haystack.components.preprocessors import DocumentSplitter
|
|
|
|
|
|
def merge_documents(documents):
|
|
"""Merge a list of doc chunks into a single doc by concatenating their content, eliminating overlapping content."""
|
|
sorted_docs = sorted(documents, key=lambda doc: doc.meta["split_idx_start"])
|
|
merged_text = ""
|
|
last_idx_end = 0
|
|
for doc in sorted_docs:
|
|
start = doc.meta["split_idx_start"] # start of the current content
|
|
|
|
# if the start of the current content is before the end of the last appended content, adjust it
|
|
if start < last_idx_end:
|
|
start = last_idx_end
|
|
|
|
# append the non-overlapping part to the merged text
|
|
merged_text = merged_text.strip()
|
|
merged_text += doc.content[start - doc.meta["split_idx_start"] :]
|
|
|
|
# update the last end index
|
|
last_idx_end = doc.meta["split_idx_start"] + len(doc.content)
|
|
|
|
return merged_text
|
|
|
|
|
|
class TestDocumentSplitter:
|
|
def test_non_text_document(self):
|
|
with pytest.raises(
|
|
ValueError, match="DocumentSplitter only works with text documents but content for document ID"
|
|
):
|
|
splitter = DocumentSplitter()
|
|
splitter.run(documents=[Document()])
|
|
|
|
def test_single_doc(self):
|
|
with pytest.raises(TypeError, match="DocumentSplitter expects a List of Documents as input."):
|
|
splitter = DocumentSplitter()
|
|
splitter.run(documents=Document())
|
|
|
|
def test_empty_list(self):
|
|
splitter = DocumentSplitter()
|
|
res = splitter.run(documents=[])
|
|
assert res == {"documents": []}
|
|
|
|
def test_unsupported_split_by(self):
|
|
with pytest.raises(ValueError, match="split_by must be one of 'word', 'sentence', 'page' or 'passage'."):
|
|
DocumentSplitter(split_by="unsupported")
|
|
|
|
def test_unsupported_split_length(self):
|
|
with pytest.raises(ValueError, match="split_length must be greater than 0."):
|
|
DocumentSplitter(split_length=0)
|
|
|
|
def test_unsupported_split_overlap(self):
|
|
with pytest.raises(ValueError, match="split_overlap must be greater than or equal to 0."):
|
|
DocumentSplitter(split_overlap=-1)
|
|
|
|
def test_split_by_word(self):
|
|
splitter = DocumentSplitter(split_by="word", split_length=10)
|
|
result = splitter.run(
|
|
documents=[
|
|
Document(
|
|
content="This is a text with some words. There is a second sentence. And there is a third sentence."
|
|
)
|
|
]
|
|
)
|
|
assert len(result["documents"]) == 2
|
|
assert result["documents"][0].content == "This is a text with some words. There is a "
|
|
assert result["documents"][1].content == "second sentence. And there is a third sentence."
|
|
|
|
def test_split_by_word_with_threshold(self):
|
|
splitter = DocumentSplitter(split_by="word", split_length=15, split_threshold=10)
|
|
result = splitter.run(
|
|
documents=[
|
|
Document(
|
|
content="This is a text with some words. There is a second sentence. And there is a third sentence."
|
|
)
|
|
]
|
|
)
|
|
assert len(result["documents"]) == 1
|
|
assert (
|
|
result["documents"][0].content
|
|
== "This is a text with some words. There is a second sentence. And there is a third sentence."
|
|
)
|
|
|
|
def test_split_by_word_multiple_input_docs(self):
|
|
splitter = DocumentSplitter(split_by="word", split_length=10)
|
|
result = splitter.run(
|
|
documents=[
|
|
Document(
|
|
content="This is a text with some words. There is a second sentence. And there is a third sentence."
|
|
),
|
|
Document(
|
|
content="This is a different text with some words. There is a second sentence. And there is a third sentence. And there is a fourth sentence."
|
|
),
|
|
]
|
|
)
|
|
assert len(result["documents"]) == 5
|
|
assert result["documents"][0].content == "This is a text with some words. There is a "
|
|
assert result["documents"][1].content == "second sentence. And there is a third sentence."
|
|
assert result["documents"][2].content == "This is a different text with some words. There is "
|
|
assert result["documents"][3].content == "a second sentence. And there is a third sentence. And "
|
|
assert result["documents"][4].content == "there is a fourth sentence."
|
|
|
|
def test_split_by_sentence(self):
|
|
splitter = DocumentSplitter(split_by="sentence", split_length=1)
|
|
result = splitter.run(
|
|
documents=[
|
|
Document(
|
|
content="This is a text with some words. There is a second sentence. And there is a third sentence."
|
|
)
|
|
]
|
|
)
|
|
assert len(result["documents"]) == 3
|
|
assert result["documents"][0].content == "This is a text with some words."
|
|
assert result["documents"][1].content == " There is a second sentence."
|
|
assert result["documents"][2].content == " And there is a third sentence."
|
|
|
|
def test_split_by_passage(self):
|
|
splitter = DocumentSplitter(split_by="passage", split_length=1)
|
|
result = splitter.run(
|
|
documents=[
|
|
Document(
|
|
content="This is a text with some words. There is a second sentence.\n\nAnd there is a third sentence.\n\n And another passage."
|
|
)
|
|
]
|
|
)
|
|
assert len(result["documents"]) == 3
|
|
assert result["documents"][0].content == "This is a text with some words. There is a second sentence.\n\n"
|
|
assert result["documents"][1].content == "And there is a third sentence.\n\n"
|
|
assert result["documents"][2].content == " And another passage."
|
|
|
|
def test_split_by_page(self):
|
|
splitter = DocumentSplitter(split_by="page", split_length=1)
|
|
result = splitter.run(
|
|
documents=[
|
|
Document(
|
|
content="This is a text with some words. There is a second sentence.\f And there is a third sentence.\f And another passage."
|
|
)
|
|
]
|
|
)
|
|
assert len(result["documents"]) == 3
|
|
assert result["documents"][0].content == "This is a text with some words. There is a second sentence.\x0c"
|
|
assert result["documents"][1].content == " And there is a third sentence.\x0c"
|
|
assert result["documents"][2].content == " And another passage."
|
|
|
|
def test_split_by_word_with_overlap(self):
|
|
splitter = DocumentSplitter(split_by="word", split_length=10, split_overlap=2)
|
|
result = splitter.run(
|
|
documents=[
|
|
Document(
|
|
content="This is a text with some words. There is a second sentence. And there is a third sentence."
|
|
)
|
|
]
|
|
)
|
|
assert len(result["documents"]) == 2
|
|
assert result["documents"][0].content == "This is a text with some words. There is a "
|
|
assert result["documents"][1].content == "is a second sentence. And there is a third sentence."
|
|
|
|
def test_source_id_stored_in_metadata(self):
|
|
splitter = DocumentSplitter(split_by="word", split_length=10)
|
|
doc1 = Document(content="This is a text with some words.")
|
|
doc2 = Document(content="This is a different text with some words.")
|
|
result = splitter.run(documents=[doc1, doc2])
|
|
assert result["documents"][0].meta["source_id"] == doc1.id
|
|
assert result["documents"][1].meta["source_id"] == doc2.id
|
|
|
|
def test_copy_metadata(self):
|
|
splitter = DocumentSplitter(split_by="word", split_length=10)
|
|
documents = [
|
|
Document(content="Text.", meta={"name": "doc 0"}),
|
|
Document(content="Text.", meta={"name": "doc 1"}),
|
|
]
|
|
result = splitter.run(documents=documents)
|
|
assert len(result["documents"]) == 2
|
|
assert result["documents"][0].id != result["documents"][1].id
|
|
for doc, split_doc in zip(documents, result["documents"]):
|
|
assert doc.meta.items() <= split_doc.meta.items()
|
|
assert split_doc.content == "Text."
|
|
|
|
def test_add_page_number_to_metadata_with_no_overlap_word_split(self):
|
|
splitter = DocumentSplitter(split_by="word", split_length=2)
|
|
doc1 = Document(content="This is some text.\f This text is on another page.")
|
|
doc2 = Document(content="This content has two.\f\f page brakes.")
|
|
result = splitter.run(documents=[doc1, doc2])
|
|
|
|
expected_pages = [1, 1, 2, 2, 2, 1, 1, 3]
|
|
for doc, p in zip(result["documents"], expected_pages):
|
|
assert doc.meta["page_number"] == p
|
|
|
|
def test_add_page_number_to_metadata_with_no_overlap_sentence_split(self):
|
|
splitter = DocumentSplitter(split_by="sentence", split_length=1)
|
|
doc1 = Document(content="This is some text.\f This text is on another page.")
|
|
doc2 = Document(content="This content has two.\f\f page brakes.")
|
|
result = splitter.run(documents=[doc1, doc2])
|
|
|
|
expected_pages = [1, 1, 1, 1]
|
|
for doc, p in zip(result["documents"], expected_pages):
|
|
assert doc.meta["page_number"] == p
|
|
|
|
def test_add_page_number_to_metadata_with_no_overlap_passage_split(self):
|
|
splitter = DocumentSplitter(split_by="passage", split_length=1)
|
|
doc1 = Document(
|
|
content="This is a text with some words.\f There is a second sentence.\n\nAnd there is a third sentence.\n\nAnd more passages.\n\n\f And another passage."
|
|
)
|
|
result = splitter.run(documents=[doc1])
|
|
|
|
expected_pages = [1, 2, 2, 2]
|
|
for doc, p in zip(result["documents"], expected_pages):
|
|
assert doc.meta["page_number"] == p
|
|
|
|
def test_add_page_number_to_metadata_with_no_overlap_page_split(self):
|
|
splitter = DocumentSplitter(split_by="page", split_length=1)
|
|
doc1 = Document(
|
|
content="This is a text with some words. There is a second sentence.\f And there is a third sentence.\f And another passage."
|
|
)
|
|
result = splitter.run(documents=[doc1])
|
|
expected_pages = [1, 2, 3]
|
|
for doc, p in zip(result["documents"], expected_pages):
|
|
assert doc.meta["page_number"] == p
|
|
|
|
splitter = DocumentSplitter(split_by="page", split_length=2)
|
|
doc1 = Document(
|
|
content="This is a text with some words. There is a second sentence.\f And there is a third sentence.\f And another passage."
|
|
)
|
|
result = splitter.run(documents=[doc1])
|
|
expected_pages = [1, 3]
|
|
|
|
for doc, p in zip(result["documents"], expected_pages):
|
|
assert doc.meta["page_number"] == p
|
|
|
|
def test_add_page_number_to_metadata_with_overlap_word_split(self):
|
|
splitter = DocumentSplitter(split_by="word", split_length=3, split_overlap=1)
|
|
doc1 = Document(content="This is some text. And\f this text is on another page.")
|
|
doc2 = Document(content="This content has two.\f\f page brakes.")
|
|
result = splitter.run(documents=[doc1, doc2])
|
|
|
|
expected_pages = [1, 1, 1, 2, 2, 1, 1, 3]
|
|
for doc, p in zip(result["documents"], expected_pages):
|
|
assert doc.meta["page_number"] == p
|
|
|
|
def test_add_page_number_to_metadata_with_overlap_sentence_split(self):
|
|
splitter = DocumentSplitter(split_by="sentence", split_length=2, split_overlap=1)
|
|
doc1 = Document(content="This is some text. And this is more text.\f This text is on another page. End.")
|
|
doc2 = Document(content="This content has two.\f\f page brakes. More text.")
|
|
result = splitter.run(documents=[doc1, doc2])
|
|
|
|
expected_pages = [1, 1, 1, 2, 1, 1]
|
|
for doc, p in zip(result["documents"], expected_pages):
|
|
assert doc.meta["page_number"] == p
|
|
|
|
def test_add_page_number_to_metadata_with_overlap_passage_split(self):
|
|
splitter = DocumentSplitter(split_by="passage", split_length=2, split_overlap=1)
|
|
doc1 = Document(
|
|
content="This is a text with some words.\f There is a second sentence.\n\nAnd there is a third sentence.\n\nAnd more passages.\n\n\f And another passage."
|
|
)
|
|
result = splitter.run(documents=[doc1])
|
|
|
|
expected_pages = [1, 2, 2]
|
|
for doc, p in zip(result["documents"], expected_pages):
|
|
assert doc.meta["page_number"] == p
|
|
|
|
def test_add_page_number_to_metadata_with_overlap_page_split(self):
|
|
splitter = DocumentSplitter(split_by="page", split_length=2, split_overlap=1)
|
|
doc1 = Document(
|
|
content="This is a text with some words. There is a second sentence.\f And there is a third sentence.\f And another passage."
|
|
)
|
|
result = splitter.run(documents=[doc1])
|
|
expected_pages = [1, 2, 3]
|
|
|
|
for doc, p in zip(result["documents"], expected_pages):
|
|
assert doc.meta["page_number"] == p
|
|
|
|
def test_add_split_overlap_information(self):
|
|
splitter = DocumentSplitter(split_length=10, split_overlap=5, split_by="word")
|
|
doc = Document(content="This is a text with some words. There is a second sentence. And a third sentence.")
|
|
docs = splitter.run(documents=[doc])
|
|
|
|
# check split_overlap is added to all the documents
|
|
assert len(docs["documents"]) == 3
|
|
for d in docs["documents"]:
|
|
assert "_split_overlap" in d.meta
|
|
|
|
# reconstruct the original document content from the split documents
|
|
assert doc.content == merge_documents(docs["documents"])
|