Spaces:
Sleeping
Sleeping
| """ | |
| Document processor module for handling various document types and chunking methods. | |
| Integrates PDF loading and chunking functionality for the Gradio interface. | |
| """ | |
| import logging | |
| from typing import Dict, List, Optional | |
| import os | |
| from pathlib import Path | |
| # Import existing modules | |
| from pdf_loader import PDFLoader | |
| from chunking_methods import ChunkingManager, FixedChunking, SentenceChunking, AdaptiveChunking | |
| import settings | |
| logger = logging.getLogger(__name__) | |
| class DocumentProcessor: | |
| """ | |
| Handles document loading and processing with multiple chunking methods. | |
| """ | |
| def __init__(self): | |
| """Initialize the document processor with available chunking methods.""" | |
| self.chunking_manager = ChunkingManager() | |
| self.pdf_loader = PDFLoader() | |
| # Register available chunking methods | |
| self.chunking_strategies = { | |
| 'fixed': FixedChunking(), | |
| 'sentence': SentenceChunking(), | |
| 'adaptive': AdaptiveChunking() | |
| } | |
| logger.info(f"DocumentProcessor initialized with {len(self.chunking_strategies)} chunking strategies") | |
| def load_document(self, file_path: str) -> Optional[str]: | |
| """ | |
| Load document content based on file type. | |
| Args: | |
| file_path: Path to the document file | |
| Returns: | |
| Document text content or None if loading fails | |
| """ | |
| try: | |
| if not os.path.exists(file_path): | |
| logger.error(f"File not found: {file_path}") | |
| return None | |
| file_extension = Path(file_path).suffix.lower() | |
| if file_extension == '.pdf': | |
| # Use PDFLoader to load the file directly | |
| try: | |
| content = self.pdf_loader.load_pdf(Path(file_path)) | |
| return content | |
| except Exception as e: | |
| logger.error(f"Error loading PDF {file_path}: {str(e)}") | |
| return None | |
| elif file_extension == '.txt': | |
| with open(file_path, 'r', encoding='utf-8') as f: | |
| return f.read() | |
| elif file_extension in ['.doc', '.docx']: | |
| try: | |
| from docx import Document | |
| doc = Document(file_path) | |
| content = [] | |
| for paragraph in doc.paragraphs: | |
| content.append(paragraph.text) | |
| return '\n'.join(content) | |
| except ImportError: | |
| logger.error("python-docx not installed. Cannot process DOCX files.") | |
| return None | |
| else: | |
| logger.error(f"Unsupported file type: {file_extension}") | |
| return None | |
| except Exception as e: | |
| logger.error(f"Error loading document {file_path}: {str(e)}") | |
| return None | |
| def process_document(self, file_path: str, chunking_methods: List[str]) -> Dict[str, List[str]]: | |
| """ | |
| Process a document with specified chunking methods. | |
| Args: | |
| file_path: Path to the document file | |
| chunking_methods: List of chunking method names to apply | |
| Returns: | |
| Dictionary mapping method names to lists of chunks | |
| """ | |
| try: | |
| # Load document content | |
| content = self.load_document(file_path) | |
| if not content: | |
| logger.error(f"Failed to load content from {file_path}") | |
| return {} | |
| logger.info(f"Loaded document with {len(content)} characters") | |
| # Process with each requested chunking method | |
| chunks_by_method = {} | |
| for method_name in chunking_methods: | |
| if method_name not in self.chunking_strategies: | |
| logger.warning(f"Unknown chunking method: {method_name}") | |
| continue | |
| try: | |
| strategy = self.chunking_strategies[method_name] | |
| chunks = strategy.create_chunks(content) | |
| chunks_by_method[method_name] = chunks | |
| logger.info(f"Generated {len(chunks)} chunks using {method_name} method") | |
| except Exception as e: | |
| logger.error(f"Error processing with {method_name} method: {str(e)}") | |
| continue | |
| return chunks_by_method | |
| except Exception as e: | |
| logger.error(f"Error processing document {file_path}: {str(e)}") | |
| return {} | |
| def get_supported_extensions(self) -> List[str]: | |
| """ | |
| Get list of supported file extensions. | |
| Returns: | |
| List of supported file extensions | |
| """ | |
| return ['.pdf', '.txt', '.docx'] | |
| def validate_file(self, file_path: str) -> bool: | |
| """ | |
| Validate if a file can be processed. | |
| Args: | |
| file_path: Path to the file to validate | |
| Returns: | |
| True if file can be processed, False otherwise | |
| """ | |
| if not os.path.exists(file_path): | |
| return False | |
| extension = Path(file_path).suffix.lower() | |
| return extension in self.get_supported_extensions() |