feat!: align page_number metadata across splitters and fix drift in overlapping splits (#13006)

Co-authored-by: Sebastian Husch Lee <[email protected]>
This commit is contained in:
bogdankostic
2026-10-01 15:36:11 +02:00
committed by GitHub
co-authored by Sebastian Husch Lee
parent 10e49fa904
commit 472da278db
14 changed files with 535 additions and 86 deletions
@@ -33,7 +33,7 @@ description: "`DocumentSplitter` divides a list of text documents into a list of
When `split_by="token"`, strings such as `<|endoftext|>` in the source document are tokenized as ordinary text. They are preserved in the chunks and count toward the token limit.
A field `"source_id"` is added to each document's `meta` data to keep track of the original document that was split. Another meta field `"page_number"` is added to each document to keep track of the page it belonged to in the original document. Other metadata are copied from the original document.
A field `"source_id"` is added to each document's `meta` data to keep track of the original document that was split. Another meta field `"page_number"` is added to each document with the page the chunk starts on, counted from the form feed (`\f`) characters in the original document. A chunk that spans a page break reports the page its text begins on. Other metadata are copied from the original document.
The DocumentSplitter is compatible with the following DocumentStores:
@@ -28,7 +28,7 @@ This component splits documents based on embedding similarity using cosine dista
It first splits text into sentences, optionally groups them, calculates embeddings for each group, and then uses cosine
distance between sequential embeddings to determine split points. Any distance above the specified percentile is treated
as a break point. The component also tracks page numbers based on form feed characters (`\f`) in the original document.
as a break point. The component also tracks page numbers based on form feed characters (`\f`) in the original document, reporting the page each chunk starts on.
This component is inspired by [5 Levels of Text Splitting](https://github.com/FullStackRetrieval-com/RetrievalTutorials/blob/main/tutorials/LevelsOfTextSplitting/5_Levels_Of_Text_Splitting.ipynb) by Greg Kamradt.
@@ -45,7 +45,7 @@ Parameters you can set when initializing the component:
Each output document's metadata includes:
- `source_id`: ID of the original document.
- `page_number`: Page number. Updated when `page_break_character` is found.
- `page_number`: The page the chunk starts on, counted from the `page_break_character` occurrences that precede its text.
- `split_id`: Index of the chunk within its parent.
- `header`: The header text for this chunk.
- `parent_headers`: List of parent header texts in hierarchy order.
@@ -33,6 +33,8 @@ The `RecursiveDocumentSplitter` expects a list of documents as input and returns
When `split_unit="token"`, strings such as `<|endoftext|>` in the source document are tokenized as ordinary text. They are preserved in the chunks and count toward the token limit.
A meta field `"page_number"` is added to each chunk with the page it starts on, counted from the form feed (`\f`) characters in the original document. A chunk that spans a page break reports the page its text begins on, matching [`DocumentSplitter`](documentsplitter.mdx).
The separators are applied in the same order as they are defined in the list. The first separator is used on the text; any resulting chunk that is within the specified `chunk_size` is retained. For chunks that exceed the defined `chunk_size`, the next separator in the list is applied. If all separators are used and the chunk still exceeds the `chunk_size`, a hard split occurs based on the `chunk_size`, taking into account whether words or characters are used as counting units. This process is repeated until all chunks are within the limits of the specified `chunk_size`.
## Usage
@@ -0,0 +1,38 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
"""
The `page_number` convention shared by the splitters.
Splitters record the page a chunk came from in its `page_number` metadata field, counting page breaks
(form feed, "\f") in the source document. The page reported is the one the chunk's first non-page-break
character lies on: a chunk that spans a page break belongs to the page it starts on, and page breaks a
chunk opens with belong to that chunk's own page rather than to the page it was carved out of.
"""
def _leading_page_breaks(text: str, page_break_character: str = "\f") -> int:
"""
Count how many pages a chunk's text starts past the page its first character is on.
This is the run of page breaks the chunk opens with, except for a chunk made up entirely of page
breaks: that chunk is an empty page and stays on the page it starts on, rather than being pushed onto
the following one. `split_by="page"` emits such chunks for every blank page.
:param text: The chunk to inspect.
:param page_break_character: The character sequence marking a page break.
:returns: The number of pages to advance, 0 if the chunk starts with anything but a page break or
consists only of page breaks.
"""
if not page_break_character:
return 0
offset = 0
while text.startswith(page_break_character, offset):
offset += len(page_break_character)
# Nothing but page breaks: an empty page, which belongs to the page it starts on.
if offset >= len(text):
return 0
return offset // len(page_break_character)
@@ -9,6 +9,7 @@ from typing import Any, Literal
from more_itertools import windowed
from haystack import Document, component, logging
from haystack.components.preprocessors._page_numbers import _leading_page_breaks
from haystack.components.preprocessors.sentence_tokenizer import Language, SentenceSplitter, nltk_imports
from haystack.core.serialization import default_from_dict, default_to_dict
from haystack.lazy_imports import LazyImport
@@ -199,7 +200,8 @@ class DocumentSplitter:
:returns: A dictionary with the following key:
- `documents`: List of documents with the split texts. Each document includes:
- A metadata field `source_id` to track the original document.
- A metadata field `page_number` to track the original page number.
- A metadata field `page_number` with the page the chunk starts on, counting form feed
("\f") characters in the original document.
- All other metadata copied from the original document.
:raises TypeError: if the input is not a list of Documents.
@@ -450,7 +452,7 @@ class DocumentSplitter:
for i, (txt, split_idx) in enumerate(zip(text_splits, splits_start_idxs, strict=True)):
copied_meta = deepcopy(meta)
copied_meta["page_number"] = splits_pages[i]
copied_meta["page_number"] = splits_pages[i] + _leading_page_breaks(txt)
copied_meta["split_id"] = i
copied_meta["split_idx_start"] = split_idx
doc = Document(content=txt, meta=copied_meta)
@@ -12,6 +12,7 @@ import numpy as np
from haystack import Document, component, logging
from haystack.components.embedders.types import DocumentEmbedder
from haystack.components.preprocessors._page_numbers import _leading_page_breaks
from haystack.components.preprocessors.sentence_tokenizer import Language, SentenceSplitter
from haystack.core.serialization import component_to_dict, default_from_dict, default_to_dict
from haystack.utils.async_utils import _execute_component_async
@@ -179,7 +180,8 @@ class EmbeddingBasedDocumentSplitter:
- A metadata field `source_id` to track the original document.
- A metadata field `split_id` to track the split number.
- A metadata field `split_idx_start` with the character offset of the chunk in the original document.
- A metadata field `page_number` to track the original page number.
- A metadata field `page_number` with the page the chunk starts on, counting form feed
("\f") characters in the original document.
- All other metadata copied from the original document.
:raises RuntimeError: If the component wasn't warmed up.
@@ -220,7 +222,8 @@ class EmbeddingBasedDocumentSplitter:
- A metadata field `source_id` to track the original document.
- A metadata field `split_id` to track the split number.
- A metadata field `split_idx_start` with the character offset of the chunk in the original document.
- A metadata field `page_number` to track the original page number.
- A metadata field `page_number` with the page the chunk starts on, counting form feed
("\f") characters in the original document.
- All other metadata copied from the original document.
:raises RuntimeError: If the component wasn't warmed up.
@@ -549,12 +552,11 @@ class EmbeddingBasedDocumentSplitter:
split_meta["split_id"] = i
split_meta["split_idx_start"] = current_char_pos
# Calculate page number for this split
# Count page breaks in the split itself
page_breaks_in_split = split_text.count("\f")
# Calculate the page number for this split
split_meta["page_number"] = current_page
# current_page is the page this split's first character is on; a split that opens with page
# breaks starts its text on a later page.
split_meta["page_number"] = current_page + _leading_page_breaks(split_text)
doc = Document(content=split_text, meta=split_meta)
documents.append(doc)
@@ -8,6 +8,7 @@ from typing import Literal
from haystack import Document, component, logging
from haystack.components.preprocessors import DocumentSplitter
from haystack.components.preprocessors._page_numbers import _leading_page_breaks
logger = logging.getLogger(__name__)
@@ -294,8 +295,12 @@ class MarkdownHeaderSplitter:
secondary_content_start_idx = header_match.end()
content_for_splitting = doc.content[secondary_content_start_idx:]
# The page this header chunk starts on; its splits are numbered relative to it.
chunk_start_page = doc.meta.get("page_number", 1)
# The page this header chunk's first character is on; its splits are numbered relative to it.
# doc.meta holds the page the chunk's *text* starts on, which already counts the breaks the chunk
# opens with, and page_break_ends below counts those same breaks again - so take them back off.
chunk_start_page = doc.meta.get("page_number", 1) - _leading_page_breaks(
doc.content, self.page_break_character
)
clean_meta = {k: v for k, v in doc.meta.items() if k != "split_id"}
@@ -312,7 +317,11 @@ class MarkdownHeaderSplitter:
while page_break_count < len(page_break_ends) and page_break_ends[page_break_count] <= split_start_idx:
page_break_count += 1
split.meta["page_number"] = chunk_start_page + page_break_count
split.meta["page_number"] = (
chunk_start_page
+ page_break_count
+ _leading_page_breaks(split.content or "", self.page_break_character)
)
split.meta["split_id"] = current_split_id
if "source_id" in doc.meta:
split.meta["source_id"] = doc.meta["source_id"]
@@ -358,8 +367,11 @@ class MarkdownHeaderSplitter:
)
for split_idx, split in enumerate(splits):
meta = deepcopy(doc.meta) if doc.meta else {}
chunk_start_page = document_start_page + doc.content.count(
self.page_break_character, 0, split["source_start_idx"]
chunk_start_page = (
document_start_page
+ doc.content.count(self.page_break_character, 0, split["source_start_idx"])
# breaks the chunk opens with end the previous page, so its text is on a later one
+ _leading_page_breaks(split["content"], self.page_break_character)
)
meta.update({"source_id": doc.id, "page_number": chunk_start_page, "split_id": split_idx})
from_header_split = split["from_header_split"]
@@ -386,7 +398,8 @@ class MarkdownHeaderSplitter:
:returns: A dictionary with the following key:
- `documents`: List of documents with the split texts. Each document includes:
- A metadata field `source_id` to track the original document.
- A metadata field `page_number` to track the original page number.
- A metadata field `page_number` with the page the chunk starts on, counting
`page_break_character` occurrences in the original document.
- A metadata field `split_id` to identify the split chunk index within its parent document.
- All other metadata copied from the original document.
:raises ValueError: If a document has `None` content.
@@ -7,6 +7,7 @@ from copy import deepcopy
from typing import Any, Literal
from haystack import Document, component, default_to_dict, logging
from haystack.components.preprocessors._page_numbers import _leading_page_breaks
from haystack.lazy_imports import LazyImport
with LazyImport("Run 'pip install tiktoken'") as tiktoken_imports:
@@ -82,6 +83,8 @@ class RecursiveDocumentSplitter:
If no separators are provided, the default separators ["\\n\\n", "sentence", "\\n", " "] are used.
:param sentence_splitter_params: Optional parameters to pass to the sentence tokenizer.
See: haystack.components.preprocessors.sentence_tokenizer.SentenceSplitter for more information.
The chunks keep the whitespace between sentences whether `keep_white_spaces` is True or False; it only
changes the rules the tokenizer uses to decide where a sentence ends.
:raises ValueError: If the overlap is greater than or equal to the chunk size or if the overlap is negative, or
if any separator is not a string.
@@ -151,11 +154,6 @@ class RecursiveDocumentSplitter:
:returns:
A tuple containing the current chunk and the remaining chunk.
"""
if self.split_units == "word":
words = current_chunk.split()
current_chunk = " ".join(words[: self.split_length])
remaining_words = words[self.split_length :]
return current_chunk, " ".join(remaining_words)
if self.split_units == "char":
text = current_chunk
current_chunk = text[: self.split_length]
@@ -172,9 +170,10 @@ class RecursiveDocumentSplitter:
"""
Applies an overlap between consecutive chunks if the chunk_overlap attribute is greater than zero.
Works for both word- and character-level splitting. It trims the last chunk if it exceeds the split_length and
adds the trimmed content to the next chunk. If the last chunk is still too long after trimming, it splits it
and adds the first chunk to the list. This process continues until the last chunk is within the split_length.
Works for character- and token-level splitting, words are handled by `_apply_word_overlap`. It trims the last
chunk if it exceeds the split_length and adds the trimmed content to the next chunk. If the last chunk is still
too long after trimming, it splits it and adds the first chunk to the list. This process continues until the
last chunk is within the split_length.
:param chunks: A list of text chunks.
:returns:
@@ -203,9 +202,7 @@ class RecursiveDocumentSplitter:
if self._chunk_length(current_chunk) > self.split_length:
current_chunk, remaining_text = self._split_chunk(current_chunk)
if idx < len(chunks) - 1:
if self.split_units == "word":
chunks[idx + 1] = remaining_text + " " + chunks[idx + 1]
elif self.split_units == "token":
if self.split_units == "token":
# For token-based splitting, combine at token level
# at this point we know that the tokenizer is already initialized
remaining_tokens = self.tiktoken_tokenizer.encode_ordinary(remaining_text) # type: ignore
@@ -246,9 +243,7 @@ class RecursiveDocumentSplitter:
return overlapped_chunks
def _create_chunk_starting_with_overlap(self, chunk: str, overlap: str) -> str:
if self.split_units == "word":
current_chunk = overlap + " " + chunk
elif self.split_units == "token":
if self.split_units == "token":
# For token-based splitting, combine at token level
# at this point we know that the tokenizer is already initialized
overlap_tokens = self.tiktoken_tokenizer.encode_ordinary(overlap) # type: ignore
@@ -263,10 +258,7 @@ class RecursiveDocumentSplitter:
prev_chunk = overlapped_chunks[-1]
overlap_start = max(0, self._chunk_length(prev_chunk) - self.split_overlap)
if self.split_units == "word":
word_chunks = prev_chunk.split()
overlap = " ".join(word_chunks[overlap_start:])
elif self.split_units == "token":
if self.split_units == "token":
# For token-based splitting, handle overlap at token level
# at this point we know that the tokenizer is already initialized
tokens = self.tiktoken_tokenizer.encode_ordinary(prev_chunk) # type: ignore
@@ -277,6 +269,58 @@ class RecursiveDocumentSplitter:
return overlap, prev_chunk
def _apply_word_overlap(self, chunks: list[str]) -> list[tuple[int, int]]:
"""
Applies an overlap between consecutive chunks when splitting by words.
Follows `_apply_overlap`: each chunk starts with the last `split_overlap` words of the chunk before it, and a
chunk that grows past `split_length` words is trimmed, its remaining words moving on to the next chunk. The
chunks are cut out of the text rather than rebuilt from their words, so they keep its whitespace and page
breaks. Words are counted as in `_chunk_length`.
:param chunks: A list of text chunks, as returned by `_chunk_text`.
:returns:
The `(start, end)` character offsets of each chunk with the overlap applied, in the concatenated `chunks`.
"""
# The words of every chunk, as offsets into the concatenated chunks. A chunk of only spaces has no words and
# joins the chunk before it (or after it, at the start of the text).
words: list[tuple[int, int]] = []
chunk_word_ends: list[int] = []
chunk_starts: list[int] = []
offset = 0
for chunk in chunks:
chunk_words = [(offset + match.start(), offset + match.end()) for match in re.finditer(r"[^ ]+", chunk)]
if chunk_words:
words.extend(chunk_words)
chunk_word_ends.append(len(words))
chunk_starts.append(offset)
offset += len(chunk)
if not words:
return [(0, offset)]
# a chunk's text runs up to where the next one starts
chunk_ends = chunk_starts[1:] + [offset]
spans = [(0, chunk_ends[0])]
# the words the latest chunk spans, as indices into `words`
first_word, end_word = 0, chunk_word_ends[0]
for idx in range(1, len(chunk_word_ends)):
if end_word - first_word <= self.split_overlap:
logger.warning(
"Overlap is the same as the previous chunk. "
"Consider increasing the `split_length` parameter or decreasing the `split_overlap` parameter."
)
# only the last chunk emits more than one chunk; the others hand their remaining words to the next one
while True:
first_word = max(first_word, end_word - self.split_overlap)
end_word = min(first_word + self.split_length, chunk_word_ends[idx])
# a trimmed chunk ends after its last word, and the next chunk picks up from its overlap
end = chunk_ends[idx] if end_word == chunk_word_ends[idx] else words[end_word - 1][1]
spans.append((words[first_word][0], end))
if idx < len(chunk_word_ends) - 1 or end_word == chunk_word_ends[idx]:
break
return spans
def _chunk_length(self, text: str) -> int:
"""
Get the length of the chunk in the specified units (words, characters, or tokens).
@@ -310,10 +354,19 @@ class RecursiveDocumentSplitter:
return [text]
for curr_separator in self.separators:
# text after the last split that the splits leave out
trailing_text = ""
if curr_separator == "sentence":
# re. ignore: correct SentenceSplitter initialization is checked at the initialization of the component
sentence_with_spans = self.nltk_tokenizer.split_sentences(text) # type: ignore
splits = [sentence["sentence"] for sentence in sentence_with_spans]
# Each split runs from where its sentence starts to where the next one does, so the splits keep the
# whitespace between sentences even where the tokenizer leaves it out (keep_white_spaces=False).
# Leaving it out would glue sentences together, lose page breaks and shift every later chunk's offset.
split_bounds = [sentence["start"] for sentence in sentence_with_spans]
if split_bounds:
split_bounds.append(sentence_with_spans[-1]["end"])
trailing_text = text[split_bounds[-1] :]
splits = [text[start:end] for start, end in zip(split_bounds, split_bounds[1:], strict=False)]
else:
# add escape "\" to the separator and wrapped it in a group so that it's included in the splits as well
escaped_separator = re.escape(curr_separator)
@@ -368,6 +421,9 @@ class RecursiveDocumentSplitter:
chunks.append("".join(current_chunk))
if chunks:
# like the whitespace left over by _fall_back_to_fixed_chunking, attach it to the last chunk rather
# than making it a chunk of its own
chunks[-1] += trailing_text
return chunks
# if no separator worked, fall back to word- or character-level chunking
@@ -438,20 +494,34 @@ class RecursiveDocumentSplitter:
)
def _run_one(self, doc: Document) -> list[Document]:
chunks = self._chunk_text(doc.content) # type: ignore # the caller already check for a non-empty doc.content
content = doc.content or "" # run() skips documents without content, so this is a non-empty string
chunks = self._chunk_text(content)
chunks = chunks[:-1] if len(chunks[-1]) == 0 else chunks # remove last empty chunk if it exists
# apply the overlap once, on the fully chunked list, so that chunks produced
# at inner recursion levels don't get the overlap applied a second time
if self.split_overlap > 0:
chunks = self._apply_overlap(chunks)
current_position = 0
current_page = 1
if self.split_overlap > 0 and self.split_units == "word":
text = "".join(chunks)
spans = self._apply_word_overlap(chunks)
chunks = [text[start:end] for start, end in spans]
chunk_starts = [start for start, _ in spans]
else:
if self.split_overlap > 0:
chunks = self._apply_overlap(chunks)
# each chunk starts where the previous one ends, minus the overlap it repeats
chunk_starts = [0]
for chunk in chunks[:-1]:
overlap = self._get_overlap([chunk])[0] if self.split_overlap > 0 else ""
chunk_starts.append(chunk_starts[-1] + len(chunk) - len(overlap))
new_docs: list[Document] = []
# Chunk starts never decrease, so the page breaks before each chunk are counted on from the previous one
# rather than from the start of the document every time.
page_breaks_before = 0
for split_nr, chunk in enumerate(chunks):
for split_nr, (chunk, current_position) in enumerate(zip(chunks, chunk_starts, strict=True)):
if split_nr > 0:
page_breaks_before += content.count("\f", chunk_starts[split_nr - 1], current_position)
meta = deepcopy(doc.meta)
meta["source_id"] = doc.id
meta["parent_id"] = doc.id
@@ -464,29 +534,11 @@ class RecursiveDocumentSplitter:
if split_nr > 0 and self.split_overlap > 0:
self._add_overlap_info(current_position, new_doc, new_docs)
# count page breaks in the chunk
current_page += chunk.count("\f")
# The page the chunk's first non-page-break character is on: breaks before the chunk, plus the
# ones it opens with.
new_doc.meta["page_number"] = 1 + page_breaks_before + _leading_page_breaks(chunk)
# if there are consecutive page breaks at the end with no more text, adjust the page number
# e.g: "text\f\f\f" -> 3 page breaks, but current_page should be 1
consecutive_page_breaks = len(chunk) - len(chunk.rstrip("\f"))
if consecutive_page_breaks > 0:
new_doc.meta["page_number"] = current_page - consecutive_page_breaks
else:
new_doc.meta["page_number"] = current_page
# keep the new chunk doc and update the current position
new_docs.append(new_doc)
# Advance current_position by chunk length minus overlap.
# split_overlap is in split_units, not chars, so get the actual
# overlap string from _get_overlap() and use its char length.
if self.split_overlap > 0 and split_nr < len(chunks) - 1:
overlap_str, _ = self._get_overlap([doc.content for doc in new_docs]) # type: ignore[misc]
overlap_char_len = len(overlap_str)
else:
overlap_char_len = 0
current_position += len(chunk) - overlap_char_len
return new_docs
@@ -498,7 +550,8 @@ class RecursiveDocumentSplitter:
:param documents: List of Documents to split.
:returns:
A dictionary containing a key "documents" with a List of Documents with smaller chunks of text corresponding
to the input documents.
to the input documents. Each chunk carries a metadata field `page_number` with the page the chunk
starts on, counting form feed ("\f") characters in the original document.
"""
if not self._is_warmed_up and ("sentence" in self.separators or self.split_units == "token"):
self.warm_up()
@@ -0,0 +1,51 @@
---
upgrade:
- |
The ``page_number`` metadata field now means the same thing in every splitter: the page the chunk's
text starts on. Two cases change:
- ``RecursiveDocumentSplitter`` reported the page a chunk *ended* on. A chunk spanning a page break
now reports the page it starts on, matching ``DocumentSplitter`` and ``MarkdownHeaderSplitter``.
- ``DocumentSplitter``, ``MarkdownHeaderSplitter`` and ``EmbeddingBasedDocumentSplitter`` reported the
previous page for a chunk that *began* with one or more form feeds (``\f``), even though all of the
chunk's text was on the following page. Such chunks now report the page their text is on. A chunk
made up entirely of page breaks is an empty page and keeps the page it starts on, so blank pages
emitted by ``split_by="page"`` still get their own number.
If you display ``page_number`` to users or store it alongside embeddings, re-index documents that
contain form feeds to pick up the corrected values.
- |
``RecursiveDocumentSplitter`` with ``split_unit="word"`` and ``split_overlap`` greater than 0 now returns
chunks copied from the document as they are. Previously it rebuilt them from their words joined by single
spaces, which added spaces the document does not have, collapsed line breaks and repeated spaces, and could
drop form feeds (``\f``). The overlap also counts words the way ``split_length`` does, as text separated by
spaces, so where words are separated only by line or page breaks, chunks can end at different words than
before. Re-index documents split this way if you need their chunks to match the source text.
- |
``RecursiveDocumentSplitter`` now keeps the whitespace between sentences when ``sentence_splitter_params``
sets ``keep_white_spaces`` to ``False``. That is the sentence tokenizer's default, so this applies whenever
``sentence_splitter_params`` is passed without ``"keep_white_spaces": True``. The tokenizer leaves that
whitespace out, and the chunks used to glue their sentences together, as in ``"One two.Three four."``: the
glued words counted as one toward ``split_length``, page breaks (``\f``) between sentences were lost, and
``split_idx_start`` and ``page_number`` fell behind the chunk's real start. Chunks are now copied from the
document whatever ``keep_white_spaces`` is set to; it only changes the rules the tokenizer uses to decide where
a sentence ends. Re-index documents split with this setting to pick up the corrected chunks.
fixes:
- |
Fixed ``RecursiveDocumentSplitter`` assigning ``page_number`` values that drifted upwards, and could
exceed the number of pages in the document, when ``split_overlap`` was greater than 0. A page break
repeated in an overlapping tail was counted once per chunk it appeared in. ``page_number`` is now
derived from the chunk's offset in the source document instead of a running counter, so a repeated
page break cannot be counted twice.
- |
Fixed ``RecursiveDocumentSplitter`` reporting ``split_idx_start`` values past the chunk's real start when
``split_unit="word"`` and ``split_overlap`` was greater than 0, which also pushed ``page_number`` onto pages
the chunk had not reached yet. ``split_idx_start`` now points at the chunk's start in the source document,
and the ``_split_overlap`` ranges cover exactly the text the two chunks share.
- |
Fixed ``RecursiveDocumentSplitter`` dropping the whitespace after the last sentence of each piece of text it
split into sentences, such as the ``"\n\n"`` ending a paragraph longer than ``split_length``. This lost text,
and moved the ``split_idx_start`` of every later chunk before its real start, which could put a chunk on
the page before the one it starts on. The whitespace is now kept at the end of the chunk it follows, so
with no overlap the chunks add up to the document again. Such a chunk can exceed ``split_length`` by that
trailing whitespace.
@@ -444,9 +444,17 @@ class TestSplittingByFunctionOrCharacterRegex:
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, strict=True):
assert doc.meta["page_number"] == p
# The 2nd and 4th chunks open with the page break(s) that end the preceding page, so their text
# starts on the following page.
expected = [
("This is some text.", 1),
("\f This text is on another page.", 2),
("This content has two.", 1),
("\f\f page brakes.", 3),
]
for doc, (content, page) in zip(result["documents"], expected, strict=True):
assert doc.content == content
assert doc.meta["page_number"] == page
def test_add_page_number_to_metadata_with_no_overlap_passage_split(self):
splitter = DocumentSplitter(split_by="passage", split_length=1)
@@ -456,9 +464,16 @@ class TestSplittingByFunctionOrCharacterRegex:
)
result = splitter.run(documents=[doc1])
expected_pages = [1, 2, 2, 2]
for doc, p in zip(result["documents"], expected_pages, strict=True):
assert doc.meta["page_number"] == p
# The last chunk opens with the page break, so its text starts on page 3.
expected = [
("This is a text with some words.\f There is a second sentence.\n\n", 1),
("And there is a third sentence.\n\n", 2),
("And more passages.\n\n", 2),
("\f And another passage.", 3),
]
for doc, (content, page) in zip(result["documents"], expected, strict=True):
assert doc.content == content
assert doc.meta["page_number"] == page
def test_add_page_number_to_metadata_with_no_overlap_page_split(self):
splitter = DocumentSplitter(split_by="page", split_length=1)
@@ -482,6 +497,18 @@ class TestSplittingByFunctionOrCharacterRegex:
for doc, p in zip(result["documents"], expected_pages, strict=True):
assert doc.meta["page_number"] == p
def test_add_page_number_to_metadata_for_blank_pages(self):
splitter = DocumentSplitter(split_by="page", split_length=1)
result = splitter.run(documents=[Document(content="Hello\f\fWorld")])
# The middle chunk is a blank page: it is page 2 itself, not a chunk pushed onto page 3.
assert [doc.content for doc in result["documents"]] == ["Hello\f", "\f", "World"]
assert [doc.meta["page_number"] for doc in result["documents"]] == [1, 2, 3]
result = splitter.run(documents=[Document(content="\fHello")])
assert [doc.content for doc in result["documents"]] == ["\f", "Hello"]
assert [doc.meta["page_number"] for doc in result["documents"]] == [1, 2]
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.")
@@ -501,9 +528,18 @@ class TestSplittingByFunctionOrCharacterRegex:
# No overlap-only trailing chunks: " End." is fully contained in doc1's previous chunk and
# " More text." in doc2's previous chunk, so both are skipped instead of creating redundant
# chunks (the latter even carried a wrong page number).
expected_pages = [1, 1, 1, 1, 1]
for doc, p in zip(result["documents"], expected_pages, strict=True):
assert doc.meta["page_number"] == p
# The 3rd and 5th chunks open with the page break(s) that end the preceding page, so their text
# starts on the following page.
expected = [
("This is some text. And this is more text.", 1),
(" And this is more text.\f This text is on another page.", 1),
("\f This text is on another page. End.", 2),
("This content has two.\f\f page brakes.", 1),
("\f\f page brakes. More text.", 3),
]
for doc, (content, page) in zip(result["documents"], expected, strict=True):
assert doc.content == content
assert doc.meta["page_number"] == page
def test_add_page_number_to_metadata_with_overlap_passage_split(self):
splitter = DocumentSplitter(split_by="passage", split_length=2, split_overlap=1)
@@ -1157,7 +1193,9 @@ class TestSplittingByTokenIntegration:
assert chunk.meta["split_id"] == split_id
start = chunk.meta["split_idx_start"]
assert text[start : start + len(chunk.content)] == chunk.content
assert chunk.meta["page_number"] == 1 + text[:start].count("\f")
# page_number is the page the chunk's text starts on, so breaks it opens with count toward it
leading_breaks = len(chunk.content) - len(chunk.content.lstrip("\f"))
assert chunk.meta["page_number"] == 1 + text[: start + leading_breaks].count("\f")
def test_basic_chunking(self):
splitter = DocumentSplitter(split_by="token", split_length=5, split_overlap=0)
@@ -1223,7 +1261,9 @@ class TestSplittingByTokenIntegration:
assert len(docs) > 1
assert docs[0].meta["page_number"] == 1
for d in docs:
expected_page = 1 + text[: d.meta["split_idx_start"]].count("\f")
assert d.content is not None
leading_breaks = len(d.content) - len(d.content.lstrip("\f"))
expected_page = 1 + text[: d.meta["split_idx_start"] + leading_breaks].count("\f")
assert d.meta["page_number"] == expected_page
assert docs[-1].meta["page_number"] == 3
@@ -1234,6 +1274,8 @@ class TestSplittingByTokenIntegration:
docs = splitter.run(documents=[doc])["documents"]
assert len(docs) > 1
for d in docs:
expected_page = 1 + text[: d.meta["split_idx_start"]].count("\f")
assert d.content is not None
leading_breaks = len(d.content) - len(d.content.lstrip("\f"))
expected_page = 1 + text[: d.meta["split_idx_start"] + leading_breaks].count("\f")
assert d.meta["page_number"] == expected_page
assert docs[-1].meta["page_number"] == 3
@@ -293,9 +293,25 @@ class TestEmbeddingBasedDocumentSplitter:
assert documents[0].content == "Page 1 content.\f"
assert documents[0].meta["page_number"] == 1
assert documents[1].content == "Page 2 content.\f\f\f"
# Should be page 2, not 4, because consecutive page breaks at the end are adjusted
# Page 2, not 4: the breaks end the chunk, so they are not part of the page its text starts on
assert documents[1].meta["page_number"] == 2
def test_create_documents_from_splits_leading_page_breaks(self):
mock_embedder = Mock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
original_doc = Document(content="Page 1 content.\f\fPage 3 content.", meta={"key": "value"})
splits = ["Page 1 content.", "\f\fPage 3 content."]
documents = splitter._create_documents_from_splits(splits, original_doc)
assert len(documents) == 2
assert documents[0].content == "Page 1 content."
assert documents[0].meta["page_number"] == 1
assert documents[1].content == "\f\fPage 3 content."
# The breaks open the chunk, so all of its text is on page 3
assert documents[1].meta["page_number"] == 3
def test_create_documents_from_splits_split_idx_start(self):
"""_create_documents_from_splits must set split_idx_start to the character offset of each chunk."""
mock_embedder = Mock()
@@ -1008,6 +1008,44 @@ def test_custom_page_break_character_in_secondary_splitting():
assert [doc.meta["page_number"] for doc in docs] == [1, 2]
def test_page_number_of_split_starting_with_page_break():
# The break ends the preceding page, so this split's own text starts on the page after it.
text = "# H1\nSentence one.\fSentence two."
splitter = MarkdownHeaderSplitter(secondary_split="period", split_length=1)
docs = splitter.run(documents=[Document(content=text)])["documents"]
assert [doc.content for doc in docs] == ["# H1\nSentence one.", "\fSentence two."]
assert [doc.meta["page_number"] for doc in docs] == [1, 2]
reference = DocumentSplitter(split_by="period", split_length=1).run(documents=[Document(content=text)])["documents"]
assert [doc.meta["page_number"] for doc in docs] == [doc.meta["page_number"] for doc in reference]
def test_page_number_of_secondary_splits_under_a_chunk_starting_with_a_page_break():
# The content before the first header opens with a page break, so that break is counted both by the
# header chunk's own page and by the page-break scan the secondary splits are numbered against.
text = "\fAAA. BBB.\n# H1\nCCC."
docs = MarkdownHeaderSplitter(secondary_split="period", split_length=1).run(documents=[Document(content=text)])[
"documents"
]
assert [doc.content for doc in docs] == ["\fAAA.", " BBB.", "\n", "# H1\nCCC."]
# the text holds a single page break, so everything after it is on page 2 - never page 3
assert [doc.meta["page_number"] for doc in docs] == [2, 2, 2, 2]
def test_page_number_of_split_starting_with_multi_character_page_break():
text = "# H1\nw1 w2.<PAGE>w3 w4."
splitter = MarkdownHeaderSplitter(page_break_character="<PAGE>", secondary_split="period", split_length=1)
docs = splitter.run(documents=[Document(content=text)])["documents"]
assert [doc.content for doc in docs] == ["# H1\nw1 w2.", "<PAGE>w3 w4."]
assert [doc.meta["page_number"] for doc in docs] == [1, 2]
def test_page_break_in_removed_header_is_counted():
text = "# H1<PAGE>\nw1 w2 w3 w4"
splitter = MarkdownHeaderSplitter(
@@ -9,6 +9,7 @@ import pytest
from pytest import LogCaptureFixture
from haystack import Document, Pipeline
from haystack.components.preprocessors.document_splitter import DocumentSplitter
from haystack.components.preprocessors.recursive_splitter import RecursiveDocumentSplitter
from haystack.components.preprocessors.sentence_tokenizer import SentenceSplitter
from haystack.components.retrievers.sentence_window_retriever import SentenceWindowRetriever
@@ -137,7 +138,7 @@ AI technology is widely used throughout industry, government, and science. Some
chunks[1].content
== "AI, in its broadest sense, is intelligence exhibited by machines, particularly computer systems.\n"
)
assert chunks[2].content == "AI technology is widely used throughout industry, government, and science."
assert chunks[2].content == "AI technology is widely used throughout industry, government, and science. "
assert (
chunks[3].content
== "Some high-profile applications include advanced web search engines (e.g., Google Search); recommendation "
@@ -145,6 +146,7 @@ AI technology is widely used throughout industry, government, and science. Some
"Siri, and Alexa); autonomous vehicles (e.g., Waymo); generative and creative tools (e.g., ChatGPT and "
"AI art); and superhuman play and analysis in strategy games (e.g., chess and Go)."
)
assert "".join(chunk.content for chunk in chunks if chunk.content is not None) == text
def test_run_split_by_dot_count_page_breaks_split_unit_char() -> None:
@@ -213,7 +215,7 @@ def test_run_split_by_word_count_page_breaks_split_unit_char():
assert doc_chunks[1].meta["split_idx_start"] == text.index(doc_chunks[1].content)
assert doc_chunks[2].content == "another page. \f "
assert doc_chunks[2].meta["page_number"] == 3
assert doc_chunks[2].meta["page_number"] == 2
assert doc_chunks[2].meta["split_id"] == 2
assert doc_chunks[2].meta["split_idx_start"] == text.index(doc_chunks[2].content)
@@ -365,6 +367,200 @@ def test_run_split_by_sentence_count_page_breaks_split_unit_char() -> None:
assert chunks_docs[6].meta["split_idx_start"] == text.index(chunks_docs[6].content)
def test_run_page_number_is_the_page_the_chunk_starts_on():
text = "aa bb\fcc dd ee ff\fgg hh"
splitter = RecursiveDocumentSplitter(split_length=4, split_overlap=0, separators=[" "], split_unit="word")
chunks = splitter.run(documents=[Document(content=text)])["documents"]
contents: list[str] = []
for chunk in chunks:
assert chunk.content is not None
contents.append(chunk.content)
assert contents == ["aa bb\fcc dd ee ", "ff\fgg hh"]
assert "".join(contents) == text
assert [chunk.meta["split_idx_start"] for chunk in chunks] == [0, 15]
# Both chunks straddle a page break and are reported on the page their text starts on.
assert [chunk.meta["page_number"] for chunk in chunks] == [1, 2]
# DocumentSplitter defines the convention; the two must not drift apart.
reference = DocumentSplitter(split_by="word", split_length=4, split_overlap=0).run(
documents=[Document(content=text)]
)["documents"]
assert [chunk.content for chunk in chunks] == [doc.content for doc in reference]
assert [chunk.meta["page_number"] for chunk in chunks] == [doc.meta["page_number"] for doc in reference]
def test_run_page_number_does_not_drift_with_overlap():
# The page break sits inside the overlapping tail, so it appears in two consecutive chunks.
text = "This is page one.\fThis is page two, it is longer."
splitter = RecursiveDocumentSplitter(split_length=20, split_overlap=5, separators=["\n"], split_unit="char")
chunks = splitter.run(documents=[Document(content=text)])["documents"]
assert [chunk.content for chunk in chunks] == [
"This is page one.\fTh",
"e.\fThis is page two,",
" two, it is longer.",
]
assert [chunk.meta["split_idx_start"] for chunk in chunks] == [0, 15, 30]
# the text contains one page break, so it spans two pages
assert [chunk.meta["page_number"] for chunk in chunks] == [1, 1, 2]
@pytest.mark.parametrize(
"text, split_length, split_overlap, expected_contents, expected_pages",
[
pytest.param(
# every chunk is exactly split_length words, and the overlap is followed by a page break, not a space
"\f".join(f"p{page}a p{page}b p{page}c." for page in range(1, 9)),
4,
1,
["p1a p1b p1c."] + [f"p{page}c.\fp{page + 1}a p{page + 1}b p{page + 1}c." for page in range(1, 8)],
[1, 1, 2, 3, 4, 5, 6, 7],
id="page_break_after_overlap",
),
pytest.param(
# the second chunk is trimmed after "ee", and the page break it holds stays in it
"aa bb cc.\fdd ee ff.",
3,
1,
["aa bb cc.", "cc.\fdd ee", "ee ff."],
[1, 1, 2],
id="trimmed_chunk_with_page_break",
),
pytest.param(
# no separator matches, so _fall_back_to_fixed_chunking cuts the text before the overlap is applied
"aa bb\fcc dd\fee ff gg\fhh",
3,
1,
["aa bb\fcc", "bb\fcc dd\fee ff", "ff gg\fhh"],
[1, 1, 3],
id="fixed_size_fallback_with_page_breaks",
),
],
)
def test_run_word_overlap_keeps_source_text_and_page_numbers(
text, split_length, split_overlap, expected_contents, expected_pages
):
splitter = RecursiveDocumentSplitter(
split_length=split_length, split_overlap=split_overlap, separators=["."], split_unit="word"
)
chunks = splitter.run(documents=[Document(content=text)])["documents"]
assert [chunk.content for chunk in chunks] == expected_contents
assert [chunk.meta["page_number"] for chunk in chunks] == expected_pages
reconstructed = ""
for split_id, chunk in enumerate(chunks):
assert chunk.content is not None
assert chunk.meta["split_id"] == split_id
start = chunk.meta["split_idx_start"]
assert text[start : start + len(chunk.content)] == chunk.content
reconstructed += chunk.content[len(reconstructed) - start :]
assert reconstructed == text
for prev_chunk, chunk in zip(chunks, chunks[1:], strict=False):
assert prev_chunk.content is not None
overlap_length = prev_chunk.meta["split_idx_start"] + len(prev_chunk.content) - chunk.meta["split_idx_start"]
assert {"doc_id": chunk.id, "range": (0, overlap_length)} in prev_chunk.meta["_split_overlap"]
prev_range = (len(prev_chunk.content) - overlap_length, len(prev_chunk.content))
assert {"doc_id": prev_chunk.id, "range": prev_range} in chunk.meta["_split_overlap"]
def test_run_sentence_separator_keeps_whitespace_after_last_sentence():
# Both paragraphs are longer than split_length, so the default separators split them into sentences, and the
# sentence tokenizer leaves out the whitespace after the last sentence of each. Losing the first "\n\n" would
# shift the chunks after it before the page break, onto page 1.
text = "One two three. Four five six.\n\nalpha.\fomega omega omega omega omega.\n\n"
splitter = RecursiveDocumentSplitter(split_length=4, split_unit="word")
chunks = splitter.run(documents=[Document(content=text)])["documents"]
contents: list[str] = []
for chunk in chunks:
assert chunk.content is not None
contents.append(chunk.content)
assert contents == ["One two three. ", "Four five six.\n\n", "alpha.\f", "omega omega omega omega ", "omega.\n\n"]
assert "".join(contents) == text
assert [chunk.meta["split_idx_start"] for chunk in chunks] == [0, 15, 31, 38, 62]
assert [chunk.meta["page_number"] for chunk in chunks] == [1, 1, 1, 2, 2]
@pytest.mark.parametrize(
"split_unit, split_length, split_overlap, expected_contents, expected_starts, expected_pages",
[
pytest.param(
"word",
6,
0,
[
"One two three. Four five six.\f",
"Seven eight nine. Ten eleven twelve.\f",
"Thirteen fourteen fifteen.",
],
[0, 32, 71],
[1, 2, 3],
id="word_no_overlap",
),
pytest.param(
"word",
6,
2,
[
"One two three. Four five six.\f",
"five six.\fSeven eight nine. Ten",
"nine. Ten eleven twelve.\fThirteen fourteen",
"Thirteen fourteen fifteen.",
],
[0, 22, 44, 71],
[1, 1, 2, 3],
id="word_overlap",
),
pytest.param(
"char",
40,
10,
[
"One two three. Four five six.\f",
"five six.\fSeven eight nine. Ten eleven",
"Ten eleven twelve.\fThirteen fourteen fif",
"urteen fifteen.",
],
[0, 22, 52, 82],
[1, 1, 2, 3],
id="char_overlap",
),
],
)
def test_run_sentence_separator_without_white_spaces_keeps_text_between_sentences(
split_unit, split_length, split_overlap, expected_contents, expected_starts, expected_pages
):
# With keep_white_spaces=False the sentence tokenizer leaves out the whitespace between sentences, page breaks
# included. The chunks still keep it, rather than gluing their sentences together as "three.Four five six."
text = "One two three. Four five six.\fSeven eight nine. Ten eleven twelve.\fThirteen fourteen fifteen."
splitter = RecursiveDocumentSplitter(
split_length=split_length,
split_overlap=split_overlap,
split_unit=split_unit,
separators=["sentence"],
sentence_splitter_params={"keep_white_spaces": False},
)
chunks = splitter.run(documents=[Document(content=text)])["documents"]
assert [chunk.content for chunk in chunks] == expected_contents
assert [chunk.meta["split_id"] for chunk in chunks] == list(range(len(expected_contents)))
assert [chunk.meta["split_idx_start"] for chunk in chunks] == expected_starts
assert [chunk.meta["page_number"] for chunk in chunks] == expected_pages
for chunk in chunks:
assert chunk.content is not None
start = chunk.meta["split_idx_start"]
assert text[start : start + len(chunk.content)] == chunk.content
def test_run_split_document_with_overlap_character_unit():
splitter = RecursiveDocumentSplitter(split_length=20, split_overlap=10, separators=["."], split_unit="char")
text = """A simple sentence1. A bright sentence2. A clever sentence3"""
@@ -563,7 +759,7 @@ def test_run_split_by_word_count_page_breaks_word_unit():
assert doc_chunks[1].meta["split_idx_start"] == text.index(doc_chunks[1].content)
assert doc_chunks[2].content == "on another page. \f "
assert doc_chunks[2].meta["page_number"] == 3
assert doc_chunks[2].meta["page_number"] == 2
assert doc_chunks[2].meta["split_id"] == 2
assert doc_chunks[2].meta["split_idx_start"] == text.index(doc_chunks[2].content)
@@ -1145,9 +1341,6 @@ def test_run_multiple_separators_with_overlap_applies_overlap_only_once():
def test_run_multiple_separators_with_overlap_keeps_chunks_contiguous(split_unit, split_length, split_overlap):
"""
The overlap must be applied only once for every split unit, not just for "char".
Whitespace is normalised before comparing because the "word" and "token" units rejoin the
overlap with a single space in ``_create_chunk_starting_with_overlap``.
"""
text = MULTI_SEPARATOR_TEXT
splitter = RecursiveDocumentSplitter(
@@ -1156,10 +1349,9 @@ def test_run_multiple_separators_with_overlap_keeps_chunks_contiguous(split_unit
chunks = splitter.run([Document(content=text)])["documents"]
assert len(chunks) > 1, "the text must actually be split, otherwise the overlap is never applied"
normalized_text = " ".join(text.split())
for chunk in chunks:
assert chunk.content is not None
assert " ".join(chunk.content.split()) in normalized_text
assert text[chunk.meta["split_idx_start"] :].startswith(chunk.content)
assert "Overview\nOverview" not in chunk.content