ChatGPT_Retrieval_Plugin_QA.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# KDBAI for Q&A with ChatGPT Retrieval Plugin"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. Install Python 3.10 if not already installed.\n",
    "\n",
    "2. Clone the `retrieval-app` repository:\n",
    "\n",
    "```\n",
    "git clone git@github.com:openai/retrieval-app.git\n",
    "```\n",
    "\n",
    "3. Navigate to the app directory:\n",
    "\n",
    "```\n",
    "cd /path/to/retrieval-app\n",
    "```\n",
    "\n",
    "4. Install `poetry`:\n",
    "\n",
    "```\n",
    "pip install poetry\n",
    "```\n",
    "\n",
    "5. Create a new virtual environment:\n",
    "\n",
    "```\n",
    "poetry env use python3.10\n",
    "```\n",
    "\n",
    "6. Install the `retrieval-app` dependencies:\n",
    "\n",
    "    Before installing the dependencies with poetry, we will need to add pyKX and a local KDBAI python package as a dependence in the `pyproject.toml` file. This can be done by pointing the KDBAI package under `[tool.poetry.dependencies]` by adding the following lines.\n",
    "\n",
    "    - __pykx = \"^1.6.0\"__\n",
    "    - __kdbai = { path = \"path/to/kdbaiPackage.tar.gz\", develop = false }__\n",
    "\n",
    "```\n",
    "poetry install\n",
    "```\n",
    "\n",
    "7. Set app environment variables:\n",
    "\n",
    "* `BEARER_TOKEN`: Secret token used by the app to authorize incoming requests. The token can be generated using [jwt.io](https://jwt.io/).\n",
    "\n",
    "* `OPENAI_API_KEY`: The OpenAI API key used for generating embeddings with the `text-embedding-ada-002` model. [Get an API key here](https://platform.openai.com/account/api-keys)!\n",
    "\n",
    "8. Set `KDBAI`-specific environment variables:\n",
    "\n",
    "* `DATASTORE`: set to `kdbai`.\n",
    "\n",
    "8. Run the app with:\n",
    "\n",
    "```\n",
    "poetry run start\n",
    "```\n",
    "When the application is up and running a FastAPI server exposes the ChatGPT plugin's endpoints for `upserting`, `querying`, and `deleting` documents.\n",
    "\n",
    "During the start up of the application `KDBAI` initiates a new HNSW index named `openaiEmbeddings` with default parameters {\"efConstruction\": 8, \"efSearch\": 8, \"M\": 32}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/alexg/Applications/miniforge3/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import requests\n",
    "from tqdm.auto import tqdm\n",
    "from datasets import load_dataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If the bearer token is not set as an environment variable, you can set it manually here\n",
    "BEARER_TOKEN = os.environ.get(\"BEARER_TOKEN\")\n",
    "headers = {\"Authorization\": f\"Bearer {BEARER_TOKEN}\"}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Dataset from Hugging Face\n",
    "\n",
    "The Adversarial_QA dataset is chosen for this demo.\n",
    "The adversarial human annotation paradigm ensures that these datasets consist of questions that current state-of-the-art models find challenging."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset adversarial_qa (/Users/alexg/.cache/huggingface/datasets/adversarial_qa/adversarialQA/1.0.0/92356be07b087c5c6a543138757828b8d61ca34de8a87807d40bbc0e6c68f04b)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of unique contexts: 2648\n"
     ]
    },
    {
     "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",
       "      <th>metadata</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>7ba1e8f4261d3170fcf42e84a81dd749116fae95</td>\n",
       "      <td>Brain</td>\n",
       "      <td>Another approach to brain function is to exami...</td>\n",
       "      <td>What sare the benifts of the blood brain barrir?</td>\n",
       "      <td>{'text': ['isolated from the bloodstream'], 'a...</td>\n",
       "      <td>{'split': 'train', 'model_in_the_loop': 'Combi...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>936a8460bfffe437b54cf3ec1e825a3b7b5627a1</td>\n",
       "      <td>Brain</td>\n",
       "      <td>Motor systems are areas of the brain that are ...</td>\n",
       "      <td>What do you think with?</td>\n",
       "      <td>{'text': ['brain'], 'answer_start': [467]}</td>\n",
       "      <td>{'split': 'train', 'model_in_the_loop': 'Combi...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24</th>\n",
       "      <td>e40737d487964dbcd26a223f2799cf56390a98a8</td>\n",
       "      <td>Brain</td>\n",
       "      <td>The brain is an organ that serves as the cente...</td>\n",
       "      <td>How are neurons connected?</td>\n",
       "      <td>{'text': ['synapses'], 'answer_start': [602]}</td>\n",
       "      <td>{'split': 'train', 'model_in_the_loop': 'Combi...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>37</th>\n",
       "      <td>a0f8e785a10f6e21e24207d24ba2823162383062</td>\n",
       "      <td>Brain</td>\n",
       "      <td>The SCN projects to a set of areas in the hypo...</td>\n",
       "      <td>The body's central biological clock is contain...</td>\n",
       "      <td>{'text': ['SCN'], 'answer_start': [4]}</td>\n",
       "      <td>{'split': 'train', 'model_in_the_loop': 'Combi...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>53</th>\n",
       "      <td>6d753d4a8878b5f5bc496d9a369b8c0b212079a0</td>\n",
       "      <td>Brain</td>\n",
       "      <td>The brain contains several motor areas that pr...</td>\n",
       "      <td>What is at the highest level?</td>\n",
       "      <td>{'text': ['the primary motor cortex'], 'answer...</td>\n",
       "      <td>{'split': 'train', 'model_in_the_loop': 'Combi...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                          id  title  \\\n",
       "0   7ba1e8f4261d3170fcf42e84a81dd749116fae95  Brain   \n",
       "12  936a8460bfffe437b54cf3ec1e825a3b7b5627a1  Brain   \n",
       "24  e40737d487964dbcd26a223f2799cf56390a98a8  Brain   \n",
       "37  a0f8e785a10f6e21e24207d24ba2823162383062  Brain   \n",
       "53  6d753d4a8878b5f5bc496d9a369b8c0b212079a0  Brain   \n",
       "\n",
       "                                              context  \\\n",
       "0   Another approach to brain function is to exami...   \n",
       "12  Motor systems are areas of the brain that are ...   \n",
       "24  The brain is an organ that serves as the cente...   \n",
       "37  The SCN projects to a set of areas in the hypo...   \n",
       "53  The brain contains several motor areas that pr...   \n",
       "\n",
       "                                             question  \\\n",
       "0    What sare the benifts of the blood brain barrir?   \n",
       "12                            What do you think with?   \n",
       "24                         How are neurons connected?   \n",
       "37  The body's central biological clock is contain...   \n",
       "53                      What is at the highest level?   \n",
       "\n",
       "                                              answers  \\\n",
       "0   {'text': ['isolated from the bloodstream'], 'a...   \n",
       "12         {'text': ['brain'], 'answer_start': [467]}   \n",
       "24      {'text': ['synapses'], 'answer_start': [602]}   \n",
       "37             {'text': ['SCN'], 'answer_start': [4]}   \n",
       "53  {'text': ['the primary motor cortex'], 'answer...   \n",
       "\n",
       "                                             metadata  \n",
       "0   {'split': 'train', 'model_in_the_loop': 'Combi...  \n",
       "12  {'split': 'train', 'model_in_the_loop': 'Combi...  \n",
       "24  {'split': 'train', 'model_in_the_loop': 'Combi...  \n",
       "37  {'split': 'train', 'model_in_the_loop': 'Combi...  \n",
       "53  {'split': 'train', 'model_in_the_loop': 'Combi...  "
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = load_dataset(\"adversarial_qa\", 'adversarialQA', split=\"train\").to_pandas()\n",
    "data = data.drop_duplicates(subset=[\"context\"])\n",
    "print(f\"Number of unique contexts: {len(data)}\")\n",
    "data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'id': '7ba1e8f4261d3170fcf42e84a81dd749116fae95',\n",
       "  'text': 'Another approach to brain function is to examine the consequences of damage to specific brain areas. Even though it is protected by the skull and meninges, surrounded by cerebrospinal fluid, and isolated from the bloodstream by the blood–brain barrier, the delicate nature of the brain makes it vulnerable to numerous diseases and several types of damage. In humans, the effects of strokes and other types of brain damage have been a key source of information about brain function. Because there is no ability to experimentally control the nature of the damage, however, this information is often difficult to interpret. In animal studies, most commonly involving rats, it is possible to use electrodes or locally injected chemicals to produce precise patterns of damage and then examine the consequences for behavior.',\n",
       "  'metadata': {'title': 'Brain'}},\n",
       " {'id': '936a8460bfffe437b54cf3ec1e825a3b7b5627a1',\n",
       "  'text': 'Motor systems are areas of the brain that are directly or indirectly involved in producing body movements, that is, in activating muscles. Except for the muscles that control the eye, which are driven by nuclei in the midbrain, all the voluntary muscles in the body are directly innervated by motor neurons in the spinal cord and hindbrain. Spinal motor neurons are controlled both by neural circuits intrinsic to the spinal cord, and by inputs that descend from the brain. The intrinsic spinal circuits implement many reflex responses, and contain pattern generators for rhythmic movements such as walking or swimming. The descending connections from the brain allow for more sophisticated control.',\n",
       "  'metadata': {'title': 'Brain'}},\n",
       " {'id': 'e40737d487964dbcd26a223f2799cf56390a98a8',\n",
       "  'text': \"The brain is an organ that serves as the center of the nervous system in all vertebrate and most invertebrate animals. Only a few invertebrates such as sponges, jellyfish, adult sea squirts and starfish do not have a brain; diffuse or localised nerve nets are present instead. The brain is located in the head, usually close to the primary sensory organs for such senses as vision, hearing, balance, taste, and smell. The brain is the most complex organ in a vertebrate's body. In a typical human, the cerebral cortex (the largest part) is estimated to contain 15–33 billion neurons, each connected by synapses to several thousand other neurons. These neurons communicate with one another by means of long protoplasmic fibers called axons, which carry trains of signal pulses called action potentials to distant parts of the brain or body targeting specific recipient cells.\",\n",
       "  'metadata': {'title': 'Brain'}}]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Formatting the documents for the retrieval plugin requests\n",
    "documents = [\n",
    "    {\n",
    "        'id': r['id'],\n",
    "        'text': r['context'],\n",
    "        'metadata': {\n",
    "            'title': r['title']\n",
    "        }\n",
    "    } for r in data.to_dict(orient='records')\n",
    "]\n",
    "documents[:3]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Insert data to the HNSW Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 27/27 [00:38<00:00,  1.44s/it]\n"
     ]
    }
   ],
   "source": [
    "batchSize = 100\n",
    "s = requests.Session()\n",
    "\n",
    "for i in tqdm(range(0, len(documents), batchSize)):\n",
    "    i_end = min(len(documents), i+batchSize)\n",
    "    res = s.post(\n",
    "        \"http://localhost:8000/upsert\",\n",
    "        headers=headers,\n",
    "        json={\n",
    "            \"documents\": documents[i:i_end]\n",
    "        }\n",
    "    )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Query the HNSW index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'query': 'Who went into orbit following Gagarin?'}, {'query': 'Who was responsible for being the best in the Soveit space prgram?'}]\n"
     ]
    }
   ],
   "source": [
    "# Select questions from the AdversarialQA dataset\n",
    "queries = data['question'].tolist()\n",
    "\n",
    "queries = [{'query': queries[i]} for i in range(len(queries))]\n",
    "searchQueries = queries[150:152]\n",
    "print(searchQueries)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Response [200]>"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "results = requests.post(\n",
    "    \"http://0.0.0.0:8000/query\",\n",
    "    headers=headers,\n",
    "    json={\n",
    "        'queries': searchQueries\n",
    "    }\n",
    ")\n",
    "results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------------\n",
      "Who went into orbit following Gagarin?\n",
      "\n",
      "0.35: On May 1, 1960, a U.S. one-man U-2 spy plane was reportedly shot down at high altitude over Soviet Union airspace. The flight was made to gain photo intelligence before the scheduled opening of an East–West summit conference, which had been scheduled in Paris, 15 days later. Captain Francis Gary Powers had bailed out of his aircraft and was captured after parachuting down onto Russian soil. Four days after Powers disappeared, the Eisenhower Administration had NASA issue a very detailed press release noting that an aircraft had \"gone missing\" north of Turkey. It speculated that the pilot might have fallen unconscious while the autopilot was still engaged, and falsely claimed that \"the pilot reported over the emergency frequency that he was experiencing oxygen difficulties.\n",
      "0.46: In 1955 American nuclear arms policy became one aimed primarily at arms control as opposed to disarmament. The failure of negotiations over arms until 1955 was due mainly to the refusal of the Russians to permit any sort of inspections. In talks located in London that year, they expressed a willingness to discuss inspections; the tables were then turned on Eisenhower, when he responded with an unwillingness on the part of the U.S. to permit inspections. In May of that year the Russians agreed to sign a treaty giving independence to Austria, and paved the way for a Geneva summit with the U.S., U.K. and France. At the Geneva Conference Eisenhower presented a proposal called \"Open Skies\" to facilitate disarmament, which included plans for Russia and the U.S. to provide mutual access to each other's skies for open surveillance of military infrastructure. Russian leader Nikita Khrushchev dismissed the proposal out of hand.\n",
      "0.5: On 6 August 1940, Estonia was annexed by the Soviet Union as the Estonian SSR. The provisions in the Estonian constitution requiring a popular referendum to decide on joining a supra-national body were ignored. Instead the vote to join the Soviet Union was taken by those elected in the elections held the previous month. Additionally those who had failed to do their \"political duty\" of voting Estonia into the USSR, specifically those who had failed to have their passports stamped for voting, were condemned to death by Soviet tribunals. The repressions followed with the mass deportations carried out by the Soviets in Estonia on 14 June 1941. Many of the country's political and intellectual leaders were killed or deported to remote areas of the USSR by the Soviet authorities in 1940–1941. Repressive actions were also taken against thousands of ordinary people.\n",
      "----------------------------------------------------------------------\n",
      "\n",
      "\n",
      "----------------------------------------------------------------------\n",
      "Who was responsible for being the best in the Soveit space prgram?\n",
      "\n",
      "0.46: From 1920, Eisenhower served under a succession of talented generals – Fox Conner, John J. Pershing, Douglas MacArthur and George Marshall. He first became executive officer to General Conner in the Panama Canal Zone, where, joined by Mamie, he served until 1924. Under Conner's tutelage, he studied military history and theory (including Carl von Clausewitz's On War), and later cited Conner's enormous influence on his military thinking, saying in 1962 that \"Fox Conner was the ablest man I ever knew.\" Conner's comment on Eisenhower was, \"[He] is one of the most capable, efficient and loyal officers I have ever met.\" On Conner's recommendation, in 1925–26 he attended the Command and General Staff College at Fort Leavenworth, Kansas, where he graduated first in a class of 245 officers.\n",
      "0.47: On May 1, 1960, a U.S. one-man U-2 spy plane was reportedly shot down at high altitude over Soviet Union airspace. The flight was made to gain photo intelligence before the scheduled opening of an East–West summit conference, which had been scheduled in Paris, 15 days later. Captain Francis Gary Powers had bailed out of his aircraft and was captured after parachuting down onto Russian soil. Four days after Powers disappeared, the Eisenhower Administration had NASA issue a very detailed press release noting that an aircraft had \"gone missing\" north of Turkey. It speculated that the pilot might have fallen unconscious while the autopilot was still engaged, and falsely claimed that \"the pilot reported over the emergency frequency that he was experiencing oxygen difficulties.\n",
      "0.49: RCA/Jive Label Group CEO Barry Weiss left the company in March 2011 to become the new CEO of Island Def Jam and Universal Republic, which were both part of Universal Music Group. Weiss had been the RCA/Jive Label Group CEO since 2008 and was head of Jive Records since 1991.\n",
      "----------------------------------------------------------------------\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for query_response in results.json()['results']:\n",
    "    query = query_response['query']\n",
    "    answers = []\n",
    "    scores = []\n",
    "    for result in query_response['results']:\n",
    "        answers.append(result['text'])\n",
    "        scores.append(round(result['score'], 2))\n",
    "    # Print the scores and corresponding answers\n",
    "    print(\"-\"*70+\"\\n\"+query+\"\\n\\n\"+\"\\n\".join([f\"{s}: {a}\" for a, s in zip(answers, scores)])+\"\\n\"+\"-\"*70+\"\\n\\n\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.10.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}