3-Retrieval_Evaluation_Week1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7d49f27a-15ba-4176-bf6d-211051a8cef6",
   "metadata": {},
   "source": [
    "#### Week 1: Vector Search Applications w/ LLMs.  Authored by Chris Sanchez."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b8a48d4e-11e4-45b5-8cfc-f595d61e11bd",
   "metadata": {},
   "source": [
    "# Week 1 - Notebook 3\n",
    "\n",
    "# Overview\n",
    "***\n",
    "Welcome to the final notebook for Week 1! Take a look at all the ground we've covered so far:\n",
    "- Chunking/splitting\n",
    "- Vectorization of text\n",
    "- Combining with metadata\n",
    "- Saving to disk\n",
    "- Class Configuration\n",
    "- Data Indexing\n",
    "- Keyword search\n",
    "- Vector search\n",
    "- OPTIONAL: Searching with Filters\n",
    "\n",
    "We are now prepared to move on to a very important topic, **Retrieval Evaluation**.  I hope you've noticed that the search results will differ (sometimes slightly, sometimes by a lot) depeding on which search method you used: `keyword_search` or `vector_search`.  As humans, it's fairly easy for us to determine whether the returned search results are relevant to the query that was submitted, (though even here there will be differing opinions on result relevance).  But how do we systematically determine which search method is better in general?  And how do we measure the relative performance of our retrieval system if we change one of it's parameters...for example, changing our embedding model? What about measuring system performance over time as more documents are added to our datastore?\n",
    "\n",
    "We need a way to evaluate our retrieval system, and this notebook will show you \"one way\" of doing that.  I say \"one way\" because there are many ways to approach this problem, and the method I'm showing you is not perfect (if anything it's a bit too conservative).  Ultimately, measuring retrieval performance is hard because it requires a lot of time and effort, and absent any user [click-data](https://en.wikipedia.org/wiki/Click_tracking), requires some form of data labeling.  With the advent of powerful generative LLMs the process of measuring retrieval performance has become much easier. Let's take a look at how that works.\n",
    "\n",
    "# Retrieval Evaluation - Process\n",
    "***\n",
    "Here's a high-level overview of how the Retrieval Evaluation process in this notebook works:\n",
    "\n",
    "1. Generate a \"golden dataset\" of query-context pairs.  I used a pseudo-LlamaIndex implementation for this step.  I say \"pseudo\" implementation because I used LlamaIndex as the backbone, but I had to rewrite significant portions of the dataset generation code because of the opinionated way that LlamaIndex is built. 100 document chunks (contexts) were randomly selected from the Impact Theory corpus and those chunks were then submitted to the `gpt-3.5-turbo` model which generated a query that could be answered by the context.  The output was 100 query-context pairs along with associated doc_ids. \n",
    "   - **Assumptions**:\n",
    "     - The generated query-context pairs are, in fact, relevant to one another i.e. the query can be answered by the context that it's paired with\n",
    "     - The generated queries are simliar in style and length to the type of queries that end users would ask\n",
    "2. The golden dataset consists of three primary keys: `corpus`, `relevant_docs`, and `queries`\n",
    "     - The `corpus` is the original text context/chunk with it's associated `doc_id`\n",
    "     - The `queries` are the LLM generated queries, one (or more) for each entry in the `corpus`\n",
    "     - The `relevant_docs` is a simple lookup table linking the `corpus` docs to the generated `queries`\n",
    "3. We pass the golden dataset into a retrieval evluation function which does the following:\n",
    "   - Takes in a `retriever` arg (`WeaviateClient`) and a few other configuration params\n",
    "   - Iterates over all queries in the golden dataset and retrieves search results for each query from Weaviate datastore\n",
    "   - Extracts all `doc_id` values from the retrieved results\n",
    "   - Extracts the `doc_id` from the associated `relevant_docs` for each query\n",
    "   - Checks if the relevant doc_id is in the list of retrieved result doc_ids\n",
    "   - After all queries are completed a `hit_rate` score and `mrr` score are calculated for the entire golden dataset\n",
    "   - Writes results to an `eval_results` folder\n",
    "\n",
    "#### In a Nutshell\n",
    "Ulitmately, given a golden dataset consisting of queries, relevant docs, and their associated doc_ids, the `retrieval_evaluation` function is checking if the relevant doc_id is found in the list of retrieved results doc_ids, for each query.\n",
    "\n",
    "#### Problems with this Approach\n",
    "The problems with this approach are many, I'll cover a few here:\n",
    "- The **Assumptions** (see section 1 above) about the golden dataset must hold true.  Given that the pairs are generated by `gpt-3.5-turbo`, I think the first assumption will generally be true.  When reviewing the dataset I did find a few questions that were not answerable given the context, but for the most part they were.  The 2nd assumption though, is going to be dependent on your particular search use case.  I think for the purposes of this course, the questions generated are a decent reflection of how someone would query this dataset, and therefore do the job of measuring retriever performance.  But I would always check a real-world query distribution before using an approach like the one presented here.\n",
    "- This approach is conversative in that there is only \"one\" right answer.  Either the relevant `doc_id` is in the results list or it isn't.  In reality, there are going to be several documents that could potentially answer the generated query, but we have no way to account for these other relevant documents, unless of course, we want to manually add doc_ids to the golden dataset (and depending on your business case, you may actually want to do that).\n",
    "- We aren't measuring recall or precision because we aren't classifying other documents as \"negatives\".  As was just mentioned, the other documents in the results list may or may not be good matches, we just don't know.  Because we don't know, we can't really classify the other documents as \"negatives\".  So for this approach, we are measuring the [\"hit rate\"](https://uplimit.com/course/vector-search-apps/admin2/content/session_cln9hzpkl00721aah4hbz06fc/module/module_clo3hmyh0006p12cb3bmygky4) which is simply a count of the number of times that we found a relevant `doc_id` match in the results list and [Mean Reciprocal Rank (MRR)](https://uplimit.com/course/vector-search-apps/admin2/content/session_cln9hzpkl00721aah4hbz06fc/module/module_clo3hmyh0006p12cb3bmygky4).  We're using MRR over other metrics such as Mean Average Precision (MAP) because we are only looking at a [single relevant answer](https://stats.stackexchange.com/questions/127041/mean-average-precision-vs-mean-reciprocal-rank).  Hit rate is a good enough metric for determining if our retriever is retrieving quality results, and MRR will become more important later on when we add a Reranker to the mix.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "aa6a6f95-91c6-4a08-a3a0-2aafc2d56b77",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-11T00:14:38.607876300Z",
     "start_time": "2023-12-11T00:14:38.504138400Z"
    }
   },
   "outputs": [],
   "source": [
    "#standard library imports\n",
    "from typing import List, Tuple, Dict, Any\n",
    "import time\n",
    "import os\n",
    "\n",
    "# utilities\n",
    "from tqdm.notebook import tqdm\n",
    "from rich import print\n",
    "from dotenv import load_dotenv\n",
    "env = load_dotenv('./.env', override=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e4522b4-5248-4936-af87-b0a2220bdb6d",
   "metadata": {},
   "source": [
    "# Assignment 1.3\n",
    "***\n",
    "#### Instructions:\n",
    "* Import the `golden_100.json` dataset using the `from_json` method of the LlamaIndex `EmbeddingQAFinetuneDataset` Class\n",
    "  - **side note: The `EmbeddingQAFinetuneDataset` Class is the same class used for creating fine-tuning datasets\n",
    "* Instantiate a new Weaviate Client (Retriever) and set the `class_name` of the Class that you created in Notebook 2\n",
    "* Evaluate your retriever results using the `retrieval_evaluation` function\n",
    "* Submit your results in the form of a text file to Uplimit (the function autogenerates a report in the `dir_outpath` directory)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c9968b62-3814-40b9-bb2a-8c76968825c0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-11T00:20:01.796516500Z",
     "start_time": "2023-12-11T00:20:00.245672600Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "Num queries in Golden Dataset: \u001B[1;36m100\u001B[0m\n",
      "text/html": "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Num queries in Golden Dataset: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">100</span>\n</pre>\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from retrieval_evaluation import calc_hit_rate_scores, calc_mrr_scores, record_results, add_params\n",
    "from llama_index.finetuning import EmbeddingQAFinetuneDataset\n",
    "from weaviate_interface import WeaviateClient\n",
    "\n",
    "#################\n",
    "##  START CODE ##\n",
    "#################\n",
    "\n",
    "# Load QA dataset\n",
    "golden_dataset = EmbeddingQAFinetuneDataset.from_json('golden_100.json')\n",
    "\n",
    "# should see 100 queries\n",
    "print(f'Num queries in Golden Dataset: {len(golden_dataset.queries)}')\n",
    "\n",
    "### Instantiate Weaviate client and set Class name\n",
    "# read env vars from local .env file\n",
    "api_key = os.environ['WEAVIATE_API_KEY']\n",
    "url = os.environ['WEAVIATE_ENDPOINT']\n",
    "\n",
    "#instantiate client\n",
    "client = WeaviateClient(api_key, url)\n",
    "\n",
    "#check if WCS instance is live and ready\n",
    "client.is_live(), client.is_ready()\n",
    "\n",
    "class_name = \"Impact_theory_minilm_256\"\n",
    "\n",
    "#check if WCS instance is live and ready\n",
    "\n",
    "\n",
    "#################\n",
    "##  END CODE   ##\n",
    "#################"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2dbc77d4-07d3-4e3f-8330-76aefea948c6",
   "metadata": {},
   "source": [
    "#### Once your golden dataset is loaded in memory, you can view its content using dot notation like so: `golden_dataset.queries`, `golden_dataset.corpus`, `golden_dataset.relevant_docs`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "946e13a3-9350-40d9-8bfe-0cd37f2ade2f",
   "metadata": {},
   "source": [
    "# Project 1: Retrieval Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "255bdabd-770e-4b12-8bd8-54e5cdc64f77",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-11T00:21:04.787792300Z",
     "start_time": "2023-12-11T00:21:04.775804200Z"
    }
   },
   "outputs": [],
   "source": [
    "def retrieval_evaluation(dataset: EmbeddingQAFinetuneDataset, \n",
    "                         class_name: str, \n",
    "                         retriever: WeaviateClient,\n",
    "                         retrieve_limit: int=5,\n",
    "                         chunk_size: int=256,\n",
    "                         hnsw_config_keys: List[str]=['maxConnections', 'efConstruction', 'ef'],\n",
    "                         display_properties: List[str]=['doc_id', 'guest', 'content'],\n",
    "                         dir_outpath: str='./eval_results',\n",
    "                         include_miss_info: bool=False,\n",
    "                         user_def_params: Dict[str,Any]=None\n",
    "                         ) -> Dict[str, str|int|float]:\n",
    "    '''\n",
    "    Given a dataset and a retriever evaluate the performance of the retriever. Returns a dict of kw and vector\n",
    "    hit rates and mrr scores. If inlude_miss_info is True, will also return a list of kw and vector responses \n",
    "    and their associated queries that did not return a hit, for deeper analysis. Text file with results output\n",
    "    is automatically saved in the dir_outpath directory.\n",
    "\n",
    "    Args:\n",
    "    -----\n",
    "    dataset: EmbeddingQAFinetuneDataset\n",
    "        Dataset to be used for evaluation\n",
    "    class_name: str\n",
    "        Name of Class on Weaviate host to be used for retrieval\n",
    "    retriever: WeaviateClient\n",
    "        WeaviateClient object to be used for retrieval \n",
    "    retrieve_limit: int=5\n",
    "        Number of documents to retrieve from Weaviate host\n",
    "    chunk_size: int=256\n",
    "        Number of tokens used to chunk text. This value is purely for results \n",
    "        recording purposes and does not affect results. \n",
    "    display_properties: List[str]=['doc_id', 'content']\n",
    "        List of properties to be returned from Weaviate host for display in response\n",
    "    dir_outpath: str='./eval_results'\n",
    "        Directory path for saving results.  Directory will be created if it does not\n",
    "        already exist. \n",
    "    include_miss_info: bool=False\n",
    "        Option to include queries and their associated kw and vector response values\n",
    "        for queries that are \"total misses\"\n",
    "    user_def_params : dict=None\n",
    "        Option for user to pass in a dictionary of user-defined parameters and their values.\n",
    "    '''\n",
    "\n",
    "    results_dict = {'n':retrieve_limit, \n",
    "                    'Retriever': retriever.model_name_or_path, \n",
    "                    'chunk_size': chunk_size,\n",
    "                    'kw_hit_rate': 0,\n",
    "                    'kw_mrr': 0,\n",
    "                    'vector_hit_rate': 0,\n",
    "                    'vector_mrr': 0,\n",
    "                    'total_misses': 0,\n",
    "                    'total_questions':0\n",
    "                    }\n",
    "    #add hnsw configs and user defined params (if any)\n",
    "    results_dict = add_params(client, class_name, results_dict, user_def_params, hnsw_config_keys)\n",
    "    \n",
    "    start = time.perf_counter()\n",
    "    miss_info = []\n",
    "    for query_id, q in tqdm(dataset.queries.items(), 'Queries'):\n",
    "        results_dict['total_questions'] += 1\n",
    "        hit = False\n",
    "        #make Keyword, Vector, and Hybrid calls to Weaviate host\n",
    "        try:\n",
    "            kw_response = retriever.keyword_search(request=q, class_name=class_name, limit=retrieve_limit, display_properties=display_properties)\n",
    "            vector_response = retriever.vector_search(request=q, class_name=class_name, limit=retrieve_limit, display_properties=display_properties)\n",
    "            \n",
    "            #collect doc_ids and position of doc_ids to check for document matches\n",
    "            kw_doc_ids = {result['doc_id']:i for i, result in enumerate(kw_response, 1)}\n",
    "            vector_doc_ids = {result['doc_id']:i for i, result in enumerate(vector_response, 1)}\n",
    "            \n",
    "            #extract doc_id for scoring purposes\n",
    "            doc_id = dataset.relevant_docs[query_id][0]\n",
    "     \n",
    "            #increment hit_rate counters and mrr scores\n",
    "            if doc_id in kw_doc_ids:\n",
    "                results_dict['kw_hit_rate'] += 1\n",
    "                results_dict['kw_mrr'] += 1/kw_doc_ids[doc_id]\n",
    "                hit = True\n",
    "            if doc_id in vector_doc_ids:\n",
    "                results_dict['vector_hit_rate'] += 1\n",
    "                results_dict['vector_mrr'] += 1/vector_doc_ids[doc_id]\n",
    "                hit = True\n",
    "                \n",
    "            # if no hits, let's capture that\n",
    "            if not hit:\n",
    "                results_dict['total_misses'] += 1\n",
    "                miss_info.append({'query': q, 'kw_response': kw_response, 'vector_response': vector_response})\n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "            continue\n",
    "    \n",
    "\n",
    "    #use raw counts to calculate final scores\n",
    "    calc_hit_rate_scores(results_dict)\n",
    "    calc_mrr_scores(results_dict)\n",
    "    \n",
    "    end = time.perf_counter() - start\n",
    "    print(f'Total Processing Time: {round(end/60, 2)} minutes')\n",
    "    record_results(results_dict, chunk_size, dir_outpath=dir_outpath, as_text=True)\n",
    "    \n",
    "    if include_miss_info:\n",
    "        return results_dict, miss_info\n",
    "    return results_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cd5ed71-d150-4312-9d41-0b43ae43bd72",
   "metadata": {},
   "source": [
    "### Run evaluation over golden dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d42568cd-d0e3-4640-8865-37b6dff80307",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-11T00:21:16.823864800Z",
     "start_time": "2023-12-11T00:21:06.590218200Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "Queries:   0%|          | 0/100 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "ef73d91f0a144cc2a940be0181d72b42"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "Total Processing Time: \u001B[1;36m0.17\u001B[0m minutes\n",
      "text/html": "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Total Processing Time: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.17</span> minutes\n</pre>\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#################\n",
    "##  START CODE ##\n",
    "#################\n",
    "results = retrieval_evaluation(golden_dataset, class_name, client)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "outputs": [
    {
     "data": {
      "text/plain": "\u001B[1m{\u001B[0m\n    \u001B[32m'n'\u001B[0m: \u001B[1;36m5\u001B[0m,\n    \u001B[32m'Retriever'\u001B[0m: \u001B[32m'sentence-transformers/all-MiniLM-L6-v2'\u001B[0m,\n    \u001B[32m'chunk_size'\u001B[0m: \u001B[1;36m256\u001B[0m,\n    \u001B[32m'kw_hit_rate'\u001B[0m: \u001B[1;36m0.72\u001B[0m,\n    \u001B[32m'kw_mrr'\u001B[0m: \u001B[1;36m0.59\u001B[0m,\n    \u001B[32m'vector_hit_rate'\u001B[0m: \u001B[1;36m0.37\u001B[0m,\n    \u001B[32m'vector_mrr'\u001B[0m: \u001B[1;36m0.29\u001B[0m,\n    \u001B[32m'total_misses'\u001B[0m: \u001B[1;36m25\u001B[0m,\n    \u001B[32m'total_questions'\u001B[0m: \u001B[1;36m100\u001B[0m,\n    \u001B[32m'maxConnections'\u001B[0m: \u001B[1;36m32\u001B[0m,\n    \u001B[32m'efConstruction'\u001B[0m: \u001B[1;36m128\u001B[0m,\n    \u001B[32m'ef'\u001B[0m: \u001B[1;36m64\u001B[0m\n\u001B[1m}\u001B[0m\n",
      "text/html": "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">{</span>\n    <span style=\"color: #008000; text-decoration-color: #008000\">'n'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,\n    <span style=\"color: #008000; text-decoration-color: #008000\">'Retriever'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'sentence-transformers/all-MiniLM-L6-v2'</span>,\n    <span style=\"color: #008000; text-decoration-color: #008000\">'chunk_size'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">256</span>,\n    <span style=\"color: #008000; text-decoration-color: #008000\">'kw_hit_rate'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.72</span>,\n    <span style=\"color: #008000; text-decoration-color: #008000\">'kw_mrr'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.59</span>,\n    <span style=\"color: #008000; text-decoration-color: #008000\">'vector_hit_rate'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.37</span>,\n    <span style=\"color: #008000; text-decoration-color: #008000\">'vector_mrr'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.29</span>,\n    <span style=\"color: #008000; text-decoration-color: #008000\">'total_misses'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">25</span>,\n    <span style=\"color: #008000; text-decoration-color: #008000\">'total_questions'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">100</span>,\n    <span style=\"color: #008000; text-decoration-color: #008000\">'maxConnections'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">32</span>,\n    <span style=\"color: #008000; text-decoration-color: #008000\">'efConstruction'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">128</span>,\n    <span style=\"color: #008000; text-decoration-color: #008000\">'ef'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">64</span>\n<span style=\"font-weight: bold\">}</span>\n</pre>\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "print(results)\n",
    "#rsults to  json\n",
    "import json\n",
    "with open('results.json', 'w') as f:\n",
    "    json.dump(results, f)\n",
    "    "
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-12-11T00:25:37.284388300Z",
     "start_time": "2023-12-11T00:25:37.250603100Z"
    }
   },
   "id": "f651b76c65cbe5d7"
  },
  {
   "cell_type": "markdown",
   "id": "50d20ebe-be43-43fc-b43c-1657ae9569f9",
   "metadata": {},
   "source": [
    "# Conclusion\n",
    "***\n",
    "\n",
    "We now have a way of measuring the performance of our system.  This will allow you to make tweaks/changes to the system and then be able to objectively tell whether or not the tweak/change improved or degraded its performance.  Here are a few things to consider going forward:  \n",
    "\n",
    "- Keep in mind what the ulitmate goal of the system is that you are building.  For this course, we are trying to retrieve the most relevant documents as possible that will effectively address a user query, assuming the information is found within the corpus.  This means that we don't need pages and pages of relevant results, we actually only need the top 3-5, just enough to allow our Reader (the OpenAI LLM) to answer the user query.  This is an important point to be thinking about as you are making changes to the retrieval system.\n",
    "- Feel free to set the `include_miss_info` param to `True`.  Doing so will return a list of both the keyword and vector responses that did not contain the relevant `doc_id` (a \"total_miss\" means the `doc_id` was not present in either the `kw_doc_ids` or the `vector_doc_ids`).  Take a look at the style of the queries being asked and compare them with the returned responses.  Why are those responses being returned?  Are they close to the intent of the query?\n",
    "- Last but not least, you are now free to make changes to your system to improve the `hit_rate` and `mrr` scores.  If it were me, I'd start with switching out to a more performant [embedding model](https://huggingface.co/spaces/mteb/leaderboard).  There will be more opportunities to pick up some low hanging fruit, but we'll have to wait until the following week when hybrid search and Rerankers are introduced.  Whatever you do though, don't change params for the `SentenceSplitter` that you use for chunking the corpus.  Due to the way the golden dataset is derived, it's unfortunately dependent on those original `SentenceSplitter` settings remaining the same across evaluations. That is, of course, unless you want to build out your own golden dataset...."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7f49f46-a414-4cce-9e33-812be531e46d",
   "metadata": {},
   "source": [
    "# **OPTIONAL but encouraged... --> Fine-Tuning an Embedding Model\n",
    "***\n",
    "#### This exercise won't cost you anything except some time...\n",
    "\n",
    "#### It is recommended to run this specific section either locally on your machine or on Google Colab. As a proxy, it will take less than 10 minutes to run in a Macbook Air M1.\n",
    "\n",
    "Aside from switching out your emebdding model to improve your retrieval results, you could also try fine tuning your embedding model (or better yet, switch out your model and then fine tune the new one...👊).  For the longest time, the problem with fine-tuning sentence embedding models was the lack of access to high quality training data.  Generative LLMs can save you days/weeks of time, depending on how large of a dataset you want to create, by automating the process of generating high quality query/context pairs.  In this section we'll go over the step-by-step process of fine-tuning our `all-MiniLM-L6-v2` embedding model from a pre-generated training dataset consisting of only 300 question-context pairs, and then comparing it's retrieval results to our baseline retrieval scores.  I highly encourage trying this method out, I saw a 10+ point jump in `vector_hit_rate` after fine-tuning the baseline model.\n",
    "\n",
    "### Fine-tune Walkthrough\n",
    "\n",
    "1. Get baseline retrieval scores (vector Hit Rate, MRR, and total misses) using out-of-the-box baseline model.  You won't know objectively if fine-tuning had any effect if you don't measure the baseline results first.  I know this goes without saying it, but practitioners sometimes want to jump straight into model improvement without first considering their starting point.\n",
    "2. Collect a training and validation dataset.  This step has already been completed for you, courtesy of `gpt-3.5-turbo`.  LlamaIndex has a great out-of-the-box solution for generating query/context embedding pairs, but it isn't exactly plug and play, so I had to rewrite the function to achieve comptability for our course.  The training dataset consists of queries generated by the LLM that can be answered from the associated context (text chunk).  These pairs were generated using a prompt specifically written for the Impact Theory corpus so the training and validation data (for the most part) are high quality and contextually relevant. \n",
    "3. Instantiate a `SentenceTransformersFinetuneEngine` Class written by LlamaIndex which does a great job of abstracting away most of the details invovled in fine-tuning a Sentence Transformer model.\n",
    "4. Fit the model and set a path where the new model will reside.  I creaed a `models/` directory in the course repo, and included the directory in the `.gitignore` file so that models aren't being pushed with every commit.\n",
    "5. Create a new dataset (as you learned in Notebook 1) but this time create the embeddings using the new fine-tuned model.\n",
    "6. Create a new index on Weaviate using the new dataset you just created.\n",
    "7. Run the `retrieval_evaluation` function again, but this time instantiate your Weaviate client with the new fine-tuned model, but hold all other parameters constant (i.e. don't change any other parameter from the baseline run).\n",
    "8. Compare the fine-tuned retrieval results to the baseline results 🥳"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65048057-c7b1-4838-88d1-8717f2541fc6",
   "metadata": {},
   "source": [
    "### Import Training + Valid datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "73235f35-982c-446b-bc2c-7decc6cb7145",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_path = './data/training_data_300.json'\n",
    "valid_path = './data/validation_data_100.json'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "62c71fbb-95f4-40bd-a76c-f136629e6acf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"># Training Samples: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">300</span>\n",
       "# Validation Samples: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">100</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "# Training Samples: \u001B[1;36m300\u001B[0m\n",
       "# Validation Samples: \u001B[1;36m100\u001B[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "training_set = EmbeddingQAFinetuneDataset.from_json(training_path)\n",
    "valid_set = EmbeddingQAFinetuneDataset.from_json(valid_path)\n",
    "num_training_examples = len(training_set.queries)\n",
    "num_valid_examples = len(valid_set.queries)\n",
    "print(f'# Training Samples: {num_training_examples}\\n# Validation Samples: {num_valid_examples}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7407014d-ace1-4ab3-b62f-92dd0fad0246",
   "metadata": {},
   "source": [
    "### Wrangle Model output path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "7b9c2eb1-f487-44be-894e-b5bf105c702d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#always a good idea to name your fine-tuned so that you can easily identify it,\n",
    "#especially if you plan on doing multiple training runs with different params\n",
    "#also probably a good idea to include the # of training samples you are using in the name\n",
    "\n",
    "model_id = client.model_name_or_path\n",
    "model_ext = model_id.split('/')[1]\n",
    "models_dir = './models'\n",
    "if not os.path.exists('./models'):\n",
    "    os.makedirs('./models') \n",
    "else:\n",
    "    print(f'{models_dir} already exists')\n",
    "ft_model_name = f'finetuned-{model_ext}-{num_training_examples}'\n",
    "model_outpath = os.path.join(models_dir, ft_model_name)\n",
    "\n",
    "print(f'Model ID: {model_id}')\n",
    "print(f'Model Outpath: {model_outpath}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce4f8049-27d8-49fd-be18-360d3d8c6acb",
   "metadata": {},
   "source": [
    "### Instantiate your Fine-tune engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "f118bf22-9c0d-47ad-aa83-e7f61ce07a7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.finetuning import SentenceTransformersFinetuneEngine\n",
    "\n",
    "finetune_engine = SentenceTransformersFinetuneEngine(\n",
    "    training_set,\n",
    "    batch_size=32,\n",
    "    model_id=model_id,\n",
    "    model_output_path=model_outpath,\n",
    "    val_dataset=valid_set,\n",
    "    epochs=10\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2aaa840-de9e-4e17-8a9d-7c412cd873bb",
   "metadata": {},
   "source": [
    "### Fit the embedding model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5068ae60-3a73-4170-9310-4c74f9cd9563",
   "metadata": {},
   "outputs": [],
   "source": [
    "# finetune_engine.finetune()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e08e8aad-5a91-4c25-ab85-f1af821b250c",
   "metadata": {},
   "source": [
    "The `finetune` method will automatically generate the model directory using the `model_output_path` that you define.  Inside the directory will be a copy of the model itself (`pytorch_model.bin`) along with all the other files it needs.  Also in that folder, assuming you provide a `val_dataset` will be an evaluation report in the `eval` directory.  The evaluation report contains several IR metrics that may or may not be useful to you, but it does allow you to compare score improvements with each training epoch.  The new fine-tuned model is loaded through the `SentenceTransformer` class just like any other HuggingFace repo model. "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "language": "python",
   "display_name": "Python 3 (ipykernel)"
  },
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}