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
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:
@@ -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 (1–6) 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}
|
||||
Reference in New Issue
Block a user