Retrieval-augmented generation pipeline

Chunk a document, embed it with a real model in your browser, retrieve the best passages for a question, and have a small chat model answer from them.

IntermediateExplained in Retrieval-augmented generation

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.

Example questions

1 Split the document into chunks

Whole sentences are packed together up to the chunk size.

40
  1. #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.
  2. #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.
  3. #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.
  4. #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.
  5. #5Bikes left more than 60 days after a repair is finished may be donated to the community bike program. Staff benefits.
  6. #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.
  7. #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.

Recorded from all-MiniLM-L6-v2

Loading the model's recorded outputs for the example questions and the handbook at every chunk size.

About the download

Embed the chunks and your question to find the most relevant passages. This uses all-MiniLM-L6-v2, a small sentence-embedding model that maps text to 384 numbers.

The first time, your browser downloads about 24 MB of model files from Hugging Face, plus a 14 MB runtime from this site. Both are cached for later visits. Everything runs on your device; nothing you type is sent anywhere.

Manage downloads

Loading the recorded results

Chunks7

Challenges

0 of 4 complete

  1. Shrink the chunks until the best match for your question is cut off partway through a sentence. (not yet)
  2. 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)
  3. 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)
  4. 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:

Try "embedding", "softmax", "overfitting", or "backpropagation".