""" Semantic Chunker for Word Documents using Docling HybridChunker. This module provides document chunking capabilities for Word documents using: - HybridChunker from docling_core for intelligent content segmentation - Provider-aware tokenizer (tiktoken for Azure, HF AutoTokenizer for vLLM) - Custom serializers for tables and images - Contextualization for each chunk Architecture: 1. DoclingDocument → HybridChunker → Contextualized Chunks 2. Preserves tables in markdown format within chunks 3. Includes image captions and descriptions 4. Token budgeting via the configured embedding provider's tokenizer Pattern follows: /tmp/pdf_parser_service.py """ from typing import Any, Dict, List from uuid import UUID import structlog from docling_core.transforms.chunker.base import BaseChunk from docling_core.transforms.chunker.hierarchical_chunker import ( DocChunk, ChunkingDocSerializer, ChunkingSerializerProvider, ) from docling_core.transforms.chunker.hybrid_chunker import HybridChunker from docling_core.transforms.serializer.markdown import ( MarkdownTableSerializer, MarkdownPictureSerializer, ) from src.llm.factory import get_docling_tokenizer from docling_core.types.doc.document import ( PictureClassificationData, PictureDescriptionData, PictureItem, DoclingDocument, ) from docling_core.transforms.serializer.base import ( BaseDocSerializer, SerializationResult, ) from docling_core.transforms.serializer.common import create_ser_result logger = structlog.get_logger(__name__) class AnnotationPictureSerializer(MarkdownPictureSerializer): """Custom serializer for picture annotations that includes captions and descriptions.""" def serialize( self, *, item: PictureItem, doc_serializer: BaseDocSerializer, doc: DoclingDocument, **kwargs: Any, ) -> SerializationResult: """Serialize picture item with captions and annotations. Args: item: Picture item to serialize doc_serializer: Document serializer for post-processing doc: Parent Docling document **kwargs: Additional serialization arguments Returns: SerializationResult with formatted text and span source """ text_parts: list[str] = [] # Add caption if available caption = item.caption_text(doc=doc) if caption: text_parts.append(f"[CAPTION]: {caption}") # Add annotations for annotation in item.annotations: if isinstance(annotation, PictureClassificationData): predicted_class = ( annotation.predicted_classes[0].class_name if annotation.predicted_classes else None ) if predicted_class is not None: text_parts.append(f"[PICTURE TYPE]: {predicted_class}") elif isinstance(annotation, PictureDescriptionData): text_parts.append(f"[PICTURE DESCRIPTION] {annotation.text} [END OF PICTURE DESCRIPTION]") text_res = "\n".join(text_parts) text_res = doc_serializer.post_process(text=text_res) return create_ser_result(text=text_res, span_source=item) class ImgTableAnnotationSerializerProvider(ChunkingSerializerProvider): """Custom serializer provider for chunking with image and table support.""" def get_serializer(self, doc: DoclingDocument): """Get serializer with custom picture and table serializers. Args: doc: Docling document to serialize Returns: ChunkingDocSerializer configured with custom serializers """ return ChunkingDocSerializer( doc=doc, picture_serializer=AnnotationPictureSerializer(), table_serializer=MarkdownTableSerializer(), ) class SemanticChunker: """ Semantic chunker for Word documents using HybridChunker. Uses: - Provider-aware tokenizer (tiktoken for Azure, HF AutoTokenizer for vLLM) - HybridChunker for intelligent document segmentation - Custom serializers for tables and images - Contextualization for enhanced chunk quality Example: chunker = SemanticChunker() chunks = chunker.chunk_document(docling_doc, file_id, "report.docx") """ def __init__(self, max_tokens: int = 8000): """ Initialize semantic chunker with tokenizer and hybrid chunker. Args: max_tokens: Maximum tokens per chunk (default: 8000, matching typical 8K-context embedding models). """ self.logger = structlog.get_logger() self.tokenizer = get_docling_tokenizer(max_tokens) # Initialize HybridChunker with custom serializers self.chunker = HybridChunker( tokenizer=self.tokenizer, serializer_provider=ImgTableAnnotationSerializerProvider(), ) self.logger.info( "SemanticChunker initialized", max_tokens=max_tokens, tokenizer=type(self.tokenizer).__name__, ) def chunk_document( self, doc: DoclingDocument, file_id: UUID, original_filename: str, ) -> List[Dict[str, Any]]: """ Chunk a Docling document into semantic units with contextualization. Args: doc: DoclingDocument from Docling conversion file_id: UUID of the file record original_filename: Original filename (e.g., "report.docx") Returns: List of chunk records with text, metadata, and contextualization """ self.logger.info( "Starting document chunking", filename=original_filename, file_id=str(file_id) ) # Apply HybridChunker chunk_iter = self.chunker.chunk(dl_doc=doc) chunks = list(chunk_iter) self.logger.info(f"Created {len(chunks)} raw chunks") # Convert chunks to records with contextualization chunk_records = [] for chunk in chunks: record = self._build_chunk_record( chunk=chunk, file_id=file_id, original_filename=original_filename ) if record: chunk_records.append(record) self.logger.info( f"Chunking complete", total_chunks=len(chunk_records), filename=original_filename ) return chunk_records def _build_chunk_record( self, chunk: BaseChunk, file_id: UUID, original_filename: str, ) -> Dict[str, Any] | None: """ Build a chunk record from a Docling chunk with metadata. Args: chunk: BaseChunk from HybridChunker file_id: UUID of the file record original_filename: Original filename Returns: Dictionary with chunk data or None if chunk is empty """ # Get contextualized text text = self.chunker.contextualize(chunk=chunk) # Skip empty chunks if not text.strip(): return None # Validate chunk as DocChunk to access metadata doc_chunk = DocChunk.model_validate(chunk) # Extract page number from doc items (for multi-sheet/multi-page docs) page_numbers = [] for item in doc_chunk.meta.doc_items: if hasattr(item, 'prov') and item.prov: for prov_item in item.prov: if hasattr(prov_item, 'page_no') and prov_item.page_no is not None: page_numbers.append(prov_item.page_no) # Build metadata metadata = { 'page_number': min(page_numbers) if page_numbers else None, 'source_filename': original_filename, 'unique_filename': original_filename, 'contains_table': self._contains_markdown_table(text), } # Build chunk record record = { 'file_id': str(file_id), 'text': text, 'metadata': metadata, 'element_type': 'Text', } return record def _contains_markdown_table(self, text: str) -> bool: """ Check if text contains a markdown table. Args: text: Text content to check Returns: True if markdown table detected, False otherwise """ # Simple heuristic: look for table separator lines (e.g., "| --- | --- |") lines = text.split('\n') for line in lines: # Check for markdown table separator pattern if '|' in line and '---' in line: return True return False