Retrieval-augmented generation, step by step
Answer questions about a document the model has never seen: split it, embed it, find the relevant pieces, and put them in the prompt.
1 Split the document into chunks
Whole sentences are packed together up to the chunk size.
- #1Northwind Cycles staff handbook. This shop and handbook are fictional and exist only for this lesson. Opening hours. The shop opens at 9 am and closes at 6 pm from Tuesday to Saturday. It is closed on Sundays and Mondays.
- #2On public holidays the shop opens from 10 am to 3 pm. Returns. Customers may return an unused bike within 30 days of purchase for a full refund. Used bikes can be exchanged within 14 days but not refunded.
- #3Clothing and helmets must be returned in their original packaging. Warranty. Every new frame carries a five-year warranty against manufacturing defects. Wheels, brakes, and other parts are covered for one year.
- #4The warranty does not cover damage from crashes, racing, or ordinary wear. Repairs. A basic tune-up costs 45 dollars and takes one working day. Customers can book a repair online or at the counter.
- #5Bikes left more than 60 days after a repair is finished may be donated to the community bike program. Staff benefits.
- #6Employees receive a 25 percent discount on bikes and a 40 percent discount on parts and clothing after their first three months. Staff may borrow a demo bike for up to one week, twice a year. Bike fitting.
- #7A professional fitting takes about 90 minutes and costs 120 dollars. Customers who buy a bike within two weeks of a fitting get the fitting fee back as store credit.
Loading the recorded results
Chunks7
Challenges
0 of 4 complete
- Shrink the chunks until the best match for your question is cut off partway through a sentence. (not yet)
- With one of the example questions, find a chunk size where the chunk holding the answer is not ranked first, and set k so it still makes it into the prompt. (not yet)
- Ask how much a tune-up is, then find a chunk size where the chunk with the price still comes first but scores below 0.30. (not yet)
- Find a question with a number for an answer that the model gets right from the retrieved chunks but not from the question alone. (not yet)
Try this
- Change the chunk size and watch which chunk wins for the same question.
- Ask a question the handbook cannot answer, such as “Do you sell skateboards?”, and read what gets retrieved anyway.
- Write the answer to How much is a tune-up? from the prompt, then from the question alone.
A tiny retrieval pipeline in codePython with NumPy
import re
import numpy as np
doc = ("Returns. Customers may return an unused bike within 30 days. "
"Used bikes can be exchanged within 14 days. "
"Repairs. A basic tune-up costs 45 dollars and takes one day.")
# 1. Chunk: pack whole sentences into chunks of at most 10 words.
chunks, current = [], []
for sentence in re.split(r"(?<=[.!?])\s+", doc):
if current and len(" ".join(current + [sentence]).split()) > 10:
chunks.append(" ".join(current))
current = []
current.append(sentence)
chunks.append(" ".join(current))
# 2. Embed: the demo uses MiniLM. Here, averaging made-up word vectors
# (topics: returning something, repairs, money) stands in for it.
topics = {"return": [1, 0, 0], "returns": [1, 0, 0], "exchanged": [1, 0, 0],
"back": [1, 0, 0], "repairs": [0, 1, 0], "tune": [0, 1, 0],
"costs": [0, 0, 1], "dollars": [0, 0, 1], "much": [0, 0, 1]}
def embed(text):
vecs = [topics[w] for w in re.findall(r"[a-z]+", text.lower()) if w in topics]
v = np.mean(vecs, axis=0)
return v / np.linalg.norm(v)
# 3. Retrieve: rank chunks by cosine similarity to the question, keep the top k.
question = "Can I bring back a bike?"
scores = np.array([embed(c) @ embed(question) for c in chunks])
top = np.argsort(-scores, kind="stable")[:2]
for i, c in enumerate(chunks):
print(f"{scores[i]:.2f} {c}")
# 4. Assemble the prompt a language model would receive.
context = "\n".join(f"[{n}] {chunks[i]}" for n, i in enumerate(top, start=1))
print("\nAnswer the question using only the context below. If the context does")
print("not contain the answer, say that you do not know.\n")
print(f"Context:\n{context}\n\nQuestion: {question}\nAnswer:")1.00 Returns. Customers may return an unused bike within 30 days.
0.71 Used bikes can be exchanged within 14 days. Repairs.
0.00 A basic tune-up costs 45 dollars and takes one day.
Answer the question using only the context below. If the context does
not contain the answer, say that you do not know.
Context:
[1] Returns. Customers may return an unused bike within 30 days.
[2] Used bikes can be exchanged within 14 days. Repairs.
Question: Can I bring back a bike?
Answer: