demo-boosting-a-small-llm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "48cd79d5",
   "metadata": {},
   "source": [
    "# RAGBooster\n",
    "\n",
    "We detail how to improve the performance of LLMs for question answering with retrieval augmentation and learned data pruning. We showcase how these techniques allow a small open source LLM with 3 billion parameters to perform on par with a large commercial LLM which has 175 billion parameters.\n",
    "\n",
    "Our library \"RAGBooster\" for learned data pruning is available as open source.\n",
    "\n",
    "### Setup\n",
    "\n",
    "In order to run this demo, we need access to GPT3.5, the Bing web API and a deployed version of the RedPajama-INCITE-Instruct-3B-v1 model from Together. Note that we implement caching and can serve the vast majority of requests from our cache for this particular demo setup.\n",
    "\n",
    " 1. **Access to GPT3.5 from OpenAI**: This demo notebook leverages GPT3.5 via the OpenAI API. It requires you to make your [OpenAI API key](https://platform.openai.com/account/api-keys) available as an environment variable via the following command:<br/><br/>`export OPENAI_API_KEY=your_secret_openai_key`<br/><br/>\n",
    "\n",
    " 1. **Access to RedPajama-INCITE-Instruct-3B-v1**: This demo also uses the 3B param language model [RedPajama-INCITE-Instruct-3B-v1](https://huggingface.co/togethercomputer/RedPajama-INCITE-Instruct-3B-v1). This model should be made available via REST API through the [manifest project](https://github.com/HazyResearch/manifest#local-huggingface-models) as follows:<br/> \n",
    "<br/> `python -m manifest.api.app \\`<br/>\n",
    "`   --model_type huggingface \\`<br/>\n",
    "`   --model_name_or_path togethercomputer/RedPajama-INCITE-Instruct-3B-v1 \\`<br/>\n",
    "`   --model_generation_type text-generation`<br/>\n",
    "`   --device 0`<br/><br/>\n",
    " \n",
    " 1. **Access to the Bing Websearch API**: Furthermore, we will query the web via [Microsoft Bing's websearch API](https://www.microsoft.com/en-us/bing/apis/bing-web-search-api). You need to make your Bing API key available as an environment variable via the following command:<br/><br/>`export BING_SUBSCRIPTION_KEY=your_secret_bing_key`<br/><br/>\n",
    " \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b48947c5",
   "metadata": {},
   "source": [
    "## Question Answering with Large Language Models\n",
    "\n",
    "The scenario for this demo is question answering with Large Language Models (LLMs). We use a dataset of questions about the place of birth of various people from the Wikifact dataset in Stanford's [HELM benchmark](https://crfm.stanford.edu/helm/latest/). We work with a sample of 500 questions from the data as final test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a46e57bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ragbooster import Generator, BingRetriever, RetrievalAugmentedModel, RAGBooster, score\n",
    "from ragbooster.demo import load_wikifact_questions\n",
    "\n",
    "questions = load_wikifact_questions('demo_data/wikifact_place_of_birth_helm.json')\n",
    "\n",
    "validation_questions = questions[:500]\n",
    "test_questions = questions[500:1000]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8c7307d",
   "metadata": {},
   "source": [
    "An example question is about the birth place of the Slovak ice hockey player Martin Kulha:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9fca76f9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Question(text='Martin Kulha was born in', correct_answers=['Poprad'], metadata={})"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "example_question = questions[5]\n",
    "example_question"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6b9ba85",
   "metadata": {},
   "source": [
    "### Question Answering with GPT3.5\n",
    "\n",
    "Let's see how well OpenAI's `'text-davinci-003'` model from the [GPT3.5 family](https://platform.openai.com/docs/models/gpt-3-5) is doing on these questions. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0f356246",
   "metadata": {},
   "outputs": [],
   "source": [
    "from manifest import Manifest \n",
    "\n",
    "gpt35_client = Manifest(client_name=\"openai\", engine=\"text-davinci-003\",\n",
    "                        cache_name=\"sqlite\", cache_connection=\"demo_data/gpt35-cache.sqlite\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "462a1c75",
   "metadata": {},
   "source": [
    "We can leverage GPT3.5 by extending the `Generator` class. We write a couple of lines of Python to define how create our prompt from the question and some few-shot examples, and how to parse the answer returned by GPT3.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "83a38b0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re \n",
    "\n",
    "class PlaceOfBirthGenerator(Generator):\n",
    "    \n",
    "    FEW_SHOT_PROMPT = \"Brown was born in England\\n\\n\"+\\\n",
    "        \"Jerry Beck was born in New York City\\n\\n\"+\\\n",
    "        \"Werner Lorenz was born in Ludwigshafen\\n\\n\"+\\\n",
    "        \"Moritz Retzsch was born in Dresden\\n\\n\"+\\\n",
    "        \"Roni Rosadi was born in Bandar Lampung\\n\\n\"      \n",
    "    \n",
    "    def __init__(self, llm):\n",
    "        super().__init__(llm=llm, max_tokens=10)    \n",
    "    \n",
    "    def _create_prompt(self, question, params):        \n",
    "        return f\"{self.FEW_SHOT_PROMPT}\\n\\n{question.text}\"            \n",
    "\n",
    "    def _extract_answer(self, response):\n",
    "        answer = response.get_response()          \n",
    "        answer = re.sub(r'[0-9]+', '', answer)\n",
    "        answer = answer.strip()   \n",
    "\n",
    "        for sep in ['\\n', ',', '.']:\n",
    "            if sep in answer:\n",
    "                answer = answer.split(sep)[0]\n",
    "\n",
    "        return answer.strip()  \n",
    "    \n",
    "gpt35 = PlaceOfBirthGenerator(llm=gpt35_client)    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b8aec60",
   "metadata": {},
   "source": [
    "Unfortunately, GPT3.5 gives us the wrong answer to the example question!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bfdad076",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPT3.5 answers \"Martin Kulha was born in\" with \"Prague\"\n",
      "The Correct answer is: Poprad\n"
     ]
    }
   ],
   "source": [
    "print(f'GPT3.5 answers \"{example_question.text}\" with \"{gpt35.generate(example_question)}\"')\n",
    "print(f'The Correct answer is: {example_question.correct_answers[0]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb7ebb63",
   "metadata": {},
   "source": [
    "We can also evaluate GPT3.5 on all our 500 test questions and find that it only answers 14% of the questions correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a13f3e1f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "'The accuracy of GPT3.5 on the test set is 0.14'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "accuracy = score(test_questions, gpt35)\n",
    "\n",
    "f'The accuracy of GPT3.5 on the test set is {accuracy}'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "245a1d27",
   "metadata": {},
   "source": [
    "## Question Answering with RedPajama-INCITE-Instruct-3B-v1\n",
    "\n",
    "Let's see how the smaller model RedPajama-INCITE-Instruct-3B-v1 is doing on this task. We connect to our local instance of this model as follows (Please adjust the code if you use a different port)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "15fdf826",
   "metadata": {},
   "outputs": [],
   "source": [
    "redpajama_port=5291\n",
    "redpajama_client = Manifest(client_name = \"huggingface\", client_connection = f\"http://127.0.0.1:{redpajama_port}\",\n",
    "                            cache_name='sqlite', cache_connection=\"demo_data/rp3b-cache.sqlite\")\n",
    "\n",
    "redpajama = PlaceOfBirthGenerator(llm=redpajama_client)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bc929ec",
   "metadata": {},
   "source": [
    "_Due a [caching bug](https://github.com/HazyResearch/manifest/issues/103) in manifest, we need to apply the following hack to get performant caching. This code will become unnecessary as soon as the bug is fixed._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6f670e99",
   "metadata": {},
   "outputs": [],
   "source": [
    "from manifest.clients.huggingface import HuggingFaceClient\n",
    "import types\n",
    "\n",
    "client = redpajama_client.client_pool.get_current_client()\n",
    "redpajama_model_params = client.get_model_params()\n",
    "\n",
    "def cached_params(self):\n",
    "    return redpajama_model_params\n",
    "\n",
    "client.get_model_params = types.MethodType(cached_params, client)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09f9fe25",
   "metadata": {},
   "source": [
    "RedPajama-INCITE-Instruct-3B-v1 also gives us the wrong answer to the example question!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ee607918",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RedPajama3B answers \"Martin Kulha was born in\" with \"Prague\"\n",
      "The correct answer is: Poprad\n"
     ]
    }
   ],
   "source": [
    "print(f'RedPajama3B answers \"{example_question.text}\" with \"{redpajama.generate(example_question)}\"')\n",
    "print(f'The correct answer is: {example_question.correct_answers[0]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b18e2c4",
   "metadata": {},
   "source": [
    "Our task is also difficult for RedPajama-INCITE-Instruct-3B-v1: it answers only 10% of the test questions correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "daa98e02",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "'The accuracy of RedPajama3B on the test set is 0.1'"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "accuracy = score(test_questions, redpajama)\n",
    "\n",
    "f'The accuracy of RedPajama3B on the test set is {accuracy}'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a52d7ec3",
   "metadata": {},
   "source": [
    "## Retrieval Augmentation with Bing Websearch\n",
    "\n",
    "We can improve the performance of our LLMs by providing them with some external data to answer the questions, for example from the web. This is called retrieval augmentation, and we use Microsoft's Bing websearch API for that by extending the `BingRetriever` class and defining how to create a query from the question text. In our case, we can just use the question text as the query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "de7a04e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyBingWebsearch(BingRetriever):\n",
    "    \n",
    "    def __init__(self, cache_path):\n",
    "        super().__init__(cache_path)\n",
    "    \n",
    "    def create_query(self, question):\n",
    "        return question.text\n",
    "    \n",
    "bing_websearch = MyBingWebsearch('demo_data/bing-cache.pkl')    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "622ec208",
   "metadata": {},
   "source": [
    "Here is the information that we find via Bing for our example question about Martin Kulha. Note that the top results already contain the correct answer 'Poprad' in the text!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b2cc462e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "https://www.wikilogy.com/biography/martin-kulha/ - Martin Kulha is an Ice Hockey Player. He was born in Poprad on August 07, 1976. Want to more about Him? In this article, we covered Martin Kulha's net worth, wiki, bio, career, height, weight, pics, family, affairs, car, salary, age, facts, and other details in 2023. Continue reading to discover who is Martin Kulha. \n",
      "\n",
      "https://www.celebsagewiki.com/martin-kulha - Martin Kulha was born on 7 August, 1976 in Poprad, Slovakia. Discover Martin Kulha's Biography, Age, Height, Physical Stats, Dating/Affairs, Family and career updates. Learn How rich is She in this year and how She spends money? Also learn how She earned most of networth at the age of 44 years old? \n",
      "\n",
      "https://icehockey.fandom.com/wiki/Martin_Kulha - Martin Kulha (born August 7, 1976) is a Slovak professional ice hockey player who formerly played with Sangliers Arvernes de Clermont in the FFHG Division 1. He is now a member of the Lyon Club in the French Division 3. Kulha had previously played in the Slovak Extraliga with HK Poprad, HC Slovan Bratislava and HK 36 Skalica Biographical information and career statistics from NHL.com, or ... \n",
      "\n"
     ]
    }
   ],
   "source": [
    "retrieved = bing_websearch.retrieve(example_question)\n",
    "for snippet, url in retrieved[:3]:\n",
    "    print(url, '-', snippet, '\\n')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5ae934e",
   "metadata": {},
   "source": [
    "Next, we write a new `Generator` which uses a different prompt tailored for retrievals and the retrieved text from Bing to generate answers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "48f294f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PlaceOfBirthGeneratorWithContext(Generator):\n",
    "    \n",
    "    RETRIEVAL_PROMPT = \"\\nJerry Beck (born February 9, 1955, in New York City) is an American animation historian,\" +\\\n",
    "        \" author, blogger, and video producer.Beck wrote or edited several books on classic\" +\\\n",
    "        \" American animation and classic characters.\\nJerry Beck was born in New York\\n\\n\" +\\\n",
    "        \"Ettore Maria Fizzarotti (1916–1985) was an Italian film director and screenwriter.\" +\\\n",
    "        \" Born in Naples, the son of the director Armando, he debuted as assistant director\" +\\\n",
    "        \" in the films of his father.\\nEttore Maria Fizzarotti was born in Naples\\n\"       \n",
    "    \n",
    "    def __init__(self, llm):\n",
    "        super().__init__(llm=llm, max_tokens=10)    \n",
    "    \n",
    "    def _create_prompt(self, question, params):        \n",
    "        retrieved_text = params['retrieved_context']\n",
    "        return f\"{self.RETRIEVAL_PROMPT}\\n\\n{retrieved_text}\\n\\n{question.text}\"               \n",
    "\n",
    "    def _extract_answer(self, response):\n",
    "        answer = response.get_response()          \n",
    "\n",
    "        answer = re.sub(r'[0-9]+', '', answer)\n",
    "        answer = answer.strip()   \n",
    "\n",
    "        for sep in ['\\n', ',', '.']:\n",
    "            if sep in answer:\n",
    "                answer = answer.split(sep)[0]\n",
    "\n",
    "        return answer.strip()  \n",
    "    \n",
    "gpt35_ctx = PlaceOfBirthGeneratorWithContext(llm=gpt35_client)   \n",
    "redpajama_ctx = PlaceOfBirthGeneratorWithContext(llm=redpajama_client)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a47ee1d",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "If we provide GPT3.5 with the retrieved extra information from Bing, it generates the correct answer in the majority of cases:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "aa80429a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPT3.5 gives the answer \"Poprad on August\" based on https://www.wikilogy.com/biography/martin-kulha/\n",
      "GPT3.5 gives the answer \"Poprad\" based on https://www.celebsagewiki.com/martin-kulha\n",
      "GPT3.5 gives the answer \"Poprad\" based on https://icehockey.fandom.com/wiki/Martin_Kulha\n",
      "GPT3.5 gives the answer \"Poprad\" based on https://biogossipy.com/martin-kulha/\n",
      "GPT3.5 gives the answer \"Poprad\" based on https://popularbio.com/martin-kulha/\n",
      "GPT3.5 gives the answer \"Poprad\" based on https://www.hockeydb.com/ihdb/stats/pdisplay.php?pid=57405\n",
      "GPT3.5 gives the answer \"Pohoří\" based on https://www.myheritage.com/names/martin_kulha\n",
      "GPT3.5 gives the answer \"Slovakia\" based on http://www.vipfaq.com/Martin%20Kulha.html\n",
      "GPT3.5 gives the answer \"Poprad\" based on https://networthmask.com/martin-kulha/\n",
      "GPT3.5 gives the answer \"August th\" based on https://en.wikipedia.org/wiki/Martin_Kulha\n"
     ]
    }
   ],
   "source": [
    "for snippet, url in retrieved[:10]:\n",
    "    answer = gpt35_ctx.generate(example_question, {'retrieved_context': snippet})\n",
    "    print(f'GPT3.5 gives the answer \"{answer}\" based on {url}')    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52e0d729",
   "metadata": {},
   "source": [
    "In order to leverage this finding, we implement a `RetrievalAugmentedModel`, which generates the final answer via a majority vote over the top-10 generated answers from GPT3.5 based on the data from Bing.\n",
    "\n",
    "This model gives us the correct answer:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "6485a015",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPT3.5 with retrieval augmentation gives the correct answer \"Poprad\"\n"
     ]
    }
   ],
   "source": [
    "rag = RetrievalAugmentedModel(bing_websearch, gpt35_ctx, k=10)\n",
    "\n",
    "print(f'GPT3.5 with retrieval augmentation gives the correct answer \"{rag.generate(example_question)}\"')   "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "021d787c",
   "metadata": {},
   "source": [
    "Retrieval augmentation is a powerful technique, even a single retrieved webpage (`k=1`) improves the accuracy of our \n",
    "LLMs by a factor of 3 to 4:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "265da01c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The accuracy of GPT3.5 with retrieval augmentation and k=1 on the test set is 0.336\n",
      "The accuracy of RedPajama3B with retrieval augmentation and k=1 on the test set is 0.41\n"
     ]
    }
   ],
   "source": [
    "gpt35_rag1 = RetrievalAugmentedModel(bing_websearch, gpt35_ctx, k=1)\n",
    "redpajama_rag1 = RetrievalAugmentedModel(bing_websearch, redpajama_ctx, k=1)\n",
    "\n",
    "accuracy_gpt35_rag1 = score(test_questions, gpt35_rag1)\n",
    "accuracy_redpajama_rag1 = score(test_questions, redpajama_rag1)\n",
    "\n",
    "print(f'The accuracy of GPT3.5 with retrieval augmentation and k=1 on the test set is {accuracy_gpt35_rag1}\\n'+\\\n",
    "f'The accuracy of RedPajama3B with retrieval augmentation and k=1 on the test set is {accuracy_redpajama_rag1}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a67ea6cd",
   "metadata": {},
   "source": [
    "Using `k=10` further improves the performance and makes the small 6B model on par with 175B parameter model from OpenAI. Both models now answer about half of the test questions correctly!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "a63d8f1b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The accuracy of GPT3.5 with retrieval augmentation and k=10 on the test set is 0.498\n",
      "The accuracy of RedPajama3B with retrieval augmentation and k=10 on the test set is 0.496\n"
     ]
    }
   ],
   "source": [
    "gpt35_rag10 = RetrievalAugmentedModel(bing_websearch, gpt35_ctx, k=10)\n",
    "redpajama_rag10 = RetrievalAugmentedModel(bing_websearch, redpajama_ctx, k=10)\n",
    "\n",
    "accuracy_gpt35_rag10 = score(test_questions, gpt35_rag10)\n",
    "accuracy_redpajama_rag10 = score(test_questions, redpajama_rag10)\n",
    "\n",
    "print(f'The accuracy of GPT3.5 with retrieval augmentation and k=10 on the test set is {accuracy_gpt35_rag10}\\n'+\\\n",
    "f'The accuracy of RedPajama3B with retrieval augmentation and k=10 on the test set is {accuracy_redpajama_rag10}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34d61c99",
   "metadata": {},
   "source": [
    "## Improving the performance further with RAGBooster\n",
    "\n",
    "We can further improve the performance of our retrieval-augmented models by learning the data importance of the retrieval sources (web domains in our case) and pruning the retrieval corpus accordingly. Checkout our recent paper on **Improving Retrieval-Augmented Large Language Models with Data-Centric Refinement** (TODO need arxiv version) for details on the algorithm behind this.\n",
    "\n",
    "We can \"boost\" the performance of our models via the `RAGBooster` class and an additional set of validation questions as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "bd6279bd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "gpt35_rag_boosted = RAGBooster(gpt35_rag10, validation_questions, \n",
    "                               learning_rate=10, num_epochs=100, n_jobs=-1)\n",
    "\n",
    "redpajama_rag_boosted = RAGBooster(redpajama_rag10, validation_questions, \n",
    "                                   learning_rate=10, num_epochs=100, n_jobs=-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3233210f",
   "metadata": {},
   "source": [
    "We find that **RAGBooster improves the accuracy of both our LLMs by approximately 3%** and makes them both answer about 53% percent of the questions correctly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "21cd1336",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "'RAGBooster boosted the accuracy of GPT3.5 with retrieval augmentation from 0.498 to 0.532!'"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "accuracy_gpt35_rag10_boosted = score(test_questions, gpt35_rag_boosted)\n",
    "\n",
    "f'RAGBooster boosted the accuracy of GPT3.5 with retrieval augmentation'+\\\n",
    "f' from {accuracy_gpt35_rag10} to {accuracy_gpt35_rag10_boosted}!'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "310a9e96",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "'RAGBooster boosted the accuracy of RedPajama3B with retrieval augmentation from 0.496 to 0.528!'"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "accuracy_redpajama_rag10_boosted = score(test_questions, redpajama_rag_boosted)\n",
    "\n",
    "f'RAGBooster boosted the accuracy of RedPajama3B with retrieval augmentation'+\\\n",
    "f' from {accuracy_redpajama_rag10} to {accuracy_redpajama_rag10_boosted}!'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80c40769",
   "metadata": {},
   "source": [
    "### Important retrieval sources\n",
    "\n",
    "Internally, RAGBooster learns an importance weight for each data source (web domain in our case) as well as a pruning threshold. We can inspect these importances via the `weights` attribute of RAGBooster:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "c5025331",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('allengelhard.com', 0.7000000000000002),\n",
       " ('ancient-origins.net', 0.7000000000000002),\n",
       " ('badmintonbites.com', 0.7000000000000002),\n",
       " ('britishmuseum.org', 0.7000000000000002),\n",
       " ('elanka.com.au', 0.7000000000000002),\n",
       " ('jewage.org', 0.7000000000000002),\n",
       " ('jukebugs.com', 0.7000000000000002),\n",
       " ('masterbond.com', 0.7000000000000002),\n",
       " ('mormonwiki.com', 0.7000000000000002),\n",
       " ('nwasianweekly.com', 0.7000000000000002),\n",
       " ('playmakerstats.com', 0.7000000000000002),\n",
       " ('thedailygardener.org', 0.7000000000000002),\n",
       " ('thestar.co.uk', 0.7000000000000002),\n",
       " ('mathrubhumi.com', 0.6997466764155842),\n",
       " ('tribuneindia.com', 0.6997466764155842),\n",
       " ('ceeol.com', 0.6997357590579518),\n",
       " ('cartoonia.ru', 0.6997044999828725),\n",
       " ('ww2db.com', 0.6997010228277173),\n",
       " ('apumone.com', 0.6996936785098945),\n",
       " ('aussiecelebs.com.au', 0.6996936785098945),\n",
       " ('namecensus.com', 0.6996899787712634),\n",
       " ('getsol.app', 0.6996804444452616),\n",
       " ('ed.ac.uk', 0.6996703510006601),\n",
       " ('lindahall.org', 0.6996703510006601),\n",
       " ('raynatours.com', 0.6996655370417912)]"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "domains_and_weights = redpajama_rag_boosted.weights\n",
    "domains_and_weights_sorted = sorted(domains_and_weights.items(), key=lambda x:x[1], reverse=True)\n",
    "\n",
    "domains_and_weights_sorted[:25]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fe1df9e",
   "metadata": {},
   "source": [
    "We find that RAGBooster identifies some very interesting data sources, for example:\n",
    "\n",
    " * [ancient-origins.net](ancient-origins.net), a website dedicated to archaeology and ancient history\n",
    " * [mormonwiki.com](mormonwiki.com), an encyclopedia about mormons\n",
    " * [britishmuseum.org](https://www.britishmuseum.org), the website of the british museum\n",
    " * [badmintonbites.com](https://badmintonbites.com), a website about important badminton players"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1494718",
   "metadata": {},
   "source": [
    "We can also compute how much RAGBooster prunes the retrieval corpus. It turns out, it uses only about 20% of the domains it saw for the validation corpus."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "34a9e57f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Pruning retrieval corpus from 2374 to 475 sources, based on learned weight threshold of 0.5476'"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "num_data_sources = len(redpajama_rag_boosted.weights)\n",
    "threshold = redpajama_rag_boosted.tuning_result.best_threshold\n",
    "\n",
    "after_pruning = [domain for domain, weight in redpajama_rag_boosted.weights.items() if weight >= threshold]\n",
    "\n",
    "num_data_sources_after_pruning = len(after_pruning)\n",
    "\n",
    "f'Pruning retrieval corpus from {num_data_sources} to {num_data_sources_after_pruning} sources, '+\\\n",
    "f'based on learned weight threshold of {threshold:.4f}'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae5fb731",
   "metadata": {},
   "source": [
    "#### Disclaimer\n",
    "\n",
    "_We would like to note that it is not always the case that small LLMs performs so well, and that retrieval augmentation as we do it results in increased inference cost. Check out our paper at https://arxiv.org/abs/2307.03027 for details._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57d0751e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}