# SPDX-FileCopyrightText: 2022-present deepset GmbH # # SPDX-License-Identifier: Apache-2.0 import re from copy import deepcopy from typing import Any, Literal from haystack import Document, component, logging from haystack.lazy_imports import LazyImport with LazyImport("Run 'pip install tiktoken'") as tiktoken_imports: import tiktoken logger = logging.getLogger(__name__) @component class RecursiveDocumentSplitter: """ Recursively chunk text into smaller chunks. This component is used to split text into smaller chunks, it does so by recursively applying a list of separators to the text. The separators are applied in the order they are provided, typically this is a list of separators that are applied in a specific order, being the last separator the most specific one. Each separator is applied to the text, it then checks each of the resulting chunks, it keeps the chunks that are within the split_length, for the ones that are larger than the split_length, it applies the next separator in the list to the remaining text. This is done until all chunks are smaller than the split_length parameter. Example: ```python from haystack import Document from haystack.components.preprocessors import RecursiveDocumentSplitter chunker = RecursiveDocumentSplitter(split_length=260, split_overlap=0, separators=["\\n\\n", "\\n", ".", " "]) text = ('''Artificial intelligence (AI) - Introduction AI, in its broadest sense, is intelligence exhibited by machines, particularly computer systems. AI technology is widely used throughout industry, government, and science. Some high-profile applications include advanced web search engines; recommendation systems; interacting via human speech; autonomous vehicles; generative and creative tools; and superhuman play and analysis in strategy games.''') doc = Document(content=text) doc_chunks = chunker.run([doc]) print(doc_chunks["documents"]) # [ # Document(id=..., content: 'Artificial intelligence (AI) - Introduction\\n\\n', meta: {'original_id': '...', 'split_id': 0, 'split_idx_start': 0, '_split_overlap': []}) # Document(id=..., content: 'AI, in its broadest sense, is intelligence exhibited by machines, particularly computer systems.\\n', meta: {'original_id': '...', 'split_id': 1, 'split_idx_start': 45, '_split_overlap': []}) # Document(id=..., content: 'AI technology is widely used throughout industry, government, and science.', meta: {'original_id': '...', 'split_id': 2, 'split_idx_start': 142, '_split_overlap': []}) # Document(id=..., content: ' Some high-profile applications include advanced web search engines; recommendation systems; interac...', meta: {'original_id': '...', 'split_id': 3, 'split_idx_start': 216, '_split_overlap': []}) # ] ``` """ # noqa: E501 def __init__( self, *, split_length: int = 200, split_overlap: int = 0, split_unit: Literal["word", "char", "token"] = "word", separators: list[str] | None = None, sentence_splitter_params: dict[str, Any] | None = None, ) -> None: """ Initializes a RecursiveDocumentSplitter. :param split_length: The maximum length of each chunk by default in words, but can be in characters or tokens. See the `split_units` parameter. :param split_overlap: The number of characters to overlap between consecutive chunks. :param split_unit: The unit of the split_length parameter. It can be either "word", "char", or "token". If "token" is selected, the text will be split into tokens using the tiktoken tokenizer (o200k_base). :param separators: An optional list of separator strings to use for splitting the text. The string separators will be treated as regular expressions unless the separator is "sentence", in that case the text will be split into sentences using a custom sentence tokenizer based on NLTK. See: haystack.components.preprocessors.sentence_tokenizer.SentenceSplitter. If no separators are provided, the default separators ["\\n\\n", "sentence", "\\n", " "] are used. :param sentence_splitter_params: Optional parameters to pass to the sentence tokenizer. See: haystack.components.preprocessors.sentence_tokenizer.SentenceSplitter for more information. :raises ValueError: If the overlap is greater than or equal to the chunk size or if the overlap is negative, or if any separator is not a string. """ self.split_length = split_length self.split_overlap = split_overlap self.split_units = split_unit self.separators = separators if separators else ["\n\n", "sentence", "\n", " "] # default separators self._check_params() self.nltk_tokenizer = None self.sentence_splitter_params = ( {"keep_white_spaces": True} if sentence_splitter_params is None else sentence_splitter_params ) self.tiktoken_tokenizer: "tiktoken.Encoding" | None = None self._is_warmed_up = False def warm_up(self) -> None: """ Warm up the sentence tokenizer and tiktoken tokenizer if needed. """ if self._is_warmed_up: return if "sentence" in self.separators: self.nltk_tokenizer = self._get_custom_sentence_tokenizer(self.sentence_splitter_params) if self.split_units == "token": tiktoken_imports.check() self.tiktoken_tokenizer = tiktoken.get_encoding("o200k_base") self._is_warmed_up = True def _check_params(self) -> None: if self.split_length < 1: raise ValueError("Split length must be at least 1 character.") if self.split_overlap < 0: raise ValueError("Overlap must be greater than zero.") if self.split_overlap >= self.split_length: raise ValueError("Overlap cannot be greater than or equal to the chunk size.") if not all(isinstance(separator, str) for separator in self.separators): raise ValueError("All separators must be strings.") @staticmethod def _get_custom_sentence_tokenizer(sentence_splitter_params: dict[str, Any]) -> Any: from haystack.components.preprocessors.sentence_tokenizer import SentenceSplitter return SentenceSplitter(**sentence_splitter_params) def _split_chunk(self, current_chunk: str) -> tuple[str, str]: """ Splits a chunk based on the split_length and split_units attribute. :param current_chunk: The current chunk to be split. :returns: A tuple containing the current chunk and the remaining chunk. """ if self.split_units == "word": words = current_chunk.split() current_chunk = " ".join(words[: self.split_length]) remaining_words = words[self.split_length :] return current_chunk, " ".join(remaining_words) if self.split_units == "char": text = current_chunk current_chunk = text[: self.split_length] remaining_chars = text[self.split_length :] return current_chunk, remaining_chars # at this point we know that the tokenizer is already initialized tokens = self.tiktoken_tokenizer.encode(current_chunk) # type: ignore current_tokens = tokens[: self.split_length] remaining_tokens = tokens[self.split_length :] return self.tiktoken_tokenizer.decode(current_tokens), self.tiktoken_tokenizer.decode(remaining_tokens) # type: ignore def _apply_overlap(self, chunks: list[str]) -> list[str]: """ Applies an overlap between consecutive chunks if the chunk_overlap attribute is greater than zero. Works for both word- and character-level splitting. It trims the last chunk if it exceeds the split_length and adds the trimmed content to the next chunk. If the last chunk is still too long after trimming, it splits it and adds the first chunk to the list. This process continues until the last chunk is within the split_length. :param chunks: A list of text chunks. :returns: A list of text chunks with the overlap applied. """ overlapped_chunks: list[str] = [] for idx, chunk in enumerate(chunks): if idx == 0: overlapped_chunks.append(chunk) continue # get the overlap between the current and previous chunk overlap, prev_chunk = self._get_overlap(overlapped_chunks) if overlap == prev_chunk: logger.warning( "Overlap is the same as the previous chunk. " "Consider increasing the `split_length` parameter or decreasing the `split_overlap` parameter." ) current_chunk = self._create_chunk_starting_with_overlap(chunk, overlap) # if this new chunk exceeds 'split_length', trim it and move the remaining text to the next chunk # if this is the last chunk, another new chunk will contain the trimmed text preceded by the overlap # of the last chunk if self._chunk_length(current_chunk) > self.split_length: current_chunk, remaining_text = self._split_chunk(current_chunk) if idx < len(chunks) - 1: if self.split_units == "word": chunks[idx + 1] = remaining_text + " " + chunks[idx + 1] elif self.split_units == "token": # For token-based splitting, combine at token level # at this point we know that the tokenizer is already initialized remaining_tokens = self.tiktoken_tokenizer.encode(remaining_text) # type: ignore next_chunk_tokens = self.tiktoken_tokenizer.encode(chunks[idx + 1]) # type: ignore chunks[idx + 1] = self.tiktoken_tokenizer.decode(remaining_tokens + next_chunk_tokens) # type: ignore else: # char chunks[idx + 1] = remaining_text + chunks[idx + 1] elif remaining_text: # create a new chunk with the trimmed text preceded by the overlap of the last chunk overlapped_chunks.append(current_chunk) chunk = remaining_text overlap, _ = self._get_overlap(overlapped_chunks) current_chunk = self._create_chunk_starting_with_overlap(chunk, overlap) overlapped_chunks.append(current_chunk) # it can still be that the new last chunk exceeds the 'split_length' # continue splitting until the last chunk is within 'split_length' if idx == len(chunks) - 1 and self._chunk_length(current_chunk) > self.split_length: last_chunk = overlapped_chunks.pop() first_chunk, remaining_chunk = self._split_chunk(last_chunk) overlapped_chunks.append(first_chunk) while remaining_chunk: # combine overlap with remaining chunk overlap, _ = self._get_overlap(overlapped_chunks) current = self._create_chunk_starting_with_overlap(remaining_chunk, overlap) # if it fits within split_length we are done if self._chunk_length(current) <= self.split_length: overlapped_chunks.append(current) break # otherwise split it again first_chunk, remaining_chunk = self._split_chunk(current) overlapped_chunks.append(first_chunk) return overlapped_chunks def _create_chunk_starting_with_overlap(self, chunk: str, overlap: str) -> str: if self.split_units == "word": current_chunk = overlap + " " + chunk elif self.split_units == "token": # For token-based splitting, combine at token level # at this point we know that the tokenizer is already initialized overlap_tokens = self.tiktoken_tokenizer.encode(overlap) # type: ignore chunk_tokens = self.tiktoken_tokenizer.encode(chunk) # type: ignore current_chunk = self.tiktoken_tokenizer.decode(overlap_tokens + chunk_tokens) # type: ignore else: # char current_chunk = overlap + chunk return current_chunk def _get_overlap(self, overlapped_chunks: list[str]) -> tuple[str, str]: """Get the previous overlapped chunk instead of the original chunk.""" prev_chunk = overlapped_chunks[-1] overlap_start = max(0, self._chunk_length(prev_chunk) - self.split_overlap) if self.split_units == "word": word_chunks = prev_chunk.split() overlap = " ".join(word_chunks[overlap_start:]) elif self.split_units == "token": # For token-based splitting, handle overlap at token level # at this point we know that the tokenizer is already initialized tokens = self.tiktoken_tokenizer.encode(prev_chunk) # type: ignore overlap_tokens = tokens[overlap_start:] overlap = self.tiktoken_tokenizer.decode(overlap_tokens) # type: ignore else: # char overlap = prev_chunk[overlap_start:] return overlap, prev_chunk def _chunk_length(self, text: str) -> int: """ Get the length of the chunk in the specified units (words, characters, or tokens). :param text: The text to measure. :returns: The length of the text in the specified units. """ if self.split_units == "word": words = [word for word in text.split(" ") if word] return len(words) if self.split_units == "char": return len(text) # token # at this point we know that the tokenizer is already initialized return len(self.tiktoken_tokenizer.encode(text)) # type: ignore def _chunk_text(self, text: str) -> list[str]: """ Recursive chunking algorithm that divides text into smaller chunks based on a list of separator characters. It starts with a list of separator characters (e.g., ["\n\n", "sentence", "\n", " "]) and attempts to divide the text using the first separator. If the resulting chunks are still larger than the specified chunk size, it moves to the next separator in the list. This process continues recursively, progressively applying each specific separator until the chunks meet the desired size criteria. :param text: The text to be split into chunks. :returns: A list of text chunks. """ if self._chunk_length(text) <= self.split_length: return [text] for curr_separator in self.separators: if curr_separator == "sentence": # re. ignore: correct SentenceSplitter initialization is checked at the initialization of the component sentence_with_spans = self.nltk_tokenizer.split_sentences(text) # type: ignore splits = [sentence["sentence"] for sentence in sentence_with_spans] else: # add escape "\" to the separator and wrapped it in a group so that it's included in the splits as well escaped_separator = re.escape(curr_separator) escaped_separator = f"({escaped_separator})" # split the text and merge every two consecutive splits, i.e.: the text and the separator after it splits = re.split(escaped_separator, text) splits = [ "".join([splits[i], splits[i + 1]]) if i < len(splits) - 1 else splits[i] for i in range(0, len(splits), 2) ] # remove last split if it's empty splits = splits[:-1] if splits[-1] == "" else splits if len(splits) == 1: # go to next separator, if current separator not found in the text continue chunks = [] current_chunk: list[str] = [] current_length = 0 # check splits, if any is too long, recursively chunk it, otherwise add to current chunk for split in splits: split_text = split # if adding this split exceeds chunk_size, process current_chunk if current_length + self._chunk_length(split_text) > self.split_length: # process current_chunk if current_chunk: # keep the good splits chunks.append("".join(current_chunk)) current_chunk = [] current_length = 0 # recursively handle splits that are too large if self._chunk_length(split_text) > self.split_length: if curr_separator == self.separators[-1]: # tried last separator, can't split further, do a fixed-split based on word/character/token fall_back_chunks = self._fall_back_to_fixed_chunking(split_text, self.split_units) chunks.extend(fall_back_chunks) else: chunks.extend(self._chunk_text(split_text)) else: current_chunk.append(split_text) current_length += self._chunk_length(split_text) else: current_chunk.append(split_text) current_length += self._chunk_length(split_text) if current_chunk: chunks.append("".join(current_chunk)) if self.split_overlap > 0: chunks = self._apply_overlap(chunks) if chunks: return chunks # if no separator worked, fall back to word- or character-level chunking chunks = self._fall_back_to_fixed_chunking(text, self.split_units) if self.split_overlap > 0: chunks = self._apply_overlap(chunks) return chunks def _fall_back_to_fixed_chunking(self, text: str, split_units: Literal["word", "char", "token"]) -> list[str]: """ Fall back to a fixed chunking approach if no separator works for the text. Splits the text into smaller chunks based on the split_length and split_units attributes, either by words, characters, or tokens. :param text: The text to be split into chunks. :param split_units: The unit of the split_length parameter. It can be either "word", "char", or "token". :returns: A list of text chunks. """ chunks = [] if split_units == "word": words = re.findall(r"\S+|\s+", text) current_chunk = [] current_length = 0 for word in words: if word != " ": current_chunk.append(word) current_length += 1 if current_length == self.split_length and current_chunk: chunks.append("".join(current_chunk)) current_chunk = [] current_length = 0 else: current_chunk.append(word) if current_chunk: chunks.append("".join(current_chunk)) elif split_units == "char": for i in range(0, self._chunk_length(text), self.split_length): chunks.append(text[i : i + self.split_length]) else: # token # at this point we know that the tokenizer is already initialized tokens = self.tiktoken_tokenizer.encode(text) # type: ignore for i in range(0, len(tokens), self.split_length): chunk_tokens = tokens[i : i + self.split_length] chunks.append(self.tiktoken_tokenizer.decode(chunk_tokens)) # type: ignore return chunks def _add_overlap_info(self, curr_pos: int, new_doc: Document, new_docs: list[Document]) -> None: prev_doc = new_docs[-1] # curr_pos and split_idx_start are character offsets, so measure the # overlap and range in characters too (not via _chunk_length, which returns a word/token count). prev_doc_length = len(prev_doc.content) # type: ignore overlap_length = prev_doc_length - (curr_pos - prev_doc.meta["split_idx_start"]) if overlap_length > 0: prev_doc.meta["_split_overlap"].append({"doc_id": new_doc.id, "range": (0, overlap_length)}) new_doc.meta["_split_overlap"].append( {"doc_id": prev_doc.id, "range": (prev_doc_length - overlap_length, prev_doc_length)} ) def _run_one(self, doc: Document) -> list[Document]: chunks = self._chunk_text(doc.content) # type: ignore # the caller already check for a non-empty doc.content chunks = chunks[:-1] if len(chunks[-1]) == 0 else chunks # remove last empty chunk if it exists current_position = 0 current_page = 1 new_docs: list[Document] = [] for split_nr, chunk in enumerate(chunks): meta = deepcopy(doc.meta) meta["parent_id"] = doc.id meta["split_id"] = split_nr meta["split_idx_start"] = current_position meta["_split_overlap"] = [] if self.split_overlap > 0 else None new_doc = Document(content=chunk, meta=meta) # add overlap information to the previous and current doc if split_nr > 0 and self.split_overlap > 0: self._add_overlap_info(current_position, new_doc, new_docs) # count page breaks in the chunk current_page += chunk.count("\f") # if there are consecutive page breaks at the end with no more text, adjust the page number # e.g: "text\f\f\f" -> 3 page breaks, but current_page should be 1 consecutive_page_breaks = len(chunk) - len(chunk.rstrip("\f")) if consecutive_page_breaks > 0: new_doc.meta["page_number"] = current_page - consecutive_page_breaks else: new_doc.meta["page_number"] = current_page # keep the new chunk doc and update the current position new_docs.append(new_doc) # Advance current_position by chunk length minus overlap. # split_overlap is in split_units, not chars, so get the actual # overlap string from _get_overlap() and use its char length. if self.split_overlap > 0 and split_nr < len(chunks) - 1: overlap_str, _ = self._get_overlap([doc.content for doc in new_docs]) # type: ignore[misc] overlap_char_len = len(overlap_str) else: overlap_char_len = 0 current_position += len(chunk) - overlap_char_len return new_docs @component.output_types(documents=list[Document]) def run(self, documents: list[Document]) -> dict[str, list[Document]]: """ Split a list of documents into documents with smaller chunks of text. :param documents: List of Documents to split. :returns: A dictionary containing a key "documents" with a List of Documents with smaller chunks of text corresponding to the input documents. """ if not self._is_warmed_up and ("sentence" in self.separators or self.split_units == "token"): self.warm_up() docs = [] for doc in documents: if not doc.content or doc.content == "": logger.warning("Document ID {doc_id} has an empty content. Skipping this document.", doc_id=doc.id) continue docs.extend(self._run_one(doc)) return {"documents": docs}