{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "942fa22a-c776-4a44-bde9-75b7cb4202ba",
   "metadata": {},
   "source": [
    "## Outline\n",
    "\n",
    "1. We collect a dataset consisting of (user_question, answer_context, dialogue_history -> answer)\n",
    "2. We duplicate a small portion of dataset, where we remove answer_context\n",
    "2. We augment 'answer_context' with (non_answer) picked by a reasonably-performing QA system: variable ordering, consistent number of answers\n",
    "3. We train the model for exact-match generation \n",
    "- Also evaluate the exact-match ratio\n",
    "- Separately evaluate with full-context questions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "766c4c50-6e72-41b2-b6d7-1e4c3c309a68",
   "metadata": {},
   "source": [
    "### 1. Positive contexts collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "33d57a85-c079-4cf1-b9ad-3b00ce916720",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0434a258-27ca-4cec-bb85-60673fea2b16",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration default-8d557d41fc795903\n",
      "Found cached dataset json (/home/xstefan3/.cache/huggingface/datasets/json/default-8d557d41fc795903/0.0.0/0f7e3662623656454fcd2b650f34e886a7db4b9104504885bd462096cc7a9f51)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "366db7856ce341a6854a08c244aa5db1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "canard_train = datasets.load_dataset(\"json\", data_files=\"datasets/CANARD_Release/train.json\")[\"train\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "73cf434e-8d36-4680-a80d-a9304ef801f2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['History', 'QuAC_dialog_id', 'Question', 'Question_no', 'Rewrite'],\n",
       "    num_rows: 31526\n",
       "})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "canard_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "02eba563-5810-4b0a-b130-920d163a54ac",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'History': ['Johnny Unitas', '1964 MVP season'],\n",
       " 'QuAC_dialog_id': 'C_2ba58216460d43aa986fc0e897537239_0',\n",
       " 'Question': 'what team did unitas play for',\n",
       " 'Question_no': 1,\n",
       " 'Rewrite': 'what team did Johnny Unitas play for?'}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "canard_train[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b73c5e59-2430-4b00-aa8b-0f926729ada1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset quac (/home/xstefan3/.cache/huggingface/datasets/quac/plain_text/1.1.0/4170258e7e72d7c81bd6441b3f3489ea1544f0ff226ce61e22bb00c6e9d01fb6)\n"
     ]
    }
   ],
   "source": [
    "quac_train = datasets.load_dataset(\"quac\", split=\"train\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "21e21544-b65c-433d-86c1-30d4507088e7",
   "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>wikipedia_page_title</th>\n",
       "      <th>background</th>\n",
       "      <th>section_title</th>\n",
       "      <th>context</th>\n",
       "      <th>turn_ids</th>\n",
       "      <th>questions</th>\n",
       "      <th>followups</th>\n",
       "      <th>yesnos</th>\n",
       "      <th>answers</th>\n",
       "      <th>orig_answers</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>dialogue_id</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>C_69758fcdfc1f46baba0e92c0f3b0919c_1</th>\n",
       "      <td>Malayali</td>\n",
       "      <td>The Malayali people or Keralite people (also s...</td>\n",
       "      <td>Geographic distribution and population</td>\n",
       "      <td>According to the Indian census of 2001, there ...</td>\n",
       "      <td>[C_69758fcdfc1f46baba0e92c0f3b0919c_1_q#0, C_6...</td>\n",
       "      <td>[Where is Malayali located?, What other langua...</td>\n",
       "      <td>[2, 1, 1, 1, 1, 1, 1]</td>\n",
       "      <td>[2, 2, 2, 2, 2, 0, 2]</td>\n",
       "      <td>{'texts': [['30,803,747 speakers of Malayalam ...</td>\n",
       "      <td>{'texts': ['30,803,747 speakers of Malayalam i...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>C_69758fcdfc1f46baba0e92c0f3b0919c_0</th>\n",
       "      <td>Malayali</td>\n",
       "      <td>The Malayali people or Keralite people (also s...</td>\n",
       "      <td>Language and literature</td>\n",
       "      <td>Malayalam is the language spoken by the Malaya...</td>\n",
       "      <td>[C_69758fcdfc1f46baba0e92c0f3b0919c_0_q#0, C_6...</td>\n",
       "      <td>[what language do they speak?, Do they speak a...</td>\n",
       "      <td>[0, 0, 0, 0, 0, 0, 0]</td>\n",
       "      <td>[2, 2, 2, 2, 2, 2, 2]</td>\n",
       "      <td>{'texts': [['Malayalam is the language spoken ...</td>\n",
       "      <td>{'texts': ['Malayalam is the language spoken b...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                     wikipedia_page_title  \\\n",
       "dialogue_id                                                 \n",
       "C_69758fcdfc1f46baba0e92c0f3b0919c_1             Malayali   \n",
       "C_69758fcdfc1f46baba0e92c0f3b0919c_0             Malayali   \n",
       "\n",
       "                                                                             background  \\\n",
       "dialogue_id                                                                               \n",
       "C_69758fcdfc1f46baba0e92c0f3b0919c_1  The Malayali people or Keralite people (also s...   \n",
       "C_69758fcdfc1f46baba0e92c0f3b0919c_0  The Malayali people or Keralite people (also s...   \n",
       "\n",
       "                                                               section_title  \\\n",
       "dialogue_id                                                                    \n",
       "C_69758fcdfc1f46baba0e92c0f3b0919c_1  Geographic distribution and population   \n",
       "C_69758fcdfc1f46baba0e92c0f3b0919c_0                 Language and literature   \n",
       "\n",
       "                                                                                context  \\\n",
       "dialogue_id                                                                               \n",
       "C_69758fcdfc1f46baba0e92c0f3b0919c_1  According to the Indian census of 2001, there ...   \n",
       "C_69758fcdfc1f46baba0e92c0f3b0919c_0  Malayalam is the language spoken by the Malaya...   \n",
       "\n",
       "                                                                               turn_ids  \\\n",
       "dialogue_id                                                                               \n",
       "C_69758fcdfc1f46baba0e92c0f3b0919c_1  [C_69758fcdfc1f46baba0e92c0f3b0919c_1_q#0, C_6...   \n",
       "C_69758fcdfc1f46baba0e92c0f3b0919c_0  [C_69758fcdfc1f46baba0e92c0f3b0919c_0_q#0, C_6...   \n",
       "\n",
       "                                                                              questions  \\\n",
       "dialogue_id                                                                               \n",
       "C_69758fcdfc1f46baba0e92c0f3b0919c_1  [Where is Malayali located?, What other langua...   \n",
       "C_69758fcdfc1f46baba0e92c0f3b0919c_0  [what language do they speak?, Do they speak a...   \n",
       "\n",
       "                                                  followups  \\\n",
       "dialogue_id                                                   \n",
       "C_69758fcdfc1f46baba0e92c0f3b0919c_1  [2, 1, 1, 1, 1, 1, 1]   \n",
       "C_69758fcdfc1f46baba0e92c0f3b0919c_0  [0, 0, 0, 0, 0, 0, 0]   \n",
       "\n",
       "                                                     yesnos  \\\n",
       "dialogue_id                                                   \n",
       "C_69758fcdfc1f46baba0e92c0f3b0919c_1  [2, 2, 2, 2, 2, 0, 2]   \n",
       "C_69758fcdfc1f46baba0e92c0f3b0919c_0  [2, 2, 2, 2, 2, 2, 2]   \n",
       "\n",
       "                                                                                answers  \\\n",
       "dialogue_id                                                                               \n",
       "C_69758fcdfc1f46baba0e92c0f3b0919c_1  {'texts': [['30,803,747 speakers of Malayalam ...   \n",
       "C_69758fcdfc1f46baba0e92c0f3b0919c_0  {'texts': [['Malayalam is the language spoken ...   \n",
       "\n",
       "                                                                           orig_answers  \n",
       "dialogue_id                                                                              \n",
       "C_69758fcdfc1f46baba0e92c0f3b0919c_1  {'texts': ['30,803,747 speakers of Malayalam i...  \n",
       "C_69758fcdfc1f46baba0e92c0f3b0919c_0  {'texts': ['Malayalam is the language spoken b...  "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "quac_train_df = quac_train.to_pandas().set_index(\"dialogue_id\", drop=True)\n",
    "quac_train_df.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "01b2994d-3ba3-4ce0-9531-20e1858ee878",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([array(['what team did unitas play for',\n",
       "              'how many games did the colts win',\n",
       "              'who did they play in the playoffs', 'did they win the super bowl',\n",
       "              'who did they play in the super bowl', 'what were unitas stats'],\n",
       "             dtype=object)                                                       ,\n",
       "       {'texts': array([array(['The Colts'], dtype=object),\n",
       "              array(['the Colts ran off 10 straight victories to finish with a 12-2 record.'],\n",
       "                    dtype=object)                                                             ,\n",
       "              array(['Cleveland Browns'], dtype=object),\n",
       "              array(['losing 27-0.'], dtype=object),\n",
       "              array(['the Packers.'], dtype=object),\n",
       "              array(['Gary Cuozzo also suffered a season-ending injury the following'],\n",
       "                    dtype=object)                                                      ],\n",
       "             dtype=object), 'answer_starts': array([array([920], dtype=int32), array([142], dtype=int32),\n",
       "              array([552], dtype=int32), array([604], dtype=int32),\n",
       "              array([1487], dtype=int32), array([1292], dtype=int32)],\n",
       "             dtype=object)}                                                                              ],\n",
       "      dtype=object)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "quac_train_df.loc['C_2ba58216460d43aa986fc0e897537239_0'][[\"questions\", \"answers\"]].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e680d1ec-ca04-452c-ad44-eae3b43559cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def answer_for_question(questions: dict, answers: list, question: str) -> str:\n",
    "    answers = [anss[0] for anss in answers[\"texts\"]]\n",
    "    # print(questions)\n",
    "    # print(question)\n",
    "    assert question in questions\n",
    "    assert len(answers) == len(questions)\n",
    "    \n",
    "    return next(a for i, a in enumerate(answers) if questions[i] == question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9a58cc80-a4b8-4b59-b5c3-e8b095e2c281",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "786095c0ea3e432dac7dd1912cb3832d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/31526 [00:00<?, ?ex/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "canard_train = canard_train.map(lambda row: \n",
    "{\n",
    "    \"true_contexts\": quac_train_df.loc[row[\"QuAC_dialog_id\"]][\"context\"],\n",
    "    \"true_page\": quac_train_df.loc[row[\"QuAC_dialog_id\"]][\"wikipedia_page_title\"],\n",
    "    \"answer\": answer_for_question(*quac_train_df.loc[row[\"QuAC_dialog_id\"]][[\"questions\", \"answers\"]].values, row[\"Question\"])\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d894b2ca-2d74-412d-bdb3-fc42637dea18",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'History': ['Johnny Unitas', '1964 MVP season'],\n",
       " 'QuAC_dialog_id': 'C_2ba58216460d43aa986fc0e897537239_0',\n",
       " 'Question': 'what team did unitas play for',\n",
       " 'Question_no': 1,\n",
       " 'Rewrite': 'what team did Johnny Unitas play for?',\n",
       " 'true_contexts': \"The 1964 season would see the Colts return to the top of the Western Conference. After dropping their season opener to the Minnesota Vikings, the Colts ran off 10 straight victories to finish with a 12-2 record. The season was one of Unitas' best as he finished with 2,824 yards passing, a league-best 9.26 yards per pass attempt, 19 touchdown passes and only 6 interceptions. He was named the NFL's Most Valuable Player by the AP and UPI for a second time. However, the season would end on a disappointing note for the Colts as they were upset by the Cleveland Browns in the 1964 NFL Championship Game, losing 27-0.  Unitas resumed his torrid passing in 1965, as he threw for 2,530 yards, 23 touchdowns and finished with a league-high and career best 97.1 passer rating. But he was lost for the balance of the season due to a knee injury in a week 12 loss to the Bears. More postseason heartbreak would follow in 1965. The Colts and Packers finished in a tie for first place in the Western Conference and a one-game playoff was played in Green Bay to decide who would be the conference representative in the 1965 NFL Championship Game. The Colts lost in overtime 13-10 due in large part to a game-tying field goal by Don Chandler that many say was incorrectly ruled good. Backup quarterback Gary Cuozzo also suffered a season-ending injury the following week and it would be running back Tom Matte who filled in as the emergency QB for the regular-season finale and the playoff loss to the Packers.  Unitas, healthy once more, threw for 2748 yards and 22 touchdowns in 1966 in a return to Pro Bowl form. However, he posted a league-high 24 interceptions. CANNOTANSWER\",\n",
       " 'true_page': 'Johnny Unitas',\n",
       " 'answer': 'The Colts'}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "canard_train[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0dfa9659-93b9-4264-9693-08b3410c869e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(7881, 31526)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import random\n",
    "\n",
    "canard_negative_subsample = canard_train.select(random.sample(list(range(len(canard_train))), k=len(canard_train)//4))\n",
    "\n",
    "len(canard_negative_subsample), len(canard_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "50ed6cc6-0d60-4cd5-b284-899f179420e4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9ab7324375c04ead8d6f97164a56b8f7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/7881 [00:00<?, ?ex/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "canard_negative_subsample = canard_negative_subsample.map(lambda row: {\"true_contexts\": \"\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "ee19b78e-cdf0-404c-87a3-37d4e179d988",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'History': [\"Dinesh D'Souza\",\n",
       "  \"Hillary's America: The Secret History of the Democratic Party\",\n",
       "  \"Is Hillary's America a documentary?\",\n",
       "  \"On July 25, 2016, D'Souza released the documentary film Hillary's America:\",\n",
       "  'Was it released in theaters?',\n",
       "  \"I don't know.\",\n",
       "  'What was the documentary about?',\n",
       "  'The film criticizes the Democratic Party and Hillary Clinton,'],\n",
       " 'QuAC_dialog_id': 'C_31bfdcd402d44289a6206d9b34765869_0',\n",
       " 'Question': 'How did the critics feel about it?',\n",
       " 'Question_no': 4,\n",
       " 'Rewrite': \"How did the critics feel about the film Hillary's America?\",\n",
       " 'true_contexts': '',\n",
       " 'true_page': \"Dinesh D'Souza\",\n",
       " 'answer': 'The film was universally panned by professional film critics.'}"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "canard_negative_subsample[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "4f3869f7-3b9d-4efc-b194-4e2f2cb38798",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "39407"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "canard_train = datasets.concatenate_datasets([canard_train, canard_negative_subsample])\n",
    "\n",
    "len(canard_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fd2f319-764a-4b31-baad-ea9a4b037e1d",
   "metadata": {},
   "source": [
    "### 2. Negative contexts collection\n",
    "\n",
    "We use BM25 to collect a realistic set of retrieves given by the IR search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "92b7428e-8975-49b0-9482-539a454a5f9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from BM25_irsystem import BM25PlusSystem, SimpleDocProcessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "35e76ccb-a2ca-4cfe-b8a1-71c665e237f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pv211_utils.trec.entities import TrecDocumentBase, TrecQueryBase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "92d0e8e0-e7ff-48ad-a35c-543ffdd76d7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "documents = {str(i): TrecDocumentBase(document_id=i, body=context) for i, context in enumerate(quac_train_df.context)}\n",
    "\n",
    "irsystem = BM25PlusSystem(documents, preprocessing=SimpleDocProcessing())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "0e906f2f-b2b6-4f78-bd98-254cb1279181",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_negative_question_responses(question: str, num_responses: 5):\n",
    "    # TODO: add contexts' titles\n",
    "    unique_responses = []\n",
    "    # question = \"What team did Johnny Unitas play for?\"\n",
    "\n",
    "    for response_doc in irsystem.search(TrecQueryBase(query_id=0, title=\"\", body=question, narrative=\"\")):\n",
    "        if response_doc.body not in unique_responses:\n",
    "            unique_responses.append(response_doc.body)\n",
    "        if len(unique_responses) >= num_responses:\n",
    "            break\n",
    "\n",
    "    return unique_responses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "f3587fed-f244-45f0-af28-6e5e36ac15b9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5318a11fc3484b188ef6c15a6b952d95",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/39407 [00:00<?, ?ex/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/39407 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "canard_train_augm = canard_train.map(\n",
    "    lambda row: {\"retrieved_contexts\": get_negative_question_responses(row[\"Question\"], num_responses=4) \n",
    "                                                  if row[\"true_contexts\"] else get_negative_question_responses(row[\"Question\"], num_responses=5)},\n",
    "    # keep_in_memory=True,\n",
    "    # num_proc=60\n",
    ")\n",
    "canard_train_augm.save_to_disk(\"canard_train_augm_full.hf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "b244a1b0-58c9-408f-91c1-10a29ccd43ce",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['History', 'QuAC_dialog_id', 'Question', 'Question_no', 'Rewrite', 'true_contexts', 'true_page', 'answer', 'retrieved_contexts'],\n",
       "    num_rows: 39407\n",
       "})"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "datasets.load_from_disk(\"canard_train_augm_full.hf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c4c5925-9abb-4d95-941b-6634f1bdb597",
   "metadata": {},
   "source": [
    "## Test dataset generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "4c5f5447-38fc-4f7d-8b46-292de36bba8a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration default-a7ce477a9c57a36e\n",
      "Found cached dataset json (/home/xstefan3/.cache/huggingface/datasets/json/default-a7ce477a9c57a36e/0.0.0/0f7e3662623656454fcd2b650f34e886a7db4b9104504885bd462096cc7a9f51)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ece835047b44454a9f31ff14f6986640",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset quac (/home/xstefan3/.cache/huggingface/datasets/quac/plain_text/1.1.0/4170258e7e72d7c81bd6441b3f3489ea1544f0ff226ce61e22bb00c6e9d01fb6)\n"
     ]
    }
   ],
   "source": [
    "import datasets\n",
    "\n",
    "# make sure that we test with conversations that the model has not seen before\n",
    "canard_test = datasets.load_dataset(\"json\", data_files=\"datasets/CANARD_Release/test.json\")[\"train\"]\n",
    "quac_test = datasets.load_dataset(\"quac\", split=\"validation\")\n",
    "quac_test_df = quac_test.to_pandas().set_index(\"dialogue_id\", drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "68f61770-cb34-43f5-aa51-6e1a68e0d362",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'History': ['Ursula K. Le Guin',\n",
       "  'Influences',\n",
       "  'what influenced her?',\n",
       "  'Le Guin was influenced by fantasy writers,',\n",
       "  'who were they?',\n",
       "  'J. R. R. Tolkien, by science fiction writers,',\n",
       "  'how did they influence her?',\n",
       "  'her influences, she replied: Once I learned to read, I read everything. I read all the famous fantasies'],\n",
       " 'QuAC_dialog_id': 'C_420bfcf5d8b344a480ac654f08ee55ad_1',\n",
       " 'Question': 'which other fantasy writer influenced her?',\n",
       " 'Question_no': 4,\n",
       " 'Rewrite': 'Besides J. R. R. Tolkien which other fantasy writer influenced Le Guin?'}"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check the match on QuAC_dialog_id\n",
    "canard_test[102]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "8f21d07e-9d94-4d29-bb6f-a4bbd5ef833e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "wikipedia_page_title                                    Ursula K. Le Guin\n",
       "background              Ursula Kroeber Le Guin (; October 21, 1929 - J...\n",
       "section_title                                                  Influences\n",
       "context                 Le Guin was influenced by fantasy writers, inc...\n",
       "turn_ids                [C_420bfcf5d8b344a480ac654f08ee55ad_1_q#0, C_4...\n",
       "questions               [what influenced her?, who were they?, how did...\n",
       "followups                                     [0, 0, 0, 0, 1, 0, 0, 0, 1]\n",
       "yesnos                                        [2, 2, 2, 2, 2, 0, 2, 2, 2]\n",
       "answers                 {'texts': [['Le Guin was influenced by fantasy...\n",
       "orig_answers            {'texts': ['Le Guin was influenced by fantasy ...\n",
       "Name: C_420bfcf5d8b344a480ac654f08ee55ad_1, dtype: object"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "quac_test_df.loc[\"C_420bfcf5d8b344a480ac654f08ee55ad_1\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "798ad24f-1cb3-4b2d-8fe0-d92e4aa97ff5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at /home/xstefan3/.cache/huggingface/datasets/json/default-a7ce477a9c57a36e/0.0.0/0f7e3662623656454fcd2b650f34e886a7db4b9104504885bd462096cc7a9f51/cache-e93159eee5490cc9.arrow\n"
     ]
    }
   ],
   "source": [
    "canard_test = canard_test.map(lambda row: \n",
    "{\n",
    "    \"true_contexts\": quac_test_df.loc[row[\"QuAC_dialog_id\"]][\"context\"],\n",
    "    \"true_page\": quac_test_df.loc[row[\"QuAC_dialog_id\"]][\"wikipedia_page_title\"],\n",
    "    \"answer\": answer_for_question(*quac_test_df.loc[row[\"QuAC_dialog_id\"]][[\"questions\", \"answers\"]].values, row[\"Question\"])\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc8fda7c-c140-438d-ac16-7df672d492eb",
   "metadata": {},
   "source": [
    "### 2. Negative contexts collection\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "500d28a4-2dde-45ca-a12b-48d6ceefb7fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We initialize a new IR system for response - pesimistic scenario\n",
    "from BM25_irsystem import BM25PlusSystem, SimpleDocProcessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "de060c5b-00eb-4d8a-b645-fedfd4ec0b29",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pv211_utils.trec.entities import TrecDocumentBase, TrecQueryBase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "63b93efa-5e2d-4bb9-a02d-3cdc23ee1e42",
   "metadata": {},
   "outputs": [],
   "source": [
    "documents = {str(i): TrecDocumentBase(document_id=i, body=context) for i, context in enumerate(quac_test_df.context)}\n",
    "\n",
    "irsystem = BM25PlusSystem(documents, preprocessing=SimpleDocProcessing())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "923bfc60-6ac0-4c9c-9f78-c2852541a7fc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5a77c28c178243888d2672ac3f133078",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5571 [00:00<?, ?ex/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "db0ef84b89d841d7b75942cb31a3c3ac",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/5571 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "canard_test_augm = canard_test.map(\n",
    "    lambda row: {\"retrieved_contexts\": get_negative_question_responses(row[\"Question\"], num_responses=4) \n",
    "                                                  if row[\"true_contexts\"] else get_negative_question_responses(row[\"Question\"], num_responses=5)},\n",
    "    # keep_in_memory=True,\n",
    "    # num_proc=60\n",
    ")\n",
    "canard_test_augm.save_to_disk(\"canard_test_augm_full.hf\")"
   ]
  }
 ],
 "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}