Vector Search with FAISS

Description

In this capsule you'll implement the Vector Search component using FAISS (Facebook AI Similarity Search), specifically the HNSW (Hierarchical Navigable Small World) algorithm for efficient high-dimensional search.

FAISS enables vector search 100-1000x faster than brute-force cosine similarity, while keeping ~95% precision. It's critical for production-ready RAG where you have millions of chunks.

By the end you'll have a production-ready FAISSIndex with save/load, batch indexing, and optimization for dot product (equivalent to cosine with normalization).

Estimated duration: 40-50 minutes

---

Objectives

By completing this capsule, you'll be able to:

• ✅ Install and configure FAISS
• ✅ Implement a FAISS HNSW index
• ✅ Normalize embeddings for dot product ≈ cosine
• ✅ Do efficient batch indexing
• ✅ Search for the top-K nearest neighbors
• ✅ Persist the index (save/load)
• ✅ Optimize HNSW parameters

---

Why FAISS?

Problem: Brute-force is slow

# Brute-force cosine similarity
def search_brute_force(query_emb, all_embeddings, k=5):
    """O(n) - Computes similarity with ALL vectors"""
    similarities = []
    for emb in all_embeddings:  # 1M vectors = 1M operations
        sim = cosine_similarity(query_emb, emb)
        similarities.append(sim)
    
    # Sort and return top-K
    top_k_indices = np.argsort(similarities)[-k:]
    return top_k_indices

# For 1M vectors: ~2-3 seconds per query ❌

Solution: FAISS HNSW

# FAISS HNSW: Approximate Nearest Neighbors
index.search(query_emb, k=5)
# For 1M vectors: ~5-10ms per query ✅ (200-400x faster)

Trade-off:

• Brute-force: 100% precision, slow
• FAISS HNSW: ~95-98% precision, 100-1000x faster

---

Step 1: Install FAISS

1.1: Installation

# CPU version (most common)
pip install faiss-cpu

# GPU version (if you have CUDA)
pip install faiss-gpu

Verify the installation:

import faiss
print(f"FAISS version: {faiss.__version__}")
# Output: FAISS version: 1.7.4

---

Step 2: Implement FAISSIndex

2.1: Base class with HNSW

Create src/search/faiss_index.py:

"""
FAISS Index - RAG System
Vector search with FAISS HNSW for efficient search
"""

import faiss
import numpy as np
import pickle
from pathlib import Path
from typing import List, Dict, Tuple
import logging


