L2_Retrieval_Augmented_Generation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "43ead32b-d4ab-431e-b64d-7a59b8962fae",
   "metadata": {},
   "source": [
    "# Lesson 2 - Retrieval Augmented Generation (RAG)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a229d9f9",
   "metadata": {},
   "source": [
    "### Import  the Needed Packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c74d8975-53e1-41a3-a846-1fed3361a4b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1382dbd-5b93-419e-ae16-735aa4e0f22a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from openai import OpenAI\n",
    "from pinecone import Pinecone, ServerlessSpec\n",
    "from tqdm.auto import tqdm\n",
    "from DLAIUtils import Utils\n",
    "\n",
    "import ast\n",
    "import os\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "955e9783-9f33-4c09-8b10-ebee6b4a236e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get api key\n",
    "utils = Utils()\n",
    "PINECONE_API_KEY = utils.get_pinecone_api_key()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c79cdc21",
   "metadata": {},
   "source": [
    "### Setup Pinecone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73e4c29d-79ee-4d6c-8a78-c7d29077fe5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "pinecone = Pinecone(api_key=PINECONE_API_KEY)\n",
    "\n",
    "utils = Utils()\n",
    "INDEX_NAME = utils.create_dlai_index_name('dl-ai')\n",
    "if INDEX_NAME in [index.name for index in pinecone.list_indexes()]:\n",
    "  pinecone.delete_index(INDEX_NAME)\n",
    "\n",
    "pinecone.create_index(name=INDEX_NAME, dimension=1536, metric='cosine',\n",
    "  spec=ServerlessSpec(cloud='aws', region='us-west-2'))\n",
    "\n",
    "index = pinecone.Index(INDEX_NAME)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94f00979",
   "metadata": {},
   "source": [
    "### Load the Dataset\n",
    "\n",
    "**Note:** To access the dataset outside of this course, just copy the following two lines of code and run it (remember to uncomment them first before executing):\n",
    "\n",
    "#!wget -q -O lesson2-wiki.csv.zip \"https://www.dropbox.com/scl/fi/yxzmsrv2sgl249zcspeqb/lesson2-wiki.csv.zip?rlkey=paehnoxjl3s5x53d1bedt4pmc&dl=0\"\n",
    "\n",
    "#!unzip lesson2-wiki.csv.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b059ccc3-f6ba-49dd-a42c-bd9a5bcfc4d3",
   "metadata": {},
   "source": [
    "<p style=\"background-color:#fff1d7; padding:15px; \"> <b>(Note: <code>max_articles_num = 500</code>):</b> To achieve a more comprehensive context for the Language Learning Model, a larger number of articles is generally more beneficial. In this lab, we've initially set <code>max_articles_num</code> to 500 for speedier results, allowing you to observe the outcomes faster. Once you've done an initial run, consider increasing this value to 750 or 1,000. You'll likely notice that the context provided to the LLM becomes richer and better. You can experiment by gradually raising this variable for different queries to observe the improvements in the LLM's contextual understanding.</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a372beb7-2322-4f6c-bb80-a9aa2c74fdaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_articles_num = 500\n",
    "df = pd.read_csv('./data/wiki.csv', nrows=max_articles_num)\n",
    "df.head()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa307525",
   "metadata": {},
   "source": [
    "### Prepare the Embeddings and Upsert to Pinecone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27ac4cb4-e900-4d9e-bfc1-cd0ada96ed80",
   "metadata": {},
   "outputs": [],
   "source": [
    "prepped = []\n",
    "\n",
    "for i, row in tqdm(df.iterrows(), total=df.shape[0]):\n",
    "    meta = ast.literal_eval(row['metadata'])\n",
    "    prepped.append({'id':row['id'], \n",
    "                    'values':ast.literal_eval(row['values']), \n",
    "                    'metadata':meta})\n",
    "    if len(prepped) >= 250:\n",
    "        index.upsert(prepped)\n",
    "        prepped = []\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7180b6b2-ad7c-4a59-ab6c-046d21957343",
   "metadata": {},
   "outputs": [],
   "source": [
    "index.describe_index_stats()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "617fe45f",
   "metadata": {},
   "source": [
    "### Connect to OpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20f1922a-122e-4f3d-a9b4-6ab1e5a67195",
   "metadata": {},
   "outputs": [],
   "source": [
    "OPENAI_API_KEY = utils.get_openai_api_key()\n",
    "openai_client = OpenAI(api_key=OPENAI_API_KEY)\n",
    "\n",
    "def get_embeddings(articles, model=\"text-embedding-ada-002\"):\n",
    "   return openai_client.embeddings.create(input = articles, model=model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f629efde",
   "metadata": {},
   "source": [
    "### Run Your Query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f790629-f7e4-40e8-9f96-cfdd46cfde9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"what is the berlin wall?\"\n",
    "\n",
    "embed = get_embeddings([query])\n",
    "res = index.query(vector=embed.data[0].embedding, top_k=3, include_metadata=True)\n",
    "text = [r['metadata']['text'] for r in res['matches']]\n",
    "print('\\n'.join(text))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26ee483a",
   "metadata": {},
   "source": [
    "### Build the Prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92cdddce-9a39-4ffd-823c-e91b7af48cf5",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"write an article titled: what is the berlin wall?\"\n",
    "embed = get_embeddings([query])\n",
    "res = index.query(vector=embed.data[0].embedding, top_k=3, include_metadata=True)\n",
    "\n",
    "contexts = [\n",
    "    x['metadata']['text'] for x in res['matches']\n",
    "]\n",
    "\n",
    "prompt_start = (\n",
    "    \"Answer the question based on the context below.\\n\\n\"+\n",
    "    \"Context:\\n\"\n",
    ")\n",
    "\n",
    "prompt_end = (\n",
    "    f\"\\n\\nQuestion: {query}\\nAnswer:\"\n",
    ")\n",
    "\n",
    "prompt = (\n",
    "    prompt_start + \"\\n\\n---\\n\\n\".join(contexts) + \n",
    "    prompt_end\n",
    ")\n",
    "\n",
    "print(prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c457d658",
   "metadata": {},
   "source": [
    "### Get the Summary "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cd1b8ed-99b6-4c47-a784-da4192aa6e6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = openai_client.completions.create(\n",
    "    model=\"gpt-3.5-turbo-instruct\",\n",
    "    prompt=prompt,\n",
    "    temperature=0,\n",
    "    max_tokens=636,\n",
    "    top_p=1,\n",
    "    frequency_penalty=0,\n",
    "    presence_penalty=0,\n",
    "    stop=None\n",
    ")\n",
    "print('-' * 80)\n",
    "print(res.choices[0].text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e985a43-5d35-4173-9ae1-3efef399c379",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}