Chunker abstraction

This commit is contained in:
Yiorgis Gozadinos 2025-11-14 15:20:20 +02:00
parent 541552215e
commit de5bb117fe
No known key found for this signature in database
7 changed files with 143 additions and 23 deletions

View file

@ -0,0 +1,31 @@
"""Document chunker abstraction for haiku.rag."""
from haiku.rag.chunkers.base import DocumentChunker
from haiku.rag.config import AppConfig, Config
__all__ = ["DocumentChunker", "get_chunker"]
def get_chunker(config: AppConfig = Config) -> DocumentChunker:
"""Get a document chunker instance based on configuration.
Args:
config: Configuration to use. Defaults to global Config.
Returns:
DocumentChunker instance configured according to the config.
Raises:
ValueError: If the chunker provider is not recognized.
"""
if config.processing.chunker == "docling-local":
from haiku.rag.chunkers.docling_local import DoclingLocalChunker
return DoclingLocalChunker(config)
if config.processing.chunker == "docling-serve":
from haiku.rag.chunkers.docling_serve import DoclingServeChunker
return DoclingServeChunker(config)
raise ValueError(f"Unsupported chunker: {config.processing.chunker}")

View file

@ -0,0 +1,28 @@
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument
class DocumentChunker(ABC):
"""Abstract base class for document chunkers.
Document chunkers split DoclingDocuments into smaller text chunks suitable
for embedding and retrieval, respecting document structure and semantic boundaries.
"""
@abstractmethod
async def chunk(self, document: "DoclingDocument") -> list[str]:
"""Split a document into chunks.
Args:
document: The DoclingDocument to chunk.
Returns:
List of text chunks with semantic boundaries preserved.
Raises:
ValueError: If chunking fails.
"""
pass

View file

@ -1,38 +1,38 @@
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from haiku.rag.config import Config from haiku.rag.chunkers.base import DocumentChunker
from haiku.rag.config import AppConfig, Config
if TYPE_CHECKING: if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument from docling_core.types.doc.document import DoclingDocument
class Chunker: class DoclingLocalChunker(DocumentChunker):
"""A class that chunks text into smaller pieces for embedding and retrieval. """Local document chunker using docling's HybridChunker.
Uses docling's structure-aware chunking to create semantically meaningful chunks Uses docling's structure-aware chunking to create semantically meaningful chunks
that respect document boundaries. that respect document boundaries. Chunking is performed locally using the
HuggingFace tokenizer specified in configuration.
Args: Args:
chunk_size: The maximum size of a chunk in tokens. config: Application configuration.
tokenizer_name: HuggingFace model name for tokenization.
""" """
def __init__( def __init__(self, config: AppConfig = Config):
self,
chunk_size: int = Config.processing.chunk_size,
tokenizer_name: str = Config.processing.chunking_tokenizer,
):
from docling_core.transforms.chunker.hybrid_chunker import HybridChunker from docling_core.transforms.chunker.hybrid_chunker import HybridChunker
from docling_core.transforms.chunker.tokenizer.huggingface import ( from docling_core.transforms.chunker.tokenizer.huggingface import (
HuggingFaceTokenizer, HuggingFaceTokenizer,
) )
from transformers import AutoTokenizer from transformers import AutoTokenizer
self.chunk_size = chunk_size self.config = config
self.tokenizer_name = tokenizer_name self.chunk_size = config.processing.chunk_size
self.tokenizer_name = config.processing.chunking_tokenizer
hf_tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) hf_tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name)
tokenizer = HuggingFaceTokenizer(tokenizer=hf_tokenizer, max_tokens=chunk_size) tokenizer = HuggingFaceTokenizer(
tokenizer=hf_tokenizer, max_tokens=self.chunk_size
)
self.chunker = HybridChunker(tokenizer=tokenizer) self.chunker = HybridChunker(tokenizer=tokenizer)
@ -51,6 +51,3 @@ class Chunker:
# Chunk using docling's hybrid chunker # Chunk using docling's hybrid chunker
chunks = list(self.chunker.chunk(document)) chunks = list(self.chunker.chunk(document))
return [self.chunker.contextualize(chunk) for chunk in chunks] return [self.chunker.contextualize(chunk) for chunk in chunks]
chunker = Chunker()

View file

