Upload 2 files

app.py

@@ -1,145 +1,97 @@
 import os
 import tempfile
 import streamlit as st
-import numpy as np
-from sentence_transformers import SentenceTransformer
-from sklearn.metrics.pairwise import cosine_similarity
+import PyPDF2
 from document_processor import DocumentProcessor
+import logging
 
-# Initialize session state variables if they don't exist
-if 'documents' not in st.session_state:
-    st.session_state.documents = []
-if 'embeddings' not in st.session_state:
-    st.session_state.embeddings = []
-if 'search_results' not in st.session_state:
-    st.session_state.search_results = []
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+def extract_text_from_pdf(file_path):
+    """Extract text from PDF with robust error handling"""
+    text = ""
+    try:
+        with open(file_path, "rb") as f:
+            pdf_reader = PyPDF2.PdfReader(f)
+            
+            # Get the number of pages
+            num_pages = len(pdf_reader.pages)
+            logger.info(f"PDF has {num_pages} pages")
+            
+            # Extract text from each page
+            for page_num in range(num_pages):
+                page = pdf_reader.pages[page_num]
+                try:
+                    page_text = page.extract_text()
+                    if page_text:
+                        text += page_text + "\n\n"
+                except Exception as e:
+                    logger.error(f"Error extracting text from page {page_num}: {e}")
+    except Exception as e:
+        logger.error(f"Error reading PDF file {file_path}: {e}")
+    
+    logger.info(f"Extracted {len(text)} characters from PDF")
+    
+    # If text is very short, it might be a scanned PDF that needs OCR
+    if len(text.strip()) < 100:
+        logger.warning(f"Very little text extracted from {file_path}. This may be a scanned PDF requiring OCR.")
+    
+    return text
 
-# Function to process uploaded files
 def process_uploaded_files(uploaded_files):
-    """Process each uploaded file and add chunks to documents."""
+    if not uploaded_files:
+        return None
+        
+    # Create a temporary directory to store processed files
     temp_dir = tempfile.mkdtemp()
-    processor = DocumentProcessor(embedding_model="sentence-transformers/all-MiniLM-L6-v2",chunk_size=500, chunk_overlap=50)
+    logger.info(f"Created temporary directory: {temp_dir}")
+    
+    # Initialize document processor with embedding model
+    processor = DocumentProcessor(
+        embedding_model="sentence-transformers/all-MiniLM-L6-v2",
+        chunk_size=500, 
+        chunk_overlap=50
+    )
     
+    # Process each file
     for uploaded_file in uploaded_files:
-        # Create a temporary file path
-        file_path = os.path.join(temp_dir, uploaded_file.name)
+        file_name = uploaded_file.name
+        logger.info(f"Processing file: {file_name}")
         
-        # Write the uploaded file to the temp directory
+        # Save the uploaded file to disk temporarily
+        file_path = os.path.join(temp_dir, file_name)
         with open(file_path, "wb") as f:
             f.write(uploaded_file.getbuffer())
         
-        try:
-            # Process the file
-            chunks = processor.process_document(file_path)
-            
-            # Add chunks to session state
-            st.session_state.documents.extend(chunks)
-            st.success(f"Processed {uploaded_file.name} - Added {len(chunks)} chunks")
-        except Exception as e:
-            st.error(f"Error processing {uploaded_file.name}: {str(e)}")
+        # Extract text based on file type
+        if file_name.lower().endswith('.pdf'):
+            text = extract_text_from_pdf(file_path)
+        elif file_name.lower().endswith(('.txt', '.md')):
+            with open(file_path, "r", encoding="utf-8", errors="replace") as f:
+                text = f.read()
+        else:
+            st.warning(f"Unsupported file type: {file_name}")
+            continue
+        
+        if not text or len(text.strip()) == 0:
+            st.warning(f"No text could be extracted from {file_name}")
+            continue
             
