diff --git a/haiku_rag_slim/haiku/rag/chunkers/docling_local.py b/haiku_rag_slim/haiku/rag/chunkers/docling_local.py index 40dd020c..7fa6d39b 100644 --- a/haiku_rag_slim/haiku/rag/chunkers/docling_local.py +++ b/haiku_rag_slim/haiku/rag/chunkers/docling_local.py @@ -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. diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 8fb114d4..47a0e599 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -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): diff --git a/tests/test_chunker.py b/tests/test_chunker.py index c65f617d..f9bf5dde 100644 --- a/tests/test_chunker.py +++ b/tests/test_chunker.py @@ -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