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