Vector Search con FAISS

Descripción

En esta cápsula implementarás el componente de Vector Search usando FAISS (Facebook AI Similarity Search), específicamente el algoritmo HNSW (Hierarchical Navigable Small World) para búsqueda eficiente en alta dimensionalidad.

FAISS permite búsqueda de vectores 100-1000x más rápida que brute-force cosine similarity, manteniendo ~95% de precisión. Es crítico para RAG production-ready donde tienes millones de chunks.

Al final tendrás un FAISSIndex production-ready con save/load, batch indexing, y optimización para dot product (equivalente a cosine con normalización).

Duración estimada: 40-50 minutos

---

Objetivos

Al completar esta cápsula, serás capaz de:

• ✅ Instalar y configurar FAISS
• ✅ Implementar FAISS HNSW index
• ✅ Normalizar embeddings para dot product ≈ cosine
• ✅ Hacer batch indexing eficiente
• ✅ Buscar top-K vecinos más cercanos
• ✅ Persistir índice (save/load)
• ✅ Optimizar parámetros HNSW

---

¿Por qué FAISS?

Problema: Brute-force es lento

# Brute-force cosine similarity
def search_brute_force(query_emb, all_embeddings, k=5):
    """O(n) - Calcula similaridad con TODOS los vectores"""
    similarities = []
    for emb in all_embeddings:  # 1M vectores = 1M operaciones
        sim = cosine_similarity(query_emb, emb)
        similarities.append(sim)
    
    # Ordenar y retornar top-K
    top_k_indices = np.argsort(similarities)[-k:]
    return top_k_indices

# Para 1M vectores: ~2-3 segundos por query ❌

Solución: FAISS HNSW

# FAISS HNSW: Approximate Nearest Neighbors
index.search(query_emb, k=5)
# Para 1M vectores: ~5-10ms por query ✅ (200-400x más rápido)

Trade-off:

• Brute-force: 100% precisión, lento
• FAISS HNSW: ~95-98% precisión, 100-1000x más rápido

---

Paso 1: Instalar FAISS

1.1: Instalación

# CPU version (más común)
pip install faiss-cpu

# GPU version (si tienes CUDA)
pip install faiss-gpu

Verificar instalación:

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

---

Paso 2: Implementar FAISSIndex

2.1: Clase base con HNSW

Crear src/search/faiss_index.py:

"""
FAISS Index - RAG System
Vector search con FAISS HNSW para búsqueda eficiente
"""

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


