Coverage for haystack/components/preprocessors/embedding_based_document_splitter.py: 95%
230 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-10-09 14:34 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5from asyncio import gather
6from collections.abc import Awaitable
7from copy import deepcopy
8from itertools import chain
9from typing import Any
11import numpy as np
13from haystack import Document, component, logging
14from haystack.components.embedders.types import DocumentEmbedder
15from haystack.components.preprocessors._page_numbers import _leading_page_breaks
16from haystack.components.preprocessors.sentence_tokenizer import Language, SentenceSplitter
17from haystack.core.serialization import component_to_dict, default_from_dict, default_to_dict
18from haystack.utils.async_utils import _execute_component_async
19from haystack.utils.deserialization import deserialize_component_inplace
21logger = logging.getLogger(__name__)
24@component
25class EmbeddingBasedDocumentSplitter:
26 """
27 Splits documents based on embedding similarity using cosine distances between sequential sentence groups.
29 This component first splits text into sentences, optionally groups them, calculates embeddings for each group,
30 and then uses cosine distance between sequential embeddings to determine split points. Any distance above
31 the specified percentile is treated as a break point. The component also tracks page numbers based on form feed
32 characters (`\f`) in the original document.
34 This component is inspired by [5 Levels of Text Splitting](
35 https://github.com/FullStackRetrieval-com/RetrievalTutorials/blob/main/tutorials/LevelsOfTextSplitting/5_Levels_Of_Text_Splitting.ipynb
36 ) by Greg Kamradt.
38 ### Usage example
40 ```python
41 from haystack import Document
42 from haystack.components.embedders import OpenAIDocumentEmbedder
43 from haystack.components.preprocessors import EmbeddingBasedDocumentSplitter
45 # Create a document with content that has a clear topic shift
46 doc = Document(
47 content="This is a first sentence. This is a second sentence. This is a third sentence. "
48 "Completely different topic. The same completely different topic."
49 )
51 # Initialize the embedder to calculate semantic similarities
52 embedder = OpenAIDocumentEmbedder()
54 # Configure the splitter with parameters that control splitting behavior
55 splitter = EmbeddingBasedDocumentSplitter(
56 document_embedder=embedder,
57 sentences_per_group=2, # Group 2 sentences before calculating embeddings
58 percentile=0.95, # Split when cosine distance exceeds 95th percentile
59 min_length=50, # Merge splits shorter than 50 characters
60 max_length=1000 # Further split chunks longer than 1000 characters
61 )
62 result = splitter.run(documents=[doc])
64 # The result contains a list of Document objects, each representing a semantic chunk
65 # Each split document includes metadata: source_id, split_id, and page_number
66 print(f"Original document split into {len(result['documents'])} chunks")
67 for i, split_doc in enumerate(result['documents']):
68 print(f"Chunk {i}: {split_doc.content[:50]}...")
69 ```
70 """
72 def __init__(
73 self,
74 *,
75 document_embedder: DocumentEmbedder,
76 sentences_per_group: int = 3,
77 percentile: float = 0.95,
78 min_length: int = 50,
79 max_length: int = 1000,
80 language: Language = "en",
81 use_split_rules: bool = True,
82 extend_abbreviations: bool = True,
83 ) -> None:
84 """
85 Initialize EmbeddingBasedDocumentSplitter.
87 :param document_embedder: The DocumentEmbedder to use for calculating embeddings.
88 :param sentences_per_group: Number of sentences to group together before embedding.
89 :param percentile: Percentile threshold for cosine distance. Distances above this percentile
90 are treated as break points.
91 :param min_length: Minimum length of splits in characters. Splits below this length will be merged.
92 :param max_length: Maximum length of splits in characters. Splits above this length will be recursively split.
93 :param language: Language for sentence tokenization.
94 :param use_split_rules: Whether to use additional split rules for sentence tokenization. Applies additional
95 split rules from SentenceSplitter to the sentence spans.
96 :param extend_abbreviations: If True, the abbreviations used by NLTK's PunktTokenizer are extended by a list
97 of curated abbreviations. Currently supported languages are: en, de.
98 If False, the default abbreviations are used.
99 """
100 self.document_embedder = document_embedder
102 if sentences_per_group <= 0:
103 raise ValueError("sentences_per_group must be greater than 0.")
104 self.sentences_per_group = sentences_per_group
106 if not 0.0 <= percentile <= 1.0:
107 raise ValueError("percentile must be between 0.0 and 1.0.")
108 self.percentile = percentile
110 if min_length < 0:
111 raise ValueError("min_length must be greater than or equal to 0.")
112 self.min_length = min_length
114 if max_length <= min_length:
115 raise ValueError("max_length must be greater than min_length.")
116 self.max_length = max_length
118 self.language = language
119 self.use_split_rules = use_split_rules
120 self.extend_abbreviations = extend_abbreviations
121 self.sentence_splitter: SentenceSplitter | None = None
123 def warm_up(self) -> None:
124 """
125 Warm up the component by initializing the sentence splitter and the document embedder.
126 """
127 if self.sentence_splitter is None:
128 self.sentence_splitter = SentenceSplitter(
129 language=self.language,
130 use_split_rules=self.use_split_rules,
131 extend_abbreviations=self.extend_abbreviations,
132 keep_white_spaces=True,
133 )
134 if hasattr(self.document_embedder, "warm_up"):
135 self.document_embedder.warm_up()
137 async def warm_up_async(self) -> None:
138 """
139 Warm up the component on the serving event loop.
141 Initializes the sentence splitter and warms up the document embedder using its async warm-up path when
142 available, falling back to the synchronous one otherwise.
143 """
144 if self.sentence_splitter is None:
145 self.sentence_splitter = SentenceSplitter(
146 language=self.language,
147 use_split_rules=self.use_split_rules,
148 extend_abbreviations=self.extend_abbreviations,
149 keep_white_spaces=True,
150 )
151 if hasattr(self.document_embedder, "warm_up_async"):
152 await self.document_embedder.warm_up_async()
153 elif hasattr(self.document_embedder, "warm_up"):
154 self.document_embedder.warm_up()
156 def close(self) -> None:
157 """
158 Release the document embedder's resources.
159 """
160 if hasattr(self.document_embedder, "close"):
161 self.document_embedder.close()
163 async def close_async(self) -> None:
164 """
165 Release the document embedder's async resources.
166 """
167 if hasattr(self.document_embedder, "close_async"):
168 await self.document_embedder.close_async()
169 elif hasattr(self.document_embedder, "close"):
170 self.document_embedder.close()
172 @component.output_types(documents=list[Document])
173 def run(self, documents: list[Document]) -> dict[str, list[Document]]:
174 """
175 Split documents based on embedding similarity.
177 :param documents: The documents to split.
178 :returns: A dictionary with the following key:
179 - `documents`: List of documents with the split texts. Each document includes:
180 - A metadata field `source_id` to track the original document.
181 - A metadata field `split_id` to track the split number.
182 - A metadata field `split_idx_start` with the character offset of the chunk in the original document.
183 - A metadata field `page_number` with the page the chunk starts on, counting form feed
184 ("\f") characters in the original document.
185 - All other metadata copied from the original document.
187 :raises RuntimeError: If the component wasn't warmed up.
188 :raises TypeError: If the input is not a list of Documents.
189 :raises ValueError: If the document content is None or empty.
190 """
191 self.warm_up()
193 if not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
194 raise TypeError("EmbeddingBasedDocumentSplitter expects a List of Documents as input.")
196 split_docs: list[Document] = []
197 for doc in documents:
198 if doc.content is None:
199 raise ValueError(
200 f"EmbeddingBasedDocumentSplitter only works with text documents but content for "
201 f"document ID {doc.id} is None."
202 )
203 if doc.content == "":
204 logger.warning("Document ID {doc_id} has an empty content. Skipping this document.", doc_id=doc.id)
205 continue
207 doc_splits = self._split_document(doc=doc)
208 split_docs.extend(doc_splits)
210 return {"documents": split_docs}
212 @component.output_types(documents=list[Document])
213 async def run_async(self, documents: list[Document]) -> dict[str, list[Document]]:
214 """
215 Asynchronously split documents based on embedding similarity.
217 This is the asynchronous version of the `run` method with the same parameters and return values.
219 :param documents: The documents to split.
220 :returns: A dictionary with the following key:
221 - `documents`: List of documents with the split texts. Each document includes:
222 - A metadata field `source_id` to track the original document.
223 - A metadata field `split_id` to track the split number.
224 - A metadata field `split_idx_start` with the character offset of the chunk in the original document.
225 - A metadata field `page_number` with the page the chunk starts on, counting form feed
226 ("\f") characters in the original document.
227 - All other metadata copied from the original document.
229 :raises RuntimeError: If the component wasn't warmed up.
230 :raises TypeError: If the input is not a list of Documents.
231 :raises ValueError: If the document content is None or empty.
232 """
233 await self.warm_up_async()
235 if not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
236 raise TypeError("EmbeddingBasedDocumentSplitter expects a List of Documents as input.")
238 tasks: list[Awaitable[list[Document]]] = []
239 for doc in documents:
240 if doc.content is None:
241 raise ValueError(
242 f"EmbeddingBasedDocumentSplitter only works with text documents but content for "
243 f"document ID {doc.id} is None."
244 )
245 if doc.content == "":
246 logger.warning("Document ID {doc_id} has an empty content. Skipping this document.", doc_id=doc.id)
247 continue
249 tasks.append(self._split_document_async(doc=doc))
251 return {"documents": [*chain.from_iterable(await gather(*tasks))]}
253 def _split_document(self, doc: Document) -> list[Document]:
254 """
255 Split a single document based on embedding similarity.
256 """
257 # Create an initial split of the document content into smaller chunks
258 # doc.content is validated in `run`
259 splits = self._split_text(text=doc.content) # type: ignore[arg-type]
261 # Merge splits smaller than min_length
262 merged_splits = self._merge_small_splits(splits=splits)
264 # Recursively split splits larger than max_length
265 final_splits = self._split_large_splits(splits=merged_splits)
267 # Create Document objects from the final splits
268 return EmbeddingBasedDocumentSplitter._create_documents_from_splits(splits=final_splits, original_doc=doc)
270 async def _split_document_async(self, doc: Document) -> list[Document]:
271 """
272 Split a single document based on embedding similarity.
273 """
274 # Create an initial split of the document content into smaller chunks
275 # doc.content is validated in `run`
276 splits = await self._split_text_async(text=doc.content) # type: ignore[arg-type]
278 # Merge splits smaller than min_length
279 merged_splits = self._merge_small_splits(splits=splits)
281 # Recursively split splits larger than max_length
282 final_splits = await self._split_large_splits_async(splits=merged_splits)
284 # Create Document objects from the final splits
285 return EmbeddingBasedDocumentSplitter._create_documents_from_splits(splits=final_splits, original_doc=doc)
287 def _prepare_sentence_groups(self, text: str) -> list[str]:
288 """Preprocess raw text into grouped sentences ready for embedding."""
289 # NOTE: `self.sentence_splitter.split_sentences` strips all white space types (e.g. new lines, page breaks,
290 # etc.) at the end of the provided text. So to not lose them, we need keep track of them and add them back to
291 # the last sentence.
292 rstripped_text = text.rstrip()
293 trailing_whitespaces = text[len(rstripped_text) :]
295 # Split the text into sentences
296 sentences_result = self.sentence_splitter.split_sentences(rstripped_text) # type: ignore[union-attr]
298 # Add back the stripped white spaces to the last sentence
299 if sentences_result and trailing_whitespaces:
300 sentences_result[-1]["sentence"] += trailing_whitespaces
301 sentences_result[-1]["end"] += len(trailing_whitespaces)
303 sentences = [sentence["sentence"] for sentence in sentences_result]
304 return self._group_sentences(sentences=sentences)
306 def _split_text(self, text: str) -> list[str]:
307 """
308 Split a text into smaller chunks based on embedding similarity.
309 """
310 sentence_groups = self._prepare_sentence_groups(text=text)
311 embeddings = self._calculate_embeddings(sentence_groups=sentence_groups)
312 split_points = self._find_split_points(embeddings=embeddings)
313 return self._create_splits_from_points(sentence_groups=sentence_groups, split_points=split_points)
315 async def _split_text_async(self, text: str) -> list[str]:
316 """
317 Asynchronously split a text into smaller chunks based on embedding similarity.
318 """
319 sentence_groups = self._prepare_sentence_groups(text=text)
320 embeddings = await self._calculate_embeddings_async(sentence_groups=sentence_groups)
321 split_points = self._find_split_points(embeddings=embeddings)
322 return self._create_splits_from_points(sentence_groups=sentence_groups, split_points=split_points)
324 def _group_sentences(self, sentences: list[str]) -> list[str]:
325 """
326 Group sentences into groups of sentences_per_group.
327 """
328 if self.sentences_per_group == 1:
329 return sentences
331 groups = []
332 for i in range(0, len(sentences), self.sentences_per_group):
333 group = sentences[i : i + self.sentences_per_group]
334 groups.append("".join(group))
336 return groups
338 def _calculate_embeddings(self, sentence_groups: list[str]) -> list[list[float]]:
339 """
340 Calculate embeddings for each sentence group using the DocumentEmbedder.
341 """
342 # Create Document objects for each group
343 group_docs = [Document(content=group) for group in sentence_groups]
344 result = self.document_embedder.run(group_docs)
345 embedded_docs = result["documents"]
346 return [doc.embedding for doc in embedded_docs]
348 async def _calculate_embeddings_async(self, sentence_groups: list[str]) -> list[list[float]]:
349 """
350 Asynchronously Calculate embeddings for each sentence group using the DocumentEmbedder.
351 """
352 # Create Document objects for each group
353 group_docs = [Document(content=group) for group in sentence_groups]
354 result = await _execute_component_async(self.document_embedder, documents=group_docs)
355 embedded_docs = result["documents"]
356 return [doc.embedding for doc in embedded_docs]
358 def _find_split_points(self, embeddings: list[list[float]]) -> list[int]:
359 """
360 Find split points based on cosine distances between sequential embeddings.
361 """
362 if len(embeddings) <= 1:
363 return []
365 # Calculate cosine distances between sequential pairs
366 distances = []
367 for i in range(len(embeddings) - 1):
368 distance = EmbeddingBasedDocumentSplitter._cosine_distance(
369 embedding1=embeddings[i], embedding2=embeddings[i + 1]
370 )
371 distances.append(distance)
373 # Calculate threshold based on percentile
374 threshold = np.percentile(distances, self.percentile * 100)
376 # Find indices where distance exceeds threshold
377 split_points = []
378 for i, distance in enumerate(distances):
379 if distance > threshold:
380 split_points.append(i + 1) # +1 because we want to split after this point
382 return split_points
384 @staticmethod
385 def _cosine_distance(embedding1: list[float], embedding2: list[float]) -> float:
386 """
387 Calculate cosine distance between two embeddings.
388 """
389 vec1 = np.array(embedding1)
390 vec2 = np.array(embedding2)
392 norm1 = float(np.linalg.norm(vec1))
393 norm2 = float(np.linalg.norm(vec2))
395 if norm1 == 0 or norm2 == 0:
396 return 1.0
398 cosine_sim = float(np.dot(vec1, vec2) / (norm1 * norm2))
400 return 1.0 - cosine_sim
402 @staticmethod
403 def _create_splits_from_points(sentence_groups: list[str], split_points: list[int]) -> list[str]:
404 """
405 Create splits based on split points.
406 """
407 if not split_points:
408 return ["".join(sentence_groups)]
410 splits = []
411 start = 0
413 for point in split_points:
414 split_text = "".join(sentence_groups[start:point])
415 if split_text:
416 splits.append(split_text)
417 start = point
419 # Add the last split
420 if start < len(sentence_groups):
421 split_text = "".join(sentence_groups[start:])
422 if split_text:
423 splits.append(split_text)
425 return splits
427 def _merge_small_splits(self, splits: list[str]) -> list[str]:
428 """
429 Merge splits that are below min_length.
430 """
431 if not splits:
432 return splits
434 merged = []
435 current_split = splits[0]
437 for split in splits[1:]:
438 # We merge splits that are smaller than min_length but only if the newly merged split is still below
439 # max_length.
440 if len(current_split) < self.min_length and len(current_split) + len(split) < self.max_length:
441 # Merge with next split
442 current_split += split
443 else:
444 # Current split is long enough, save it and start a new one
445 merged.append(current_split)
446 current_split = split
448 # Don't forget the last split
449 merged.append(current_split)
451 # The loop only merges forward, so the final split can still be below min_length. Merge it backwards,
452 # subject to the same max_length limit as forward merges.
453 if (
454 len(merged) > 1
455 and len(merged[-1]) < self.min_length
456 and len(merged[-2]) + len(merged[-1]) < self.max_length
457 ):
458 trailing_split = merged.pop()
459 merged[-1] += trailing_split
461 return merged
463 def _split_large_splits(self, splits: list[str]) -> list[str]:
464 """
465 Recursively split splits that are above max_length.
467 This method checks each split and if it exceeds max_length, it attempts to split it further using the same
468 embedding-based approach. This is done recursively until all splits are within the max_length limit or no
469 further splitting is possible.
471 This works because the threshold for splits is calculated dynamically based on the provided of embeddings.
473 Keep in sync with `_split_large_splits_async`.
474 """
475 final_splits = []
477 for split in splits:
478 if len(split) <= self.max_length:
479 final_splits.append(split)
480 else:
481 # Recursively split large splits
482 # We can reuse the same _split_text method to split the text into smaller chunks because the threshold
483 # for splits is calculated dynamically based on embeddings from `split`.
484 sub_splits = self._split_text(text=split)
486 # Stop splitting if no further split is possible or continue with recursion
487 if len(sub_splits) == 1:
488 logger.warning(
489 "Could not split a chunk further below max_length={max_length}. "
490 "Returning chunk of length {length}.",
491 max_length=self.max_length,
492 length=len(split),
493 )
494 final_splits.append(split)
495 else:
496 final_splits.extend(self._split_large_splits(splits=sub_splits))
498 return final_splits
500 async def _split_large_splits_async(self, splits: list[str]) -> list[str]:
501 """
502 Asynchronously and recursively split splits that are above max_length.
504 Mirrors `_split_large_splits`, but embeds through the async path so that the recursion does not block the
505 event loop with synchronous embedder calls.
507 Keep in sync with `_split_large_splits`, which also documents why re-running the splitter on an oversized
508 chunk can split it further.
509 """
510 final_splits = []
512 for split in splits:
513 if len(split) <= self.max_length:
514 final_splits.append(split)
515 else:
516 # Recursively split large splits
517 # We can reuse the same _split_text_async method to split the text into smaller chunks because the
518 # threshold for splits is calculated dynamically based on embeddings from `split`.
519 sub_splits = await self._split_text_async(text=split)
521 # Stop splitting if no further split is possible or continue with recursion
522 if len(sub_splits) == 1:
523 logger.warning(
524 "Could not split a chunk further below max_length={max_length}. "
525 "Returning chunk of length {length}.",
526 max_length=self.max_length,
527 length=len(split),
528 )
529 final_splits.append(split)
530 else:
531 final_splits.extend(await self._split_large_splits_async(splits=sub_splits))
533 return final_splits
535 @staticmethod
536 def _create_documents_from_splits(splits: list[str], original_doc: Document) -> list[Document]:
537 """
538 Create Document objects from splits.
539 """
540 documents = []
541 metadata = deepcopy(original_doc.meta)
542 metadata["source_id"] = original_doc.id
544 # Track page number and character offset across splits.
545 # Chunks are contiguous substrings of the original text (all joins are pure string concatenation),
546 # so the character offset is simply accumulated by adding the length of each chunk.
547 current_page = 1
548 current_char_pos = 0
550 for i, split_text in enumerate(splits):
551 split_meta = deepcopy(metadata)
552 split_meta["split_id"] = i
553 split_meta["split_idx_start"] = current_char_pos
555 page_breaks_in_split = split_text.count("\f")
557 # current_page is the page this split's first character is on; a split that opens with page
558 # breaks starts its text on a later page.
559 split_meta["page_number"] = current_page + _leading_page_breaks(split_text)
561 doc = Document(content=split_text, meta=split_meta)
562 documents.append(doc)
564 # Advance position and page counter for the next split
565 current_char_pos += len(split_text)
566 current_page += page_breaks_in_split
568 return documents
570 def to_dict(self) -> dict[str, Any]:
571 """
572 Serializes the component to a dictionary.
574 :returns:
575 Serialized dictionary representation of the component.
576 """
577 return default_to_dict(
578 self,
579 document_embedder=component_to_dict(obj=self.document_embedder, name="document_embedder"),
580 sentences_per_group=self.sentences_per_group,
581 percentile=self.percentile,
582 min_length=self.min_length,
583 max_length=self.max_length,
584 language=self.language,
585 use_split_rules=self.use_split_rules,
586 extend_abbreviations=self.extend_abbreviations,
587 )
589 @classmethod
590 def from_dict(cls, data: dict[str, Any]) -> "EmbeddingBasedDocumentSplitter":
591 """
592 Deserializes the component from a dictionary.
594 :param data:
595 The dictionary to deserialize and create the component.
597 :returns:
598 The deserialized component.
599 """
600 deserialize_component_inplace(data["init_parameters"], key="document_embedder")
601 return default_from_dict(cls, data)