chore: import upstream snapshot with attribution

haystack/components/preprocessors/embedding_based_document_splitter.py

@@ -0,0 +1,545 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from asyncio import gather
+from collections.abc import Awaitable
+from copy import deepcopy
+from itertools import chain
+from typing import Any
+
+import numpy as np
+
+from haystack import Document, component, logging
+from haystack.components.embedders.types import DocumentEmbedder
+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
+from haystack.utils.deserialization import deserialize_component_inplace
+
+logger = logging.getLogger(__name__)
+
+
+@component
+class EmbeddingBasedDocumentSplitter:
+    """
+    Splits documents based on embedding similarity using cosine distances between sequential sentence groups.
+
+    This component 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.
+
+    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.
+
+    ### Usage example
+
+    ```python
+    from haystack import Document
+    from haystack.components.embedders import OpenAIDocumentEmbedder
+    from haystack.components.preprocessors import EmbeddingBasedDocumentSplitter
+
+    # Create a document with content that has a clear topic shift
+    doc = Document(
+        content="This is a first sentence. This is a second sentence. This is a third sentence. "
+        "Completely different topic. The same completely different topic."
+    )
+
+    # Initialize the embedder to calculate semantic similarities
+    embedder = OpenAIDocumentEmbedder()
+
+    # Configure the splitter with parameters that control splitting behavior
+    splitter = EmbeddingBasedDocumentSplitter(
+        document_embedder=embedder,
+        sentences_per_group=2,      # Group 2 sentences before calculating embeddings
+        percentile=0.95,            # Split when cosine distance exceeds 95th percentile
+        min_length=50,              # Merge splits shorter than 50 characters
+        max_length=1000             # Further split chunks longer than 1000 characters
+    )
+    result = splitter.run(documents=[doc])
+
+    # The result contains a list of Document objects, each representing a semantic chunk
+    # Each split document includes metadata: source_id, split_id, and page_number
+    print(f"Original document split into {len(result['documents'])} chunks")
+    for i, split_doc in enumerate(result['documents']):
+        print(f"Chunk {i}: {split_doc.content[:50]}...")
+    ```
+    """
+
+    def __init__(
+        self,
+        *,
+        document_embedder: DocumentEmbedder,
+        sentences_per_group: int = 3,
+        percentile: float = 0.95,
+        min_length: int = 50,
+        max_length: int = 1000,
+        language: Language = "en",
+        use_split_rules: bool = True,
+        extend_abbreviations: bool = True,
+    ) -> None:
+        """
+        Initialize EmbeddingBasedDocumentSplitter.
+
+        :param document_embedder: The DocumentEmbedder to use for calculating embeddings.
+        :param sentences_per_group: Number of sentences to group together before embedding.
+        :param percentile: Percentile threshold for cosine distance. Distances above this percentile
+            are treated as break points.
+        :param min_length: Minimum length of splits in characters. Splits below this length will be merged.
+        :param max_length: Maximum length of splits in characters. Splits above this length will be recursively split.
+        :param language: Language for sentence tokenization.
+        :param use_split_rules: Whether to use additional split rules for sentence tokenization. Applies additional
+            split rules from SentenceSplitter to the sentence spans.
+        :param extend_abbreviations: If True, the abbreviations used by NLTK's PunktTokenizer are extended by a list
+            of curated abbreviations. Currently supported languages are: en, de.
+            If False, the default abbreviations are used.
+        """
+        self.document_embedder = document_embedder
+
+        if sentences_per_group <= 0:
+            raise ValueError("sentences_per_group must be greater than 0.")
+        self.sentences_per_group = sentences_per_group
+
+        if not 0.0 <= percentile <= 1.0:
+            raise ValueError("percentile must be between 0.0 and 1.0.")
+        self.percentile = percentile
+
+        if min_length < 0:
+            raise ValueError("min_length must be greater than or equal to 0.")
+        self.min_length = min_length
+
+        if max_length <= min_length:
+            raise ValueError("max_length must be greater than min_length.")
+        self.max_length = max_length
+
+        self.language = language
+        self.use_split_rules = use_split_rules
+        self.extend_abbreviations = extend_abbreviations
+        self.sentence_splitter: SentenceSplitter | None = None
+
+    def warm_up(self) -> None:
+        """
+        Warm up the component by initializing the sentence splitter and the document embedder.
+        """
+        if self.sentence_splitter is None:
+            self.sentence_splitter = SentenceSplitter(
+                language=self.language,
+                use_split_rules=self.use_split_rules,
+                extend_abbreviations=self.extend_abbreviations,
+                keep_white_spaces=True,
+            )
+        if hasattr(self.document_embedder, "warm_up"):
+            self.document_embedder.warm_up()
+
+    async def warm_up_async(self) -> None:
+        """
+        Warm up the component on the serving event loop.
+
+        Initializes the sentence splitter and warms up the document embedder using its async warm-up path when
+        available, falling back to the synchronous one otherwise.
+        """
+        if self.sentence_splitter is None:
+            self.sentence_splitter = SentenceSplitter(
+                language=self.language,
+                use_split_rules=self.use_split_rules,
+                extend_abbreviations=self.extend_abbreviations,
+                keep_white_spaces=True,
+            )
+        if hasattr(self.document_embedder, "warm_up_async"):
+            await self.document_embedder.warm_up_async()
+        elif hasattr(self.document_embedder, "warm_up"):
+            self.document_embedder.warm_up()
+
+    def close(self) -> None:
+        """
+        Release the document embedder's resources.
+        """
+        if hasattr(self.document_embedder, "close"):
+            self.document_embedder.close()
+
+    async def close_async(self) -> None:
+        """
+        Release the document embedder's async resources.
+        """
+        if hasattr(self.document_embedder, "close_async"):
+            await self.document_embedder.close_async()
+        elif hasattr(self.document_embedder, "close"):
+            self.document_embedder.close()
+
+    @component.output_types(documents=list[Document])
+    def run(self, documents: list[Document]) -> dict[str, list[Document]]:
+        """
+        Split documents based on embedding similarity.
+
+        :param documents: The documents to split.
+        :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 `split_id` to track the split number.
+                - A metadata field `page_number` to track the original page number.
+                - All other metadata copied from the original document.
+
+        :raises RuntimeError: If the component wasn't warmed up.
+        :raises TypeError: If the input is not a list of Documents.
+        :raises ValueError: If the document content is None or empty.
+        """
+        self.warm_up()
+
+        if not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
+            raise TypeError("EmbeddingBasedDocumentSplitter expects a List of Documents as input.")
+
+        split_docs: list[Document] = []
+        for doc in documents:
+            if doc.content is None:
+                raise ValueError(
+                    f"EmbeddingBasedDocumentSplitter only works with text documents but content for "
+                    f"document ID {doc.id} is None."
+                )
+            if doc.content == "":
+                logger.warning("Document ID {doc_id} has an empty content. Skipping this document.", doc_id=doc.id)
+                continue
+
+            doc_splits = self._split_document(doc=doc)
+            split_docs.extend(doc_splits)
+
+        return {"documents": split_docs}
+
+    @component.output_types(documents=list[Document])
+    async def run_async(self, documents: list[Document]) -> dict[str, list[Document]]:
+        """
+        Asynchronously split documents based on embedding similarity.
+
+        This is the asynchronous version of the `run` method with the same parameters and return values.
+
+        :param documents: The documents to split.
+        :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 `split_id` to track the split number.
+                - A metadata field `page_number` to track the original page number.
+                - All other metadata copied from the original document.
+
+        :raises RuntimeError: If the component wasn't warmed up.
+        :raises TypeError: If the input is not a list of Documents.
+        :raises ValueError: If the document content is None or empty.
+        """
+        await self.warm_up_async()
+
+        if not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
+            raise TypeError("EmbeddingBasedDocumentSplitter expects a List of Documents as input.")
+
+        tasks: list[Awaitable[list[Document]]] = []
+        for doc in documents:
+            if doc.content is None:
+                raise ValueError(
+                    f"EmbeddingBasedDocumentSplitter only works with text documents but content for "
+                    f"document ID {doc.id} is None."
+                )
+            if doc.content == "":
+                logger.warning("Document ID {doc_id} has an empty content. Skipping this document.", doc_id=doc.id)
+                continue
+
+            tasks.append(self._split_document_async(doc=doc))
+
+        return {"documents": [*chain.from_iterable(await gather(*tasks))]}
+
+    def _split_document(self, doc: Document) -> list[Document]:
+        """
+        Split a single document based on embedding similarity.
+        """
+        # Create an initial split of the document content into smaller chunks
+        # doc.content is validated in `run`
+        splits = self._split_text(text=doc.content)  # type: ignore[arg-type]
+
+        # Merge splits smaller than min_length
+        merged_splits = self._merge_small_splits(splits=splits)
+
+        # Recursively split splits larger than max_length
+        final_splits = self._split_large_splits(splits=merged_splits)
+
+        # Create Document objects from the final splits
+        return EmbeddingBasedDocumentSplitter._create_documents_from_splits(splits=final_splits, original_doc=doc)
+
+    async def _split_document_async(self, doc: Document) -> list[Document]:
+        """
+        Split a single document based on embedding similarity.
+        """
+        # Create an initial split of the document content into smaller chunks
+        # doc.content is validated in `run`
+        splits = await self._split_text_async(text=doc.content)  # type: ignore[arg-type]
+
+        # Merge splits smaller than min_length
+        merged_splits = self._merge_small_splits(splits=splits)
+
+        # Recursively split splits larger than max_length
+        final_splits = self._split_large_splits(splits=merged_splits)
+
+        # Create Document objects from the final splits
+        return EmbeddingBasedDocumentSplitter._create_documents_from_splits(splits=final_splits, original_doc=doc)
+
+    def _prepare_sentence_groups(self, text: str) -> list[str]:
+        """Preprocess raw text into grouped sentences ready for embedding."""
+        # NOTE: `self.sentence_splitter.split_sentences` strips all white space types (e.g. new lines, page breaks,
+        # etc.) at the end of the provided text. So to not lose them, we need keep track of them and add them back to
+        # the last sentence.
+        rstripped_text = text.rstrip()
+        trailing_whitespaces = text[len(rstripped_text) :]
+
+        # Split the text into sentences
+        sentences_result = self.sentence_splitter.split_sentences(rstripped_text)  # type: ignore[union-attr]
+
+        # Add back the stripped white spaces to the last sentence
+        if sentences_result and trailing_whitespaces:
+            sentences_result[-1]["sentence"] += trailing_whitespaces
+            sentences_result[-1]["end"] += len(trailing_whitespaces)
+
+        sentences = [sentence["sentence"] for sentence in sentences_result]
+        return self._group_sentences(sentences=sentences)
+
+    def _split_text(self, text: str) -> list[str]:
+        """
+        Split a text into smaller chunks based on embedding similarity.
+        """
+        sentence_groups = self._prepare_sentence_groups(text=text)
+        embeddings = self._calculate_embeddings(sentence_groups=sentence_groups)
+        split_points = self._find_split_points(embeddings=embeddings)
+        return self._create_splits_from_points(sentence_groups=sentence_groups, split_points=split_points)
+
+    async def _split_text_async(self, text: str) -> list[str]:
+        """
+        Asynchronously split a text into smaller chunks based on embedding similarity.
+        """
+        sentence_groups = self._prepare_sentence_groups(text=text)
+        embeddings = await self._calculate_embeddings_async(sentence_groups=sentence_groups)
+        split_points = self._find_split_points(embeddings=embeddings)
+        return self._create_splits_from_points(sentence_groups=sentence_groups, split_points=split_points)
+
+    def _group_sentences(self, sentences: list[str]) -> list[str]:
+        """
+        Group sentences into groups of sentences_per_group.
+        """
+        if self.sentences_per_group == 1:
+            return sentences
+
+        groups = []
+        for i in range(0, len(sentences), self.sentences_per_group):
+            group = sentences[i : i + self.sentences_per_group]
+            groups.append("".join(group))
+
+        return groups
+
+    def _calculate_embeddings(self, sentence_groups: list[str]) -> list[list[float]]:
+        """
+        Calculate embeddings for each sentence group using the DocumentEmbedder.
+        """
+        # Create Document objects for each group
+        group_docs = [Document(content=group) for group in sentence_groups]
+        result = self.document_embedder.run(group_docs)
+        embedded_docs = result["documents"]
+        return [doc.embedding for doc in embedded_docs]
+
+    async def _calculate_embeddings_async(self, sentence_groups: list[str]) -> list[list[float]]:
+        """
+        Asynchronously Calculate embeddings for each sentence group using the DocumentEmbedder.
+        """
+        # Create Document objects for each group
+        group_docs = [Document(content=group) for group in sentence_groups]
+        result = await _execute_component_async(self.document_embedder, documents=group_docs)
+        embedded_docs = result["documents"]
+        return [doc.embedding for doc in embedded_docs]
+
+    def _find_split_points(self, embeddings: list[list[float]]) -> list[int]:
+        """
+        Find split points based on cosine distances between sequential embeddings.
+        """
+        if len(embeddings) <= 1:
+            return []
+
+        # Calculate cosine distances between sequential pairs
+        distances = []
+        for i in range(len(embeddings) - 1):
+            distance = EmbeddingBasedDocumentSplitter._cosine_distance(
+                embedding1=embeddings[i], embedding2=embeddings[i + 1]
+            )
+            distances.append(distance)
+
+        # Calculate threshold based on percentile
+        threshold = np.percentile(distances, self.percentile * 100)
+
+        # Find indices where distance exceeds threshold
+        split_points = []
+        for i, distance in enumerate(distances):
+            if distance > threshold:
+                split_points.append(i + 1)  # +1 because we want to split after this point
+
+        return split_points
+
+    @staticmethod
+    def _cosine_distance(embedding1: list[float], embedding2: list[float]) -> float:
+        """
+        Calculate cosine distance between two embeddings.
+        """
+        vec1 = np.array(embedding1)
+        vec2 = np.array(embedding2)
+
+        norm1 = float(np.linalg.norm(vec1))
+        norm2 = float(np.linalg.norm(vec2))
+
+        if norm1 == 0 or norm2 == 0:
+            return 1.0
+
+        cosine_sim = float(np.dot(vec1, vec2) / (norm1 * norm2))
+
+        return 1.0 - cosine_sim
+
+    @staticmethod
+    def _create_splits_from_points(sentence_groups: list[str], split_points: list[int]) -> list[str]:
+        """
+        Create splits based on split points.
+        """
+        if not split_points:
+            return ["".join(sentence_groups)]
+
+        splits = []
+        start = 0
+
+        for point in split_points:
+            split_text = "".join(sentence_groups[start:point])
+            if split_text:
+                splits.append(split_text)
+            start = point
+
+        # Add the last split
+        if start < len(sentence_groups):
+            split_text = "".join(sentence_groups[start:])
+            if split_text:
+                splits.append(split_text)
+
+        return splits
+
+    def _merge_small_splits(self, splits: list[str]) -> list[str]:
+        """
+        Merge splits that are below min_length.
+        """
+        if not splits:
+            return splits
+
+        merged = []
+        current_split = splits[0]
+
+        for split in splits[1:]:
+            # We merge splits that are smaller than min_length but only if the newly merged split is still below
+            # max_length.
+            if len(current_split) < self.min_length and len(current_split) + len(split) < self.max_length:
+                # Merge with next split
+                current_split += split
+            else:
+                # Current split is long enough, save it and start a new one
+                merged.append(current_split)
+                current_split = split
+
+        # Don't forget the last split
+        merged.append(current_split)
+
+        return merged
+
+    def _split_large_splits(self, splits: list[str]) -> list[str]:
+        """
+        Recursively split splits that are above max_length.
+
+        This method checks each split and if it exceeds max_length, it attempts to split it further using the same
+        embedding-based approach. This is done recursively until all splits are within the max_length limit or no
+        further splitting is possible.
+
+        This works because the threshold for splits is calculated dynamically based on the provided of embeddings.
+        """
+        final_splits = []
+
+        for split in splits:
+            if len(split) <= self.max_length:
+                final_splits.append(split)
+            else:
+                # Recursively split large splits
+                # We can reuse the same _split_text method to split the text into smaller chunks because the threshold
+                # for splits is calculated dynamically based on embeddings from `split`.
+                sub_splits = self._split_text(text=split)
+
+                # Stop splitting if no further split is possible or continue with recursion
+                if len(sub_splits) == 1:
+                    logger.warning(
+                        "Could not split a chunk further below max_length={max_length}. "
+                        "Returning chunk of length {length}.",
+                        max_length=self.max_length,
+                        length=len(split),
+                    )
+                    final_splits.append(split)
+                else:
+                    final_splits.extend(self._split_large_splits(splits=sub_splits))
+
+        return final_splits
+
+    @staticmethod
+    def _create_documents_from_splits(splits: list[str], original_doc: Document) -> list[Document]:
+        """
+        Create Document objects from splits.
+        """
+        documents = []
+        metadata = deepcopy(original_doc.meta)
+        metadata["source_id"] = original_doc.id
+
+        # Calculate page numbers for each split
+        current_page = 1
+
+        for i, split_text in enumerate(splits):
+            split_meta = deepcopy(metadata)
+            split_meta["split_id"] = i
+
+            # 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
+
+            doc = Document(content=split_text, meta=split_meta)
+            documents.append(doc)
+
+            # Update page counter for next split
+            current_page += page_breaks_in_split
+
+        return documents
+
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Serializes the component to a dictionary.
+
+        :returns:
+            Serialized dictionary representation of the component.
+        """
+        return default_to_dict(
+            self,
+            document_embedder=component_to_dict(obj=self.document_embedder, name="document_embedder"),
+            sentences_per_group=self.sentences_per_group,
+            percentile=self.percentile,
+            min_length=self.min_length,
+            max_length=self.max_length,
+            language=self.language,
+            use_split_rules=self.use_split_rules,
+            extend_abbreviations=self.extend_abbreviations,
+        )
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "EmbeddingBasedDocumentSplitter":
+        """
+        Deserializes the component from a dictionary.
+
+        :param data:
+            The dictionary to deserialize and create the component.
+
+        :returns:
+            The deserialized component.
+        """
+        deserialize_component_inplace(data["init_parameters"], key="document_embedder")
+        return default_from_dict(cls, data)