Re-ranking Search Results with Cohere for Better RAG

Retrieval-Augmented Generation (RAG) has become the standard architecture for building LLM-powered applications that rely on private or up-to-date data. However, the quality of a RAG system is heavily dependent on the relevance of the context retrieved from the database. If the search engine feeds irrelevant documents to the LLM, the generated response will be inaccurate, hallucinated, or incomplete. This is where re-ranking comes into play.

What is Re-ranking?

In a standard RAG pipeline, retrieval is usually handled by a vector database using dense embeddings. This is a "bi-encoder" approach: the query and the documents are embedded independently, and their similarity is calculated using cosine similarity. While this is extremely fast and scales well to millions of documents, it lacks deep semantic understanding because the query and document never interact during the embedding process.

Re-ranking introduces a second, more computationally intensive stage using a "cross-encoder." A cross-encoder takes the query and a candidate document together as a single input, allowing the model to deeply analyze the relationship between the two. It outputs a highly accurate relevance score. Because cross-encoders are slower, you only run them on a small subset of documents (e.g., the top 20 retrieved by the vector database) to find the absolute best matches.

Why Re-ranking Matters for RAG

Implementing a re-ranking step provides several critical benefits for RAG applications:

Improved Accuracy: By using a cross-encoder, the system understands nuanced relationships between the query and the document, surfacing the most relevant context.
Reduced Hallucinations: When the LLM is provided with highly relevant context, it is less likely to guess or hallucinate information.
Token Efficiency: Re-ranking allows you to pass only the top 3-5 most relevant chunks to the LLM, saving token costs and reducing the chance of the LLM getting confused by "lost in the middle" phenomena where it ignores context placed in the middle of a large prompt.

Why Choose Cohere for Re-ranking?

Cohere offers a dedicated Rerank API that is widely considered the industry standard for improving search relevance. Unlike general-purpose LLMs, Cohere's rerank models (like rerank-english-v3.0) are specifically fine-tuned for the task of scoring document relevance. They are highly optimized, fast, and support features like multilingual re-ranking and document truncation, making them incredibly easy to drop into an existing RAG pipeline.

How to Use Cohere Rerank in Your RAG Pipeline

Integrating Cohere Rerank involves a simple two-step retrieval process. First, you retrieve a broad set of documents using your vector database. Second, you pass those documents and the user's query to the Cohere Rerank API to get the refined, top results.

Prerequisites

To follow along with the code examples, you will need a Cohere API key and the Cohere Python SDK. You can install the SDK using pip:

pip install cohere

Step 1: Initial Document Retrieval

For this example, we will simulate the initial retrieval step. In a real application, this would be a query to a vector database like Pinecone, Weaviate, or Milvus. We will retrieve a broad set of documents—say, the top 10 results—knowing that some might be only loosely related to the query.

import cohere

# Initialize the Cohere client
co = cohere.Client("YOUR_COHERE_API_KEY")

# The user's question
user_query = "How does Cohere improve RAG pipelines?"

# Simulated initial retrieval from a vector database (Top 10 results)
initial_documents = [
    {"id": "doc1", "text": "Cohere is an AI company specializing in natural language processing."},
    {"id": "doc2", "text": "RAG stands for Retrieval-Augmented Generation, a technique to ground LLMs."},
    {"id": "doc3", "text": "Vector databases store embeddings for fast similarity search."},
    {"id": "doc4", "text": "Cross-encoders are more accurate than bi-encoders but are computationally heavier."},
    {"id": "doc5", "text": "Cohere's rerank model significantly improves search relevance by scoring query-document pairs."},
    {"id": "doc6", "text": "Python is a popular programming language for machine learning."},
    {"id": "doc7", "text": "Chunking strategies affect the performance of RAG systems."},
    {"id": "doc8", "text": "Cohere offers generation, embedding, and rerank APIs for developers."},
    {"id": "doc9", "text": "The transformer architecture was introduced in the 'Attention is All You Need' paper."},
    {"id": "doc10", "text": "Using a reranker reduces hallucinations by providing better context to the LLM."}
]

print(f"Retrieved {len(initial_documents)} initial documents.")

Step 2: Re-ranking with Cohere

Now we will pass these 10 documents to the Cohere Rerank API. We will ask the API to return the top 3 most relevant documents. The API will use a cross-encoder to evaluate the query against each document and return them sorted by relevance score.

# Extract just the text from the documents for the API call
docs_text = [doc["text"] for doc in initial_documents]

# Call the Cohere Rerank API
response = co.rerank(
    model="rerank-english-v3.0",
    query=user_query,
    documents=docs_text,
    top_n=3
)

# Reconstruct the top documents based on the rerank results
top_docs = []
for result in response.results:
    # result.index gives us the position in the original list
    top_docs.append(initial_documents[result.index])

print("Top documents after re-ranking:")
for i, doc in enumerate(top_docs):
    print(f"{i+1}. {doc['text']}")

Step 3: Integrating with the LLM (Generation)

Once you have your highly relevant, re-ranked documents, you can construct the final prompt to send to your LLM (whether it's Cohere's Command model, OpenAI's GPT, or an open-source model). Because the context is now highly accurate, you can confidently use a smaller number of chunks.

# Construct the context string from the top re-ranked documents
context = "\n".join([doc["text"] for doc in top_docs])

# Create the prompt for the LLM
prompt = f"""
Context information is below.
---------------------
{context}
---------------------
Given the context information and not prior knowledge, answer the query.
Query: {user_query}
Answer:
"""

# Generate the response using Cohere's Command model
gen_response = co.chat(
    message=prompt,
    model="command-r-plus"
)

print("\nLLM Response:")
print(gen_response.text)

Best Practices for Re-ranking

To get the most out of Cohere Rerank in your RAG pipeline, consider the following best practices:

Retrieve Broadly, Re-rank Narrowly: A common pattern is to retrieve 20 to 50 documents from your vector database, and then re-rank them down to the top 3 to 5. This gives the cross-encoder a large enough pool to find the hidden gems that the bi-encoder might have ranked lower, while keeping LLM token costs low.
Use the Latest Models: Cohere frequently updates its models. Ensure you are using the latest version, such as rerank-english-v3.0 or rerank-multilingual-v3.0, to benefit from improved accuracy and speed.
Optimize Your Chunking: Re-ranking works best on coherent chunks of text. If your chunks are too small, they might lack context. If they are too large, the cross-encoder might struggle to pinpoint the relevant section. Aim for chunks of 200-500 words.
Handle Long Documents Properly: If you are re-ranking long documents, be aware of the model's token limits. Cohere's API allows you to specify a max_tokens_per_doc parameter, which can truncate documents to fit the model's constraints, ensuring the API call doesn't fail.
Cache Initial Retrieval Results: Since the initial vector search and the subsequent re-ranking are separate steps, you can cache the initial retrieval results for common queries and only run the re-ranker if the user's query changes slightly, saving compute and reducing latency.

Conclusion

Re-ranking is a transformative step for any production-grade RAG application. By bridging the gap between fast, approximate vector search and slow, highly accurate cross-encoding, you ensure that your LLM receives the best possible context. Cohere's Rerank API makes this two-stage retrieval process incredibly simple to implement. By retrieving broadly, re-ranking narrowly with Cohere, and adhering to best practices around chunking and model selection, developers can significantly reduce hallucinations, improve answer accuracy, and build highly trustworthy AI systems.