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:
parent
de5bb117fe
commit
75cbdb47a8
3 changed files with 135 additions and 11 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue