mirror of
https://github.com/deepset-ai/haystack.git
synced 2026-10-12 02:12:50 +00:00
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:
co-authored by
Sebastian Husch Lee
parent
10e49fa904
commit
472da278db
@@ -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:
|
||||
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user