CoolFace
Apppublic

aniket47/document-intelligence-chatbot

sourceHugging Facemitupdated 11mo agoView on Hugging Face
0likes
document_processor.py171 linesDownload Raw Back to components
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