{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Creating a synthetic Q&A dataset\n",
    "We use [`davinci-instruct-beta-v2`](https://beta.openai.com/docs/engines/instruct-series-beta), a model specialized in following instructions, to create questions based on the given context. Then we also use [`davinci-instruct-beta-v2`](https://beta.openai.com/docs/engines/instruct-series-beta) to answer those questions, given the same context. \n",
    "\n",
    "This is expensive, and will also take a long time, as we call the davinci engine for each section. You can simply download the final dataset instead.\n",
    "\n",
    "We're using the dataset created using the [previous notebook](olympics-1-collect-data.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.1 Read in the data, and create a context\n",
    "Create a context by concatenating the title, the heading and the content of that section"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "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>title</th>\n",
       "      <th>heading</th>\n",
       "      <th>content</th>\n",
       "      <th>tokens</th>\n",
       "      <th>context</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2020 Summer Olympics</td>\n",
       "      <td>Summary</td>\n",
       "      <td>The 2020 Summer Olympics (Japanese: 2020年夏季オリン...</td>\n",
       "      <td>713</td>\n",
       "      <td>2020 Summer Olympics\\nSummary\\n\\nThe 2020 Summ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2020 Summer Olympics</td>\n",
       "      <td>Host city selection</td>\n",
       "      <td>The International Olympic Committee (IOC) vote...</td>\n",
       "      <td>126</td>\n",
       "      <td>2020 Summer Olympics\\nHost city selection\\n\\nT...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2020 Summer Olympics</td>\n",
       "      <td>Impact of the COVID-19 pandemic</td>\n",
       "      <td>In January 2020, concerns were raised about th...</td>\n",
       "      <td>369</td>\n",
       "      <td>2020 Summer Olympics\\nImpact of the COVID-19 p...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2020 Summer Olympics</td>\n",
       "      <td>Qualifying event cancellation and postponement</td>\n",
       "      <td>Concerns about the pandemic began to affect qu...</td>\n",
       "      <td>298</td>\n",
       "      <td>2020 Summer Olympics\\nQualifying event cancell...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2020 Summer Olympics</td>\n",
       "      <td>Effect on doping tests</td>\n",
       "      <td>Mandatory doping tests were being severely res...</td>\n",
       "      <td>163</td>\n",
       "      <td>2020 Summer Olympics\\nEffect on doping tests\\n...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                  title                                         heading  \\\n",
       "0  2020 Summer Olympics                                         Summary   \n",
       "1  2020 Summer Olympics                             Host city selection   \n",
       "2  2020 Summer Olympics                 Impact of the COVID-19 pandemic   \n",
       "3  2020 Summer Olympics  Qualifying event cancellation and postponement   \n",
       "4  2020 Summer Olympics                          Effect on doping tests   \n",
       "\n",
       "                                             content  tokens  \\\n",
       "0  The 2020 Summer Olympics (Japanese: 2020年夏季オリン...     713   \n",
       "1  The International Olympic Committee (IOC) vote...     126   \n",
       "2  In January 2020, concerns were raised about th...     369   \n",
       "3  Concerns about the pandemic began to affect qu...     298   \n",
       "4  Mandatory doping tests were being severely res...     163   \n",
       "\n",
       "                                             context  \n",
       "0  2020 Summer Olympics\\nSummary\\n\\nThe 2020 Summ...  \n",
       "1  2020 Summer Olympics\\nHost city selection\\n\\nT...  \n",
       "2  2020 Summer Olympics\\nImpact of the COVID-19 p...  \n",
       "3  2020 Summer Olympics\\nQualifying event cancell...  \n",
       "4  2020 Summer Olympics\\nEffect on doping tests\\n...  "
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "df = pd.read_csv('olympics-data/olympics_sections.csv')\n",
    "df['context'] = df.title + \"\\n\" + df.heading + \"\\n\\n\" + df.content\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.2 Create questions based on the context\n",
    "Use davinci-instruct to generate a number of plausible questions relating to the Wikipedia section contents.\n",
    "\n",
    "Note: We have used temperature=0, but it may be beneficial to experiment with a higher temperature to get a higher diversity of questions.\n",
    "\n",
    "<span style=\"color:orange\">**WARNING: This step will last a long time, and consume a lot of tokens, as it calls davinci-instruct for every section to generate a number of questions.**</span>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1. What is the 2020 Summer Olympics?\n",
      "2. When did the 2020 Summer Olympics take place?\n",
      "3. Who won the most medals at the 2020 Summer Olympics?\n",
      "4. Who won the most gold medals at the 2020 Summer Olympics?\n",
      "5. Who won the most medals at the 2020 Summer Olympics?\n"
     ]
    }
   ],
   "source": [
    "import openai\n",
    "\n",
    "def get_questions(context):\n",
    "    try:\n",
    "        response = openai.Completion.create(\n",
    "            engine=\"davinci-instruct-beta-v2\",\n",
    "            prompt=f\"Write questions based on the text below\\n\\nText: {context}\\n\\nQuestions:\\n1.\",\n",
    "            temperature=0,\n",
    "            max_tokens=257,\n",
    "            top_p=1,\n",
    "            frequency_penalty=0,\n",
    "            presence_penalty=0,\n",
    "            stop=[\"\\n\\n\"]\n",
    "        )\n",
    "        return response['choices'][0]['text']\n",
    "    except:\n",
    "        return \"\"\n",
    "\n",
    "\n",
    "df['questions']= df.context.apply(get_questions)\n",
    "df['questions'] = \"1.\" + df.questions\n",
    "print(df[['questions']].values[0][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The prompt is designed to generate a number of questions. Example questions above were generated based on the summary section of the 2020 Summer Olympics page.\n",
    "\n",
    "We can observe that the questions 3 and 5 above repeat. Sometimes the generated questions could be ambiguous without the context. We will show that even despite these limitations we can create a successful model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The 2020 Summer Olympics (Japanese: 2020年夏季オリンピック, Hepburn: Nisen Nijū-nen Kaki Orinpikku), officially the Games of the XXXII Olympiad (第三十二回オリンピック競技大会, Dai Sanjūni-kai Orinpikku Kyōgi Taikai) and branded as Tokyo 2020 (東京2020, Tōkyō Nii Zero Nii Zero), was an international multi-sport event held from 23 July to 8 August 2021 in Tokyo, Japan, with some preliminary events that began on 21 July.\n",
      "Tokyo was selected as the host city during the 125th IOC Session in Buenos Aires, Argentina, on 7 September 2013. Originally scheduled to take place from 24 July to 9 August 2020, the event was postponed to 2021 in March 2020 as a result of the COVID-19 pandemic, the first such instance in the history of the Olympic Games (previous games had been cancelled but not rescheduled). However, the event retained the Tokyo 2020 name for marketing and branding purposes. It was largely held behind closed doors with no public spectators permitted due to the declaration of a state of emergency in the Greater Tokyo Area in response to the pandemic. The Summer Paralympics were held between 24 August and 5 September 2021, 16 days after the completion of the Olympics.The 2020 Games were the fourth Olympic Games to be held in Japan, following the Tokyo 1964 (Summer), Sapporo 1972 (Winter) and Nagano 1998 (Winter) games. Tokyo is the first city in Asia to hold the Summer Games twice. The 2020 Games were the second of three consecutive Olympics to be held in East Asia, following the 2018 Winter Olympics in Pyeongchang, South Korea and preceding the 2022 Winter Olympics in Beijing, China.\n",
      "New events were introduced in existing sports for 2020, including 3x3 basketball, freestyle BMX and mixed gender team events in a number of existing sports, as well as the return of madison cycling for men and an introduction of the same event for women. New IOC policies also allowed the host organizing committee to add new sports to the Olympic program for just one Games. The disciplines added by the Japanese Olympic Committee were baseball and softball, karate, sport climbing, surfing and skateboarding, the last four of which made their Olympic debuts, and the last three of which will remain on the Olympic program.The United States topped the medal count by both total golds (39) and total medals (113), with China finishing second by both respects (38 and 88). Host nation Japan finished third, setting a record for the most gold medals and total medals ever won by their delegation at an Olympic Games with 27 and 58. Great Britain finished fourth, with a total of 22 gold and 65 medals, becoming the first nation at the Summer Olympics to increase or equal their total medals won in the two Games subsequent to hosting them. The Russian delegation competing as the ROC (not to be confused with the Republic of China (Taiwan) which competed as Chinese Taipei, not ROC) finished fifth with 20 gold medals and third in the overall medal count, with 71 medals. Bermuda, the Philippines and Qatar won their first-ever Olympic gold medals. Burkina Faso, San Marino and Turkmenistan won their first-ever Olympic medals.\n"
     ]
    }
   ],
   "source": [
    "print(df.content.values[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.3 Create answers based on the context\n",
    "Use davinci-instruct to answer the questions given the relevant Wikipedia section contents\n",
    "\n",
    "Note: We have used temperature=0, but it may be beneficial to experiment with a higher temperature to get a higher diversity of questions.\n",
    "\n",
    "<span style=\"color:orange\">**WARNING: This step will last a long time, and consume a lot of tokens, as it calls davinci-instruct for every section to answer all the questions.**</span>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1. The 2020 Summer Olympics is an international multi-sport event held from 23 July to 8 August 2021 in Tokyo, Japan.\n",
      "2. The 2020 Summer Olympics took place from 23 July to 8 August 2021.\n",
      "3. The United States topped the medal count by both total golds (39) and total medals (113), with China finishing second by both respects (38 and 88).\n",
      "4. The United States topped the medal count by both total golds (39) and total medals (113), with China finishing second by both respects (38 and 88).\n",
      "5. The United States topped the medal count by both total golds (39) and total medals (113), with China finishing second by both respects (38 and 88).\n"
     ]
    }
   ],
   "source": [
    "def get_answers(row):\n",
    "    try:\n",
    "        response = openai.Completion.create(\n",
    "            engine=\"davinci-instruct-beta-v2\",\n",
    "            prompt=f\"Write questions based on the text below\\n\\nText: {row.context}\\n\\nQuestions:\\n{row.questions}\\n\\nAnswers:\\n1.\",\n",
    "            temperature=0,\n",
    "            max_tokens=257,\n",
    "            top_p=1,\n",
    "            frequency_penalty=0,\n",
    "            presence_penalty=0\n",
    "        )\n",
    "        return response['choices'][0]['text']\n",
    "    except Exception as e:\n",
    "        print (e)\n",
    "        return \"\"\n",
    "\n",
    "\n",
    "df['answers']= df.apply(get_answers, axis=1)\n",
    "df['answers'] = \"1.\" + df.answers\n",
    "df = df.dropna().reset_index().drop('index',axis=1)\n",
    "print(df[['answers']].values[0][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These are the answers to the questions above based on the context around the host city selection. \n",
    "\n",
    "We can see that answers 3-5 contain the correct answer, but instead of answering the question directly, the answer is a verbatim extraction. Despite these occasional lower quality answers, we will show that the model can learn the task reasonably well, given a high number of examples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.4 Save the Olympics Q&A dataset based on Wikipedia sections\n",
    "We save the file for use in the [next notebook](olympics-3-train-qa.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "df.to_csv('olympics-data/olympics_qa.csv', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.5 Search file (DEPRECATED)\n",
    "We create a search file ([API reference](https://beta.openai.com/docs/api-reference/files/list)), which can be used to retrieve the relevant context when a question is asked.\n",
    "\n",
    "<span style=\"color:orange\">**DEPRECATED: The /search endpoint is deprecated in favour of using embeddings. Embeddings are cheaper, faster and can support a better search experience. See [Question Answering Guide](https://github.com/openai/openai-cookbook/blob/main/examples/Question_answering_using_embeddings.ipynb) for a search implementation using the embeddings**</span>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df[df.tokens<2000]\n",
    "df[['context', 'tokens']].rename(columns={'context':'text','tokens':'metadata'}).to_json('olympics-data/olympics_search.jsonl', orient='records', lines=True)\n",
    "\n",
    "search_file = openai.File.create(\n",
    "  file=open(\"olympics-data/olympics_search.jsonl\"),\n",
    "  purpose='search'\n",
    ")\n",
    "olympics_search_fileid = search_file['id']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.6 Answer questions based on the context provided\n",
    "\n",
    "We will use a simple implementation of the answers endpoint. This works by simply using the [/search endpoint](https://beta.openai.com/docs/api-reference/searches), which searches over an indexed file to obtain the relevant sections which can be included in the context, following by a question and answering prompt given a specified model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Athletics at the 2020 Summer Olympics – Women's 4 × 100 metres relay\n",
      "Summary\n",
      "\n",
      "The women's 4 × 100 metres relay event at the 2020 Summer Olympics took place on 5 and 6 August 2021 at the Japan National Stadium. There were 16 competing relay teams, with each team having 5 members from which 4 were selected in each round.\n",
      "\n",
      "###\n",
      "\n",
      "Athletics at the 2020 Summer Olympics – Men's 4 × 100 metres relay\n",
      "Qualification\n",
      "\n",
      "National Olympic Committees (NOCs) could qualify one relay team in one of three following ways:\n",
      "The top 8 NOCs at the 2019 World Athletics Championships qualified a relay team.\n",
      "The top 8 NOCs at the 2021 World Athletics Relays qualified a relay team.\n",
      "Where an NOC placed in the top 8 at both the 2019 World Championships and the 2021 World Relays, the quota place was allocated to the world top list as of 29 June 2021. In this case, 4 teams did so, so there are 4 places available through the world rankings.A total of five athletes may be entered for a relay team. Should a NOC have also entered individual athletes in the corresponding individual event (100 m), the entered individual athletes must be included in the total of five (5) athletes entered for the relay event. In addition of five, NOCs can nominate a maximum of one alternate athlete for each team.\n",
      "The qualifying period was originally from 1 May 2019 to 29 June 2020. Due to the COVID-19 pandemic, the period was suspended from 6 April 2020 to 30 November 2020, with the end date extended to 29 June 2021. The qualifying time standards could be obtained in various meets during the given period that have the approval of the IAAF. Both indoor and outdoor meets are eligible. The most recent Area Championships may be counted in the ranking, even if not during the qualifying period.\n"
     ]
    }
   ],
   "source": [
    "from answers_with_ft import create_context, answer_question\n",
    "print(create_context(\"Where did women's 4 x 100 metres relay event take place during the 2020 Summer Olympics?\", olympics_search_fileid, max_len=400))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "' Japan National Stadium'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "answer_question(olympics_search_fileid, \"davinci-instruct-beta-v2\", \n",
    "            \"Where did women's 4 x 100 metres relay event take place during the 2020 Summer Olympics?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After we fine-tune the model for Q&A we'll be able to use it instead of [`davinci-instruct-beta-v2`](https://beta.openai.com/docs/engines/instruct-series-beta), to obtain better answers when the question can't be answered based on the context. We see a downside of [`davinci-instruct-beta-v2`](https://beta.openai.com/docs/engines/instruct-series-beta), which always attempts to answer the question, regardless of the relevant context being present or not. (Note the second question is asking about a future event, set in 2024.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "' Japan National Stadium'"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "answer_question(olympics_search_fileid, \"davinci-instruct-beta-v2\", \n",
    "            \"Where did women's 4 x 100 metres relay event take place during the 2048 Summer Olympics?\", max_len=1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that davinci has a tendency to answer the question, even if the question can't be answered given the context provided. Note the question asked regarding 2048 Summer Olympics, which didn't happen yet, and the retrieved content has only returned results for 2020."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.7 (Optional) Investigation into how likely the search endpoint is to return the relevant context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0, 58)\n"
     ]
    }
   ],
   "source": [
    "def check_context(title, heading, question, max_len=1800, search_model='ada', max_rerank=10):\n",
    "    \"\"\"\n",
    "    Evaluate the performance of the search model in retrieving the correct context\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    title: str\n",
    "        The title of the Wikipedia page\n",
    "    heading: str\n",
    "        The heading of the Wikipedia section\n",
    "    qusetion: str\n",
    "        The question\n",
    "    max_len: int\n",
    "        The maximum length of the context\n",
    "    search_model: str\n",
    "        The search model to use - `ada` is most cost effective\n",
    "    max_rerank: int\n",
    "        The maximum number of reranking documents to use the search model on\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    rank: int\n",
    "        The rank of the correct context\n",
    "    token_length: int\n",
    "        The number of tokens needed to obtain the correct context\n",
    "    \"\"\"\n",
    "    \n",
    "    try:\n",
    "        results = openai.Engine(search_model).search(\n",
    "            search_model=search_model, \n",
    "            query=question, \n",
    "            max_rerank=max_rerank,\n",
    "            file=olympics_search_fileid,\n",
    "            return_metadata=True\n",
    "        )\n",
    "        index=-1\n",
    "        returns = []\n",
    "        cur_len = 0\n",
    "        for result in results['data']:\n",
    "            cur_len += int(result['metadata']) + 4 # we add 4 tokens for the separator `\\n\\n###\\n\\n`\n",
    "            if cur_len > max_len:\n",
    "                break\n",
    "            returns.append(result['text'])\n",
    "            res = result['text'].split('\\n')\n",
    "            if res[0] == title and res[1] == heading:\n",
    "                index = len(returns) - 1\n",
    "                break\n",
    "        return index, cur_len\n",
    "    except Exception as e:\n",
    "        #print (e)\n",
    "        return []\n",
    "print(check_context(\"Athletics at the 2020 Summer Olympics – Women's 4 × 100 metres relay\", \"Summary\", \"Where did women's 4 x 100 metres relay event take place during the 2020 Summer Olympics?\", max_len=10000))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We utilize the generated questions based on context to estimate how often we can retrieve the original context. These questions are noisy, so this is not a perfect estimate.\n",
    "\n",
    "Our questions and answers are prefixed with numbered bullet points, however due to the way they were generated, they are missing the first number, hence we add \"1.\" to the list of questions (and answers).\n",
    "\n",
    "We calculate the rank of the section retrieved using ada search, and the number of tokens in the context needed to retrieve the relevant section in full."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0    [(132, 27104), (-1, 22939), (8, 2151), (2, 121...\n",
       "1    [(4, 1737), (0, 130), (8, 744), (96, 17208), (...\n",
       "2          [(0, 373), (0, 373), (-1, 40610), (1, 570)]\n",
       "3            [(0, 302), (0, 302), (5, 968), (8, 1425)]\n",
       "4                      [(0, 167), (0, 167), (2, 1442)]\n",
       "Name: ada, dtype: object"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ada_results = df.apply(lambda x: [\n",
    "                    check_context( x.title, \n",
    "                                   x.heading, \n",
    "                                   q[3:],     # remove the number prefix\n",
    "                                   max_len=1000000, # set a large number to get the full context \n",
    "                                   search_model='ada', \n",
    "                                   max_rerank=200,\n",
    "                                 ) \n",
    "                    for q in (x.questions).split('\\n') # split the questions\n",
    "                    if len(q) >10 # remove the empty questions\n",
    "                ], axis=1)\n",
    "ada_results.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "out = pd.concat([ada_results], axis=1)\n",
    "out.columns = ['ada']\n",
    "out.to_csv('olympics-data/search_engine_results.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def expand_lists(out):\n",
    "    \"\"\"\n",
    "    Expand a pandas series containing lists into a series, where each list element becomes a value on its own\n",
    "\n",
    "    Input is a row per paragraph, which has multiple questions\n",
    "    Output is a row per question\n",
    "    \"\"\"\n",
    "    cols = [pd.DataFrame(out[name].tolist()).stack().reset_index(level=1, drop=True).rename(name) for name in out.columns] \n",
    "    return pd.concat(cols, axis=1)\n",
    "\n",
    "out_expanded = expand_lists(out)\n",
    "out_expanded['rank'] = out_expanded.ada.apply(lambda x: x[0] if x != [] else -2)\n",
    "out_expanded['tokens'] = out_expanded.ada.apply(lambda x: x[1] if x != [] else -2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "74.3% of relevant paragraphs are retrieved within the first 2k tokens\n"
     ]
    }
   ],
   "source": [
    "within_2k = (out_expanded.tokens < 2000).mean()\n",
    "print(f\"{within_2k*100:.1f}% of relevant paragraphs are retrieved within the first 2k tokens\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The relevant context can be obtained 74% of the time on this dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7.4% of relevant paragraphs are not retrieved within the first 200 results\n"
     ]
    }
   ],
   "source": [
    "outside_200 = (out_expanded['rank'] == -1).mean()\n",
    "print(f\"{outside_200*100:.1f}% of relevant paragraphs are not retrieved within the first 200 results\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "7.4% of the time, this is due to the keyword search part of the search algorithm not retrieving the relevant context within the first 200 results.\n",
    "18.3% of the time this is due to the semantic search not placing the relevant context within the first 2000 tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 11664 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# plot a histogram, and add axis descriptions and title\n",
    "out_expanded[(out_expanded['rank'] >=0)&(out_expanded['rank'] <30)]['rank'].hist(bins=29)\n",
    "plt.xlabel('rank')\n",
    "plt.ylabel('count')\n",
    "plt.title('Histogram of ranks of retrieved paragraphs')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 10804 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "out_expanded[(out_expanded.tokens>=0)&(out_expanded.tokens < 2000)]['tokens'].hist(bins=29)\n",
    "plt.xlabel('tokens')\n",
    "plt.ylabel('count')\n",
    "plt.title('Histogram of the number of minimum tokens needed')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can observe that the context is most likely to be returned as one of the first results, and most likely to be returned within the first 200-500 tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-2     0.000063\n",
       "-1     0.074428\n",
       " 0     0.453420\n",
       " 1     0.089515\n",
       " 2     0.047146\n",
       " 3     0.032437\n",
       " 4     0.024139\n",
       " 5     0.019676\n",
       " 6     0.015967\n",
       " 7     0.013452\n",
       " 8     0.011189\n",
       " 9     0.009869\n",
       " 10    0.009178\n",
       "Name: rank, dtype: float64"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# normalized value_counts\n",
    "out_expanded['rank'].value_counts(normalize=True).sort_index()[:13]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "probabilities of the relevant context being returned at each rank. (-2 means a processing error, -1 means the rank is >200)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.9 64-bit ('3.9.9')",
   "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.9"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "cb9817b186a29e4e9713184d901f26c1ee05ad25243d878baff7f31bb1fef480"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