class FAISSIndex:
    """
    FAISS index with the HNSW algorithm
    
    Features:
    - HNSW (Hierarchical Navigable Small World)
    - L2 normalization (dot product ≈ cosine)
    - Batch indexing
    - Persistent storage (save/load)
    - Metadata tracking
    
    Example:
        index = FAISSIndex(dim=1536)
        index.add(embeddings, chunks)
        results = index.search(query_embedding, k=5)
    """
    
    def __init__(
        self,
        dim: int = 1536,
        m: int = 32,
        ef_construction: int = 200,
        ef_search: int = 64
    ):
        """
        Initialize the FAISS HNSW index
        
        Args:
            dim: Embedding dimensionality (1536 for OpenAI)
            m: Number of connections per layer (default: 32)
                - Higher = better recall, more memory
            ef_construction: Size of the dynamic candidate list during construction
                - Higher = better quality, slower construction
            ef_search: Size of the dynamic candidate list during search
                - Higher = better recall, slower search
        """
        self.dim = dim
        self.m = m
        self.ef_construction = ef_construction
        self.ef_search = ef_search
        
        # Create the HNSW index
        # IndexHNSWFlat: HNSW with flat (brute-force) at the last layer
        self.index = faiss.IndexHNSWFlat(dim, m)
        
        # Configure the parameters
        self.index.hnsw.efConstruction = ef_construction
        self.index.hnsw.efSearch = ef_search
        
        # Storage for chunks (metadata)
        self.chunks = []
        
        self.logger = logging.getLogger(__name__)
    
    def add(
        self,
        embeddings: np.ndarray,
        chunks: List,
        normalize: bool = True
    ) -> None:
        """
        Add embeddings to the index
        
        Args:
            embeddings: Array of embeddings (shape: [n, dim])
            chunks: List of Chunk objects or dicts
            normalize: If True, normalizes embeddings (L2 norm)
        """
        # Validate dimensions
        if embeddings.shape[1] != self.dim:
            raise ValueError(
                f"Embedding dim {embeddings.shape[1]} != index dim {self.dim}"
            )
        
        # Convert to float32 (required by FAISS)
        embeddings = np.array(embeddings).astype('float32')
        
        # Normalize for dot product ≈ cosine similarity
        if normalize:
            faiss.normalize_L2(embeddings)
        
        # Add to the index
        self.index.add(embeddings)
        
        # Save chunks (metadata)
        self.chunks.extend(chunks)
        
        self.logger.info(
            f"✅ Indexed {len(chunks)} chunks (total: {self.index.ntotal})"
        )
    
    def add_batch(
        self,
        embeddings: np.ndarray,
        chunks: List,
        batch_size: int = 1000
    ) -> None:
        """
        Add embeddings in batches (for large datasets)
        
        Args:
            embeddings: Array of embeddings
            chunks: List of chunks
            batch_size: Batch size
        """
        total = len(embeddings)
        
        for i in range(0, total, batch_size):
            end = min(i + batch_size, total)
            batch_embeddings = embeddings[i:end]
            batch_chunks = chunks[i:end]
            
            self.add(batch_embeddings, batch_chunks)
            
            self.logger.info(f"Batch {i//batch_size + 1}: {end}/{total} chunks")
    
    def search(
        self,
        query_embedding: np.ndarray,
        k: int = 5,
        normalize: bool = True
    ) -> List[Dict]:
        """
        Search for the top-K nearest neighbors
        
        Args:
            query_embedding: Query embedding (shape: [dim])
            k: Number of results
            normalize: If True, normalizes the query
        
        Returns:
            List of dicts with chunk and score
        """
        # Validate that there are indexed vectors
        if self.index.ntotal == 0:
            raise ValueError("Index is empty. Add vectors first.")
        
        # Prepare the query
        query = np.array([query_embedding]).astype('float32')
        
        if normalize:
            faiss.normalize_L2(query)
        
        # Search
        distances, indices = self.index.search(query, k)
        
        # Format the results
        results = []
        for i, idx in enumerate(indices[0]):
            if idx == -1:  # FAISS returns -1 if it doesn't find enough neighbors
                break
            
            # Convert the distance to a score (dot product)
            # With L2 normalization, dot product ∈ [0, 1]
            score = float(distances[0][i])
            
            chunk = self.chunks[idx]
            
            results.append({
                'chunk': chunk,
                'score': score,
                'index': int(idx)
            })
        
        return results
    
    def save(self, path: str) -> None:
        """
        Save the index to disk
        
        Args:
            path: Directory to save to
        """
        path_obj = Path(path)
        path_obj.mkdir(parents=True, exist_ok=True)
        
        # Save the FAISS index
        index_path = path_obj / "faiss.index"
        faiss.write_index(self.index, str(index_path))
        
        # Save chunks (metadata)
        chunks_path = path_obj / "chunks.pkl"
        with open(chunks_path, 'wb') as f:
            pickle.dump(self.chunks, f)
        
        # Save config
        config = {
            'dim': self.dim,
            'm': self.m,
            'ef_construction': self.ef_construction,
            'ef_search': self.ef_search,
            'total_vectors': self.index.ntotal
        }
        config_path = path_obj / "config.pkl"
        with open(config_path, 'wb') as f:
            pickle.dump(config, f)
        
        self.logger.info(f"✅ Saved index to {path}")
    
    def load(self, path: str) -> None:
        """
        Load the index from disk
        
        Args:
            path: Directory where it's saved
        """
        path_obj = Path(path)
        
        if not path_obj.exists():
            raise FileNotFoundError(f"Index path not found: {path}")
        
        # Load the FAISS index
        index_path = path_obj / "faiss.index"
        self.index = faiss.read_index(str(index_path))
        
        # Load chunks
        chunks_path = path_obj / "chunks.pkl"
        with open(chunks_path, 'rb') as f:
            self.chunks = pickle.load(f)
        
        # Load config
        config_path = path_obj / "config.pkl"
        if config_path.exists():
            with open(config_path, 'rb') as f:
                config = pickle.load(f)
            
            self.dim = config['dim']
            self.m = config['m']
            self.ef_construction = config['ef_construction']
            self.ef_search = config['ef_search']
        
        self.logger.info(
            f"✅ Loaded index from {path} ({self.index.ntotal} vectors)"
        )
    
    def get_stats(self) -> Dict:
        """
        Get index statistics
        
        Returns:
            Dict with statistics
        """
        return {
            'total_vectors': self.index.ntotal,
            'dim': self.dim,
            'm': self.m,
            'ef_construction': self.ef_construction,
            'ef_search': self.ef_search,
            'index_size_mb': self.estimate_size_mb()
        }
    
    def estimate_size_mb(self) -> float:
        """
        Estimate the index size in MB
        
        Returns:
            Estimated size in MB
        """
        # Rough estimate:
        # - Each vector: dim * 4 bytes (float32)
        # - HNSW overhead: ~m * 8 bytes per vector
        vector_size = self.dim * 4
        hnsw_overhead = self.m * 8
        total_bytes = self.index.ntotal * (vector_size + hnsw_overhead)
        
        return total_bytes / (1024 ** 2)


# Demo
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    # Create the index
    index = FAISSIndex(dim=1536, m=32)
    
    # Example data
    embeddings = np.random.rand(1000, 1536).astype('float32')
    chunks = [{'id': i, 'text': f"Chunk {i}"} for i in range(1000)]
    
    # Index
    index.add(embeddings, chunks)
    
    # Search
    query_emb = np.random.rand(1536).astype('float32')
    results = index.search(query_emb, k=5)
    
    print(f"\n✅ Found {len(results)} results")
    for r in results:
        print(f"  Score: {r['score']:.4f} - {r['chunk']['text']}")
    
    # Save
    index.save("./faiss_index")
    
    # Load
    new_index = FAISSIndex(dim=1536)
    new_index.load("./faiss_index")
    
    print(f"\n✅ Loaded index with {new_index.index.ntotal} vectors")

