Spaces:
Sleeping
Sleeping
File size: 4,367 Bytes
c0f31c1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | """
PageChunker.py
A module for page-level document chunking with token counting and preprocessing.
Features:
- Page-based document splitting
- Content validation
- Blank page detection
- Document metadata enrichment
"""
from typing import List, Optional
import logging
from langchain_core.documents import Document
from core.BaseChunker import BaseChunker
logger = logging.getLogger(__name__)
class PageChunker(BaseChunker):
"""Handles document chunking at the page level."""
def __init__(self, model_name=None, embedding_model=None):
"""
Initialize page chunker with specified models.
Args:
model_name: Name of the model for tokenization
embedding_model: Model for generating embeddings
"""
super().__init__(model_name, embedding_model)
self.page_stats = []
def _is_blank_page(self, text: str) -> bool:
"""Check if page is blank or contains only whitespace/special characters."""
cleaned_text = text.strip().replace('\n', '').replace('\r', '').replace('\t', '')
return len(cleaned_text) < self.BLANK_THRESHOLD
def _process_single_page(self, content: str, page_number: int, preprocess: bool) -> Optional[Document]:
"""
Process a single page with optional preprocessing and analysis.
Args:
content: The page content
page_number: The page number
preprocess: Whether to preprocess the text
Returns:
Document object with processed content and metadata, or None if page is blank
"""
if self._is_blank_page(content):
self.page_stats.append(f"Page {page_number} is blank.")
return None
# Optionally preprocess the text
if preprocess:
content = self.preprocess_text(content)
# Analyze the page and generate metadata
stats = self.analyze_text(content)
metadata = {
"page": page_number,
"char_count": stats["char_count"],
"token_count": stats["token_count"],
"sentence_count": stats["sentence_count"],
"word_count": stats["word_count"],
"has_ocr": str(stats.get("has_content", True)),
"is_blank": "false"
}
return Document(page_content=content, metadata=metadata)
def page_process_document(self, file_path: str, preprocess: bool = False) -> List[Document]:
"""
Process PDF document page by page with analysis and optional preprocessing.
Args:
file_path: Path to the PDF file
preprocess: Whether to preprocess page text
Returns:
List of Document objects, one per non-blank page
"""
try:
self.page_stats = [] # Reset stats for this document
raw_pages = self.load_document(file_path)
processed_pages = []
logger.info(f"Processing document with {len(raw_pages)} pages")
for idx, page in enumerate(raw_pages):
processed_page = self._process_single_page(page.page_content, idx + 1, preprocess)
if processed_page:
processed_pages.append(processed_page)
# Output skipped pages for transparency
if self.page_stats:
logger.info("\n".join(self.page_stats))
logger.info(f"Processed {len(processed_pages)} non-blank pages")
return processed_pages
except Exception as e:
logger.error(f"Error in page_process_document: {e}")
raise
def process_document(self, file_path: str, preprocess: bool = True) -> List[Document]:
"""
Process document using page chunking strategy (implements abstract method).
Args:
file_path: Path to the PDF file
preprocess: Whether to preprocess page text
Returns:
List of Document objects, one per non-blank page
"""
return self.page_process_document(file_path, preprocess) |