Spaces:
Sleeping
Sleeping
| """ | |
| pdf_chunker.py — Section- and table-aware PDF parser for the enterprise manual. | |
| Uses PyMuPDF (fitz) for high-fidelity extraction with layout awareness. | |
| """ | |
| import re | |
| import logging | |
| from dataclasses import dataclass, field | |
| from typing import List, Optional | |
| from pathlib import Path | |
| import fitz # PyMuPDF | |
| logger = logging.getLogger(__name__) | |
| class DocumentChunk: | |
| """A chunk of text from the PDF with metadata.""" | |
| chunk_id: int | |
| text: str | |
| page_number: int | |
| section_title: str = "" | |
| chunk_type: str = "text" # "text", "table", "heading", "list" | |
| start_char: int = 0 | |
| end_char: int = 0 | |
| word_count: int = 0 | |
| def __post_init__(self): | |
| self.word_count = len(self.text.split()) | |
| class ParsedDocument: | |
| """Complete parsed document.""" | |
| chunks: List[DocumentChunk] | |
| total_pages: int | |
| total_chunks: int | |
| filename: str | |
| sections: List[str] = field(default_factory=list) | |
| class PDFChunker: | |
| """ | |
| Section- and table-aware PDF parser. | |
| Features: | |
| - Heading/section detection | |
| - Table extraction | |
| - Smart chunking with overlap | |
| - Metadata preservation | |
| """ | |
| # Patterns for section detection | |
| _HEADING_PATTERNS = [ | |
| re.compile(r'^(\d+\.?\d*\.?\d*)\s+(.+)$'), # "1.2.3 Title" | |
| re.compile(r'^(Chapter|Section|Part|Appendix)\s+\w+', re.IGNORECASE), | |
| re.compile(r'^[A-Z][A-Z\s]{5,}$'), # ALL CAPS headings | |
| re.compile(r'^(Introduction|Overview|Summary|Conclusion|References|Glossary|FAQ)', re.IGNORECASE), | |
| ] | |
| # Table indicators | |
| _TABLE_INDICATORS = [ | |
| re.compile(r'\|.*\|.*\|'), # Pipe-delimited tables | |
| re.compile(r'\t.*\t'), # Tab-delimited | |
| re.compile(r'^\s*[-+]+\s*$'), # Separator lines | |
| ] | |
| def __init__( | |
| self, | |
| chunk_size: int = 512, | |
| chunk_overlap: int = 64, | |
| min_chunk_size: int = 50, | |
| ): | |
| self.chunk_size = chunk_size | |
| self.chunk_overlap = chunk_overlap | |
| self.min_chunk_size = min_chunk_size | |
| def parse(self, pdf_path: str) -> ParsedDocument: | |
| """ | |
| Parse a PDF file into structured chunks. | |
| Args: | |
| pdf_path: Path to the PDF file. | |
| Returns: | |
| ParsedDocument with all chunks and metadata. | |
| """ | |
| pdf_path = Path(pdf_path) | |
| if not pdf_path.exists(): | |
| raise FileNotFoundError(f"PDF not found: {pdf_path}") | |
| logger.info(f"Parsing PDF: {pdf_path}") | |
| doc = fitz.open(str(pdf_path)) | |
| all_blocks = [] | |
| sections_found = [] | |
| for page_num in range(len(doc)): | |
| page = doc[page_num] | |
| # Extract text blocks with position info | |
| blocks = page.get_text("dict", flags=fitz.TEXTFLAGS_TEXT)["blocks"] | |
| page_text_parts = [] | |
| for block in blocks: | |
| if block.get("type", 0) == 0: # Text block | |
| lines_text = [] | |
| for line in block.get("lines", []): | |
| spans_text = "" | |
| max_font_size = 0 | |
| is_bold = False | |
| for span in line.get("spans", []): | |
| spans_text += span.get("text", "") | |
| font_size = span.get("size", 12) | |
| if font_size > max_font_size: | |
| max_font_size = font_size | |
| if "bold" in span.get("font", "").lower(): | |
| is_bold = True | |
| line_text = spans_text.strip() | |
| if line_text: | |
| lines_text.append({ | |
| "text": line_text, | |
| "font_size": max_font_size, | |
| "is_bold": is_bold, | |
| }) | |
| for line_info in lines_text: | |
| page_text_parts.append({ | |
| "text": line_info["text"], | |
| "page": page_num + 1, | |
| "font_size": line_info["font_size"], | |
| "is_bold": line_info["is_bold"], | |
| }) | |
| # Also extract any tables | |
| try: | |
| tables = page.find_tables() | |
| if tables and tables.tables: | |
| for table in tables.tables: | |
| table_text = self._format_table(table) | |
| if table_text.strip(): | |
| page_text_parts.append({ | |
| "text": table_text, | |
| "page": page_num + 1, | |
| "font_size": 12, | |
| "is_bold": False, | |
| "is_table": True, | |
| }) | |
| except Exception: | |
| pass # Table extraction may not be available in all versions | |
| all_blocks.extend(page_text_parts) | |
| total_pages = len(doc) | |
| doc.close() | |
| # Build sections and chunk | |
| chunks = self._build_chunks(all_blocks, sections_found) | |
| result = ParsedDocument( | |
| chunks=chunks, | |
| total_pages=total_pages, | |
| total_chunks=len(chunks), | |
| filename=pdf_path.name, | |
| sections=sections_found, | |
| ) | |
| logger.info(f"Parsed {result.total_pages} pages into {result.total_chunks} chunks") | |
| return result | |
| def _is_heading(self, text: str, font_size: float = 12, is_bold: bool = False) -> bool: | |
| """Detect if a line is a section heading.""" | |
| if font_size > 14 and is_bold: | |
| return True | |
| if font_size > 16: | |
| return True | |
| for pattern in self._HEADING_PATTERNS: | |
| if pattern.match(text.strip()): | |
| return True | |
| return False | |
| def _format_table(self, table) -> str: | |
| """Format a table into readable text.""" | |
| try: | |
| rows = table.extract() | |
| if not rows: | |
| return "" | |
| lines = [] | |
| for row in rows: | |
| cells = [str(cell) if cell is not None else "" for cell in row] | |
| lines.append(" | ".join(cells)) | |
| return "\n".join(lines) | |
| except Exception: | |
| return "" | |
| def _build_chunks( | |
| self, | |
| blocks: List[dict], | |
| sections_found: List[str], | |
| ) -> List[DocumentChunk]: | |
| """Build overlapping chunks from extracted blocks.""" | |
| chunks = [] | |
| current_section = "Document Start" | |
| current_text = "" | |
| current_page = 1 | |
| chunk_id = 0 | |
| for block in blocks: | |
| text = block["text"] | |
| page = block.get("page", 1) | |
| font_size = block.get("font_size", 12) | |
| is_bold = block.get("is_bold", False) | |
| is_table = block.get("is_table", False) | |
| # Check if this is a heading | |
| if self._is_heading(text, font_size, is_bold) and not is_table: | |
| # Flush current buffer as a chunk | |
| if len(current_text.strip()) >= self.min_chunk_size: | |
| chunks.append(DocumentChunk( | |
| chunk_id=chunk_id, | |
| text=current_text.strip(), | |
| page_number=current_page, | |
| section_title=current_section, | |
| chunk_type="text", | |
| )) | |
| chunk_id += 1 | |
| current_section = text.strip() | |
| sections_found.append(current_section) | |
| # Start new chunk with overlap from previous | |
| overlap_text = current_text[-self.chunk_overlap:] if len(current_text) > self.chunk_overlap else "" | |
| current_text = overlap_text + "\n" + text + "\n" | |
| current_page = page | |
| continue | |
| # Handle tables as separate chunks | |
| if is_table: | |
| if len(current_text.strip()) >= self.min_chunk_size: | |
| chunks.append(DocumentChunk( | |
| chunk_id=chunk_id, | |
| text=current_text.strip(), | |
| page_number=current_page, | |
| section_title=current_section, | |
| chunk_type="text", | |
| )) | |
| chunk_id += 1 | |
| current_text = "" | |
| chunks.append(DocumentChunk( | |
| chunk_id=chunk_id, | |
| text=text.strip(), | |
| page_number=page, | |
| section_title=current_section, | |
| chunk_type="table", | |
| )) | |
| chunk_id += 1 | |
| continue | |
| # Accumulate text | |
| current_text += " " + text | |
| current_page = page | |
| # Check if we exceeded chunk size | |
| if len(current_text.split()) >= self.chunk_size: | |
| words = current_text.split() | |
| chunk_words = words[:self.chunk_size] | |
| overlap_words = words[max(0, self.chunk_size - self.chunk_overlap):] | |
| chunks.append(DocumentChunk( | |
| chunk_id=chunk_id, | |
| text=" ".join(chunk_words).strip(), | |
| page_number=current_page, | |
| section_title=current_section, | |
| chunk_type="text", | |
| )) | |
| chunk_id += 1 | |
| current_text = " ".join(overlap_words) | |
| # Flush remaining text | |
| if len(current_text.strip()) >= self.min_chunk_size: | |
| chunks.append(DocumentChunk( | |
| chunk_id=chunk_id, | |
| text=current_text.strip(), | |
| page_number=current_page, | |
| section_title=current_section, | |
| chunk_type="text", | |
| )) | |
| return chunks | |