deepset-ai--haystack/haystack/components/preprocessors/embedding_based_document_splitter.py

# 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)