{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "984169ca",
   "metadata": {},
   "source": [
    "# SQL Question Answering Benchmarking: Chinook\n",
    "\n",
    "Here we go over how to benchmark performance on a question answering task over a SQL database.\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": 28,
   "id": "44874486",
   "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": "0f66405e",
   "metadata": {},
   "source": [
    "## Loading the data\n",
    "\n",
    "First, let's load the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0df1393f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b220d07ee5d14909bc842b4545cdc0de",
       "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--sql-qa-chinook to /Users/harrisonchase/.cache/huggingface/datasets/LangChainDatasets___json/LangChainDatasets--sql-qa-chinook-7528565d2d992b47/0.0.0/0f7e3662623656454fcd2b650f34e886a7db4b9104504885bd462096cc7a9f51...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e89e3c8ef76f49889c4b39c624828c71",
       "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": "a8421df6c26045e8978c7086cb418222",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/1.44k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d1fb6becc3324a85bf039a53caf30924",
       "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--sql-qa-chinook-7528565d2d992b47/0.0.0/0f7e3662623656454fcd2b650f34e886a7db4b9104504885bd462096cc7a9f51. Subsequent calls will reuse this data.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9d68ad1b3e4a4bd79f92597aac4d3cc9",
       "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(\"sql-qa-chinook\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ab44d504",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'question': 'How many employees are there?', 'answer': '8'}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a16b75d",
   "metadata": {},
   "source": [
    "## Setting up a chain\n",
    "This uses the example Chinook database.\n",
    "To set it up follow the instructions on https://database.guide/2-sample-databases-sqlite/, placing the `.db` file in a notebooks folder at the root of this repository.\n",
    "\n",
    "Note that here we load a simple chain. If you want to experiment with more complex chains, or an agent, just create the `chain` object in a different way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5b2d5e98",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import OpenAI, SQLDatabase, SQLDatabaseChain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "33cdcbfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "db = SQLDatabase.from_uri(\"sqlite:///../../../notebooks/Chinook.db\")\n",
    "llm = OpenAI(temperature=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0b5d8f6",
   "metadata": {},
   "source": [
    "Now we can create a SQL database chain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "8843cb0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = SQLDatabaseChain.from_llm(llm, db, input_key=\"question\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c0062e7",
   "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": 27,
   "id": "d28c5e7d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'question': 'How many employees are there?',\n",
       " 'answer': '8',\n",
       " 'result': ' There are 8 employees.'}"
      ]
     },
     "execution_count": 27,
     "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. Note that we add a try-except because this chain can sometimes error (if SQL is written incorrectly, etc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "24b4c66e",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions = []\n",
    "predicted_dataset = []\n",
    "error_dataset = []\n",
    "for data in dataset:\n",
    "    try:\n",
    "        predictions.append(chain(data))\n",
    "        predicted_dataset.append(data)\n",
    "    except:\n",
    "        error_dataset.append(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4783344b",
   "metadata": {},
   "source": [
    "## Evaluate performance\n",
    "Now we can evaluate the predictions. We can use a language model to score them programatically"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "d0a9341d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.evaluation.qa import QAEvalChain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "1612dec1",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(temperature=0)\n",
    "eval_chain = QAEvalChain.from_llm(llm)\n",
    "graded_outputs = eval_chain.evaluate(predicted_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": 23,
   "id": "2a689df5",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, prediction in enumerate(predictions):\n",
    "    prediction['grade'] = graded_outputs[i]['text']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "27b61215",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Counter({' CORRECT': 3, ' INCORRECT': 4})"
      ]
     },
     "execution_count": 24,
     "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": 25,
   "id": "47c692a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "incorrect = [pred for pred in predictions if pred['grade'] == \" INCORRECT\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "0ef976c1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'question': 'How many employees are also customers?',\n",
       " 'answer': 'None',\n",
       " 'result': ' 59 employees are also customers.',\n",
       " 'grade': ' INCORRECT'}"
      ]
     },
     "execution_count": 26,
     "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.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}