squadv2-rag.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c88d6fd8-81d6-4165-a408-e379f8f7192a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#!pip install -U transformers rank_bm25 evaluate unstructured bitsandbytes rouge_score python-dotenv --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "342e0db2-8c92-4f27-9f5a-600c1acfd4ad",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'4.38.2'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import transformers\n",
    "\n",
    "transformers.logging.disable_progress_bar()\n",
    "transformers.logging.set_verbosity_error()\n",
    "\n",
    "transformers.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9b870cd9-c280-48eb-847c-f920998a629b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[nltk_data] Downloading package wordnet to /home/chkei001/nltk_data...\n",
      "[nltk_data]   Package wordnet is already up-to-date!\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from model import EncoderModel, DecoderModel, BM25Model\n",
    "from store import VectorStore\n",
    "from tqdm import tqdm\n",
    "import torch\n",
    "from sklearn.metrics import ndcg_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8fefbe59-f3a0-4116-b9dc-0b43653ad905",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Token will not been saved to git credential helper. Pass `add_to_git_credential=True` if you want to set the git credential as well.\n",
      "Token is valid (permission: read).\n",
      "Your token has been saved to /home/chkei001/.cache/huggingface/token\n",
      "Login successful\n"
     ]
    }
   ],
   "source": [
    "from huggingface_hub import login\n",
    "import os\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "hf_auth = os.getenv(\"HF\")\n",
    "\n",
    "login(token=hf_auth)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd0657f5-75f8-43ea-9a86-840909870695",
   "metadata": {},
   "source": [
    "# Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d8097976-b930-407b-a56c-ba32b008cfa6",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['id', 'title', 'context', 'question', 'answers'],\n",
       "        num_rows: 130319\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['id', 'title', 'context', 'question', 'answers'],\n",
       "        num_rows: 11873\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds = load_dataset(\"squad_v2\")\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "57daaaab-380b-4627-8371-c2c1a7f3034d",
   "metadata": {
    "tags": []
   },
   "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>context</th>\n",
       "      <th>question</th>\n",
       "      <th>answers</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>The Normans (Norman: Nourmands; French: Norman...</td>\n",
       "      <td>In what country is Normandy located?</td>\n",
       "      <td>{'text': ['France', 'France', 'France', 'Franc...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>The Normans (Norman: Nourmands; French: Norman...</td>\n",
       "      <td>When were the Normans in Normandy?</td>\n",
       "      <td>{'text': ['10th and 11th centuries', 'in the 1...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>The Normans (Norman: Nourmands; French: Norman...</td>\n",
       "      <td>From which countries did the Norse originate?</td>\n",
       "      <td>{'text': ['Denmark, Iceland and Norway', 'Denm...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                             context   \n",
       "0  The Normans (Norman: Nourmands; French: Norman...  \\\n",
       "1  The Normans (Norman: Nourmands; French: Norman...   \n",
       "2  The Normans (Norman: Nourmands; French: Norman...   \n",
       "\n",
       "                                        question   \n",
       "0           In what country is Normandy located?  \\\n",
       "1             When were the Normans in Normandy?   \n",
       "2  From which countries did the Norse originate?   \n",
       "\n",
       "                                             answers  \n",
       "0  {'text': ['France', 'France', 'France', 'Franc...  \n",
       "1  {'text': ['10th and 11th centuries', 'in the 1...  \n",
       "2  {'text': ['Denmark, Iceland and Norway', 'Denm...  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "df_val = ds[\"validation\"].to_pandas()[[\"context\", \"question\", \"answers\"]]\n",
    "display(df_val.head(3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f58e5156-021e-4673-b1e8-67dca3712110",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# extract first answer of answer list\n",
    "extract_answers = lambda answer: \"\" if len(answer['text']) == 0 else answer['text'][0]\n",
    "v_extract_answers = np.vectorize(extract_answers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8758e3de-9c95-4296-b224-1d069d2c1e46",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5945/11873\n"
     ]
    }
   ],
   "source": [
    "df_val[\"answers\"] = v_extract_answers(df_val[\"answers\"].values)\n",
    "\n",
    "print(f\"{df_val[df_val['answers'] == ''].shape[0]}/{df_val.shape[0]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "23012a5b-e1d8-4932-9218-afaeac47a04d",
   "metadata": {
    "tags": []
   },
   "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>context</th>\n",
       "      <th>question</th>\n",
       "      <th>answers</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>6719</th>\n",
       "      <td>According to PolitiFact the top 400 richest Am...</td>\n",
       "      <td>What did the richest 400 Americans have as chi...</td>\n",
       "      <td>grew up in substantial privilege</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11420</th>\n",
       "      <td>The British failures in North America, combine...</td>\n",
       "      <td>How many of the Pitt's planned expeditions wer...</td>\n",
       "      <td>Two of the expeditions were successful, with F...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7963</th>\n",
       "      <td>At the same time the Mongols imported Central ...</td>\n",
       "      <td>Who did the Mongols send to Bukhara as adminis...</td>\n",
       "      <td>Han Chinese and Khitans</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9256</th>\n",
       "      <td>The other third of the water flows through the...</td>\n",
       "      <td>Where does the Nederrijn change it's name?</td>\n",
       "      <td>Wijk bij Duurstede</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6749</th>\n",
       "      <td>In Marxian analysis, capitalist firms increasi...</td>\n",
       "      <td>What do capitalist firms substitute equipment ...</td>\n",
       "      <td>labor inputs</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4613</th>\n",
       "      <td>The Very high-speed Backbone Network Service (...</td>\n",
       "      <td>What were select locations connected to?</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>257</th>\n",
       "      <td>When considering computational problems, a pro...</td>\n",
       "      <td>What is a string over a Greek number when cons...</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>233</th>\n",
       "      <td>Closely related fields in theoretical computer...</td>\n",
       "      <td>What is the process that asks a more specific ...</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4784</th>\n",
       "      <td>A variety of alternatives to the Y. pestis hav...</td>\n",
       "      <td>In what year was Scott and Duncan's research p...</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>319</th>\n",
       "      <td>A Turing machine is a mathematical model of a ...</td>\n",
       "      <td>What is a scientific device that manipulates s...</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>500 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                 context   \n",
       "6719   According to PolitiFact the top 400 richest Am...  \\\n",
       "11420  The British failures in North America, combine...   \n",
       "7963   At the same time the Mongols imported Central ...   \n",
       "9256   The other third of the water flows through the...   \n",
       "6749   In Marxian analysis, capitalist firms increasi...   \n",
       "...                                                  ...   \n",
       "4613   The Very high-speed Backbone Network Service (...   \n",
       "257    When considering computational problems, a pro...   \n",
       "233    Closely related fields in theoretical computer...   \n",
       "4784   A variety of alternatives to the Y. pestis hav...   \n",
       "319    A Turing machine is a mathematical model of a ...   \n",
       "\n",
       "                                                question   \n",
       "6719   What did the richest 400 Americans have as chi...  \\\n",
       "11420  How many of the Pitt's planned expeditions wer...   \n",
       "7963   Who did the Mongols send to Bukhara as adminis...   \n",
       "9256          Where does the Nederrijn change it's name?   \n",
       "6749   What do capitalist firms substitute equipment ...   \n",
       "...                                                  ...   \n",
       "4613           What were select locations connected to?    \n",
       "257    What is a string over a Greek number when cons...   \n",
       "233    What is the process that asks a more specific ...   \n",
       "4784   In what year was Scott and Duncan's research p...   \n",
       "319    What is a scientific device that manipulates s...   \n",
       "\n",
       "                                                 answers  \n",
       "6719                    grew up in substantial privilege  \n",
       "11420  Two of the expeditions were successful, with F...  \n",
       "7963                             Han Chinese and Khitans  \n",
       "9256                                  Wijk bij Duurstede  \n",
       "6749                                        labor inputs  \n",
       "...                                                  ...  \n",
       "4613                                                      \n",
       "257                                                       \n",
       "233                                                       \n",
       "4784                                                      \n",
       "319                                                       \n",
       "\n",
       "[500 rows x 3 columns]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# sample 400 answerable examples and 100 unanswerable examples \n",
    "test_set_answerable = df_val[df_val['answers'] != ''].sample(n=400, random_state=1)\n",
    "test_set_not_answerable = df_val[df_val['answers'] == ''].sample(n=100, random_state=1)\n",
    "test_set = pd.concat([test_set_answerable, test_set_not_answerable])\n",
    "test_set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0c025253-4477-483b-b184-8704b6236466",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# calcualte true binary relevance for ndcg\n",
    "def true_binary_relevance(result_idxs, original_id):\n",
    "    return [1 if i == original_id else 0 for i in result_idxs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49b69f12-bf9d-493b-b637-100ba46ad3dd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "retriever_models = [\n",
    "    \"sentence-transformers/all-MiniLM-L6-v2\",\n",
    "    \"BAAI/bge-base-en-v1.5\",\n",
    "    \"WhereIsAI/UAE-Large-V1\",\n",
    "    \"BAAI/bge-m3\"\n",
    "]\n",
    "causal_models = [\n",
    "    \"HuggingFaceH4/zephyr-7b-beta\",\n",
    "    \"mistralai/Mistral-7B-Instruct-v0.2\",\n",
    "    \"HuggingFaceH4/zephyr-7b-gemma-v0.1\",\n",
    "    \"mistralai/Mixtral-8x7B-Instruct-v0.1\",\n",
    "    \"google/gemma-2b-it\",\n",
    "    \"google/gemma-7b-it\",\n",
    "]\n",
    "\n",
    "retriever_results = []\n",
    "\n",
    "# causal models loop\n",
    "for causal_id in causal_models:\n",
    "    causal_lm = DecoderModel(causal_id, device=\"auto\")\n",
    "    \n",
    "    # retriever models loop\n",
    "    for retriever_id in retriever_models:\n",
    "        \n",
    "        # retriever setup loop\n",
    "        for hybrid in [True, False]:\n",
    "            # init new vector store with retriever\n",
    "            db = VectorStore(retriever_id, hybrid)\n",
    "            # embed documents\n",
    "            db.add_documents(test_set[\"context\"].values.tolist(), test_set.index.tolist(), batch_size=50)\n",
    "            \n",
    "            is_hybrid = 'yes' if hybrid else 'no'\n",
    "            causal_lm_results = []\n",
    "            \n",
    "            print(f\"Retriever: {retriever_id} - Causal LM: {causal_id} - hybrid: {is_hybrid}\")\n",
    "            \n",
    "            with tqdm(total=len(test_set.question.values)) as pbar:\n",
    "                # loop through dataset\n",
    "                for document_id, (_, query, correct_answer) in test_set.iterrows():\n",
    "                    \n",
    "                    best_contexts = \"\"\n",
    "                    best_ndcg = -1000000\n",
    "                    \n",
    "                    # loop distance metrics\n",
    "                    for distance_metric in [\"cosine\", \"ip\", \"l2\"]:\n",
    "                        # retrieve documents\n",
    "                        results = db.search(query)\n",
    "                        \n",
    "                        # unpack results\n",
    "                        idxs = [result[\"id\"] for result in results]\n",
    "                        scores = [result[\"score\"] for result in results]\n",
    "                        contexts = [result[\"document\"] for result in results]\n",
    "                        \n",
    "                        # retriever results\n",
    "                        true_relevance = true_binary_relevance(idxs, document_id)\n",
    "                        ndcg = ndcg_score([true_relevance], [scores])\n",
    "                        \n",
    "                        # Only save results for examples for which a context could be found\n",
    "                        if correct_answer != \"\":\n",
    "                            retriever_results.append({\n",
    "                                \"model\": retriever_id,\n",
    "                                \"ndcg\": ndcg,\n",
    "                                \"metric\": distance_metric,\n",
    "                                \"hybrid\": is_hybrid\n",
    "                            })\n",
    "                        \n",
    "                        \n",
    "                        # caching to give generator best possible context\n",
    "                        if ndcg > best_ndcg:\n",
    "                            best_ndcg = ndcg\n",
    "                            best_contexts = contexts\n",
    "\n",
    "                    # concatenate list of contexts to one string\n",
    "                    context_input = \"\\n\\n\".join(best_contexts)\n",
    "                    \n",
    "                    # generate an answer\n",
    "                    answer = causal_lm(query, context_input)\n",
    "\n",
    "                    causal_lm_results.append(\n",
    "                        {\n",
    "                            \"model\": causal_id,\n",
    "                            \"question\": query,\n",
    "                            \"answer\": answer,\n",
    "                            \"context\": context_input,\n",
    "                            \"correct_answer\": correct_answer if correct_answer != \"\" else \"Not answerable from the given context.\"\n",
    "                        }\n",
    "                    )\n",
    "                    pbar.update(1)\n",
    "                del db\n",
    "                torch.cuda.empty_cache()\n",
    "                pd.DataFrame(causal_lm_results).to_csv(f\"./results/{causal_id.replace('/', '-')}_({retriever_id.replace('/', '-')})_{is_hybrid}.csv\")\n",
    "    del causal_lm\n",
    "    torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eda6c70b-c643-4d8b-be7b-c4c99aab1889",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "pd.DataFrame(retriever_results).to_csv(\"retriever_results_v2.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e6eb33a-4bcc-4afc-9062-bc40ab1381d8",
   "metadata": {},
   "source": [
    "# Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0d56c20b-082e-4939-baa9-6a2d8bb2432a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "retriever_results = pd.read_csv(\"retriever_results_v2.csv\", index_col=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adf91860-b6b6-4acd-946e-5d75f2426a45",
   "metadata": {},
   "source": [
    "## Retriever"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a7558597-64d2-4f75-887c-1c0b9eead975",
   "metadata": {
    "tags": []
   },
   "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></th>\n",
       "      <th></th>\n",
       "      <th>ndcg</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>model</th>\n",
       "      <th>metric</th>\n",
       "      <th>hybrid</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"6\" valign=\"top\">BAAI/bge-base-en-v1.5</th>\n",
       "      <th rowspan=\"2\" valign=\"top\">cosine</th>\n",
       "      <th>no</th>\n",
       "      <td>0.855504</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>yes</th>\n",
       "      <td>0.869621</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">ip</th>\n",
       "      <th>no</th>\n",
       "      <td>0.855504</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>yes</th>\n",
       "      <td>0.869621</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">l2</th>\n",
       "      <th>no</th>\n",
       "      <td>0.855504</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>yes</th>\n",
       "      <td>0.869621</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"6\" valign=\"top\">BAAI/bge-m3</th>\n",
       "      <th rowspan=\"2\" valign=\"top\">cosine</th>\n",
       "      <th>no</th>\n",
       "      <td>0.795663</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>yes</th>\n",
       "      <td>0.859185</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">ip</th>\n",
       "      <th>no</th>\n",
       "      <td>0.795663</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>yes</th>\n",
       "      <td>0.859185</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">l2</th>\n",
       "      <th>no</th>\n",
       "      <td>0.795663</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>yes</th>\n",
       "      <td>0.859185</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"6\" valign=\"top\">WhereIsAI/UAE-Large-V1</th>\n",
       "      <th rowspan=\"2\" valign=\"top\">cosine</th>\n",
       "      <th>no</th>\n",
       "      <td>0.857350</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>yes</th>\n",
       "      <td>0.870450</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">ip</th>\n",
       "      <th>no</th>\n",
       "      <td>0.857350</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>yes</th>\n",
       "      <td>0.870450</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">l2</th>\n",
       "      <th>no</th>\n",
       "      <td>0.857350</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>yes</th>\n",
       "      <td>0.870450</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"6\" valign=\"top\">sentence-transformers/all-MiniLM-L6-v2</th>\n",
       "      <th rowspan=\"2\" valign=\"top\">cosine</th>\n",
       "      <th>no</th>\n",
       "      <td>0.822904</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>yes</th>\n",
       "      <td>0.873533</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">ip</th>\n",
       "      <th>no</th>\n",
       "      <td>0.822904</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>yes</th>\n",
       "      <td>0.873533</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">l2</th>\n",
       "      <th>no</th>\n",
       "      <td>0.822904</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>yes</th>\n",
       "      <td>0.873533</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                          ndcg\n",
       "model                                  metric hybrid          \n",
       "BAAI/bge-base-en-v1.5                  cosine no      0.855504\n",
       "                                              yes     0.869621\n",
       "                                       ip     no      0.855504\n",
       "                                              yes     0.869621\n",
       "                                       l2     no      0.855504\n",
       "                                              yes     0.869621\n",
       "BAAI/bge-m3                            cosine no      0.795663\n",
       "                                              yes     0.859185\n",
       "                                       ip     no      0.795663\n",
       "                                              yes     0.859185\n",
       "                                       l2     no      0.795663\n",
       "                                              yes     0.859185\n",
       "WhereIsAI/UAE-Large-V1                 cosine no      0.857350\n",
       "                                              yes     0.870450\n",
       "                                       ip     no      0.857350\n",
       "                                              yes     0.870450\n",
       "                                       l2     no      0.857350\n",
       "                                              yes     0.870450\n",
       "sentence-transformers/all-MiniLM-L6-v2 cosine no      0.822904\n",
       "                                              yes     0.873533\n",
       "                                       ip     no      0.822904\n",
       "                                              yes     0.873533\n",
       "                                       l2     no      0.822904\n",
       "                                              yes     0.873533"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "retriever_results.groupby(['model', 'metric', 'hybrid']).mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2b25177-f031-4e65-aa72-04848dfdb611",
   "metadata": {},
   "source": [
    "# Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "8bd3f482-3de1-4069-9ba8-d35171676382",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from nltk.translate.bleu_score import sentence_bleu\n",
    "from rouge_score import rouge_scorer\n",
    "\n",
    "scorer = rouge_scorer.RougeScorer(['rouge1', 'rougeL'], use_stemmer=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "ae04ec86-1148-431f-ba2b-546e6c17ab71",
   "metadata": {
    "tags": []
   },
   "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>model</th>\n",
       "      <th>bleu</th>\n",
       "      <th>rouge_1_precision</th>\n",
       "      <th>rouge_1_recall</th>\n",
       "      <th>rouge_1_fmeasure</th>\n",
       "      <th>rouge_L_precision</th>\n",
       "      <th>rouge_L_recall</th>\n",
       "      <th>rouge_L_fmeasure</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>google/gemma-2b-it</td>\n",
       "      <td>9.452660e-232</td>\n",
       "      <td>0.090218</td>\n",
       "      <td>0.394386</td>\n",
       "      <td>0.133061</td>\n",
       "      <td>0.080232</td>\n",
       "      <td>0.364439</td>\n",
       "      <td>0.118982</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>mistralai/Mistral-7B-Instruct-v0.2</td>\n",
       "      <td>8.661046e-232</td>\n",
       "      <td>0.067992</td>\n",
       "      <td>0.510451</td>\n",
       "      <td>0.107381</td>\n",
       "      <td>0.063086</td>\n",
       "      <td>0.476818</td>\n",
       "      <td>0.099192</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>HuggingFaceH4/zephyr-7b-beta</td>\n",
       "      <td>7.555362e-232</td>\n",
       "      <td>0.040825</td>\n",
       "      <td>0.559209</td>\n",
       "      <td>0.071184</td>\n",
       "      <td>0.037134</td>\n",
       "      <td>0.526515</td>\n",
       "      <td>0.064916</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>google/gemma-7b-it</td>\n",
       "      <td>1.036681e-231</td>\n",
       "      <td>0.000559</td>\n",
       "      <td>0.001250</td>\n",
       "      <td>0.000444</td>\n",
       "      <td>0.000559</td>\n",
       "      <td>0.001250</td>\n",
       "      <td>0.000444</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>mistralai/Mixtral-8x7B-Instruct-v0.1</td>\n",
       "      <td>9.063576e-232</td>\n",
       "      <td>0.076472</td>\n",
       "      <td>0.522214</td>\n",
       "      <td>0.121311</td>\n",
       "      <td>0.070003</td>\n",
       "      <td>0.483617</td>\n",
       "      <td>0.110831</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>HuggingFaceH4/zephyr-7b-gemma-v0.1</td>\n",
       "      <td>7.211556e-232</td>\n",
       "      <td>0.028043</td>\n",
       "      <td>0.474737</td>\n",
       "      <td>0.049678</td>\n",
       "      <td>0.026126</td>\n",
       "      <td>0.453578</td>\n",
       "      <td>0.046420</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                  model           bleu  rouge_1_precision   \n",
       "0                    google/gemma-2b-it  9.452660e-232           0.090218  \\\n",
       "1    mistralai/Mistral-7B-Instruct-v0.2  8.661046e-232           0.067992   \n",
       "2          HuggingFaceH4/zephyr-7b-beta  7.555362e-232           0.040825   \n",
       "3                    google/gemma-7b-it  1.036681e-231           0.000559   \n",
       "4  mistralai/Mixtral-8x7B-Instruct-v0.1  9.063576e-232           0.076472   \n",
       "5    HuggingFaceH4/zephyr-7b-gemma-v0.1  7.211556e-232           0.028043   \n",
       "\n",
       "   rouge_1_recall  rouge_1_fmeasure  rouge_L_precision  rouge_L_recall   \n",
       "0        0.394386          0.133061           0.080232        0.364439  \\\n",
       "1        0.510451          0.107381           0.063086        0.476818   \n",
       "2        0.559209          0.071184           0.037134        0.526515   \n",
       "3        0.001250          0.000444           0.000559        0.001250   \n",
       "4        0.522214          0.121311           0.070003        0.483617   \n",
       "5        0.474737          0.049678           0.026126        0.453578   \n",
       "\n",
       "   rouge_L_fmeasure  \n",
       "0          0.118982  \n",
       "1          0.099192  \n",
       "2          0.064916  \n",
       "3          0.000444  \n",
       "4          0.110831  \n",
       "5          0.046420  "
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "results_path = Path(\"./results/\")\n",
    "\n",
    "result_rows = []\n",
    "\n",
    "for file in results_path.glob(\"*.csv\"):\n",
    "    df = pd.read_csv(file, index_col=0)\n",
    "    model = df.model.values[0]\n",
    "    model_result = {\n",
    "        \"bleu\": [],\n",
    "        \"rouge_1_precision\": [],\n",
    "        \"rouge_1_recall\": [],\n",
    "        \"rouge_1_fmeasure\": [],\n",
    "        \"rouge_L_precision\": [],\n",
    "        \"rouge_L_recall\": [],\n",
    "        \"rouge_L_fmeasure\": [],\n",
    "    }\n",
    "    for i, (model_id, question, answer, context, correct_answer) in df.iterrows():\n",
    "        bleu = sentence_bleu(\n",
    "            references=correct_answer,\n",
    "            hypothesis=answer\n",
    "        )\n",
    "    \n",
    "        scores = scorer.score(correct_answer, answer)\n",
    "        precision, recall, fmeasure = scores[\"rouge1\"]\n",
    "        precision_L, recall_L, fmeasure_L = scores[\"rougeL\"]\n",
    "        \n",
    "        model_result[\"bleu\"].append(bleu)\n",
    "        model_result[\"rouge_1_precision\"].append(precision)\n",
    "        model_result[\"rouge_1_recall\"].append(recall)\n",
    "        model_result[\"rouge_1_fmeasure\"].append(fmeasure)\n",
    "        model_result[\"rouge_L_precision\"].append(precision_L)\n",
    "        model_result[\"rouge_L_recall\"].append(recall_L)\n",
    "        model_result[\"rouge_L_fmeasure\"].append(fmeasure_L)\n",
    "        \n",
    "    result_rows.append(\n",
    "        {\n",
    "            \"model\": model,\n",
    "            \"bleu\": np.mean(model_result[\"bleu\"]),\n",
    "            \"rouge_1_precision\": np.mean(model_result[\"rouge_1_precision\"]),\n",
    "            \"rouge_1_recall\": np.mean(model_result[\"rouge_1_recall\"]),\n",
    "            \"rouge_1_fmeasure\": np.mean(model_result[\"rouge_1_fmeasure\"]),\n",
    "            \"rouge_L_precision\": np.mean(model_result[\"rouge_L_precision\"]),\n",
    "            \"rouge_L_recall\": np.mean(model_result[\"rouge_L_recall\"]),\n",
    "            \"rouge_L_fmeasure\": np.mean(model_result[\"rouge_L_fmeasure\"])\n",
    "        }\n",
    "    )\n",
    "pd.DataFrame(result_rows)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcd42c9a-bfa2-46e7-a506-3f1d3af5c9ef",
   "metadata": {},
   "source": [
    "# Judging LLM-as-a-Judge"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5856ae7-d590-4a69-a1e1-b2cd66dd916b",
   "metadata": {},
   "source": [
    "Results can be found in `llm-as-a-judge.ipynb`."
   ]
  }
 ],
 "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}