From 756b00ca970d149ff3b20e30c5850a8014f41ffb Mon Sep 17 00:00:00 2001 From: shivsin25 Date: Thu, 24 Sep 2026 16:35:02 +0530 Subject: [PATCH] feat: support split_by="token" in HierarchicalDocumentSplitter Add "token" to split_by and a keyword-only tokenizer_encoding parameter, both passed to the inner DocumentSplitter instances, so block sizes can be counted in LLM tokens. Add warm_up() to load the tiktoken encoding or the sentence tokenizer when the pipeline starts instead of on the first run(). Closes #12917 Co-Authored-By: Claude Opus 5.5 (1M context) --- .../hierarchicaldocumentsplitter.mdx | 5 +- .../hierarchical_document_splitter.py | 29 +++++- ...splitter-token-split-62a86907ebb795f5.yaml | 8 ++ .../test_hierarchical_doc_splitter.py | 88 ++++++++++++++++++- 4 files changed, 122 insertions(+), 8 deletions(-) create mode 100644 releasenotes/notes/hierarchical-splitter-token-split-62a86907ebb795f5.yaml diff --git a/docs-website/docs/pipeline-components/preprocessors/hierarchicaldocumentsplitter.mdx b/docs-website/docs/pipeline-components/preprocessors/hierarchicaldocumentsplitter.mdx index bad1eb53488..59f28be8996 100644 --- a/docs-website/docs/pipeline-components/preprocessors/hierarchicaldocumentsplitter.mdx +++ b/docs-website/docs/pipeline-components/preprocessors/hierarchicaldocumentsplitter.mdx @@ -40,8 +40,9 @@ This descending order of sizes builds the hierarchy. These additional parameters can be set when the component is initialized: -- `split_by` can be `"word"` (default), `"sentence"`, `"passage"`, `"page"`. -- `split_overlap` is an integer indicating the number of overlapping words, sentences, or passages between chunks, 0 being the default. +- `split_by` can be `"word"` (default), `"sentence"`, `"passage"`, `"page"`, or `"token"`. With `"token"`, block sizes are counted in LLM tokens using tiktoken, so every block fits a known context budget. This requires `pip install tiktoken`. +- `split_overlap` is an integer indicating the number of overlapping words, sentences, passages, or tokens between chunks, 0 being the default. +- `tokenizer_encoding` is the tiktoken encoding used when `split_by="token"`. It defaults to `"o200k_base"`, used by current OpenAI models. Set it to match your model, for example `"cl100k_base"` for OpenAI's `text-embedding-3` embedding models. ## Usage diff --git a/haystack/components/preprocessors/hierarchical_document_splitter.py b/haystack/components/preprocessors/hierarchical_document_splitter.py index e008a05ec43..8a635e6e437 100644 --- a/haystack/components/preprocessors/hierarchical_document_splitter.py +++ b/haystack/components/preprocessors/hierarchical_document_splitter.py @@ -40,14 +40,20 @@ def __init__( self, block_sizes: set[int], split_overlap: int = 0, - split_by: Literal["word", "sentence", "page", "passage"] = "word", + split_by: Literal["word", "sentence", "page", "passage", "token"] = "word", + *, + tokenizer_encoding: str = "o200k_base", ) -> 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. + :param split_by: The unit for splitting your documents: `word`, `sentence`, `page`, `passage`, or `token`. + `token` splits by token count using tiktoken (requires `pip install tiktoken`), so that block sizes + match an LLM's context budget. + :param tokenizer_encoding: The tiktoken encoding to use when `split_by="token"`. Defaults to `"o200k_base"` + (current OpenAI models). Only used when `split_by="token"`. :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`. """ @@ -69,8 +75,16 @@ def __init__( self.splitters: dict[int, DocumentSplitter] = {} self.split_overlap = split_overlap self.split_by = split_by + self.tokenizer_encoding = tokenizer_encoding self._build_block_sizes() + def warm_up(self) -> None: + """ + Warm up the splitters by loading the sentence tokenizer or the tiktoken encoding they need. + """ + for splitter in self.splitters.values(): + splitter.warm_up() + @component.output_types(documents=list[Document]) def run(self, documents: list[Document]) -> dict[str, list[Document]]: """ @@ -87,7 +101,10 @@ def run(self, documents: list[Document]) -> dict[str, list[Document]]: 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 + split_length=block_size, + split_overlap=self.split_overlap, + split_by=self.split_by, + tokenizer_encoding=self.tokenizer_encoding, ) @staticmethod @@ -141,7 +158,11 @@ def to_dict(self) -> dict[str, Any]: 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 + self, + block_sizes=self.block_sizes, + split_overlap=self.split_overlap, + split_by=self.split_by, + tokenizer_encoding=self.tokenizer_encoding, ) @classmethod diff --git a/releasenotes/notes/hierarchical-splitter-token-split-62a86907ebb795f5.yaml b/releasenotes/notes/hierarchical-splitter-token-split-62a86907ebb795f5.yaml new file mode 100644 index 00000000000..5a8d4748d94 --- /dev/null +++ b/releasenotes/notes/hierarchical-splitter-token-split-62a86907ebb795f5.yaml @@ -0,0 +1,8 @@ +--- +enhancements: + - | + Add ``split_by="token"`` to ``HierarchicalDocumentSplitter``, so block sizes are counted in LLM tokens using + tiktoken, as in ``DocumentSplitter``. The encoding defaults to ``"o200k_base"`` and can be changed with the new + ``tokenizer_encoding`` parameter. ``HierarchicalDocumentSplitter`` also gets a ``warm_up()`` method that loads the + sentence tokenizer or tiktoken encoding up front instead of on the first ``run()``. Requires ``pip install tiktoken`` + for ``split_by="token"``. diff --git a/test/components/preprocessors/test_hierarchical_doc_splitter.py b/test/components/preprocessors/test_hierarchical_doc_splitter.py index 1c1f0bf7963..40b7bedc7d2 100644 --- a/test/components/preprocessors/test_hierarchical_doc_splitter.py +++ b/test/components/preprocessors/test_hierarchical_doc_splitter.py @@ -2,10 +2,13 @@ # # SPDX-License-Identifier: Apache-2.0 +from unittest.mock import Mock + import pytest from haystack import Document, Pipeline from haystack.components.preprocessors import HierarchicalDocumentSplitter +from haystack.components.preprocessors import document_splitter as document_splitter_module from haystack.components.writers import DocumentWriter @@ -43,7 +46,12 @@ def test_to_dict(self): expected = builder.to_dict() assert expected == { "type": "haystack.components.preprocessors.hierarchical_document_splitter.HierarchicalDocumentSplitter", - "init_parameters": {"block_sizes": [300, 200, 100], "split_overlap": 25, "split_by": "word"}, + "init_parameters": { + "block_sizes": [300, 200, 100], + "split_overlap": 25, + "split_by": "word", + "tokenizer_encoding": "o200k_base", + }, } def test_from_dict(self): @@ -56,6 +64,8 @@ def test_from_dict(self): assert builder.block_sizes == [10, 5, 2] assert builder.split_overlap == 0 assert builder.split_by == "word" + # dicts serialized before tokenizer_encoding existed get the default + assert builder.tokenizer_encoding == "o200k_base" def test_run(self): builder = HierarchicalDocumentSplitter(block_sizes={10, 5, 2}, split_overlap=0, split_by="word") @@ -144,7 +154,12 @@ def test_to_dict_in_pipeline(self, in_memory_doc_store): assert expected["components"]["hierarchical_doc_splitter"] == { "type": "haystack.components.preprocessors.hierarchical_document_splitter.HierarchicalDocumentSplitter", - "init_parameters": {"block_sizes": [10, 5, 2], "split_overlap": 0, "split_by": "word"}, + "init_parameters": { + "block_sizes": [10, 5, 2], + "split_overlap": 0, + "split_by": "word", + "tokenizer_encoding": "o200k_base", + }, } def test_from_dict_in_pipeline(self): @@ -276,3 +291,72 @@ def test_hierarchical_splitter_multiple_block_sizes(self): for child in children: assert child.meta["__parent_id"] == doc.id assert child.meta["__level"] == doc.meta["__level"] + 1 + + def test_init_with_token_split(self): + splitter = HierarchicalDocumentSplitter( + block_sizes={100, 20}, split_by="token", tokenizer_encoding="cl100k_base" + ) + assert splitter.split_by == "token" + assert splitter.tokenizer_encoding == "cl100k_base" + for block_size, inner_splitter in splitter.splitters.items(): + assert inner_splitter.split_by == "token" + assert inner_splitter.split_length == block_size + assert inner_splitter.tokenizer_encoding == "cl100k_base" + + def test_serialization_round_trip_with_token_split(self): + splitter = HierarchicalDocumentSplitter( + block_sizes={100, 20}, split_overlap=5, split_by="token", tokenizer_encoding="cl100k_base" + ) + restored = HierarchicalDocumentSplitter.from_dict(splitter.to_dict()) + assert restored.block_sizes == [100, 20] + assert restored.split_overlap == 5 + assert restored.split_by == "token" + assert restored.tokenizer_encoding == "cl100k_base" + assert restored.splitters[20].tokenizer_encoding == "cl100k_base" + + def test_warm_up_loads_the_encoding_for_every_block_size_once(self, monkeypatch): + encoding = Mock() + get_encoding = Mock(return_value=encoding) + monkeypatch.setattr(document_splitter_module.tiktoken, "get_encoding", get_encoding) + + splitter = HierarchicalDocumentSplitter( + block_sizes={100, 20}, split_by="token", tokenizer_encoding="cl100k_base" + ) + splitter.warm_up() + splitter.warm_up() + + # one call per block size, the second warm_up() is a no-op + assert get_encoding.call_count == 2 + get_encoding.assert_called_with("cl100k_base") + assert all(inner_splitter._tiktoken_tokenizer is encoding for inner_splitter in splitter.splitters.values()) + + +@pytest.mark.integration +class TestHierarchicalDocumentSplitterByTokenIntegration: + """Integration tests for split_by="token" that load a real tiktoken encoding.""" + + @pytest.mark.parametrize("encoding", ["o200k_base", "cl100k_base"]) + def test_run_split_by_token(self, encoding): + text = ( + "Haystack pipelines connect retrievers, rankers and generators. Hierarchical chunks let an " + "AutoMergingRetriever return a parent block when enough of its children match.\f" + "Token-based block sizes keep every chunk inside the context budget of the model." + ) + splitter = HierarchicalDocumentSplitter(block_sizes={20, 8}, split_by="token", tokenizer_encoding=encoding) + splitter.warm_up() + + documents = splitter.run([Document(content=text)])["documents"] + + tokenizer = splitter.splitters[20]._tiktoken_tokenizer + assert tokenizer is not None + docs_by_id = {doc.id: doc for doc in documents} + assert {doc.meta["__level"] for doc in documents} == {0, 1, 2} + for doc in documents[1:]: + assert doc.content is not None + assert len(tokenizer.encode_ordinary(doc.content)) <= doc.meta["__block_size"] + assert doc.id in docs_by_id[doc.meta["__parent_id"]].meta["__children_ids"] + # without overlap, the children of every block add up to exactly the block's text + for doc in documents: + if doc.meta["__children_ids"]: + children_text = "".join(str(docs_by_id[child_id].content) for child_id in doc.meta["__children_ids"]) + assert children_text == doc.content