Add: enhance embedding process with document cleaning and batch validation

src/embedding.py

@@ -2,7 +2,7 @@ from typing import List, Any, Optional
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from sentence_transformers import SentenceTransformer
 import numpy as np
-from src.data_loader import load_all_documents
+from src.cleaner import get_default_cleaner
 
 class EmbeddingPipeline:
     def __init__(self, model_name: str = 'all-MiniLM-L6-v2', chunk_size: int = 1000, chunk_overlap: int = 200):
@@ -10,26 +10,75 @@ class EmbeddingPipeline:
         self.chunk_size = chunk_size
         self.chunk_overlap = chunk_overlap
         self.model = SentenceTransformer(model_name)
+        self.cleaner = get_default_cleaner()
         print(f"[INFO] Loaded embedding model: {model_name}")
         
     def chunk_documents(self, documents: List[Any]) -> List[Any]:
+        # 1. Clean documents first
+        cleaned_docs = self.cleaner.clean_documents(documents)
+        
+        # 2. Split into chunks
         splitter = RecursiveCharacterTextSplitter(
             chunk_size=self.chunk_size,
             chunk_overlap=self.chunk_overlap,
             length_function=len,
-            separators=["\n\n", "\n", " ", ""]
+            separators=["\n\n", "\n", " "]
         )
         
-        chunks = splitter.split_documents(documents)
-        print(f"[INFO] Split {len(documents)} documents into {len(chunks)} chunks.")
+        chunks = splitter.split_documents(cleaned_docs)
+        print(f"[INFO] Split {len(cleaned_docs)} documents into {len(chunks)} chunks.")
         return chunks
     
-    def embed_chunks(self, chunks: List[Any]) -> np.ndarray:
-        texts = [chunk.page_content for chunk in chunks]
-        print(f"[INFO] Generating embeddings for {len(texts)} chunks...")
-        embeddings = self.model.encode(texts, show_progress_bar = True)
-        print(f"[INFO] Embeddings shape : {embeddings.shape}")
-        return embeddings
+    def embed_chunks(self, chunks: List[Any], batch_size: int = 500) -> tuple[np.ndarray, List[Any]]:
+        """
+        Embeds chunks in batches for stability and better error reporting.
+        Returns a tuple of (embeddings_array, valid_chunks_list).
+        """
+        valid_chunks = []
+        all_embeddings = []
+        
+        print(f"[INFO] Generating embeddings for {len(chunks)} chunks in batches of {batch_size}...")
+        
+        for i in range(0, len(chunks), batch_size):
+            batch = chunks[i:i+batch_size]
+            batch_texts = []
+            current_batch_chunks = []
+            
+            # Additional validation per chunk
+            for chunk in batch:
+                content = getattr(chunk, 'page_content', None)
+                if content is not None and isinstance(content, str) and content.strip():
+                    batch_texts.append(content)
+                    current_batch_chunks.append(chunk)
+            
+            if not batch_texts:
+                continue
+                
+            try:
+                batch_embeddings = self.model.encode(batch_texts, show_progress_bar=False)
+                all_embeddings.append(batch_embeddings)
+                valid_chunks.extend(current_batch_chunks)
+            except Exception as e:
+                print(f"[ERROR] Failed to embed batch starting at index {i}. Error: {e}")
+                print("[INFO] Attempting to identify problematic chunk in batch...")
+                for j, text in enumerate(batch_texts):
+                    try:
+                        self.model.encode([text], show_progress_bar=False)
+                    except Exception as ex:
+                        print(f"[ERROR] Problematic chunk found at original index {i+j}!")
+                        # Safely encode for windows terminal printing to avoid crashes
+                        safe_text = text[:100].encode('ascii', 'replace').decode('ascii')
+                        print(f"[DEBUG] Content snippet: {safe_text}...")
+                        # We skip this specific chunk and continue
+                        continue
+        
+        if not all_embeddings:
+            print("[WARNING] No embeddings were generated.")
+            return np.array([]), []
+            
+        final_embeddings = np.vstack(all_embeddings)
+        print(f"[INFO] Total valid embeddings: {final_embeddings.shape[0]} / {len(chunks)}")
+        return final_embeddings, valid_chunks
     
 if __name__ == "__main__":
     docs = load_all_documents('Research/data/pdf')