aniket47/document-intelligence-chatbot
0
1import PyPDF22import re3from typing import List, Dict4import io5 6class DocumentProcessor:7 """8 Handles PDF document processing and intelligent text chunking9 """10 11 def __init__(self, chunk_size: int = 1000, chunk_overlap: int = 200):12 self.chunk_size = chunk_size13 self.chunk_overlap = chunk_overlap14 15 def extract_text_from_pdf(self, pdf_file) -> Dict[str, any]:16 """17 Extract text from PDF file and preserve metadata18 19 Args:20 pdf_file: Uploaded PDF file object21 22 Returns:23 Dict containing extracted text, metadata, and page information24 """25 try:26 # Read PDF using PyPDF227 pdf_reader = PyPDF2.PdfReader(pdf_file)28 29 # Extract metadata30 metadata = {31 'filename': pdf_file.name,32 'num_pages': len(pdf_reader.pages),33 'title': pdf_reader.metadata.get('/Title', '') if pdf_reader.metadata else '',34 'author': pdf_reader.metadata.get('/Author', '') if pdf_reader.metadata else '',35 'subject': pdf_reader.metadata.get('/Subject', '') if pdf_reader.metadata else ''36 }37 38 # Extract text from each page39 pages_text = []40 full_text = ""41 42 for page_num, page in enumerate(pdf_reader.pages):43 try:44 page_text = page.extract_text()45 if page_text.strip(): # Only add non-empty pages46 pages_text.append({47 'page_number': page_num + 1,48 'text': page_text.strip()49 })50 full_text += f"\n\n[Page {page_num + 1}]\n{page_text.strip()}"51 except Exception as e:52 print(f"Error extracting text from page {page_num + 1}: {str(e)}")53 continue54 55 return {56 'full_text': full_text.strip(),57 'pages': pages_text,58 'metadata': metadata59 }60 61 except Exception as e:62 raise Exception(f"Error processing PDF: {str(e)}")63 64 def clean_text(self, text: str) -> str:65 """66 Clean and normalize extracted text67 68 Args:69 text: Raw extracted text70 71 Returns:72 Cleaned text73 """74 # Remove excessive whitespace75 text = re.sub(r'\s+', ' ', text)76 77 # Remove special characters but keep punctuation78 text = re.sub(r'[^\w\s\.\,\;\:\!\?\-\(\)\[\]\"\'\/]', '', text)79 80 # Fix common PDF extraction issues81 text = text.replace('�', '') # Remove replacement characters82 text = re.sub(r'([a-z])([A-Z])', r'\1 \2', text) # Add space between words83 84 return text.strip()85 86 def chunk_text(self, text: str, metadata: Dict) -> List[Dict]:87 """88 Split text into overlapping chunks for better retrieval89 90 Args:91 text: Full document text92 metadata: Document metadata93 94 Returns:95 List of text chunks with metadata96 """97 # Clean the text first98 cleaned_text = self.clean_text(text)99 100 # Split into sentences for better chunking101 sentences = re.split(r'(?<=[.!?])\s+', cleaned_text)102 103 chunks = []104 current_chunk = ""105 current_length = 0106 chunk_id = 0107 108 for sentence in sentences:109 sentence_length = len(sentence)110 111 # If adding this sentence would exceed chunk size, save current chunk112 if current_length + sentence_length > self.chunk_size and current_chunk:113 chunks.append({114 'chunk_id': chunk_id,115 'text': current_chunk.strip(),116 'metadata': {117 **metadata,118 'chunk_size': len(current_chunk),119 'chunk_index': chunk_id120 }121 })122 chunk_id += 1123 124 # Start new chunk with overlap125 if self.chunk_overlap > 0:126 # Take last few sentences for overlap127 overlap_sentences = current_chunk.split('. ')[-2:]128 current_chunk = '. '.join(overlap_sentences) + '. ' + sentence129 current_length = len(current_chunk)130 else:131 current_chunk = sentence132 current_length = sentence_length133 else:134 # Add sentence to current chunk135 if current_chunk:136 current_chunk += " " + sentence137 else:138 current_chunk = sentence139 current_length += sentence_length140 141 # Add the last chunk if it exists142 if current_chunk.strip():143 chunks.append({144 'chunk_id': chunk_id,145 'text': current_chunk.strip(),146 'metadata': {147 **metadata,148 'chunk_size': len(current_chunk),149 'chunk_index': chunk_id150 }151 })152 153 return chunks154 155 def process_document(self, pdf_file) -> List[Dict]:156 """157 Complete document processing pipeline158 159 Args:160 pdf_file: Uploaded PDF file161 162 Returns:163 List of processed text chunks with metadata164 """165 # Extract text and metadata166 doc_data = self.extract_text_from_pdf(pdf_file)167 168 # Create chunks169 chunks = self.chunk_text(doc_data['full_text'], doc_data['metadata'])170 171 return chunks