notebooks/Vanilla_RAG.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mz33G3t6gbOl"
      },
      "source": [
        "# RAG\n",
        "\n",
        "Retrieval-Augmented Generation (RAG) is a technique that combines the strengths of pre-trained language models with the ability to retrieve information from a large corpus of documents. RAG **enables the language model to produce more informed, accurate, and contextually relevant answers** than by relying on its pre-trained knowledge alone.\n",
        "\n",
        "At Cohere, all RAG calls come with... **precise citations**! 🎉\n",
        "The model cites which groups of words, in the RAG chunks, were used to generate the final answer.  \n",
        "These citations make it easy to check where the model’s generated response claims are coming from and they help users gain visibility into the model reasoning.  \n",
        "\n",
        "RAG consists of 3 steps:\n",
        "- Step 1: Indexing and given a user query, retrieve the relevant chunks from the index\n",
        "- Step 2: Optionally, rerank the retrieved chunks\n",
        "- Step 3: Generate the model final answer with **precise citations**, given the retrieved and reranked chunks\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nSB0pnt0gbOo"
      },
      "source": [
        "## Step 0 - Imports & Getting some data\n",
        "\n",
        "In this example, we'll use a recent piece of text, that wasn't in the training data: the Wikipedia page of the movie \"Dune 2\".   \n",
        "\n",
        "In practice, you would typically do RAG on much longer text, that doesn't fit in the context window of the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "H787BXXYvD0a",
        "outputId": "04ef5e04-7760-4d40-deeb-663536b38f20"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m52.8/52.8 kB\u001b[0m \u001b[31m1.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.1/3.1 MB\u001b[0m \u001b[31m33.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ],
      "source": [
        "# we'll use Cohere to cover all building blocks of RAG\n",
        "# TODO: upgrade to \"cohere>5\"\n",
        "%pip install \"cohere<5\" --quiet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rACbepFGgbOo"
      },
      "outputs": [],
      "source": [
        "import cohere\n",
        "API_KEY = \"...\" # fill in your Cohere API key here\n",
        "co = cohere.Client(API_KEY)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QdvbqfFrgbOq",
        "outputId": "3882c95c-46bf-4dcc-99a2-453b3c2fc7c4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for wikipedia (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
          ]
        }
      ],
      "source": [
        "# we'll get some wikipedia data\n",
        "!pip install wikipedia --quiet\n",
        "import wikipedia"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xP-bWt9XgbOq",
        "outputId": "72276fb2-0d6b-415d-af74-452a013ae84b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "The text has roughly 5323 words.\n"
          ]
        }
      ],
      "source": [
        "# let's get the wikipedia article about Dune Part Two\n",
        "article = wikipedia.page('Dune Part Two')\n",
        "text = article.content\n",
        "print(f\"The text has roughly {len(text.split())} words.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-1aJ7hKGgbOr"
      },
      "source": [
        "## Step 1 - Indexing and given a user query, retrieve the relevant chunks from the index\n",
        "\n",
        "We index the document in a vector database. This requires getting the documents, chunking them, embedding, and indexing them in a vector database. Then we retrieved relevant results based on the users' query.\n",
        "\n",
        "### We split the document into chunks of roughly 512 words"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZUph1JX41665",
        "outputId": "6c63a93f-6999-47af-e704-d4a88727bc75"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m256.9/256.9 kB\u001b[0m \u001b[31m6.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m66.6/66.6 kB\u001b[0m \u001b[31m8.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m138.5/138.5 kB\u001b[0m \u001b[31m14.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ],
      "source": [
        "# For chunking let's use langchain to help us split the text\n",
        "%pip install -qU langchain-text-splitters --quiet\n",
        "from langchain_text_splitters import RecursiveCharacterTextSplitter"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uhXW7iHC1-Q6",
        "outputId": "d68ac348-4b73-4c6a-a445-6c510bdb0881"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "The text has been broken down in 91 chunks.\n"
          ]
        }
      ],
      "source": [
        "# Create basic configurations to chunk the text\n",
        "text_splitter = RecursiveCharacterTextSplitter(\n",
        "    chunk_size=512,\n",
        "    chunk_overlap=50,\n",
        "    length_function=len,\n",
        "    is_separator_regex=False,\n",
        ")\n",
        "\n",
        "# Split the text into chunks with some overlap\n",
        "chunks_ = text_splitter.create_documents([text])\n",
        "chunks = [c.page_content for c in chunks_]\n",
        "print(f\"The text has been broken down in {len(chunks)} chunks.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P8g0sE2hgbOs"
      },
      "source": [
        "### Embed every text chunk\n",
        "\n",
        "Cohere embeddings are state-of-the-art."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KEarMPEqgbOs",
        "outputId": "7da0e06d-f637-4470-8e01-6de8249be64b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "We just computed 91 embeddings.\n"
          ]
        }
      ],
      "source": [
        "# Because the texts being embedded are the chunks we are searching over, we set the input type as search_doc\n",
        "model=\"embed-english-v3.0\"\n",
        "response = co.embed(\n",
        "    texts= chunks,\n",
        "    model=model,\n",
        "    input_type=\"search_document\",\n",
        "    embedding_types=['float']\n",
        ")\n",
        "embeddings = response.embeddings.float\n",
        "print(f\"We just computed {len(embeddings)} embeddings.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HM6vKeypgbOs"
      },
      "source": [
        "### Store the embeddings in a vector database\n",
        "\n",
        "We use the simplest vector database ever: a python dictionary using `np.array()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sdW7M8HLvB-9"
      },
      "outputs": [],
      "source": [
        "# We use the simplest vector database ever: a python dictionary\n",
        "!pip install numpy --quiet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H2srFH-IgbOs"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "vector_database = {i: np.array(embedding) for i, embedding in enumerate(embeddings)}\n",
        "# { 0: array([...]), 1: array([...]), 2: array([...]), ..., 10: array([...]) }"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q6NGVurZgbOs"
      },
      "source": [
        "## Given a user query, retrieve the relevant chunks from the vector database\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eC05yJQ7jlek"
      },
      "source": [
        "### Define the user question"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y2HTxspKgbOs"
      },
      "outputs": [],
      "source": [
        "query = \"Name everyone involved in writing the script, directing, and producing 'Dune: Part Two'?\"\n",
        "\n",
        "# Note: the relevant passage in the wikipedia page we're looking for is:\n",
        "# \"[...] Dune: Part Two was originally scheduled to be released on October 20, 2023, but was delayed to November 17, 2023, before moving forward two weeks to November 3, 2023, to adjust to changes in release schedules from other studios. It was later postponed by over four months to March 15, 2024, due to the 2023 Hollywood labor disputes. After the strikes were resolved, the film moved once more up two weeks to March 1, 2024. [...]\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9oULg1tOjjOW"
      },
      "source": [
        "### Embed the user question\n",
        "\n",
        "Cohere embeddings are state-of-the-art."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yrUuS6vXgbOs",
        "outputId": "0c64a930-f817-43c2-d775-1d9145cb304e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "query_embedding:  [-0.068603516, -0.02947998, -0.06274414, -0.015449524, -0.033294678, 0.0056877136, -0.047210693, 0.04714966, -0.024871826, 0.008148193, 0.0770874, 0.023880005, -0.058685303, -0.052520752, 0.012832642, 0.024398804, 0.0053215027, 0.035491943, 0.02961731, -0.0069847107, 0.01083374, -0.0011358261, -0.002199173, 0.018417358, 0.027389526, -0.002691269, -0.026535034, 0.015197754, 0.024368286, 0.03729248, 0.0057754517, -0.02229309, -0.014694214, 0.019989014, -0.0036315918, -0.013793945, 0.02835083, 0.006011963, 0.011428833, 0.008682251, 0.046142578, -0.040039062, -0.032196045, -0.002653122, -0.012580872, -0.0041618347, 0.03111267, -0.016799927, 0.014801025, -0.00030636787, -0.033050537, 0.033966064, -0.016021729, -0.025009155, -0.007534027, -0.017074585, 0.008415222, -0.10620117, 0.019195557, -0.015686035, -0.0043182373, -0.045440674, 0.05404663, 0.030776978, -0.014129639, -0.01499939, -0.007286072, 0.009933472, 0.06390381, 0.02444458, -0.010345459, 0.041931152, 0.032989502, -0.04522705, 0.056610107, 0.0068893433, -0.008911133, 0.012489319, 0.01675415, 0.020065308, 0.018753052, 0.022659302, -0.051849365, -0.04925537, 0.046325684, -0.005268097, 0.0026874542, -0.036712646, 0.009437561, -0.0037841797, -0.01473999, -0.034179688, -0.0011606216, 0.05026245, 0.0020771027, -0.016021729, -0.0044898987, 0.04168701, -0.015205383, 0.019210815, -0.012374878, -0.031311035, 0.03111267, -0.040100098, -0.016479492, 0.020446777, 0.010192871, 0.0037841797, -0.0023765564, 0.015220642, -0.016571045, -0.006454468, 0.037384033, -0.044555664, -0.008262634, 0.019546509, 0.009460449, 0.014701843, 0.02658081, -0.02078247, 0.015571594, 0.013153076, -0.010375977, 0.047912598, 0.005393982, -0.007911682, -0.019378662, 0.023529053, -0.0033550262, -0.04598999, -0.0052871704, 0.040252686, 0.011375427, 0.01550293, -0.004508972, 0.006515503, 0.003370285, -0.022766113, 0.00062561035, -0.0007596016, -0.0015277863, 0.0149002075, 0.061401367, 8.261204e-05, 0.06359863, -0.01537323, 0.007446289, 0.018814087, 0.02507019, 0.024215698, 0.006122589, 0.005886078, -0.03829956, 0.029037476, 0.07720947, 0.016921997, 0.022109985, 0.005958557, 0.028793335, 0.019485474, 0.015174866, 0.026153564, 0.032318115, 0.034210205, 0.027145386, -0.019515991, -0.018661499, 0.020477295, 0.008598328, -0.06573486, -0.037109375, 0.04043579, 0.030471802, -0.0010843277, 0.009757996, 0.026947021, 0.037017822, -0.018234253, -0.0115356445, 0.099365234, 0.027816772, -0.019927979, 0.0020961761, 0.013198853, -0.019073486, 2.7656555e-05, 0.041259766, 0.029510498, -0.016204834, 0.028137207, 0.039489746, 0.034698486, -0.03918457, -0.029418945, 0.02041626, 0.0073432922, -0.018569946, -0.009849548, 0.002861023, 0.030319214, -0.012886047, 0.014671326, -0.035827637, 0.007247925, -0.027709961, -0.022079468, 0.0012960434, 0.015426636, -0.01725769, 0.01525116, 0.025360107, -0.0077400208, -0.039916992, 0.029037476, -0.011154175, 0.007736206, -0.041748047, 0.05343628, 0.007286072, 0.0435791, 0.034301758, -0.047210693, 0.03552246, -0.015327454, 0.029922485, -0.018859863, 0.013053894, -0.028060913, 0.07757568, -0.020462036, 0.070739746, -0.010223389, 0.03604126, 0.02758789, -0.023284912, 0.012184143, 0.029144287, 0.023880005, -0.019378662, -0.0051116943, 0.0048675537, 0.01864624, -0.04397583, -0.007598877, 0.0713501, 0.0115737915, 0.002922058, 0.011619568, 0.017364502, 0.031921387, -0.0019664764, -0.008575439, 0.003484726, -0.09466553, 0.03475952, 0.026611328, -0.039520264, -0.0104522705, -0.005443573, -0.008392334, 0.012908936, 0.0043792725, -0.002456665, -0.028396606, -0.02027893, -0.0005569458, 0.027786255, 0.03427124, -0.0062332153, -0.018203735, 0.019241333, 0.07244873, -0.0028057098, 0.01234436, -0.0018787384, -0.027496338, 0.0015287399, -0.004032135, -0.013748169, -0.01878357, 0.0018053055, -0.01159668, 0.028213501, 0.004776001, 0.042388916, 0.0024280548, 0.017471313, -0.038085938, 0.026321411, 0.02973938, 0.06213379, 0.006401062, 0.036102295, -0.028121948, -0.00869751, -0.016693115, 0.029190063, 0.016784668, -0.008628845, 0.0039634705, -0.0035381317, 0.019500732, 0.025009155, -0.04547119, -0.003572464, 0.05215454, 0.067871094, -0.04257202, -0.02293396, -0.027175903, 0.05340576, 0.019226074, 0.039978027, 0.056121826, -0.028320312, -0.020217896, -0.035003662, 0.03225708, 0.028656006, 0.062347412, 0.12915039, -0.0137786865, 0.0022201538, -0.057434082, -0.04397583, -0.049865723, -0.013160706, -0.03353882, 0.006427765, -0.014823914, -0.008201599, -0.036346436, -0.037353516, -0.010528564, -0.015930176, -0.027572632, 0.0074272156, 0.004547119, -0.024414062, -0.018859863, -0.020095825, 0.029632568, -0.00067043304, -0.044036865, -0.0043411255, -0.005256653, -0.019195557, 0.022262573, -0.00020956993, -0.013877869, -0.011108398, -0.020324707, -0.015808105, -0.025039673, -0.009498596, 0.05090332, 0.0046195984, -0.017150879, 0.04309082, -0.029067993, 0.002670288, -0.00026249886, -0.032409668, -0.053100586, 0.012481689, -0.014633179, 0.0013475418, -0.034332275, 0.038330078, 0.014892578, -0.046936035, 0.021591187, -0.020385742, -0.0052604675, 0.02796936, 0.0014333725, 0.012077332, -0.0118255615, -0.005569458, 0.008491516, 0.009841919, 0.0031318665, -0.003408432, -0.007144928, 0.040374756, -0.0038928986, 0.005279541, -0.008415222, 0.031707764, 0.0140686035, -0.015029907, -0.02810669, -0.0078125, -0.030853271, -0.03201294, 0.021316528, -0.036193848, -0.0423584, 0.0072784424, 0.014801025, 0.0019607544, -0.012367249, -0.009056091, -0.021438599, -0.02645874, 0.038726807, -0.007549286, 0.0049591064, 0.019012451, 0.017791748, -0.009185791, 0.04006958, 0.003107071, -0.0075302124, -0.010375977, -0.009246826, -0.02130127, -0.0056762695, -0.0076789856, 0.010009766, -0.010536194, 0.041107178, 0.0021133423, 0.029891968, 0.01626587, 0.042236328, -0.02784729, -0.032836914, 0.0317688, 0.045715332, 0.000116825104, 0.028030396, 0.007205963, 0.012512207, -0.035583496, -0.048034668, -0.023529053, -0.04953003, 0.0345459, -0.048339844, -0.060272217, -0.004512787, 0.04425049, 0.0076141357, 0.029510498, 0.007396698, 0.003353119, -0.038726807, 0.07183838, -0.026901245, -0.023529053, -0.038085938, 0.068725586, 0.018096924, -0.013534546, 0.05883789, -0.016113281, 0.017944336, 0.041046143, 0.022918701, 0.036499023, 0.015296936, -0.04916382, 0.0075683594, -0.011390686, 0.009735107, -0.0070152283, 0.003129959, -0.032562256, 0.0003478527, -0.0036640167, -0.006893158, -0.016098022, -0.034332275, 0.037750244, -0.010269165, 0.016494751, -0.02394104, 0.03753662, -0.022644043, -0.0008234978, 0.001001358, -0.048217773, 0.04989624, 0.0078125, 0.0044937134, 0.027038574, 0.04736328, -0.02973938, -0.011726379, 0.01348114, 0.021408081, 0.00844574, -0.03741455, -0.015686035, -0.040893555, 0.001452446, -0.025405884, 0.07348633, 0.038238525, -0.019958496, 0.023071289, -0.016403198, -0.08105469, 0.0071029663, -0.019088745, 5.8174133e-05, -0.005569458, 0.01399231, 0.02255249, 0.011222839, 0.00028824806, 0.0066184998, 0.0017499924, -0.009864807, -0.0115737915, 0.053100586, 0.0065231323, 0.001865387, -0.026428223, 0.03692627, 0.025390625, 0.022613525, 0.018722534, 0.007675171, -0.03439331, 0.041625977, -0.01789856, -0.041046143, 0.0051460266, 0.04144287, 0.048553467, 0.054595947, -0.01108551, -0.033935547, -0.026275635, -0.0118255615, -0.021362305, -0.009841919, -0.00724411, 0.028900146, 0.009887695, -0.023803711, 0.016311646, 0.018798828, -0.03668213, 0.046844482, 0.010696411, -0.014717102, -0.008110046, -0.004589081, -0.0028076172, -0.050811768, -0.017196655, -0.03491211, 0.0074005127, -0.038909912, 0.032440186, -0.034362793, -0.008682251, 0.032928467, -0.04626465, -0.009666443, 0.018951416, 0.031951904, -0.003791809, 0.02015686, -0.05532837, -0.005683899, -0.00054216385, -0.0034332275, 0.008659363, 0.02130127, -0.038879395, -0.0033397675, -0.03866577, -0.0049934387, 0.017944336, 0.001496315, 0.019485474, -0.004348755, 0.00046491623, 0.0007157326, 0.035614014, -0.027694702, 0.03692627, -0.008491516, 0.0524292, -0.016662598, -0.0017795563, -0.021575928, -0.018753052, -0.049346924, -0.06652832, 0.04272461, 0.03186035, 0.0011978149, 0.03463745, 0.024002075, 0.02607727, 0.020446777, 0.0256958, 0.026855469, 0.0074005127, -0.067993164, 0.017944336, -0.0039482117, 0.05496216, -0.041412354, 0.014175415, 0.02444458, -0.026412964, 0.057403564, -0.026779175, 0.023254395, 0.03945923, 0.033569336, -0.030258179, -0.039093018, -0.036468506, 0.017105103, 0.009635925, 0.025497437, 0.04156494, -0.02571106, -0.0010414124, -0.005630493, -0.016448975, -0.026733398, 0.001326561, -0.042022705, 0.0012521744, -0.041259766, -0.12182617, -0.03857422, 0.12548828, -0.005947113, -0.020736694, -0.0033855438, 0.03778076, -0.033813477, 0.038970947, 0.003921509, 0.011810303, 0.031982422, -0.032562256, -0.002653122, -0.025009155, -0.03805542, -0.016998291, 0.018173218, 0.0158844, 0.0011739731, 0.048217773, -0.020401001, 0.044708252, -0.017318726, 0.014457703, -0.041809082, 0.010543823, 0.041931152, 0.076293945, -0.054779053, 0.060272217, -0.046936035, 0.02949524, 0.00554657, 0.041534424, -0.013046265, -0.056152344, 0.010406494, 0.02973938, -0.023727417, -0.022476196, -0.024734497, -0.013168335, 0.060424805, 0.011787415, 0.018997192, -0.043426514, -0.00077724457, -0.010154724, 0.017150879, -0.01171875, -0.022476196, 0.0034255981, -0.0026454926, 0.004837036, -0.0043296814, 0.02619934, -0.021560669, -0.039733887, -0.022415161, -0.06817627, -0.023223877, -0.018585205, -0.015319824, 0.012588501, 0.0064353943, -0.013748169, 0.043304443, 0.002626419, -0.029373169, -0.016784668, -0.026184082, 0.05847168, 0.034179688, 0.03842163, -0.05493164, -0.017486572, 0.016540527, 0.03164673, 0.089904785, 0.013534546, -0.07684326, -0.024108887, 0.07434082, 0.030395508, 0.007091522, 0.07373047, 0.012527466, -0.010856628, -0.01828003, -0.045196533, 0.00065279007, -0.0637207, 0.010726929, 0.023880005, -0.0030708313, -0.012298584, 0.027236938, -0.04928589, 0.023071289, 0.008674622, -0.023529053, -0.015838623, -0.010543823, 0.012168884, 0.014854431, -0.05834961, -0.06088257, -0.012313843, 0.035461426, 0.02027893, 0.019348145, -0.014602661, -0.02104187, -0.0309906, 0.001405716, -0.019973755, -0.00157547, -0.003944397, 0.0009326935, -0.02078247, -0.015731812, -0.044433594, 0.03390503, 0.057159424, 0.018585205, -0.023895264, -0.0057029724, 0.0049552917, 0.013412476, 0.022399902, 0.010154724, 0.0519104, 0.06591797, 0.018341064, 0.012161255, -0.05810547, -0.043304443, -0.031173706, 0.0023860931, -0.003944397, 0.11425781, -0.031036377, 0.019989014, -0.038635254, -0.025939941, 0.035064697, 0.041168213, 0.03161621, -0.069885254, -0.04537964, 0.028945923, -0.023162842, 0.019226074, -0.028442383, 0.015594482, -0.019256592, -0.0046463013, 0.034240723, 0.009124756, 0.05718994, 0.031219482, 0.02154541, 0.009590149, 0.00076818466, 0.04849243, -0.029129028, -0.03375244, -0.023391724, -0.028381348, -0.029708862, -0.0132369995, 0.010353088, 0.020263672, -0.030807495, 0.01007843, -0.03704834, 0.023376465, -0.03665161, 0.03741455, 0.015144348, 0.057281494, 0.03137207, 0.048431396, 0.021194458, 0.008110046, -0.03540039, -0.015312195, 0.022384644, 0.0065956116, 0.008056641, 0.0018348694, -0.009246826, 0.030380249, 0.0003862381, 0.0051841736, 0.04486084, 0.017807007, 0.0026130676, 0.07977295, 0.05419922, 0.062194824, 0.02633667, 0.024841309, -0.041625977, -0.005897522, 0.04031372, -0.055908203, 0.0026226044, -0.05340576, -0.05496216, 0.011474609, -0.006954193, -0.013122559, 0.019714355, -0.07159424, 0.031173706, 0.0034255981, -0.0034103394, 0.0440979, 0.011779785, -0.007827759, -0.03173828, -0.020950317, -0.030166626, -0.035308838, 0.030792236, 0.04525757, -0.028701782, -0.011100769, -0.02331543, -0.0357666, -0.025680542, 0.0011911392, 0.01940918, 0.05706787, 0.028381348, 0.007133484, -0.07733154, -0.007686615, 0.03869629, 0.0066833496, 0.008842468, 0.03439331, -0.014282227, 0.0357666, -0.004737854, -0.039794922, -0.0070381165, 0.02670288, 0.0107421875, 0.016189575, -0.06555176, -0.0138549805, 0.0008363724, -0.016693115, 0.006904602, -0.020263672, -0.030426025, 0.008453369, -0.046173096, -0.01802063, -0.013595581, -0.0044288635, -0.0039978027, -0.0044898987, 0.0007619858, 0.003921509, 0.0053977966, 0.020385742, -0.012329102, -0.023803711, -0.0057525635, 0.038330078, -0.014549255, -0.06298828, -0.047607422, 0.039245605, -0.06781006, -0.035217285, -0.009056091, 0.019927979, -0.003932953, -0.020309448, -0.017044067, 0.018127441, -8.624792e-05, -0.043182373, 0.009590149, 0.035308838, 0.031951904, 0.0011615753, -0.042022705, 0.079956055, 0.026687622, 0.013542175, -0.0074157715, -0.00983429, -0.0022563934, 0.07373047, 0.059387207, 0.03488159, 0.0071372986, -0.06427002, -0.0546875, -0.02482605, 0.11071777, -0.021072388, 0.01626587, -0.049713135, 0.061553955, -0.016860962, 0.051971436, -0.012962341, -0.0011711121, -0.014198303, -0.0061149597, -0.005836487, 0.00022387505, -0.027618408, 0.019836426, 0.009933472, 0.02368164, -0.020309448, -0.0049591064, -0.008628845, -0.03253174, -0.017684937, 0.02468872, -0.0023498535, 0.01448822, 0.061920166, 0.031707764, -0.0026416779, -0.040985107, -0.06335449, -0.036071777, 0.05404663, -0.0044136047, -0.0146102905, -0.0033416748, 0.028671265, -0.012771606, -0.0016565323, -0.0038909912, -0.02407837, -0.009857178, 0.0014467239, -0.008720398, -0.006011963, 0.032073975, -0.033325195, 0.014862061, -0.017227173, -0.018753052, -0.0060424805, 0.022567749, -0.017654419, -0.017562866, -0.07244873, -0.0881958, 0.050476074, 0.02609253, -0.032409668, 0.07458496, 0.009399414, 0.009117126, -0.031051636, -0.03451538, -0.004219055, -0.05718994, 0.020080566, -0.025421143, -0.010948181, 0.06341553, -0.009231567, -0.021697998, -0.009719849, 0.012802124, -0.020370483, 0.0034389496, 0.018859863, -0.025680542, 0.0013141632, 0.068603516, -0.021026611, 0.021881104, -0.0395813, -0.0019073486, 0.0056037903, -0.032348633]\n"
          ]
        }
      ],
      "source": [
        "# Because the text being embedded is the search query, we set the input type as search_query\n",
        "response = co.embed(\n",
        "    texts=[query],\n",
        "    model=model,\n",
        "    input_type=\"search_query\",\n",
        "    embedding_types=['float']\n",
        ")\n",
        "query_embedding = response.embeddings.float[0]\n",
        "print(\"query_embedding: \", query_embedding)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8K8B87CGgbOt"
      },
      "source": [
        "### Retrieve the most relevant chunks from the vector database\n",
        "\n",
        "We use cosine similarity to find the most similar chunks"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nik3es32gbOt",
        "outputId": "a1c30024-52e1-42c7-8836-a2c590559aca"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "similarity scores:  [0.6953257882233425, 0.3713410510180273, 0.46501499776898964, 0.5448546916785195, 0.4014738351361969, 0.3231420292334584, 0.3179003053384008, 0.42799691553367775, 0.18882594531435784, 0.36868801306504106, 0.3404040737300553, 0.3852837621219358, 0.2600249419491577, 0.3723244353775111, 0.3631492691137214, 0.47574774051439606, 0.40415422750911745, 0.4149923346201023, 0.5014741934381444, 0.3549433331883204, 0.32072714802512714, 0.14770872479410424, 0.585277816615252, 0.6999636953772764, 0.7722295084104617, 0.4895347049465806, 0.5170096485954725, 0.7137817366881455, 0.5224900699612323, 0.5914632581598285, 0.2657897083381463, 0.6462342489537262, 0.6317222315431096, 0.5982303530756702, 0.5138265091630297, 0.41385121172723643, 0.4293941094100836, 0.4173182546482015, 0.42621236706314475, 0.4428474375355954, 0.35058541576139896, 0.3578709652019502, 0.3930157841938308, 0.3564608202848675, 0.23016661533167404, 0.4933441863421645, 0.41037089239250985, 0.39993051898770193, 0.3119997063424595, 0.2677143729521374, 0.3700866951454496, 0.46727994925061545, 0.4393343280374425, 0.42111290117172434, 0.4485349189824285, 0.4710573736688592, 0.24169956903740436, 0.3840442910806355, 0.14284631817675886, 0.5381588054138154, 0.431113882725076, 0.5189547209048608, 0.3950667224233914, 0.32429768756510174, 0.4370358125161736, 0.18727062244331039, 0.5206375682478743, 0.5175737635701252, 0.5326043981628349, 0.45586923626994363, 0.21667338125532032, 0.16459878595959285, 0.22236726481673777, 0.5187259906958807, 0.2884444442338396, 0.286407544555338, 0.2313840178160818, 0.2057731158935257, 0.5973876998341746, 0.42904243401792086, 0.4081217538000544, 0.5330523063972133, 0.45080561486977405, 0.414703452285757, 0.2569028899107211, 0.5087916806929323, 0.14159076456040554, 0.46505779053352697, 0.556364222182839, 0.35464351181035236, 0.40174477023626]\n",
            "Here are the indices of the top 10 chunks after retrieval:  [24 27 23  0 31 32 33 78 29 22]\n",
            "Here are the top 10 chunks after retrieval: \n",
            "== stunt coordinator. Dune: Part Two was produced by Villeneuve, Mary Parent, and Cale Boyter, with Tanya Lapointe, Brian Herbert, Byron Merritt, Kim Herbert, Thomas Tull, Jon Spaihts, Richard P. Rubinstein, John Harrison, and Herbert W. Gain serving as executive producers and Kevin J. Anderson as creative consultant. Legendary CEO Joshua Grode confirmed in April 2019 that they plan to make a sequel, adding that \"there's a logical place to stop the [first] movie before the book is over\".In December 2020,\n",
            "== that.\"On October 26, 2021, Legendary officially greenlit Dune: Part Two, with a spokesperson for the company stating, \"We would not have gotten to this point without the extraordinary vision of Denis and the amazing work of his talented crew, the writers, our stellar cast, our partners at Warner Bros., and of course the fans! Here's to more Dune.\" Production work had occurred back-to-back with the first film, as Villeneuve and his wife Lapointe immediately took a flight to Budapest in order to begin\n",
            "== series. Villeneuve ultimately secured a two-film deal with Warner Bros. Pictures, in the same style as the two-part adaption of Stephen King's It in 2017 and 2019. In January 2019, Joe Walker was confirmed to be serving as the film's editor. Other crew included Brad Riker as supervising art director, Patrice Vermette as production designer, Paul Lambert as visual effects supervisor, Gerd Nefzer as special effects supervisor, and Thomas Struthers as stunt coordinator. Dune: Part Two was produced by\n",
            "== Dune: Part Two is a 2024 American epic science fiction film directed and produced by Denis Villeneuve, who co-wrote the screenplay with Jon Spaihts. The sequel to Dune (2021) adapts the 1965 novel Dune by Frank Herbert and follows Paul Atreides as he unites with the Fremen people of the desert planet Arrakis to wage war against House Harkonnen. Timothée Chalamet, Rebecca Ferguson, Josh Brolin, Stellan Skarsgård, Dave Bautista, Zendaya, Charlotte Rampling, and Javier Bardem reprise their roles from the first\n",
            "== Eric Roth was hired to co-write the screenplay in April 2017 for the Dune films, and Jon Spaihts was later confirmed to be co-writing the script alongside Roth and Villeneuve. Game of Thrones language creator David Peterson was confirmed to be developing languages for the film in April 2019. Villeneuve and Peterson had created the Chakobsa language, which was used by actors on set. In November 2019, Spaihts stepped down as showrunner for Dune: Prophecy to focus on Dune: Part Two. In June 2020, Greig Fraser\n",
            "== on Dune: Part Two. In June 2020, Greig Fraser said, \"It's a fully formed story in itself with places to go. It's a fully standalone epic film that people will get a lot out of when they see it\". Between the release of Dune and the confirmation of Dune: Part Two, Villeneuve started working the script in a way that production could begin immediately once the film was greenlit. By February 2021, Roth created a full treatment for the sequel, with writing beginning that August. He confirmed that Feyd-Rautha\n",
            "== that August. He confirmed that Feyd-Rautha would appear in the film, and stated he will be a \"very important character\". In March 2022, Villeneuve had mostly finished writing the screenplay. Craig Mazin and Roth wrote additional literary material for the film.Villeneuve stated that the film would continue directly from the first, and specifically described it as being the \"second part.\" He described the film as being an \"epic war movie\", adding that while the first film was more \"contemplative\", the second\n",
            "== On the review aggregator website Rotten Tomatoes, 93% of 378 critics' reviews are positive, with an average rating of 8.4/10. The website's consensus reads: \"Visually thrilling and narratively epic, Dune: Part Two continues Denis Villeneuve's adaptation of the beloved sci-fi series in spectacular form.\" Metacritic, which uses a weighted average, assigned the film a score of 79 out of 100, based on 62 critics, indicating \"generally favorable\" reviews. Audiences surveyed by CinemaScore gave the film an\n",
            "== theatrical experience is at the very heart of the cinematic language for me.\" With Dune: Part Two being greenlit, Villeneuve said that his primary concern was to complete the filming as soon as possible, with the earliest he expected to start in the last quarter of 2022. However, he noted that production would be facilitated by the work already established on the first film, which would help expedite production.\n",
            "== By November 2016, Legendary Pictures had obtained the film and TV rights for the Dune franchise, based on the eponymous 1965 novel by Frank Herbert. Vice chair of worldwide production for Legendary Mary Parent began discussing with Denis Villeneuve about directing a film adaptation, quickly hiring him after realizing his passion for Dune. By February 2018, Villeneuve was confirmed to be hired as director, and intended to adapt the novel as a two-part film series. Villeneuve ultimately secured a two-film\n"
          ]
        }
      ],
      "source": [
        "def cosine_similarity(a, b):\n",
        "    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))\n",
        "\n",
        "# Calculate similarity between the user question & each chunk\n",
        "similarities = [cosine_similarity(query_embedding, chunk) for chunk in embeddings]\n",
        "print(\"similarity scores: \", similarities)\n",
        "\n",
        "# Get indices of the top 10 most similar chunks\n",
        "sorted_indices = np.argsort(similarities)[::-1]\n",
        "\n",
        "# Keep only the top 10 indices\n",
        "top_indices = sorted_indices[:10]\n",
        "print(\"Here are the indices of the top 10 chunks after retrieval: \", top_indices)\n",
        "\n",
        "# Retrieve the top 10 most similar chunks\n",
        "top_chunks_after_retrieval = [chunks[i] for i in top_indices]\n",
        "print(\"Here are the top 10 chunks after retrieval: \")\n",
        "for t in top_chunks_after_retrieval:\n",
        "    print(\"== \" + t)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qzcpds3VgbOt"
      },
      "source": [
        "## Step 2 - Rerank the chunks retrieved from the vector database\n",
        "\n",
        "We rerank the 10 chunks retrieved from the vector database. Reranking boosts retrieval accuracy.\n",
        "\n",
        "Reranking lets us go from 10 chunks retrieved from the vector database, to the 3 most relevant chunks."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2J4LywVygbOt",
        "outputId": "7a4c89bf-fc5e-409f-9304-fce006b9d8bf"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Here are the top 3 chunks after rerank: \n",
            "== Dune: Part Two is a 2024 American epic science fiction film directed and produced by Denis Villeneuve, who co-wrote the screenplay with Jon Spaihts. The sequel to Dune (2021) adapts the 1965 novel Dune by Frank Herbert and follows Paul Atreides as he unites with the Fremen people of the desert planet Arrakis to wage war against House Harkonnen. Timothée Chalamet, Rebecca Ferguson, Josh Brolin, Stellan Skarsgård, Dave Bautista, Zendaya, Charlotte Rampling, and Javier Bardem reprise their roles from the first\n",
            "== stunt coordinator. Dune: Part Two was produced by Villeneuve, Mary Parent, and Cale Boyter, with Tanya Lapointe, Brian Herbert, Byron Merritt, Kim Herbert, Thomas Tull, Jon Spaihts, Richard P. Rubinstein, John Harrison, and Herbert W. Gain serving as executive producers and Kevin J. Anderson as creative consultant. Legendary CEO Joshua Grode confirmed in April 2019 that they plan to make a sequel, adding that \"there's a logical place to stop the [first] movie before the book is over\".In December 2020,\n",
            "== series. Villeneuve ultimately secured a two-film deal with Warner Bros. Pictures, in the same style as the two-part adaption of Stephen King's It in 2017 and 2019. In January 2019, Joe Walker was confirmed to be serving as the film's editor. Other crew included Brad Riker as supervising art director, Patrice Vermette as production designer, Paul Lambert as visual effects supervisor, Gerd Nefzer as special effects supervisor, and Thomas Struthers as stunt coordinator. Dune: Part Two was produced by\n"
          ]
        }
      ],
      "source": [
        "response = co.rerank(\n",
        "    query=query,\n",
        "    documents=top_chunks_after_retrieval,\n",
        "    top_n=3,\n",
        "    model=\"rerank-english-v2.0\",\n",
        ")\n",
        "\n",
        "top_chunks_after_rerank = [result.document['text'] for result in response]\n",
        "print(\"Here are the top 3 chunks after rerank: \")\n",
        "for t in top_chunks_after_rerank:\n",
        "    print(\"== \" + t)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KuPL0VUXgbOt"
      },
      "source": [
        "## Step 3 - Generate the model final answer, given the retrieved and reranked chunks"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oCNXWH8GgbOt"
      },
      "outputs": [],
      "source": [
        "# preamble containing instructions about the task and the desired style for the output.\n",
        "preamble = \"\"\"\n",
        "## Task & Context\n",
        "You help people answer their questions and other requests interactively. You will be asked a very wide array of requests on all kinds of topics. You will be equipped with a wide range of search engines or similar tools to help you, which you use to research your answer. You should focus on serving the user's needs as best you can, which will be wide-ranging.\n",
        "\n",
        "## Style Guide\n",
        "Unless the user asks for a different style of answer, you should answer in full sentences, using proper grammar and spelling.\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BevatShtgbOt",
        "outputId": "af71f4a9-787a-4ee3-9598-20692fb3bf16"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Final answer:\n",
            "Here are the key crew members involved in 'Dune: Part Two':\n",
            "\n",
            "- Denis Villeneuve: director and producer\n",
            "- Jon Spaihts: co-writer of the screenplay\n",
            "- Mary Parent and Cale Boyter: producers \n",
            "- Tanya Lapointe, Brian Herbert, Byron Merritt, Kim Herbert, Thomas Tull, Richard P. Rubinstein, John Harrison, Herbert W. Gain and Kevin J. Anderson: executive producers \n",
            "- Joe Walker: editor\n",
            "- Brad Riker: supervising art director\n",
            "- Patrice Vermette: production designer\n",
            "- Paul Lambert: visual effects supervisor\n",
            "- Gerd Nefzer: special effects supervisor\n",
            "- Thomas Struthers: stunt coordinator. \n",
            "\n",
            "A number of crew members from the first film returned for the sequel, which is set to be released in 2024.\n"
          ]
        }
      ],
      "source": [
        "# retrieved documents\n",
        "documents = [\n",
        "    {\"title\": \"chunk 0\", \"snippet\": top_chunks_after_rerank[0]},\n",
        "    {\"title\": \"chunk 1\", \"snippet\": top_chunks_after_rerank[1]},\n",
        "    {\"title\": \"chunk 2\", \"snippet\": top_chunks_after_rerank[2]},\n",
        "  ]\n",
        "\n",
        "# get model response\n",
        "response = co.chat(\n",
        "  message=query,\n",
        "  documents=documents,\n",
        "  preamble=preamble,\n",
        "  model=\"command-r\",\n",
        "  temperature=0.3\n",
        ")\n",
        "\n",
        "print(\"Final answer:\")\n",
        "print(response.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "20wcn-EjlXZd"
      },
      "source": [
        "Note: this is indeed the answer you'd expect, and here was the passage of text in wikipedia explaining it!\n",
        "\n",
        "\" [...] Dune: Part Two was originally scheduled to be released on October 20, 2023, but was delayed to November 17, 2023, before moving forward two weeks to November 3, 2023, to adjust to changes in release schedules from other studios. It was later postponed by over four months to March 15, 2024, due to the 2023 Hollywood labor disputes. After the strikes were resolved, the film moved once more up two weeks to March 1, 2024. [...]\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RoSVDXSsgbOt"
      },
      "source": [
        "## Bonus: Citations come for free with Cohere! 🎉\n",
        "\n",
        "At Cohere, all RAG calls come with... precise citations! 🎉\n",
        "The model cites which groups of words, in the RAG chunks, were used to generate the final answer.  \n",
        "These citations make it easy to check where the model’s generated response claims are coming from.  \n",
        "They help users gain visibility into the model reasoning, and sanity check the final model generation.  \n",
        "These citations are optional — you can decide to ignore them.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BVTuQdmDgbOt",
        "outputId": "f843b262-d8bb-45ba-cbfb-9915da104eda"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Citations that support the final answer:\n",
            "{'start': 63, 'end': 79, 'text': 'Denis Villeneuve', 'document_ids': ['doc_0']}\n",
            "{'start': 81, 'end': 102, 'text': 'director and producer', 'document_ids': ['doc_0']}\n",
            "{'start': 105, 'end': 116, 'text': 'Jon Spaihts', 'document_ids': ['doc_0']}\n",
            "{'start': 118, 'end': 145, 'text': 'co-writer of the screenplay', 'document_ids': ['doc_0']}\n",
            "{'start': 148, 'end': 159, 'text': 'Mary Parent', 'document_ids': ['doc_1']}\n",
            "{'start': 164, 'end': 175, 'text': 'Cale Boyter', 'document_ids': ['doc_1']}\n",
            "{'start': 177, 'end': 186, 'text': 'producers', 'document_ids': ['doc_1']}\n",
            "{'start': 190, 'end': 204, 'text': 'Tanya Lapointe', 'document_ids': ['doc_1']}\n",
            "{'start': 206, 'end': 219, 'text': 'Brian Herbert', 'document_ids': ['doc_1']}\n",
            "{'start': 221, 'end': 234, 'text': 'Byron Merritt', 'document_ids': ['doc_1']}\n",
            "{'start': 236, 'end': 247, 'text': 'Kim Herbert', 'document_ids': ['doc_1']}\n",
            "{'start': 249, 'end': 260, 'text': 'Thomas Tull', 'document_ids': ['doc_1']}\n",
            "{'start': 262, 'end': 283, 'text': 'Richard P. Rubinstein', 'document_ids': ['doc_1']}\n",
            "{'start': 285, 'end': 298, 'text': 'John Harrison', 'document_ids': ['doc_1']}\n",
            "{'start': 300, 'end': 315, 'text': 'Herbert W. Gain', 'document_ids': ['doc_1']}\n",
            "{'start': 320, 'end': 337, 'text': 'Kevin J. Anderson', 'document_ids': ['doc_1']}\n",
            "{'start': 339, 'end': 358, 'text': 'executive producers', 'document_ids': ['doc_1']}\n",
            "{'start': 362, 'end': 372, 'text': 'Joe Walker', 'document_ids': ['doc_2']}\n",
            "{'start': 374, 'end': 380, 'text': 'editor', 'document_ids': ['doc_2']}\n",
            "{'start': 383, 'end': 393, 'text': 'Brad Riker', 'document_ids': ['doc_2']}\n",
            "{'start': 395, 'end': 419, 'text': 'supervising art director', 'document_ids': ['doc_2']}\n",
            "{'start': 422, 'end': 438, 'text': 'Patrice Vermette', 'document_ids': ['doc_2']}\n",
            "{'start': 440, 'end': 459, 'text': 'production designer', 'document_ids': ['doc_2']}\n",
            "{'start': 462, 'end': 474, 'text': 'Paul Lambert', 'document_ids': ['doc_2']}\n",
            "{'start': 476, 'end': 501, 'text': 'visual effects supervisor', 'document_ids': ['doc_2']}\n",
            "{'start': 504, 'end': 515, 'text': 'Gerd Nefzer', 'document_ids': ['doc_2']}\n",
            "{'start': 517, 'end': 543, 'text': 'special effects supervisor', 'document_ids': ['doc_2']}\n",
            "{'start': 546, 'end': 562, 'text': 'Thomas Struthers', 'document_ids': ['doc_2']}\n",
            "{'start': 564, 'end': 582, 'text': 'stunt coordinator.', 'document_ids': ['doc_2']}\n",
            "{'start': 686, 'end': 691, 'text': '2024.', 'document_ids': ['doc_0']}\n"
          ]
        }
      ],
      "source": [
        "print(\"Citations that support the final answer:\")\n",
        "for cite in response.citations:\n",
        "  print(cite)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IueXaIJggbOu",
        "outputId": "c816af51-74be-42c9-e94e-9820bbf95f79"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Here are the key crew members involved in 'Dune: Part Two':\n",
            "\n",
            "- **Denis Villeneuve**[1]: **director and producer**[1]\n",
            "- **Jon Spaihts**[1]: **co-writer of the screenplay**[1]\n",
            "- **Mary Parent**[2] and **Cale Boyter**[2]: **producers**[2] \n",
            "- **Tanya Lapointe**[2], **Brian Herbert**[2], **Byron Merritt**[2], **Kim Herbert**[2], **Thomas Tull**[2], **Richard P. Rubinstein**[2], **John Harrison**[2], **Herbert W. Gain**[2] and **Kevin J. Anderson**[2]: **executive producers**[2] \n",
            "- **Joe Walker**[3]: **editor**[3]\n",
            "- **Brad Riker**[3]: **supervising art director**[3]\n",
            "- **Patrice Vermette**[3]: **production designer**[3]\n",
            "- **Paul Lambert**[3]: **visual effects supervisor**[3]\n",
            "- **Gerd Nefzer**[3]: **special effects supervisor**[3]\n",
            "- **Thomas Struthers**[3]: **stunt coordinator.**[3] \n",
            "\n",
            "A number of crew members from the first film returned for the sequel, which is set to be released in **2024.**[1]\n",
            "\n",
            "[1] source: \"Dune: Part Two is a 2024 American epic science fiction film directed and produced by Denis Villeneuve, who co-wrote the screenplay with Jon Spaihts. The sequel to Dune (2021) adapts the 1965 novel Dune by Frank Herbert and follows Paul Atreides as he unites with the Fremen people of the desert planet Arrakis to wage war against House Harkonnen. Timothée Chalamet, Rebecca Ferguson, Josh Brolin, Stellan Skarsgård, Dave Bautista, Zendaya, Charlotte Rampling, and Javier Bardem reprise their roles from the first\"\n",
            "[2] source: \"stunt coordinator. Dune: Part Two was produced by Villeneuve, Mary Parent, and Cale Boyter, with Tanya Lapointe, Brian Herbert, Byron Merritt, Kim Herbert, Thomas Tull, Jon Spaihts, Richard P. Rubinstein, John Harrison, and Herbert W. Gain serving as executive producers and Kevin J. Anderson as creative consultant. Legendary CEO Joshua Grode confirmed in April 2019 that they plan to make a sequel, adding that \"there's a logical place to stop the [first] movie before the book is over\".In December 2020,\"\n",
            "[3] source: \"series. Villeneuve ultimately secured a two-film deal with Warner Bros. Pictures, in the same style as the two-part adaption of Stephen King's It in 2017 and 2019. In January 2019, Joe Walker was confirmed to be serving as the film's editor. Other crew included Brad Riker as supervising art director, Patrice Vermette as production designer, Paul Lambert as visual effects supervisor, Gerd Nefzer as special effects supervisor, and Thomas Struthers as stunt coordinator. Dune: Part Two was produced by\"\n"
          ]
        }
      ],
      "source": [
        "def insert_citations_in_order(text, citations):\n",
        "    \"\"\"\n",
        "    A helper function to pretty print citations.\n",
        "    \"\"\"\n",
        "    offset = 0\n",
        "    document_id_to_number = {}\n",
        "    citation_number = 0\n",
        "    modified_citations = []\n",
        "\n",
        "    # Process citations, assigning numbers based on unique document_ids\n",
        "    for citation in citations:\n",
        "        citation_numbers = []\n",
        "        for document_id in citation[\"document_ids\"]:\n",
        "            if document_id not in document_id_to_number:\n",
        "                citation_number = int(document_id.split('_')[1]) # extract the document id\n",
        "                document_id_to_number[document_id] = citation_number\n",
        "            citation_numbers.append(document_id_to_number[document_id])\n",
        "\n",
        "        # Adjust start/end with offset\n",
        "        start, end = citation['start'] + offset, citation['end'] + offset\n",
        "        placeholder = ''.join([f'[{number}]' for number in citation_numbers])\n",
        "        # Bold the cited text and append the placeholder\n",
        "        modification = f'**{text[start:end]}**{placeholder}'\n",
        "        # Replace the cited text with its bolded version + placeholder\n",
        "        text = text[:start] + modification + text[end:]\n",
        "        # Update the offset for subsequent replacements\n",
        "        offset += len(modification) - (end - start)\n",
        "\n",
        "    # Prepare citations for listing at the bottom, ensuring unique document_ids are listed once\n",
        "    unique_citations = {number: doc_id for doc_id, number in document_id_to_number.items()}\n",
        "    citation_list = '\\n'.join([f'[{doc_id}] source: \"{documents[doc_id][\"snippet\"]}\"' for doc_id, number in in dict(sorted(unique_citations.items(), key=lambda item: item[1])).items()])\n",
        "    text_with_citations = f'{text}\\n\\n{citation_list}'\n",
        "\n",
        "    return text_with_citations\n",
        "\n",
        "\n",
        "print(insert_citations_in_order(response.text, response.citations))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Kp4c_HkYIEn_"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "hackathon_docs_3",
      "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": 0
}