RAG-with-Gemini-workshop.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5098c5d4-e2df-4ee7-993d-26bdfe07b6da",
   "metadata": {},
   "source": [
    "# RAG and AI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a3b331b-5882-48b7-a6d1-f639e512dc7e",
   "metadata": {},
   "source": [
    "## Embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8922b9f-eff8-4348-9be0-b82d23b33b00",
   "metadata": {},
   "source": [
    "Embeddings can be generated locally using a number of python modules (most commonly `sentence-transformers`), however depending on device this can be slow. For now we'll use Gemini's embedding generation functions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acc33f77-df80-4a34-8ec1-a40f4569ecb7",
   "metadata": {},
   "source": [
    "This workbook requires a Google Gemini API key:\n",
    "- Google provide Gemini with a free tier available to anyone with a Google account.\n",
    "- To get your own key visit [aistudio.google.com/app/apikey](https://aistudio.google.com/app/apikey) and click \"Create API key\"\n",
    "\n",
    "This API key needs to be saved to your project's .env file. Open your .env and add *GEMINI_API_KEY=\"MY-API-KEY-HERE\"* or run\n",
    "```\n",
    "# Linux\n",
    "echo 'GEMINI_API_KEY=\"MY-API-KEY-HERE\"' >> .env\n",
    "\n",
    "# Windows\n",
    "echo GEMINI_API_KEY=\"MY-API-KEY-HERE\"> .env\n",
    "```\n",
    "\n",
    "This workbook also requires the Gemini python module, google-generativeai"
   ]
  },
  {
   "cell_type": "code",
   "id": "3dfe260f-5f4a-4804-8ce6-b7759772c778",
   "metadata": {},
   "source": [
    "import os\n",
    "import google.generativeai as genai\n",
    "from dotenv import load_dotenv                  # Allow us to load environment variables"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "c25cf831-0c6f-45c8-a0a4-05172602955d",
   "metadata": {},
   "source": [
    "load_dotenv()\n",
    "gemini_api_key = os.environ[\"GEMINI_API_KEY\"]\n",
    "genai.configure(api_key=gemini_api_key)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "45c45ae4-b20a-401e-9793-18ae8f94b948",
   "metadata": {},
   "source": [
    "### Generating a single embedding"
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Google provide a free embeddings API, rate limited to 1500 calls per minute. For our purposes this is perfect.",
   "id": "d067edc026db70c4"
  },
  {
   "cell_type": "code",
   "id": "bbfaa332-624c-4dbf-b4ad-549059071139",
   "metadata": {},
   "source": [
    "result = genai.embed_content(\n",
    "    model=\"models/text-embedding-004\",          # Embedding model to use\n",
    "    content=\"What is the meaning of life?\",     # Document text to create an embedding of\n",
    "    task_type=\"retrieval_document\",             # Task type (\"retreival_query\", \"retreival_document\", \"clustering\", \"semantic_similarity\", \"classification\")\n",
    "    title=\"Embedding of a single string\"        # Parameter for retreival_document tasks. Ostensibly the document title\n",
    ")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "b45e73eb-7807-4145-b9d5-9f6e7be7731e",
   "metadata": {},
   "source": [
    "result[\"embedding\"][:10]    # Shortened to the first 10 elements"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "644d290a-726f-42d8-b693-a96e41309b9c",
   "metadata": {},
   "source": [
    "# Loading and chunking documents"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c288dee-d38b-4806-8b28-c1773e940a10",
   "metadata": {},
   "source": [
    "### Loading pdfs\n",
    "\n",
    "This can be acomplished with a number of tools, including OCR, but we'll be using `pypdf`"
   ]
  },
  {
   "cell_type": "code",
   "id": "cbe9885d-e51b-4881-9cda-f435466fd296",
   "metadata": {},
   "source": [
    "from pypdf import PdfReader    # `pip install pypdf`\n",
    "from os import listdir\n",
    "from os.path import isfile, join\n",
    "import re"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "2a85f76e-4691-4fed-ac7b-9361e0146528",
   "metadata": {},
   "source": [
    "#### Loading a single document"
   ]
  },
  {
   "cell_type": "code",
   "id": "4272f3af-8a26-4a4a-869b-ca7b39d4788b",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "reader = PdfReader('./data/sample/ResearchBazaarQueensland2024.pdf')\n",
    "print(f\"File contains {str(len(reader.pages))} pages\")\n",
    "for page in reader.pages:\n",
    "    print(page.extract_text(0))"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "7beb5a30-4d21-48fd-b3c1-416693976fa8",
   "metadata": {},
   "source": [
    "#### Approaching cleaning"
   ]
  },
  {
   "cell_type": "code",
   "id": "d57b89d2-3f8a-4e14-9b72-bfcef2c11525",
   "metadata": {},
   "source": [
    "def clean_text(text: str) -> str:\n",
    "    pass"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "21ec83e3-6eed-4d75-8bd9-b7c7110ce273",
   "metadata": {},
   "source": [
    "#### Loading all documents"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8622e9bf-4563-4cf2-b5cb-62e4d195c9b0",
   "metadata": {},
   "source": [
    "Load all documents and clean the text in preperation for chunking"
   ]
  },
  {
   "cell_type": "code",
   "id": "15b91f7e-c1ce-4866-a9c6-b039ce3fe18e",
   "metadata": {},
   "source": [
    "data_dir = './data/ethics/'\n",
    "documents = {}"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "16441b39-5d56-488c-bc83-4254bdf46762",
   "metadata": {},
   "source": "# Iterate through all documents in our data directory, read them with PdfReader, and clean the text of each page\n",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "284a5876a6582877",
   "metadata": {},
   "source": [
    "#### Perform Chunking"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c313757-80b2-4abc-9999-386e4224ba3c",
   "metadata": {},
   "source": [
    "Determining paragraphs from a PDF is [*hard*](https://pypdf.readthedocs.io/en/stable/user/extract-text.html#why-text-extraction-is-hard). We're going to assume two sucessive newlines means a new paragraph.<br>\n",
    "The paragraph problem is more easily solvable with OCR, however that is beyond the scope of this workbook.\n",
    "\n",
    "Determining suitable chunk size is almost as hard, and just as situation specific. Embeddings models also enforce limits on the size of the input, further complicating things.\n",
    "For our purposes we'll use 128 tokens as a chunk size, as this represents a moderately sized paragraph; roughly the resolution we want to be able to trace back to.\n",
    "\n",
    "Embedding quality can be improved by actions such as prepending the tail of the previous chunk to the new one in order to provide more context, or even including the title of the document. As Gemini's text embedding API provides a specific parameter for passing the title of a document, we won't worry about providing additional context for the moment."
   ]
  },
  {
   "cell_type": "code",
   "id": "f6c76570a4fbab53",
   "metadata": {},
   "source": [
    "chunk_token_limit = 128\n",
    "chunks = {}"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "b8cf675e-e286-470a-91df-b255e5befd8b",
   "metadata": {},
   "source": "# Iterate through each document's text, adding each word to a chunk as we go. Chunks should still be identified by document IDs.\n",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "fb474688eae21db9",
   "metadata": {},
   "source": [
    "# Generating and Storing embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48959cb08111fe5a",
   "metadata": {},
   "source": [
    "#### Generating a single embedding\n",
    "As shown earlier, generating a single embedding is a simple task, so how do we scale?"
   ]
  },
  {
   "cell_type": "code",
   "id": "ff6e0181d7482764",
   "metadata": {
    "scrolled": true
   },
   "source": "genai.embed_content(model=\"models/text-embedding-004\", content=\"What is the meaning of life?\", task_type=\"retrieval_document\", title=\"Embedding of a single string\")",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "6afcf3464ff2f665",
   "metadata": {},
   "source": [
    "## How to store embeddings?\n",
    "\n",
    "There are many vector-compatible databases, as well as databases aimed soley at storing and retrieving vectors. One of the most lightweight options is Chroma DB."
   ]
  },
  {
   "cell_type": "code",
   "id": "4ec53cdacb84dd97",
   "metadata": {},
   "source": [
    "import chromadb\n",
    "chroma = chromadb.PersistentClient(path=\"./chroma-db\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "aa27c85634775ea7",
   "metadata": {},
   "source": [
    "# Delete the collection if it already exists\n",
    "chroma.delete_collection(name=\"gemini-embeddings\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a3d6e3ca6aabaead",
   "metadata": {},
   "source": [
    "#### Generating embeddings for all of our documents"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96c873d0-03b0-4b4f-ae3f-0fe5fc7c9c29",
   "metadata": {},
   "source": [
    "Create a Chroma DB collection, with the embedding generation function set to use Gemini"
   ]
  },
  {
   "cell_type": "code",
   "id": "b0045e77bb6bddff",
   "metadata": {},
   "source": [
    "from chromadb.utils.embedding_functions.google_embedding_function import GoogleGenerativeAiEmbeddingFunction\n",
    "db = chroma.get_or_create_collection(name=\"gemini-embeddings\", embedding_function=GoogleGenerativeAiEmbeddingFunction(api_key=gemini_api_key, model_name=\"models/text-embedding-004\", task_type=\"RETRIEVAL_DOCUMENT\"))"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "05db1af0-f1f7-4b03-b6c9-445964c69861",
   "metadata": {},
   "source": "Add each chunk to the database (this may take a while)"
  },
  {
   "cell_type": "code",
   "id": "549ab4621c4a7aea",
   "metadata": {},
   "source": "# Iterate through the chunks (dictionary of lists) clean the text and add each chunk to the database\n",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "8e77863e00d642de",
   "metadata": {},
   "source": [
    "#### Querying our documents\n",
    "Now that we have a database full of embeddings, we can get to querying"
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "What do we want to ask?",
   "id": "d8493b458453b9cb"
  },
  {
   "cell_type": "code",
   "id": "7f6a49fd3322674f",
   "metadata": {},
   "source": "question = \"What are user's perceptions of their twitter data usage?\"",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Create an embedding of the question. Note the changed \"task_type\" from \"retrieval_document\" to \"retrieval_query\"",
   "id": "7239ca19795d13dd"
  },
  {
   "cell_type": "code",
   "id": "81a4c01990d1ef67",
   "metadata": {},
   "source": [
    "question_embedding = genai.embed_content(model=\"models/text-embedding-004\", content=question, task_type=\"retrieval_query\")['embedding']\n",
    "results = db.query(query_embeddings=question_embedding, n_results=5)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "We have results! The results dictionary provides different useful pieces of information, including the matched ids, the original text \"documents\", as well as distance between the matches (the higher the score the better the match).",
   "id": "9af6b001ad651dd6"
  },
  {
   "cell_type": "markdown",
   "id": "d3e1ae4559d122a3",
   "metadata": {},
   "source": [
    "#### Passing documents to Gemini\n",
    "\n",
    "Now that we have our matched documents, how can we use them for discussions?"
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "The code below defines the role of the AI. Try changing the prompt, you'll likely get quite different responses.",
   "id": "b11d253b0865ce2c"
  },
  {
   "cell_type": "code",
   "id": "242c6c1259230b3e",
   "metadata": {},
   "source": [
    "prompt = f\"\"\"\n",
    "You are a informative bot that answers questions using only the document texts provided. You will be provided multiple, use all as needed.\n",
    "If you do not know the answer from the provided context and documents, do not use your prior knowledge, and tell the user that you do not know.\n",
    "Generate at least one paragraph.\n",
    "Be comprehensive, helpful, human-readable, and provide detailed background information in your answer.\n",
    "QUESTION: {question}\n",
    "TEXTS: {results['documents'][0]}\n",
    "\"\"\""
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "99b3975fc18af318",
   "metadata": {},
   "source": [
    "model = genai.GenerativeModel(model_name=\"gemini-1.5-flash\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Print out the result of our question",
   "id": "32cf728eef7e0b6e"
  },
  {
   "cell_type": "code",
   "id": "be3f5f3af8a5fb8a",
   "metadata": {},
   "source": [
    "result = model.generate_content(prompt).text\n",
    "print(result)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "9ea5018228cb5590",
   "metadata": {},
   "source": [
    "#### But which document did that come from?\n",
    "\n",
    "We know it came from one of the (up to) 5 supplied documents, but can we get more detailed than that?"
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Rearrange the results data structure to include the document IDs",
   "id": "136838e6b1ca60dc"
  },
  {
   "cell_type": "code",
   "id": "1297cd0b7d5c1c2c",
   "metadata": {},
   "source": [
    "results_with_ids = {\n",
    "    results['ids'][0][i]: results['documents'][0][i] for i in range(len(results['ids'][0]))\n",
    "}"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "b26d90c2f6cd7ff1",
   "metadata": {},
   "source": [
    "results_with_ids"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "question = \"What are user's perceptions of their twitter data usage?\"",
   "id": "e6c5598eb0eff772",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "c1b4d86c941f728d",
   "metadata": {},
   "source": [
    "prompt = f\"\"\"\n",
    "You are a informative bot that answers questions using only the document texts provided. You will be provided multiple, use all as needed.\n",
    "If you do not know the answer from the provided context and documents, do not use your prior knowledge, and tell the user that you do not know.\n",
    "Generate at least one paragraph.\n",
    "Be comprehensive, helpful, human-readable, and provide detailed background information in your answer. Provide at least one paragraph.\n",
    "Finish each response with \"SOURCE:\" followed by a python list of all document IDs that you used to answer the question.\n",
    "QUESTION: {question}\n",
    "TEXTS: {results_with_ids}\n",
    "\"\"\""
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "c70d51dc1cdb4d16",
   "metadata": {},
   "source": [
    "result = model.generate_content(prompt).text\n",
    "print(result)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "We have a new response, this time with the IDs of the documents used*. We can now look up the document containing the answer in our Chroma DB instance.\n",
    "\n",
    "*Keep in mind AI is fallible, and may not have correctly identified the exact document it is using"
   ],
   "id": "5a2a577d0bc4a05c"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "a83d924638f0484a",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "#### What about follow up questions?",
   "id": "10f1d2db3c89e528"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Using the chat API",
   "id": "bf95c60b6837e831"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "chat = model.start_chat(\n",
    "    history=[\n",
    "        {\"role\": \"user\", \"parts\": prompt},\n",
    "        {\"role\": \"model\", \"parts\": result}\n",
    "    ]\n",
    ")\n",
    "\n",
    "chat_response = chat.send_message(\"\")"
   ],
   "id": "9dc328ae95e0b2f6",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Further improvements\n",
    "\n",
    "- Include the page number and paragraph number in the Chroma DB entry\n",
    "- Open a pdf viewer and highlight the matched text\n",
    "- Looping discussion with additional documents being added as needed"
   ],
   "id": "a6b457ec84b6d75b"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "",
   "id": "a12aea4af97f8269"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}