Retrieve the right context

The R in RAG. Before the model answers, we find the handful of chunks most likely to contain the answer — so the prompt carries just the relevant facts, not all seven (or, in a real system, all seventy thousand). If you took the semantic-search course this is old friend: embed the question, score it against every chunk with cosine similarity, take the top few.

Add a retrieve function to rag.py:

import numpy as np

def retrieve(query, k=3):
  query_embedding = embedder.encode(query, normalize_embeddings=True)
  scores = doc_embeddings @ query_embedding      # cosine similarity, all chunks
  top_k = np.argsort(-scores)[:k]                # indices of the k best
  return [documents[i] for i in top_k]

Because the embeddings are normalized, that one @ (matrix–vector product) is cosine similarity against every chunk at once; np.argsort(-scores)[:k] grabs the highest few. Same retriever you've built before — here it's a component, not the whole app.

Try it

Ask the question from Chapter 2 and see what it pulls:

for chunk in retrieve("How much does the Pro plan cost?"):
  print("-", chunk)

python rag.py

- Nimbus Notes Pro costs $6 per month and adds unlimited notebooks, version history, and offline sync.
- The Nimbus Notes free plan includes 3 notebooks and 500 MB of storage.
- All paid Nimbus Notes plans come with a 30-day money-back guarantee.

The exact fact needed — "$6 per month" — comes back first, even though the question said "Pro plan" and never said "cost per month" the way the document does. Semantic retrieval matched on meaning, and the two runner-up chunks are plausibly relevant too.

That top-of-list chunk is the ground truth the model was missing in Chapter 2. Now we hand it over — and watch the hallucination disappear.