Reranking Search Results for Better Accuracy

Welcome back! In our previous lesson, we explored hybrid retrieval, which combines metadata and vector search to enhance the accuracy and relevance of search results. Today, we will focus on reranking, a technique that further refines search results by ordering them based on relevance. This lesson will build on your existing knowledge of vector search and introduce you to the practical implementation of reranking using cosine similarity.

Reranking is an essential step in semantic search systems. While initial search results may be relevant, reranking ensures that the most pertinent results appear at the top. This is particularly important in scenarios where users expect quick and accurate answers. By leveraging cosine similarity, we can measure the relevance of each document to the query and reorder the results accordingly. Let's dive into how this works.

As you already know, cosine similarity is a metric that measures the similarity between two vectors, with scores ranging from -1 (completely opposite) to 1 (identical). It helps us assess the relevance of documents to a query.

Now, let's implement reranking using cosine similarity. We'll use the rerank_results function to reorder search results based on their relevance to a query. Here's the code provided in src/main.py:

from sklearn.metrics.pairwise import cosine_similarity
from sentence_transformers import SentenceTransformer

def rerank_results(query_embedding, retrieved_docs, doc_embeddings):
    """
    Rerank results by computing cosine similarity and sorting by relevance.
    """
    scores = cosine_similarity([query_embedding], doc_embeddings)[0]
    reranked = sorted(zip(retrieved_docs, scores), key=lambda x: x[1], reverse=True)
    return reranked

In this function, we first compute the cosine similarity scores between the query embedding and each document embedding. These scores indicate how relevant each document is to the query. We then sort the documents based on these scores in descending order, ensuring that the most relevant documents appear first.

Let's walk through a practical example using the provided code. Suppose we have a query about "Latest AI trends" and a set of retrieved documents. We want to rerank these documents to ensure the most relevant ones are at the top.

# Load embedding model
model_name = "sentence-transformers/all-MiniLM-L6-v2"
model = SentenceTransformer(model_name)

# Example: Perform reranking
query_embedding = model.encode(["Latest AI trends"])[0]
retrieved_docs = ["AI in healthcare", "AI in finance", "AI in education"]
doc_embeddings = model.encode(retrieved_docs)

reranked_results = rerank_results(query_embedding, retrieved_docs, doc_embeddings)
for doc, score in reranked_results:
    print(f"Doc: {doc}, Score: {score:.4f}")

Output:

Doc: AI in healthcare, Score: 0.5153
Doc: AI in finance, Score: 0.4803
Doc: AI in education, Score: 0.4689

In this example, we first load the embedding model and encode the query and documents into embeddings. We then call the rerank_results function to reorder the documents based on their cosine similarity scores. The output displays each document along with its relevance score, allowing us to see which documents are most pertinent to the query.