-    # Return the temp directory for later cleanup
-    return temp_dir
-
-# Function to create embeddings
-@st.cache_resource
-def load_embedding_model():
-    """Load the sentence transformer model."""
-    return SentenceTransformer('all-MiniLM-L6-v2')
-
-def create_embeddings():
-    """Create embeddings for all documents in the session state."""
-    model = load_embedding_model()
-    
-    # Extract texts from documents
-    texts = [doc["text"] for doc in st.session_state.documents]
-    
-    # Create embeddings
-    with st.spinner('Creating embeddings...'):
-        embeddings = model.encode(texts, show_progress_bar=True)
-    
-    st.session_state.embeddings = embeddings
-    st.success(f"Created embeddings for {len(embeddings)} document chunks")
-
-# Search function
-def search_documents(query, top_k=5):
-    """Search documents based on the query and return top k results."""
-    model = load_embedding_model()
-    
-    # Create query embedding
-    query_embedding = model.encode([query])[0]
-    
-    # Calculate similarities
-    similarities = cosine_similarity(
-        [query_embedding], 
-        st.session_state.embeddings
-    )[0]
-    
-    # Get top-k indices
-    top_indices = np.argsort(similarities)[-top_k:][::-1]
-    
-    # Create results
-    results = []
-    for idx in top_indices:
-        document = st.session_state.documents[idx]
-        results.append({
-            "document": document,
-            "similarity": similarities[idx],
-            "index": idx
-        })
-    
-    st.session_state.search_results = results
-    return results
-
-# Main App UI
-st.title("Document Search with Embeddings")
-
-# Sidebar for uploading files and controlling the app
-with st.sidebar:
-    st.header("Document Control")
-    
-    # File uploader
-    uploaded_files = st.file_uploader("Upload documents", accept_multiple_files=True, 
-                                     type=["txt", "md", "html", "pdf"])
-    
-    # Process uploaded files when the user clicks the button
-    if st.button("Process Files") and uploaded_files:
-        temp_dir = process_uploaded_files(uploaded_files)
+        # Create metadata
+        metadata = {
+            "file_name": file_name,
+            "file_path": file_path,
+        }
         
-    # Create embeddings button
-    if st.button("Create Embeddings") and st.session_state.documents:
-        create_embeddings()
+        # Process the document
+        doc_ids = processor.process_document(text, metadata)
+        st.info(f"Processed {file_name} - Added {len(doc_ids)} chunks")
     
-    # Display stats
-    st.subheader("Stats")
-    st.write(f"Documents: {len(st.session_state.documents)}")
-    st.write(f"Embeddings: {len(st.session_state.embeddings)}")
+    # Save the processor state
+    processor.save(os.path.join(temp_dir, "processor_state"))
     
-    # Instructions and tips
-    st.subheader("Processing Tips")
-    st.info("Supported file types: TXT, MD, HTML, PDF")
-    st.info("1. Upload your documents\n2. Process the files\n3. Create embeddings\n4. Search!")
+    return temp_dir
 
-# Main area - Search interface
-if st.session_state.embeddings is not None and len(st.session_state.embeddings) > 0:
-    st.subheader("Document Search")
-    
-    # Search query input
-    query = st.text_input("Enter your search query")
-    num_results = st.slider("Number of results", min_value=1, max_value=10, value=3)
-    
-    # Search button
-    if st.button("Search") and query:
-        results = search_documents(query, top_k=num_results)
-        
-    # Display search results
-    if st.session_state.search_results:
-        st.subheader("Search Results")
-        for i, result in enumerate(st.session_state.search_results):
-            with st.expander(f"Result {i+1} - {result['document']['source']} (Score: {result['similarity']:.4f})"):
-                st.markdown(f"**Source:** {result['document']['source']}")
-                st.markdown(f"**Chunk ID:** {result['document']['chunk_id']}")
-                st.markdown("**Content:**")
-                st.markdown(result['document']['text'])
-else:
-    st.info("Please upload documents and create embeddings to enable search.")
+# The rest of your app.py code remains the same...

document_processor.py

@@ -1,210 +1,327 @@
-from typing import List, Dict, Any, Optional, Union, Callable
+from typing import List, Dict, Any, Optional, Tuple, Union
+import os
 import re
 import uuid
+import json
 import numpy as np
-from dataclasses import dataclass
-
-from embedding_model import EmbeddingModel
-from memory_db import VectorDatabase
-
-@dataclass
-class ProcessedChunk:
-    """Represents a processed text chunk with metadata"""
-    text: str
-    metadata: Dict[str, Any]
-    doc_id: str
-    embedding: Optional[np.ndarray] = None
+from tqdm import tqdm
+from memory_db import VectorDatabase, Document
+from transformers import AutoTokenizer, AutoModel
+import torch
+import logging
 
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
 class DocumentProcessor:
