NLP & LLMs

Rag And Agents

Advanced RAG Techniques

Chapter 2's pipeline works, but has real limitations: pure semantic (embedding) similarity can miss exact keyword matches, the top-k retrieved chunks aren't alw

JrCodex·6 min read

Jr Codex NLP & LLM Notes

Level: Advanced Prerequisites: Chapter 2: Building a RAG Pipeline Time to complete: ~25 minutes


Table of Contents

  1. Where the Basic Pipeline Falls Short
  2. Hybrid Search
  3. Re-Ranking
  4. Query Transformation
  5. Contextual Chunk Retrieval
  6. Handling "I Don't Know" Correctly
  7. A More Robust RAG Pipeline
  8. Summary & Next Steps

1. Where the Basic Pipeline Falls Short

Chapter 2's pipeline works, but has real limitations: pure semantic (embedding) similarity can miss exact keyword matches, the top-k retrieved chunks aren't always genuinely the MOST relevant, and a poorly-phrased user query can retrieve irrelevant chunks entirely. This chapter covers techniques addressing each.


Hybrid search combines Chapter 2's semantic (embedding) search with keyword search (directly related to Module 1, Chapter 3's TF-IDF), since each catches different kinds of relevant matches.

Why NEITHER Approach Alone Is Sufficient
─────────────────────────────────────────
  Semantic search (Ch.2):     excellent at finding
                                  CONCEPTUALLY similar text, even
                                  with DIFFERENT wording — but
                                  can MISS an exact, specific
                                  term (a product SKU, an exact
                                  error code) that doesn't carry
                                  much "semantic" weight on its
                                  own
  Keyword search (TF-IDF,         excellent at finding EXACT
  Module 1, Ch.3):                   term matches — but MISSES
                                        conceptually related text
                                        using DIFFERENT words
                                        (Module 1, Ch.3, Section
                                        2's bag-of-words
                                        limitation)
─────────────────────────────────────────
from rank_bm25 import BM25Okapi      # a standard, refined keyword-search algorithm,
                                          # related to TF-IDF (Module 1, Ch.3)
 
def hybrid_search(query, chunks, chunk_embeddings, embedding_model, alpha=0.5, top_k=5):
    # Semantic scores (Chapter 2)
    query_embedding = embedding_model.encode([query])
    semantic_scores = (chunk_embeddings @ query_embedding.T).flatten()
 
    # Keyword scores (BM25 — an improved TF-IDF variant, Module 1, Ch.3)
    tokenized_chunks = [chunk.split() for chunk in chunks]
    bm25 = BM25Okapi(tokenized_chunks)
    keyword_scores = bm25.get_scores(query.split())
 
    # Normalize BOTH score sets to [0, 1] before COMBINING
    semantic_scores_norm = (semantic_scores - semantic_scores.min()) / (semantic_scores.max() - semantic_scores.min() + 1e-10)
    keyword_scores_norm = (keyword_scores - keyword_scores.min()) / (keyword_scores.max() - keyword_scores.min() + 1e-10)
 
    combined_scores = alpha * semantic_scores_norm + (1 - alpha) * keyword_scores_norm
 
    top_indices = combined_scores.argsort()[-top_k:][::-1]
    return [chunks[i] for i in top_indices]

3. Re-Ranking

Re-ranking adds a second, more computationally expensive but more accurate scoring pass over an initial set of candidates — directly analogous to DL Notes, Module 4, Chapter 4's two-stage object detection (propose candidates cheaply, THEN score them more carefully).

from sentence_transformers import CrossEncoder
 
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
 
def retrieve_with_reranking(query, chunks, chunk_embeddings, embedding_model, initial_k=20, final_k=3):
    # Stage 1: CHEAP, initial retrieval (Chapter 2) — cast a WIDE net
    query_embedding = embedding_model.encode([query])
    scores = (chunk_embeddings @ query_embedding.T).flatten()
    top_indices = scores.argsort()[-initial_k:][::-1]
    candidates = [chunks[i] for i in top_indices]
 
    # Stage 2: EXPENSIVE, precise re-ranking — narrow down to the BEST few
    pairs = [[query, candidate] for candidate in candidates]
    rerank_scores = reranker.predict(pairs)
 
    reranked_indices = rerank_scores.argsort()[-final_k:][::-1]
    return [candidates[i] for i in reranked_indices]
