feat: Add tiktoken tokenizers support to HybridChunker (#240)

* feat: Add tiktoken tokenizers support to HybridChunker

Signed-off-by: ruizguille <[email protected]>

* separate OpenAI tokenizer

Signed-off-by: Panos Vagenas <[email protected]>

---------

Signed-off-by: ruizguille <[email protected]>
Signed-off-by: Panos Vagenas <[email protected]>
Co-authored-by: Panos Vagenas <[email protected]>

Author: ruizguille

docling_core/transforms/chunker/hybrid_chunker.py

@@ -8,27 +8,21 @@
 from functools import cached_property
 from typing import Any, Iterable, Iterator, Optional, Union
 
-from pydantic import (
-    BaseModel,
-    ConfigDict,
-    PositiveInt,
-    TypeAdapter,
-    computed_field,
-    model_validator,
-)
-from typing_extensions import Self
+from pydantic import BaseModel, ConfigDict, Field, computed_field, model_validator
 
 from docling_core.transforms.chunker.hierarchical_chunker import (
     ChunkingSerializerProvider,
 )
+from docling_core.transforms.chunker.tokenizer.base import BaseTokenizer
 
 try:
     import semchunk
-    from transformers import AutoTokenizer, PreTrainedTokenizerBase
 except ImportError:
     raise RuntimeError(
-        "Module requires 'chunking' extra; to install, run: "
-        "`pip install 'docling-core[chunking]'`"
+        "Extra required by module: 'chunking' by default (or 'chunking-openai' if "
+        "specifically using OpenAI tokenization); to install, run: "
+        "`pip install 'docling-core[chunking]'` or "
+        "`pip install 'docling-core[chunking-openai]'`"
     )
 
 from docling_core.experimental.serializer.base import (
@@ -45,6 +39,16 @@
 from docling_core.types import DoclingDocument
 
 
+def _get_default_tokenizer():
+    from docling_core.transforms.chunker.tokenizer.huggingface import (
+        HuggingFaceTokenizer,
+    )
+
+    return HuggingFaceTokenizer.from_pretrained(
+        model_name="sentence-transformers/all-MiniLM-L6-v2"
+    )
+
+
 class HybridChunker(BaseChunker):
     r"""Chunker doing tokenization-aware refinements on top of document layout chunking.
 
@@ -58,26 +62,40 @@ class HybridChunker(BaseChunker):
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    tokenizer: Union[PreTrainedTokenizerBase, str] = (
-        "sentence-transformers/all-MiniLM-L6-v2"
-    )
-    max_tokens: int = None  # type: ignore[assignment]
+    tokenizer: BaseTokenizer = Field(default_factory=_get_default_tokenizer)
     merge_peers: bool = True
 
     serializer_provider: BaseSerializerProvider = ChunkingSerializerProvider()
 
-    @model_validator(mode="after")
-    def _patch_tokenizer_and_max_tokens(self) -> Self:
-        self._tokenizer = (
-            self.tokenizer
-            if isinstance(self.tokenizer, PreTrainedTokenizerBase)
-            else AutoTokenizer.from_pretrained(self.tokenizer)
-        )
-        if self.max_tokens is None:
-            self.max_tokens = TypeAdapter(PositiveInt).validate_python(
-                self._tokenizer.model_max_length
-            )
-        return self
+    @model_validator(mode="before")
+    @classmethod
+    def _patch(cls, data: Any) -> Any:
+        if isinstance(data, dict) and (tokenizer := data.get("tokenizer")):
+            max_tokens = data.get("max_tokens")
+            if isinstance(tokenizer, BaseTokenizer):
+                pass
+            else:
+                from docling_core.transforms.chunker.tokenizer.huggingface import (
+                    HuggingFaceTokenizer,
+                )
+
+                if isinstance(tokenizer, str):
+                    data["tokenizer"] = HuggingFaceTokenizer.from_pretrained(
+                        model_name=tokenizer,
+                        max_tokens=max_tokens,
+                    )
+                else:
+                    # migrate previous HF-based tokenizers
+                    kwargs = {"tokenizer": tokenizer}
+                    if max_tokens is not None:
+                        kwargs["max_tokens"] = max_tokens
+                    data["tokenizer"] = HuggingFaceTokenizer(**kwargs)
+        return data
+
+    @property
+    def max_tokens(self) -> int:
+        """Get maximum number of tokens allowed."""
+        return self.tokenizer.get_max_tokens()
 
     @computed_field  # type: ignore[misc]
     @cached_property
@@ -92,7 +110,7 @@ def _count_text_tokens(self, text: Optional[Union[str, list[str]]]):
             for t in text:
                 total += self._count_text_tokens(t)
             return total
-        return len(self._tokenizer.tokenize(text))
+        return self.tokenizer.count_tokens(text=text)
 
     class _ChunkLengthInfo(BaseModel):
         total_len: int
@@ -101,7 +119,7 @@ class _ChunkLengthInfo(BaseModel):
 
     def _count_chunk_tokens(self, doc_chunk: DocChunk):
         ser_txt = self.contextualize(chunk=doc_chunk)
-        return len(self._tokenizer.tokenize(text=ser_txt))
+        return self.tokenizer.count_tokens(text=ser_txt)
 
     def _doc_chunk_length(self, doc_chunk: DocChunk):
         text_length = self._count_text_tokens(doc_chunk.text)
@@ -198,7 +216,7 @@ def _split_using_plain_text(
             # captions:
             available_length = self.max_tokens - lengths.other_len
             sem_chunker = semchunk.chunkerify(
-                self._tokenizer, chunk_size=available_length
+                self.tokenizer.get_tokenizer(), chunk_size=available_length
             )
             if available_length <= 0:
                 warnings.warn(

docling_core/transforms/chunker/tokenizer/__init__.py

@@ -0,0 +1 @@
+"""Define the tokenizer types."""

docling_core/transforms/chunker/tokenizer/base.py

@@ -0,0 +1,25 @@
+"""Define base classes for tokenization."""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class BaseTokenizer(BaseModel, ABC):
+    """Base tokenizer class."""
+
+    @abstractmethod
+    def count_tokens(self, text: str) -> int:
+        """Get number of tokens for given text."""
+        ...
+
+    @abstractmethod
+    def get_max_tokens(self) -> int:
+        """Get maximum number of tokens allowed."""
+        ...
+
+    @abstractmethod
+    def get_tokenizer(self) -> Any:
+        """Get underlying tokenizer object."""
+        ...

docling_core/transforms/chunker/tokenizer/huggingface.py

@@ -0,0 +1,70 @@
+"""HuggingFace tokenization."""
+
+import sys
+from os import PathLike
+from typing import Optional, Union
+
+from pydantic import ConfigDict, PositiveInt, TypeAdapter, model_validator
+from typing_extensions import Self
+
+from docling_core.transforms.chunker.tokenizer.base import BaseTokenizer
+
+try:
+    from transformers import AutoTokenizer, PreTrainedTokenizerBase
+except ImportError:
+    raise RuntimeError(
+        "Module requires 'chunking' extra; to install, run: "
+        "`pip install 'docling-core[chunking]'`"
+    )
+
+
+class HuggingFaceTokenizer(BaseTokenizer):
+    """HuggingFace tokenizer."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    tokenizer: PreTrainedTokenizerBase
+    max_tokens: int = None  # type: ignore[assignment]
+
+    @model_validator(mode="after")
+    def _patch(self) -> Self:
+        if hasattr(self.tokenizer, "model_max_length"):
+            model_max_tokens: PositiveInt = TypeAdapter(PositiveInt).validate_python(
+                self.tokenizer.model_max_length
+            )
+            user_max_tokens = self.max_tokens or sys.maxsize
+            self.max_tokens = min(model_max_tokens, user_max_tokens)
+        elif self.max_tokens is None:
+            raise ValueError(
+                "max_tokens must be defined as model does not define model_max_length"
+            )
+        return self
+
+    def count_tokens(self, text: str):
+        """Get number of tokens for given text."""
+        return len(self.tokenizer.tokenize(text=text))
+
+    def get_max_tokens(self):
+        """Get maximum number of tokens allowed."""
+        return self.max_tokens
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_name: Union[str, PathLike],
+        max_tokens: Optional[int] = None,
+        **kwargs,
+    ) -> Self:
+        """Create tokenizer from model name."""
+        my_kwargs = {
+            "tokenizer": AutoTokenizer.from_pretrained(
+                pretrained_model_name_or_path=model_name, **kwargs
+            ),
+        }
+        if max_tokens is not None:
+            my_kwargs["max_tokens"] = max_tokens
+        return cls(**my_kwargs)
+
+    def get_tokenizer(self):
+        """Get underlying tokenizer object."""
+        return self.tokenizer

docling_core/transforms/chunker/tokenizer/openai.py

@@ -0,0 +1,34 @@
+"""OpenAI tokenization."""
+
+from pydantic import ConfigDict
+
+from docling_core.transforms.chunker.hybrid_chunker import BaseTokenizer
+
+try:
+    import tiktoken
+except ImportError:
+    raise RuntimeError(
+        "Module requires 'chunking-openai' extra; to install, run: "
+        "`pip install 'docling-core[chunking-openai]'`"
+    )
+
+
+class OpenAITokenizer(BaseTokenizer):
+    """OpenAI tokenizer."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    tokenizer: tiktoken.Encoding
+    max_tokens: int
+
+    def count_tokens(self, text: str):
+        """Get number of tokens for given text."""
+        return len(self.tokenizer.encode(text=text))
+
+    def get_max_tokens(self):
+        """Get maximum number of tokens allowed."""
+        return self.max_tokens
+
+    def get_tokenizer(self):
+        """Get underlying tokenizer object."""
+        return self.tokenizer