learn/generation/langchain/handbook/08-langchain-retrieval-agent.ipynb

{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "TXC2wBpCU9f7"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pinecone-io/examples/blob/master/learn/generation/langchain/handbook/08-langchain-retrieval-agent.ipynb) [![Open nbviewer](https://raw.githubusercontent.com/pinecone-io/examples/master/assets/nbviewer-shield.svg)](https://nbviewer.org/github/pinecone-io/examples/blob/master/learn/generation/langchain/handbook/08-langchain-retrieval-agent.ipynb)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "bhWwrfbbVGOA"
      },
      "source": [
        "#### [LangChain Handbook](https://pinecone.io/learn/langchain)\n",
        "\n",
        "# Retrieval Agents\n",
        "\n",
        "We've seen in previous chapters how powerful [retrieval augmentation](https://www.pinecone.io/learn/langchain-retrieval-augmentation/) and [conversational agents](https://www.pinecone.io/learn/langchain-agents/) can be. They become even more impressive when we begin using them together.\n",
        "\n",
        "Conversational agents can struggle with data freshness, knowledge about specific domains, or accessing internal documentation. By coupling agents with retrieval augmentation tools we no longer have these problems.\n",
        "\n",
        "One the other side, using \"naive\" retrieval augmentation without the use of an agent means we will retrieve contexts with *every* query. Again, this isn't always ideal as not every query requires access to external knowledge.\n",
        "\n",
        "Merging these methods gives us the best of both worlds. In this notebook we'll learn how to do this.\n",
        "\n",
        "To begin, we must install the prerequisite libraries that we will be using in this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pva9ehKXUpU2",
        "outputId": "3bbcf2dd-1889-412f-d45a-f56945ac4f2f"
      },
      "outputs": [],
      "source": [
        "!pip install -qU \\\n",
        "    openai==1.6.1 \\\n",
        "    pinecone-client==3.1.0 \\\n",
        "    langchain==0.1.1 \\\n",
        "    langchain-community==0.0.13 \\\n",
        "    tiktoken==0.5.2 \\\n",
        "    datasets==2.12.0"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "ZTgrOQziXUto"
      },
      "source": [
        "## Building the Knowledge Base"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "qNyRsz0ZXXaq"
      },
      "source": [
        "We start by constructing our knowledge base. We'll use a mostly prepared dataset called **S**tanford **Qu**estion-**A**nswering **D**ataset (SQuAD) hosted on Hugging Face *Datasets*. We download it like so:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "laSDMjqQXuj-",
        "outputId": "5272df99-eb4b-4ec2-c513-504e067be2b6"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "79d95a944b44423d9a2dd0082cdd80b0",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Downloading readme:   0%|          | 0.00/7.83k [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "488d023c133946f18514d9eab88d18a6",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Downloading data files:   0%|          | 0/2 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "b516de7e929a48ebaed4eda6bc8e0697",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Downloading data:   0%|          | 0.00/14.5M [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "98435a29bc384a218fd34eed7cc593cf",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Downloading data:   0%|          | 0.00/1.82M [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "9733b641eea7468c95fb4d013a42c08e",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Extracting data files:   0%|          | 0/2 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "ee1fb919fdb44f258cdee51bd469e1b6",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Generating train split:   0%|          | 0/87599 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "72b2cf694b7942358c4d8ebb974b6c95",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Generating validation split:   0%|          | 0/10570 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/plain": [
              "Dataset({\n",
              "    features: ['id', 'title', 'context', 'question', 'answers'],\n",
              "    num_rows: 87599\n",
              "})"
            ]
          },
          "execution_count": 1,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "from datasets import load_dataset\n",
        "\n",
        "data = load_dataset('squad', split='train')\n",
        "data"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "8casaLEpX18U"
      },
      "source": [
        "The dataset does contain duplicate contexts, which we can remove like so:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 511
        },
        "id": "JnWZTcJiXzor",
        "outputId": "42659cec-23c3-4349-b007-676da99d834d"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>id</th>\n",
              "      <th>title</th>\n",
              "      <th>context</th>\n",
              "      <th>question</th>\n",
              "      <th>answers</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>5733be284776f41900661182</td>\n",
              "      <td>University_of_Notre_Dame</td>\n",
              "      <td>Architecturally, the school has a Catholic cha...</td>\n",
              "      <td>To whom did the Virgin Mary allegedly appear i...</td>\n",
              "      <td>{'text': ['Saint Bernadette Soubirous'], 'answ...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>5733be284776f4190066117f</td>\n",
              "      <td>University_of_Notre_Dame</td>\n",
              "      <td>Architecturally, the school has a Catholic cha...</td>\n",
              "      <td>What is in front of the Notre Dame Main Building?</td>\n",
              "      <td>{'text': ['a copper statue of Christ'], 'answe...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>5733be284776f41900661180</td>\n",
              "      <td>University_of_Notre_Dame</td>\n",
              "      <td>Architecturally, the school has a Catholic cha...</td>\n",
              "      <td>The Basilica of the Sacred heart at Notre Dame...</td>\n",
              "      <td>{'text': ['the Main Building'], 'answer_start'...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>5733be284776f41900661181</td>\n",
              "      <td>University_of_Notre_Dame</td>\n",
              "      <td>Architecturally, the school has a Catholic cha...</td>\n",
              "      <td>What is the Grotto at Notre Dame?</td>\n",
              "      <td>{'text': ['a Marian place of prayer and reflec...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>5733be284776f4190066117e</td>\n",
              "      <td>University_of_Notre_Dame</td>\n",
              "      <td>Architecturally, the school has a Catholic cha...</td>\n",
              "      <td>What sits on top of the Main Building at Notre...</td>\n",
              "      <td>{'text': ['a golden statue of the Virgin Mary'...</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "                         id                     title  \\\n",
              "0  5733be284776f41900661182  University_of_Notre_Dame   \n",
              "1  5733be284776f4190066117f  University_of_Notre_Dame   \n",
              "2  5733be284776f41900661180  University_of_Notre_Dame   \n",
              "3  5733be284776f41900661181  University_of_Notre_Dame   \n",
              "4  5733be284776f4190066117e  University_of_Notre_Dame   \n",
              "\n",
              "                                             context  \\\n",
              "0  Architecturally, the school has a Catholic cha...   \n",
              "1  Architecturally, the school has a Catholic cha...   \n",
              "2  Architecturally, the school has a Catholic cha...   \n",
              "3  Architecturally, the school has a Catholic cha...   \n",
              "4  Architecturally, the school has a Catholic cha...   \n",
              "\n",
              "                                            question  \\\n",
              "0  To whom did the Virgin Mary allegedly appear i...   \n",
              "1  What is in front of the Notre Dame Main Building?   \n",
              "2  The Basilica of the Sacred heart at Notre Dame...   \n",
              "3                  What is the Grotto at Notre Dame?   \n",
              "4  What sits on top of the Main Building at Notre...   \n",
              "\n",
              "                                             answers  \n",
              "0  {'text': ['Saint Bernadette Soubirous'], 'answ...  \n",
              "1  {'text': ['a copper statue of Christ'], 'answe...  \n",
              "2  {'text': ['the Main Building'], 'answer_start'...  \n",
              "3  {'text': ['a Marian place of prayer and reflec...  \n",
              "4  {'text': ['a golden statue of the Virgin Mary'...  "
            ]
          },
          "execution_count": 2,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "data = data.to_pandas()\n",
        "data.head()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 528
        },
        "id": "er2nM4vsYD4R",
        "outputId": "d58b5640-be25-495a-b90c-585940c851bf"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>id</th>\n",
              "      <th>title</th>\n",
              "      <th>context</th>\n",
              "      <th>question</th>\n",
              "      <th>answers</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>5733be284776f41900661182</td>\n",
              "      <td>University_of_Notre_Dame</td>\n",
              "      <td>Architecturally, the school has a Catholic cha...</td>\n",
              "      <td>To whom did the Virgin Mary allegedly appear i...</td>\n",
              "      <td>{'text': ['Saint Bernadette Soubirous'], 'answ...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>5</th>\n",
              "      <td>5733bf84d058e614000b61be</td>\n",
              "      <td>University_of_Notre_Dame</td>\n",
              "      <td>As at most other universities, Notre Dame's st...</td>\n",
              "      <td>When did the Scholastic Magazine of Notre dame...</td>\n",
              "      <td>{'text': ['September 1876'], 'answer_start': [...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>10</th>\n",
              "      <td>5733bed24776f41900661188</td>\n",
              "      <td>University_of_Notre_Dame</td>\n",
              "      <td>The university is the major seat of the Congre...</td>\n",
              "      <td>Where is the headquarters of the Congregation ...</td>\n",
              "      <td>{'text': ['Rome'], 'answer_start': [119]}</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>15</th>\n",
              "      <td>5733a6424776f41900660f51</td>\n",
              "      <td>University_of_Notre_Dame</td>\n",
              "      <td>The College of Engineering was established in ...</td>\n",
              "      <td>How many BS level degrees are offered in the C...</td>\n",
              "      <td>{'text': ['eight'], 'answer_start': [487]}</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>20</th>\n",
              "      <td>5733a70c4776f41900660f64</td>\n",
              "      <td>University_of_Notre_Dame</td>\n",
              "      <td>All of Notre Dame's undergraduate students are...</td>\n",
              "      <td>What entity provides help with the management ...</td>\n",
              "      <td>{'text': ['Learning Resource Center'], 'answer...</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "                          id                     title  \\\n",
              "0   5733be284776f41900661182  University_of_Notre_Dame   \n",
              "5   5733bf84d058e614000b61be  University_of_Notre_Dame   \n",
              "10  5733bed24776f41900661188  University_of_Notre_Dame   \n",
              "15  5733a6424776f41900660f51  University_of_Notre_Dame   \n",
              "20  5733a70c4776f41900660f64  University_of_Notre_Dame   \n",
              "\n",
              "                                              context  \\\n",
              "0   Architecturally, the school has a Catholic cha...   \n",
              "5   As at most other universities, Notre Dame's st...   \n",
              "10  The university is the major seat of the Congre...   \n",
              "15  The College of Engineering was established in ...   \n",
              "20  All of Notre Dame's undergraduate students are...   \n",
              "\n",
              "                                             question  \\\n",
              "0   To whom did the Virgin Mary allegedly appear i...   \n",
              "5   When did the Scholastic Magazine of Notre dame...   \n",
              "10  Where is the headquarters of the Congregation ...   \n",
              "15  How many BS level degrees are offered in the C...   \n",
              "20  What entity provides help with the management ...   \n",
              "\n",
              "                                              answers  \n",
              "0   {'text': ['Saint Bernadette Soubirous'], 'answ...  \n",
              "5   {'text': ['September 1876'], 'answer_start': [...  \n",
              "10          {'text': ['Rome'], 'answer_start': [119]}  \n",
              "15         {'text': ['eight'], 'answer_start': [487]}  \n",
              "20  {'text': ['Learning Resource Center'], 'answer...  "
            ]
          },
          "execution_count": 3,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "data.drop_duplicates(subset='context', keep='first', inplace=True)\n",
        "data.head()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "B2_Pt7N6Zg2X"
      },
      "source": [
        "### Initialize the Embedding Model and Vector DB"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "bGoS84KYZnSK"
      },
      "source": [
        "We'll be using OpenAI's `text-embedding-ada-002` model initialize via LangChain and the Pinecone vector DB. We start by initializing the embedding model, for this we need an [OpenAI API key](https://platform.openai.com/).\n",
        "\n",
        "*(Note that OpenAI is a paid service and so running the remainder of this notebook may incur some small cost)*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "U57x2_87YSpb",
        "outputId": "cf4c99af-24b7-4bf5-97ea-71091c8fc2ce"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from getpass import getpass\n",
        "from langchain.embeddings.openai import OpenAIEmbeddings\n",
        "\n",
        "# get API key from top-right dropdown on OpenAI website\n",
        "OPENAI_API_KEY = os.getenv(\"OPENAI_API_KEY\") or getpass(\"Enter your OpenAI API key: \")\n",
        "model_name = 'text-embedding-ada-002'\n",
        "\n",
        "embed = OpenAIEmbeddings(\n",
        "    model=model_name,\n",
        "    openai_api_key=OPENAI_API_KEY\n",
        ")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "JQTfOTR6aBRS"
      },
      "source": [
        "Now we create our vector DB to store our vectors. For this we need to get a [free Pinecone API key](https://app.pinecone.io) — the API key can be found in the \"API Keys\" button found in the left navbar of the Pinecone dashboard."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [],
      "source": [
        "from pinecone import Pinecone\n",
        "\n",
        "# initialize connection to pinecone (get API key at app.pinecone.io)\n",
        "api_key = os.getenv(\"PINECONE_API_KEY\") or getpass(\"Enter your Pinecone API key: \")\n",
        "\n",
        "# configure client\n",
        "pc = Pinecone(api_key=api_key)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now we setup our index specification, this allows us to define the cloud provider and region where we want to deploy our index. You can find a list of all [available providers and regions here](https://docs.pinecone.io/docs/projects)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [],
      "source": [
        "from pinecone import ServerlessSpec\n",
        "\n",
        "spec = ServerlessSpec(\n",
        "    cloud=\"aws\", region=\"us-east-1\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Creating an index, we set `dimension` equal to to dimensionality of Ada-002 (`1536`), and use a `metric` also compatible with Ada-002 (this can be either `cosine` or `dotproduct`). We also pass our `spec` to index initialization."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "C3wrG-9yaJel",
        "outputId": "842bf46f-fd0f-4322-8d90-a7fd769b7687"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'dimension': 1536,\n",
              " 'index_fullness': 0.0,\n",
              " 'namespaces': {},\n",
              " 'total_vector_count': 0}"
            ]
          },
          "execution_count": 8,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import time\n",
        "\n",
        "index_name = \"langchain-retrieval-agent\"\n",
        "existing_indexes = [\n",
        "    index_info[\"name\"] for index_info in pc.list_indexes()\n",
        "]\n",
        "\n",
        "# check if index already exists (it shouldn't if this is first time)\n",
        "if index_name not in existing_indexes:\n",
        "    # if does not exist, create index\n",
        "    pc.create_index(\n",
        "        index_name,\n",
        "        dimension=1536,  # dimensionality of ada 002\n",
        "        metric='dotproduct',\n",
        "        spec=spec\n",
        "    )\n",
        "    # wait for index to be initialized\n",
        "    while not pc.describe_index(index_name).status['ready']:\n",
        "        time.sleep(1)\n",
        "\n",
        "# connect to index\n",
        "index = pc.Index(index_name)\n",
        "time.sleep(1)\n",
        "# view index stats\n",
        "index.describe_index_stats()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "AD5IGOoLaVx7"
      },
      "source": [
        "We should see that the new Pinecone index has a `total_vector_count` of `0`, as we haven't added any vectors yet.\n",
        "\n",
        "## Indexing\n",
        "\n",
        "We can perform the indexing task using the LangChain vector store object. But for now it is much faster to do it via the Pinecone python client directly. We will do this in batches of `100` or more."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "321ebcd192fc41e0a7a2ba9c0e7b8556",
            "7c898f43b39545879013df760aa4f6f1",
            "1b5fe655986842a6afd4a2eb2c0f47fc",
            "76872ca737b64616aeacc90f2f7fc655",
            "0679733cce4f4b5eb25f8f56dc7ebe22",
            "71551fc297b54fb4a21d787d117e01a1",
            "b90283522431447b925b226431bb4fa6",
            "a03d372e8b4f456e8663e54c818e0fe9",
            "6fc773ccbc594fddbc1f2f0adaa96186",
            "07b47869f727455fb5ba67e99767c04c",
            "4b2e12a33584489690aaf69f51cd4a3b"
          ]
        },
        "id": "AhDcbRGTaWPi",
        "outputId": "dca5a4d0-5d19-4542-c26c-3073bfa13f2a"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "862314956bfc40af91a2f2e8d68a0fdd",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "  0%|          | 0/189 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "from tqdm.auto import tqdm\n",
        "\n",
        "batch_size = 100\n",
        "\n",
        "texts = []\n",
        "metadatas = []\n",
        "\n",
        "for i in tqdm(range(0, len(data), batch_size)):\n",
        "    # get end of batch\n",
        "    i_end = min(len(data), i+batch_size)\n",
        "    batch = data.iloc[i:i_end]\n",
        "    # first get metadata fields for this record\n",
        "    metadatas = [{\n",
        "        'title': record['title'],\n",
        "        'text': record['context']\n",
        "    } for j, record in batch.iterrows()]\n",
        "    # get the list of contexts / documents\n",
        "    documents = batch['context']\n",
        "    # create document embeddings\n",
        "    embeds = embed.embed_documents(documents)\n",
        "    # get IDs\n",
        "    ids = batch['id']\n",
        "    # add everything to pinecone\n",
        "    index.upsert(vectors=zip(ids, embeds, metadatas))"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "jDUnLdy1b7G1"
      },
      "source": [
        "We've indexed everything, now we can check the number of vectors in our index like so:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SiccGZKAb_Qo",
        "outputId": "5cafd8c8-771f-46d0-e261-12c9b5b1b058"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'dimension': 1536,\n",
              " 'index_fullness': 0.0,\n",
              " 'namespaces': {'': {'vector_count': 18891}},\n",
              " 'total_vector_count': 18891}"
            ]
          },
          "execution_count": 10,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "index.describe_index_stats()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "b-3oolT5cCR8"
      },
      "source": [
        "## Creating a Vector Store and Querying"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "DcZ12U06cCH5"
      },
      "source": [
        "Now that we've build our index we can switch back over to LangChain. We start by initializing a vector store using the same index we just built. We do that like so:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "0MBJ477-cFNw"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/Users/jamesbriggs/opt/anaconda3/envs/ml/lib/python3.9/site-packages/langchain_community/vectorstores/pinecone.py:74: UserWarning: Passing in `embedding` as a Callable is deprecated. Please pass in an Embeddings object instead.\n",
            "  warnings.warn(\n"
          ]
        }
      ],
      "source": [
        "from langchain.vectorstores import Pinecone\n",
        "\n",
        "text_field = \"text\"  # the metadata field that contains our text\n",
        "\n",
        "# initialize the vector store object\n",
        "vectorstore = Pinecone(\n",
        "    index, embed.embed_query, text_field\n",
        ")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "3K3xRthWcXzW"
      },
      "source": [
        "As in previous examples, we can use the `similarity_search` method to do a pure semantic search (without the generation component)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uITMZtzschJF",
        "outputId": "88fae02f-6f82-47d7-a153-21b3e0615709"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[Document(page_content=\"In 1919 Father James Burns became president of Notre Dame, and in three years he produced an academic revolution that brought the school up to national standards by adopting the elective system and moving away from the university's traditional scholastic and classical emphasis. By contrast, the Jesuit colleges, bastions of academic conservatism, were reluctant to move to a system of electives. Their graduates were shut out of Harvard Law School for that reason. Notre Dame continued to grow over the years, adding more colleges, programs, and sports teams. By 1921, with the addition of the College of Commerce, Notre Dame had grown from a small college to a university with five colleges and a professional law school. The university continued to expand and add new residence halls and buildings with each subsequent president.\", metadata={'title': 'University_of_Notre_Dame'}),\n",
              " Document(page_content='The College of Engineering was established in 1920, however, early courses in civil and mechanical engineering were a part of the College of Science since the 1870s. Today the college, housed in the Fitzpatrick, Cushing, and Stinson-Remick Halls of Engineering, includes five departments of study – aerospace and mechanical engineering, chemical and biomolecular engineering, civil engineering and geological sciences, computer science and engineering, and electrical engineering – with eight B.S. degrees offered. Additionally, the college offers five-year dual degree programs with the Colleges of Arts and Letters and of Business awarding additional B.A. and Master of Business Administration (MBA) degrees, respectively.', metadata={'title': 'University_of_Notre_Dame'}),\n",
              " Document(page_content='Since 2005, Notre Dame has been led by John I. Jenkins, C.S.C., the 17th president of the university. Jenkins took over the position from Malloy on July 1, 2005. In his inaugural address, Jenkins described his goals of making the university a leader in research that recognizes ethics and building the connection between faith and studies. During his tenure, Notre Dame has increased its endowment, enlarged its student body, and undergone many construction projects on campus, including Compton Family Ice Arena, a new architecture hall, additional residence halls, and the Campus Crossroads, a $400m enhancement and expansion of Notre Dame Stadium.', metadata={'title': 'University_of_Notre_Dame'})]"
            ]
          },
          "execution_count": 12,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "query = \"when was the college of engineering in the University of Notre Dame established?\"\n",
        "\n",
        "vectorstore.similarity_search(\n",
        "    query,  # our search query\n",
        "    k=3  # return 3 most relevant docs\n",
        ")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "-zGF6YsgczqT"
      },
      "source": [
        "Looks like we're getting good results. Let's take a look at how we can begin integrating this into a conversational agent."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "tFsIOm73dcOI"
      },
      "source": [
        "## Initializing the Conversational Agent"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "XMv6TXWkdfNR"
      },
      "source": [
        "Our conversational agent needs a Chat LLM, conversational memory, and a `RetrievalQA` chain to initialize. We create these using:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "zMRs9Klic5-Y"
      },
      "outputs": [],
      "source": [
        "from langchain.chat_models import ChatOpenAI\n",
        "from langchain.chains.conversation.memory import ConversationBufferWindowMemory\n",
        "from langchain.chains import RetrievalQA\n",
        "\n",
        "# chat completion llm\n",
        "llm = ChatOpenAI(\n",
        "    openai_api_key=OPENAI_API_KEY,\n",
        "    model_name='gpt-3.5-turbo',\n",
        "    temperature=0.0\n",
        ")\n",
        "# conversational memory\n",
        "conversational_memory = ConversationBufferWindowMemory(\n",
        "    memory_key='chat_history',\n",
        "    k=5,\n",
        "    return_messages=True\n",
        ")\n",
        "# retrieval qa chain\n",
        "qa = RetrievalQA.from_chain_type(\n",
        "    llm=llm,\n",
        "    chain_type=\"stuff\",\n",
        "    retriever=vectorstore.as_retriever()\n",
        ")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "-ySfWyZLdboX"
      },
      "source": [
        "Using these we can generate an answer using the `run` method:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "LaYSq0V-dxHw",
        "outputId": "a85730a3-ab15-49b6-8aba-6da774fd3e3f"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'The College of Engineering at the University of Notre Dame was established in 1920.'"
            ]
          },
          "execution_count": 14,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "qa.run(query)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "DtSXR5RXdyU0"
      },
      "source": [
        "But this isn't yet ready for our conversational agent. For that we need to convert this retrieval chain into a tool. We do that like so:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "FwCYrS4duqBW"
      },
      "outputs": [],
      "source": [
        "from langchain.agents import Tool\n",
        "\n",
        "tools = [\n",
        "    Tool(\n",
        "        name='Knowledge Base',\n",
        "        func=qa.run,\n",
        "        description=(\n",
        "            'use this tool when answering general knowledge queries to get '\n",
        "            'more information about the topic'\n",
        "        )\n",
        "    )\n",
        "]"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "wXi_0ipTvM_l"
      },
      "source": [
        "Now we can initialize the agent like so:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "JaKTzPUEvOoy"
      },
      "outputs": [],
      "source": [
        "from langchain.agents import initialize_agent\n",
        "\n",
        "agent = initialize_agent(\n",
        "    agent='chat-conversational-react-description',\n",
        "    tools=tools,\n",
        "    llm=llm,\n",
        "    verbose=True,\n",
        "    max_iterations=3,\n",
        "    early_stopping_method='generate',\n",
        "    memory=conversational_memory\n",
        ")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "WbXl-AzVvszB"
      },
      "source": [
        "With that our retrieval augmented conversational agent is ready and we can begin using it."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "IlxUBWKcvzeP"
      },
      "source": [
        "### Using the Conversational Agent"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "ZZapCP4Pv2kz"
      },
      "source": [
        "To make queries we simply call the `agent` directly."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RJoAhy76vzAB",
        "outputId": "31ae7589-e66a-4835-d1d8-555bb59c962c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
            "\u001b[32;1m\u001b[1;3m{\n",
            "    \"action\": \"Knowledge Base\",\n",
            "    \"action_input\": \"When was the College of Engineering in the University of Notre Dame established?\"\n",
            "}\u001b[0m\n",
            "Observation: \u001b[36;1m\u001b[1;3mThe College of Engineering at the University of Notre Dame was established in 1920.\u001b[0m\n",
            "Thought:\u001b[32;1m\u001b[1;3m{\n",
            "    \"action\": \"Final Answer\",\n",
            "    \"action_input\": \"The College of Engineering at the University of Notre Dame was established in 1920.\"\n",
            "}\u001b[0m\n",
            "\n",
            "\u001b[1m> Finished chain.\u001b[0m\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "{'input': 'when was the college of engineering in the University of Notre Dame established?',\n",
              " 'chat_history': [],\n",
              " 'output': 'The College of Engineering at the University of Notre Dame was established in 1920.'}"
            ]
          },
          "execution_count": 17,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "agent(query)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "YcMqa9Va2hU6"
      },
      "source": [
        "Looks great, now what if we ask it a non-general knowledge question?"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "85vipqC02deV",
        "outputId": "345c724e-eaea-4a20-9445-99794bc743fe"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
            "\u001b[32;1m\u001b[1;3m{\n",
            "    \"action\": \"Final Answer\",\n",
            "    \"action_input\": \"The product of 2 multiplied by 7 is 14.\"\n",
            "}\u001b[0m\n",
            "\n",
            "\u001b[1m> Finished chain.\u001b[0m\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "{'input': 'what is 2 * 7?',\n",
              " 'chat_history': [HumanMessage(content='when was the college of engineering in the University of Notre Dame established?'),\n",
              "  AIMessage(content='The College of Engineering at the University of Notre Dame was established in 1920.')],\n",
              " 'output': 'The product of 2 multiplied by 7 is 14.'}"
            ]
          },
          "execution_count": 18,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "agent(\"what is 2 * 7?\")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "gR_b0IN32rQ9"
      },
      "source": [
        "Perfect, the agent is able to recognize that it doesn't need to refer to it's general knowledge tool for that question. Let's try some more questions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mQeicHTj2pmY",
        "outputId": "9810647d-0846-4ca8-ae2e-1b249ca4ad3d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
            "\u001b[32;1m\u001b[1;3m{\n",
            "    \"action\": \"Knowledge Base\",\n",
            "    \"action_input\": \"University of Notre Dame\"\n",
            "}\u001b[0m\n",
            "Observation: \u001b[36;1m\u001b[1;3mThe University of Notre Dame is a Catholic research university located in South Bend, Indiana, in the United States. It is known for its strong academic programs, including undergraduate colleges in Arts and Letters, Science, Engineering, Business, and the Architecture School. The university also has a graduate program with over 50 master's, doctoral, and professional degree programs. Notre Dame is recognized as one of the top universities in the United States and has a strong alumni network. It is also known for its iconic landmarks, such as the Golden Dome and the Basilica. The university is committed to research and has various institutes dedicated to different fields of study. Notre Dame is also home to the Notre Dame Global Adaptation Index, which ranks countries based on their vulnerability to climate change.\u001b[0m\n",
            "Thought:\u001b[32;1m\u001b[1;3m{\n",
            "    \"action\": \"Final Answer\",\n",
            "    \"action_input\": \"The University of Notre Dame is a Catholic research university located in South Bend, Indiana. It offers strong academic programs in various fields, including Arts and Letters, Science, Engineering, Business, and Architecture. Notre Dame is known for its academic excellence, iconic landmarks like the Golden Dome and the Basilica, and its commitment to research. It is also home to the Notre Dame Global Adaptation Index, which ranks countries based on their vulnerability to climate change.\"\n",
            "}\u001b[0m\n",
            "\n",
            "\u001b[1m> Finished chain.\u001b[0m\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "{'input': 'can you tell me some facts about the University of Notre Dame?',\n",
              " 'chat_history': [HumanMessage(content='when was the college of engineering in the University of Notre Dame established?'),\n",
              "  AIMessage(content='The College of Engineering at the University of Notre Dame was established in 1920.'),\n",
              "  HumanMessage(content='what is 2 * 7?'),\n",
              "  AIMessage(content='The product of 2 multiplied by 7 is 14.')],\n",
              " 'output': 'The University of Notre Dame is a Catholic research university located in South Bend, Indiana. It offers strong academic programs in various fields, including Arts and Letters, Science, Engineering, Business, and Architecture. Notre Dame is known for its academic excellence, iconic landmarks like the Golden Dome and the Basilica, and its commitment to research. It is also home to the Notre Dame Global Adaptation Index, which ranks countries based on their vulnerability to climate change.'}"
            ]
          },
          "execution_count": 19,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "agent(\"can you tell me some facts about the University of Notre Dame?\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "G93vLXso3B5Z",
        "outputId": "cb345147-5630-4a5a-bb9f-242ad91d9968"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
            "\u001b[32;1m\u001b[1;3m{\n",
            "    \"action\": \"Final Answer\",\n",
            "    \"action_input\": \"The University of Notre Dame is a Catholic research university located in South Bend, Indiana. It offers strong academic programs and is known for its iconic landmarks and commitment to research.\"\n",
            "}\u001b[0m\n",
            "\n",
            "\u001b[1m> Finished chain.\u001b[0m\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "{'input': 'can you summarize these facts in two short sentences',\n",
              " 'chat_history': [HumanMessage(content='when was the college of engineering in the University of Notre Dame established?'),\n",
              "  AIMessage(content='The College of Engineering at the University of Notre Dame was established in 1920.'),\n",
              "  HumanMessage(content='what is 2 * 7?'),\n",
              "  AIMessage(content='The product of 2 multiplied by 7 is 14.'),\n",
              "  HumanMessage(content='can you tell me some facts about the University of Notre Dame?'),\n",
              "  AIMessage(content='The University of Notre Dame is a Catholic research university located in South Bend, Indiana. It offers strong academic programs in various fields, including Arts and Letters, Science, Engineering, Business, and Architecture. Notre Dame is known for its academic excellence, iconic landmarks like the Golden Dome and the Basilica, and its commitment to research. It is also home to the Notre Dame Global Adaptation Index, which ranks countries based on their vulnerability to climate change.')],\n",
              " 'output': 'The University of Notre Dame is a Catholic research university located in South Bend, Indiana. It offers strong academic programs and is known for its iconic landmarks and commitment to research.'}"
            ]
          },
          "execution_count": 20,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "agent(\"can you summarize these facts in two short sentences\")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "PWivmw9F3bCw"
      },
      "source": [
        "Looks great! We're also able to ask questions that refer to previous interactions in the conversation and the agent is able to refer to the conversation history to as a source of information.\n",
        "\n",
        "That's all for this example of building a retrieval augmented conversational agent with OpenAI and Pinecone (the OP stack) and LangChain.\n",
        "\n",
        "Once finished, we delete the Pinecone index to save resources:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "Pa1whr8V3Wfm"
      },
      "outputs": [],
      "source": [
        "pc.delete_index(index_name)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "Ykg5TYA033yR"
      },
      "source": [
        "---"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.9.12"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "0679733cce4f4b5eb25f8f56dc7ebe22": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "07b47869f727455fb5ba67e99767c04c": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "1b5fe655986842a6afd4a2eb2c0f47fc": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_a03d372e8b4f456e8663e54c818e0fe9",
            "max": 189,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_6fc773ccbc594fddbc1f2f0adaa96186",
            "value": 189
          }
        },
        "321ebcd192fc41e0a7a2ba9c0e7b8556": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_7c898f43b39545879013df760aa4f6f1",
              "IPY_MODEL_1b5fe655986842a6afd4a2eb2c0f47fc",
              "IPY_MODEL_76872ca737b64616aeacc90f2f7fc655"
            ],
            "layout": "IPY_MODEL_0679733cce4f4b5eb25f8f56dc7ebe22"
          }
        },
        "4b2e12a33584489690aaf69f51cd4a3b": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "6fc773ccbc594fddbc1f2f0adaa96186": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "71551fc297b54fb4a21d787d117e01a1": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "76872ca737b64616aeacc90f2f7fc655": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_07b47869f727455fb5ba67e99767c04c",
            "placeholder": "​",
            "style": "IPY_MODEL_4b2e12a33584489690aaf69f51cd4a3b",
            "value": " 189/189 [03:23&lt;00:00,  1.12s/it]"
          }
        },
        "7c898f43b39545879013df760aa4f6f1": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_71551fc297b54fb4a21d787d117e01a1",
            "placeholder": "​",
            "style": "IPY_MODEL_b90283522431447b925b226431bb4fa6",
            "value": "100%"
          }
        },
        "a03d372e8b4f456e8663e54c818e0fe9": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "b90283522431447b925b226431bb4fa6": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}