Why Re-Ranking Uses a DIFFERENT Model Than Initial Retrieval
─────────────────────────────────────────
  A CROSS-ENCODER (used for re-ranking) processes the QUERY
  and CANDIDATE together, through FULL self-attention (DL
  Notes, Module 6, Ch.2) between them — much MORE accurate at
  judging TRUE relevance, but far too SLOW to run against
  MILLIONS of chunks directly. A BI-ENCODER (used for initial
  retrieval, Chapter 2) embeds the query and EVERY chunk
  SEPARATELY, enabling fast vector search, at some cost to
  precision. Combining both — cheap bi-encoder retrieval, THEN
  expensive cross-encoder re-ranking on a SMALL candidate set
  — gets the BEST of both.
─────────────────────────────────────────

4. Query Transformation

Sometimes the user's RAW query isn't the best possible search query — query transformation uses the LLM itself to improve the query before retrieval.

def transform_query(client, original_query):
    prompt = f"""Rewrite the following question to be more specific and suitable
for searching a knowledge base. Return ONLY the rewritten question.
 
Original question: {original_query}"""
 
    response = client.chat.completions.create(
        model="gpt-4", messages=[{"role": "user", "content": prompt}],
    )
    return response.choices[0].message.content
 
# "What's up with the return thing?" -> "What is the company's product return policy?"
A Related Technique: HyDE (Hypothetical Document Embeddings)
─────────────────────────────────────────
  Instead of embedding the QUESTION directly, first ask the
  LLM to generate a HYPOTHETICAL ANSWER (which may not even
  be fully accurate), then embed THAT hypothetical answer and
  use IT for retrieval — since a plausible ANSWER often
  resembles the ACTUAL relevant document more closely (in
  embedding space) than the ORIGINAL question does.
─────────────────────────────────────────

5. Contextual Chunk Retrieval

Chapter 2, Section 2's fixed-size chunks lose surrounding context. One refinement: retrieve based on a SMALL chunk's embedding, but provide the LLM with a LARGER surrounding window of text.

def retrieve_with_expanded_context(query, chunks, chunk_embeddings, embedding_model, window=1):
    query_embedding = embedding_model.encode([query])
    scores = (chunk_embeddings @ query_embedding.T).flatten()
    best_idx = scores.argmax()
 
    # Retrieve the SURROUNDING chunks too, not just the single best match
    start = max(0, best_idx - window)
    end = min(len(chunks), best_idx + window + 1)
    return chunks[start:end]      # provides MORE context than the single matching chunk alone

6. Handling "I Don't Know" Correctly

A genuinely important, often-overlooked detail: the prompt (Chapter 2, Section 6) should explicitly instruct the model to say it doesn't know, rather than Chapter 1's hallucination filling the gap when retrieval fails to find genuinely relevant information.

robust_rag_prompt = """Answer the question using ONLY the context provided below.
 
IMPORTANT: If the context does not contain enough information to answer the
question, respond with "I don't have enough information to answer this
question" rather than guessing or using outside knowledge.
 
Context:
{context}
 
Question: {question}"""

This single instruction, tested and verified, directly determines whether a RAG system fails gracefully (admitting uncertainty) or fails dangerously (confidently hallucinating an answer) — a critical distinction for any production RAG deployment (Module 9).


7. A More Robust RAG Pipeline

def robust_rag_query(query, chunks, chunk_embeddings, embedding_model, reranker, llm_client):
    transformed_query = transform_query(llm_client, query)      # Section 4
 
    candidates = hybrid_search(transformed_query, chunks, chunk_embeddings, embedding_model, top_k=20)      # Section 2
    reranked = retrieve_with_reranking(transformed_query, candidates, ..., final_k=3)      # Section 3
 
    context = "\n\n".join(reranked)
 
    prompt = robust_rag_prompt.format(context=context, question=query)      # Section 6
    response = llm_client.chat.completions.create(model="gpt-4", messages=[{"role": "user", "content": prompt}])
    return response.choices[0].message.content

8. Summary & Next Steps

Key Takeaways

  • Hybrid search combines semantic embedding search with keyword-based search (BM25), since each catches different kinds of relevant matches.
  • Re-ranking uses an expensive but accurate cross-encoder to refine an initial, cheap set of candidates from a fast bi-encoder retrieval step.
  • Query transformation uses the LLM itself to rewrite vague queries into more effective search queries before retrieval.
  • Explicitly instructing the model to admit uncertainty when context is insufficient is essential for preventing hallucination from silently masking retrieval failures.

Concept Check

  1. Why does hybrid search combine semantic and keyword search rather than relying on either alone?
  2. Why is a cross-encoder used for re-ranking rather than for the initial retrieval step across all chunks?
  3. Why is explicitly prompting the model to say "I don't know" an important safeguard in a RAG system?

Next Chapter

Chapter 4: Agents & Tool Use


Jr Codex — 1-on-1 Personalized Coaching | Back to Module Index | Back to NLP & LLM Index