-    """Handles document splitting and processing for vector storage"""
+    """Process documents for semantic search"""
     
     def __init__(
-        self,
-        embedding_model: EmbeddingModel,
-        chunk_size: int = 500,
+        self, 
+        embedding_model: Optional[str] = "sentence-transformers/all-MiniLM-L6-v2",
+        chunk_size: int = 500, 
         chunk_overlap: int = 50,
-        min_chunk_size: int = 100,
-        metadata_extractor: Optional[Callable[[str], Dict[str, Any]]] = None
+        device: Optional[str] = None
     ):
         """
-        Initialize document processor
+        Initialize the document processor
         
         Args:
-            embedding_model: Model used to create embeddings
-            chunk_size: Target size of chunks in characters
+            embedding_model: Model name or path for generating embeddings
+            chunk_size: Size of text chunks in characters
             chunk_overlap: Overlap between chunks in characters
-            min_chunk_size: Minimum size for chunks
-            metadata_extractor: Optional function to extract metadata from documents
+            device: Device to run model on ('cpu', 'cuda', etc.)
         """
-        self.embedding_model = embedding_model
         self.chunk_size = chunk_size
         self.chunk_overlap = chunk_overlap
-        self.min_chunk_size = min_chunk_size
-        self.metadata_extractor = metadata_extractor or (lambda _: {})
-    
-    def process_document(
-        self, 
-        document: str, 
-        doc_id: Optional[str] = None,
-        metadata: Optional[Dict[str, Any]] = None
-    ) -> List[ProcessedChunk]:
+        
+        # Initialize embedding model
+        if embedding_model is not None:
+            logger.info(f"Loading embedding model: {embedding_model}")
+            self.tokenizer = AutoTokenizer.from_pretrained(embedding_model)
+            self.model = AutoModel.from_pretrained(embedding_model)
+            
+            # Set device (use CUDA if available)
+            if device is None:
+                self.device = "cuda" if torch.cuda.is_available() else "cpu"
+            else:
+                self.device = device
+                
+            logger.info(f"Using device: {self.device}")
+            self.model.to(self.device)
+            self.embedding_dim = self.model.config.hidden_size
+        else:
+            self.tokenizer = None
+            self.model = None
+            self.device = "cpu"
+            self.embedding_dim = 384  # Default dimension
+        
+        # Initialize database
+        self.db = VectorDatabase(dimension=self.embedding_dim)
+        
+    def _generate_embedding(self, text: str) -> np.ndarray:
         """
-        Process a document into chunks with embeddings
+        Generate embedding for a text using the model
         
         Args:
-            document: Document text to process
-            doc_id: Optional document ID (will be auto-generated if None)
-            metadata: Optional metadata to associate with all chunks
+            text: Input text
             
         Returns:
-            List of processed chunks
+            Embedding vector
         """
-        # Generate document ID if not provided
-        if doc_id is None:
-            doc_id = str(uuid.uuid4())
+        if self.model is None:
+            raise ValueError("No embedding model provided")
+        
+        if not text or text.isspace():
+            logger.warning("Attempted to generate embedding for empty text")
+            # Return zero vector as fallback
+            return np.zeros((1, self.embedding_dim))
             
-        # Extract metadata if not provided
-        if metadata is None:
-            metadata = self.metadata_extractor(document)
+        # Truncate text if too long for model
+        max_length = self.tokenizer.model_max_length
         
-        # Split the document into chunks
-        chunks = self.split_text(document)
+        # Tokenize and prepare for model
+        inputs = self.tokenizer(
+            text, 
+            max_length=max_length, 
+            padding=True, 
+            truncation=True,
+            return_tensors="pt"
+        ).to(self.device)
         
-        # Process each chunk
-        processed_chunks = []
-        for i, chunk_text in enumerate(chunks):
-            # Skip chunks that are too small
-            if len(chunk_text) < self.min_chunk_size:
-                continue
-                
-            # Create chunk metadata (including position info)
-            chunk_metadata = {
-                **metadata,
-                "chunk_index": i,
-                "total_chunks": len(chunks),
-                "document_id": doc_id
-            }
-            
-            # Generate chunk ID
-            chunk_id = f"{doc_id}_{i}"
-            
-            # Create embedding
-            embedding = self.embedding_model.embed_text(chunk_text)
-            
-            # Create processed chunk
-            chunk = ProcessedChunk(
-                text=chunk_text,
-                metadata=chunk_metadata,
-                doc_id=chunk_id,
-                embedding=embedding
-            )
-            
-            processed_chunks.append(chunk)
-            
-        return processed_chunks
-    
-    def split_text(self, text: str) -> List[str]:
+        # Generate embeddings
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            
+        # Get CLS token embedding (first token)
+        embedding = outputs.last_hidden_state[:, 0, :].cpu().numpy()
+        
+        return embedding[0]  # Return the actual vector, not a batch
+        
+    def chunk_text(self, text: str) -> List[str]:
         """
-        Split text into overlapping chunks
+        Split text into chunks of approximately chunk_size with overlap
         
         Args:
-            text: Text to split
+            text: Text to split into chunks
             
         Returns:
             List of text chunks
         """
-        # Clean and normalize text
-        text = self._normalize_text(text)
+        # Handle empty text
+        if not text or len(text.strip()) == 0:
+            logger.warning("Empty text provided for chunking")
+            return []
+            
+        # Clean text - remove excessive whitespace
+        text = re.sub(r'\s+', ' ', text).strip()
+        
+        logger.info(f"Chunking text of length {len(text)} characters")
         
-        # Simple character-based chunking
         chunks = []
         start = 0
+        text_length = len(text)
         
-        while start < len(text):
-            # Determine end position for this chunk
-            end = start + self.chunk_size
+        while start < text_length:
+            # Get chunk end position
+            end = min(start + self.chunk_size, text_length)
             
-            # Adjust to end at sentence boundary if possible
-            if end < len(text):
-                # Try to find sentence boundary near the end
-                sentence_end = self._find_sentence_boundary(text, end)
-                if sentence_end > start + self.min_chunk_size:
-                    end = sentence_end
-            else:
-                end = len(text)
+            # If not at the end of text, try to find a good break point
+            if end < text_length:
+                # Try to break at sentence or punctuation
+                break_chars = '.!?。，,;；'
+                # Find the last break character in the chunk
+                last_break = max([text.rfind(char, start, end) for char in break_chars])
+                
+                # If found a break point, use it
+                if last_break > start:
+                    end = last_break + 1
+                else:
+                    # If no good break, try to break at space
+                    last_space = text.rfind(' ', start, end)
+                    if last_space > start:
+                        end = last_space + 1
             
-            # Add the chunk
-            chunks.append(text[start:end])
+            # Add chunk
+            chunk_text = text[start:end].strip()
+            if chunk_text:  # Only add non-empty chunks
+                chunks.append(chunk_text)
             
-            # Move start position for next chunk
-            start = end - self.chunk_overlap
-            if start < 0 or start >= len(text):
-                break
+            # Move to next chunk with overlap
+            start = end - self.chunk_overlap if end < text_length else text_length
         
+        logger.info(f"Created {len(chunks)} chunks")
         return chunks
