Cache the chunking tokenizer to avoid HF Hub 429

This commit is contained in:
Yiorgis Gozadinos 2026-05-13 13:32:24 +03:00
parent d512f9fbd1
commit 0f4c4af495
No known key found for this signature in database
3 changed files with 34 additions and 2 deletions

View file

@ -19,6 +19,7 @@
- **HTML text ingest path picks up converter options.** `convert_text(format="html"/"md")` previously used a bare `DoclingDocConverter()` with zero format options — the wix corpus ingest path. It now uses the same shared `_build_format_options()` helper as the file path.
- **Relative `<img>` paths resolve during URL ingest.** `HaikuRAG.convert()` and the converter `convert_file` / `convert_text` methods now thread a `source_uri` through to `HTMLBackendOptions.source_uri` / `MarkdownBackendOptions.source_uri`. URL ingest uses the originating URL; file ingest uses `file://`; raw text accepts an optional override. docling-serve accepts the kwarg as a no-op (its API has no equivalent option).
- **CLI tracebacks no longer dump per-frame locals.** The Typer app now passes `pretty_exceptions_show_locals=False`, so exceptions involving a `DoclingDocument` (or any large object) print readable rich tracebacks instead of pages of inline base64 image URIs. Set `_TYPER_STANDARD_TRACEBACK=1` for plain Python tracebacks.
- **Batch ingest no longer hits HF Hub's 429 rate limit.** The chunking tokenizer is now loaded once per process via `@functools.cache` instead of once per chunker instance.
### Documentation

View file

@ -1,3 +1,4 @@
from functools import cache
from typing import TYPE_CHECKING, cast
from haiku.rag.chunkers.base import DocumentChunker
@ -9,6 +10,16 @@ if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument
@cache
def _get_tokenizer(name: str):
# `AutoTokenizer.from_pretrained` triggers an HF Hub `model_info` request
# per call. Batch ingest builds one chunker per document, so without this
# cache HF rate-limits at 1000 requests / 5 minutes.
from transformers import AutoTokenizer
return AutoTokenizer.from_pretrained(name)
def _create_markdown_serializer_provider(use_markdown_tables: bool = True):
"""Create a markdown serializer provider with configurable table rendering.
@ -64,7 +75,6 @@ class DoclingLocalChunker(DocumentChunker):
from docling_core.transforms.chunker.tokenizer.huggingface import (
HuggingFaceTokenizer,
)
from transformers import AutoTokenizer
self.config = config
self.chunk_size = config.processing.chunk_size
@ -72,7 +82,7 @@ class DoclingLocalChunker(DocumentChunker):
self.tokenizer_name = config.processing.chunking_tokenizer
if self.chunker_type == "hybrid":
hf_tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name)
hf_tokenizer = _get_tokenizer(self.tokenizer_name)
tokenizer = HuggingFaceTokenizer(
tokenizer=hf_tokenizer, max_tokens=self.chunk_size
)

View file

@ -84,6 +84,27 @@ def test_get_chunker_invalid():
get_chunker(config)
def test_tokenizer_cached_across_chunker_instances():
"""Repeated DoclingLocalChunker instantiations share one loaded tokenizer.
Each `AutoTokenizer.from_pretrained` call triggers an `HfApi.model_info`
HTTP request to check for revision drift. Batch ingest creates one
chunker per document, which without caching hits HF Hub's 1000-per-5min
limit and crashes with HTTP 429.
"""
from haiku.rag.chunkers.docling_local import _get_tokenizer
_get_tokenizer.cache_clear()
DoclingLocalChunker()
DoclingLocalChunker()
DoclingLocalChunker()
info = _get_tokenizer.cache_info()
assert info.misses == 1
assert info.hits == 2
@pytest.mark.asyncio
async def test_local_chunker_hierarchical(qa_corpus: Dataset):
"""Test DoclingLocalChunker with hierarchical chunking."""