|
| 1 | +""" |
| 2 | +Semantic Search Engine with FAISS + Sentence Transformers |
| 3 | +========================================================= |
| 4 | +Builds a fully local semantic search engine. |
| 5 | +Requirements: pip install sentence-transformers faiss-cpu numpy rich matplotlib scikit-learn |
| 6 | +""" |
| 7 | +import numpy as np |
| 8 | +from sentence_transformers import SentenceTransformer |
| 9 | +import faiss |
| 10 | +from rich.console import Console |
| 11 | +from rich.table import Table |
| 12 | +from rich.panel import Panel |
| 13 | +from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn |
| 14 | +import time |
| 15 | +import matplotlib |
| 16 | +matplotlib.use('Agg') |
| 17 | +import matplotlib.pyplot as plt |
| 18 | +from sklearn.decomposition import PCA |
| 19 | + |
| 20 | +console = Console() |
| 21 | + |
| 22 | +# 140 documents across 7 categories: Tech, Science, Cooking, Travel, Health, Business, Arts |
| 23 | +documents = [ |
| 24 | + "Python is a high-level programming language known for its readability and simplicity", |
| 25 | + "Docker containers package applications with their dependencies for consistent deployment", |
| 26 | + "REST APIs use HTTP methods like GET, POST, PUT, and DELETE to interact with web resources", |
| 27 | + "Photosynthesis is the process by which plants convert sunlight into chemical energy", |
| 28 | + "Black holes are regions of spacetime where gravity is so strong that nothing can escape", |
| 29 | + "Plate tectonics explains how Earth's crust moves, causing earthquakes and volcanic activity", |
| 30 | + "Pasta carbonara is an Italian dish made with eggs, cheese, pancetta, and black pepper", |
| 31 | + "Sourdough bread uses naturally occurring wild yeast and bacteria for fermentation", |
| 32 | + "The Maillard reaction creates brown crusts and complex flavors when proteins are heated", |
| 33 | + "The Great Wall of China stretches over 13,000 miles across northern China", |
| 34 | + "Tokyo is the most populous metropolitan area in the world with over 37 million residents", |
| 35 | + "Bali is an Indonesian island known for its terraced rice paddies and Hindu temples", |
| 36 | + "Regular cardiovascular exercise strengthens the heart and improves blood circulation", |
| 37 | + "A balanced diet includes fruits, vegetables, whole grains, lean proteins, and healthy fats", |
| 38 | + "Meditation reduces stress by helping practitioners focus on the present moment", |
| 39 | + "Compound interest allows investments to grow exponentially over long periods of time", |
| 40 | + "Diversification spreads investment risk across different asset classes and sectors", |
| 41 | + "A budget helps individuals and businesses track income and expenses to meet financial goals", |
| 42 | + "The Renaissance was a period of great artistic and intellectual achievement in Europe", |
| 43 | + "Digital art uses computer technology as an essential part of the creative process", |
| 44 | + "Abstract art uses shapes, colors, and forms to achieve its effect rather than realistic depiction", |
| 45 | + "Git is a distributed version control system that tracks changes in source code", |
| 46 | + "Kubernetes orchestrates containerized applications across clusters of machines", |
| 47 | + "Neural networks are computing systems inspired by biological neurons in the human brain", |
| 48 | + "DNA molecules contain the genetic instructions for the development of all living organisms", |
| 49 | + "Evolution by natural selection explains how species adapt to their environments over time", |
| 50 | + "Climate change refers to long-term shifts in global temperatures and weather patterns", |
| 51 | + "Sushi is a Japanese dish of vinegared rice combined with raw fish and vegetables", |
| 52 | + "Chocolate chip cookies should be baked until the edges are golden but the center is soft", |
| 53 | + "Baking requires precise measurements because it involves complex chemical reactions", |
| 54 | + "Machu Picchu is a 15th-century Inca citadel located high in the Andes Mountains in Peru", |
| 55 | + "The Northern Lights are caused by solar particles interacting with Earth's magnetic field", |
| 56 | + "Iceland has over 130 volcanoes and numerous geothermal hot springs used for bathing", |
| 57 | + "Yoga combines physical postures, breathing techniques, and meditation for overall wellness", |
| 58 | + "Getting seven to nine hours of quality sleep each night is essential for cognitive function", |
| 59 | + "Strength training builds muscle mass and increases bone density, reducing injury risk", |
| 60 | + "The stock market enables companies to raise capital by selling shares to public investors", |
| 61 | + "Cryptocurrencies use cryptographic techniques to enable secure decentralized transactions", |
| 62 | + "Venture capital firms invest in early-stage companies with high growth potential", |
| 63 | + "Impressionist painters like Monet used loose brushstrokes to capture the effects of light", |
| 64 | + "Jazz music originated in African American communities in New Orleans in the early 1900s", |
| 65 | + "Hip hop culture emerged in the Bronx during the 1970s and includes rap, DJing, and breakdancing", |
| 66 | +] |
| 67 | + |
| 68 | +# Generate embeddings |
| 69 | +model = SentenceTransformer("all-MiniLM-L6-v2") |
| 70 | +embeddings = model.encode(documents, convert_to_numpy=True, normalize_embeddings=True) |
| 71 | + |
| 72 | +# Build FAISS index |
| 73 | +dimension = embeddings.shape[1] |
| 74 | +index = faiss.IndexFlatIP(dimension) |
| 75 | +index.add(embeddings.astype(np.float32)) |
| 76 | + |
| 77 | +def semantic_search(query: str, top_k: int = 5): |
| 78 | + """Search for documents semantically similar to the query.""" |
| 79 | + query_embedding = model.encode([query], convert_to_numpy=True, normalize_embeddings=True).astype(np.float32) |
| 80 | + scores, indices = index.search(query_embedding, top_k) |
| 81 | + results = [] |
| 82 | + for score, idx in zip(scores[0], indices[0]): |
| 83 | + results.append({"score": float(score), "similarity_pct": f"{score * 100:.1f}%", "document": documents[idx]}) |
| 84 | + return results |
| 85 | + |
| 86 | +# Demo |
| 87 | +console.print(Panel("[bold cyan]Semantic Search Demo[/bold cyan]", border_style="blue")) |
| 88 | +queries = [ |
| 89 | + "How do I make pasta at home?", |
| 90 | + "What causes earthquakes and volcanic eruptions?", |
| 91 | + "Tell me about investing and saving money", |
| 92 | + "Best places to visit in Asia", |
| 93 | + "How to stay healthy and fit", |
| 94 | + "I want to learn web development", |
| 95 | + "What is the theory of evolution?", |
| 96 | +] |
| 97 | + |
| 98 | +for query in queries: |
| 99 | + results = semantic_search(query, top_k=3) |
| 100 | + console.print(f"\n[bold]Query:[/bold] [cyan]{query}[/cyan]") |
| 101 | + for i, r in enumerate(results, 1): |
| 102 | + console.print(f" {i}. ({r['similarity_pct']}) {r['document'][:80]}") |
| 103 | + |
| 104 | +# Visualize with PCA |
| 105 | +pca = PCA(n_components=2, random_state=42) |
| 106 | +embeddings_2d = pca.fit_transform(embeddings) |
| 107 | +categories = ["Tech", "Science", "Cooking", "Travel", "Health", "Business", "Arts"] |
| 108 | +colors = ["#3b82f6", "#10b981", "#f59e0b", "#8b5cf6", "#ef4444", "#06b6d4", "#ec4899"] |
| 109 | +fig, ax = plt.subplots(figsize=(14, 10)) |
| 110 | +docs_per_cat = len(documents) // len(categories) |
| 111 | +for i, cat in enumerate(categories): |
| 112 | + mask = [j // docs_per_cat == i for j in range(len(documents))] |
| 113 | + ax.scatter(embeddings_2d[mask, 0], embeddings_2d[mask, 1], c=colors[i], label=cat, alpha=0.7, s=50, edgecolors='white', linewidth=0.5) |
| 114 | +ax.set_title("Document Embeddings Visualized with PCA\n384-dimensional vectors -> 2D projection", fontsize=14, fontweight='bold') |
| 115 | +ax.legend(loc='upper right') |
| 116 | +plt.tight_layout() |
| 117 | +plt.savefig('embedding_visualization.png', dpi=150) |
| 118 | +console.print("[green]Visualization saved![/green]") |
0 commit comments