-    
-    def _normalize_text(self, text: str) -> str:
-        """Normalize text by removing extra whitespace"""
-        # Replace multiple whitespace with single space
-        text = re.sub(r'\s+', ' ', text)
-        # Trim whitespace from start and end
-        return text.strip()
-    
-    def _find_sentence_boundary(self, text: str, position: int) -> int:
+        
+    def process_document(
+        self, 
+        text: str, 
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> List[str]:
         """
-        Find the nearest sentence boundary after the given position
+        Process a document and add to database
         
         Args:
-            text: Text to search in
-            position: Position to start searching from
+            text: Document text
+            metadata: Optional metadata for the document
             
         Returns:
-            Index of sentence boundary
+            List of document IDs created
         """
-        # Look for typical sentence end patterns after position
-        for i in range(position, min(position + 100, len(text))):
-            if i >= len(text):
-                return len(text)
+        metadata = metadata or {}
+        doc_ids = []
+        
+        # Log text details
+        logger.info(f"Processing document: {len(text)} chars, metadata: {metadata}")
+        
+        if not text or len(text.strip()) == 0:
+            logger.warning("Empty document text provided")
+            return doc_ids
+        
+        # Chunk text
+        chunks = self.chunk_text(text)
+        
+        if not chunks:
+            logger.warning("No chunks were created from the document")
+            return doc_ids
+        
+        # Generate embeddings and add to database
+        for i, chunk in enumerate(chunks):
+            # Skip empty chunks
+            if not chunk or chunk.isspace():
+                continue
                 
-            # Check for sentence endings
-            if text[i] in '.!?' and (i == len(text) - 1 or text[i+1] == ' '):
-                return i + 1
+            try:
+                # Generate embedding
+                embedding = self._generate_embedding(chunk)
                 
-        # If no sentence boundary found, just return the original position
-        return position
-    
-    def add_to_database(
+                # Create chunk metadata
+                chunk_metadata = metadata.copy()
+                chunk_metadata["chunk_index"] = i
+                chunk_metadata["chunk_count"] = len(chunks)
+                
+                # Add to database
+                doc_id = self.db.add_document(
+                    text=chunk,
+                    embedding=embedding,
+                    metadata=chunk_metadata
+                )
+                
+                doc_ids.append(doc_id)
+                
+            except Exception as e:
+                logger.error(f"Error processing chunk {i}: {e}")
+                continue
+            
+        logger.info(f"Added {len(doc_ids)} chunks to database")
+        return doc_ids
+        
+    def process_documents(
         self, 
-        document: str, 
-        database: VectorDatabase,
-        doc_id: Optional[str] = None,
-        metadata: Optional[Dict[str, Any]] = None
-    ) -> List[str]:
+        documents: List[Tuple[str, Dict[str, Any]]]
+    ) -> Dict[str, List[str]]:
+        """
+        Process multiple documents and add to database
+        
+        Args:
+            documents: List of (text, metadata) tuples
+            
+        Returns:
+            Dictionary mapping metadata keys to lists of document IDs
+        """
+        result = {}
+        
+        for text, metadata in tqdm(documents, desc="Processing documents"):
+            doc_id = str(uuid.uuid4())
+            metadata["document_id"] = doc_id
+            
+            chunk_ids = self.process_document(text, metadata)
+            
+            # Store IDs by document ID
+            result[doc_id] = chunk_ids
+            
+        return result
+        
+    def search(
+        self, 
+        query: str, 
+        top_k: int = 5,
+        filter_func: Optional[callable] = None
+    ) -> List[Tuple[Document, float]]:
+        """
+        Search for documents similar to query
+        
+        Args:
+            query: Search query
+            top_k: Number of results to return
+            filter_func: Optional function to filter results
+            
+        Returns:
+            List of (document, score) tuples
+        """
+        # Generate embedding for query
+        query_embedding = self._generate_embedding(query)
+        
+        # Search in database
+        results = self.db.search(
+            query_embedding=query_embedding,
+            top_k=top_k,
+            filter_func=filter_func
+        )
+        
+        return results
+        
+    def save(self, directory: str) -> None:
+        """
+        Save processor state to directory
+        
+        Args:
+            directory: Directory to save to
+        """
+        os.makedirs(directory, exist_ok=True)
+        
+        # Save database
+        self.db.save(os.path.join(directory, "database"))
+        
+        # Save config
+        config = {
+            "chunk_size": self.chunk_size,
+            "chunk_overlap": self.chunk_overlap,
+            "embedding_dim": self.embedding_dim
+        }
+        
+        with open(os.path.join(directory, "config.json"), "w") as f:
+            json.dump(config, f)
+            
+    @classmethod
+    def load(
+        cls, 
+        directory: str,
+        embedding_model: Optional[str] = "sentence-transformers/all-MiniLM-L6-v2"
+    ) -> 'DocumentProcessor':
         """
-        Process a document and add its chunks to the database
+        Load processor from directory
         
         Args:
-            document: Document text to process
-            database: Vector database to add chunks to
-            doc_id: Optional document ID
-            metadata: Optional metadata to associate with chunks
+            directory: Directory to load from
+            embedding_model: Model to use for embeddings
             
         Returns:
-            List of chunk IDs added to the database
-        """
-        # Process document
-        chunks = self.process_document(document, doc_id, metadata)
-        
-        # Add each chunk to the database
-        chunk_ids = []
-        for chunk in chunks:
-            database.add_document(
-                doc_id=chunk.doc_id,
-                text=chunk.text,
-                metadata=chunk.metadata,
-                embedding=chunk.embedding
-            )
-            chunk_ids.append(chunk.doc_id)
-            
-        return chunk_ids
+            Loaded DocumentProcessor
+        """
+        # Load config
+        with open(os.path.join(directory, "config.json"), "r") as f:
+            config = json.load(f)
+            
+        # Create processor
+        processor = cls(
+            embedding_model=embedding_model,
+            chunk_size=config["chunk_size"],
+            chunk_overlap=config["chunk_overlap"]
+        )
+        
+        # Load database
+        processor.db = VectorDatabase.load(os.path.join(directory, "database"))
+        
+        return processor