Chunker abstraction
This commit is contained in:
parent
541552215e
commit
de5bb117fe
7 changed files with 143 additions and 23 deletions
31
haiku_rag_slim/haiku/rag/chunkers/__init__.py
Normal file
31
haiku_rag_slim/haiku/rag/chunkers/__init__.py
Normal 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}")
|
||||||
28
haiku_rag_slim/haiku/rag/chunkers/base.py
Normal file
28
haiku_rag_slim/haiku/rag/chunkers/base.py
Normal 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
|
||||||
|
|
@ -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()
|
|
||||||
31
haiku_rag_slim/haiku/rag/chunkers/docling_serve.py
Normal file
31
haiku_rag_slim/haiku/rag/chunkers/docling_serve.py
Normal 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")
|
||||||
|
|
@ -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"
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue