Spaces:
Sleeping
Sleeping
File size: 4,012 Bytes
a4538e5 d2029d7 a4538e5 d2029d7 a4538e5 d2029d7 a4538e5 d2029d7 a4538e5 d2029d7 a4538e5 d2029d7 a4538e5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | from typing import List, Any, Optional
from langchain_text_splitters import RecursiveCharacterTextSplitter
from sentence_transformers import SentenceTransformer
import numpy as np
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):
self.model_name = model_name
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", " "]
)
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], 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')
emb_pipe = EmbeddingPipeline()
chunks = emb_pipe.chunk_documents(docs)
embeddings = emb_pipe.embed_chunks(chunks)
print(f"[INFO] Example embeddings:", embeddings[0] if len(embeddings) > 0 else None)
|