{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Module 15: Retrieval-Augmented Generation"
      ],
      "id": "6e9858b4-c914-4610-86c7-767c7e82f8ce"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [],
      "id": "e92ca0a5-b154-4448-ad35-ab2ad0932732"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "b5b60493-6efe-4e84-978b-a6d1fedff97d"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Introduction\n",
        "\n",
        "Your model (m06–m08) knows exactly what was in its training data, frozen\n",
        "at a cutoff date, blended into billions of weights. Ask it about a\n",
        "document it never saw — your company’s wiki, a paper from last week, a\n",
        "fact it half-remembers — and it will confidently make something up.\n",
        "**Retrieval-Augmented Generation (RAG)** fixes this without touching the\n",
        "weights: at query time, *search* an external corpus for relevant\n",
        "passages and paste them into the prompt. The model then answers from\n",
        "fresh, specific, **checkable** text instead of parametric memory alone.\n",
        "\n",
        "The engine underneath is **nearest-neighbor search in an embedding\n",
        "space**. Encode every document as a vector; encode the question the same\n",
        "way; retrieve the documents whose vectors point most nearly the same\n",
        "direction (highest **cosine similarity**); stuff them into the prompt;\n",
        "generate. This module builds that whole pipeline from scratch.\n",
        "\n",
        "To keep it fully runnable with no trained encoder, we encode with\n",
        "classic **TF-IDF** vectors — term frequency times inverse document\n",
        "frequency, computable by hand. Modern RAG swaps TF-IDF for *dense*\n",
        "embeddings from a trained encoder (the m04 idea, learned end-to-end),\n",
        "but **the pipeline is identical** — only the `encode` step changes.\n",
        "\n",
        "Why it matters for LLMs:\n",
        "\n",
        "- **Knowledge without retraining.** New facts arrive by adding documents\n",
        "  to the index, not by a training run. The model stays fixed; the corpus\n",
        "  is live.\n",
        "- **Grounding and citations.** The answer is conditioned on retrieved\n",
        "  text you can show the user — the standard defense against\n",
        "  hallucination.\n",
        "- **It is everywhere.** Chat-with-your-docs, coding assistants over a\n",
        "  repo, search copilots — almost every applied LLM system is a RAG\n",
        "  system.\n",
        "\n",
        "### What You’ll Learn\n",
        "\n",
        "After this module, you can:\n",
        "\n",
        "- Explain the RAG pipeline — **encode → search → retrieve → augment →\n",
        "  generate** — and why it beats parametric memory for fresh or private\n",
        "  knowledge.\n",
        "- Build **TF-IDF** document vectors from scratch (`tf`, smoothed `idf`).\n",
        "- Rank documents by **cosine similarity** and retrieve the top-k.\n",
        "- **Chunk** long documents, and de-duplicate results with **Maximal\n",
        "  Marginal Relevance (MMR)**.\n",
        "- Assemble a grounded prompt and wrap it all in a `Retriever`.\n",
        "\n",
        "### Prerequisites\n",
        "\n",
        "This module requires familiarity with:\n",
        "\n",
        "- [Module 04: Embeddings](../m04_embeddings/lesson.qmd) — vectors as\n",
        "  meaning; dense retrieval replaces TF-IDF with *learned* embeddings of\n",
        "  exactly this kind.\n",
        "- [Module 08: Generation](../m08_generation/lesson.qmd) — RAG conditions\n",
        "  generation on retrieved context; the decoding is unchanged.\n",
        "\n",
        "## Intuition: The RAG Pipeline\n",
        "\n",
        "RAG is five steps, and only the middle three are new — the model at the\n",
        "end is the same one you already built. Step through what each stage\n",
        "produces:"
      ],
      "id": "fd25666a-e81a-412d-a122-8f0af925ba8f"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "21eb215c-21bd-4b26-b392-09da9cda9088"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "56469c37-25a7-40f0-a677-b99b3f13f333"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "ef9f4fcd-b297-4e89-a303-f12bd7d04b50"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Key Insight**\n",
        ">\n",
        "> Only **Encode → Search → Retrieve → Augment** is retrieval;\n",
        "> **Generate** is the model you already have. RAG is not a new kind of\n",
        "> model — it is a way of *choosing what goes in the prompt*. Everything\n",
        "> hard is in ranking documents by relevance.\n",
        "\n",
        "## The Math: TF-IDF and Cosine Similarity\n",
        "\n",
        "To search by meaning we need documents as vectors. The classic recipe\n",
        "weights each term by how often it appears in a document (**term\n",
        "frequency**) against how rare it is across the corpus (**inverse\n",
        "document frequency**), so shared *rare* words count and boilerplate like\n",
        "“the” does not:\n",
        "\n",
        "$$\\text{tf}(t, d) = \\text{count of } t \\text{ in } d, \\qquad\n",
        "\\text{idf}(t) = \\ln\\!\\frac{1 + N}{1 + \\text{df}(t)} + 1, \\qquad\n",
        "\\text{tfidf}(t, d) = \\text{tf}(t, d)\\cdot\\text{idf}(t),$$\n",
        "\n",
        "where $N$ is the number of documents and $\\text{df}(t)$ how many contain\n",
        "$t$. The $+1$ smoothing keeps every idf positive; a term in *every*\n",
        "document gets idf $= 1$.\n",
        "\n",
        "Relevance is the angle between vectors, not their length — a long\n",
        "document is not more relevant just for being long — so we compare with\n",
        "**cosine similarity**:\n",
        "\n",
        "$$\\cos(\\mathbf{q}, \\mathbf{d}) = \\frac{\\mathbf{q}\\cdot\\mathbf{d}}{\\lVert\\mathbf{q}\\rVert\\,\\lVert\\mathbf{d}\\rVert}.$$\n",
        "\n",
        "Normalize both vectors to unit length and cosine is just a dot product.\n",
        "`retrieval.py` implements `tfidf_matrix`, `cosine_similarity`, and\n",
        "`retrieve`.\n",
        "\n",
        "## Code: Build a Retriever\n",
        "\n",
        "The `Retriever` class fits a corpus once, then answers queries by cosine\n",
        "top-k:"
      ],
      "id": "90bd26d3-b371-409a-88a0-49b35800266c"
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[1] 0.233  Attention computes a weighted sum of value vectors using que...\n",
            "[0] 0.185  The transformer architecture uses self-attention to relate e..."
          ]
        }
      ],
      "source": [
        "from retrieval import Retriever\n",
        "\n",
        "corpus = [\n",
        "    \"The transformer architecture uses self-attention to relate every token to every other token.\",\n",
        "    \"Attention computes a weighted sum of value vectors using query-key similarity scores.\",\n",
        "    \"Byte-pair encoding builds a subword vocabulary by merging the most frequent adjacent pairs.\",\n",
        "    \"Photosynthesis lets plants convert sunlight, water, and carbon dioxide into glucose and oxygen.\",\n",
        "    \"Mount Everest is the tallest mountain above sea level, on the border of Nepal and Tibet.\",\n",
        "]\n",
        "\n",
        "retriever = Retriever().fit(corpus)\n",
        "for idx, score, doc in retriever.search(\"how does attention work?\", k=2):\n",
        "    print(f\"[{idx}] {score:.3f}  {doc[:60]}...\")"
      ],
      "id": "766a984d"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The attention documents win because they *share the rare, informative\n",
        "words* of the query. Now assemble the retrieved passages into a grounded\n",
        "prompt — the “augment” step:"
      ],
      "id": "9ba9a71d-6c5e-4231-a196-1aeb1935a613"
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Answer the question using only the context below. If the answer is not in the context, say you don't know.\n",
            "\n",
            "Context:\n",
            "[1] Attention computes a weighted sum of value vectors using query-key similarity scores.\n",
            "[2] The transformer architecture uses self-attention to relate every token to every other token.\n",
            "\n",
            "Question: how does attention work?\n",
            "Answer:"
          ]
        }
      ],
      "source": [
        "from retrieval import build_prompt\n",
        "\n",
        "query = \"how does attention work?\"\n",
        "hits = retriever.search(query, k=2)\n",
        "prompt = build_prompt(query, [doc for _, _, doc in hits])\n",
        "print(prompt)"
      ],
      "id": "c8a01814"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "That string is exactly what you would hand to `generate()` from m08. The\n",
        "model now answers *from the retrieved context* — and because you have\n",
        "the passages, you can show them as citations. Swap the TF-IDF `encode`\n",
        "for a trained dense encoder and nothing else in this pipeline changes.\n",
        "\n",
        "## Interactive: Watch Retrieval Rank the Corpus\n",
        "\n",
        "Pick a question and see cosine similarity score every document. The\n",
        "**top-k** (here k = 3) are what gets retrieved and pasted into the\n",
        "prompt; everything else is ignored. Notice how the scores concentrate on\n",
        "the documents that share the query’s informative words."
      ],
      "id": "4c1c4d39-1cf6-430e-9684-d92ebacd257f"
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [],
      "source": [],
      "id": "2b603c6c"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "c92ad41e-92b0-46b4-8631-0eb80ba185a5"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "51a02f5b-b6a4-4f7f-84b2-afaaff0656b7"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "b3701b59-93e3-4444-82a6-0b2045e600c8"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "abab5641-b597-4702-843f-39c7ee55a165"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Try This**\n",
        ">\n",
        "> 1.  Switch between the attention query and the photosynthesis query.\n",
        ">     Watch the mass of similarity jump to a completely different pair\n",
        ">     of documents — the corpus is the same; only the query moved.\n",
        "> 2.  Note the *ignored* documents still get a small nonzero score from\n",
        ">     incidental shared words (“the”, “of”). idf is what keeps those\n",
        ">     from dominating.\n",
        "\n",
        "## Chunking and Diversity\n",
        "\n",
        "Two practical problems break naïve retrieval, and `retrieval.py` handles\n",
        "both.\n",
        "\n",
        "**Chunking.** You do not index whole documents — a 50-page PDF has one\n",
        "vector that means nothing specific. You split it into small overlapping\n",
        "windows so each retrievable passage is focused, with a little overlap so\n",
        "a fact spanning a boundary survives in at least one chunk:"
      ],
      "id": "141c6521-1875-4bf9-9c5d-346cc70702fc"
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['Retrieval augmented generation searches a', 'searches a corpus then conditions', 'then conditions the model on', 'model on the results']"
          ]
        }
      ],
      "source": [
        "from retrieval import chunk_text\n",
        "\n",
        "passage = \"Retrieval augmented generation searches a corpus then conditions the model on the results\"\n",
        "print(chunk_text(passage, chunk_size=5, overlap=2))"
      ],
      "id": "350dfb79"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Diversity.** Plain top-k happily returns three near-identical\n",
        "passages, wasting the context window. **Maximal Marginal Relevance\n",
        "(MMR)** picks each next passage for relevance *minus* similarity to what\n",
        "is already chosen:\n",
        "\n",
        "$$\\text{next} = \\arg\\max_{c}\\;\\Big[\\lambda\\cdot\\text{rel}(c) - (1-\\lambda)\\max_{s\\in S}\\text{sim}(c, s)\\Big].$$"
      ],
      "id": "28b7f538-1ee6-4083-bac5-890f72dd0715"
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "plain top-2: [0, 1]\n",
            "MMR top-2:   [0, 2]"
          ]
        }
      ],
      "source": [
        "from retrieval import tfidf_matrix, encode_query, retrieve, mmr\n",
        "\n",
        "docs = [\n",
        "    \"transformer attention model\",\n",
        "    \"transformer attention model network\",   # near-duplicate of doc 0\n",
        "    \"retrieval augmented generation search\",\n",
        "]\n",
        "matrix, vocab, idf = tfidf_matrix(docs)\n",
        "qv = encode_query(\"transformer attention\", vocab, idf)\n",
        "\n",
        "print(\"plain top-2:\", [i for i, _ in retrieve(qv, matrix, k=2)])   # two near-duplicates\n",
        "print(\"MMR top-2:  \", mmr(qv, matrix, k=2, lambda_=0.3))            # swaps in the different doc"
      ],
      "id": "e469d4d0"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "With $\\lambda = 1$, MMR is just top-k; lower $\\lambda$ trades a little\n",
        "relevance for coverage.\n",
        "\n",
        "## Common Pitfalls\n",
        "\n",
        "1.  **Chunk too big or too small.** Whole-document chunks retrieve\n",
        "    nothing specific; one-sentence chunks lose context. A few hundred\n",
        "    words with light overlap is the usual sweet spot.\n",
        "2.  **Lexical retrieval misses paraphrases.** TF-IDF matches *words*, so\n",
        "    “car” won’t retrieve a passage about “automobiles.” This is exactly\n",
        "    why production RAG uses **dense** embeddings — they match meaning,\n",
        "    not surface form.\n",
        "3.  **Retrieving duplicates.** Top-k over a corpus with repeats fills\n",
        "    the context with the same fact. De-duplicate (MMR) or the model sees\n",
        "    no new information.\n",
        "4.  **Stuffing too much context.** More passages is not better —\n",
        "    irrelevant context distracts the model and costs tokens. Retrieve\n",
        "    few, retrieve well.\n",
        "5.  **Trusting retrieval blindly.** If the answer isn’t in the corpus, a\n",
        "    grounded prompt should let the model *say so* — the `build_prompt`\n",
        "    instruction asks for exactly that. RAG reduces hallucination; it\n",
        "    does not abolish it.\n",
        "\n",
        "## Exercises\n",
        "\n",
        "### Exercise 1: idf by hand"
      ],
      "id": "defbefb0-99d5-4aeb-89fa-01974e2800fc"
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [],
      "source": [
        "from retrieval import build_vocab, compute_idf\n",
        "\n",
        "# For docs [\"the cat sat\", \"the dog ran\", \"the cat ran\"], compute idf by hand\n",
        "# (N=3) for \"the\" (df=3), \"cat\" (df=2), \"sat\" (df=1). Confirm compute_idf agrees\n",
        "# and that the rarest term has the highest idf.\n",
        "\n",
        "# Your implementation here:"
      ],
      "id": "a033704c"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Exercise 2: retrieve then augment"
      ],
      "id": "589f8c7e-0b3c-4a34-936d-fb28f812bcc0"
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [],
      "source": [
        "from retrieval import Retriever, build_prompt\n",
        "\n",
        "# Fit a Retriever on 4–5 facts of your own, retrieve the top-2 for a question, and\n",
        "# build_prompt them. Print the prompt and check the right facts were pulled in.\n",
        "\n",
        "# Your implementation here:"
      ],
      "id": "4028840c"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Exercise 3: when does MMR help?"
      ],
      "id": "2e5508ac-5d28-4490-8560-39e8cb7b516f"
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [],
      "source": [
        "from retrieval import tfidf_matrix, encode_query, mmr\n",
        "\n",
        "# Build a corpus with two near-duplicate relevant docs and one different relevant\n",
        "# doc. Show that lambda=1.0 returns both duplicates but lambda=0.2 returns a\n",
        "# diverse pair.\n",
        "\n",
        "# Your implementation here:"
      ],
      "id": "fc197580"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Summary\n",
        "\n",
        "Key takeaways:\n",
        "\n",
        "1.  **RAG conditions generation on retrieved text** — encode → search →\n",
        "    retrieve → augment → generate — so the model answers from a live\n",
        "    corpus, not just frozen weights.\n",
        "2.  **Retrieval is nearest-neighbor search** — encode documents and the\n",
        "    query into one vector space and rank by **cosine similarity**.\n",
        "3.  **TF-IDF is the from-scratch encoder** — term frequency times\n",
        "    smoothed idf makes shared *rare* words count; dense embeddings are\n",
        "    the same pipeline with a learned encoder.\n",
        "4.  **Chunk long documents** into small overlapping passages so each\n",
        "    retrieved unit is focused.\n",
        "5.  **MMR buys diversity** — pick passages for relevance minus\n",
        "    redundancy so the context isn’t three copies of one fact.\n",
        "6.  **RAG’s win is grounding** — fresh, private, citable knowledge with\n",
        "    no retraining, and a real dent in hallucination.\n",
        "\n",
        "## What’s Next\n",
        "\n",
        "You can now ground a model in an external corpus. The natural next step\n",
        "is letting the model *act*: **tool use and agents**, where retrieval\n",
        "becomes one action among many (search, run code, call an API) inside a\n",
        "plan-act-observe loop — the same “put the right thing in the context”\n",
        "idea, now driven by the model itself.\n",
        "\n",
        "### Going Deeper\n",
        "\n",
        "**Core Papers:**\n",
        "\n",
        "- [Retrieval-Augmented Generation for Knowledge-Intensive NLP\n",
        "  Tasks](https://arxiv.org/abs/2005.11401) — Lewis et al. (2020), the\n",
        "  paper that named RAG and combined a retriever with a generator\n",
        "  end-to-end.\n",
        "- [Dense Passage Retrieval for Open-Domain Question\n",
        "  Answering](https://arxiv.org/abs/2004.04906) — Karpukhin et\n",
        "  al. (2020), learned dense embeddings that beat BM25 by 9–19% on\n",
        "  retrieval.\n",
        "- [REALM: Retrieval-Augmented Language Model\n",
        "  Pre-Training](https://arxiv.org/abs/2002.08909) — Guu et al. (2020),\n",
        "  retrieval baked into pre-training.\n",
        "\n",
        "**Practical Resources:**\n",
        "\n",
        "- [The Probabilistic Relevance Framework: BM25 and\n",
        "  Beyond](https://www.staff.city.ac.uk/~sbrp622/papers/foundations_bm25_review.pdf)\n",
        "  — Robertson & Zaragoza (2009), the classic lexical retriever TF-IDF\n",
        "  leads toward.\n",
        "- [FAISS](https://github.com/facebookresearch/faiss) — the library that\n",
        "  makes cosine top-k fast over millions of vectors."
      ],
      "id": "5611fca0-4bbb-4f05-9e5b-b05d92e4e5a9"
    }
  ],
  "nbformat": 4,
  "nbformat_minor": 5,
  "metadata": {
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "path": "/opt/buildhome/.asdf/installs/python/3.11.14/share/jupyter/kernels/python3"
    },
    "language_info": {
      "name": "python",
      "codemirror_mode": {
        "name": "ipython",
        "version": "3"
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.14"
    }
  }
}