RAG-Assistant1 / vector_store.py
Rixhot's picture
Upload 5 files
afa96cc verified
Raw
History Blame Contribute Delete
1.92 kB
import faiss # Facebook AI Similarity Search
import numpy as np # for working with arrays of numbers
def build_index(chunks):
"""
Builds a FAISS index from all chunk embeddings.
Input : chunks β†’ list of chunk dicts, each with an "embedding" field
Output : a FAISS index ready for searching
"""
# Stack all embeddings into a 2D array
# Shape: (number_of_chunks, 384)
# astype("float32") because FAISS requires 32-bit floats
embeddings = np.array(
[chunk["embedding"] for chunk in chunks]
).astype("float32")
dimension = embeddings.shape[1] # size of each vector (384 for MiniLM)
index = faiss.IndexFlatL2(dimension) # create the index with L2 distance
index.add(embeddings) # add all chunk vectors to the index
print(f"FAISS index built with {index.ntotal} vectors.")
return index
def search_index(index, chunks, query_embedding, top_k=3):
"""
Searches the FAISS index for the most relevant chunks.
Input:
index β†’ the FAISS index we built
chunks β†’ our original list of chunk dicts
query_embedding β†’ the vector of the user question
top_k β†’ how many results to return (default: 3)
Output : list of the top_k most relevant chunk dicts
"""
# FAISS expects a 2D array even for a single query
# So we wrap it in another list and convert to float32
query_vec = np.array([query_embedding]).astype("float32")
# Search the index
# D β†’ distances (how far each result is from the query)
# I β†’ indices (which chunk numbers are the closest matches)
D, I = index.search(query_vec, top_k)
results = []
for idx in I[0]: # I[0] because we have just one query
if idx != -1: # FAISS returns -1 if no result was found
results.append(chunks[idx])
return results