Copy_of_Day_2_Document_Q&A_with_RAG.ipynb

{
  "metadata": {
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3",
      "language": "python"
    },
    "language_info": {
      "name": "python",
      "version": "3.10.14",
      "mimetype": "text/x-python",
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "pygments_lexer": "ipython3",
      "nbconvert_exporter": "python",
      "file_extension": ".py"
    },
    "colab": {
      "toc_visible": true,
      "provenance": [],
      "include_colab_link": true
    },
    "kaggle": {
      "accelerator": "none",
      "dataSources": [],
      "dockerImageVersionId": 30786,
      "isInternetEnabled": true,
      "language": "python",
      "sourceType": "notebook",
      "isGpuEnabled": false
    }
  },
  "nbformat_minor": 0,
  "nbformat": 4,
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/nf15/handson-ml/blob/master/Copy_of_Day_2_Document_Q%26A_with_RAG.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "##### Copyright 2024 Google LLC."
      ],
      "metadata": {
        "id": "Tce3stUlHN0L"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# @title 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."
      ],
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu",
        "jupyter": {
          "source_hidden": true
        }
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Day 2 - Document Q&A with RAG using Chroma\n",
        "\n",
        "Welcome back to the Kaggle 5-day Generative AI course!\n",
        "\n",
        "**NOTE**: The Day 1 notebook contains lots of information for getting set up with Kaggle Notebooks. If you are having any issues, please [check out the troubleshooting steps there](https://www.kaggle.com/code/markishere/day-1-prompting#Get-started-with-Kaggle-notebooks).\n",
        "\n",
        "Two big limitations of LLMs are 1) that they only \"know\" the information that they were trained on, and 2) that they have limited input context windows. A way to address both of these limitations is to use a technique called Retrieval Augmented Generation, or RAG. A RAG system has three stages:\n",
        "\n",
        "1. Indexing\n",
        "2. Retrieval\n",
        "3. Generation\n",
        "\n",
        "Indexing happens ahead of time, and allows you to quickly look up relevant information at query-time. When a query comes in, you retrieve relevant documents, combine them with your instructions and the user's query, and have the LLM generate a tailored answer in natural language using the supplied information. This allows you to provide information that the model hasn't seen before, such as product-specific knowledge or live weather updates.\n",
        "\n",
        "In this notebook you will use the Gemini API to create a vector database, retrieve answers to questions from the database and generate a final answer. You will use [Chroma](https://docs.trychroma.com/), an open-source vector database. With Chroma, you can store embeddings alongside metadata, embed documents and queries, and search your documents.\n",
        "\n",
        "## For help\n",
        "\n",
        "**Common issues are covered in the [FAQ and troubleshooting guide](https://www.kaggle.com/code/markishere/day-0-troubleshooting-and-faqs).**\n"
      ],
      "metadata": {
        "id": "CsVPnR8VbXE6"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Setup\n",
        "\n",
        "First, install ChromaDB and the Gemini API Python SDK."
      ],
      "metadata": {
        "id": "akuOzK4dJl3j"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%pip install -U -q \"google-generativeai>=0.8.3\" chromadb"
      ],
      "metadata": {
        "id": "JbXe7Oodc5dP",
        "execution": {
          "iopub.status.busy": "2024-11-15T16:23:00.772568Z",
          "iopub.execute_input": "2024-11-15T16:23:00.773108Z",
          "iopub.status.idle": "2024-11-15T16:23:17.733735Z",
          "shell.execute_reply.started": "2024-11-15T16:23:00.773061Z",
          "shell.execute_reply": "2024-11-15T16:23:17.731974Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "You may see output containing `ERROR: pip's dependency resolver does not currently take into account all the packages that are installed` - this is OK, the packages are still installed and compatible for this codelab. Also note that you do not have to restart the kernel."
      ],
      "metadata": {
        "id": "KHI1h1DO44nc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import google.generativeai as genai\n",
        "from IPython.display import Markdown"
      ],
      "metadata": {
        "id": "muuhsDmmKdHi",
        "execution": {
          "iopub.status.busy": "2024-11-15T16:23:17.747093Z",
          "iopub.execute_input": "2024-11-15T16:23:17.747551Z",
          "iopub.status.idle": "2024-11-15T16:23:17.758182Z",
          "shell.execute_reply.started": "2024-11-15T16:23:17.747508Z",
          "shell.execute_reply": "2024-11-15T16:23:17.756784Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Set up your API key\n",
        "\n",
        "To run the following cell, your API key must be stored it in a [Kaggle secret](https://www.kaggle.com/discussions/product-feedback/114053) named `GOOGLE_API_KEY`.\n",
        "\n",
        "If you don't already have an API key, you can grab one from [AI Studio](https://aistudio.google.com/app/apikey). You can find [detailed instructions in the docs](https://ai.google.dev/gemini-api/docs/api-key).\n",
        "\n",
        "To make the key available through Kaggle secrets, choose `Secrets` from the `Add-ons` menu and follow the instructions to add your key or enable it for this notebook."
      ],
      "metadata": {
        "id": "FQOGMejVu-6D"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from kaggle_secrets import UserSecretsClient\n",
        "\n",
        "GOOGLE_API_KEY = UserSecretsClient().get_secret(\"GOOGLE_API_KEY\")\n",
        "genai.configure(api_key=GOOGLE_API_KEY)"
      ],
      "metadata": {
        "id": "ysayz8skEfBW",
        "execution": {
          "iopub.status.busy": "2024-11-15T16:24:36.468912Z",
          "iopub.execute_input": "2024-11-15T16:24:36.469508Z",
          "iopub.status.idle": "2024-11-15T16:24:36.708671Z",
          "shell.execute_reply.started": "2024-11-15T16:24:36.469462Z",
          "shell.execute_reply": "2024-11-15T16:24:36.707423Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "If you received an error response along the lines of `No user secrets exist for kernel id ...`, then you need to add your API key via `Add-ons`, `Secrets` **and** enable it.\n",
        "\n",
        "![Screenshot of the checkbox to enable GOOGLE_API_KEY secret](https://storage.googleapis.com/kaggle-media/Images/5gdai_sc_3.png)"
      ],
      "metadata": {
        "id": "52b101760a45"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Explore available models\n",
        "\n",
        "You will be using the [`embedContent`](https://ai.google.dev/api/embeddings#method:-models.embedcontent) API method to calculate embeddings in this guide. Find a model that supports it through the [`models.list`](https://ai.google.dev/api/models#method:-models.list) endpoint. You can also find more information about the embedding models on [the models page](https://ai.google.dev/gemini-api/docs/models/gemini#text-embedding).\n",
        "\n",
        "`text-embedding-004` is the most recent embedding model, so you will use it for this exercise."
      ],
      "metadata": {
        "id": "fegnGFpMS4AI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "for m in genai.list_models():\n",
        "    if \"embedContent\" in m.supported_generation_methods:\n",
        "        print(m.name)"
      ],
      "metadata": {
        "id": "Km5d13_FS2Q_",
        "execution": {
          "iopub.status.busy": "2024-11-15T16:28:56.509583Z",
          "iopub.execute_input": "2024-11-15T16:28:56.510119Z",
          "iopub.status.idle": "2024-11-15T16:28:56.915206Z",
          "shell.execute_reply.started": "2024-11-15T16:28:56.510077Z",
          "shell.execute_reply": "2024-11-15T16:28:56.913997Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Data\n",
        "\n",
        "Here is a small set of documents you will use to create an embedding database."
      ],
      "metadata": {
        "id": "3XWKXoXwOGxS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "DOCUMENT1 = \"Operating the Climate Control System  Your Googlecar has a climate control system that allows you to adjust the temperature and airflow in the car. To operate the climate control system, use the buttons and knobs located on the center console.  Temperature: The temperature knob controls the temperature inside the car. Turn the knob clockwise to increase the temperature or counterclockwise to decrease the temperature. Airflow: The airflow knob controls the amount of airflow inside the car. Turn the knob clockwise to increase the airflow or counterclockwise to decrease the airflow. Fan speed: The fan speed knob controls the speed of the fan. Turn the knob clockwise to increase the fan speed or counterclockwise to decrease the fan speed. Mode: The mode button allows you to select the desired mode. The available modes are: Auto: The car will automatically adjust the temperature and airflow to maintain a comfortable level. Cool: The car will blow cool air into the car. Heat: The car will blow warm air into the car. Defrost: The car will blow warm air onto the windshield to defrost it.\"\n",
        "DOCUMENT2 = 'Your Googlecar has a large touchscreen display that provides access to a variety of features, including navigation, entertainment, and climate control. To use the touchscreen display, simply touch the desired icon.  For example, you can touch the \"Navigation\" icon to get directions to your destination or touch the \"Music\" icon to play your favorite songs.'\n",
        "DOCUMENT3 = \"Shifting Gears Your Googlecar has an automatic transmission. To shift gears, simply move the shift lever to the desired position.  Park: This position is used when you are parked. The wheels are locked and the car cannot move. Reverse: This position is used to back up. Neutral: This position is used when you are stopped at a light or in traffic. The car is not in gear and will not move unless you press the gas pedal. Drive: This position is used to drive forward. Low: This position is used for driving in snow or other slippery conditions.\"\n",
        "\n",
        "documents = [DOCUMENT1, DOCUMENT2, DOCUMENT3]"
      ],
      "metadata": {
        "id": "k8nsbhFJKmG-",
        "execution": {
          "iopub.status.busy": "2024-11-15T16:33:49.63319Z",
          "iopub.execute_input": "2024-11-15T16:33:49.633732Z",
          "iopub.status.idle": "2024-11-15T16:33:49.642729Z",
          "shell.execute_reply.started": "2024-11-15T16:33:49.633679Z",
          "shell.execute_reply": "2024-11-15T16:33:49.640928Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Creating the embedding database with ChromaDB\n",
        "\n",
        "Create a [custom function](https://docs.trychroma.com/guides/embeddings#custom-embedding-functions) to generate embeddings with the Gemini API. In this task, you are implementing a retrieval system, so the `task_type` for generating the *document* embeddings is `retrieval_document`. Later, you will use `retrieval_query` for the *query* embeddings. Check out the [API reference](https://ai.google.dev/api/embeddings#v1beta.TaskType) for the full list of supported tasks.\n"
      ],
      "metadata": {
        "id": "yDzxArLeOexD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from chromadb import Documents, EmbeddingFunction, Embeddings\n",
        "from google.api_core import retry\n",
        "\n",
        "\n",
        "class GeminiEmbeddingFunction(EmbeddingFunction):\n",
        "    # Specify whether to generate embeddings for documents, or queries\n",
        "    document_mode = True\n",
        "\n",
        "    def __call__(self, input: Documents) -> Embeddings:\n",
        "        if self.document_mode:\n",
        "            embedding_task = \"retrieval_document\"\n",
        "        else:\n",
        "            embedding_task = \"retrieval_query\"\n",
        "\n",
        "        retry_policy = {\"retry\": retry.Retry(predicate=retry.if_transient_error)}\n",
        "\n",
        "        response = genai.embed_content(\n",
        "            model=\"models/text-embedding-004\",\n",
        "            content=input,\n",
        "            task_type=embedding_task,\n",
        "            request_options=retry_policy,\n",
        "        )\n",
        "        return response[\"embedding\"]"
      ],
      "metadata": {
        "id": "mF7Uu1kCQsT0",
        "execution": {
          "iopub.status.busy": "2024-11-15T16:38:39.505652Z",
          "iopub.execute_input": "2024-11-15T16:38:39.507417Z",
          "iopub.status.idle": "2024-11-15T16:38:39.515817Z",
          "shell.execute_reply.started": "2024-11-15T16:38:39.507367Z",
          "shell.execute_reply": "2024-11-15T16:38:39.514453Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now create a [Chroma database client](https://docs.trychroma.com/getting-started) that uses the `GeminiEmbeddingFunction` and populate the database with the documents you defined above."
      ],
      "metadata": {
        "id": "HrDWLyopPNBf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import chromadb\n",
        "\n",
        "DB_NAME = \"googlecardb\"\n",
        "embed_fn = GeminiEmbeddingFunction()\n",
        "embed_fn.document_mode = True\n",
        "\n",
        "chroma_client = chromadb.Client()\n",
        "db = chroma_client.get_or_create_collection(name=DB_NAME, embedding_function=embed_fn)\n",
        "\n",
        "db.add(documents=documents, ids=[str(i) for i in range(len(documents))])"
      ],
      "metadata": {
        "id": "OITXgxZlLoXU",
        "execution": {
          "iopub.status.busy": "2024-11-15T16:40:12.217574Z",
          "iopub.execute_input": "2024-11-15T16:40:12.218104Z",
          "iopub.status.idle": "2024-11-15T16:40:12.705847Z",
          "shell.execute_reply.started": "2024-11-15T16:40:12.218059Z",
          "shell.execute_reply": "2024-11-15T16:40:12.704593Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Confirm that the data was inserted by looking at the database."
      ],
      "metadata": {
        "id": "2QbwFgfXp-fL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "db.count()\n",
        "# You can peek at the data too.\n",
        "# db.peek(1)"
      ],
      "metadata": {
        "id": "kQ9PHUL_l-hf",
        "execution": {
          "iopub.status.busy": "2024-11-15T16:40:39.821101Z",
          "iopub.execute_input": "2024-11-15T16:40:39.822269Z",
          "iopub.status.idle": "2024-11-15T16:40:39.834985Z",
          "shell.execute_reply.started": "2024-11-15T16:40:39.822186Z",
          "shell.execute_reply": "2024-11-15T16:40:39.833345Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Retrieval: Find relevant documents\n",
        "\n",
        "To search the Chroma database, call the `query` method. Note that you also switch to the `retrieval_query` mode of embedding generation.\n"
      ],
      "metadata": {
        "id": "Tu5zRErgsQ8u"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Switch to query mode when generating embeddings.\n",
        "embed_fn.document_mode = False\n",
        "\n",
        "# Search the Chroma DB using the specified query.\n",
        "query = \"How do you use the touchscreen to play music?\"\n",
        "\n",
        "result = db.query(query_texts=[query], n_results=1)\n",
        "[[passage]] = result[\"documents\"]\n",
        "\n",
        "Markdown(passage)"
      ],
      "metadata": {
        "id": "gQdJMbTSLtKE",
        "execution": {
          "iopub.status.busy": "2024-11-15T16:46:28.684781Z",
          "iopub.execute_input": "2024-11-15T16:46:28.685289Z",
          "iopub.status.idle": "2024-11-15T16:46:29.152201Z",
          "shell.execute_reply.started": "2024-11-15T16:46:28.685238Z",
          "shell.execute_reply": "2024-11-15T16:46:29.150859Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Augmented generation: Answer the question\n",
        "\n",
        "Now that you have found a relevant passage from the set of documents (the *retrieval* step), you can now assemble a generation prompt to have the Gemini API *generate* a final answer. Note that in this example only a single passage was retrieved. In practice, especially when the size of your underlying data is large, you will want to retrieve more than one result and let the Gemini model determine what passages are relevant in answering the question. For this reason it's OK if some retrieved passages are not directly related to the question - this generation step should ignore them."
      ],
      "metadata": {
        "id": "s8PNRMpOQkm5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "passage_oneline = passage.replace(\"\\n\", \" \")\n",
        "query_oneline = query.replace(\"\\n\", \" \")\n",
        "\n",
        "# This prompt is where you can specify any guidance on tone, or what topics the model should stick to, or avoid.\n",
        "prompt = f\"\"\"You are a helpful and informative bot that answers questions using text from the reference passage included below.\n",
        "Be sure to respond in a complete sentence, being comprehensive, including all relevant background information.\n",
        "However, you are talking to a non-technical audience, so be sure to break down complicated concepts and\n",
        "strike a friendly and converstional tone. If the passage is irrelevant to the answer, you may ignore it.\n",
        "\n",
        "QUESTION: {query_oneline}\n",
        "PASSAGE: {passage_oneline}\n",
        "\"\"\"\n",
        "print(prompt)"
      ],
      "metadata": {
        "id": "b6_Y-GOymaXu",
        "execution": {
          "iopub.status.busy": "2024-11-15T16:51:30.730209Z",
          "iopub.execute_input": "2024-11-15T16:51:30.730846Z",
          "iopub.status.idle": "2024-11-15T16:51:30.741116Z",
          "shell.execute_reply.started": "2024-11-15T16:51:30.730796Z",
          "shell.execute_reply": "2024-11-15T16:51:30.739171Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now use the `generate_content` method to to generate an answer to the question."
      ],
      "metadata": {
        "id": "VRy6yXzcPxLB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model = genai.GenerativeModel(\"gemini-1.5-flash-latest\")\n",
        "answer = model.generate_content(prompt)\n",
        "Markdown(answer.text)"
      ],
      "metadata": {
        "id": "EwfyxFM6Giy9",
        "execution": {
          "iopub.status.busy": "2024-11-15T17:15:36.614576Z",
          "iopub.execute_input": "2024-11-15T17:15:36.615079Z",
          "iopub.status.idle": "2024-11-15T17:15:37.075101Z",
          "shell.execute_reply.started": "2024-11-15T17:15:36.615036Z",
          "shell.execute_reply": "2024-11-15T17:15:37.07357Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Next steps\n",
        "\n",
        "Congrats on building a Retrieval-Augmented Generation app!\n",
        "\n",
        "To learn more about using embeddings in the Gemini API, check out the [Intro to embeddings](https://ai.google.dev/gemini-api/docs/embeddings) or to learn more fundamentals, study the [embeddings chapter](https://developers.google.com/machine-learning/crash-course/embeddings) of the Machine Learning Crash Course.\n",
        "\n",
        "For a hosted RAG system, check out the [Semantic Retrieval service](https://ai.google.dev/gemini-api/docs/semantic_retrieval) in the Gemini API. You can implement question-answering on your own documents in a single request, or host a database for even faster responses."
      ],
      "metadata": {
        "id": "ThTbjAJ7eGP5"
      }
    }
  ]
}