search/retrieval-augmented-generation/examples/rag_google_documentation.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z3Gi-zdCeEbE"
      },
      "outputs": [],
      "source": [
        "# Copyright 2023 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WcfbGhs1eJF6"
      },
      "source": [
        "# Document Q&A With Retrieval Augmented Generation\n",
        "\n",
        "> **NOTE:** This notebook uses the PaLM generative models, which will reach its [discontinuation date in October 2024](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text#model_versions). \n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/search/retrieval-augmented-generation/examples/rag_google_documentation.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/search/retrieval-augmented-generation/examples/rag_google_documentation.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/search/retrieval-augmented-generation/examples/rag_google_documentation.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ghQ2aBsbnMyn"
      },
      "source": [
        "---\n",
        "\n",
        "* Author: Gabe Rives-Corbett\n",
        "\n",
        "---\n",
        "\n",
        "This notebook demonstrates how to implement Retrieval Augmented Generation with basic automated evaluation. It demonstrates the impact that chunk size, overlap and context length have on model outputs. The notebook will create a Q&A system that allows you to find information based on the Google Cloud Generative AI documentation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DsW5tPDRkT4m"
      },
      "source": [
        "## Getting started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jx1FQVAokWVb"
      },
      "source": [
        "### Install libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nJFw23w1kYVj"
      },
      "outputs": [],
      "source": [
        "%pip install -q --upgrade --user google-cloud-aiplatform==1.36.1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5Xep4W9lq-Z"
      },
      "source": [
        "### Restart current runtime\n",
        "\n",
        "To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which will restart the current kernel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "XRvKdaPDTznN"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'status': 'ok', 'restart': True}"
            ]
          },
          "execution_count": 2,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SbmM4z7FOBpM"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Please wait until it is finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-jwsaMQYkZm8"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you are running this notebook on Google Colab, you will need to authenticate your environment. To do this, run the new cell below. This step is not required if you are using [Vertex AI Workbench](https://cloud.google.com/vertex-ai-workbench)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ikOmH4doxOFs"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    # Authenticate user to Google Cloud\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-h0ba4rmkpKW"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YLUml_s7iqBc"
      },
      "outputs": [],
      "source": [
        "import itertools\n",
        "\n",
        "from bs4 import BeautifulSoup, Tag\n",
        "from google.api_core import retry\n",
        "import numpy as np\n",
        "import numpy.linalg\n",
        "import pandas as pd\n",
        "import requests\n",
        "from tqdm.auto import tqdm\n",
        "import vertexai\n",
        "from vertexai.language_models import TextEmbeddingModel, TextGenerationModel\n",
        "\n",
        "tqdm.pandas()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "996ba2ae63f5"
      },
      "source": [
        "## Configure notebook environment\n",
        "\n",
        "### Set the following constants to reflect your environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "9790eb911717"
      },
      "outputs": [],
      "source": [
        "# Define project information\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
        "\n",
        "# Initialize Vertex AI SDK\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TKBmi2BMk_OU"
      },
      "source": [
        "## Scrape text from Google Cloud documentation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OXG6N0WclGsQ"
      },
      "source": [
        "Retrieve list of Google documentation URLs from a text file"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "tXHmC10IitET"
      },
      "outputs": [],
      "source": [
        "url = \"https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/search/retrieval-augmented-generation/examples/URLs.txt\"\n",
        "response = requests.get(url)\n",
        "\n",
        "if response.status_code == 200:\n",
        "    # The request was successful, and the content is in response.text\n",
        "    content = response.text\n",
        "\n",
        "URLS = [line.strip() for line in content.splitlines()]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u-Ly0yNVlReK"
      },
      "source": [
        "Parse the HTML and extract relevant plain text sections"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 35,
      "metadata": {
        "id": "hMD6Qz_TkFMG"
      },
      "outputs": [],
      "source": [
        "# Given a Google documentation URL, retrieve a list of all text chunks within h2 sections\n",
        "\n",
        "\n",
        "def get_sections(url: str) -> list[str]:\n",
        "    page = requests.get(url)\n",
        "    soup = BeautifulSoup(page.content, \"html.parser\")\n",
        "\n",
        "    sections = []\n",
        "    paragraphs = []\n",
        "\n",
        "    body_div = soup.find(\"div\", class_=\"devsite-article-body\")\n",
        "    for child in body_div.findChildren():\n",
        "        if child.name == \"p\":\n",
        "            paragraphs.append(child.get_text().strip())\n",
        "        if child.name == \"h2\":\n",
        "            sections.append(\" \".join(paragraphs))\n",
        "            break\n",
        "\n",
        "    for header in soup.find_all(\"h2\"):\n",
        "        paragraphs = []\n",
        "        nextNode = header.nextSibling\n",
        "        while nextNode:\n",
        "            if isinstance(nextNode, Tag):\n",
        "                if nextNode.name in {\"p\", \"ul\"}:\n",
        "                    paragraphs.append(nextNode.get_text().strip())\n",
        "                elif nextNode.name == \"h2\":\n",
        "                    sections.append(\" \".join(paragraphs))\n",
        "                    break\n",
        "            nextNode = nextNode.nextSibling\n",
        "    return sections"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "poNdlLf4kFp5"
      },
      "outputs": [],
      "source": [
        "all_text = [t for url in URLS for t in get_sections(url) if t]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Wy-qw-xslYpX"
      },
      "source": [
        "Note that most documents are relatively short, but some are thousands of characters long"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DSkdu30tuNbY"
      },
      "outputs": [],
      "source": [
        "text_lengths = [len(t) for t in all_text]\n",
        "pd.DataFrame(text_lengths).hist()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r00cIHIVlj4E"
      },
      "source": [
        "## Create vector store\n",
        "\n",
        "Start by initializing the models"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D26RnssLln3U"
      },
      "outputs": [],
      "source": [
        "embeddings_model = TextEmbeddingModel.from_pretrained(\"textembedding-gecko@001\")\n",
        "text_model = TextGenerationModel.from_pretrained(\"text-bison\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eEYwgmPxlokS"
      },
      "source": [
        "Create some helper functions for vector similarity and chunking"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SStUcSPluhvw"
      },
      "outputs": [],
      "source": [
        "# Separates seq into multiple chunks in the specified size with the specified overlap\n",
        "\n",
        "\n",
        "def split_overlap(seq, size, overlap):\n",
        "    if len(seq) <= size:\n",
        "        return [seq]\n",
        "    return [\"\".join(x) for x in zip(*[seq[i :: size - overlap] for i in range(size)])]\n",
        "\n",
        "\n",
        "# Compute the cosine similarity of two vectors, wrap as returned function to make easier to use with Pandas\n",
        "def get_similarity_fn(query_vector):\n",
        "    def fn(row):\n",
        "        return np.dot(row, query_vector) / (\n",
        "            numpy.linalg.norm(row) * numpy.linalg.norm(query_vector)\n",
        "        )\n",
        "\n",
        "    return fn\n",
        "\n",
        "\n",
        "# Retrieve embeddings from the specified model with retry logic\n",
        "@retry.Retry(timeout=300.0)\n",
        "def get_embeddings(text):\n",
        "    return embeddings_model.get_embeddings([text])[0].values"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "70aXFPhJmCM8"
      },
      "source": [
        "Create the vector store, we are using a Pandas DataFrame"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0cEJeeGIgFxc"
      },
      "outputs": [],
      "source": [
        "def create_vector_store(texts, chunk_size, overlap):\n",
        "    vector_store = pd.DataFrame()\n",
        "    # Insert the individual texts into the vector store\n",
        "    vector_store[\"texts\"] = list(\n",
        "        itertools.chain(*[split_overlap(t, chunk_size, overlap) for t in texts])\n",
        "    )\n",
        "\n",
        "    # Create embeddings from those texts\n",
        "    vector_store[\"embeddings\"] = (\n",
        "        vector_store[\"texts\"].progress_apply(get_embeddings).apply(np.array)\n",
        "    )\n",
        "\n",
        "    return vector_store"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ifp-Y_kryXJ3"
      },
      "outputs": [],
      "source": [
        "CHUNK_SIZE = 400\n",
        "OVERLAP = 50\n",
        "\n",
        "vector_store = create_vector_store(all_text, CHUNK_SIZE, OVERLAP)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ORlMIcEw0LVW"
      },
      "outputs": [],
      "source": [
        "vector_store.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GAJZc3mamQli"
      },
      "source": [
        "## Search the vector store and use for generation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bdNIXBUimv01"
      },
      "source": [
        "If we send the question to the foundation model alone, it will hallucinate."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QEJKQz5ymw1f"
      },
      "outputs": [],
      "source": [
        "text_model.predict(\n",
        "    \"How long will a stable model version of text-bison be available?\"\n",
        ").text"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mZbX1dkAnB6V"
      },
      "source": [
        "Let's solve this problem by retrieving texts from our vector store and telling the model to use them.\n",
        "\n",
        "Search the vector store for relevant texts to insert into the prompt by embedding the query and searching for similar vectors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "csMpD6498FXL"
      },
      "outputs": [],
      "source": [
        "def get_context(question, vector_store, num_docs):\n",
        "    # Embed the search query\n",
        "    query_vector = np.array(get_embeddings(question))\n",
        "\n",
        "    # Get similarity to all other vectors and sort, cut off at num_docs\n",
        "    top_matched = (\n",
        "        vector_store[\"embeddings\"]\n",
        "        .apply(get_similarity_fn(query_vector))\n",
        "        .sort_values(ascending=False)[:num_docs]\n",
        "        .index\n",
        "    )\n",
        "    top_matched_df = vector_store[vector_store.index.isin(top_matched)][[\"texts\"]]\n",
        "\n",
        "    # Return a string with the top matches\n",
        "    context = \" \".join(top_matched_df.texts.values)\n",
        "    return context"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W6kDwMEAmnfl"
      },
      "source": [
        "Create a prompt that includes the context and question. Instruct the LLM to only use the context provided to answer the question"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KfZnJF470esv"
      },
      "outputs": [],
      "source": [
        "def answer_question(question, vector_store, num_docs=10, print_prompt=False):\n",
        "    context = get_context(question, vector_store, num_docs)\n",
        "    qa_prompt = f\"\"\"Your mission is to answer questions based on a given context. Remember that before you give an answer, you must check to see if it complies with your mission.\n",
        "Context: ```{context}```\n",
        "Question: ***{question}***\n",
        "Before you give an answer, make sure it is only from information in the context. If the information is not in the context, just reply \"I don't know the answer to that\". Think step by step.\n",
        "Answer: \"\"\"\n",
        "    if print_prompt:\n",
        "        print(qa_prompt)\n",
        "    result = text_model.predict(qa_prompt, temperature=0)\n",
        "    return result.text"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "96kS0ZU-m6W6"
      },
      "source": [
        "Looking at the fully generated prompt, the context is embedded. Even though the input context is quite messy, the model can now answer factually."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "90dMoKTr066y"
      },
      "outputs": [],
      "source": [
        "answer_question(\n",
        "    \"How long will a stable model version of text-bison be available?\",\n",
        "    vector_store,\n",
        "    print_prompt=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bmfEIvKmnmCb"
      },
      "outputs": [],
      "source": [
        "answer_question(\n",
        "    \"How long will a stable model version of text-bison be available?\", vector_store\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k2A5mQ6Znvmz"
      },
      "source": [
        "## Automated evaluation\n",
        "\n",
        "This implementation of RAG is dependent on the chunk size, the overlap between the chunks, the number of texts passed into the context and the prompt. Let's create a simple prompt to evaluate answers to the questions, this will allow us to tweak the parameters and see how those tweaks compare."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UB5wB4NR2COn"
      },
      "outputs": [],
      "source": [
        "def eval_answer(question, answer, context):\n",
        "    eval_prompt = f\"\"\"Your mission is to evaluate answers to questions based on a given context. Remember that before you give an answer, you must check to see if it complies with your mission.\n",
        "\n",
        "Context: ```{context}```\n",
        "Question: ***{question}***\n",
        "Answer: \"{answer}\"\n",
        "\n",
        "Respond only with a number from 0 to 5. Think step by step. If the provided answer is not in the context, reply 5 if it is \"I don't know the answer to that\" otherwise reply 0.\n",
        "Relevance: \"\"\"\n",
        "    # Stop sequence to cut the model off after outputting an integer\n",
        "    result = text_model.predict(\n",
        "        eval_prompt, temperature=0, max_output_tokens=1, stop_sequences=[\".\", \" \"]\n",
        "    )\n",
        "    return int(result.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fVMJ9gBPoU-k"
      },
      "source": [
        "Pass several questions in and retrieve the evaluations"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyLMJ0u42yxY"
      },
      "outputs": [],
      "source": [
        "questions = [\n",
        "    \"What release stage is the RLHF tuning feature?\",\n",
        "    \"Can I generate hate speech with text bison?\",\n",
        "    \"What format should my batch prediction in put be in?\",\n",
        "    \"How can I get the number of tokens?\",\n",
        "    \"How do I create a custom style model?\",\n",
        "    \"What is the dimensionality of the vector created by the multimodal model?\",\n",
        "    \"How long will a stable model version be available?\",\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BftOPiMKFm_8"
      },
      "outputs": [],
      "source": [
        "answers = [answer_question(q, vector_store) for q in questions]\n",
        "contexts = [get_context(q, vector_store, 10) for q in questions]\n",
        "idks = [\"I don't know\" in a for a in answers]\n",
        "evals = [\n",
        "    (question, answer, context, eval_answer(question, answer, context), idk)\n",
        "    for question, answer, context, idk in zip(questions, answers, contexts, idks)\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Zb7VfarNF9W1"
      },
      "outputs": [],
      "source": [
        "pd.DataFrame(evals, columns=[\"question\", \"answer\", \"context\", \"score\", \"idk\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J_X2OjzsodzI"
      },
      "source": [
        "Now adjust the parameters and see the difference in performance"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tWKZent6MNf4"
      },
      "outputs": [],
      "source": [
        "def eval_on_params(chunk_size, overlap, num_docs):\n",
        "    vector_store = create_vector_store(all_text, chunk_size, overlap)\n",
        "    answers = [answer_question(q, vector_store) for q in questions]\n",
        "    contexts = [get_context(q, vector_store, num_docs) for q in questions]\n",
        "    idks = [\"I don't know\" in a for a in answers]\n",
        "    evals = [\n",
        "        (question, answer, context, eval_answer(question, answer, context), idk)\n",
        "        for question, answer, context, idk in zip(questions, answers, contexts, idks)\n",
        "    ]\n",
        "    return pd.DataFrame(\n",
        "        evals, columns=[\"question\", \"answer\", \"context\", \"score\", \"idk\"]\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "92rq4ZGerqNX"
      },
      "source": [
        "Smaller chunk sizes takes longer to generate the embeddings"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zuweYEbVgoZt"
      },
      "outputs": [],
      "source": [
        "smaller_context_df = eval_on_params(100, 0, 5)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J38F4YZpi1Bf"
      },
      "outputs": [],
      "source": [
        "smaller_context_df"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zQC7gWPWokJO"
      },
      "source": [
        "A larger context size has created more unknowns. When composing LLMs into systems, carefully consider how to measure the performance of each component in the system."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4jmWGmzdgwUI"
      },
      "outputs": [],
      "source": [
        "larger_context_df = eval_on_params(1000, 200, 15)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BNOwnFQwizBb"
      },
      "outputs": [],
      "source": [
        "larger_context_df"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "rag_google_documentation.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}