class FAISSIndex:
    """
    FAISS index con 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
    ):
        """
        Inicializar FAISS HNSW index
        
        Args:
            dim: Dimensionalidad de embeddings (1536 para OpenAI)
            m: Número de conexiones por layer (default: 32)
                - Más alto = mejor recall, más memoria
            ef_construction: Size de dynamic candidate list durante construcción
                - Más alto = mejor calidad, construcción más lenta
            ef_search: Size de dynamic candidate list durante búsqueda
                - Más alto = mejor recall, búsqueda más lenta
        """
        self.dim = dim
        self.m = m
        self.ef_construction = ef_construction
        self.ef_search = ef_search
        
        # Crear índice HNSW
        # IndexHNSWFlat: HNSW con flat (brute-force) en último layer
        self.index = faiss.IndexHNSWFlat(dim, m)
        
        # Configurar parámetros
        self.index.hnsw.efConstruction = ef_construction
        self.index.hnsw.efSearch = ef_search
        
        # Storage para chunks (metadata)
        self.chunks = []
        
        self.logger = logging.getLogger(__name__)
    
    def add(
        self,
        embeddings: np.ndarray,
        chunks: List,
        normalize: bool = True
    ) -> None:
        """
        Agregar embeddings al índice
        
        Args:
            embeddings: Array de embeddings (shape: [n, dim])
            chunks: Lista de Chunk objects o dicts
            normalize: Si True, normaliza embeddings (L2 norm)
        """
        # Validar dimensiones
        if embeddings.shape[1] != self.dim:
            raise ValueError(
                f"Embedding dim {embeddings.shape[1]} != index dim {self.dim}"
            )
        
        # Convertir a float32 (requerido por FAISS)
        embeddings = np.array(embeddings).astype('float32')
        
        # Normalizar para dot product ≈ cosine similarity
        if normalize:
            faiss.normalize_L2(embeddings)
        
        # Agregar al índice
        self.index.add(embeddings)
        
        # Guardar 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:
        """
        Agregar embeddings en batches (para datasets grandes)
        
        Args:
            embeddings: Array de embeddings
            chunks: Lista de chunks
            batch_size: Tamaño de batch
        """
        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]:
        """
        Buscar top-K vecinos más cercanos
        
        Args:
            query_embedding: Embedding del query (shape: [dim])
            k: Cantidad de resultados
            normalize: Si True, normaliza query
        
        Returns:
            Lista de dicts con chunk y score
        """
        # Validar que hay vectores indexados
        if self.index.ntotal == 0:
            raise ValueError("Index is empty. Add vectors first.")
        
        # Preparar query
        query = np.array([query_embedding]).astype('float32')
        
        if normalize:
            faiss.normalize_L2(query)
        
        # Buscar
        distances, indices = self.index.search(query, k)
        
        # Formatear resultados
        results = []
        for i, idx in enumerate(indices[0]):
            if idx == -1:  # FAISS retorna -1 si no encuentra suficientes vecinos
                break
            
            # Convertir distancia a score (dot product)
            # Con 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:
        """
        Guardar índice a disco
        
        Args:
            path: Directorio donde guardar
        """
        path_obj = Path(path)
        path_obj.mkdir(parents=True, exist_ok=True)
        
        # Guardar índice FAISS
        index_path = path_obj / "faiss.index"
        faiss.write_index(self.index, str(index_path))
        
        # Guardar chunks (metadata)
        chunks_path = path_obj / "chunks.pkl"
        with open(chunks_path, 'wb') as f:
            pickle.dump(self.chunks, f)
        
        # Guardar 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:
        """
        Cargar índice desde disco
        
        Args:
            path: Directorio donde está guardado
        """
        path_obj = Path(path)
        
        if not path_obj.exists():
            raise FileNotFoundError(f"Index path not found: {path}")
        
        # Cargar índice FAISS
        index_path = path_obj / "faiss.index"
        self.index = faiss.read_index(str(index_path))
        
        # Cargar chunks
        chunks_path = path_obj / "chunks.pkl"
        with open(chunks_path, 'rb') as f:
            self.chunks = pickle.load(f)
        
        # Cargar 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:
        """
        Obtener estadísticas del índice
        
        Returns:
            Dict con estadísticas
        """
        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:
        """
        Estimar tamaño del índice en MB
        
        Returns:
            Tamaño estimado en MB
        """
        # Estimación aproximada:
        # - Cada vector: dim * 4 bytes (float32)
        # - HNSW overhead: ~m * 8 bytes por 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)
    
    # Crear índice
    index = FAISSIndex(dim=1536, m=32)
    
    # Datos de ejemplo
    embeddings = np.random.rand(1000, 1536).astype('float32')
    chunks = [{'id': i, 'text': f"Chunk {i}"} for i in range(1000)]
    
    # Indexar
    index.add(embeddings, chunks)
    
    # Buscar
    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']}")
    
    # Guardar
    index.save("./faiss_index")
    
    # Cargar
    new_index = FAISSIndex(dim=1536)
    new_index.load("./faiss_index")
    
    print(f"\n✅ Loaded index with {new_index.index.ntotal} vectors")

---

Paso 3: Optimización de parámetros HNSW

3.1: Parámetros clave

# M (connections per layer)
# - Valor típico: 16-64
# - M=16: Menos memoria, recall ~90%
# - M=32: Balance (recomendado)
# - M=64: Más memoria, recall ~98%

# ef_construction (construction time)
# - Valor típico: 40-500
# - ef=40: Construcción rápida, recall ~90%
# - ef=200: Balance (recomendado)
# - ef=500: Construcción lenta, recall ~99%

# ef_search (search time)
# - Valor típico: 16-512
# - ef=16: Búsqueda muy rápida, recall ~85%
# - ef=64: Balance (recomendado)
# - ef=128: Búsqueda más lenta, recall ~95%

3.2: Benchmark de parámetros

def benchmark_parameters():
    """Benchmark diferentes configuraciones"""
    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)
        
        # Indexar
        start = time.time()
        index.add(embeddings, chunks)
        index_time = time.time() - start
        
        # Buscar
        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")

---

Paso 4: Integración con pipeline

4.1: Pipeline completo con FAISS

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

# Crear pipeline
pipeline = RAGPipeline(chunk_size=500, overlap=50)

# Procesar documentos
chunks, embeddings = pipeline.process_directory("./data/documents")

# Crear índice FAISS
index = FAISSIndex(dim=embeddings.shape[1], m=32)
index.add(embeddings, chunks)

# Guardar índice
index.save("./rag_index")

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

# Buscar
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

Problema 1: RuntimeError: Error in faiss::write_index

Causa: Path no existe o sin permisos

Solución:

from pathlib import Path

# Crear directorio si no existe
path = Path("./faiss_index")
path.mkdir(parents=True, exist_ok=True)

index.save(str(path))

---

Problema 2: Búsqueda retorna scores negativos

Causa: No normalizaste embeddings

Solución:

# SIEMPRE normalizar para dot product ≈ cosine
faiss.normalize_L2(embeddings)  # Durante indexing
faiss.normalize_L2(query)        # Durante búsqueda

---

Problema 3: Recall muy bajo (~70%)

Causa: Parámetros HNSW muy agresivos

Solución:

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

# O recrear con M más alto
index = FAISSIndex(dim=1536, m=64)  # Default: 32

---

Problema 4: Index muy grande (>1GB)

Causa: HNSW usa mucha memoria para M alto

Solución:

# Opción 1: Reducir M
index = FAISSIndex(dim=1536, m=16)  # Menos memoria

# Opción 2: Usar IVF (Inverted File Index) para datasets >1M
index = faiss.IndexIVFFlat(quantizer, dim, nlist)

---

Resumen

En esta cápsula implementaste:

• ✅ FAISSIndex con algoritmo HNSW
• ✅ L2 normalization para dot product ≈ cosine
• ✅ Batch indexing para datasets grandes
• ✅ Save/load persistente
• ✅ Optimización de parámetros (M, ef_construction, ef_search)
• ✅ Integración con pipeline RAG completo

Próxima cápsula: Evaluation Framework - Implementar métricas (nDCG, MRR, Recall@K) y A/B testing.

---

Recursos Adicionales

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

---

Módulo 8 - Cápsula 04