@ -0,0 +1,31 @@
from typing import TYPE_CHECKING
from haiku.rag.chunkers.base import DocumentChunker
from haiku.rag.config import AppConfig
if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument
class DoclingServeChunker(DocumentChunker):
"""Remote document chunker using docling-serve API.
Placeholder - will be implemented in a future commit.
"""
def __init__(self, config: AppConfig):
raise NotImplementedError("DoclingServeChunker not yet implemented")
async def chunk(self, document: "DoclingDocument") -> list[str]:
"""Split the document into chunks via docling-serve.
Args:
document: The DoclingDocument to be split into chunks.
Returns:
A list of text chunks with semantic boundaries.
Raises:
NotImplementedError: This chunker is not yet implemented.
"""
raise NotImplementedError("DoclingServeChunker not yet implemented")

View file

@ -55,6 +55,7 @@ class ProcessingConfig(BaseModel):
context_chunk_radius: int = 0 context_chunk_radius: int = 0
markdown_preprocessor: str = "" markdown_preprocessor: str = ""
converter: str = "docling-local" converter: str = "docling-local"
chunker: str = "docling-local"
chunking_tokenizer: str = "Qwen/Qwen3-Embedding-0.6B" chunking_tokenizer: str = "Qwen/Qwen3-Embedding-0.6B"

View file

@ -150,9 +150,11 @@ class ChunkRepository:
) -> list[Chunk]: ) -> list[Chunk]:
"""Create chunks and embeddings for a document from DoclingDocument.""" """Create chunks and embeddings for a document from DoclingDocument."""
# Lazy imports to avoid loading docling during module import # Lazy imports to avoid loading docling during module import
from haiku.rag.chunker import chunker from haiku.rag.chunkers import get_chunker
from haiku.rag.converters import get_converter from haiku.rag.converters import get_converter
chunker = get_chunker(self.store._config)
# Optionally preprocess markdown before chunking # Optionally preprocess markdown before chunking
processed_document = document processed_document = document
preprocessor_path = self.store._config.processing.markdown_preprocessor preprocessor_path = self.store._config.processing.markdown_preprocessor

View file

@ -2,14 +2,16 @@ import pytest
from datasets import Dataset from datasets import Dataset
from transformers import AutoTokenizer from transformers import AutoTokenizer
from haiku.rag.chunker import Chunker from haiku.rag.chunkers import get_chunker
from haiku.rag.config import Config from haiku.rag.chunkers.docling_local import DoclingLocalChunker
from haiku.rag.config import AppConfig, Config
from haiku.rag.converters import get_converter from haiku.rag.converters import get_converter
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_chunker(qa_corpus: Dataset): async def test_local_chunker(qa_corpus: Dataset):
chunker = Chunker() """Test DoclingLocalChunker with real document."""
chunker = DoclingLocalChunker()
doc_text = qa_corpus[0]["document_extracted"] doc_text = qa_corpus[0]["document_extracted"]
# Convert text to DoclingDocument # Convert text to DoclingDocument
@ -43,3 +45,31 @@ async def test_chunker(qa_corpus: Dataset):
# Due to structure-aware chunking, we might have some variation in token count # Due to structure-aware chunking, we might have some variation in token count
# but it should be reasonable # but it should be reasonable
assert abs(total_tokens - original_tokens) <= original_tokens * 0.1 assert abs(total_tokens - original_tokens) <= original_tokens * 0.1
@pytest.mark.asyncio
async def test_local_chunker_custom_config():
"""Test DoclingLocalChunker with custom configuration."""
config = AppConfig()
config.processing.chunk_size = 128
config.processing.chunking_tokenizer = "Qwen/Qwen3-Embedding-0.6B"
chunker = DoclingLocalChunker(config)
assert chunker.chunk_size == 128
assert chunker.tokenizer_name == "Qwen/Qwen3-Embedding-0.6B"
def test_get_chunker_docling_local():
"""Test factory returns DoclingLocalChunker for docling-local."""
config = AppConfig()
config.processing.chunker = "docling-local"
chunker = get_chunker(config)
assert isinstance(chunker, DoclingLocalChunker)
def test_get_chunker_invalid():
"""Test factory raises error for invalid chunker."""
config = AppConfig()
config.processing.chunker = "invalid-chunker"
with pytest.raises(ValueError, match="Unsupported chunker"):
get_chunker(config)