Coverage for src/qdrant_loader/core/chunking/docling/tokenizer.py: 66%

1"""Builds the docling tokenizer that bounds chunk size. 

2 

3:class:`TokenizerFactory` is the single place that turns our engine-neutral 

4:class:`~.config.TokenizerConfig` into a docling ``BaseTokenizer`` — the chunking 

5analogue of the conversion option builder. It is a small pure class holding the 

6injected config and imports docling lazily, so importing this package costs nothing 

7until a chunk is actually sized. 

8 

9Aligning the chunk budget to the *embedding* model's tokenizer is the whole point of 

10this layer; a mismatched tokenizer is the character-budget mistake it exists to fix. 

11""" 

12 

13from __future__ import annotations 

14 

15from typing import TYPE_CHECKING 

16 

17from qdrant_loader.utils.logging import LoggingConfig 

18 

19from .config import TokenizerConfig, TokenizerKind 

20from .exceptions import TokenizerUnavailableError 

21 

22if TYPE_CHECKING: 

23 from docling_core.transforms.chunker.tokenizer.base import BaseTokenizer 

24 

25logger = LoggingConfig.get_logger(__name__) 

26 

27 

28class TokenizerFactory: 

29 """Config -> docling ``BaseTokenizer``. The only place tokenizers are built.""" 

30 

31 def __init__(self, config: TokenizerConfig) -> None: 

32 self._config = config 

33 

34 def build(self) -> BaseTokenizer: 

35 """Construct the tokenizer for the configured kind. 

36 

37 Raises :class:`TokenizerUnavailableError` (loud, not a silent char-count 

38 fallback) when the backing library or model cannot be loaded. When the config 

39 is an approximate counting proxy (``embedding.tokenizer="none"``), warns once 

40 here so the choice is observable rather than silent. 

41 """ 

42 if self._config.approximate: 

43 logger.warning( 

44 "Chunk token counts are approximate: the embedding model exposes no " 

45 "tokenizer, so this proxy is used only for counting. Set a chunking " 

46 "tokenizer matching your embedding model for exact token budgeting.", 

47 proxy_tokenizer=self._config.model, 

48 ) 

49 match self._config.kind: 

50 case TokenizerKind.OPENAI: 

51 return self._build_openai() 

52 case TokenizerKind.HUGGINGFACE: 

53 return self._build_huggingface() 

54 raise TokenizerUnavailableError( 

55 str(self._config.kind), "unknown tokenizer kind" 

56 ) 

57 

58 def _build_openai(self) -> BaseTokenizer: 

59 try: 

60 import tiktoken 

61 from docling_core.transforms.chunker.tokenizer.openai import ( 

62 OpenAITokenizer, 

63 ) 

64 

65 encoding = tiktoken.get_encoding(self._config.model) 

66 except Exception as error: # missing dep / unknown encoding / no cached BPE 

67 raise TokenizerUnavailableError("openai", str(error)) from error 

68 return OpenAITokenizer(tokenizer=encoding, max_tokens=self._config.max_tokens) 

69 

70 def _build_huggingface(self) -> BaseTokenizer: 

71 try: 

72 from docling_core.transforms.chunker.tokenizer.huggingface import ( 

73 HuggingFaceTokenizer, 

74 ) 

75 

76 return HuggingFaceTokenizer.from_pretrained( 

77 model_name=self._config.model, 

78 max_tokens=self._config.max_tokens, 

79 ) 

80 except Exception as error: # missing dep / model not present / offline hub 

81 raise TokenizerUnavailableError("huggingface", str(error)) from error