Spaces:
Sleeping
Sleeping
File size: 3,669 Bytes
dff2c1e | 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 | """Quick test to verify the updated embedding pipeline works correctly."""
import sys
import os
sys.path.insert(0, os.getcwd())
from langchain_core.documents import Document
from src.embedding import EmbeddingPipeline
from src.vectorstore import FaissVectorStore
from src.cleaner import DocumentCleaner
import numpy as np
def test_pipeline():
print("=" * 60)
print("QUICK PIPELINE VERIFICATION TEST")
print("=" * 60)
# 1. Test DocumentCleaner
print("\n--- Test 1: DocumentCleaner ---")
cleaner = DocumentCleaner(min_length=10)
test_docs = [
Document(page_content="This is a valid document with real content about databases."),
Document(page_content="Short"), # too short
Document(page_content=""), # empty
Document(page_content=None), # None - this would crash old code
Document(page_content="Another valid document explaining SQL queries and joins."),
Document(page_content=" \n\n\t "), # whitespace only
]
cleaned = cleaner.clean_documents(test_docs)
print(f" Input: {len(test_docs)} docs -> Cleaned: {len(cleaned)} docs")
assert len(cleaned) == 2, f"Expected 2 valid docs, got {len(cleaned)}"
print(" ✅ DocumentCleaner works correctly!")
# 2. Test EmbeddingPipeline with cleaning + batching
print("\n--- Test 2: EmbeddingPipeline (clean + chunk + embed) ---")
# Create realistic documents like what PyPDFLoader returns
docs = [
Document(page_content="Database Management Systems provide an organized way to store and manage data. " * 5),
Document(page_content="SQL is a standard language for accessing and manipulating databases. " * 5),
Document(page_content=None), # Simulates a bad PDF page
Document(page_content="Machine Learning is a branch of artificial intelligence focused on algorithms. " * 5),
Document(page_content=""), # Empty page
Document(page_content="x"), # Too short - should be cleaned
]
pipe = EmbeddingPipeline(chunk_size=200, chunk_overlap=50)
chunks = pipe.chunk_documents(docs)
print(f" Chunks created: {len(chunks)}")
embeddings, valid_chunks = pipe.embed_chunks(chunks, batch_size=2)
print(f" Embeddings shape: {embeddings.shape}")
print(f" Valid chunks: {len(valid_chunks)}")
assert embeddings.shape[0] == len(valid_chunks), "Embeddings and chunks count must match!"
assert embeddings.shape[0] > 0, "Should have some embeddings!"
print(" ✅ EmbeddingPipeline works correctly!")
# 3. Test FaissVectorStore integration
print("\n--- Test 3: FaissVectorStore build + query ---")
test_store_dir = "test_faiss_store"
store = FaissVectorStore(persist_dir=test_store_dir, chunk_size=200, chunk_overlap=50)
store.build_from_documents(docs)
if store.index is not None:
results = store.query("What is a database?", top_k=2)
print(f" Query returned {len(results)} results")
for r in results:
snippet = r['metadata']['texts'][:80] if r.get('metadata') and r['metadata'].get('texts') else "None"
print(f" Distance: {r['distance']:.4f} | {snippet}...")
print(" ✅ FaissVectorStore works correctly!")
else:
print(" ⚠️ Vector store index is empty")
# Cleanup
import shutil
if os.path.exists(test_store_dir):
shutil.rmtree(test_store_dir)
print("\n" + "=" * 60)
print("🎉 ALL TESTS PASSED! Pipeline is working correctly.")
print("=" * 60)
print("\nYou can now run: python main.py")
if __name__ == "__main__":
test_pipeline()
|