{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "984169ca",
   "metadata": {},
   "source": [
    "# Question Answering Benchmarking: State of the Union Address\n",
    "\n",
    "Here we go over how to benchmark performance on a question answering task over a state of the union address.\n",
    "\n",
    "It is highly reccomended that you do any evaluation/benchmarking with tracing enabled. See [here](https://langchain.readthedocs.io/en/latest/tracing.html) for an explanation of what tracing is and how to set it up."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "f127fb04",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Comment this out if you are NOT using tracing\n",
    "import os\n",
    "os.environ[\"LANGCHAIN_HANDLER\"] = \"langchain\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a16b75d",
   "metadata": {},
   "source": [
    "## Loading the data\n",
    "First, let's load the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5b2d5e98",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5d66c27b9b4744989843142f08f5c1b4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading readme:   0%|          | 0.00/21.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading and preparing dataset json/LangChainDatasets--question-answering-state-of-the-union to /Users/harrisonchase/.cache/huggingface/datasets/LangChainDatasets___json/LangChainDatasets--question-answering-state-of-the-union-a7e5a3b2db4f440d/0.0.0/0f7e3662623656454fcd2b650f34e886a7db4b9104504885bd462096cc7a9f51...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9e21e2ab96a0491ea5e252720d7dfa26",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data files:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c883830e068c42d39da8406ab38574c4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/2.90k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3b085715e52e49948d2a59d27e004eba",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Extracting data files:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset json downloaded and prepared to /Users/harrisonchase/.cache/huggingface/datasets/LangChainDatasets___json/LangChainDatasets--question-answering-state-of-the-union-a7e5a3b2db4f440d/0.0.0/0f7e3662623656454fcd2b650f34e886a7db4b9104504885bd462096cc7a9f51. Subsequent calls will reuse this data.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ee900d35e27d4843b42b31811b43212b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from langchain.evaluation.loading import load_dataset\n",
    "dataset = load_dataset(\"question-answering-state-of-the-union\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ab6a716",
   "metadata": {},
   "source": [
    "## Setting up a chain\n",
    "Now we need to create some pipelines for doing question answering. Step one in that is creating an index over the data in question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c18680b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.document_loaders import TextLoader\n",
    "loader = TextLoader(\"../../modules/state_of_the_union.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7f0de2b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.indexes import VectorstoreIndexCreator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ef84ff99",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running Chroma using direct local API.\n",
      "Using DuckDB in-memory for database. Data will be transient.\n"
     ]
    }
   ],
   "source": [
    "vectorstore = VectorstoreIndexCreator().from_loaders([loader]).vectorstore"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0b5d8f6",
   "metadata": {},
   "source": [
    "Now we can create a question answering chain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8843cb0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import VectorDBQA\n",
    "from langchain.llms import OpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "573719a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = VectorDBQA.from_chain_type(llm=OpenAI(), chain_type=\"stuff\", vectorstore=vectorstore, input_key=\"question\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37d669e9",
   "metadata": {},
   "source": [
    "## Make a prediction\n",
    "\n",
    "First, we can make predictions one datapoint at a time. Doing it at this level of granularity allows use to explore the outputs in detail, and also is a lot cheaper than running over multiple datapoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "3089e409",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'question': 'What is the purpose of the NATO Alliance?',\n",
       " 'answer': 'The purpose of the NATO Alliance is to secure peace and stability in Europe after World War 2.',\n",
       " 'result': ' The NATO Alliance was created to secure peace and stability in Europe after World War 2.'}"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain(dataset[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0c16cd7",
   "metadata": {},
   "source": [
    "## Make many predictions\n",
    "Now we can make predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "24b4c66e",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions = chain.apply(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49d969fb",
   "metadata": {},
   "source": [
    "## Evaluate performance\n",
    "Now we can evaluate the predictions. The first thing we can do is look at them by eye."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1d583f03",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'question': 'What is the purpose of the NATO Alliance?',\n",
       " 'answer': 'The purpose of the NATO Alliance is to secure peace and stability in Europe after World War 2.',\n",
       " 'result': ' The purpose of the NATO Alliance is to secure peace and stability in Europe after World War 2.'}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predictions[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4783344b",
   "metadata": {},
   "source": [
    "Next, we can use a language model to score them programatically"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d0a9341d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.evaluation.qa import QAEvalChain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1612dec1",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(temperature=0)\n",
    "eval_chain = QAEvalChain.from_llm(llm)\n",
    "graded_outputs = eval_chain.evaluate(dataset, predictions, question_key=\"question\", prediction_key=\"result\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79587806",
   "metadata": {},
   "source": [
    "We can add in the graded output to the `predictions` dict and then get a count of the grades."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2a689df5",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, prediction in enumerate(predictions):\n",
    "    prediction['grade'] = graded_outputs[i]['text']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "27b61215",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Counter({' CORRECT': 7, ' INCORRECT': 4})"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from collections import Counter\n",
    "Counter([pred['grade'] for pred in predictions])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12fe30f4",
   "metadata": {},
   "source": [
    "We can also filter the datapoints to the incorrect examples and look at them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "47c692a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "incorrect = [pred for pred in predictions if pred['grade'] == \" INCORRECT\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "0ef976c1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'question': 'What is the U.S. Department of Justice doing to combat the crimes of Russian oligarchs?',\n",
       " 'answer': 'The U.S. Department of Justice is assembling a dedicated task force to go after the crimes of Russian oligarchs.',\n",
       " 'result': ' The U.S. Department of Justice is assembling a dedicated task force to go after the crimes of Russian oligarchs and is naming a chief prosecutor for pandemic fraud.',\n",
       " 'grade': ' INCORRECT'}"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "incorrect[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7710401a",
   "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}