Specialty · AILesson 6 of 6Free to read

Embeddings and retrieval · AI Foundations

How does an AI assistant find the right document to answer your question?

Why you need it

A language model only knows what it saw in training: not your company's refund policy, and not last week's changes. Retrieval-augmented generation (RAG) fixes that by looking up relevant documents first and putting them in the prompt. The lookup runs on embeddings: vectors that place texts with similar meaning close together. It's the same nearest-neighbour idea as lesson 4, with a different way to measure "close".

The intuition

An embedding model turns a text into a list of numbers, a vector. Texts with similar meaning get vectors that point in similar directions. In this lesson we use tiny made-up 3-number vectors so you can see what's happening; real embeddings come from a model, have many more numbers, and no single number has a readable meaning.

To compare directions, use cosine similarity:

cosine(a, b) = (a · b) / (|a| × |b|)

a · b (the dot product) multiplies matching numbers and adds them; |a| is the vector's length. The result is 1 for the same direction, 0 for unrelated, and −1 for opposite. Dividing by the lengths matters: a long document shouldn't win just because its vector is long.

RAG in three steps:

embed the question
rank the documents by cosine similarity and keep the top k, dropping any below a minimum score
paste them into the prompt: "Answer using these documents: … Question: …"

Definition

Embedding

A vector produced by a model so that texts with similar meaning are close together.

Cosine similarity

How closely two vectors point the same way: the dot product divided by both lengths.

RAG

Retrieval-augmented generation: find relevant documents and give them to the language model with the question.

Examples

Cosine similarity

Try it yourself

cat and kitten point the same way; car points elsewhere.

import math

def cosine(a, b):
    dot = sum(x * y for x, y in zip(a, b))
    return dot / (math.hypot(*a) * math.hypot(*b))

cat = [0.9, 0.8, 0.1]
kitten = [0.85, 0.75, 0.2]
car = [0.1, 0.2, 0.95]
print(round(cosine(cat, kitten), 3))
print(round(cosine(cat, car), 3))

Output

0.996
0.293

Step by step

L4Multiply matching numbers and add: the dot product.
L5The vector's length, √(x² + y² + z²).
L10Very close to 1: nearly the same meaning.
L11Much lower: unrelated.

Retrieve, then prompt

Try it yourself

Rank the documents for a question and build the prompt.

import math

def cosine(a, b):
    dot = sum(x * y for x, y in zip(a, b))
    return dot / (math.hypot(*a) * math.hypot(*b))

docs = {
    "Refund policy": [0.9, 0.1, 0.0],
    "Shipping times": [0.1, 0.9, 0.1],
    "Reset your password": [0.0, 0.2, 0.9],
}
question = "How do I get my money back?"
q = [0.8, 0.2, 0.1]
ranked = sorted(docs, key=lambda t: cosine(q, docs[t]), reverse=True)
print(ranked)
prompt = f"Answer using this document: {ranked[0]}\nQuestion: {question}"
print(prompt)

Output

['Refund policy', 'Shipping times', 'Reset your password']
Answer using this document: Refund policy
Question: How do I get my money back?

Step by step

L13The question's embedding (made up here; a model would produce it).
L14Most similar first.
L16In a real system the document's text goes in, and the prompt is sent to the language model.

Common mistakes

Dot product without dividing by length

Long documents get long vectors and win every search. Use cosine, or normalise vectors to length 1 first.

Always taking the top k

If nothing is relevant, the top k are still returned and the model answers from the wrong documents. Set a minimum score.

Mixing embedding models

Vectors from different models live in different spaces. Embed the question and the documents with the same model.

Quick quiz

1. What does this print?

a = [1, 0]
b = [0, 1]
c = [2, 0]
dot = lambda u, v: sum(x * y for x, y in zip(u, v))
print(dot(a, b), dot(a, c))

A
```
0 2
```
B
```
1 1
```
C
```
0 1
```

Show the answer

Answer: A

a and b share no direction (0); c is a in the same direction but twice as long (2).

2. Why rank documents by cosine similarity rather than the raw dot product?

ACosine compares direction (meaning) and ignores length
BCosine is faster to compute
CThe dot product only works in 2D

Show the answer

Answer: A

Otherwise a long vector wins just by being long.

3. What does RAG give the language model?

ARelevant documents in the prompt, so it can use facts it never learned
BNew training that permanently updates its weights
CA faster way to generate text

Show the answer

Answer: A

Retrieval happens at question time; the model itself doesn't change.

Practice

3 exercises in three steps: predict what a program prints, change a program to meet a new requirement, then write one from scratch. Your code runs in the browser and is checked the moment you submit.

PredictLength one
Type exactly what this prints.
ModifyThe FAQ that wins everything
A long, rambling FAQ page has a huge vector and wins every search. Change score to cosine similarity, so only the direction counts.
WriteRetrieve
Write retrieve(query, docs, k, min_score). docs maps a title to its vector. Return a list of up to k titles:
- only documents whose cosine similarity with query is at least min_score
- most similar first; equal scores in alphabetical order of title
Vectors are never all zeros. You'll need your own cosine function too.

Practise this lesson

Open the interactive lesson to edit and run these examples and solve the exercises with instant feedback. Free to start, no sign-up.

Open the interactive lesson Course map

Interactive lessons open in order. If you have not done the earlier ones yet, you start at the course map.

Terms in this lesson

Term	Chinese	Definition
Embedding	嵌入	A vector that places similar meanings close together.
Cosine similarity	余弦相似度	Dot product divided by both lengths.
RAG	检索增强生成	Retrieving documents and adding them to a model's prompt.