Add chunker configuration options

- Add `chunker_type` config field to choose between "hybrid" (default) and "hierarchical" chunking strategies
  - Add `chunking_merge_peers` config field for HybridChunker (default: true)
  - Add `chunking_use_markdown_tables` config field to control table serialization format (default: false, matching docling's default)
This commit is contained in:
Yiorgis Gozadinos 2025-11-14 16:51:04 +02:00
parent de5bb117fe
commit 75cbdb47a8
No known key found for this signature in database
3 changed files with 135 additions and 11 deletions

View file

@ -7,18 +7,57 @@ if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument
class DoclingLocalChunker(DocumentChunker):
"""Local document chunker using docling's HybridChunker.
def _create_markdown_serializer_provider(use_markdown_tables: bool = True):
"""Create a markdown serializer provider with configurable table rendering.
Uses docling's structure-aware chunking to create semantically meaningful chunks
that respect document boundaries. Chunking is performed locally using the
HuggingFace tokenizer specified in configuration.
This function creates a custom serializer provider that extends ChunkingSerializerProvider
from docling-core. It's implemented as a factory function to avoid importing
docling-core at module level.
Args:
use_markdown_tables: If True, use MarkdownTableSerializer for rendering tables as
markdown. If False, use default TripletTableSerializer for narrative format.
"""
from docling_core.transforms.chunker.hierarchical_chunker import (
ChunkingDocSerializer,
ChunkingSerializerProvider,
)
from docling_core.transforms.serializer.markdown import MarkdownTableSerializer
class MDTableSerializerProvider(ChunkingSerializerProvider):
"""Serializer provider for markdown table output."""
def __init__(self, use_markdown_tables: bool = True):
self.use_markdown_tables = use_markdown_tables
def get_serializer(self, doc):
if self.use_markdown_tables:
return ChunkingDocSerializer(
doc=doc,
table_serializer=MarkdownTableSerializer(),
)
else:
# Use default ChunkingDocSerializer (TripletTableSerializer)
return ChunkingDocSerializer(doc=doc)
return MDTableSerializerProvider(use_markdown_tables=use_markdown_tables)
class DoclingLocalChunker(DocumentChunker):
"""Local document chunker using docling's chunkers.
Supports both hybrid (structure-aware) and hierarchical chunking strategies.
Chunking is performed locally using the HuggingFace tokenizer specified in
configuration.
Args:
config: Application configuration.
"""
def __init__(self, config: AppConfig = Config):
from docling_core.transforms.chunker.hierarchical_chunker import (
HierarchicalChunker,
)
from docling_core.transforms.chunker.hybrid_chunker import HybridChunker
from docling_core.transforms.chunker.tokenizer.huggingface import (
HuggingFaceTokenizer,
@ -27,14 +66,32 @@ class DoclingLocalChunker(DocumentChunker):
self.config = config
self.chunk_size = config.processing.chunk_size
self.chunker_type = config.processing.chunker_type
self.tokenizer_name = config.processing.chunking_tokenizer
hf_tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name)
tokenizer = HuggingFaceTokenizer(
tokenizer=hf_tokenizer, max_tokens=self.chunk_size
)
self.chunker = HybridChunker(tokenizer=tokenizer)
if self.chunker_type == "hybrid":
hf_tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name)
tokenizer = HuggingFaceTokenizer(
tokenizer=hf_tokenizer, max_tokens=self.chunk_size
)
serializer_provider = _create_markdown_serializer_provider(
use_markdown_tables=config.processing.chunking_use_markdown_tables
)
self.chunker = HybridChunker(
tokenizer=tokenizer,
merge_peers=config.processing.chunking_merge_peers,
serializer_provider=serializer_provider,
)
elif self.chunker_type == "hierarchical":
serializer_provider = _create_markdown_serializer_provider(
use_markdown_tables=config.processing.chunking_use_markdown_tables
)
self.chunker = HierarchicalChunker(serializer_provider=serializer_provider)
else:
raise ValueError(
f"Unsupported chunker_type: {self.chunker_type}. "
"Must be 'hybrid' or 'hierarchical'."
)
async def chunk(self, document: "DoclingDocument") -> list[str]:
"""Split the document into chunks using docling's structure-aware chunking.

View file

@ -56,7 +56,10 @@ class ProcessingConfig(BaseModel):
markdown_preprocessor: str = ""
converter: str = "docling-local"
chunker: str = "docling-local"
chunker_type: str = "hybrid"
chunking_tokenizer: str = "Qwen/Qwen3-Embedding-0.6B"
chunking_merge_peers: bool = True
chunking_use_markdown_tables: bool = False
class OllamaConfig(BaseModel):

View file

@ -73,3 +73,67 @@ def test_get_chunker_invalid():
config.processing.chunker = "invalid-chunker"
with pytest.raises(ValueError, match="Unsupported chunker"):
get_chunker(config)
@pytest.mark.asyncio
async def test_local_chunker_hierarchical(qa_corpus: Dataset):
"""Test DoclingLocalChunker with hierarchical chunking."""
config = AppConfig()
config.processing.chunker_type = "hierarchical"
chunker = DoclingLocalChunker(config)
doc_text = qa_corpus[0]["document_extracted"]
converter = get_converter(Config)
doc = converter.convert_text(doc_text, name="test.md")
chunks = await chunker.chunk(doc)
# Hierarchical chunker should produce chunks
assert len(chunks) > 0
# Each chunk should be non-empty
for chunk in chunks:
assert len(chunk.strip()) > 0
def test_local_chunker_invalid_type():
"""Test DoclingLocalChunker raises error for invalid chunker_type."""
config = AppConfig()
config.processing.chunker_type = "invalid-type"
with pytest.raises(ValueError, match="Unsupported chunker_type"):
DoclingLocalChunker(config)
@pytest.mark.asyncio
async def test_local_chunker_markdown_tables():
"""Test DoclingLocalChunker with markdown table serialization."""
markdown_with_table = """# Test Document
| Column 1 | Column 2 |
|----------|----------|
| Value A | Value B |
| Value D | Value E |
"""
converter = get_converter(Config)
doc = converter.convert_text(markdown_with_table, name="test.md")
# Test with markdown tables enabled
config_md = AppConfig()
config_md.processing.chunking_use_markdown_tables = True
chunker_md = DoclingLocalChunker(config_md)
chunks_md = await chunker_md.chunk(doc)
# Should contain markdown table format
assert any("|" in chunk for chunk in chunks_md)
assert any("Column 1" in chunk for chunk in chunks_md)
# Test with markdown tables disabled (narrative format)
config_narrative = AppConfig()
config_narrative.processing.chunking_use_markdown_tables = False
chunker_narrative = DoclingLocalChunker(config_narrative)
chunks_narrative = await chunker_narrative.chunk(doc)
# Should contain narrative format (no pipe characters in table)
table_content = [chunk for chunk in chunks_narrative if "Value" in chunk][0]
# Narrative format uses commas, not pipes for table structure
assert "," in table_content and "|" not in table_content