{
  "cells": [
    {
      "attachments": {
        "image-2.png": {
          "image/png": "[encoded data, 173972 characters]"
        },
        "image.png": {
          "image/png": "[encoded data, 131288 characters]"
        }
      },
      "cell_type": "markdown",
      "metadata": {
        "id": "ELUZeeM6BshL"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/mongodb-developer/GenAI-Showcase/blob/main/notebooks/rag/rag_mongodb_llama3_huggingface_open_source.ipynb) \n",
        "\n",
        "# Implementing RAG pipelines with MongoDB and Llama3 and Open models From Hugging Face\n",
        "\n",
        "This notebook is designed to demonstrate how to integrate and utilize Hugging Face's open-source models, specifically Llama3, with MongoDB to implement Retrieval-Augmented Generation (RAG) pipelines for enhanced question answering capabilities.\n",
        "\n",
        "The process involves preparing a dataset of arXiv papers, transforming their data for effective retrieval, setting up a MongoDB database with vector search capabilities, and using llama3 model for generating answers based on the retrieved documents.\n",
        "\n",
        "Key Highlights:\n",
        "- Usage of Hugging Face open-source models and MongoDB for creating RAG pipelines.\n",
        "- Steps include dataset preparation, database setup, data ingestion, and query processing.\n",
        "- Detailed guidance on setting up MongoDB collections and vector search indexes.\n",
        "- Integration with the Llama3 model from Hugging Face for answering complex queries.\n",
        "\n",
        "Follow the following instruction to set up a MongoDB database and enable vector search:\n",
        "1. [Register a free Atlas account](https://account.mongodb.com/account/register?utm_campaign=devrel&utm_source=community&utm_medium=cta&utm_content=GitHub%20Cookbook&utm_term=richmond.alake)\n",
        " or sign in to your existing Atlas account.\n",
        "\n",
        "2. [Follow the instructions](https://www.mongodb.com/docs/atlas/tutorial/deploy-free-tier-cluster/)\n",
        " (select Atlas UI as the procedure) to deploy your first cluster, which distributes your data across multiple servers for improved performance and redundancy.\n",
        " \n",
        " ![image.png](attachment:image.png)\n",
        "\n",
        "3. For a free Cluster, be sure to select \"Shared\" option when creating your new cluster. See image below for details\n",
        "\n",
        "![image-2.png](attachment:image-2.png)\n",
        "\n",
        "4. Create the database: `knowledge_base`, and collection `research_papers`\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BPgc75ATDIvA"
      },
      "source": [
        "## Import Libraries\n",
        "\n",
        "Import libraries into development environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UUlChXcYO43e"
      },
      "outputs": [],
      "source": [
        "!pip install datasets pandas pymongo sentence_transformers\n",
        "!pip install -U transformers\n",
        "# Install the library below if using GPU, if using CPU, please comment out below\n",
        "!pip install accelerate"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JZHLRiGvDSgB"
      },
      "source": [
        "## Dataset Loading and Preparation\n",
        "\n",
        "Load the dataset from HuggingFace.\n",
        "\n",
        "Only using the first 100 datapoint for demo purposes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uRvntN3jO8j1"
      },
      "outputs": [],
      "source": [
        "# Load Dataset\n",
        "from datasets import load_dataset\n",
        "import pandas as pd\n",
        "import os\n",
        "\n",
        "# Make sure you have an Hugging Face token(HF_TOKEN) in your development environment before runing the code below\n",
        "# How to get a token: https://huggingface.co/docs/hub/en/security-tokens\n",
        "# Dataset Location: https://huggingface.co/datasets/MongoDB/subset_arxiv_papers_with_embeddings\n",
        "os.environ[\"HF_TOKEN\"] = \"place_hugging_face_access_token here\" # Do not use this in production environment, use a .env file instead\n",
        "\n",
        "dataset = load_dataset(\"MongoDB/subset_arxiv_papers_with_embeddings\")\n",
        "\n",
        "# Convert the dataset to a pandas dataframe\n",
        "dataset_df = pd.DataFrame(dataset['train'])\n",
        "\n",
        "dataset_df.head(5)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LUsc7qu6O_Nz"
      },
      "outputs": [],
      "source": [
        "# Data Preparation\n",
        "\n",
        "# Only use the first 100 for demo and POC purposes\n",
        "dataset_df = dataset_df.head(100)\n",
        "\n",
        "# Remove the embedding from each data point in the dataset as we are going to create new embeddings with an open source embedding model from Hugging Face\n",
        "dataset_df = dataset_df.drop(columns=['embedding'])\n",
        "dataset_df.head(5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sX3xd3ijDvNJ"
      },
      "source": [
        "## Generate Embeddings"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 677
        },
        "id": "PgGL45kgPB9r",
        "outputId": "67c74bdc-b731-44bb-adf7-077cf84f6dcc"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "summary": "{\n  \"name\": \"dataset_df\",\n  \"rows\": 100,\n  \"fields\": [\n    {\n      \"column\": \"id\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 0.0029011491975882153,\n        \"min\": 704.0001,\n        \"max\": 704.01,\n        \"num_unique_values\": 100,\n        \"samples\": [\n          704.0084,\n          704.0054,\n          704.0071\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"submitter\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 93,\n        \"samples\": [\n          \"A. M. Kamchatnov\",\n          \"Mikhail Kostylev\",\n          \"Tom\\\\'a\\\\v{s} Vina\\\\v{r}\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"authors\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 96,\n        \"samples\": [\n          \"Partha Mukhopadhyay\",\n          \"Tetsuo Kurosaki, Miki Wadati\",\n          \"Marc Mars, Filipe C. Mena, Raul Vera\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"title\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 100,\n        \"samples\": [\n          \"Formation of density singularities in ideal hydrodynamics of freely\\n  cooling inelastic gases: a family of exact solutions\",\n          \"The Hardy-Lorentz Spaces $H^{p,q}(R^n)$\",\n          \"Pairwise comparisons of typological profiles (of languages)\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"comments\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 83,\n        \"samples\": [\n          \"Resubmit with new results on the upper bound of the number of steady\\n  states. 20 pages, 2 figures, See\\n  http://www.math.rutgers.edu/~sontag/PUBDIR/index.html for online preprints\\n  and reprints of related work\",\n          \"37 pages, 15 figures; published version\",\n          \"4 pages, 2 figures; mistakes due to an erroneous electron-phonon\\n  coupling constant have been corrected; mode splitting is larger than\\n  originally expected\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"journal-ref\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 52,\n        \"samples\": [\n          \"J.Opt.Soc.Am.A 23(10): 2592-2601 (2006)\",\n          \"Int.J.Mod.Phys.A22:1953-1982,2007\",\n          \"Comm. Math. Helv. 84 (2009), No. 2, pp. 259--296.\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"doi\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 50,\n        \"samples\": [\n          \"10.1103/PhysRevD.75.094018\",\n          \"10.1088/0264-9381/24/14/008\",\n          \"10.1103/PhysRevLett.99.126405\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"report-no\",\n      \"properties\": {\n        \"dtype\": \"category\",\n        \"num_unique_values\": 8,\n        \"samples\": [\n          \"IGPG-07/03-2\",\n          \"DAMTP-2007-30\",\n          \"ANL-HEP-PR-07-12\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"categories\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 66,\n        \"samples\": [\n          \"math.PR\",\n          \"math.GR math.DG\",\n          \"hep-ph\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"license\",\n      \"properties\": {\n        \"dtype\": \"category\",\n        \"num_unique_values\": 2,\n        \"samples\": [\n          \"http://creativecommons.org/licenses/by-nc-nd/4.0/\",\n          \"http://arxiv.org/licenses/nonexclusive-distrib/1.0/\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"abstract\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 100,\n        \"samples\": [\n          \"  We employ granular hydrodynamics to investigate a paradigmatic problem of\\nclustering of particles in a freely cooling dilute granular gas. We consider\\nlarge-scale hydrodynamic motions where the viscosity and heat conduction can be\\nneglected, and one arrives at the equations of ideal gas dynamics with an\\nadditional term describing bulk energy losses due to inelastic collisions. We\\nemploy Lagrangian coordinates and derive a broad family of exact non-stationary\\nanalytical solutions that depend only on one spatial coordinate. These\\nsolutions exhibit a new type of singularity, where the gas density blows up in\\na finite time when starting from smooth initial conditions. The density blowups\\nsignal formation of close-packed clusters of particles. As the density blow-up\\ntime $t_c$ is approached, the maximum density exhibits a power law $\\\\sim\\n(t_c-t)^{-2}$. The velocity gradient blows up as $\\\\sim - (t_c-t)^{-1}$ while\\nthe velocity itself remains continuous and develops a cusp (rather than a shock\\ndiscontinuity) at the singularity. The gas temperature vanishes at the\\nsingularity, and the singularity follows the isobaric scenario: the gas\\npressure remains finite and approximately uniform in space and constant in time\\nclose to the singularity. An additional exact solution shows that the density\\nblowup, of the same type, may coexist with an \\\"ordinary\\\" shock, at which the\\nhydrodynamic fields are discontinuous but finite. We confirm stability of the\\nexact solutions with respect to small one-dimensional perturbations by solving\\nthe ideal hydrodynamic equations numerically. Furthermore, numerical solutions\\nshow that the local features of the density blowup hold universally,\\nindependently of details of the initial and boundary conditions.\\n\",\n          \"  In this paper we consider the Hardy-Lorentz spaces $H^{p,q}(R^n)$, with\\n$0<p\\\\le 1$, $0<q\\\\le \\\\infty$. We discuss the atomic decomposition of the\\nelements in these spaces, their interpolation properties, and the behavior of\\nsingular integrals and other operators acting on them.\\n\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"versions\",\n      \"properties\": {\n        \"dtype\": \"object\",\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"update_date\",\n      \"properties\": {\n        \"dtype\": \"date\",\n        \"min\": \"2007-05-23 00:00:00\",\n        \"max\": \"2022-03-31 00:00:00\",\n        \"num_unique_values\": 47,\n        \"samples\": [\n          \"2008-09-05 00:00:00\",\n          \"2008-01-19 00:00:00\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"authors_parsed\",\n      \"properties\": {\n        \"dtype\": \"object\",\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"embedding\",\n      \"properties\": {\n        \"dtype\": \"object\",\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    }\n  ]\n}",
              "type": "dataframe",
              "variable_name": "dataset_df"
            },
            "text/html": [
              "\n",
              "  <div id=\"df-2d3a27db-c5ee-46da-afce-99284beacd5d\" class=\"colab-df-container\">\n",
              "    <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>submitter</th>\n",
              "      <th>authors</th>\n",
              "      <th>title</th>\n",
              "      <th>comments</th>\n",
              "      <th>journal-ref</th>\n",
              "      <th>doi</th>\n",
              "      <th>report-no</th>\n",
              "      <th>categories</th>\n",
              "      <th>license</th>\n",
              "      <th>abstract</th>\n",
              "      <th>versions</th>\n",
              "      <th>update_date</th>\n",
              "      <th>authors_parsed</th>\n",
              "      <th>embedding</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>704.0001</td>\n",
              "      <td>Pavel Nadolsky</td>\n",
              "      <td>C. Bal\\'azs, E. L. Berger, P. M. Nadolsky, C.-...</td>\n",
              "      <td>Calculation of prompt diphoton production cross...</td>\n",
              "      <td>37 pages, 15 figures; published version</td>\n",
              "      <td>Phys.Rev.D76:013009,2007</td>\n",
              "      <td>10.1103/PhysRevD.76.013009</td>\n",
              "      <td>ANL-HEP-PR-07-12</td>\n",
              "      <td>hep-ph</td>\n",
              "      <td>None</td>\n",
              "      <td>A fully differential calculation in perturba...</td>\n",
              "      <td>[{'version': 'v1', 'created': 'Mon, 2 Apr 2007...</td>\n",
              "      <td>2008-11-26</td>\n",
              "      <td>[[Balázs, C., ], [Berger, E. L., ], [Nadolsky,...</td>\n",
              "      <td>[-0.0073745595291256905, -0.03725249320268631,...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>704.0002</td>\n",
              "      <td>Louis Theran</td>\n",
              "      <td>Ileana Streinu and Louis Theran</td>\n",
              "      <td>Sparsity-certifying Graph Decompositions</td>\n",
              "      <td>To appear in Graphs and Combinatorics</td>\n",
              "      <td>None</td>\n",
              "      <td>None</td>\n",
              "      <td>None</td>\n",
              "      <td>math.CO cs.CG</td>\n",
              "      <td>http://arxiv.org/licenses/nonexclusive-distrib...</td>\n",
              "      <td>We describe a new algorithm, the $(k,\\ell)$-...</td>\n",
              "      <td>[{'version': 'v1', 'created': 'Sat, 31 Mar 200...</td>\n",
              "      <td>2008-12-13</td>\n",
              "      <td>[[Streinu, Ileana, ], [Theran, Louis, ]]</td>\n",
              "      <td>[0.005753430537879467, 0.007056022062897682, 0...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>704.0003</td>\n",
              "      <td>Hongjun Pan</td>\n",
              "      <td>Hongjun Pan</td>\n",
              "      <td>The evolution of the Earth-Moon system based o...</td>\n",
              "      <td>23 pages, 3 figures</td>\n",
              "      <td>None</td>\n",
              "      <td>None</td>\n",
              "      <td>None</td>\n",
              "      <td>physics.gen-ph</td>\n",
              "      <td>None</td>\n",
              "      <td>The evolution of Earth-Moon system is descri...</td>\n",
              "      <td>[{'version': 'v1', 'created': 'Sun, 1 Apr 2007...</td>\n",
              "      <td>2008-01-13</td>\n",
              "      <td>[[Pan, Hongjun, ]]</td>\n",
              "      <td>[-0.0057186526246368885, 0.022108040750026703,...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>704.0004</td>\n",
              "      <td>David Callan</td>\n",
              "      <td>David Callan</td>\n",
              "      <td>A determinant of Stirling cycle numbers counts...</td>\n",
              "      <td>11 pages</td>\n",
              "      <td>None</td>\n",
              "      <td>None</td>\n",
              "      <td>None</td>\n",
              "      <td>math.CO</td>\n",
              "      <td>None</td>\n",
              "      <td>We show that a determinant of Stirling cycle...</td>\n",
              "      <td>[{'version': 'v1', 'created': 'Sat, 31 Mar 200...</td>\n",
              "      <td>2007-05-23</td>\n",
              "      <td>[[Callan, David, ]]</td>\n",
              "      <td>[-0.02010205015540123, -0.0021757606882601976,...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>704.0005</td>\n",
              "      <td>Alberto Torchinsky</td>\n",
              "      <td>Wael Abu-Shammala and Alberto Torchinsky</td>\n",
              "      <td>From dyadic $\\Lambda_{\\alpha}$ to $\\Lambda_{\\a...</td>\n",
              "      <td>None</td>\n",
              "      <td>Illinois J. Math. 52 (2008) no.2, 681-689</td>\n",
              "      <td>None</td>\n",
              "      <td>None</td>\n",
              "      <td>math.CA math.FA</td>\n",
              "      <td>None</td>\n",
              "      <td>In this paper we show how to compute the $\\L...</td>\n",
              "      <td>[{'version': 'v1', 'created': 'Mon, 2 Apr 2007...</td>\n",
              "      <td>2013-10-15</td>\n",
              "      <td>[[Abu-Shammala, Wael, ], [Torchinsky, Alberto, ]]</td>\n",
              "      <td>[-0.0027832775376737118, 0.014300416223704815,...</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "    <div class=\"colab-df-buttons\">\n",
              "\n",
              "  <div class=\"colab-df-container\">\n",
              "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-2d3a27db-c5ee-46da-afce-99284beacd5d')\"\n",
              "            title=\"Convert this dataframe to an interactive table.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
              "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "\n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    .colab-df-buttons div {\n",
              "      margin-bottom: 4px;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "    <script>\n",
              "      const buttonEl =\n",
              "        document.querySelector('#df-2d3a27db-c5ee-46da-afce-99284beacd5d button.colab-df-convert');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      async function convertToInteractive(key) {\n",
              "        const element = document.querySelector('#df-2d3a27db-c5ee-46da-afce-99284beacd5d');\n",
              "        const dataTable =\n",
              "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                    [key], {});\n",
              "        if (!dataTable) return;\n",
              "\n",
              "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "          + ' to learn more about interactive tables.';\n",
              "        element.innerHTML = '';\n",
              "        dataTable['output_type'] = 'display_data';\n",
              "        await google.colab.output.renderOutput(dataTable, element);\n",
              "        const docLink = document.createElement('div');\n",
              "        docLink.innerHTML = docLinkHtml;\n",
              "        element.appendChild(docLink);\n",
              "      }\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "\n",
              "<div id=\"df-3a97a838-2b0e-445a-af3a-e782e6330ba2\">\n",
              "  <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-3a97a838-2b0e-445a-af3a-e782e6330ba2')\"\n",
              "            title=\"Suggest charts\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "     width=\"24px\">\n",
              "    <g>\n",
              "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
              "    </g>\n",
              "</svg>\n",
              "  </button>\n",
              "\n",
              "<style>\n",
              "  .colab-df-quickchart {\n",
              "      --bg-color: #E8F0FE;\n",
              "      --fill-color: #1967D2;\n",
              "      --hover-bg-color: #E2EBFA;\n",
              "      --hover-fill-color: #174EA6;\n",
              "      --disabled-fill-color: #AAA;\n",
              "      --disabled-bg-color: #DDD;\n",
              "  }\n",
              "\n",
              "  [theme=dark] .colab-df-quickchart {\n",
              "      --bg-color: #3B4455;\n",
              "      --fill-color: #D2E3FC;\n",
              "      --hover-bg-color: #434B5C;\n",
              "      --hover-fill-color: #FFFFFF;\n",
              "      --disabled-bg-color: #3B4455;\n",
              "      --disabled-fill-color: #666;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart {\n",
              "    background-color: var(--bg-color);\n",
              "    border: none;\n",
              "    border-radius: 50%;\n",
              "    cursor: pointer;\n",
              "    display: none;\n",
              "    fill: var(--fill-color);\n",
              "    height: 32px;\n",
              "    padding: 0;\n",
              "    width: 32px;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart:hover {\n",
              "    background-color: var(--hover-bg-color);\n",
              "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "    fill: var(--button-hover-fill-color);\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart-complete:disabled,\n",
              "  .colab-df-quickchart-complete:disabled:hover {\n",
              "    background-color: var(--disabled-bg-color);\n",
              "    fill: var(--disabled-fill-color);\n",
              "    box-shadow: none;\n",
              "  }\n",
              "\n",
              "  .colab-df-spinner {\n",
              "    border: 2px solid var(--fill-color);\n",
              "    border-color: transparent;\n",
              "    border-bottom-color: var(--fill-color);\n",
              "    animation:\n",
              "      spin 1s steps(1) infinite;\n",
              "  }\n",
              "\n",
              "  @keyframes spin {\n",
              "    0% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "      border-left-color: var(--fill-color);\n",
              "    }\n",
              "    20% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    30% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    40% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    60% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    80% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "    90% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "  }\n",
              "</style>\n",
              "\n",
              "  <script>\n",
              "    async function quickchart(key) {\n",
              "      const quickchartButtonEl =\n",
              "        document.querySelector('#' + key + ' button');\n",
              "      quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
              "      quickchartButtonEl.classList.add('colab-df-spinner');\n",
              "      try {\n",
              "        const charts = await google.colab.kernel.invokeFunction(\n",
              "            'suggestCharts', [key], {});\n",
              "      } catch (error) {\n",
              "        console.error('Error during call to suggestCharts:', error);\n",
              "      }\n",
              "      quickchartButtonEl.classList.remove('colab-df-spinner');\n",
              "      quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
              "    }\n",
              "    (() => {\n",
              "      let quickchartButtonEl =\n",
              "        document.querySelector('#df-3a97a838-2b0e-445a-af3a-e782e6330ba2 button');\n",
              "      quickchartButtonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "    })();\n",
              "  </script>\n",
              "</div>\n",
              "    </div>\n",
              "  </div>\n"
            ],
            "text/plain": [
              "         id           submitter  \\\n",
              "0  704.0001      Pavel Nadolsky   \n",
              "1  704.0002        Louis Theran   \n",
              "2  704.0003         Hongjun Pan   \n",
              "3  704.0004        David Callan   \n",
              "4  704.0005  Alberto Torchinsky   \n",
              "\n",
              "                                             authors  \\\n",
              "0  C. Bal\\'azs, E. L. Berger, P. M. Nadolsky, C.-...   \n",
              "1                    Ileana Streinu and Louis Theran   \n",
              "2                                        Hongjun Pan   \n",
              "3                                       David Callan   \n",
              "4           Wael Abu-Shammala and Alberto Torchinsky   \n",
              "\n",
              "                                               title  \\\n",
              "0  Calculation of prompt diphoton production cross...   \n",
              "1           Sparsity-certifying Graph Decompositions   \n",
              "2  The evolution of the Earth-Moon system based o...   \n",
              "3  A determinant of Stirling cycle numbers counts...   \n",
              "4  From dyadic $\\Lambda_{\\alpha}$ to $\\Lambda_{\\a...   \n",
              "\n",
              "                                  comments  \\\n",
              "0  37 pages, 15 figures; published version   \n",
              "1    To appear in Graphs and Combinatorics   \n",
              "2                      23 pages, 3 figures   \n",
              "3                                 11 pages   \n",
              "4                                     None   \n",
              "\n",
              "                                 journal-ref                         doi  \\\n",
              "0                   Phys.Rev.D76:013009,2007  10.1103/PhysRevD.76.013009   \n",
              "1                                       None                        None   \n",
              "2                                       None                        None   \n",
              "3                                       None                        None   \n",
              "4  Illinois J. Math. 52 (2008) no.2, 681-689                        None   \n",
              "\n",
              "          report-no       categories  \\\n",
              "0  ANL-HEP-PR-07-12           hep-ph   \n",
              "1              None    math.CO cs.CG   \n",
              "2              None   physics.gen-ph   \n",
              "3              None          math.CO   \n",
              "4              None  math.CA math.FA   \n",
              "\n",
              "                                             license  \\\n",
              "0                                               None   \n",
              "1  http://arxiv.org/licenses/nonexclusive-distrib...   \n",
              "2                                               None   \n",
              "3                                               None   \n",
              "4                                               None   \n",
              "\n",
              "                                            abstract  \\\n",
              "0    A fully differential calculation in perturba...   \n",
              "1    We describe a new algorithm, the $(k,\\ell)$-...   \n",
              "2    The evolution of Earth-Moon system is descri...   \n",
              "3    We show that a determinant of Stirling cycle...   \n",
              "4    In this paper we show how to compute the $\\L...   \n",
              "\n",
              "                                            versions update_date  \\\n",
              "0  [{'version': 'v1', 'created': 'Mon, 2 Apr 2007...  2008-11-26   \n",
              "1  [{'version': 'v1', 'created': 'Sat, 31 Mar 200...  2008-12-13   \n",
              "2  [{'version': 'v1', 'created': 'Sun, 1 Apr 2007...  2008-01-13   \n",
              "3  [{'version': 'v1', 'created': 'Sat, 31 Mar 200...  2007-05-23   \n",
              "4  [{'version': 'v1', 'created': 'Mon, 2 Apr 2007...  2013-10-15   \n",
              "\n",
              "                                      authors_parsed  \\\n",
              "0  [[Balázs, C., ], [Berger, E. L., ], [Nadolsky,...   \n",
              "1           [[Streinu, Ileana, ], [Theran, Louis, ]]   \n",
              "2                                 [[Pan, Hongjun, ]]   \n",
              "3                                [[Callan, David, ]]   \n",
              "4  [[Abu-Shammala, Wael, ], [Torchinsky, Alberto, ]]   \n",
              "\n",
              "                                           embedding  \n",
              "0  [-0.0073745595291256905, -0.03725249320268631,...  \n",
              "1  [0.005753430537879467, 0.007056022062897682, 0...  \n",
              "2  [-0.0057186526246368885, 0.022108040750026703,...  \n",
              "3  [-0.02010205015540123, -0.0021757606882601976,...  \n",
              "4  [-0.0027832775376737118, 0.014300416223704815,...  "
            ]
          },
          "execution_count": 5,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "from sentence_transformers import SentenceTransformer\n",
        "\n",
        "# https://huggingface.co/thenlper/gte-large\n",
        "embedding_model = SentenceTransformer('thenlper/gte-large')\n",
        "\n",
        "def get_embedding(text: str) -> list[float]:\n",
        "  if not text.strip():\n",
        "      print(\"Attempted to get embedding for empty text.\")\n",
        "      return []\n",
        "\n",
        "  embedding = embedding_model.encode(text)\n",
        "\n",
        "  return embedding.tolist()\n",
        "\n",
        "dataset_df['embedding'] = dataset_df.apply(lambda x: get_embedding(x['title'] + \" \" + x['authors'] + \" \" + x['abstract']), axis=1)\n",
        "\n",
        "dataset_df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XAkrmDaF_9kN"
      },
      "source": [
        "## Database and Collection Setup\n",
        "\n",
        "Complete the steps below if not already carried out previously.\n",
        "Creating a database and collection within MongoDB is made simple with MongoDB Atlas.\n",
        "\n",
        "1. [Register a free Atlas account](https://account.mongodb.com/account/register?utm_campaign=devrel&utm_source=community&utm_medium=cta&utm_content=GitHub%20Cookbook&utm_term=richmond.alake)\n",
        " or sign in to your existing Atlas account.\n",
        "\n",
        "2. [Follow the instructions](https://www.mongodb.com/docs/atlas/tutorial/deploy-free-tier-cluster/)\n",
        " (select Atlas UI as the procedure)  to deploy your first cluster.\n",
        "\n",
        "3. Create the database: `knowledge_base`.\n",
        "4. Within the database` knowledge_base`, create the following collections: `research_papers`\n",
        "5. Create a\n",
        "[vector search index](https://www.mongodb.com/docs/atlas/atlas-vector-search/create-index/#procedure)\n",
        " named `vector_index` for the `research_papers` collection. This index enables the RAG application to retrieve records as additional context to supplement user queries via vector search. Below is the JSON definition of the data collection vector search index.\n",
        "\n",
        "\n",
        " Below is a snipper of what the vector search index definition should looks like:\n",
        " ```\n",
        "    {\n",
        "      \"fields\": [\n",
        "        {\n",
        "          \"numDimensions\": 1024,\n",
        "          \"path\": \"embedding\",\n",
        "          \"similarity\": \"cosine\",\n",
        "          \"type\": \"vector\"\n",
        "        }\n",
        "      ]\n",
        "    }\n",
        "  ```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "F-YbBIraPFgb",
        "outputId": "37a58536-ea57-49cf-b861-8560cb506a33"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Connection to MongoDB successful\n"
          ]
        }
      ],
      "source": [
        "import pymongo\n",
        "\n",
        "def get_mongo_client(mongo_uri):\n",
        "    \"\"\"Establish connection to the MongoDB.\"\"\"\n",
        "    try:\n",
        "        client = pymongo.MongoClient(mongo_uri, appname=\"devrel.content.python\")\n",
        "        print(\"Connection to MongoDB successful\")\n",
        "        return client\n",
        "    except pymongo.errors.ConnectionFailure as e:\n",
        "        print(f\"Connection failed: {e}\")\n",
        "        return None\n",
        "\n",
        "# Ensure connection strings are placed securely in environment variables and not disclosed in production environments.\n",
        "mongo_uri = \"mongodb...pName=Cluster0\"  # Placeholder, replace with your connection string or actual environment variable fetching method.\n",
        "\n",
        "if not mongo_uri:\n",
        "    print(\"MONGO_URI not set in environment variables.\")\n",
        "\n",
        "mongo_client = get_mongo_client(mongo_uri)\n",
        "\n",
        "# Ingest data into MongoDB\n",
        "db = mongo_client['knowledge_base']\n",
        "collection = db['research_papers']\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "I0DP8D6qPGyJ",
        "outputId": "769249a6-f30e-4b62-a677-021c3ac04e70"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "DeleteResult({'n': 100, 'electionId': ObjectId('7fffffff000000000000001f'), 'opTime': {'ts': Timestamp(1713636955, 100), 't': 31}, 'ok': 1.0, '$clusterTime': {'clusterTime': Timestamp(1713636955, 100), 'signature': {'hash': b'J\\x17\\x95:(\\x9f\\xb4\\x96\\xcdv:\"\\xbc\\x0c)\\x98\\xd3\\tq\\x89', 'keyId': 7320226449804230662}}, 'operationTime': Timestamp(1713636955, 100)}, acknowledged=True)"
            ]
          },
          "execution_count": 7,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Delete any existing records in the collection\n",
        "collection.delete_many({})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g3jt3YfPD2GY"
      },
      "source": [
        "## Data Ingestion"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5o6SOG7dPHzD",
        "outputId": "22fbeac4-2532-47de-94ca-70ec300da5c3"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Data ingestion into MongoDB completed\n"
          ]
        }
      ],
      "source": [
        "documents = dataset_df.to_dict('records')\n",
        "collection.insert_many(documents)\n",
        "\n",
        "print(\"Data ingestion into MongoDB completed\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZyBMu6NFD4sD"
      },
      "source": [
        "## Vector Search"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "PRC31ofFPKiV"
      },
      "outputs": [],
      "source": [
        "\n",
        "def vector_search(user_query, collection):\n",
        "    \"\"\"\n",
        "    Perform a vector search in the MongoDB collection based on the user query.\n",
        "\n",
        "    Args:\n",
        "    user_query (str): The user's query string.\n",
        "    collection (MongoCollection): The MongoDB collection to search.\n",
        "\n",
        "    Returns:\n",
        "    list: A list of matching documents.\n",
        "    \"\"\"\n",
        "\n",
        "    # Generate embedding for the user query\n",
        "    query_embedding = get_embedding(user_query)\n",
        "\n",
        "    if query_embedding is None:\n",
        "        return \"Invalid query or embedding generation failed.\"\n",
        "\n",
        "    # Define the vector search pipeline\n",
        "    vector_search_stage = {\n",
        "        \"$vectorSearch\": {\n",
        "            \"index\": \"vector_index\",\n",
        "            \"queryVector\": query_embedding,\n",
        "            \"path\": \"embedding\",\n",
        "            \"numCandidates\": 150,  # Number of candidate matches to consider\n",
        "            \"limit\": 4  # Return top 4 matches\n",
        "        }\n",
        "    }\n",
        "\n",
        "    unset_stage = {\n",
        "        \"$unset\": \"embedding\"  # Exclude the 'embedding' field from the results\n",
        "    }\n",
        "\n",
        "    project_stage = {\n",
        "        \"$project\": {\n",
        "            \"_id\": 0,  # Exclude the _id field\n",
        "            \"fullplot\": 1,  # Include the plot field\n",
        "            \"title\": 1,  # Include the title field\n",
        "            \"genres\": 1, # Include the genres field\n",
        "            \"score\": {\n",
        "                \"$meta\": \"vectorSearchScore\"  # Include the search score\n",
        "            }\n",
        "        }\n",
        "    }\n",
        "\n",
        "    pipeline = [vector_search_stage, unset_stage, project_stage]\n",
        "\n",
        "    # Execute the search\n",
        "    results = collection.aggregate(pipeline)\n",
        "    return list(results)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "9MLEt5JtPLt6"
      },
      "outputs": [],
      "source": [
        "def get_search_result(query, collection):\n",
        "\n",
        "  get_knowledge = vector_search(query, collection)\n",
        "\n",
        "  search_result = ''\n",
        "  for result in get_knowledge:\n",
        "      search_result += f\"Title: {result.get('title', 'N/A')}, Plot: {result.get('abstract', 'N/A')}\\n\"\n",
        "\n",
        "  return search_result"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P3U3nZDOD-Dl"
      },
      "source": [
        "## Handling User Queries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jkswY16gPNLJ",
        "outputId": "e4f98b1d-d5b0-4497-b019-fe22e6336c5e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[{'role': 'system', 'content': 'You are a research assistant!'}, {'role': 'user', 'content': 'Query: Get me papers on Artificial Intelligence?\\nContinue to answer the query by using the Search Results:\\n.'}]\n"
          ]
        }
      ],
      "source": [
        "# Conduct query with retrival of sources\n",
        "query = \"Get me papers on Artificial Intelligence?\"\n",
        "source_information = get_search_result(query, collection)\n",
        "combined_information = f\"Query: {query}\\nContinue to answer the query by using the Search Results:\\n{source_information}.\"\n",
        "messages = [\n",
        "    {\"role\": \"system\", \"content\": \"You are a research assistant!\"},\n",
        "    {\"role\": \"user\", \"content\": combined_information},\n",
        "]\n",
        "print(messages)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PDfpkvNNEBXs"
      },
      "source": [
        "## Loading and Using Llama3"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 66,
          "referenced_widgets": [
            "f469e6ca90984864b1c841049a39e154",
            "ece3866a2f3242d696a374710dcdaa53",
            "28c797fe529b432787f13f4853b4ee58",
            "7fec4cd68fc94cd4b3a79e001fb967ae",
            "bbc2c01aee5442999cc82d703029e7d1",
            "078cc39fd28443659b0ec2339b749296",
            "0818fd8f9e1e4abeaf8c169b022c20f8",
            "0be57a9aa4124e4fbe2ac57baf878ad9",
            "957f547f721b4b739fa89c5a4518a1be",
            "c6ac614925f346b286ec1f6e48468d86",
            "0804231321fb4c60973f91880116e6a9"
          ]
        },
        "id": "7hKfOmElPOJF",
        "outputId": "c8186561-d3b4-41c2-98e8-da0af9133e9c"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "f469e6ca90984864b1c841049a39e154",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
        "import torch\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"meta-llama/Meta-Llama-3.1-8B-Instruct\")\n",
        "# CPU Enabled uncomment below 👇🏽\n",
        "# model = AutoModelForCausalLM.from_pretrained(\"meta-llama/Meta-Llama-3.1-8B-Instruct\")\n",
        "# GPU Enabled use below 👇🏽\n",
        "model = AutoModelForCausalLM.from_pretrained(\"meta-llama/Meta-Llama-3.1-8B-Instruct\", torch_dtype=torch.bfloat16, device_map=\"auto\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Foy-fO58-69G",
        "outputId": "c24b8009-7490-4b05-fb51-5d4575d84001"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
            "Setting `pad_token_id` to `eos_token_id`:128001 for open-end generation.\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "I'd be happy to help you with that. Here are some research papers on Artificial Intelligence that I found using a search engine:\n",
            "\n",
            "**1. \"Deep Learning\" by Yann LeCun, Yoshua Bengio, and Geoffrey Hinton (2015)**\n",
            "\n",
            "This paper is a seminal work on deep learning, a subset of AI that involves training neural networks to perform tasks such as image recognition, speech recognition, and natural language processing.\n",
            "\n",
            "Source: LeCun, Y., Bengio, Y., & Hinton, G. (2015). Deep learning. Nature, 521(7553), 436-444. doi: 10.1038/nature14539\n",
            "\n",
            "**2. \"AlphaGo: Mastering the Game of Go with Deep Neural Networks and Tree Search\" by Demis Hassabis, Shane Legg, and Joseph Modayil (2015)**\n",
            "\n",
            "This paper describes the development of AlphaGo, a computer program that uses AI to play the game of Go. AlphaGo was able to defeat a human world champion in a five-game match, marking a significant milestone in AI research.\n",
            "\n",
            "Source: Hassabis, D., Legg, S., & Modayil, J. (2015). AlphaGo: Mastering the\n"
          ]
        }
      ],
      "source": [
        "input_ids = tokenizer.apply_chat_template(\n",
        "    messages,\n",
        "    add_generation_prompt=True,\n",
        "    return_tensors=\"pt\"\n",
        ").to(model.device)\n",
        "\n",
        "terminators = [\n",
        "    tokenizer.eos_token_id,\n",
        "    tokenizer.convert_tokens_to_ids(\"<|eot_id|>\")\n",
        "]\n",
        "\n",
        "outputs = model.generate(\n",
        "    input_ids,\n",
        "    max_new_tokens=256,\n",
        "    eos_token_id=terminators,\n",
        "    do_sample=True,\n",
        "    temperature=0.6,\n",
        "    top_p=0.9,\n",
        ")\n",
        "response = outputs[0][input_ids.shape[-1]:]\n",
        "print(tokenizer.decode(response, skip_special_tokens=True))"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "machine_shape": "hm",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "078cc39fd28443659b0ec2339b749296": {
          "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
          }
        },
        "0804231321fb4c60973f91880116e6a9": {
          "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": ""
          }
        },
        "0818fd8f9e1e4abeaf8c169b022c20f8": {
          "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": ""
          }
        },
        "0be57a9aa4124e4fbe2ac57baf878ad9": {
          "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
          }
        },
        "28c797fe529b432787f13f4853b4ee58": {
          "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_0be57a9aa4124e4fbe2ac57baf878ad9",
            "max": 4,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_957f547f721b4b739fa89c5a4518a1be",
            "value": 4
          }
        },
        "7fec4cd68fc94cd4b3a79e001fb967ae": {
          "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_c6ac614925f346b286ec1f6e48468d86",
            "placeholder": "​",
            "style": "IPY_MODEL_0804231321fb4c60973f91880116e6a9",
            "value": " 4/4 [01:05&lt;00:00, 14.00s/it]"
          }
        },
        "957f547f721b4b739fa89c5a4518a1be": {
          "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": ""
          }
        },
        "bbc2c01aee5442999cc82d703029e7d1": {
          "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
          }
        },
        "c6ac614925f346b286ec1f6e48468d86": {
          "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
          }
        },
        "ece3866a2f3242d696a374710dcdaa53": {
          "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_078cc39fd28443659b0ec2339b749296",
            "placeholder": "​",
            "style": "IPY_MODEL_0818fd8f9e1e4abeaf8c169b022c20f8",
            "value": "Loading checkpoint shards: 100%"
          }
        },
        "f469e6ca90984864b1c841049a39e154": {
          "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_ece3866a2f3242d696a374710dcdaa53",
              "IPY_MODEL_28c797fe529b432787f13f4853b4ee58",
              "IPY_MODEL_7fec4cd68fc94cd4b3a79e001fb967ae"
            ],
            "layout": "IPY_MODEL_bbc2c01aee5442999cc82d703029e7d1"
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
