Step 1: Build a corpus + load the embedding model
Language models have a frozen worldview: they only know what was in their pretraining corpus. Ask GPT-2 about a 2024 event and it hallucinates. Retrieval-Augmented Generation fixes this at the cheapest possible place in the stack — you look up relevant facts and paste them into the prompt. No fine-tuning, no retraining, just embed + search + prompt.
RAG is the single most widely-deployed LLM pattern in 2024–2026. Every enterprise "chat with your docs" product, every helpdesk bot, every search-grounded answer in ChatGPT and Claude and Gemini — all RAG underneath. The Lewis et al. 2020 paper introduced the term; the whole ecosystem (LangChain, LlamaIndex, pgvector, Milvus, Weaviate, Pinecone) exists to make this one pipeline easier to build.
The pipeline in five moves
documents → [chunker] → chunks → [embedding model] → vectors (in a store)
↓
query → [embedding model] → vector → [similarity search]
↓
top-K chunks + query → [LLM] → grounded answer
Every knob matters:
- Chunk size — too small = missing context, too big = diluted signal. 256–512 tokens is typical.
- Embedding model — BGE, E5, Jina, and OpenAI's ada-002 are common. All produce ~768-1024-dim vectors.
- Retrieval K — top-3 to top-10 chunks is standard. More = more context but more noise.
- Vector store — in-memory matrix (this lab), FAISS, Milvus, pgvector. Trade-off: latency vs scale.
- Prompt template — how you inject the retrieved context affects the LLM's willingness to use it.
References to read alongside
- RAG paper (Lewis et al., 2020) — the original. Still the cleanest articulation of why retrieval + generation outperforms either alone.
- BGE embeddings (Xiao et al., 2023) — the embedding model family we'll use. Open, top of MTEB for months.
- DPR (Karpukhin et al., 2020) — Dense Passage Retrieval; trained embeddings specifically for QA retrieval.
- ColBERT (Khattab & Zaharia, 2020) — late-interaction retrieval. Better accuracy at the cost of storage.
- Lost in the Middle (Liu et al., 2023) — classic empirical finding: LLMs use info at the start and end of context, ignore the middle. Implications for RAG prompt design.
Models
- Embedding:
BAAI/bge-small-en-v1.5(~130 MB, 384-dim vectors, produced by the BGE team) - Generator:
TinyLlama-1.1B-Chat(~2 GB, used throughout this course)
Our knowledge base: 8 carefully-chosen facts about modern LLMs. This is tiny but lets you see RAG end-to-end cleanly. In production your corpus is tens of thousands to millions of documents and the retrieval step has to be fast (milliseconds).
import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM
device = torch.device('cuda')
documents = [
"Llama 3 was released by Meta in April 2024. It was trained on approximately 15 trillion tokens — roughly 5x more data than Llama 2. The 8B and 70B variants are open-weight.",
"The Transformer architecture was introduced by Vaswani et al. in the 2017 paper 'Attention Is All You Need'. It replaced recurrence and convolution with pure attention, kicking off the modern LLM era.",
"GPT-3, released by OpenAI in 2020, has 175 billion parameters. Its emergent few-shot learning abilities surprised researchers and marked a major scaling-law inflection point.",
"LoRA (Low-Rank Adaptation) was proposed by Hu et al. in 2021. It lets you fine-tune a frozen pretrained model by adding tiny low-rank matrices on top of attention layers, typically cutting trainable parameters by 10000x.",
"The BGE family of embedding models was released by the Beijing Academy of Artificial Intelligence in 2023. BGE-small has 33 million parameters and produces 384-dimensional vectors.",
"Mixture of Experts (MoE) models use many specialist subnetworks and a router that sends each token to a subset. Mixtral 8x7B (released December 2023 by Mistral) demonstrated MoE at competitive quality.",
"The Chinchilla paper (Hoffmann et al., 2022) established compute-optimal training: models should see roughly 20 tokens per parameter during pretraining. Most pre-Chinchilla models were undertrained.",
"Retrieval-Augmented Generation was formalized in the 2020 RAG paper by Lewis et al. It combines a parametric language model with a non-parametric knowledge store accessed via dense retrieval.",
]
print(f'{len(documents)} documents in corpus')
EMBED_ID = 'BAAI/bge-small-en-v1.5'
print(f'Loading embedding model {EMBED_ID} (~130 MB first time)...')
embed_tokenizer = AutoTokenizer.from_pretrained(EMBED_ID)
embed_model = AutoModel.from_pretrained(EMBED_ID).to(device)
embed_model.eval()
# The idiomatic BGE encode: tokenize, forward pass, take [CLS] or mean-pool, then L2-normalize.
# BGE's team recommends [CLS] pooling specifically (not mean pool) for this model family.
def embed(texts, max_len=256):
"""Encode a list of strings into L2-normalized 384-dim vectors on CUDA."""
enc = embed_tokenizer(texts, padding=True, truncation=True, max_length=max_len, return_tensors='pt').to(device)
with torch.no_grad():
out = embed_model(**enc)
# BGE uses [CLS] pooling: the first token's hidden state at the last layer
pooled = out.last_hidden_state[:, 0]
# L2-normalize so cosine similarity reduces to dot product
return F.normalize(pooled, dim=-1)
# Quick smoke test
v = embed(['What is Llama 3?', 'How big is Llama 3?'])
print(f'Encoded shape: {tuple(v.shape)}, norm: {torch.norm(v, dim=-1).cpu().tolist()}')
similar = (v[0] @ v[1]).item()
print(f'Cosine sim between two Llama-3 questions: {similar:.3f} (should be > 0.8 — they are semantically close)')
