Re-Ranking for Quality

pip install sentence-transformers cohere numpy

Here is the broken experience:

import numpy as np
from sentence_transformers import SentenceTransformer

# 1. The Setup: A simple knowledge base
documents = [
    "The sun is a star found at the center of the solar system.",
    "Rain boots are recommended when walking outside in a storm.", # The answer we want
    "Solar panels generate energy from the sun.",
    "The weather forecast predicts clear skies today.",
    "Umbrellas are useful for shielding against rain.", # Also a good answer
    "Wind turbines create electricity from air currents.",
    "Storm chasers drive vehicles to track tornadoes.",
    "Water is essential for all known forms of life."
]

# 2. The Query
query = "What should I wear for a storm?"

# 3. The "Old Way" - Vector Search (Bi-Encoder)
# We load a small model to simulate vector search
model = SentenceTransformer('all-MiniLM-L6-v2')

# Encode docs and query
doc_embeddings = model.encode(documents)
query_embedding = model.encode(query)

# Calculate similarity (Dot Product)
scores = np.dot(doc_embeddings, query_embedding)

# Sort results (highest score first)
# We use argsort, then reverse it to get descending order
ranked_indices = np.argsort(scores)[::-1]

print(f"Query: {query}\n")
print("--- Vector Search Results (Top 3 passed to LLM) ---")
for i in range(3):
    idx = ranked_indices[i]
    print(f"Rank #{i+1}: {documents[idx]} (Score: {scores[idx]:.4f})")

print("\n--- The Missed Opportunity ---")
# Let's look further down the list
for i in range(3, 6):
    idx = ranked_indices[i]
    print(f"Rank #{i+1}: {documents[idx]} (Score: {scores[idx]:.4f})")

Run this code. Depending on the specific model nuances, you will often see "Storm chasers drive vehicles..." or "The weather forecast..." ranked very high because they share the word "storm" or "weather."

Meanwhile, "Rain boots are recommended..." might slip to rank #4 or #5.

If your RAG pipeline only sends the top 3 chunks to the LLM, the LLM will say: "I don't know what you should wear, but I know storm chasers drive vehicles."

That is a bad user experience. We need to fix the ranking.

---

Let's Build It

We will implement a Two-Stage Retrieval system.

Step 1: Get Your API Key

We will use Cohere because they have a specialized endpoint just for this.

Step 2: Define the Data and "Wide Net" Retrieval

We will reuse the data from the problem section, but let's assume we have already run the vector search and retrieved a large candidate list (e.g., top 10).

In a real app, this list comes from your Vector Database (Chroma, Pinecone, etc). Here, we will simulate that we retrieved 10 mixed-quality results.

import cohere

# Initialize the client
# REPLACE WITH YOUR ACTUAL API KEY
co = cohere.Client('YOUR_COHERE_API_KEY')

# The user query
query = "What should I wear for a storm?"

# The "Wide Net" - Imagine we retrieved these 10 from our Vector DB
# Note: I've mixed relevant answers with irrelevant ones that share keywords
retrieved_docs = [
    "Storm chasers drive vehicles to track tornadoes.", # Keyword match "storm"
    "The weather forecast predicts clear skies today.", # Semantic match "weather"
    "Solar panels generate energy from the sun.",       # Irrelevant
    "Rain boots are recommended when walking outside in a storm.", # THE ANSWER
    "Wind turbines create electricity from air currents.", # Irrelevant
    "Umbrellas are useful for shielding against rain.", # Relevant
    "Thunderstorms can cause power outages.",           # Context, but not the answer
    "Cotton clothes are breathable in summer.",         # Clothing, but wrong context
    "Winter coats are heavy and warm.",                 # Clothing, wrong context
    "Lightning strikes are dangerous during a storm."   # Context
]

print(f"Stage 1: Retrieved {len(retrieved_docs)} documents from Vector DB.")

Step 3: The Rerank

Unlike a Bi-Encoder (which compares two frozen vectors), a Cross-Encoder takes the query and the document together as a single input and outputs a relevance score. It "reads" the query in the context of the document. It is much slower, which is why we only run it on 10 or 20 items, not the whole database.

# Call Cohere Rerank
# model="rerank-english-v3.0" is optimized for English queries
results = co.rerank(
    query=query,
    documents=retrieved_docs,
    top_n=3, # We only want the absolute best 3 for our LLM
    model="rerank-english-v3.0"
)

print("\nStage 2: Reranking complete.")

Step 4: Display the Results

print(f"\nQuery: {query}")
print("-" * 50)

for idx, result in enumerate(results.results):
    # result.index is the index in the original list
    # result.relevance_score is how confident the model is (0 to 1)
    
    original_doc = retrieved_docs[result.index]
    score = result.relevance_score
    
    print(f"Rank #{idx + 1} (Score: {score:.4f})")
    print(f"Content: {original_doc}")
    print("-" * 50)

You should see a dramatic improvement. "Rain boots..." and "Umbrellas..." should jump to the top, even if they weren't at the top of the original list. Notice the scores—they are calibrated relevance probabilities. A score of 0.99 means "I am almost certain this answers the question."

Step 5: Filtering by Score

Sometimes, even the best result in your list is bad. If you search for "How to bake a cake" in our weather dataset, the "best" match might be "Solar panels..." simply because it's the least unlike the query.

We should add a threshold. If the best score is below 0.5, we shouldn't send it to the LLM.

print("\n--- Filtering Bad Matches ---")

bad_query = "How do I bake a cake?"

bad_results = co.rerank(
    query=bad_query,
    documents=retrieved_docs,
    top_n=3,
    model="rerank-english-v3.0"
)

for result in bad_results.results:
    if result.relevance_score < 0.5:
        print(f"Skipping document (Score {result.relevance_score:.4f} is too low)")
    else:
        print(f"Keeping document: {retrieved_docs[result.index]}")

This prevents hallucinations. If the retrieval is bad, it's better to say "I don't know" than to answer using irrelevant data.

---

Now You Try

You have the basic logic. Now let's push it further.

In the real world, your documents aren't just strings; they are objects with titles, dates, and URLs.

Create a list of dictionaries:

    docs = [
        {"text": "Rain boots...", "source": "manual.pdf"},
        {"text": "Solar panels...", "source": "science_book.txt"}
    ]
    

    import time
    start = time.time()
    # ... rerank call ...
    end = time.time()
    print(f"Reranking took {end - start} seconds")
    

Compare this speed to the simple vector search from the "Problem" section. Notice the difference? This is why we don't rerank the whole database!

---

Challenge Project: The Accuracy Benchmark

Your manager wants to know if paying for Cohere Rerank is worth it. You need to prove it with data.

Create a script that compares "First Result Accuracy" between standard Vector Search and Reranked Search.

Query: "What powers a computer?"
Correct Answer: "Electricity flows through circuits."

Vector Search picked: "Electric eels swim in water." (Incorrect)
Rerank Search picked: "Electricity flows through circuits." (Correct)

Score: Vector 0/1, Rerank 1/1

* You will need to manually define what the "Correct Answer" string is for your script to check against.

---

What You Learned

Today you moved from "Basic RAG" to "Production RAG."

* Bi-Encoders (Vector Search): Fast, cheap, but sometimes shallow. Good for finding the "neighborhood" of the answer.

* Cross-Encoders (Reranking): Slower, more expensive, but deeply intelligent. Good for finding the exact address.