---

Step 3: HNSW parameter optimization

3.1: Key parameters

# M (connections per layer)
# - Typical value: 16-64
# - M=16: Less memory, recall ~90%
# - M=32: Balance (recommended)
# - M=64: More memory, recall ~98%

# ef_construction (construction time)
# - Typical value: 40-500
# - ef=40: Fast construction, recall ~90%
# - ef=200: Balance (recommended)
# - ef=500: Slow construction, recall ~99%

# ef_search (search time)
# - Typical value: 16-512
# - ef=16: Very fast search, recall ~85%
# - ef=64: Balance (recommended)
# - ef=128: Slower search, recall ~95%

3.2: Parameter benchmark

def benchmark_parameters():
    """Benchmark different configurations"""
    configs = [
        {'m': 16, 'ef_construction': 40, 'ef_search': 16},   # Fast
        {'m': 32, 'ef_construction': 200, 'ef_search': 64},  # Balanced
        {'m': 64, 'ef_construction': 500, 'ef_search': 128}, # Accurate
    ]
    
    for config in configs:
        index = FAISSIndex(**config)
        
        # Index
        start = time.time()
        index.add(embeddings, chunks)
        index_time = time.time() - start
        
        # Search
        start = time.time()
        for query in test_queries:
            results = index.search(query, k=10)
        search_time = (time.time() - start) / len(test_queries)
        
        print(f"\nConfig: M={config['m']}, ef_c={config['ef_construction']}, ef_s={config['ef_search']}")
        print(f"  Index time: {index_time:.2f}s")
        print(f"  Search time: {search_time*1000:.2f}ms/query")
        print(f"  Index size: {index.estimate_size_mb():.1f}MB")

---

Step 4: Integration with the pipeline

4.1: Complete pipeline with FAISS

from src.pipeline.rag_pipeline import RAGPipeline
from src.search.faiss_index import FAISSIndex

# Create the pipeline
pipeline = RAGPipeline(chunk_size=500, overlap=50)

# Process the documents
chunks, embeddings = pipeline.process_directory("./data/documents")

# Create the FAISS index
index = FAISSIndex(dim=embeddings.shape[1], m=32)
index.add(embeddings, chunks)

# Save the index
index.save("./rag_index")

print(f"✅ Indexed {len(chunks)} chunks")
print(f"Index size: {index.estimate_size_mb():.1f}MB")

# Search
query = "How to install Python?"
query_emb = pipeline.embedder.embed(query)
results = index.search(query_emb, k=5)

for i, r in enumerate(results, 1):
    print(f"\n{i}. Score: {r['score']:.4f}")
    print(f"   {r['chunk'].text[:200]}...")

---

Troubleshooting

Problem 1: RuntimeError: Error in faiss::write_index

Cause: The path doesn't exist or lacks permissions

Solution:

from pathlib import Path

# Create the directory if it doesn't exist
path = Path("./faiss_index")
path.mkdir(parents=True, exist_ok=True)

index.save(str(path))

---

Problem 2: Search returns negative scores

Cause: You didn't normalize the embeddings

Solution:

# ALWAYS normalize for dot product ≈ cosine
faiss.normalize_L2(embeddings)  # During indexing
faiss.normalize_L2(query)        # During search

---

Problem 3: Very low recall (~70%)

Cause: HNSW parameters that are too aggressive

Solution:

# Increase ef_search
index.index.hnsw.efSearch = 128  # Default: 64

# Or recreate with a higher M
index = FAISSIndex(dim=1536, m=64)  # Default: 32

---

Problem 4: Index too large (>1GB)

Cause: HNSW uses a lot of memory for high M

Solution:

# Option 1: Reduce M
index = FAISSIndex(dim=1536, m=16)  # Less memory

# Option 2: Use IVF (Inverted File Index) for datasets >1M
index = faiss.IndexIVFFlat(quantizer, dim, nlist)

---

Summary

In this capsule you implemented:

• ✅ FAISSIndex with the HNSW algorithm
• ✅ L2 normalization for dot product ≈ cosine
• ✅ Batch indexing for large datasets
• ✅ Persistent save/load
• ✅ Parameter optimization (M, ef_construction, ef_search)
• ✅ Integration with the complete RAG pipeline

Next capsule: Evaluation Framework - Implement metrics (nDCG, MRR, Recall@K) and A/B testing.

---

Additional Resources

1. FAISS Documentation - Official wiki
2. FAISS Tutorial - Pinecone guide
3. HNSW Paper - Original algorithm paper
4. FAISS Benchmarks - Performance comparisons
5. ANN Benchmarks - Comprehensive ANN comparison
6. FAISS Best Practices - Index selection guide
7. Vector Search at Scale - DeepLearning.AI course

---

Module 8 - Capsule 04