From 88b80ec724f491a1e7e3e04bc922b178e1105e58 Mon Sep 17 00:00:00 2001 From: shivsin25 Date: Sat, 26 Sep 2026 20:38:58 +0530 Subject: [PATCH] feat: make the tiktoken encoding configurable in RecursiveDocumentSplitter Add a keyword-only tokenizer_encoding parameter, defaulting to the previously hardcoded "o200k_base", so that token-based chunks can be sized with the encoding of the model they are meant for. Same name, default and behavior as in DocumentSplitter. Closes #12973 Co-Authored-By: Claude Opus 5.5 (1M context) --- .../preprocessors/recursivesplitter.mdx | 1 + .../preprocessors/recursive_splitter.py | 9 +++- ...r-tokenizer-encoding-5b7b308985596dcc.yaml | 7 +++ .../preprocessors/test_recursive_splitter.py | 49 +++++++++++++++++++ 4 files changed, 64 insertions(+), 2 deletions(-) create mode 100644 releasenotes/notes/recursive-splitter-tokenizer-encoding-5b7b308985596dcc.yaml diff --git a/docs-website/docs/pipeline-components/preprocessors/recursivesplitter.mdx b/docs-website/docs/pipeline-components/preprocessors/recursivesplitter.mdx index d133cf5530d..d56fdf6cfa2 100644 --- a/docs-website/docs/pipeline-components/preprocessors/recursivesplitter.mdx +++ b/docs-website/docs/pipeline-components/preprocessors/recursivesplitter.mdx @@ -30,6 +30,7 @@ The `RecursiveDocumentSplitter` expects a list of documents as input and returns - `split_unit`: The unit of the `split_length` parameter. Can be either `"word"`, `"char"`, or `"token"`. - `separators`: An optional list of separator strings to use for splitting the text. If you don’t provide any separators, the default ones are `["\n\n", "sentence", "\n", " "]`. The string separators will be treated as regular expressions. If the separator is `"sentence"`, the text will be split into sentences using a custom sentence tokenizer based on NLTK. See [SentenceSplitter](https://github.com/deepset-ai/haystack/blob/main/haystack/components/preprocessors/sentence_tokenizer.py#L116) code for more information. - `sentence_splitter_params`: Optional parameters to pass to the [SentenceSplitter](https://github.com/deepset-ai/haystack/blob/main/haystack/components/preprocessors/sentence_tokenizer.py#L116). +- `tokenizer_encoding`: The tiktoken encoding used when `split_unit="token"`. 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, so that chunks fit the model's token limit. When `split_unit="token"`, strings such as `<|endoftext|>` in the source document are tokenized as ordinary text. They are preserved in the chunks and count toward the token limit. diff --git a/haystack/components/preprocessors/recursive_splitter.py b/haystack/components/preprocessors/recursive_splitter.py index 0917d7e7999..5b777fb14fe 100644 --- a/haystack/components/preprocessors/recursive_splitter.py +++ b/haystack/components/preprocessors/recursive_splitter.py @@ -64,6 +64,7 @@ def __init__( split_unit: Literal["word", "char", "token"] = "word", separators: list[str] | None = None, sentence_splitter_params: dict[str, Any] | None = None, + tokenizer_encoding: str = "o200k_base", ) -> None: """ Initializes a RecursiveDocumentSplitter. @@ -73,7 +74,7 @@ def __init__( :param split_overlap: The number of overlapping units (words, characters, or tokens, per `split_unit`) 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). + If "token" is selected, the text will be split into tokens using tiktoken (see `tokenizer_encoding`). Special-token strings in document content are encoded as ordinary text. :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 @@ -82,6 +83,8 @@ def __init__( 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. + :param tokenizer_encoding: The tiktoken encoding to use when `split_unit="token"`. Defaults to `"o200k_base"` + (current OpenAI models). Only used when `split_unit="token"`. :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. @@ -95,6 +98,7 @@ def __init__( self.sentence_splitter_params = ( {"keep_white_spaces": True} if sentence_splitter_params is None else sentence_splitter_params ) + self.tokenizer_encoding = tokenizer_encoding self.tiktoken_tokenizer: "tiktoken.Encoding" | None = None self._is_warmed_up = False @@ -108,7 +112,7 @@ def warm_up(self) -> None: 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.tiktoken_tokenizer = tiktoken.get_encoding(self.tokenizer_encoding) self._is_warmed_up = True def to_dict(self) -> dict[str, Any]: @@ -125,6 +129,7 @@ def to_dict(self) -> dict[str, Any]: split_unit=self.split_units, separators=self.separators, sentence_splitter_params=self.sentence_splitter_params, + tokenizer_encoding=self.tokenizer_encoding, ) def _check_params(self) -> None: diff --git a/releasenotes/notes/recursive-splitter-tokenizer-encoding-5b7b308985596dcc.yaml b/releasenotes/notes/recursive-splitter-tokenizer-encoding-5b7b308985596dcc.yaml new file mode 100644 index 00000000000..e682eaf8f15 --- /dev/null +++ b/releasenotes/notes/recursive-splitter-tokenizer-encoding-5b7b308985596dcc.yaml @@ -0,0 +1,7 @@ +--- +enhancements: + - | + Add a ``tokenizer_encoding`` parameter to ``RecursiveDocumentSplitter`` to choose the tiktoken encoding used + when ``split_unit="token"``, as in ``DocumentSplitter``. It defaults to ``"o200k_base"``, the encoding used + until now. Set it to match your model, for example ``"cl100k_base"`` for OpenAI's ``text-embedding-3`` models, + so that chunks fit the model's token limit. diff --git a/test/components/preprocessors/test_recursive_splitter.py b/test/components/preprocessors/test_recursive_splitter.py index 6790082f9e0..e4049a0742e 100644 --- a/test/components/preprocessors/test_recursive_splitter.py +++ b/test/components/preprocessors/test_recursive_splitter.py @@ -955,6 +955,26 @@ def test_special_token_strings_are_split_as_literal_text(split_overlap): assert reconstructed == text +@pytest.mark.integration +@pytest.mark.parametrize("encoding", ["o200k_base", "cl100k_base"]) +def test_run_split_by_token_count_with_tokenizer_encoding(encoding): + splitter = RecursiveDocumentSplitter( + split_length=10, separators=["."], split_unit="token", tokenizer_encoding=encoding + ) + text = ( + "Haystack pipelines connect retrievers, rankers and generators into search systems. " + "Embedding models count tokens with their own encoding. Chunks must fit that limit." + ) + chunks = splitter.run([Document(content=text)])["documents"] + + assert splitter.tiktoken_tokenizer is not None + assert splitter.tiktoken_tokenizer.name == encoding + assert len(chunks) > 1 + for chunk in chunks: + assert chunk.content is not None + assert len(splitter.tiktoken_tokenizer.encode_ordinary(chunk.content)) <= 10 + + @pytest.mark.integration def test_run_split_by_token_count(): splitter = RecursiveDocumentSplitter(split_length=5, separators=["."], split_unit="token") @@ -1234,6 +1254,35 @@ def test_warm_up_is_idempotent_token(monkeypatch): assert splitter.tiktoken_tokenizer is sentinel +def test_warm_up_uses_tokenizer_encoding(monkeypatch): + import haystack.components.preprocessors.recursive_splitter as mod + + get_encoding = Mock() + monkeypatch.setattr(mod.tiktoken, "get_encoding", get_encoding) + + splitter = RecursiveDocumentSplitter(split_unit="token", split_length=10, tokenizer_encoding="cl100k_base") + splitter.warm_up() + + get_encoding.assert_called_once_with("cl100k_base") + + +def test_serialization_keeps_tokenizer_encoding(): + splitter = RecursiveDocumentSplitter(split_length=10, split_unit="token", tokenizer_encoding="cl100k_base") + pipeline = Pipeline() + pipeline.add_component("chunker", splitter) + + restored = Pipeline.loads(pipeline.dumps()).get_component("chunker") + assert isinstance(restored, RecursiveDocumentSplitter) + assert restored.tokenizer_encoding == "cl100k_base" + + # pipelines serialized before tokenizer_encoding existed get the default + data = pipeline.to_dict() + del data["components"]["chunker"]["init_parameters"]["tokenizer_encoding"] + restored = Pipeline.from_dict(data).get_component("chunker") + assert isinstance(restored, RecursiveDocumentSplitter) + assert restored.tokenizer_encoding == "o200k_base" + + def test_fallback_overlap_char_unit(): """split_overlap must be applied even when no separator matches (char unit).""" splitter = RecursiveDocumentSplitter(split_length=5, split_overlap=2, separators=["\n\n"], split_unit="char")