Index FAISS pour la recherche rapide

Quand on a des milliers de documents, calculer la similarite avec chacun devient lent. FAISS (Facebook AI Similarity Search) est une librairie optimisee pour la recherche de vecteurs similaires, meme avec des millions d'entrées.

Le principe :
1. Créer un index FAISS avec la dimension des vecteurs
2. Ajouter les vecteurs des documents dans l'index
3. Chercher les vecteurs les plus proches d'un vecteur query

import faiss
import numpy as np
dimension = 128
index = faiss.IndexFlatL2(dimension) # L2 = distance euclidienne
vectors = np.random.random((10, dimension)).astype('float32')
index.add(vectors) # Ajoute 10 vecteurs
distances, indices = index.search(query_vector, k=3) # Top 3

Attention : FAISS travaille avec des float32, pas des float64.

Écris une fonction build_faiss_index(vectors) qui prend un numpy array 2D de vecteurs, crée un index FAISS IndexFlatL2 et y ajoute les vecteurs. Renvoie l'index.

Écris aussi une fonction faiss_search(index, query_vector, k) qui cherche les k vecteurs les plus proches et renvoie un tuple (distances, indices).

Les vecteurs doivent etre convertis en float32 avant d'etre ajoutes a l'index.

Tests (1/1)

Tests

import numpy as np
vectors = np.random.random((20, 64)).astype('float32')
index = build_faiss_index(vectors)
assert index.ntotal == 20, f'L index doit contenir 20 vecteurs, en a {index.ntotal}'
query = np.random.random((1, 64)).astype('float32')
distances, indices = faiss_search(index, query, 3)
assert indices.shape == (1, 3), f'Attendu shape (1,3), obtenu {indices.shape}'
assert distances.shape == (1, 3), f'Distances shape incorrecte: {distances.shape}'
assert all(0 <= i < 20 for i in indices[0]), 'Les indices doivent etre dans [0, 19]'
query_1d = np.random.random(64).astype('float32')
d2, i2 = faiss_search(index, query_1d, 2)
assert i2.shape[1] == 2, 'La recherche avec un vecteur 1D doit aussi marcher'

Indices (3 disponibles)

Solution officielle

import faiss
import numpy as np

def build_faiss_index(vectors):
    vectors = np.array(vectors, dtype='float32')
    dimension = vectors.shape[1]
    index = faiss.IndexFlatL2(dimension)
    index.add(vectors)
    return index

def faiss_search(index, query_vector, k):
    query = np.array(query_vector, dtype='float32')
    if query.ndim == 1:
        query = query.reshape(1, -1)
    distances, indices = index.search(query, k)
    return distances, indices

← Top-K recherche de documents Combiner TF-IDF et FAISS →

solution.py

Bravo!

Index FAISS pour la recherche rapide

Tests (1/1)

Indices (3 disponibles)

Solution officielle