feat: add page_number to metadata in DocumentSplitter

haystack/components/preprocessors/document_splitter.py

@@ -1,5 +1,5 @@
 from copy import deepcopy
-from typing import List, Literal
+from typing import Dict, List, Literal, Tuple
 
 from more_itertools import windowed
 
@@ -53,7 +53,7 @@ def run(self, documents: List[Document]):
 
  :returns: A dictionary with the following key:
  - `documents`: List of documents with the split texts. A metadata field "source_id" is added to each
- document to keep track of the original document that was split. Other metadata are copied from the original
+ document to keep track of the original document that was split. Another metadata field "page_number" is added to each number to keep track of the page it belonged to in the original document. Other metadata are copied from the original
  document.
 
  :raises TypeError: if the input is not a list of Documents.
@@ -70,10 +70,12 @@ def run(self, documents: List[Document]):
  f"DocumentSplitter only works with text documents but document.content for document ID {doc.id} is None."
  )
  units = self._split_into_units(doc.content, self.split_by)
- text_splits = self._concatenate_units(units, self.split_length, self.split_overlap)
+ text_splits, splits_pages = self._concatenate_units(units, self.split_length, self.split_overlap)
  metadata = deepcopy(doc.meta)
  metadata["source_id"] = doc.id
- split_docs += [Document(content=txt, meta=metadata) for txt in text_splits]
+ split_docs += self._create_docs_from_splits(
+ text_splits=text_splits, splits_pages=splits_pages, meta=metadata
+ )
  return {"documents": split_docs}
 
  def _split_into_units(self, text: str, split_by: Literal["word", "sentence", "passage", "page"]) -> List[str]:
@@ -95,15 +97,40 @@ def _split_into_units(self, text: str, split_by: Literal["word", "sentence", "pa
  units[i] += split_at
  return units
 
- def _concatenate_units(self, elements: List[str], split_length: int, split_overlap: int) -> List[str]:
+ def _concatenate_units(
+ self, elements: List[str], split_length: int, split_overlap: int
+ ) -> Tuple[List[str], List[int]]:
  """
- Concatenates the elements into parts of split_length units.
+ Concatenates the elements into parts of split_length units keeping track of the original page number that each element belongs.
  """
  text_splits = []
+ splits_pages = []
+ 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)
  if len(txt) > 0:
  text_splits.append(txt)
- return text_splits
+ splits_pages.append(cur_page)
+ processed_units = current_units[: split_length - split_overlap]
+ 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
+
+ @staticmethod
+ def _create_docs_from_splits(text_splits: List[str], splits_pages: List[int], meta: Dict) -> List[Document]:
+ """
+ Creates Document objects from text splits enriching them with page number and the metadata of the original document.
+ """
+ documents: List[Document] = []
+
+ for i, txt in enumerate(text_splits):
+ meta = deepcopy(meta)
+ doc = Document(content=txt, meta=meta)
+ doc.meta["page_number"] = splits_pages[i]
+ documents.append(doc)
+ return documents

releasenotes/notes/add-page-number-to-document-splitter-162e9dc7443575f0.yaml

@@ -0,0 +1,7 @@
+---
+highlights: >
+ Add the "page_number" field to the metadata of all output documents.
+
+enhancements:
+ - |
+ Now the DocumentSplitter adds the "page_number" field to the metadata of all output documents to keep track of the page of the original document it belongs to.

test/components/preprocessors/test_document_splitter.py

@@ -141,3 +141,98 @@ def test_copy_metadata(self):
  for doc, split_doc in zip(documents, result["documents"]):
  assert doc.meta.items() <= split_doc.meta.items()
  assert split_doc.content == "Text."
+
+ def test_add_page_number_to_metadata_with_no_overlap_word_split(self):
+ splitter = DocumentSplitter(split_by="word", split_length=2)
+ doc1 = Document(content="This is some text.\f This text is on another page.")
+ doc2 = Document(content="This content has two.\f\f page brakes.")
+ result = splitter.run(documents=[doc1, doc2])
+
+ expected_pages = [1, 1, 2, 2, 2, 1, 1, 3]
+ for doc, p in zip(result["documents"], expected_pages):
+ assert doc.meta["page_number"] == p
+
+ def test_add_page_number_to_metadata_with_no_overlap_sentence_split(self):
+ splitter = DocumentSplitter(split_by="sentence", split_length=1)
+ doc1 = Document(content="This is some text.\f This text is on another page.")
+ doc2 = Document(content="This content has two.\f\f page brakes.")
+ result = splitter.run(documents=[doc1, doc2])
+
+ expected_pages = [1, 1, 1, 1]
+ for doc, p in zip(result["documents"], expected_pages):
+ assert doc.meta["page_number"] == p
+
+ def test_add_page_number_to_metadata_with_no_overlap_passage_split(self):
+ splitter = DocumentSplitter(split_by="passage", split_length=1)
+ doc1 = Document(
+ content="This is a text with some words.\f There is a second sentence.\n\nAnd there is a third sentence.\n\nAnd more passages.\n\n\f And another passage."
+ )
+ result = splitter.run(documents=[doc1])
+
+ expected_pages = [1, 2, 2, 2]
+ for doc, p in zip(result["documents"], expected_pages):
+ assert doc.meta["page_number"] == p
+
+ def test_add_page_number_to_metadata_with_no_overlap_page_split(self):
+ splitter = DocumentSplitter(split_by="page", split_length=1)
+ doc1 = Document(
+ content="This is a text with some words. There is a second sentence.\f And there is a third sentence.\f And another passage."
+ )
+ result = splitter.run(documents=[doc1])
+ expected_pages = [1, 2, 3]
+ for doc, p in zip(result["documents"], expected_pages):
+ assert doc.meta["page_number"] == p
+
+ splitter = DocumentSplitter(split_by="page", split_length=2)
+ doc1 = Document(
+ content="This is a text with some words. There is a second sentence.\f And there is a third sentence.\f And another passage."
+ )
+ result = splitter.run(documents=[doc1])
+ expected_pages = [1, 3]
+
+ for doc, p in zip(result["documents"], expected_pages):
+ assert doc.meta["page_number"] == p
+
+ def test_add_page_number_to_metadata_with_overlap_word_split(self):
+ splitter = DocumentSplitter(split_by="word", split_length=3, split_overlap=1)
+ doc1 = Document(content="This is some text. And\f this text is on another page.")
+ doc2 = Document(content="This content has two.\f\f page brakes.")
+ result = splitter.run(documents=[doc1, doc2])
+
+ expected_pages = [1, 1, 1, 2, 2, 1, 1, 3]
+ for doc, p in zip(result["documents"], expected_pages):
+ print(doc.content, doc.meta, p)
+ assert doc.meta["page_number"] == p
+
+ def test_add_page_number_to_metadata_with_overlap_sentence_split(self):
+ splitter = DocumentSplitter(split_by="sentence", split_length=2, split_overlap=1)
+ doc1 = Document(content="This is some text. And this is more text.\f This text is on another page. End.")
+ doc2 = Document(content="This content has two.\f\f page brakes. More text.")
+ result = splitter.run(documents=[doc1, doc2])
+
+ expected_pages = [1, 1, 1, 2, 1, 1]
+ for doc, p in zip(result["documents"], expected_pages):
+ print(doc.content, doc.meta, p)
+ assert doc.meta["page_number"] == p
+
+ def test_add_page_number_to_metadata_with_overlap_passage_split(self):
+ splitter = DocumentSplitter(split_by="passage", split_length=2, split_overlap=1)
+ doc1 = Document(
+ content="This is a text with some words.\f There is a second sentence.\n\nAnd there is a third sentence.\n\nAnd more passages.\n\n\f And another passage."
+ )
+ result = splitter.run(documents=[doc1])
+
+ expected_pages = [1, 2, 2]
+ for doc, p in zip(result["documents"], expected_pages):
+ assert doc.meta["page_number"] == p
+
+ def test_add_page_number_to_metadata_with_overlap_page_split(self):
+ splitter = DocumentSplitter(split_by="page", split_length=2, split_overlap=1)
+ doc1 = Document(
+ content="This is a text with some words. There is a second sentence.\f And there is a third sentence.\f And another passage."
+ )
+ result = splitter.run(documents=[doc1])
+ expected_pages = [1, 2, 3]
+
+ for doc, p in zip(result["documents"], expected_pages):
+ assert doc.meta["page_number"] == p