chore: import upstream snapshot with attribution
CodeQL / Analyze (python) (push) Has been cancelled
Update Platform Components Table / update (push) Has been cancelled
Docker image release / Build base image (push) Has been cancelled
Sync docs with Docusaurus / sync (push) Has been cancelled
Tests / Check if changed (push) Has been cancelled
Tests / format (push) Has been cancelled
Tests / check-imports (push) Has been cancelled
Tests / Unit / macos-latest (push) Has been cancelled
Tests / Unit / ubuntu-latest (push) Has been cancelled
Tests / Unit / windows-latest (push) Has been cancelled
Tests / mypy (push) Has been cancelled
Tests / Integration / ubuntu-latest (push) Has been cancelled
Tests / Integration / macos-latest (push) Has been cancelled
Tests / Integration / windows-latest (push) Has been cancelled
Tests / notify-slack-on-failure (push) Has been cancelled
Tests / Mark tests as completed (push) Has been cancelled

This commit is contained in:
wehub-resource-sync
2026-07-13 13:22:28 +08:00
commit c56bef871b
9296 changed files with 1854228 additions and 0 deletions
@@ -0,0 +1,38 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
import sys
from typing import TYPE_CHECKING
from lazy_imports import LazyImporter
_import_structure = {
"csv_document_cleaner": ["CSVDocumentCleaner"],
"csv_document_splitter": ["CSVDocumentSplitter"],
"document_cleaner": ["DocumentCleaner"],
"document_preprocessor": ["DocumentPreprocessor"],
"document_splitter": ["DocumentSplitter"],
"embedding_based_document_splitter": ["EmbeddingBasedDocumentSplitter"],
"hierarchical_document_splitter": ["HierarchicalDocumentSplitter"],
"markdown_header_splitter": ["MarkdownHeaderSplitter"],
"python_code_splitter": ["PythonCodeSplitter"],
"recursive_splitter": ["RecursiveDocumentSplitter"],
"text_cleaner": ["TextCleaner"],
}
if TYPE_CHECKING:
from .csv_document_cleaner import CSVDocumentCleaner as CSVDocumentCleaner
from .csv_document_splitter import CSVDocumentSplitter as CSVDocumentSplitter
from .document_cleaner import DocumentCleaner as DocumentCleaner
from .document_preprocessor import DocumentPreprocessor as DocumentPreprocessor
from .document_splitter import DocumentSplitter as DocumentSplitter
from .embedding_based_document_splitter import EmbeddingBasedDocumentSplitter as EmbeddingBasedDocumentSplitter
from .hierarchical_document_splitter import HierarchicalDocumentSplitter as HierarchicalDocumentSplitter
from .markdown_header_splitter import MarkdownHeaderSplitter as MarkdownHeaderSplitter
from .python_code_splitter import PythonCodeSplitter as PythonCodeSplitter
from .recursive_splitter import RecursiveDocumentSplitter as RecursiveDocumentSplitter
from .text_cleaner import TextCleaner as TextCleaner
else:
sys.modules[__name__] = LazyImporter(name=__name__, module_file=__file__, import_structure=_import_structure)
@@ -0,0 +1,178 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
from copy import deepcopy
from io import StringIO
from typing import Optional
from haystack import Document, component, logging
from haystack.lazy_imports import LazyImport
with LazyImport("Run 'pip install pandas'") as pandas_import:
import pandas as pd
logger = logging.getLogger(__name__)
@component
class CSVDocumentCleaner:
"""
A component for cleaning CSV documents by removing empty rows and columns.
This component processes CSV content stored in Documents, allowing
for the optional ignoring of a specified number of rows and columns before performing
the cleaning operation. Additionally, it provides options to keep document IDs and
control whether empty rows and columns should be removed.
"""
def __init__(
self,
*,
ignore_rows: int = 0,
ignore_columns: int = 0,
remove_empty_rows: bool = True,
remove_empty_columns: bool = True,
keep_id: bool = False,
) -> None:
"""
Initializes the CSVDocumentCleaner component.
:param ignore_rows: Number of rows to ignore from the top of the CSV table before processing.
:param ignore_columns: Number of columns to ignore from the left of the CSV table before processing.
:param remove_empty_rows: Whether to remove rows that are entirely empty.
:param remove_empty_columns: Whether to remove columns that are entirely empty.
:param keep_id: Whether to retain the original document ID in the output document.
Rows and columns ignored using these parameters are preserved in the final output, meaning
they are not considered when removing empty rows and columns.
"""
self.ignore_rows = ignore_rows
self.ignore_columns = ignore_columns
self.remove_empty_rows = remove_empty_rows
self.remove_empty_columns = remove_empty_columns
self.keep_id = keep_id
pandas_import.check()
@component.output_types(documents=list[Document])
def run(self, documents: list[Document]) -> dict[str, list[Document]]:
"""
Cleans CSV documents by removing empty rows and columns while preserving specified ignored rows and columns.
:param documents: List of Documents containing CSV-formatted content.
:return: A dictionary with a list of cleaned Documents under the key "documents".
Processing steps:
1. Reads each document's content as a CSV table.
2. Retains the specified number of `ignore_rows` from the top and `ignore_columns` from the left.
3. Drops any rows and columns that are entirely empty (if enabled by `remove_empty_rows` and
`remove_empty_columns`).
4. Reattaches the ignored rows and columns to maintain their original positions.
5. Returns the cleaned CSV content as a new `Document` object, with an option to retain the original
document ID.
"""
if len(documents) == 0:
return {"documents": []}
ignore_rows = self.ignore_rows
ignore_columns = self.ignore_columns
cleaned_documents = []
for document in documents:
try:
df = pd.read_csv(StringIO(document.content), header=None, dtype=object)
except Exception as e:
logger.exception(
"Error processing document {id}. Keeping it, but skipping cleaning. Error: {error}",
id=document.id,
error=e,
)
cleaned_documents.append(document)
continue
if ignore_rows > df.shape[0] or ignore_columns > df.shape[1]:
logger.warning(
"Document {id} has fewer rows {df_rows} or columns {df_cols} "
"than the number of rows {rows} or columns {cols} to ignore. "
"Keeping the entire document.",
id=document.id,
df_rows=df.shape[0],
df_cols=df.shape[1],
rows=ignore_rows,
cols=ignore_columns,
)
cleaned_documents.append(document)
continue
final_df = self._clean_df(df=df, ignore_rows=ignore_rows, ignore_columns=ignore_columns)
clean_doc = Document(
id=document.id if self.keep_id else "",
content=final_df.to_csv(index=False, header=False, lineterminator="\n"),
blob=document.blob,
meta=deepcopy(document.meta),
score=document.score,
embedding=document.embedding,
sparse_embedding=document.sparse_embedding,
)
cleaned_documents.append(clean_doc)
return {"documents": cleaned_documents}
def _clean_df(self, df: "pd.DataFrame", ignore_rows: int, ignore_columns: int) -> "pd.DataFrame":
"""
Cleans a DataFrame by removing empty rows and columns while preserving ignored sections.
:param df: The input DataFrame representing the CSV data.
:param ignore_rows: Number of top rows to ignore.
:param ignore_columns: Number of left columns to ignore.
"""
# Get ignored rows and columns
ignored_rows = self._get_ignored_rows(df=df, ignore_rows=ignore_rows)
ignored_columns = self._get_ignored_columns(df=df, ignore_columns=ignore_columns)
final_df = df.iloc[ignore_rows:, ignore_columns:]
# Drop rows that are entirely empty
if self.remove_empty_rows:
final_df = final_df.dropna(axis=0, how="all")
# Drop columns that are entirely empty
if self.remove_empty_columns:
final_df = final_df.dropna(axis=1, how="all")
# Reattach ignored rows
if ignore_rows > 0 and ignored_rows is not None:
# Keep only relevant columns
ignored_rows = ignored_rows.loc[:, final_df.columns]
final_df = pd.concat([ignored_rows, final_df], axis=0)
# Reattach ignored columns
if ignore_columns > 0 and ignored_columns is not None:
# Keep only relevant rows
ignored_columns = ignored_columns.loc[final_df.index, :]
final_df = pd.concat([ignored_columns, final_df], axis=1)
return final_df
@staticmethod
def _get_ignored_rows(df: "pd.DataFrame", ignore_rows: int) -> Optional["pd.DataFrame"]:
"""
Extracts the rows to be ignored from the DataFrame.
:param df: The input DataFrame.
:param ignore_rows: Number of rows to extract from the top.
"""
if ignore_rows > 0:
return df.iloc[:ignore_rows, :]
return None
@staticmethod
def _get_ignored_columns(df: "pd.DataFrame", ignore_columns: int) -> Optional["pd.DataFrame"]:
"""
Extracts the columns to be ignored from the DataFrame.
:param df: The input DataFrame.
:param ignore_columns: Number of columns to extract from the left.
"""
if ignore_columns > 0:
return df.iloc[:, :ignore_columns]
return None
@@ -0,0 +1,286 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
from io import StringIO
from typing import Any, Literal, get_args
from haystack import Document, component, logging
from haystack.lazy_imports import LazyImport
with LazyImport("Run 'pip install pandas'") as pandas_import:
import pandas as pd
logger = logging.getLogger(__name__)
SplitMode = Literal["threshold", "row-wise"]
@component
class CSVDocumentSplitter:
"""
A component for splitting CSV documents into sub-tables based on split arguments.
The splitter supports two modes of operation:
- identify consecutive empty rows or columns that exceed a given threshold
and uses them as delimiters to segment the document into smaller tables.
- split each row into a separate sub-table, represented as a Document.
"""
def __init__(
self,
row_split_threshold: int | None = 2,
column_split_threshold: int | None = 2,
read_csv_kwargs: dict[str, Any] | None = None,
split_mode: SplitMode = "threshold",
) -> None:
"""
Initializes the CSVDocumentSplitter component.
:param row_split_threshold: The minimum number of consecutive empty rows required to trigger a split.
:param column_split_threshold: The minimum number of consecutive empty columns required to trigger a split.
:param read_csv_kwargs: Additional keyword arguments to pass to `pandas.read_csv`.
By default, the component with options:
- `header=None`
- `skip_blank_lines=False` to preserve blank lines
- `dtype=object` to prevent type inference (e.g., converting numbers to floats).
See https://pandas.pydata.org/docs/reference/api/pandas.read_csv.html for more information.
:param split_mode:
If `threshold`, the component will split the document based on the number of
consecutive empty rows or columns that exceed the `row_split_threshold` or `column_split_threshold`.
If `row-wise`, the component will split each row into a separate sub-table.
"""
pandas_import.check()
if split_mode not in get_args(SplitMode):
raise ValueError(
f"Split mode '{split_mode}' not recognized. Choose one among: {', '.join(get_args(SplitMode))}."
)
if row_split_threshold is not None and row_split_threshold < 1:
raise ValueError("row_split_threshold must be greater than 0")
if column_split_threshold is not None and column_split_threshold < 1:
raise ValueError("column_split_threshold must be greater than 0")
if row_split_threshold is None and column_split_threshold is None:
raise ValueError("At least one of row_split_threshold or column_split_threshold must be specified.")
self.row_split_threshold = row_split_threshold
self.column_split_threshold = column_split_threshold
self.read_csv_kwargs = read_csv_kwargs or {}
self.split_mode = split_mode
@component.output_types(documents=list[Document])
def run(self, documents: list[Document]) -> dict[str, list[Document]]:
"""
Processes and splits a list of CSV documents into multiple sub-tables.
**Splitting Process:**
1. Applies a row-based split if `row_split_threshold` is provided.
2. Applies a column-based split if `column_split_threshold` is provided.
3. If both thresholds are specified, performs a recursive split by rows first, then columns, ensuring
further fragmentation of any sub-tables that still contain empty sections.
4. Sorts the resulting sub-tables based on their original positions within the document.
:param documents: A list of Documents containing CSV-formatted content.
Each document is assumed to contain one or more tables separated by empty rows or columns.
:return:
A dictionary with a key `"documents"`, mapping to a list of new `Document` objects,
each representing an extracted sub-table from the original CSV.
The metadata of each document includes:
- A field `source_id` to track the original document.
- A field `row_idx_start` to indicate the starting row index of the sub-table in the original table.
- A field `col_idx_start` to indicate the starting column index of the sub-table in the original table.
- A field `split_id` to indicate the order of the split in the original document.
- All other metadata copied from the original document.
- If a document cannot be processed, it is returned unchanged.
- The `meta` field from the original document is preserved in the split documents.
"""
if len(documents) == 0:
return {"documents": documents}
resolved_read_csv_kwargs = {"header": None, "skip_blank_lines": False, "dtype": object, **self.read_csv_kwargs}
split_documents = []
split_dfs = []
for document in documents:
try:
df = pd.read_csv(StringIO(document.content), **resolved_read_csv_kwargs)
except Exception as e:
logger.exception(
"Error processing document {document_id}. Keeping it, but skipping splitting. Error: {error}",
document_id=document.id,
error=e,
)
split_documents.append(document)
continue
if self.split_mode == "row-wise":
# each row is a separate sub-table
split_dfs = self._split_by_row(df=df)
elif self.split_mode == "threshold":
if self.row_split_threshold is not None and self.column_split_threshold is None:
# split by rows
split_dfs = self._split_dataframe(df=df, split_threshold=self.row_split_threshold, axis="row")
elif self.column_split_threshold is not None and self.row_split_threshold is None:
# split by columns
split_dfs = self._split_dataframe(df=df, split_threshold=self.column_split_threshold, axis="column")
else:
# recursive split
split_dfs = self._recursive_split(
df=df,
row_split_threshold=self.row_split_threshold, # type: ignore
column_split_threshold=self.column_split_threshold, # type: ignore
)
# check if no sub-tables were found
if len(split_dfs) == 0:
logger.warning(
"No sub-tables found while splitting CSV Document with id {doc_id}. Skipping document.",
doc_id=document.id,
)
continue
# Sort split_dfs first by row index, then by column index
split_dfs.sort(key=lambda dataframe: (dataframe.index[0], dataframe.columns[0]))
for split_id, split_df in enumerate(split_dfs):
split_documents.append(
Document(
content=split_df.to_csv(index=False, header=False, lineterminator="\n"),
meta={
**document.meta.copy(),
"source_id": document.id,
"row_idx_start": int(split_df.index[0]),
"col_idx_start": int(split_df.columns[0]),
"split_id": split_id,
},
)
)
return {"documents": split_documents}
@staticmethod
def _find_split_indices(
df: "pd.DataFrame", split_threshold: int, axis: Literal["row", "column"]
) -> list[tuple[int, int]]:
"""
Finds the indices of consecutive empty rows or columns in a DataFrame.
:param df: DataFrame to split.
:param split_threshold: Minimum number of consecutive empty rows or columns to trigger a split.
:param axis: Axis along which to find empty elements. Either "row" or "column".
:return: List of indices where consecutive empty rows or columns start.
"""
if axis == "row":
empty_elements = df[df.isnull().all(axis=1)].index.tolist()
else:
empty_elements = df.columns[df.isnull().all(axis=0)].tolist()
# If no empty elements found, return empty list
if len(empty_elements) == 0:
return []
# Identify groups of consecutive empty elements
split_indices = []
consecutive_count = 1
start_index = empty_elements[0]
for i in range(1, len(empty_elements)):
if empty_elements[i] == empty_elements[i - 1] + 1:
consecutive_count += 1
else:
if consecutive_count >= split_threshold:
split_indices.append((start_index, empty_elements[i - 1]))
consecutive_count = 1
start_index = empty_elements[i]
# Handle the last group of consecutive elements
if consecutive_count >= split_threshold:
split_indices.append((start_index, empty_elements[-1]))
return split_indices
def _split_dataframe(
self, df: "pd.DataFrame", split_threshold: int, axis: Literal["row", "column"]
) -> list["pd.DataFrame"]:
"""
Splits a DataFrame into sub-tables based on consecutive empty rows or columns exceeding `split_threshold`.
:param df: DataFrame to split.
:param split_threshold: Minimum number of consecutive empty rows or columns to trigger a split.
:param axis: Axis along which to split. Either "row" or "column".
:return: List of split DataFrames.
"""
# Find indices of consecutive empty rows or columns
split_indices = self._find_split_indices(df=df, split_threshold=split_threshold, axis=axis)
# If no split_indices are found, return the original DataFrame
if len(split_indices) == 0:
return [df]
# Split the DataFrame at identified indices
sub_tables = []
table_start_idx = 0
df_length = df.shape[0] if axis == "row" else df.shape[1]
for empty_start_idx, empty_end_idx in split_indices + [(df_length, df_length)]:
# Avoid empty splits
if empty_start_idx - table_start_idx >= 1:
if axis == "row":
sub_table = df.iloc[table_start_idx:empty_start_idx]
else:
sub_table = df.iloc[:, table_start_idx:empty_start_idx]
if not sub_table.empty:
sub_tables.append(sub_table)
table_start_idx = empty_end_idx + 1
return sub_tables
def _recursive_split(
self, df: "pd.DataFrame", row_split_threshold: int, column_split_threshold: int
) -> list["pd.DataFrame"]:
"""
Recursively splits a DataFrame.
Recursively splits a DataFrame first by empty rows, then by empty columns, and repeats the process
until no more splits are possible. Returns a list of DataFrames, each representing a fully separated sub-table.
:param df: A Pandas DataFrame representing a table (or multiple tables) extracted from a CSV.
:param row_split_threshold: The minimum number of consecutive empty rows required to trigger a split.
:param column_split_threshold: The minimum number of consecutive empty columns to trigger a split.
"""
# Step 1: Split by rows
new_sub_tables = self._split_dataframe(df=df, split_threshold=row_split_threshold, axis="row")
# Step 2: Split by columns
final_tables = []
for table in new_sub_tables:
final_tables.extend(self._split_dataframe(df=table, split_threshold=column_split_threshold, axis="column"))
# Step 3: Recursively reapply splitting checked by whether any new empty rows appear after column split
result = []
for table in final_tables:
# Check if there are consecutive rows >= row_split_threshold now present
if len(self._find_split_indices(df=table, split_threshold=row_split_threshold, axis="row")) > 0:
result.extend(
self._recursive_split(
df=table, row_split_threshold=row_split_threshold, column_split_threshold=column_split_threshold
)
)
else:
result.append(table)
return result
def _split_by_row(self, df: "pd.DataFrame") -> list["pd.DataFrame"]:
"""Split each CSV row into a separate subtable"""
split_dfs = []
for idx, row in enumerate(df.itertuples(index=False)):
split_df = pd.DataFrame(row).T
split_df.index = [idx] # Set the index of the new DataFrame to idx
split_dfs.append(split_df)
return split_dfs
@@ -0,0 +1,352 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
import re
from collections.abc import Generator
from copy import deepcopy
from functools import partial, reduce
from itertools import chain
from typing import Literal
from unicodedata import normalize
from haystack import Document, component, logging
logger = logging.getLogger(__name__)
@component
class DocumentCleaner:
"""
Cleans the text in the documents.
It removes extra whitespaces,
empty lines, specified substrings, regexes,
page headers and footers (in this order).
### Usage example:
```python
from haystack import Document
from haystack.components.preprocessors import DocumentCleaner
doc = Document(content="This is a document to clean\\n\\n\\nsubstring to remove")
cleaner = DocumentCleaner(remove_substrings = ["substring to remove"])
result = cleaner.run(documents=[doc])
assert result["documents"][0].content == "This is a document to clean "
```
"""
def __init__(
self,
remove_empty_lines: bool = True,
remove_extra_whitespaces: bool = True,
remove_repeated_substrings: bool = False,
keep_id: bool = False,
remove_substrings: list[str] | None = None,
remove_regex: str | None = None,
unicode_normalization: Literal["NFC", "NFKC", "NFD", "NFKD"] | None = None,
ascii_only: bool = False,
strip_whitespaces: bool = False,
replace_regexes: dict[str, str] | None = None,
) -> None:
"""
Initialize DocumentCleaner.
:param remove_empty_lines: If `True`, removes empty lines.
:param remove_extra_whitespaces: If `True`, removes extra whitespaces.
:param remove_repeated_substrings: If `True`, removes repeated substrings (headers and footers) from pages.
Pages must be separated by a form feed character "\\f",
which is supported by `TextFileToDocument` and `AzureOCRDocumentConverter`.
:param remove_substrings: List of substrings to remove from the text.
:param remove_regex: Regex to match and replace substrings by "".
:param keep_id: If `True`, keeps the IDs of the original documents.
:param unicode_normalization: Unicode normalization form to apply to the text.
Note: This will run before any other steps.
:param ascii_only: Whether to convert the text to ASCII only.
Will remove accents from characters and replace them with ASCII characters.
Other non-ASCII characters will be removed.
Note: This will run before any pattern matching or removal.
:param strip_whitespaces: If `True`, removes leading and trailing whitespace from the document content
using Python's `str.strip()`. Unlike `remove_extra_whitespaces`, this only affects the beginning
and end of the text, preserving internal whitespace (useful for markdown formatting).
:param replace_regexes: A dictionary mapping regex patterns to their replacement strings.
For example, `{r'\\n\\n+': '\\n'}` replaces multiple consecutive newlines with a single newline.
This is applied after `remove_regex` and allows custom replacements instead of just removal.
"""
self._validate_params(unicode_normalization=unicode_normalization)
self.remove_empty_lines = remove_empty_lines
self.remove_extra_whitespaces = remove_extra_whitespaces
self.remove_repeated_substrings = remove_repeated_substrings
self.remove_substrings = remove_substrings
self.remove_regex = remove_regex
self.keep_id = keep_id
self.unicode_normalization = unicode_normalization
self.ascii_only = ascii_only
self.strip_whitespaces = strip_whitespaces
self.replace_regexes = replace_regexes
def _validate_params(self, unicode_normalization: str | None) -> None:
"""
Validate the parameters of the DocumentCleaner.
:param unicode_normalization: Unicode normalization form to apply to the text.
:raises ValueError: if the parameters are not valid.
"""
if unicode_normalization and unicode_normalization not in ["NFC", "NFKC", "NFD", "NFKD"]:
raise ValueError("unicode_normalization must be one of 'NFC', 'NFKC', 'NFD', 'NFKD'.")
@component.output_types(documents=list[Document])
def run(self, documents: list[Document]) -> dict[str, list[Document]]:
"""
Cleans up the documents.
:param documents: List of Documents to clean.
:returns: A dictionary with the following key:
- `documents`: List of cleaned Documents.
:raises TypeError: if documents is not a list of Documents.
"""
if not isinstance(documents, list) or documents and not isinstance(documents[0], Document):
raise TypeError("DocumentCleaner expects a List of Documents as input.")
cleaned_docs = []
for doc in documents:
if doc.content is None:
logger.warning(
"DocumentCleaner only cleans text documents but document.content for document ID"
" {document_id} is None.",
document_id=doc.id,
)
cleaned_docs.append(doc)
continue
text = doc.content
if self.unicode_normalization:
text = self._normalize_unicode(text, self.unicode_normalization)
if self.ascii_only:
text = self._ascii_only(text)
if self.remove_extra_whitespaces:
text = self._remove_extra_whitespaces(text)
if self.remove_empty_lines:
text = self._remove_empty_lines(text)
if self.remove_substrings:
text = self._remove_substrings(text, self.remove_substrings)
if self.remove_regex:
text = self._remove_regex(text, self.remove_regex)
if self.replace_regexes:
text = self._replace_regexes(text, self.replace_regexes)
if self.remove_repeated_substrings:
text = self._remove_repeated_substrings(text)
if self.strip_whitespaces:
text = text.strip()
clean_doc = Document(
id=doc.id if self.keep_id else "",
content=text,
blob=doc.blob,
meta=deepcopy(doc.meta),
score=doc.score,
embedding=doc.embedding,
sparse_embedding=doc.sparse_embedding,
)
cleaned_docs.append(clean_doc)
return {"documents": cleaned_docs}
def _normalize_unicode(self, text: str, form: Literal["NFC", "NFKC", "NFD", "NFKD"]) -> str:
"""
Normalize the unicode of the text.
:param text: Text to normalize.
:param form: Unicode normalization form to apply to the text.
Options: "NFC", "NFKC", "NFD", "NFKD".
:returns: The normalized text.
"""
return normalize(form, text)
def _ascii_only(self, text: str) -> str:
"""
Convert the text to ASCII only.
Will remove accents from characters and replace them with ASCII characters.
Other non-ASCII characters will be removed.
:param text: Text to convert to ASCII only.
:returns: The text in ASCII only.
"""
# First normalize the text to NFKD to separate the characters and their diacritics
# Then encode it to ASCII and ignore any characters that can't be encoded
return self._normalize_unicode(text, "NFKD").encode("ascii", "ignore").decode("utf-8")
def _remove_empty_lines(self, text: str) -> str:
"""
Remove empty lines and lines that contain nothing but whitespaces from text.
:param text: Text to clean.
:returns: The text without empty lines.
"""
pages = text.split("\f")
cleaned_pages = ["\n".join(line for line in page.split("\n") if line.strip()) for page in pages]
return "\f".join(cleaned_pages)
def _remove_extra_whitespaces(self, text: str) -> str:
"""
Remove extra whitespaces from text.
:param text: Text to clean.
:returns: The text without extra whitespaces.
"""
texts = text.split("\f")
cleaned_text = [re.sub(r"\s\s+", " ", text).strip() for text in texts]
return "\f".join(cleaned_text)
def _remove_regex(self, text: str, regex: str) -> str:
"""
Remove substrings that match the specified regex from the text.
:param text: Text to clean.
:param regex: Regex to match and replace substrings by "".
:returns: The text without the substrings that match the regex.
"""
texts = text.split("\f")
cleaned_text = [re.sub(regex, "", text).strip() for text in texts]
return "\f".join(cleaned_text)
def _replace_regexes(self, text: str, replace_regexes: dict[str, str]) -> str:
"""
Replace substrings that match the specified regex patterns with custom replacement strings.
:param text: Text to clean.
:param replace_regexes: A dictionary mapping regex patterns to their replacement strings.
:returns: The text with the regex matches replaced by the specified strings.
"""
pages = text.split("\f")
cleaned_pages = []
for page in pages:
for pattern, replacement in replace_regexes.items():
page = re.sub(pattern, replacement, page)
cleaned_pages.append(page)
return "\f".join(cleaned_pages)
def _remove_substrings(self, text: str, substrings: list[str]) -> str:
"""
Remove all specified substrings from the text.
:param text: Text to clean.
:param substrings: Substrings to remove.
:returns: The text without the specified substrings.
"""
for substring in substrings:
text = text.replace(substring, "")
return text
def _remove_repeated_substrings(self, text: str) -> str:
"""
Remove any substrings from the text that occur repeatedly on every page. For example headers or footers.
Pages in the text need to be separated by form feed character "\f".
:param text: Text to clean.
:returns: The text without the repeated substrings.
"""
return self._find_and_remove_header_footer(
text, n_chars=300, n_first_pages_to_ignore=1, n_last_pages_to_ignore=1
)
def _find_and_remove_header_footer(
self, text: str, n_chars: int, n_first_pages_to_ignore: int, n_last_pages_to_ignore: int
) -> str:
"""
Heuristic to find footers and headers across different pages by searching for the longest common string.
Pages in the text need to be separated by form feed character "\f".
For headers, we only search in the first n_chars characters (for footer: last n_chars).
Note: This heuristic uses exact matches and therefore works well for footers like "Copyright 2019 by XXX",
but won't detect "Page 3 of 4" or similar.
:param n_chars: The number of first/last characters where the header/footer shall be searched in.
:param n_first_pages_to_ignore: The number of first pages to ignore
(e.g. TOCs often don't contain footer/header).
:param n_last_pages_to_ignore: The number of last pages to ignore.
:returns: The text without the found headers and footers.
"""
pages = text.split("\f")
# header
start_of_pages = [p[:n_chars] for p in pages[n_first_pages_to_ignore:-n_last_pages_to_ignore]]
found_header = self._find_longest_common_ngram(start_of_pages)
if found_header:
pages = [page.replace(found_header, "") for page in pages]
# footer
end_of_pages = [p[-n_chars:] for p in pages[n_first_pages_to_ignore:-n_last_pages_to_ignore]]
found_footer = self._find_longest_common_ngram(end_of_pages)
if found_footer:
pages = [page.replace(found_footer, "") for page in pages]
logger.debug(
"Removed header '{header}' and footer '{footer}' in document", header=found_header, footer=found_footer
)
return "\f".join(pages)
def _ngram(self, seq: str, n: int) -> Generator[str, None, None]:
"""
Return all ngrams of length n from a text sequence. Each ngram consists of n words split by whitespace.
:param seq: The sequence to generate ngrams from.
:param n: The length of the ngrams to generate.
:returns: A Generator generating all ngrams of length n from the given sequence.
"""
# In order to maintain the original whitespace, but still consider \n and \t for n-gram tokenization,
# we add a space here and remove it after creation of the ngrams again (see below)
seq = seq.replace("\n", " \n")
seq = seq.replace("\t", " \t")
words = seq.split(" ")
return (" ".join(words[i : i + n]).replace(" \n", "\n").replace(" \t", "\t") for i in range(len(words) - n + 1))
def _allngram(self, seq: str, min_ngram: int, max_ngram: int) -> set[str]:
"""
Generates all possible ngrams from a given sequence of text.
Considering all ngram lengths between the minimum and maximum length.
:param seq: The sequence to generate ngrams from.
:param min_ngram: The minimum length of ngram to consider.
:param max_ngram: The maximum length of ngram to consider.
:returns: A set of all ngrams from the given sequence.
"""
lengths = range(min_ngram, max_ngram) if max_ngram else range(min_ngram, len(seq))
ngrams = map(partial(self._ngram, seq), lengths)
return set(chain.from_iterable(ngrams))
def _find_longest_common_ngram(self, sequences: list[str], min_ngram: int = 3, max_ngram: int = 30) -> str:
"""
Find the longest common ngram across a list of text sequences (e.g. start of pages).
Considering all ngram lengths between the minimum and maximum length. Helpful for finding footers, headers etc.
Empty sequences are ignored.
:param sequences: The list of strings that shall be searched for common n_grams.
:param max_ngram: The maximum length of ngram to consider.
:param min_ngram: The minimum length of ngram to consider.
:returns: The longest ngram that all sequences have in common.
"""
sequences = [s for s in sequences if s] # filter empty sequences
if len(sequences) < 2:
# a single sequence has no ngram "in common" with any other; treating
# its own longest ngram as a repeated header/footer would wipe it
return ""
seqs_ngrams = map(partial(self._allngram, min_ngram=min_ngram, max_ngram=max_ngram), sequences)
intersection = reduce(set.intersection, seqs_ngrams)
longest = max(intersection, key=len, default="")
return longest if longest.strip() else ""
@@ -0,0 +1,198 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
from collections.abc import Callable
from typing import TYPE_CHECKING, Any, Literal
from haystack import Document, Pipeline, default_from_dict, default_to_dict, super_component
from haystack.components.preprocessors.document_cleaner import DocumentCleaner
from haystack.components.preprocessors.document_splitter import DocumentSplitter, Language
from haystack.utils import deserialize_callable, serialize_callable
@super_component
class DocumentPreprocessor:
"""
A SuperComponent that first splits and then cleans documents.
This component consists of a DocumentSplitter followed by a DocumentCleaner in a single pipeline.
It takes a list of documents as input and returns a processed list of documents.
Usage example:
```python
from haystack import Document
from haystack.components.preprocessors import DocumentPreprocessor
doc = Document(content="I love pizza!")
preprocessor = DocumentPreprocessor()
result = preprocessor.run(documents=[doc])
print(result["documents"])
```
"""
def __init__( # noqa: PLR0913 (too-many-arguments)
self,
*,
# --- DocumentSplitter arguments ---
split_by: Literal["function", "page", "passage", "period", "word", "line", "sentence"] = "word",
split_length: int = 250,
split_overlap: int = 0,
split_threshold: int = 0,
splitting_function: Callable[[str], list[str]] | None = None,
respect_sentence_boundary: bool = False,
language: Language = "en",
use_split_rules: bool = True,
extend_abbreviations: bool = True,
# --- DocumentCleaner arguments ---
remove_empty_lines: bool = True,
remove_extra_whitespaces: bool = True,
remove_repeated_substrings: bool = False,
keep_id: bool = False,
remove_substrings: list[str] | None = None,
remove_regex: str | None = None,
unicode_normalization: Literal["NFC", "NFKC", "NFD", "NFKD"] | None = None,
ascii_only: bool = False,
) -> None:
"""
Initialize a DocumentPreProcessor that first splits and then cleans documents.
**Splitter Parameters**:
:param split_by: The unit of splitting: "function", "page", "passage", "period", "word", "line", or "sentence".
:param split_length: The maximum number of units (words, lines, pages, and so on) in each split.
:param split_overlap: The number of overlapping units between consecutive splits.
:param split_threshold: The minimum number of units per split. If a split is smaller than this, it's merged
with the previous split.
:param splitting_function: A custom function for splitting if `split_by="function"`.
:param respect_sentence_boundary: If `True`, splits by words but tries not to break inside a sentence.
:param language: Language used by the sentence tokenizer if `split_by="sentence"` or
`respect_sentence_boundary=True`.
:param use_split_rules: Whether to apply additional splitting heuristics for the sentence splitter.
:param extend_abbreviations: Whether to extend the sentence splitter with curated abbreviations for certain
languages.
**Cleaner Parameters**:
:param remove_empty_lines: If `True`, removes empty lines.
:param remove_extra_whitespaces: If `True`, removes extra whitespaces.
:param remove_repeated_substrings: If `True`, removes repeated substrings like headers/footers across pages.
:param keep_id: If `True`, keeps the original document IDs.
:param remove_substrings: A list of strings to remove from the document content.
:param remove_regex: A regex pattern whose matches will be removed from the document content.
:param unicode_normalization: Unicode normalization form to apply to the text, for example `"NFC"`.
:param ascii_only: If `True`, converts text to ASCII only.
"""
# Store arguments for serialization
self.remove_empty_lines = remove_empty_lines
self.remove_extra_whitespaces = remove_extra_whitespaces
self.remove_repeated_substrings = remove_repeated_substrings
self.keep_id = keep_id
self.remove_substrings = remove_substrings
self.remove_regex = remove_regex
self.unicode_normalization = unicode_normalization
self.ascii_only = ascii_only
self.split_by = split_by
self.split_length = split_length
self.split_overlap = split_overlap
self.split_threshold = split_threshold
self.splitting_function = splitting_function
self.respect_sentence_boundary = respect_sentence_boundary
self.language = language
self.use_split_rules = use_split_rules
self.extend_abbreviations = extend_abbreviations
# Instantiate sub-components
splitter = DocumentSplitter(
split_by=self.split_by,
split_length=self.split_length,
split_overlap=self.split_overlap,
split_threshold=self.split_threshold,
splitting_function=self.splitting_function,
respect_sentence_boundary=self.respect_sentence_boundary,
language=self.language,
use_split_rules=self.use_split_rules,
extend_abbreviations=self.extend_abbreviations,
)
cleaner = DocumentCleaner(
remove_empty_lines=self.remove_empty_lines,
remove_extra_whitespaces=self.remove_extra_whitespaces,
remove_repeated_substrings=self.remove_repeated_substrings,
keep_id=self.keep_id,
remove_substrings=self.remove_substrings,
remove_regex=self.remove_regex,
unicode_normalization=self.unicode_normalization,
ascii_only=self.ascii_only,
)
# Build the Pipeline
pp = Pipeline()
pp.add_component("splitter", splitter)
pp.add_component("cleaner", cleaner)
# Connect the splitter output to cleaner
pp.connect("splitter.documents", "cleaner.documents")
self.pipeline = pp
# Define how pipeline inputs/outputs map to sub-component inputs/outputs
self.input_mapping = {
# The pipeline input "documents" feeds into "splitter.documents"
"documents": ["splitter.documents"]
}
# The pipeline output "documents" comes from "cleaner.documents"
self.output_mapping = {"cleaner.documents": "documents"}
if TYPE_CHECKING:
# fake method, never executed, but static analyzers will not complain about missing method
def run(self, *, documents: list[Document]) -> dict[str, list[Document]]: # noqa: D102
...
def warm_up(self) -> None: # noqa: D102
...
def to_dict(self) -> dict[str, Any]:
"""
Serialize SuperComponent to a dictionary.
:return:
Dictionary with serialized data.
"""
splitting_function = None
if self.splitting_function is not None:
splitting_function = serialize_callable(self.splitting_function)
return default_to_dict(
self,
remove_empty_lines=self.remove_empty_lines,
remove_extra_whitespaces=self.remove_extra_whitespaces,
remove_repeated_substrings=self.remove_repeated_substrings,
keep_id=self.keep_id,
remove_substrings=self.remove_substrings,
remove_regex=self.remove_regex,
unicode_normalization=self.unicode_normalization,
ascii_only=self.ascii_only,
split_by=self.split_by,
split_length=self.split_length,
split_overlap=self.split_overlap,
split_threshold=self.split_threshold,
splitting_function=splitting_function,
respect_sentence_boundary=self.respect_sentence_boundary,
language=self.language,
use_split_rules=self.use_split_rules,
extend_abbreviations=self.extend_abbreviations,
)
@classmethod
def from_dict(cls, data: dict[str, Any]) -> "DocumentPreprocessor":
"""
Deserializes the SuperComponent from a dictionary.
:param data:
Dictionary to deserialize from.
:returns:
Deserialized SuperComponent.
"""
splitting_function = data["init_parameters"].get("splitting_function", None)
if splitting_function:
data["init_parameters"]["splitting_function"] = deserialize_callable(splitting_function)
return default_from_dict(cls, data)
@@ -0,0 +1,499 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
from collections.abc import Callable
from copy import deepcopy
from typing import Any, Literal
from more_itertools import windowed
from haystack import Document, component, logging
from haystack.components.preprocessors.sentence_tokenizer import Language, SentenceSplitter, nltk_imports
from haystack.core.serialization import default_from_dict, default_to_dict
from haystack.utils import deserialize_callable, serialize_callable
logger = logging.getLogger(__name__)
# mapping of split by character, 'function' and 'sentence' don't split by character
_CHARACTER_SPLIT_BY_MAPPING = {"page": "\f", "passage": "\n\n", "period": ".", "word": " ", "line": "\n"}
@component
class DocumentSplitter:
"""
Splits long documents into smaller chunks.
This is a common preprocessing step during indexing. It helps Embedders create meaningful semantic representations
and prevents exceeding language model context limits.
The DocumentSplitter is compatible with the following DocumentStores:
- [Astra](https://docs.haystack.deepset.ai/docs/astradocumentstore)
- [Chroma](https://docs.haystack.deepset.ai/docs/chromadocumentstore) limited support, overlapping information is
not stored
- [Elasticsearch](https://docs.haystack.deepset.ai/docs/elasticsearch-document-store)
- [OpenSearch](https://docs.haystack.deepset.ai/docs/opensearch-document-store)
- [Pgvector](https://docs.haystack.deepset.ai/docs/pgvectordocumentstore)
- [Pinecone](https://docs.haystack.deepset.ai/docs/pinecone-document-store) limited support, overlapping
information is not stored
- [Qdrant](https://docs.haystack.deepset.ai/docs/qdrant-document-store)
- [Weaviate](https://docs.haystack.deepset.ai/docs/weaviatedocumentstore)
### Usage example
```python
from haystack import Document
from haystack.components.preprocessors import DocumentSplitter
doc = Document(content="Moonlight shimmered softly, wolves howled nearby, night enveloped everything.")
splitter = DocumentSplitter(split_by="word", split_length=3, split_overlap=0)
result = splitter.run(documents=[doc])
```
"""
def __init__(
self,
split_by: Literal["function", "page", "passage", "period", "word", "line", "sentence"] = "word",
split_length: int = 200,
split_overlap: int = 0,
split_threshold: int = 0,
splitting_function: Callable[[str], list[str]] | None = None,
respect_sentence_boundary: bool = False,
language: Language = "en",
use_split_rules: bool = True,
extend_abbreviations: bool = True,
*,
skip_empty_documents: bool = True,
) -> None:
"""
Initialize DocumentSplitter.
:param split_by: The unit for splitting your documents. Choose from:
- `word` for splitting by spaces (" ")
- `period` for splitting by periods (".")
- `page` for splitting by form feed ("\\f")
- `passage` for splitting by double line breaks ("\\n\\n")
- `line` for splitting each line ("\\n")
- `sentence` for splitting by NLTK sentence tokenizer
:param split_length: The maximum number of units in each split.
:param split_overlap: The number of overlapping units for each split.
:param split_threshold: The minimum number of units per split. If a split has fewer units
than the threshold, it's attached to the previous split.
:param splitting_function: Necessary when `split_by` is set to "function".
This is a function which must accept a single `str` as input and return a `list` of `str` as output,
representing the chunks after splitting.
:param respect_sentence_boundary: Choose whether to respect sentence boundaries when splitting by "word".
If True, uses NLTK to detect sentence boundaries, ensuring splits occur only between sentences.
:param language: Choose the language for the NLTK tokenizer. The default is English ("en").
:param use_split_rules: Choose whether to use additional split rules when splitting by `sentence`.
:param extend_abbreviations: Choose whether to extend NLTK's PunktTokenizer abbreviations with a list
of curated abbreviations, if available. This is currently supported for English ("en") and German ("de").
:param skip_empty_documents: Choose whether to skip documents with empty content. Default is True.
Set to False when downstream components in the Pipeline (like LLMDocumentContentExtractor) can extract text
from non-textual documents.
"""
self.split_by = split_by
self.split_length = split_length
self.split_overlap = split_overlap
self.split_threshold = split_threshold
self.splitting_function = splitting_function
self.respect_sentence_boundary = respect_sentence_boundary
self.language = language
self.use_split_rules = use_split_rules
self.extend_abbreviations = extend_abbreviations
self.skip_empty_documents = skip_empty_documents
self._init_checks(
split_by=split_by,
split_length=split_length,
split_overlap=split_overlap,
splitting_function=splitting_function,
respect_sentence_boundary=respect_sentence_boundary,
)
self._use_sentence_splitter = split_by == "sentence" or (respect_sentence_boundary and split_by == "word")
if self._use_sentence_splitter:
nltk_imports.check()
self.sentence_splitter: SentenceSplitter | None = None
def _init_checks(
self,
*,
split_by: str,
split_length: int,
split_overlap: int,
splitting_function: Callable | None,
respect_sentence_boundary: bool,
) -> None:
"""
Validates initialization parameters for DocumentSplitter.
:param split_by: The unit for splitting documents
:param split_length: The maximum number of units in each split
:param split_overlap: The number of overlapping units for each split
:param splitting_function: Custom function for splitting when split_by="function"
:param respect_sentence_boundary: Whether to respect sentence boundaries when splitting
:raises ValueError: If any parameter is invalid
"""
valid_split_by = ["function", "page", "passage", "period", "word", "line", "sentence"]
if split_by not in valid_split_by:
raise ValueError(f"split_by must be one of {', '.join(valid_split_by)}.")
if split_by == "function" and splitting_function is None:
raise ValueError("When 'split_by' is set to 'function', a valid 'splitting_function' must be provided.")
if split_length <= 0:
raise ValueError("split_length must be greater than 0.")
if split_overlap < 0:
raise ValueError("split_overlap must be greater than or equal to 0.")
if split_overlap >= split_length:
raise ValueError("split_overlap must be less than split_length.")
if respect_sentence_boundary and split_by != "word":
logger.warning(
"The 'respect_sentence_boundary' option is only supported for `split_by='word'`. "
"The option `respect_sentence_boundary` will be set to `False`."
)
self.respect_sentence_boundary = False
def warm_up(self) -> None:
"""
Warm up the DocumentSplitter by loading the sentence tokenizer.
"""
if self._use_sentence_splitter and 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,
)
@component.output_types(documents=list[Document])
def run(self, documents: list[Document]) -> dict[str, list[Document]]:
"""
Split documents into smaller parts.
Splits documents by the unit expressed in `split_by`, with a length of `split_length`
and an overlap of `split_overlap`.
: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 `page_number` to track the original page number.
- All other metadata copied from the original document.
:raises TypeError: if the input is not a list of Documents.
:raises ValueError: if the content of a document is None.
"""
if self._use_sentence_splitter and self.sentence_splitter is None:
self.warm_up()
if not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
raise TypeError("DocumentSplitter expects a List of Documents as input.")
split_docs: list[Document] = []
for doc in documents:
if doc.content is None:
raise ValueError(
f"DocumentSplitter only works with text documents but content for document ID {doc.id} is None."
)
if doc.content == "" and self.skip_empty_documents:
logger.warning("Document ID {doc_id} has an empty content. Skipping this document.", doc_id=doc.id)
continue
split_docs += self._split_document(doc)
return {"documents": split_docs}
def _split_document(self, doc: Document) -> list[Document]:
if self.split_by == "sentence" or self.respect_sentence_boundary:
return self._split_by_nltk_sentence(doc)
if self.split_by == "function" and self.splitting_function is not None:
return self._split_by_function(doc)
return self._split_by_character(doc)
def _split_by_nltk_sentence(self, doc: Document) -> list[Document]:
split_docs = []
result = self.sentence_splitter.split_sentences(doc.content) # type: ignore # None check is done in run()
units = [sentence["sentence"] for sentence in result]
if self.respect_sentence_boundary:
text_splits, splits_pages, splits_start_idxs = self._concatenate_sentences_based_on_word_amount(
sentences=units, split_length=self.split_length, split_overlap=self.split_overlap
)
else:
text_splits, splits_pages, splits_start_idxs = self._concatenate_units(
elements=units,
split_length=self.split_length,
split_overlap=self.split_overlap,
split_threshold=self.split_threshold,
)
metadata = deepcopy(doc.meta)
metadata["source_id"] = doc.id
split_docs += self._create_docs_from_splits(
text_splits=text_splits, splits_pages=splits_pages, splits_start_idxs=splits_start_idxs, meta=metadata
)
return split_docs
def _split_by_character(self, doc: Document) -> list[Document]:
split_at = _CHARACTER_SPLIT_BY_MAPPING[self.split_by]
units = doc.content.split(split_at) # type: ignore[union-attr]
# Add the delimiter back to all units except the last one
for i in range(len(units) - 1):
units[i] += split_at
text_splits, splits_pages, splits_start_idxs = self._concatenate_units(
units, self.split_length, self.split_overlap, self.split_threshold
)
metadata = deepcopy(doc.meta)
metadata["source_id"] = doc.id
return self._create_docs_from_splits(
text_splits=text_splits, splits_pages=splits_pages, splits_start_idxs=splits_start_idxs, meta=metadata
)
def _split_by_function(self, doc: Document) -> list[Document]:
# the check for None is done already in the run method
splits = self.splitting_function(doc.content) # type: ignore
docs: list[Document] = []
for s in splits:
meta = deepcopy(doc.meta)
meta["source_id"] = doc.id
docs.append(Document(content=s, meta=meta))
return docs
def _concatenate_units(
self, elements: list[str], split_length: int, split_overlap: int, split_threshold: int
) -> tuple[list[str], list[int], list[int]]:
"""
Concatenates the elements into parts of split_length units.
Keeps track of the original page number that each element belongs. If the length of the current units is less
than the pre-defined `split_threshold`, it does not create a new split. Instead, it concatenates the current
units with the last split, preventing the creation of excessively small splits.
"""
text_splits: list[str] = []
splits_pages: list[int] = []
splits_start_idxs: list[int] = []
cur_start_idx = 0
cur_page = 1
segments = windowed(elements, n=split_length, step=split_length - split_overlap)
for seg in segments:
current_units = [unit for unit in seg if unit is not None]
txt = "".join(current_units)
# check if length of current units is below split_threshold
if len(current_units) < split_threshold and len(text_splits) > 0:
# concatenate the last split with the current one
text_splits[-1] += txt
# NOTE: If skip_empty_documents is True, this line skips documents that have content=""
elif not self.skip_empty_documents or len(txt) > 0:
text_splits.append(txt)
splits_pages.append(cur_page)
splits_start_idxs.append(cur_start_idx)
processed_units = current_units[: split_length - split_overlap]
cur_start_idx += len("".join(processed_units))
if self.split_by == "page":
num_page_breaks = len(processed_units)
else:
num_page_breaks = sum(processed_unit.count("\f") for processed_unit in processed_units)
cur_page += num_page_breaks
return text_splits, splits_pages, splits_start_idxs
def _create_docs_from_splits(
self, text_splits: list[str], splits_pages: list[int], splits_start_idxs: list[int], meta: dict[str, Any]
) -> list[Document]:
"""
Creates Document objects from splits enriching them with page number and the metadata of the original document.
"""
documents: list[Document] = []
for i, (txt, split_idx) in enumerate(zip(text_splits, splits_start_idxs, strict=True)):
copied_meta = deepcopy(meta)
copied_meta["page_number"] = splits_pages[i]
copied_meta["split_id"] = i
copied_meta["split_idx_start"] = split_idx
doc = Document(content=txt, meta=copied_meta)
documents.append(doc)
if self.split_overlap <= 0:
continue
doc.meta["_split_overlap"] = []
if i == 0:
continue
doc_start_idx = splits_start_idxs[i]
previous_doc = documents[i - 1]
previous_doc_start_idx = splits_start_idxs[i - 1]
self._add_split_overlap_information(doc, doc_start_idx, previous_doc, previous_doc_start_idx)
return documents
@staticmethod
def _add_split_overlap_information(
current_doc: Document, current_doc_start_idx: int, previous_doc: Document, previous_doc_start_idx: int
) -> None:
"""
Adds split overlap information to the current and previous Document's meta.
:param current_doc: The Document that is being split.
:param current_doc_start_idx: The starting index of the current Document.
:param previous_doc: The Document that was split before the current Document.
:param previous_doc_start_idx: The starting index of the previous Document.
"""
overlapping_range = (current_doc_start_idx - previous_doc_start_idx, len(previous_doc.content)) # type: ignore
if overlapping_range[0] < overlapping_range[1]:
overlapping_str = previous_doc.content[overlapping_range[0] : overlapping_range[1]] # type: ignore
if current_doc.content.startswith(overlapping_str): # type: ignore
# add split overlap information to this Document regarding the previous Document
current_doc.meta["_split_overlap"].append({"doc_id": previous_doc.id, "range": overlapping_range})
# add split overlap information to previous Document regarding this Document
overlapping_range = (0, overlapping_range[1] - overlapping_range[0])
previous_doc.meta["_split_overlap"].append({"doc_id": current_doc.id, "range": overlapping_range})
def to_dict(self) -> dict[str, Any]:
"""
Serializes the component to a dictionary.
"""
serialized = default_to_dict(
self,
split_by=self.split_by,
split_length=self.split_length,
split_overlap=self.split_overlap,
split_threshold=self.split_threshold,
respect_sentence_boundary=self.respect_sentence_boundary,
language=self.language,
use_split_rules=self.use_split_rules,
extend_abbreviations=self.extend_abbreviations,
skip_empty_documents=self.skip_empty_documents,
)
if self.splitting_function:
serialized["init_parameters"]["splitting_function"] = serialize_callable(self.splitting_function)
return serialized
@classmethod
def from_dict(cls, data: dict[str, Any]) -> "DocumentSplitter":
"""
Deserializes the component from a dictionary.
"""
init_params = data.get("init_parameters", {})
splitting_function = init_params.get("splitting_function", None)
if splitting_function:
init_params["splitting_function"] = deserialize_callable(splitting_function)
return default_from_dict(cls, data)
@staticmethod
def _concatenate_sentences_based_on_word_amount(
sentences: list[str], split_length: int, split_overlap: int
) -> tuple[list[str], list[int], list[int]]:
"""
Groups the sentences into chunks of `split_length` words while respecting sentence boundaries.
This function is only used when splitting by `word` and `respect_sentence_boundary` is set to `True`, i.e.:
with NLTK sentence tokenizer.
:param sentences: The list of sentences to split.
:param split_length: The maximum number of words in each split.
:param split_overlap: The number of overlapping words in each split.
:returns: A tuple containing the concatenated sentences, the start page numbers, and the start indices.
"""
# chunk information
chunk_word_count = 0
chunk_starting_page_number = 1
chunk_start_idx = 0
current_chunk: list[str] = []
# output lists
split_start_page_numbers = []
list_of_splits: list[list[str]] = []
split_start_indices = []
for sentence_idx, sentence in enumerate(sentences):
current_chunk.append(sentence)
chunk_word_count += len(sentence.split())
next_sentence_word_count = (
len(sentences[sentence_idx + 1].split()) if sentence_idx < len(sentences) - 1 else 0
)
# Number of words in the current chunk plus the next sentence is larger than the split_length,
# or we reached the last sentence
if (chunk_word_count + next_sentence_word_count) > split_length or sentence_idx == len(sentences) - 1:
# Save current chunk and start a new one
list_of_splits.append(current_chunk)
split_start_page_numbers.append(chunk_starting_page_number)
split_start_indices.append(chunk_start_idx)
# Get the number of sentences that overlap with the next chunk
num_sentences_to_keep = DocumentSplitter._number_of_sentences_to_keep(
sentences=current_chunk, split_length=split_length, split_overlap=split_overlap
)
# Set up information for the new chunk
if num_sentences_to_keep > 0:
# Processed sentences are the ones that are not overlapping with the next chunk
processed_sentences = current_chunk[:-num_sentences_to_keep]
chunk_starting_page_number += sum(sent.count("\f") for sent in processed_sentences)
chunk_start_idx += len("".join(processed_sentences))
# Next chunk starts with the sentences that were overlapping with the previous chunk
current_chunk = current_chunk[-num_sentences_to_keep:]
chunk_word_count = sum(len(s.split()) for s in current_chunk)
else:
# Here processed_sentences is the same as current_chunk since there is no overlap
chunk_starting_page_number += sum(sent.count("\f") for sent in current_chunk)
chunk_start_idx += len("".join(current_chunk))
current_chunk = []
chunk_word_count = 0
# Concatenate the sentences together within each split
text_splits = []
for split in list_of_splits:
text = "".join(split)
if len(text) > 0:
text_splits.append(text)
return text_splits, split_start_page_numbers, split_start_indices
@staticmethod
def _number_of_sentences_to_keep(sentences: list[str], split_length: int, split_overlap: int) -> int:
"""
Returns the number of sentences to keep in the next chunk based on the `split_overlap` and `split_length`.
:param sentences: The list of sentences to split.
:param split_length: The maximum number of words in each split.
:param split_overlap: The number of overlapping words in each split.
:returns: The number of sentences to keep in the next chunk.
"""
# If the split_overlap is 0, we don't need to keep any sentences
if split_overlap == 0:
return 0
num_sentences_to_keep = 0
num_words = 0
# Next overlapping Document should not start exactly the same as the previous one, so we skip the first sentence
for sent in reversed(sentences[1:]):
num_words += len(sent.split())
# If the number of words is larger than the split_length then don't add any more sentences
if num_words > split_length:
break
num_sentences_to_keep += 1
if num_words > split_overlap:
break
return num_sentences_to_keep
@@ -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)
@@ -0,0 +1,156 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
from dataclasses import replace
from typing import Any, Literal
from haystack import Document, component, default_from_dict, default_to_dict
from haystack.components.preprocessors import DocumentSplitter
@component
class HierarchicalDocumentSplitter:
"""
Splits a documents into different block sizes building a hierarchical tree structure of blocks of different sizes.
The root node of the tree is the original document, the leaf nodes are the smallest blocks. The blocks in between
are connected such that the smaller blocks are children of the parent-larger blocks.
## Usage example
```python
from haystack import Document
from haystack.components.preprocessors import HierarchicalDocumentSplitter
doc = Document(content="This is a simple test document")
splitter = HierarchicalDocumentSplitter(block_sizes={3, 2}, split_overlap=0, split_by="word")
splitter.run([doc])
# >> {'documents': [Document(id=3f7..., content: 'This is a simple test document', meta: {'block_size': 0, 'parent_id': None, 'children_ids': ['5ff..', '8dc..'], 'level': 0}),
# >> Document(id=5ff.., content: 'This is a ', meta: {'block_size': 3, 'parent_id': '3f7..', 'children_ids': ['f19..', '52c..'], 'level': 1, 'source_id': '3f7..', 'page_number': 1, 'split_id': 0, 'split_idx_start': 0}),
# >> Document(id=8dc.., content: 'simple test document', meta: {'block_size': 3, 'parent_id': '3f7..', 'children_ids': ['39d..', 'e23..'], 'level': 1, 'source_id': '3f7..', 'page_number': 1, 'split_id': 1, 'split_idx_start': 10}),
# >> Document(id=f19.., content: 'This is ', meta: {'block_size': 2, 'parent_id': '5ff..', 'children_ids': [], 'level': 2, 'source_id': '5ff..', 'page_number': 1, 'split_id': 0, 'split_idx_start': 0}),
# >> Document(id=52c.., content: 'a ', meta: {'block_size': 2, 'parent_id': '5ff..', 'children_ids': [], 'level': 2, 'source_id': '5ff..', 'page_number': 1, 'split_id': 1, 'split_idx_start': 8}),
# >> Document(id=39d.., content: 'simple test ', meta: {'block_size': 2, 'parent_id': '8dc..', 'children_ids': [], 'level': 2, 'source_id': '8dc..', 'page_number': 1, 'split_id': 0, 'split_idx_start': 0}),
# >> Document(id=e23.., content: 'document', meta: {'block_size': 2, 'parent_id': '8dc..', 'children_ids': [], 'level': 2, 'source_id': '8dc..', 'page_number': 1, 'split_id': 1, 'split_idx_start': 12})]}
```
""" # noqa: E501
def __init__(
self,
block_sizes: set[int],
split_overlap: int = 0,
split_by: Literal["word", "sentence", "page", "passage"] = "word",
) -> None:
"""
Initialize HierarchicalDocumentSplitter.
:param block_sizes: Set of block sizes to split the document into. The blocks are split in descending order.
:param split_overlap: The number of overlapping units for each split.
:param split_by: The unit for splitting your documents.
:raises ValueError: If `block_sizes` is empty, if `split_overlap` is negative, or if `split_overlap` is
greater than or equal to the smallest value in `block_sizes`.
"""
if not block_sizes:
raise ValueError("block_sizes must not be empty. Provide at least one block size.")
if split_overlap < 0:
raise ValueError("split_overlap must be greater than or equal to 0.")
smallest_block_size = min(block_sizes)
if split_overlap >= smallest_block_size:
raise ValueError(
f"split_overlap ({split_overlap}) must be less than the smallest value in block_sizes "
f"({smallest_block_size}). Reduce split_overlap or increase the smallest block size."
)
self.block_sizes = sorted(set(block_sizes), reverse=True)
self.splitters: dict[int, DocumentSplitter] = {}
self.split_overlap = split_overlap
self.split_by = split_by
self._build_block_sizes()
@component.output_types(documents=list[Document])
def run(self, documents: list[Document]) -> dict[str, list[Document]]:
"""
Builds a hierarchical document structure for each document in a list of documents.
:param documents: List of Documents to split into hierarchical blocks.
:returns: List of HierarchicalDocument
"""
hierarchical_docs = []
for doc in documents:
hierarchical_docs.extend(self.build_hierarchy_from_doc(doc))
return {"documents": hierarchical_docs}
def _build_block_sizes(self) -> None:
for block_size in self.block_sizes:
self.splitters[block_size] = DocumentSplitter(
split_length=block_size, split_overlap=self.split_overlap, split_by=self.split_by
)
@staticmethod
def _add_meta_data(document: Document) -> Document:
new_meta = {**document.meta, "__block_size": 0, "__parent_id": None, "__children_ids": [], "__level": 0}
return replace(document, meta=new_meta)
def build_hierarchy_from_doc(self, document: Document) -> list[Document]:
"""
Build a hierarchical tree document structure from a single document.
Given a document, this function splits the document into hierarchical blocks of different sizes represented
as HierarchicalDocument objects.
:param document: Document to split into hierarchical blocks.
:returns:
List of HierarchicalDocument
"""
root = self._add_meta_data(document)
current_level_nodes = [root]
all_docs = []
for block in self.block_sizes:
next_level_nodes = []
for doc in current_level_nodes:
splitted_docs = self.splitters[block].run([doc])
child_docs = splitted_docs["documents"]
# if it's only one document skip
if len(child_docs) == 1:
next_level_nodes.append(doc)
continue
for child_doc in child_docs:
child_doc = self._add_meta_data(child_doc)
child_doc.meta["__level"] = doc.meta["__level"] + 1
child_doc.meta["__block_size"] = block
child_doc.meta["__parent_id"] = doc.id
all_docs.append(child_doc)
doc.meta["__children_ids"].append(child_doc.id)
next_level_nodes.append(child_doc)
current_level_nodes = next_level_nodes
return [root] + all_docs
def to_dict(self) -> dict[str, Any]:
"""
Returns a dictionary representation of the component.
:returns:
Serialized dictionary representation of the component.
"""
return default_to_dict(
self, block_sizes=self.block_sizes, split_overlap=self.split_overlap, split_by=self.split_by
)
@classmethod
def from_dict(cls, data: dict[str, Any]) -> "HierarchicalDocumentSplitter":
"""
Deserialize this component from a dictionary.
:param data:
The dictionary to deserialize and create the component.
:returns:
The deserialized component.
"""
return default_from_dict(cls, data)
@@ -0,0 +1,382 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
import re
from typing import Literal
from haystack import Document, component, logging
from haystack.components.preprocessors import DocumentSplitter
logger = logging.getLogger(__name__)
@component
class MarkdownHeaderSplitter:
"""
Split documents at ATX-style Markdown headers (#), with optional secondary splitting.
This component processes text documents by:
- Splitting them into chunks at Markdown headers (e.g., '#', '##', etc.), preserving header hierarchy as metadata.
- Optionally applying a secondary split (by word, passage, period, or line) to each chunk
(using haystack's DocumentSplitter).
- Preserving and propagating metadata such as parent headers, page numbers, and split IDs.
"""
def __init__(
self,
*,
page_break_character: str = "\f",
keep_headers: bool = True,
header_split_levels: list[int] | None = None,
secondary_split: Literal["word", "passage", "period", "line"] | None = None,
split_length: int = 200,
split_overlap: int = 0,
split_threshold: int = 0,
skip_empty_documents: bool = True,
) -> None:
"""
Initialize the MarkdownHeaderSplitter.
:param page_break_character: Character used to identify page breaks. Defaults to form feed ("\f").
:param keep_headers: If True, headers are kept in the content. If False, headers are moved to metadata.
Defaults to True.
:param header_split_levels: List of header levels (16) to split on. For example, `[1, 2]` splits only
on `#` and `##` headers, merging content under deeper headers into the preceding chunk. Defaults to
all levels `[1, 2, 3, 4, 5, 6]`.
:param secondary_split: Optional secondary split condition after header splitting.
Options are None, "word", "passage", "period", "line". Defaults to None.
:param split_length: The maximum number of units in each split when using secondary splitting. Defaults to 200.
:param split_overlap: The number of overlapping units for each split when using secondary splitting.
Defaults to 0.
:param split_threshold: The minimum number of units per split when using secondary splitting. Defaults to 0.
:param skip_empty_documents: Choose whether to skip documents with empty content. Default is True.
Set to False when downstream components in the Pipeline (like LLMDocumentContentExtractor) can extract text
from non-textual documents.
"""
if header_split_levels is None:
header_split_levels = [1, 2, 3, 4, 5, 6]
if not isinstance(header_split_levels, list) or len(header_split_levels) == 0:
raise ValueError("header_split_levels must be a non-empty list.")
invalid = [lvl for lvl in header_split_levels if not isinstance(lvl, int) or lvl < 1 or lvl > 6]
if invalid:
raise ValueError(
f"header_split_levels contains invalid values: {invalid}. All levels must be integers between 1 and 6."
)
if len(header_split_levels) != len(set(header_split_levels)):
raise ValueError("header_split_levels must not contain duplicate values.")
self.page_break_character = page_break_character
self.secondary_split = secondary_split
self.split_length = split_length
self.split_overlap = split_overlap
self.split_threshold = split_threshold
self.skip_empty_documents = skip_empty_documents
self.keep_headers = keep_headers
self.header_split_levels = header_split_levels
self._header_split_levels_set = set(header_split_levels)
self._header_pattern = re.compile(r"(?m)^(#{1,6}) (.+)$") # ATX-style .md-headers
# Matches fenced code blocks delimited by triple backticks (```) or triple tildes (~~~).
# Broken down:
# ^ - fence must start at the beginning of a line (MULTILINE)
# (?P<fence>`{3,}|~{3,})
# - named capture group "fence": three or more backticks OR three or
# more tildes. Capturing it allows the closing fence to be matched
# with a backreference, so ```-opened blocks must close with ```
# and ~~~-opened blocks must close with ~~~.
# [^\n]* - optional language identifier (e.g. "python") and any other text
# on the opening fence line, up to the newline
# \n - newline ending the opening fence line
# .*? - the code block body, matched lazily (DOTALL so . matches newlines)
# ^(?P=fence) - closing fence: must be identical to the opening fence (backreference),
# and must start at the beginning of a line
# \s*$ - optional trailing whitespace after the closing fence
self._code_block_pattern = re.compile(
r"^(?P<fence>`{3,}|~{3,})[^\n]*\n.*?^(?P=fence)\s*$", re.MULTILINE | re.DOTALL
)
self._is_warmed_up = False
# initialize secondary_splitter only if needed
if self.secondary_split:
self.secondary_splitter = DocumentSplitter(
split_by=self.secondary_split,
split_length=self.split_length,
split_overlap=self.split_overlap,
split_threshold=self.split_threshold,
)
def warm_up(self) -> None:
"""
Warm up the MarkdownHeaderSplitter.
"""
if self.secondary_split and not self._is_warmed_up:
self.secondary_splitter.warm_up()
self._is_warmed_up = True
def _code_block_spans(self, text: str) -> list[tuple[int, int]]:
"""Return the (start, end) character spans of all fenced code blocks in text."""
return [(m.start(), m.end()) for m in self._code_block_pattern.finditer(text)]
def _split_text_by_markdown_headers(self, text: str, doc_id: str) -> list[dict]:
"""Split text by ATX-style headers (#) and create chunks with appropriate metadata."""
logger.debug("Splitting text by markdown headers")
# Pre-compute fenced code block spans so that # lines inside code blocks (e.g. Python comments) are not
# mistaken for Markdown headers.
code_spans = self._code_block_spans(text)
# find headers at the configured levels only, excluding any that fall inside a code block. Content between
# skipped headers is absorbed into the preceding chunk's span since end = next_match.start().
matches = [
m
for m in re.finditer(self._header_pattern, text)
if len(m.group(1)) in self._header_split_levels_set
and not any(start <= m.start() < end for start, end in code_spans)
]
# return unsplit if no headers found
if not matches:
logger.info(
"No headers found in document {doc_id}; returning full document as single chunk.", doc_id=doc_id
)
return [{"content": text, "meta": {}}]
# process headers and build chunks
chunks: list[dict] = []
header_stack: list[str | None] = [None] * 6
pending_headers: list[str] = [] # store empty headers to prepend to next content
has_content = False # flag to track if any header has content
for i, match in enumerate(matches):
# extract header info
header_prefix = match.group(1)
header_text = match.group(2)
level = len(header_prefix)
# get content
start = match.end()
end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
content = text[start:end]
# update header stack to track nesting
header_stack[level - 1] = header_text
for j in range(level, 6):
header_stack[j] = None
# skip splits w/o content
if not content.strip(): # this strip is needed to avoid counting whitespace as content
if self.keep_headers:
header_line = f"{header_prefix} {header_text}"
pending_headers.append(header_line)
continue
has_content = True # at least one header has content
# Build parent metadata from the current header stack so the first child of a
# contentful section still inherits its full ancestor chain.
parent_headers = [h for h in header_stack[: level - 1] if h is not None]
logger.debug(
"Creating chunk for header '{header_text}' at level {level}", header_text=header_text, level=level
)
if self.keep_headers:
header_line = f"{header_prefix} {header_text}"
# add pending & current header to content
chunk_content = ""
if pending_headers:
chunk_content += "\n".join(pending_headers) + "\n"
chunk_content += f"{header_line}{content}"
chunks.append(
{"content": chunk_content, "meta": {"header": header_text, "parent_headers": parent_headers}}
)
pending_headers = [] # reset pending headers
else:
chunks.append({"content": content, "meta": {"header": header_text, "parent_headers": parent_headers}})
# return doc unchunked if no headers have content
if not has_content:
logger.info(
"Document {doc_id} contains only headers with no content; returning original document.", doc_id=doc_id
)
return [{"content": text, "meta": {}}]
return chunks
def _apply_secondary_splitting(self, documents: list[Document]) -> list[Document]:
"""
Apply secondary splitting while preserving header metadata and structure.
Ensures page counting is maintained across splits.
"""
result_docs = []
current_split_id = 0 # track split_id across all secondary splits from the same parent
for doc in documents:
if doc.content is None:
result_docs.append(doc)
continue
content_for_splitting: str = doc.content
if not self.keep_headers: # skip header extraction if keep_headers
# extract header information
header_match = re.match(self._header_pattern, doc.content)
if header_match:
content_for_splitting = doc.content[header_match.end() :]
# track page from meta
current_page = doc.meta.get("page_number", 1)
# create a clean meta dict without split_id for secondary splitting
clean_meta = {k: v for k, v in doc.meta.items() if k != "split_id"}
secondary_splits = self.secondary_splitter.run(
documents=[Document(content=content_for_splitting, meta=clean_meta)]
)["documents"]
# split processing
for i, split in enumerate(secondary_splits):
# calculate page number for this split
if i > 0 and secondary_splits[i - 1].content:
current_page = self._update_page_number_with_breaks(secondary_splits[i - 1].content, current_page)
# set page number and split_id to meta
split.meta["page_number"] = current_page
split.meta["split_id"] = current_split_id
# ensure source_id is preserved from the original document
if "source_id" in doc.meta:
split.meta["source_id"] = doc.meta["source_id"]
current_split_id += 1
# preserve header metadata if we're not keeping headers in content
if not self.keep_headers:
for key in ["header", "parent_headers"]:
if key in doc.meta:
split.meta[key] = doc.meta[key]
result_docs.append(split)
logger.debug(
"Secondary splitting complete. Final count: {final_count} documents.", final_count=len(result_docs)
)
return result_docs
def _update_page_number_with_breaks(self, content: str | None, current_page: int) -> int:
"""
Update page number based on page breaks in content.
:param content: Content to check for page breaks
:param current_page: Current page number
:return: New current page number
"""
if not isinstance(content, str):
return current_page
page_breaks = content.count(self.page_break_character)
new_page_number = current_page + page_breaks
if page_breaks > 0:
logger.debug(
"Found {page_breaks} page breaks, page number updated: {old}{new}",
page_breaks=page_breaks,
old=current_page,
new=new_page_number,
)
return new_page_number
def _split_documents_by_markdown_headers(self, documents: list[Document]) -> list[Document]:
"""Split a list of documents by markdown headers, preserving metadata."""
result_docs = []
for doc in documents:
logger.debug("Splitting document with id={doc_id}", doc_id=doc.id)
# mypy: doc.content is Optional[str], so we must check for None before passing to splitting method
if doc.content is None:
continue
splits = self._split_text_by_markdown_headers(doc.content, doc.id)
docs = []
current_page = doc.meta.get("page_number", 1) if doc.meta else 1
total_page_breaks = doc.content.count(self.page_break_character)
logger.debug(
"Processing document with id={doc_id}: starting at page {start_page}, "
"contains {page_breaks} page breaks in total",
doc_id=doc.id,
start_page=current_page,
page_breaks=total_page_breaks,
)
for split_idx, split in enumerate(splits):
meta = doc.meta.copy() if doc.meta else {}
meta.update({"source_id": doc.id, "page_number": current_page, "split_id": split_idx})
if split.get("meta"):
meta.update(split["meta"])
current_page = self._update_page_number_with_breaks(split["content"], current_page)
docs.append(Document(content=split["content"], meta=meta))
logger.debug(
"Split into {num_docs} documents for id={doc_id}, final page: {current_page}",
num_docs=len(docs),
doc_id=doc.id,
current_page=current_page,
)
result_docs.extend(docs)
return result_docs
@component.output_types(documents=list[Document])
def run(self, documents: list[Document]) -> dict[str, list[Document]]:
"""
Run the markdown header splitter with optional secondary splitting.
:param documents: List of 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 `page_number` to track the original page number.
- A metadata field `split_id` to identify the split chunk index within its parent document.
- All other metadata copied from the original document.
:raises ValueError: If a document has `None` content.
:raises TypeError: If a document's content is not a string.
"""
if self.secondary_split and not self._is_warmed_up:
self.warm_up()
# validate input documents
for doc in documents:
if doc.content is None:
raise ValueError(
"MarkdownHeaderSplitter only works with text documents but content for document ID"
f" {doc.id} is None."
)
if not isinstance(doc.content, str):
raise TypeError("MarkdownHeaderSplitter only works with text documents (str content).")
final_docs = []
for doc in documents:
# handle empty documents
if not doc.content or not doc.content.strip(): # avoid counting whitespace as content
if self.skip_empty_documents:
logger.warning("Document ID {doc_id} has an empty content. Skipping this document.", doc_id=doc.id)
continue
# keep empty documents
final_docs.append(doc)
logger.warning(
"Document ID {doc_id} has an empty content. Keeping this document as per configuration.",
doc_id=doc.id,
)
continue
# split this document by headers
header_split_docs = self._split_documents_by_markdown_headers([doc])
# apply secondary splitting if configured
if self.secondary_split:
doc_splits = self._apply_secondary_splitting(header_split_docs)
else:
doc_splits = header_split_docs
final_docs.extend(doc_splits)
return {"documents": final_docs}
@@ -0,0 +1,612 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
import ast
import math
from dataclasses import dataclass, field
from typing import Any
from haystack import Document, component, logging
from haystack.components.preprocessors.document_splitter import DocumentSplitter
logger = logging.getLogger(__name__)
@dataclass
class _CodeUnit:
"""One syntactic split unit (function, class header, method, imports block, statement, ...)."""
source: str
start_line: int
end_line: int
kind: str
name: str | None = None
class_name: str | None = None
class_signature: str | None = None
decorators: list[str] = field(default_factory=list)
docstring: str | None = None
@component
class PythonCodeSplitter:
"""
Split Python source code into syntax-aware chunks.
The component parses each source with :mod:`ast` into *units* (module docstring,
consecutive ``import`` blocks, top-level functions, class headers, methods, nested
classes, and remaining statements) and merges them greedily in source order toward
``max_effective_lines`` per chunk, where effective lines are
``ceil(len(source) / expected_chars_per_line)``. Functions and methods are kept
whole; the resulting chunks read top-to-bottom like the original file with comments
and blank lines preserved.
A function whose effective length exceeds ``oversized_factor * max_effective_lines``
is the only case where chunks may overlap: it is broken down with a line-based
secondary split (:class:`DocumentSplitter`, ``split_by="line"``) and the resulting
pieces carry ``secondary_split=True`` along with the originating function's metadata.
The primary split never adds overlap.
Per-chunk metadata: ``source_id``, ``split_id``, ``start_line``, ``end_line``,
``unit_kinds``; plus ``include_classes``, ``decorators``, and ``docstrings`` (when
``strip_docstrings=True``) where applicable. ``file_name`` and any other parent
document meta are propagated.
Usage example:
```python
from haystack import Document
from haystack.components.preprocessors import PythonCodeSplitter
source = '''
\"\"\"Example module.\"\"\"
from math import sqrt
class Circle:
def __init__(self, r: float) -> None:
self.r = r
def area(self) -> float:
return 3.14159 * self.r * self.r
'''
splitter = PythonCodeSplitter(min_effective_lines=4, max_effective_lines=6)
result = splitter.run(documents=[Document(content=source, meta={"file_name": "circle.py"})])
for chunk in result["documents"]:
print(chunk.meta["start_line"], chunk.meta["end_line"], chunk.meta.get("include_classes"))
```
Pass ``strip_docstrings=True`` to move docstrings out of the chunk content and into
each chunk's ``meta["docstrings"]`` list. This is useful for RAG when docstrings are
large: stripping shrinks the stored content while the docstring text can still
influence retrieval via ``meta_fields_to_embed=["docstrings"]`` on the embedder.
"""
def __init__(
self,
*,
min_effective_lines: int = 20,
max_effective_lines: int = 100,
expected_chars_per_line: int = 45,
oversized_factor: int = 3,
strip_docstrings: bool = False,
preserve_class_definition: bool = True,
secondary_split_overlap: int = 5,
secondary_split_length: int | None = None,
) -> None:
"""
Initialize the PythonCodeSplitter.
:param min_effective_lines: Minimum effective lines per chunk. While the running
chunk is below this threshold the splitter keeps merging in the next unit.
:param max_effective_lines: Target effective lines per chunk. Units are merged
greedily while doing so brings the running total closer to this target.
:param expected_chars_per_line: Used to convert characters into effective lines as
``ceil(len(source) / expected_chars_per_line)``; long lines count as more than one.
:param oversized_factor: A function whose effective length exceeds
``oversized_factor * max_effective_lines`` triggers the line-based secondary
split with overlap.
:param strip_docstrings: If ``True``, function/method/class docstrings are moved
from the chunk content into ``meta["docstrings"]`` (source order). The
module-level docstring is kept in place since it is itself a top-level unit.
:param preserve_class_definition: If ``True`` (default), chunks that contain class
members but not the class header are prefixed with the bare class signature
(decorators plus the ``class Foo(...):`` lines) in source order.
:param secondary_split_overlap: Line overlap for the secondary splitter; only used
in the oversized fallback. The primary AST split never adds overlap.
:param secondary_split_length: Lines per chunk for the secondary splitter.
Defaults to ``max_effective_lines`` when ``None``.
:raises ValueError: If any parameter is invalid (negative, zero where positive is
required, or ``min_effective_lines > max_effective_lines``).
"""
if min_effective_lines < 1:
raise ValueError("min_effective_lines must be at least 1.")
if max_effective_lines < 1:
raise ValueError("max_effective_lines must be at least 1.")
if min_effective_lines > max_effective_lines:
raise ValueError("min_effective_lines must not be greater than max_effective_lines.")
if expected_chars_per_line < 1:
raise ValueError("expected_chars_per_line must be at least 1.")
if oversized_factor < 1:
raise ValueError("oversized_factor must be at least 1.")
if secondary_split_overlap < 0:
raise ValueError("secondary_split_overlap must be non-negative.")
if secondary_split_length is not None and secondary_split_length < 1:
raise ValueError("secondary_split_length must be at least 1.")
self.min_effective_lines = min_effective_lines
self.max_effective_lines = max_effective_lines
self.expected_chars_per_line = expected_chars_per_line
self.oversized_factor = oversized_factor
self.strip_docstrings = strip_docstrings
self.preserve_class_definition = preserve_class_definition
self.secondary_split_overlap = secondary_split_overlap
self.secondary_split_length = secondary_split_length
def _effective_lines(self, text: str) -> int:
"""Return the number of *effective lines* for ``text`` (see class docstring)."""
if not text:
return 0
return max(1, math.ceil(len(text) / self.expected_chars_per_line))
def _is_oversized(self, unit: "_CodeUnit") -> bool:
"""Return ``True`` if ``unit`` should trigger the secondary line-based split."""
return self._effective_lines(unit.source) > self.oversized_factor * self.max_effective_lines
@staticmethod
def _slice_lines(source_lines: list[str], start: int, end: int) -> str:
"""Slice ``source_lines`` between the 1-indexed ``start`` and ``end`` (inclusive)."""
start = max(start, 1)
if end < start:
return ""
return "".join(source_lines[start - 1 : end])
@staticmethod
def _safe_unparse(node: ast.AST) -> str:
"""Return ``ast.unparse(node)`` but tolerate exotic nodes by falling back to ``repr``."""
try:
return ast.unparse(node)
except Exception: # pragma: no cover - defensive guard
return repr(node)
def _strip_docstring(
self,
node: ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef,
source_lines: list[str],
unit_start: int,
unit_end: int,
) -> tuple[str, str | None]:
"""Strip ``node``'s docstring from ``source_lines[unit_start..unit_end]`` if safely possible."""
docstring = ast.get_docstring(node)
body = node.body
if not docstring or not body:
return self._slice_lines(source_lines, unit_start, unit_end), None
first = body[0]
if not (
isinstance(first, ast.Expr) and isinstance(first.value, ast.Constant) and isinstance(first.value.value, str)
):
return self._slice_lines(source_lines, unit_start, unit_end), None
# Skip stripping when the docstring shares a line with the def/class (would
# leave broken syntax) or extends past the caller's slice (e.g. class_header).
ds_start = first.lineno
ds_end = first.end_lineno or first.lineno
if ds_start <= node.lineno or ds_end > unit_end:
return self._slice_lines(source_lines, unit_start, unit_end), None
before = source_lines[unit_start - 1 : ds_start - 1]
after = source_lines[ds_end:unit_end]
return "".join(before + after), docstring
def _emit_class_units(self, cls: ast.ClassDef, source_lines: list[str], cursor: int, units: list[_CodeUnit]) -> int:
"""Emit class header and per-method units for ``cls``; return the next cursor (1-indexed)."""
class_start = cls.decorator_list[0].lineno if cls.decorator_list else cls.lineno
class_end = cls.end_lineno or cls.lineno
class_name = cls.name
class_decorators = [self._safe_unparse(d) for d in cls.decorator_list]
# Methods, async methods, and nested classes become their own units so a
# method is never split mid-statement.
split_children_idx = [
k
for k, child in enumerate(cls.body)
if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef))
]
# Bare class signature (decorators + ``class Foo(...):`` lines) used by
# ``preserve_class_definition`` to prefix later chunks of the same class.
class_signature: str | None = None
if cls.body:
body_start = cls.body[0].lineno
if body_start > class_start:
class_signature = self._slice_lines(source_lines, class_start, body_start - 1)
# Whole class fits in one unit when there are no inner split points.
if not split_children_idx:
unit_slice = self._slice_lines(source_lines, cursor, class_end)
stripped_docstring: str | None = None
if self.strip_docstrings:
unit_slice, stripped_docstring = self._strip_docstring(cls, source_lines, cursor, class_end)
units.append(
_CodeUnit(
source=unit_slice,
start_line=class_start,
end_line=class_end,
kind="class",
name=class_name,
class_name=class_name,
class_signature=class_signature,
decorators=class_decorators,
docstring=stripped_docstring,
)
)
return class_end + 1
# Class header: from outer cursor up to (but excluding) the first split child.
first_child = cls.body[split_children_idx[0]]
if (
isinstance(first_child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef))
and first_child.decorator_list
):
first_child_start = first_child.decorator_list[0].lineno
else:
first_child_start = first_child.lineno
header_end = first_child_start - 1
header_slice = self._slice_lines(source_lines, cursor, header_end)
header_docstring: str | None = None
if self.strip_docstrings:
header_slice, header_docstring = self._strip_docstring(cls, source_lines, cursor, header_end)
units.append(
_CodeUnit(
source=header_slice,
start_line=class_start,
end_line=header_end,
kind="class_header",
name=class_name,
class_name=class_name,
class_signature=class_signature,
decorators=class_decorators,
docstring=header_docstring,
)
)
inner_cursor = header_end + 1
for idx in split_children_idx:
child = cls.body[idx]
if not isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
continue # narrowed above; kept for the type checker
child_start = child.decorator_list[0].lineno if child.decorator_list else child.lineno
child_end = child.end_lineno or child.lineno
decorators = [self._safe_unparse(d) for d in child.decorator_list]
unit_slice = self._slice_lines(source_lines, inner_cursor, child_end)
stripped_docstring = None
if self.strip_docstrings:
unit_slice, stripped_docstring = self._strip_docstring(child, source_lines, inner_cursor, child_end)
kind = "method" if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)) else "nested_class"
units.append(
_CodeUnit(
source=unit_slice,
start_line=child_start,
end_line=child_end,
kind=kind,
name=child.name,
class_name=class_name,
class_signature=class_signature,
decorators=decorators,
docstring=stripped_docstring,
)
)
inner_cursor = child_end + 1
# Append trailing class-body lines (comments / blanks after the last method).
if inner_cursor <= class_end and units:
trailing = self._slice_lines(source_lines, inner_cursor, class_end)
units[-1].source += trailing
units[-1].end_line = class_end
return class_end + 1
def _extract_units(self, source: str) -> list[_CodeUnit]:
"""Parse ``source`` and produce the ordered list of syntactic split units."""
tree = ast.parse(source)
source_lines = source.splitlines(keepends=True)
total_lines = len(source_lines)
units: list[_CodeUnit] = []
cursor = 1
body = tree.body
node_idx = 0
node_count = len(body)
while node_idx < node_count:
node = body[node_idx]
# Module docstring (only valid as the very first statement).
if (
node_idx == 0
and isinstance(node, ast.Expr)
and isinstance(node.value, ast.Constant)
and isinstance(node.value.value, str)
):
end = node.end_lineno or node.lineno
units.append(
_CodeUnit(
source=self._slice_lines(source_lines, cursor, end),
start_line=node.lineno,
end_line=end,
kind="module_docstring",
)
)
cursor = end + 1
node_idx += 1
continue
# Group consecutive imports into one unit.
if isinstance(node, (ast.Import, ast.ImportFrom)):
import_end_idx = node_idx
while import_end_idx < node_count and isinstance(body[import_end_idx], (ast.Import, ast.ImportFrom)):
import_end_idx += 1
last = body[import_end_idx - 1]
end = last.end_lineno or last.lineno
units.append(
_CodeUnit(
source=self._slice_lines(source_lines, cursor, end),
start_line=node.lineno,
end_line=end,
kind="imports",
)
)
cursor = end + 1
node_idx = import_end_idx
continue
if isinstance(node, ast.ClassDef):
cursor = self._emit_class_units(node, source_lines, cursor, units)
node_idx += 1
continue
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
start = node.decorator_list[0].lineno if node.decorator_list else node.lineno
end = node.end_lineno or node.lineno
decorators = [self._safe_unparse(d) for d in node.decorator_list]
unit_slice = self._slice_lines(source_lines, cursor, end)
stripped_docstring: str | None = None
if self.strip_docstrings:
unit_slice, stripped_docstring = self._strip_docstring(node, source_lines, cursor, end)
units.append(
_CodeUnit(
source=unit_slice,
start_line=start,
end_line=end,
kind="function",
name=node.name,
decorators=decorators,
docstring=stripped_docstring,
)
)
cursor = end + 1
node_idx += 1
continue
# Catch-all for top-level statements (assignments, conditionals, etc.).
end = node.end_lineno or node.lineno
units.append(
_CodeUnit(
source=self._slice_lines(source_lines, cursor, end),
start_line=node.lineno,
end_line=end,
kind="statement",
)
)
cursor = end + 1
node_idx += 1
# Append trailing content (comments after the last node) so the split is loss-less.
if cursor <= total_lines and units:
trailing = self._slice_lines(source_lines, cursor, total_lines)
units[-1].source += trailing
units[-1].end_line = total_lines
elif cursor <= total_lines and not units:
units.append(
_CodeUnit(
source=self._slice_lines(source_lines, cursor, total_lines),
start_line=cursor,
end_line=total_lines,
kind="statement",
)
)
return units
def _merge_units(self, units: list[_CodeUnit]) -> list[list[_CodeUnit]]:
"""Greedily merge units toward ``max_effective_lines``; oversized units become solo chunks."""
chunks: list[list[_CodeUnit]] = []
current: list[_CodeUnit] = []
current_lines = 0
target = self.max_effective_lines
def flush() -> None:
nonlocal current, current_lines
if current:
chunks.append(current)
current = []
current_lines = 0
for unit in units:
if self._is_oversized(unit):
flush()
chunks.append([unit])
continue
unit_eff = self._effective_lines(unit.source)
if not current:
current = [unit]
current_lines = unit_eff
continue
# Keep merging while below the minimum or while adding moves us closer to the target.
new_total = current_lines + unit_eff
if current_lines < self.min_effective_lines or abs(new_total - target) < abs(current_lines - target):
current.append(unit)
current_lines = new_total
else:
flush()
current = [unit]
current_lines = unit_eff
flush()
return chunks
@staticmethod
def _ordered_unique(items: list[str]) -> list[str]:
"""Return the list of unique items in their first-seen order."""
return list(dict.fromkeys(items))
def _build_chunk_meta(self, chunk: list[_CodeUnit], parent_doc: Document) -> dict[str, Any]:
"""Construct the output meta dict for a chunk of merged units."""
meta: dict[str, Any] = {}
if parent_doc.meta:
meta.update({k: v for k, v in parent_doc.meta.items() if k not in {"split_id"}})
meta["source_id"] = parent_doc.id
# Units are emitted in source order, so chunk[0]/chunk[-1] give the extremes.
meta["start_line"] = chunk[0].start_line
meta["end_line"] = chunk[-1].end_line
meta["unit_kinds"] = [u.kind for u in chunk]
include_classes = self._ordered_unique([u.class_name for u in chunk if u.class_name])
if include_classes:
meta["include_classes"] = include_classes
decorators: list[str] = []
for u in chunk:
decorators.extend(u.decorators)
decorators = self._ordered_unique(decorators)
if decorators:
meta["decorators"] = decorators
if self.strip_docstrings:
docstrings = [u.docstring for u in chunk if u.docstring]
if docstrings:
meta["docstrings"] = docstrings
return meta
def _render_chunk_content(self, chunk: list[_CodeUnit]) -> str:
"""Render chunk content, optionally prefixing class signatures for orphan members."""
body = "".join(u.source for u in chunk)
if not self.preserve_class_definition:
return body
classes_with_header = {u.class_name for u in chunk if u.kind in {"class", "class_header"} and u.class_name}
prepended: list[str] = []
seen: set[str] = set()
for u in chunk:
if (
u.class_name
and u.class_name not in classes_with_header
and u.class_name not in seen
and u.class_signature
):
prepended.append(u.class_signature)
seen.add(u.class_name)
if not prepended:
return body
return "".join(prepended) + body
def _secondary_split(self, unit: _CodeUnit, parent_doc: Document) -> list[Document]:
"""Apply a line-based fallback split with overlap to a single oversized unit."""
qualified_name = unit.name or unit.kind
if unit.class_name and unit.name:
qualified_name = f"{unit.class_name}.{unit.name}"
logger.warning(
"Oversized {kind} '{func_name}' at lines {start}-{end} ({eff} effective lines) exceeds "
"{factor}x max_effective_lines={max_effective_lines}; falling back to line-based secondary split "
"with overlap={overlap}.",
kind=unit.kind,
func_name=qualified_name,
start=unit.start_line,
end=unit.end_line,
eff=self._effective_lines(unit.source),
factor=self.oversized_factor,
max_effective_lines=self.max_effective_lines,
overlap=self.secondary_split_overlap,
)
# DocumentSplitter measures in physical lines; this approximates effective lines.
split_length = (
self.secondary_split_length if self.secondary_split_length is not None else self.max_effective_lines
)
overlap = min(self.secondary_split_overlap, max(0, split_length - 1))
splitter = DocumentSplitter(split_by="line", split_length=split_length, split_overlap=overlap)
intermediate = splitter.run(documents=[Document(content=unit.source)])["documents"]
base_meta = self._build_chunk_meta([unit], parent_doc)
results: list[Document] = []
for idx, piece in enumerate(intermediate):
meta = dict(base_meta)
meta["secondary_split"] = True
meta["secondary_split_index"] = idx
meta["secondary_split_total"] = len(intermediate)
results.append(Document(content=piece.content or "", meta=meta))
return results
@component.output_types(documents=list[Document])
def run(self, documents: list[Document]) -> dict[str, list[Document]]:
"""
Split each Python source ``Document`` into syntax-aware chunks.
:param documents: Documents whose ``content`` is Python source code. Each
document's ``meta`` is propagated onto its chunks.
:returns: ``{"documents": [...]}`` where each chunk's meta additionally carries
``source_id``, ``split_id``, ``start_line``, ``end_line``, ``unit_kinds`` and
- where applicable - ``include_classes``, ``decorators``, ``docstrings``,
``secondary_split``.
:raises ValueError: If any document's content is ``None``.
:raises TypeError: If any document's content is not a string.
:raises SyntaxError: If a document's content is not valid Python.
"""
for doc in documents:
if doc.content is None:
raise ValueError(
f"PythonCodeSplitter only works with text documents but content for document ID {doc.id} is None."
)
if not isinstance(doc.content, str):
raise TypeError("PythonCodeSplitter only works with text documents (str content).")
final_docs: list[Document] = []
for doc in documents:
assert doc.content is not None # narrowed by the loop above
if not doc.content.strip():
logger.warning("Document ID {doc_id} has empty content. Skipping this document.", doc_id=doc.id)
continue
units = self._extract_units(doc.content)
if not units:
continue
chunks = self._merge_units(units)
split_id = 0
for chunk in chunks:
if len(chunk) == 1 and self._is_oversized(chunk[0]):
for piece in self._secondary_split(chunk[0], doc):
piece.meta["split_id"] = split_id
split_id += 1
final_docs.append(piece)
continue
content = self._render_chunk_content(chunk)
meta = self._build_chunk_meta(chunk, doc)
meta["split_id"] = split_id
split_id += 1
final_docs.append(Document(content=content, meta=meta))
return {"documents": final_docs}
@@ -0,0 +1,484 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# 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}
@@ -0,0 +1,237 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
import re
from pathlib import Path
from typing import Any, Literal
from haystack import logging
from haystack.lazy_imports import LazyImport
with LazyImport("Run 'pip install nltk>=3.9.1'") as nltk_imports:
import nltk
logger = logging.getLogger(__name__)
Language = Literal[
"ru", "sl", "es", "sv", "tr", "cs", "da", "nl", "en", "et", "fi", "fr", "de", "el", "it", "no", "pl", "pt", "ml"
]
ISO639_TO_NLTK = {
"ru": "russian",
"sl": "slovene",
"es": "spanish",
"sv": "swedish",
"tr": "turkish",
"cs": "czech",
"da": "danish",
"nl": "dutch",
"en": "english",
"et": "estonian",
"fi": "finnish",
"fr": "french",
"de": "german",
"el": "greek",
"it": "italian",
"no": "norwegian",
"pl": "polish",
"pt": "portuguese",
"ml": "malayalam",
}
QUOTE_SPANS_RE = re.compile(r'"[^"]*"|\'[^\']*\'')
if nltk_imports.is_successful():
def load_sentence_tokenizer(
language: Language, keep_white_spaces: bool = False
) -> nltk.tokenize.punkt.PunktSentenceTokenizer:
"""
Utility function to load the nltk sentence tokenizer.
:param language: The language for the tokenizer.
:param keep_white_spaces: If True, the tokenizer will keep white spaces between sentences.
:returns: nltk sentence tokenizer.
"""
try:
nltk.data.find("tokenizers/punkt_tab")
except LookupError:
try:
nltk.download("punkt_tab")
except FileExistsError as error:
logger.debug("NLTK punkt tokenizer seems to be already downloaded. Error message: {error}", error=error)
language_name = ISO639_TO_NLTK.get(language)
if language_name is not None:
sentence_tokenizer = nltk.data.load(f"tokenizers/punkt_tab/{language_name}.pickle")
else:
logger.warning(
"PreProcessor couldn't find the default sentence tokenizer model for {language}. "
" Using English instead. You may train your own model and use the 'tokenizer_model_folder' parameter.",
language=language,
)
sentence_tokenizer = nltk.data.load("tokenizers/punkt_tab/english.pickle")
if keep_white_spaces:
sentence_tokenizer._lang_vars = CustomPunktLanguageVars()
return sentence_tokenizer
class CustomPunktLanguageVars(nltk.tokenize.punkt.PunktLanguageVars):
# The following adjustment of PunktSentenceTokenizer is inspired by:
# https://stackoverflow.com/questions/33139531/preserve-empty-lines-with-nltks-punkt-tokenizer
# It is needed for preserving whitespace while splitting text into sentences.
_period_context_fmt = r"""
%(SentEndChars)s # a potential sentence ending
\s* # match potential whitespace [ \t\n\x0B\f\r]
(?=(?P<after_tok>
%(NonWord)s # either other punctuation
|
(?P<next_tok>\S+) # or some other token - original version: \s+(?P<next_tok>\S+)
))"""
def period_context_re(self) -> re.Pattern:
"""
Compiles and returns a regular expression to find contexts including possible sentence boundaries.
:returns: A compiled regular expression pattern.
"""
try:
return self._re_period_context # type: ignore
except: # noqa: E722
self._re_period_context = re.compile(
self._period_context_fmt
% {
"NonWord": self._re_non_word_chars,
# SentEndChars might be followed by closing brackets, so we match them here.
"SentEndChars": self._re_sent_end_chars + r"[\)\]}]*",
},
re.UNICODE | re.VERBOSE,
)
return self._re_period_context
class SentenceSplitter:
"""
SentenceSplitter splits a text into sentences using the nltk sentence tokenizer
"""
def __init__(
self,
language: Language = "en",
use_split_rules: bool = True,
extend_abbreviations: bool = True,
keep_white_spaces: bool = False,
) -> None:
"""
Initializes the SentenceSplitter with the specified language, split rules, and abbreviation handling.
:param language: The language for the tokenizer. Default is "en".
:param use_split_rules: If True, the additional split rules are used. If False, the rules are not used.
:param extend_abbreviations: If True, the abbreviations used by NLTK's PunktTokenizer are extended by a list
of curated abbreviations if available. If False, the default abbreviations are used.
Currently supported languages are: en, de.
:param keep_white_spaces: If True, the tokenizer will keep white spaces between sentences.
"""
nltk_imports.check()
self.language = language
# after checking nltk_imports, we are sure that load_sentence_tokenizer is defined
self.sentence_tokenizer = load_sentence_tokenizer(language, keep_white_spaces=keep_white_spaces)
self.use_split_rules = use_split_rules
if extend_abbreviations:
abbreviations = SentenceSplitter._read_abbreviations(language)
self.sentence_tokenizer._params.abbrev_types.update(abbreviations)
self.keep_white_spaces = keep_white_spaces
def split_sentences(self, text: str) -> list[dict[str, Any]]:
"""
Splits a text into sentences including references to original char positions for each split.
:param text: The text to split.
:returns: list of sentences with positions.
"""
sentence_spans = list(self.sentence_tokenizer.span_tokenize(text))
if self.use_split_rules:
sentence_spans = SentenceSplitter._apply_split_rules(text, sentence_spans)
return [{"sentence": text[start:end], "start": start, "end": end} for start, end in sentence_spans]
@staticmethod
def _apply_split_rules(text: str, sentence_spans: list[tuple[int, int]]) -> list[tuple[int, int]]:
"""
Applies additional split rules to the sentence spans.
:param text: The text to split.
:param sentence_spans: The list of sentence spans to split.
:returns: The list of sentence spans after applying the split rules.
"""
new_sentence_spans = []
quote_spans = [match.span() for match in QUOTE_SPANS_RE.finditer(text)]
while sentence_spans:
span = sentence_spans.pop(0)
next_span = sentence_spans[0] if len(sentence_spans) > 0 else None
while next_span and SentenceSplitter._needs_join(text, span, next_span, quote_spans):
sentence_spans.pop(0)
span = (span[0], next_span[1])
next_span = sentence_spans[0] if len(sentence_spans) > 0 else None
start, end = span
new_sentence_spans.append((start, end))
return new_sentence_spans
@staticmethod
def _needs_join(
text: str, span: tuple[int, int], next_span: tuple[int, int], quote_spans: list[tuple[int, int]]
) -> bool:
"""
Checks if the spans need to be joined as parts of one sentence.
This method determines whether two adjacent sentence spans should be joined back together as a single sentence.
It's used to prevent incorrect sentence splitting in specific cases like quotations, numbered lists,
and parenthetical expressions.
:param text: The text containing the spans.
:param span: Tuple of (start, end) positions for the current sentence span.
:param next_span: Tuple of (start, end) positions for the next sentence span.
:param quote_spans: All quoted spans within text.
:returns:
True if the spans needs to be joined.
"""
start, end = span
next_start, next_end = next_span
# sentence. sentence"\nsentence -> no split (end << quote_end)
# sentence.", sentence -> no split (end < quote_end)
# sentence?", sentence -> no split (end < quote_end)
if any(quote_start < end < quote_end for quote_start, quote_end in quote_spans):
# sentence boundary is inside a quote
return True
# sentence." sentence -> split (end == quote_end)
# sentence?" sentence -> no split (end == quote_end)
if any(quote_start < end == quote_end and text[quote_end - 2] == "?" for quote_start, quote_end in quote_spans):
# question is cited
return True
if re.search(r"(^|\n)\s*\d{1,2}\.$", text[start:end]) is not None:
# sentence ends with a numeration
return True
# next sentence starts with a bracket or we return False
return re.search(r"^\s*[\(\[]", text[next_start:next_end]) is not None
@staticmethod
def _read_abbreviations(lang: Language) -> list[str]:
"""
Reads the abbreviations for a given language from the abbreviations file.
:param lang: The language to read the abbreviations for.
:returns: List of abbreviations.
"""
abbreviations_file = Path(__file__).parent.parent.parent / f"data/abbreviations/{lang}.txt"
if not abbreviations_file.exists():
logger.warning("No abbreviations file found for {language}. Using default abbreviations.", language=lang)
return []
return abbreviations_file.read_text().split("\n")
@@ -0,0 +1,83 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
import re
import string
from typing import Any
from haystack import component
@component
class TextCleaner:
"""
Cleans text strings.
It can remove substrings matching a list of regular expressions, convert text to lowercase,
remove punctuation, and remove numbers.
Use it to clean up text data before evaluation.
### Usage example
```python
from haystack.components.preprocessors import TextCleaner
text_to_clean = "1Moonlight shimmered softly, 300 Wolves howled nearby, Night enveloped everything."
cleaner = TextCleaner(convert_to_lowercase=True, remove_punctuation=False, remove_numbers=True)
result = cleaner.run(texts=[text_to_clean])
```
"""
def __init__(
self,
remove_regexps: list[str] | None = None,
convert_to_lowercase: bool = False,
remove_punctuation: bool = False,
remove_numbers: bool = False,
) -> None:
"""
Initializes the TextCleaner component.
:param remove_regexps: A list of regex patterns to remove matching substrings from the text.
:param convert_to_lowercase: If `True`, converts all characters to lowercase.
:param remove_punctuation: If `True`, removes punctuation from the text.
:param remove_numbers: If `True`, removes numerical digits from the text.
"""
self._remove_regexps = remove_regexps
self._convert_to_lowercase = convert_to_lowercase
self._remove_punctuation = remove_punctuation
self._remove_numbers = remove_numbers
self._regex = None
if remove_regexps:
self._regex = re.compile("|".join(remove_regexps), flags=re.IGNORECASE)
to_remove = ""
if remove_punctuation:
to_remove = string.punctuation
if remove_numbers:
to_remove += string.digits
self._translator = str.maketrans("", "", to_remove) if to_remove else None
@component.output_types(texts=list[str])
def run(self, texts: list[str]) -> dict[str, Any]:
"""
Cleans up the given list of strings.
:param texts: List of strings to clean.
:returns: A dictionary with the following key:
- `texts`: the cleaned list of strings.
"""
if self._regex:
texts = [self._regex.sub("", text) for text in texts]
if self._convert_to_lowercase:
texts = [text.lower() for text in texts]
if self._translator:
texts = [text.translate(self._translator) for text in texts]
return {"texts": texts}