{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cf85bea0-68bf-4405-96ec-37579b2e9587",
   "metadata": {},
   "source": [
    "# Homework and bakeoff: Few-shot OpenQA with DSPy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a28e9bf5-7956-4c63-9129-7f2cbc468075",
   "metadata": {},
   "outputs": [],
   "source": [
    "__author__ = \"Christopher Potts and Omar Khattab\"\n",
    "__version__ = \"CS224u, Stanford, Fall 2024\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7b5d964-a45c-496a-bb46-8f31d7b2d591",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/cgpotts/cs224u/blob/master/hw_openqa.ipynb)\n",
    "\n",
    "If Colab is opened with this badge, please **save a copy to drive** (from the File menu) before running the notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8570fc5-2ac0-4c0e-b350-71990937ebd8",
   "metadata": {},
   "source": [
    "## Overview"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4da2d82-8c54-4d41-a59d-891f83f85f6e",
   "metadata": {},
   "source": [
    "The goal of this homework is to explore retrieval-augmented in-context learning. This is an exciting area that brings together a number of recent task ideas and modeling innovations. We will use the [DSPy programming library](http://dspy.ai) to build systems in this new mode.\n",
    "\n",
    "Our core task is __open-domain question answering (OpenQA)__. In this task, all that is given by the dataset is a question text, and the task is to answer that question. By contrast, in many modern QA tasks, the dataset provides a text and a gold passage, usually with a firm guarantee that the answer will be a substring of the passage.\n",
    "\n",
    "OpenQA is substantially harder than standard QA. The usual strategy is to use a _retriever_ to find passages in a large collection of texts and train a _reader_ to find answers in those passages. This means we have no guarantee that the retrieved passage will contain the answer we need. If we don't retrieve a passage containing the answer, our reader has no hope of succeeding. Although this is challenging, it is much more realistic and widely applicable than standard QA. After all, with the right retriever, an OpenQA system could be deployed over the entire Web.\n",
    "\n",
    "The task posed by this homework is harder even than OpenQA. We are calling this task __few-shot OpenQA__. The defining feature of this task is that the reader is simply a frozen, general purpose language model. It accepts string inputs (prompts) and produces text in response. It is not trained to answer questions per se, and nothing about its structure ensures that it will respond with a substring of the prompt corresponding to anything like an answer.\n",
    "\n",
    "__Few-shot QA__ (but not OpenQA!) is explored in the famous GPT-3 paper ([Brown et al. 2020](https://arxiv.org/abs/2005.14165)). The authors are able to get traction on the problem using GPT-3, an incredible finding. Our task here – __few-shot OpenQA__ – pushes this even further by retrieving passages to use in the prompt rather than assuming that the gold passage can be used in the prompt. If we can make this work, then it should be a major step towards flexibly and easily deploying QA technologies in new domains.\n",
    "\n",
    "In summary:\n",
    "\n",
    "| Task             | Passage given | Task-specific reader training |Task-specific retriever training  | \n",
    "|-----------------:|:-------------:|:-----------------------------:|:--------------------------------:|\n",
    "| QA               | yes           | yes                           | n/a                              |\n",
    "| OpenQA           | no            | yes                           | maybe                            |\n",
    "| Few-shot QA      | yes           | no                            | n/a                              |\n",
    "| Few-shot OpenQA  | no            | no                            | maybe                            | \n",
    "\n",
    "Just to repeat: your mission is to explore the final line in this table. The core notebook and assignment don't address the issue of training the retriever in a task-specific way, but this is something you could pursue for a final project; [the ColBERT codebase](https://github.com/stanford-futuredata/ColBERT) makes easy.\n",
    "\n",
    "It is a requirement of the bake-off that a general-purpose language model be used. In particular, trained QA systems cannot be used at all, and no fine-tuning is allowed either. See the original system question at the bottom of this message for guidance on which models are allowed.\n",
    "\n",
    "Note: the models we are working with here are _big_. This poses a challenge that is increasingly common in NLP: you have to pay one way or another. You can pay to use the GPT-3 API, or you can pay to use a local model on a heavy-duty cluster computer, or you can pay with time by using a local model on a more modest computer."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cd32bb4-067f-4cd6-943f-3e5574400beb",
   "metadata": {},
   "source": [
    "## Set-up"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "149bcb0f-bc76-4277-a359-742d6dcee063",
   "metadata": {},
   "source": [
    "We have sought to make this notebook self-contained and easy to use on a personal computer, on Google Colab, and in Sagemaker Studio. For personal computer use, we assume you have already done everything in [setup.ipynb](setup.ipynb]). For cloud usage, the next few code blocks should handle all set-up steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62b983bb-a20a-4c2a-9eee-9c553dd8c070",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    # This library is our indicator that the required installs\n",
    "    # need to be done.\n",
    "    import datasets\n",
    "    root_path = '.'\n",
    "except ModuleNotFoundError:\n",
    "    !git clone https://github.com/cgpotts/cs224u/\n",
    "    !pip install -r cs224u/requirements.txt\n",
    "    root_path = 'dspy'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1afb7791-92a3-474d-8c2f-807dc2441413",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "import os\n",
    "import dspy\n",
    "import warnings\n",
    "\n",
    "from dotenv import load_dotenv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ad2651d-9868-45bf-81df-afafa832896a",
   "metadata": {},
   "source": [
    "### OpenAI set-up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a04cb488-cd40-4f9d-b884-8ff83b012042",
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9da704b-d27b-480a-93b5-e16cf7c51803",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"DSP_NOTEBOOK_CACHEDIR\"] = os.path.join(root_path, 'cache')\n",
    "\n",
    "\n",
    "# keep the API keys in a `.env` file in the local root directory\n",
    "load_dotenv()\n",
    "\n",
    "openai_key = os.getenv('OPENAI_API_KEY')  # use the .env file as it is a good practice to keep keys outside of one's code"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4465db5-4572-457e-b053-10f0185536fe",
   "metadata": {},
   "source": [
    "### ColBERT set-up"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63747138-ae9c-4fc6-b508-0ad411fd58ac",
   "metadata": {},
   "source": [
    "#### Pretrained ColBERT parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "caf06f38-3dd6-4171-8b72-f5703df140d8",
   "metadata": {},
   "source": [
    "For this set-up phase, we first download pretrained ColBERT parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4925831c-f4d2-4b6b-87ed-145125c8f668",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(os.path.join(\"data\", \"openqa\", \"colbertv2.0.tar.gz\")):\n",
    "    !mkdir -p data/openqa\n",
    "    # ColBERTv2 checkpoint trained on MS MARCO Passage Ranking (388MB compressed)\n",
    "    !wget https://downloads.cs.stanford.edu/nlp/data/colbert/colbertv2/colbertv2.0.tar.gz -P data/openqa/\n",
    "    !tar -xvzf data/openqa/colbertv2.0.tar.gz -C data/openqa/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a89bd6c-5497-4808-b2cf-9e15d9c37104",
   "metadata": {},
   "source": [
    "If something went wrong with the above, you can just download the file https://downloads.cs.stanford.edu/nlp/data/colbert/colbertv2/colbertv2.0.tar.gz, unarchive it, and move the resulting `colbertv2.0` directory into the `data/openqa` directory."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bf501ee-290b-4007-b04d-c175c3124806",
   "metadata": {},
   "source": [
    "#### Prebuilt ColBERT index"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e330d4c-3d9f-4182-9c40-37010a5e5afd",
   "metadata": {},
   "source": [
    "Then we download a prebuilt index that covers the domain of our task:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d4bc73b-6936-44e2-bc80-7aca800abd75",
   "metadata": {},
   "outputs": [],
   "source": [
    "index_home = os.path.join(\"experiments\", \"notebook\", \"indexes\", \"cs224u.collection.2bits\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11d8dd51-d0d9-4d72-ae77-f00f6151889d",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(index_home):\n",
    "    !wget https://web.stanford.edu/class/cs224u/data/cs224u.collection.2bits.tgz -P experiments/notebook/indexes\n",
    "    !tar -xvzf experiments/notebook/indexes/cs224u.collection.2bits.tgz -C experiments/notebook/indexes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a2443e3-efe6-444d-98bd-63fa2634d8ff",
   "metadata": {},
   "source": [
    "If something went wrong with the above, download the file https://web.stanford.edu/class/cs224u/data/cs224u.collection.2bits.tgz, unarchive it, and move the resulting `cs224u.collection.2bits` directory into the `experiments/notebook/indexes` directory (which you will probably need to create)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bd3e9ff-c3d4-4a02-acfa-0d9183bce133",
   "metadata": {},
   "source": [
    "#### ColBERT server"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90fde4ad-a0e2-49a2-b3ec-446c5cb95a95",
   "metadata": {},
   "source": [
    "The final step for ColBERT is to create a local retriever for use with DSPy. The preferred method for doing this is with a seperate server. \n",
    "\n",
    "__Local usage__: If you are working with this notebook locally and a having a CUDA based device, make sure you have the [CUDA Toolkit](https://developer.nvidia.com/cuda-downloads) installed to be able to serve it on GPU. Otherwise you would need to run it on CPU. The next steps are for GPU usage. \n",
    "Open a new terminal window, navigate to the same directory as this notebook, and enter the following code:\n",
    "\n",
    "```\n",
    "conda activate nlu\n",
    "git clone https://github.com/stanford-futuredata/ColBERT/ ;\n",
    "export INDEX_ROOT=experiments/notebook/indexes/cs224u.collection.2bits/ ;\n",
    "export INDEX_HOME=cs224u.collection.2bits ;\n",
    "export PORT=8888 ;\n",
    "python ColBERT/server.py \n",
    "```\n",
    "\n",
    "This will start the server and you should be all set.\n",
    "\n",
    "__Colab usage__: If you are working in a Colab _and you have a Pro account_, you can open a terminal from the lower left corner of the interface and then paste in the above code. If you don't have a Pro account, then it should work to uncomment the code in the following cell and run it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d2542a5-610f-4e77-9f00-ec86b18e1d91",
   "metadata": {},
   "outputs": [],
   "source": [
    "#os.environ[\"INDEX_ROOT\"] = \"experiments/notebook/indexes/cs224u.collection.2bits/\"\n",
    "#os.environ[\"INDEX_HOME\"] = \"cs224u.collection.2bits\"\n",
    "#os.environ[\"PORT\"] = \"8888\"\n",
    "\n",
    "#!git clone https://github.com/stanford-futuredata/ColBERT/\n",
    "\n",
    "#!nohup python ColBERT/server.py &"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea56ea89-7055-401f-8afc-8e7687da284b",
   "metadata": {},
   "source": [
    "This will take a minute to get started, and you won't get feedback on its progress, unfortunately. You can check that a server is running:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d07cae52-81d9-4ff4-9a3b-cf79cf2d74a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "!ps aux | grep server.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5988243e-6985-4ccd-a40e-9a2475566ed3",
   "metadata": {},
   "source": [
    "You should see a mention of `ColBERT/server.py` in the output. \n",
    "\n",
    "Assuming the server is now running, we can connect to it:\n",
    "\n",
    "*Note*: if the ColBERT server is running on a separate server, you can update `127.0.0.1` with the server's IP address."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "556f268f-1a71-476d-bb13-8093feadaf6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "rm = dspy.ColBERTv2(url=\"http://127.0.0.1:8888/api/search\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d9bbc99-fcb2-4b22-b09b-c35371611709",
   "metadata": {},
   "source": [
    "### DSPy set-up"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "562fc500-bbab-48b7-a704-67c6d57bb09b",
   "metadata": {},
   "source": [
    "Here we establish the Language Model `lm` and Retriever Model `rm` that we will be using. The defaults for `lm` are just for development. You may want to develop using an inexpensive model and then do your final evalautions wih an expensive one. DSPy has support for a wide range of model APIs and local models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c118b014-e13f-433d-ad60-074636c7e738",
   "metadata": {},
   "outputs": [],
   "source": [
    "lm = dspy.OpenAI(model='gpt-3.5-turbo', api_key=openai_key)\n",
    "\n",
    "dspy.settings.configure(lm=lm, rm=rm)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "711785d7-6bb9-4041-92e6-cc5f9308477e",
   "metadata": {},
   "source": [
    "Here's a command you can run to see which OpenAI models are available; OpenAI has entered into an increasingly closed mode where many older models are not available, so there are likely to be some surprises lurking here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a859fbb-e985-4031-b8ed-34f3b034db8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# client = OpenAI(api_key = openai_key)\n",
    "# models = client.models.list()\n",
    "# print(models)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f0b3dc2-87d7-4b8b-b603-ee567e008710",
   "metadata": {},
   "source": [
    "## SQuAD"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de295d35-fea5-46d2-9a01-022ad88e54cd",
   "metadata": {},
   "source": [
    "Our core development dataset is [SQuAD](https://rajpurkar.github.io/SQuAD-explorer/). We chose this dataset because it is well-known and widely used, and it is large enough to support lots of meaningful development work, without, though, being so large as to require lots of compute power."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5eaf2fd0-d060-4100-8702-f7311efd6129",
   "metadata": {},
   "outputs": [],
   "source": [
    "squad = load_dataset(\"squad\", trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36965402-e3da-4531-b7e9-4b12cebcdf30",
   "metadata": {},
   "source": [
    "The following utility just reads a SQuAD split in as a list of `dspy.Example` instances:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21ad3e0b-7662-43b8-9409-a1a57442458b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_squad_split(squad, split=\"validation\"):\n",
    "    \"\"\"\n",
    "    Use `split='train'` for the train split.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    list of dspy.Example with attributes question, answer\n",
    "\n",
    "    \"\"\"\n",
    "    data = zip(*[squad[split][field] for field in squad[split].features])\n",
    "    exs = [dspy.Example(question=q, answer=a['text'][0]).with_inputs(\"question\")\n",
    "           for eid, title, context, q, a in data]\n",
    "    return exs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3847d38-4e70-46b7-bf46-4c8b784c5ee5",
   "metadata": {},
   "source": [
    "### SQuAD train"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "051c91e1-586b-4747-be39-3092e60f182f",
   "metadata": {},
   "source": [
    "To build few-shot prompts, we will often sample SQuAD train examples, so we load that split here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66c4feba-d580-4984-a449-0b92a53ef13a",
   "metadata": {},
   "outputs": [],
   "source": [
    "squad_train = get_squad_split(squad, split=\"train\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22ab8c41-8eae-4d15-ad4d-e28b3c58eb4a",
   "metadata": {},
   "source": [
    "### SQuAD dev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37198b33-c47b-4e0e-af8b-c00860658cc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "squad_dev = get_squad_split(squad)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c40c768-57cc-4a07-a3ef-5e34262b0ace",
   "metadata": {},
   "source": [
    "### SQuAD dev sample"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "636601b5-c7ad-4177-a6d6-f3afdb0bedae",
   "metadata": {},
   "source": [
    "Evaluations are expensive in this new era! Here's a small sample to use for dev assessments:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64ccdd0e-de78-440d-bfbe-358c12eada5c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "random.seed(1)\n",
    "\n",
    "dev_exs = random.sample(squad_dev, k=200)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28265d01-890d-4f04-b518-da0e8a1cb235",
   "metadata": {},
   "source": [
    "## DSPy basics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "101b9d12-fc8c-4aae-b09e-0d72a4aa54f5",
   "metadata": {},
   "source": [
    "### LM usage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c278daac-11f9-4327-a06f-1c408a06a71d",
   "metadata": {},
   "source": [
    "Here's the most basic way to use the LM:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02364ed6-3c6d-4eaf-849a-2d9e30d84b53",
   "metadata": {},
   "outputs": [],
   "source": [
    "lm(\"What is the birthplace of the first author to win a Hugo Award for a translation?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2356d9a8-750b-4383-bc5b-4173ca5c13ac",
   "metadata": {},
   "source": [
    "Keyword arguments to the underlying LM are passed through:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19ab8deb-3b9d-4f57-bd70-7c170be294c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "lm(\"Which U.S. states border no U.S. states?\", temperature=0.9, n=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c50e8d99-6d49-420d-ab5d-cc01b53cd4a1",
   "metadata": {},
   "source": [
    "With `lm.inspect_history`, we can see the most recent language model calls:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f7cb5a5-3a3f-4e78-b9af-488fadc896ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = lm.inspect_history(n=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9a0f2df-df1f-422d-bff3-6c4a3d947f6e",
   "metadata": {},
   "source": [
    "### Signature-based prediction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0888a2a-fcaa-44b0-beef-b097c856c74b",
   "metadata": {},
   "source": [
    "In DSPy, __signatures__ are declarative statements about what we want the model to do. In the following `\"question -> answer\"` is the signature (the most basic QA signature one could write), and `dspy.Predict` is used to turn this into a complete QA system: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11f04cba-7d46-4680-a154-a0062243618e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "basic_predictor = dspy.Predict(\"question -> answer\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a994fc10-8c28-4836-bd7c-008a9a3d34e4",
   "metadata": {},
   "source": [
    "Here we use `basic_predictor`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfe102d0-c0e5-45ad-aae8-51f6ea6e70f6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "basic_predictor(question=\"What is the birthplace of the first author to win a Hugo Award for a translation?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97ad1a1b-d422-46df-8b03-4054bfec5fc1",
   "metadata": {},
   "source": [
    "And here is the prompt that was given to the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af4f86db-97d5-4742-8468-4627de12f181",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "_ = lm.inspect_history(n=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7845f62b-8af3-4fa2-af8d-f4ddd1181f30",
   "metadata": {},
   "source": [
    "In many cases, we will want more control over the prompt. Writing a small custom `dspy.Signature` class is the easiest way to accomplish this. In the following, we just just tweak the initial instruction and provide some formatting guidance for the answer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9dda770-e7ca-4682-b283-d4f4525adb68",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class BasicQASignature(dspy.Signature):\n",
    "    __doc__ = \"\"\"Answer questions with short factoid answers.\"\"\"\n",
    "\n",
    "    question = dspy.InputField()\n",
    "    answer = dspy.OutputField(desc=\"often between 1 and 5 words\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83f7297a-b174-4566-8748-6b66d515ed25",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sig_predictor = dspy.Predict(BasicQASignature)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30a02d55-fac7-43c3-851b-3e21f06df14d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sig_predictor(question=\"Which U.S. states border no U.S. states?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f613c05-ec78-4129-ac44-71bac4e253ad",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "_ = lm.inspect_history(n=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56169f83-0df6-46dc-b617-fadff1b96132",
   "metadata": {},
   "source": [
    "### Modules"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e244aed0-403e-4707-a9eb-e29fc1e4dc57",
   "metadata": {},
   "source": [
    "One of the hallmarks of DSPy is that it adopts design patterns from PyTorch. The main example of this is DSPy's use of the `Module` as the basic unit for writing simple and complex programs. Here is a very basic module for QA that makes use of `BasicQASignature` as we defined it just above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "506c9946-fee4-4dc3-a05e-767f85c25292",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class BasicQA(dspy.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.generate_answer = dspy.Predict(BasicQASignature)\n",
    "\n",
    "    def forward(self, question):\n",
    "        return self.generate_answer(question=question)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c921aaed-6351-443d-8478-6e9a64b49d45",
   "metadata": {},
   "source": [
    "As with PyTorch, the `forward` method is called when we want to make predictions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d67f65b7-0fea-40a3-ac78-795926025653",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "basic_qa_model = BasicQA()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7044ec84-f76a-4af4-93f5-a82579237868",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "basic_qa_model(question=\"What is the birthplace of the first author to win a Hugo Award for a translation?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9961c89-914a-4736-8a01-1a3cd401821a",
   "metadata": {
    "tags": []
   },
   "source": [
    "The modular design of DSPy starts to become apparent now. If you want to change the above to use chain of thought instead of regular predictions, you need only change `dspy.Predict` to `dspy.ChainOfThought`, and similarly for `dspy.ReAct`, `dspy.ProgramOfThought`, or a module you wrote yourself."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17b9dac4-9495-448d-ad4b-38ab7260915b",
   "metadata": {},
   "source": [
    "### Optimizing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5200696-808f-4061-878c-26aa20701d58",
   "metadata": {},
   "source": [
    "The QA system we've defined so far is a zero-shot system. To change it into a few-shot system, we will rely on a DSPy optimizer (__teleprompter__). This will allow us to flexibly move between the zero-shot and few-shot formulations. The following code achieves this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdab9850-1d4d-4e04-b389-6ef55fd0b425",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from dspy.teleprompt import LabeledFewShot"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72045705-2f27-46fe-b7bc-801a7d2e9ae4",
   "metadata": {},
   "source": [
    "Here we instantiate a `LabeledFewShot` teleprompter that will add three demonstrations. These will be sampled randomly from the set of train examples we provide:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "befc8c18-fc64-4947-800d-6455dee90b68",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fewshot_teleprompter = LabeledFewShot(k=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25b512d9-ec54-49ca-ab7a-900dc2eea6cc",
   "metadata": {},
   "source": [
    "And then we call `compile` on `basic_qa_model` as we defined it above. This returns a new module that we use like any other in DSPy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b327b603-2943-4bcc-b498-1fbafafefd0c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "basic_fewshot_qa_model = fewshot_teleprompter.compile(basic_qa_model, trainset=squad_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "921ec7f1-ed44-48a3-9c92-0e22911688f0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "basic_fewshot_qa_model(question=\"What is the birthplace of the first author to win a Hugo Award for a translation?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3180cdaa-dd34-4e0d-8455-e394bfde9fcc",
   "metadata": {},
   "source": [
    "With `inspect_history`, we can see that prompts now contain demonstrations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9507fe6-1cc7-4c0a-8520-cc3812a07c27",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "_ = lm.inspect_history(n=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ba48440-4a65-41e8-b397-7b79f65fa0fe",
   "metadata": {},
   "source": [
    "### Evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27e8734a-49a1-4093-a2fb-09bb7d2f2859",
   "metadata": {},
   "source": [
    "Our evaluation metric is a standard one for SQuAD and related tasks: exact match of the answer (EM)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f052a79f-ad9f-4f3e-a195-809567e2eea1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from dspy.evaluate import answer_exact_match\n",
    "from dspy.evaluate.evaluate import Evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7341a846-5158-4acf-8aa5-aca61ef40174",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "answer_exact_match(dspy.Example(answer=\"STAGE 2!\"), dspy.Prediction(answer=\"stage 2\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f640e5bf-823f-4cdb-92da-a28f5cea7760",
   "metadata": {},
   "source": [
    "In DSPy, `Evaluate` objects provide a uniform interface for running evaluations. Here are two for us to use in development. The first will evaluate on all of `dev_exs` and should provide a meaningful picture of how a system is doing. It could be expensive to use it a lot, though. The second is for debugging and is probably too small to give a reliable estimate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6a3b58b-c114-4316-b8d7-2d8049132c1d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dev_evaluater = Evaluate(\n",
    "    devset=dev_exs, # 200 examples\n",
    "    num_threads=1,\n",
    "    display_progress=True,\n",
    "    display_table=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72fd504e-4684-445d-b0cc-363cd1685b73",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "tiny_evaluater = Evaluate(\n",
    "    devset=dev_exs[: 15],\n",
    "    num_threads=1,\n",
    "    display_progress=True,\n",
    "    display_table=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47d88a6c-55b6-4994-9413-1a2301c94903",
   "metadata": {},
   "source": [
    "Here is a tiny (debugging-oriented) evaluation of our few-shot QA sytem:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f736065-39f2-4d76-a258-e852767dbb8e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "tiny_evaluater(basic_fewshot_qa_model, metric=answer_exact_match)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3f1b41d-760c-4f28-8a2d-7f037b4f9d97",
   "metadata": {},
   "source": [
    "### Retrieval"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51d4c8f4-a537-4d9b-9500-f881fceef1de",
   "metadata": {},
   "source": [
    "The final major component of our systems is retrieval. When we defined `rm`, we connected to a remote ColBERT index and retriever system that we can now use for search."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5dfd114-96cf-468a-bac3-d3d39d6f3ca6",
   "metadata": {},
   "source": [
    "The basic `dspy.retrieve` method returns only passages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8f4cafc-f5db-41c0-9561-ecde61331666",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "retriever = dspy.Retrieve(k=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "891fc391-c177-4da7-9332-ab20cdba3c0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "passages = retriever(\"What is the birthplace of the first author to win a Hugo Award for a translation?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abdac37b-b5fe-421c-826f-4fd699cb2e36",
   "metadata": {},
   "outputs": [],
   "source": [
    "passages.passages[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4e1f577-408c-4ede-9e27-65a24aafca5f",
   "metadata": {},
   "source": [
    "If we need passages with scores and other metadata, we can call `rm` directly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37f4ff3d-de41-4fc7-8943-6dfd894dd1f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "rm(\"What is the birthplace of the first author to win a Hugo Award for a translation?\", k=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db2017ee-1375-4251-a24f-7f792852ffac",
   "metadata": {},
   "source": [
    "## Few-shot OpenQA with context"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f05c0997-e624-4fc0-824d-e13066978b0a",
   "metadata": {},
   "source": [
    "Let's build on the above core concepts to define a basic retrieval-augmented generation (RAG) program. This program solves the core task of few-shot OpenQA task and will serve as the basis for the homework questions:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1714dd71-02aa-4e84-840f-807b4e501732",
   "metadata": {},
   "source": [
    "We begin with a signature that takes context into account but is otherwise just like `BasicQASignature` above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcd35bc4-9bbc-4286-ad6b-d7474b51423e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class ContextQASignature(dspy.Signature):\n",
    "    __doc__ = \"\"\"Answer questions with short factoid answers.\"\"\"\n",
    "\n",
    "    context = dspy.InputField(desc=\"may contain relevant facts\")\n",
    "    question = dspy.InputField()\n",
    "    answer = dspy.OutputField(desc=\"often between 1 and 5 words\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3de481d-d4dd-42c3-99c1-7a112c7f521f",
   "metadata": {},
   "source": [
    "And here is a complete program/system for the task:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f84481c9-6e6a-4331-a5e1-cf2b9e27b082",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class RAG(dspy.Module):\n",
    "    def __init__(self, num_passages=1):\n",
    "        super().__init__()\n",
    "        self.retrieve = dspy.Retrieve(k=num_passages)\n",
    "        self.generate_answer = dspy.Predict(ContextQASignature)\n",
    "\n",
    "    def forward(self, question):\n",
    "        context = self.retrieve(question).passages\n",
    "        prediction = self.generate_answer(context=context, question=question)\n",
    "        return dspy.Prediction(context=context, answer=prediction.answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c4c1ddd-455e-4ad5-b1b1-bf1267364660",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "rag_model = RAG(num_passages=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fce28fa-aa8a-4cec-a07a-5365a5eb32df",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "rag_model(question=\"What is the birthplace of the first author to win a Hugo Award for a translation?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "873286e4-aaa7-4359-a5f2-04f8cbcceac8",
   "metadata": {},
   "source": [
    "An optional tiny evaluation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25715550-2ba8-44bb-9a11-ca3bc3e0af56",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "tiny_evaluater(rag_model, metric=answer_exact_match)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eeb28556-d901-4b6b-8a7d-93040203294d",
   "metadata": {},
   "source": [
    "## Question 1: Optimizing RAG [2 points]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f481081-3e16-4ccd-b379-c7e0c3011286",
   "metadata": {},
   "source": [
    "We used `RAG` above as a zero-shot system. We could turn it into a few-shot system by using `LabeledFewShot` as we did in [the teleprompting section](#Teleprompting) above, but this may actually be problematic: if we randomly sample demonstrations with retrieved passages, we might be instructing the model with a lot of cases where the context passage isn't helping (and may actually be actively misleading the model). \n",
    "\n",
    "What we'd like to do is select demonstrations where the model gets the answer correct and the context passage does contain the answer. To do this, we will use the DSPy `BootstrapFewShot` optimizer. There are two steps for this: (1) defining a metric and (2) running the optimizer.\n",
    "\n",
    "__Note__: The code for this question can be found in the DSPy tutorials, and you should feel free to make use of that code. The goal is to help you understand the design patterns and overall logic of optimizing DSPy programs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff0ae027-f51a-4607-822f-2a721acde73c",
   "metadata": {},
   "source": [
    "__Task 1__: Complete `validate_context_and_answer` according to the specification in the docstring."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad09f428-1a68-4cef-9e7a-7c6faf9244a4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def validate_context_and_answer(example, pred, trace=None):\n",
    "    \"\"\"Return True if `example.answer` matches `pred.answer` according\n",
    "    to `dspy.evaluate.answer_exact_match` and `pred.context` contains\n",
    "    `example.answer` according to `dspy.evaluate.answer_passage_match`.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    example: dspy.Example \n",
    "        with attributes `answer` and `context`\n",
    "    pred: dspy.Example \n",
    "        with attributes `answer` and `context`\n",
    "    trace : None (included for dspy internal compatibility)\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    bool\n",
    "\n",
    "    \"\"\"\n",
    "    pass\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0918c5f-2864-4856-adcf-ccb94aca6108",
   "metadata": {},
   "source": [
    "A test you can use to check your implementation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddbca59c-7c11-4e4b-bec3-18a2ecde563b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def test_validate_context_and_answer(func):\n",
    "    examples = [\n",
    "        (\n",
    "            dspy.Example(question=\"Q1\", answer=\"B\"),\n",
    "            dspy.Prediction(question=\"Q1\", context=\"A B C\", answer=\"B\"),\n",
    "            True\n",
    "        ),\n",
    "        # Context doesn't contain answer, but predicted answer is correct.\n",
    "        (\n",
    "            dspy.Example(question=\"Q1\", answer=\"D\"),\n",
    "            dspy.Prediction(question=\"Q1\", context=\"A B C\", answer=\"D\"),\n",
    "            False\n",
    "        ),\n",
    "        # Context contains answer, but predicted answer is not correct.\n",
    "        (\n",
    "            dspy.Example(question=\"Q1\", answer=\"C\"),\n",
    "            dspy.Prediction(question=\"Q1\", context=\"A B C\", answer=\"D\"),\n",
    "            False\n",
    "        )\n",
    "    ]\n",
    "    errcount = 0\n",
    "    for ex, pred, result in examples:\n",
    "        predicted = func(ex, pred, trace=None)\n",
    "        if predicted != result:\n",
    "            errcount += 1\n",
    "            print(f\"Error for `{func.__name__}`: \"\n",
    "                  f\"Expected inputs\\n\\t{ex}\\n\\t{pred} to return {result}.\")\n",
    "    if errcount == 0:\n",
    "        print(f\"No errors detected for `{func.__name__}`\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38aa3141-ae8f-4bc6-a26e-64d108537612",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "test_validate_context_and_answer(validate_context_and_answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f091ca79-ccbe-4429-bfe7-c4f1de118adf",
   "metadata": {
    "tags": []
   },
   "source": [
    "__Task 2__: Complete `bootstrap_optimize` according to the specification in the docstring."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "129e7e92-033e-4f38-b309-94dac2b66d87",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from dspy.teleprompt import BootstrapFewShot\n",
    "\n",
    "def bootstrap_optimize(model):\n",
    "    \"\"\"Use `BootstrapFewShot` to optimize `model`, with the metric set\n",
    "    to `validate_context_and_answer` as defined above and default\n",
    "    values for all other keyword arguments to `BootstrapFewShot`.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    model: dspy.Module\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    dspy.Module, the optimized version of `model`\n",
    "\n",
    "    \"\"\"\n",
    "    pass\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f26af22-6d7d-4456-9e30-8f80f5b0b401",
   "metadata": {},
   "source": [
    "A test you can use to check your implementation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30dd9f56-31cc-4e45-8bc2-0fe56c3806a7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def test_bootstrap_optimize(func):\n",
    "    model = RAG()\n",
    "    compiled = func(model)\n",
    "    if not hasattr(compiled, \"_compiled\") or not compiled._compiled:\n",
    "        print(f\"Error for `{func.__name__}`: \"\n",
    "               \"The return value is not a compiled program.\")\n",
    "        return None\n",
    "    state = compiled.dump_state()\n",
    "    if not state['generate_answer']['demos']:\n",
    "        print(f\"Error for `{func.__name__}`: \"\n",
    "               \"The compiled program has no `demos`.\")\n",
    "        return None\n",
    "    print(f\"No errors detected for `{func.__name__}`\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "733c1c31-9c68-4d2a-a1c3-4eed0e1caa6f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "test_bootstrap_optimize(bootstrap_optimize)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9971ffd0-a563-4d56-a896-0735a63ae92f",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Question 2: Multi-passage summarization [2 points]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "527e4d85-05d7-4bc2-beaf-2434d4fe41da",
   "metadata": {},
   "source": [
    "The `dspy.Retrieve` layer in our `RAG` retrieves `k` passages, where `k` is under the control of the user. One hypothesis one might have is that it would be good to summarize these passages before using them as evidence. This seems especially likely to help in scenarios where the question can be answered only by synthesizing information across documents – it might be too much to ask the language model to do both synthesizing and answering in a single step.\n",
    "\n",
    "The current question maps out a basic strategy for summarization. The heart of it is a new signature called `SummarizeSignature`. This can be used on its own with a simple `dspy.Predict` call, and we'll incorporate it into a RAG program in the next question.\n",
    "\n",
    "For this question, though, your task is just to complete `SummarizeSignature`. The requirements are as follows:\n",
    "\n",
    "1. A `__doc__` value that gives an instruction that seems to work well. You can decide what to say here.\n",
    "2. A `dspy.InputField` named `context`. You can decide whether to use the `desc` parameter.\n",
    "3. A `dspy.OutputField` named `summary`. You can decide whether to use the `desc` parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7c080a6-d2e1-4d7a-ac60-5ef3a6931ad1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class SummarizeSignature(dspy.Signature):\n",
    "    pass\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6091f20-1cae-4e04-bd50-929271ae6a18",
   "metadata": {},
   "source": [
    "Here's a simple test that just checks for the required pieces in a basic way:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "428c971a-2dcd-4344-afa0-8e52938ca4ee",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def test_SummarizeSignature(sigclass):\n",
    "    fields = sigclass.fields\n",
    "    expected_fieldnames = ['context', 'summary']\n",
    "    fieldnames = sorted(fields)\n",
    "    errcount = 0\n",
    "    if expected_fieldnames != fieldnames:\n",
    "        errcount += 1\n",
    "        print(f\"Error for `{sigclass.__name__}`: \"\n",
    "              f\"Expected fieldnames {expected_fieldnames}, got {fieldnames}.\")\n",
    "    if not sigclass.__doc__:\n",
    "        errcount += 1\n",
    "        print(f\"Error for `{sigclass.__name__}`: No docstring specified.\")\n",
    "    if errcount == 0:\n",
    "        print(f\"No errors detected for `{sigclass.__name__}`\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fec39e4d-b3ca-4355-9695-c9f61a24f3d2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "test_SummarizeSignature(SummarizeSignature)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "881ee5d0-2bd6-4cb0-873c-21f963a78555",
   "metadata": {},
   "source": [
    "Here is the simplest way to use `SummarizeSignature`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e8129a3-90fb-4810-b18b-84e75525dd16",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "summarizer = dspy.Predict(SummarizeSignature)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8bfdb34-07d8-4488-a1ed-1b4788a2a119",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "summarizer(context=retriever(\"Where is Guarani spoken?\").passages)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24b7a44f-c50d-4674-9fbb-fcf216222e3a",
   "metadata": {},
   "source": [
    "## Question 3: Summarizing RAG [2 points]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c812e189-2778-442f-81b1-577c303445a8",
   "metadata": {},
   "source": [
    "Your task for this question is to modify `RAG` as defined above so that the retrieved passages are summarized before being passed to `generate_answer`. \n",
    "\n",
    "Here is the `RAG` system copied from above with the class name changed to the one we will use for this new system. Your task is to add the summarization step. This should be very straightforward given the modular design that DSPy supports and encourages!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa956847-d6ff-40cf-bf02-9043863e548e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class SummarizingRAG(dspy.Module):\n",
    "    def __init__(self, num_passages=3):\n",
    "        # Please name your summarization later `summarize` so that we\n",
    "        # can check for its presence.\n",
    "        super().__init__()\n",
    "        self.retrieve = dspy.Retrieve(k=num_passages)\n",
    "        ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "        self.generate_answer = dspy.Predict(ContextQASignature)\n",
    "\n",
    "    def forward(self, question):\n",
    "        context = self.retrieve(question).passages\n",
    "        ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "        prediction = self.generate_answer(context=context, question=question)\n",
    "        return dspy.Prediction(context=context, answer=prediction.answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f16b65d8-14ca-4dbe-a815-a0d00940b0b4",
   "metadata": {},
   "source": [
    "A simple test for this design spec:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "824f063d-1735-4bf3-9337-e9728a5b7800",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def test_SummarizingRAG(classname):\n",
    "    model = classname(num_passages=3)\n",
    "    errcount = 0\n",
    "    if not hasattr(model, \"summarize\"):\n",
    "        errcount += 1\n",
    "        print(f\"Error for `{classname.__name__}`: \"\n",
    "              f\"Expected a layer called 'summarize'\")\n",
    "    context = model.retrieve(\"What are some foods?\").passages\n",
    "    pred = model(\"What are some foods?\")\n",
    "    if context == pred.context:\n",
    "        errcount += 1\n",
    "        print(f\"Error for `{classname.__name__}`: \"\n",
    "              \"The model seems to be using raw retrieved contexts \"\n",
    "              \"for predictions rather than summarizing them.\")\n",
    "    if errcount == 0:\n",
    "        print(f\"No errors detected for `{classname.__name__}`\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e1d38c4-540c-4473-b464-0061b5b8a09c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "test_SummarizingRAG(SummarizingRAG)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9e3b6ad-b960-4b8f-8137-a9b90953b2fc",
   "metadata": {},
   "source": [
    "Model usage:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95e03778-5a7f-4119-ac62-f4965bfe9a8c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "summarizing_rag_model = SummarizingRAG()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5077370b-156f-4738-948b-86d832b2ebbd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "summarizing_rag_model(question=\"What is the birthplace of the first author to win a Hugo Award for a translation?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49ebc3e9-7baf-4e1c-81a0-8c1e3588a882",
   "metadata": {},
   "source": [
    "Note: if you decide to use `BootstrapFewShot` on this, be sure not to use the metric we defined above, which requires that the passage embeds the correct answer as a substring. Now that we are summarizing, this is unlikely to hold, even if the answers are good ones."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19655d70-007c-4a41-9f3b-e20df9c5169b",
   "metadata": {},
   "source": [
    "## Question 4: Your original system [3 points]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d8d268f-70e1-4325-a2ff-7361d73788b9",
   "metadata": {},
   "source": [
    "This question asks you to design your own few-shot OpenQA system. All of the code above can be used and modified for this, and the requirement is just that you try something new that goes beyond what we've done so far. \n",
    "\n",
    "Terms for the bake-off:\n",
    "\n",
    "* You can make free use of SQuAD and other publicly available data.\n",
    "\n",
    "* The LM must be an autoregressive language model. No trained QA components can be used. This includes general purpose LMs that have been fine-tuned for QA. (We have obviously waded into some vague territory here. The spirit of this is to make use of frozen, general-purpose models. We welcome questions about exactly how this is defined, since it could be instructive to explore this.)\n",
    "\n",
    "Here are some ideas for the original system:\n",
    "\n",
    "* We have relied almost entirely on `dspy.Predict`. Drop-in replacements include `dspy.ChainOfThought` and `dspy.ReAct`.\n",
    "\n",
    "* We have used only one retriever. DSPy supports other retrieval mechanisms, including retrieval using [You.com](https://you.com/).\n",
    "\n",
    "* DSPy includes additional optimizers. Two that are worth trying are `SignatureOptimizer` for automatic prompt exploration and `BootstrapFewShotWithRandomSearch`, which combines `LabeledFewShot` and `BootstrapFewShot`,\n",
    "\n",
    "* Our one-step summarization procedure from Question 3 doesn't change the query to the retriever. We might want it to change as we gather evidence. This is a common design principle for multi-hop OpenQA systems.\n",
    "\n",
    "__Original system instructions__:\n",
    "\n",
    "In the cell below, please provide a brief technical description of your original system, so that the teaching team can gain an understanding of what it does. This will help us to understand your code and analyze all the submissions to identify patterns and strategies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b557f3c3-ee72-480e-9d99-9095372f99c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# PLEASE MAKE SURE TO INCLUDE THE FOLLOWING BETWEEN THE START AND STOP COMMENTS:\n",
    "#   1) Textual description of your system.\n",
    "#   2) The code for your original system.\n",
    "# PLEASE MAKE SURE NOT TO DELETE OR EDIT THE START AND STOP COMMENTS\n",
    "\n",
    "# START COMMENT: Enter your system description in this cell.\n",
    "\n",
    "def foo(s):\n",
    "    return True\n",
    "\n",
    "# STOP COMMENT: Please do not remove this comment."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5b39c60-7494-46a6-b450-42b7e9fe3aad",
   "metadata": {},
   "source": [
    "## Question 5: Bakeoff entry [1 point]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cff871c1-cc38-4e2f-af38-45b3619e8329",
   "metadata": {},
   "source": [
    "For the bake-off, you simply need to be able to run your system on the file \n",
    "\n",
    "```data/openqa/cs224u-openqa-test-unlabeled.txt```\n",
    "\n",
    "The following code should download it for you if necessary:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ca87f81-556b-46eb-904f-a3df70fdacb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import wget\n",
    "\n",
    "if not os.path.exists(os.path.join(\"data\", \"openqa\", \"cs224u-openqa-test-unlabeled.txt\")):\n",
    "    os.makedirs(os.path.join('data', 'openqa'), exist_ok=True)\n",
    "    wget.download('https://web.stanford.edu/class/cs224u/data/cs224u-openqa-test-unlabeled.txt', out='data/openqa/')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68d0024b-9af7-4e3b-930e-1e7603d4d85c",
   "metadata": {},
   "source": [
    "If the above fails, you can just download https://web.stanford.edu/class/cs224u/data/cs224u-openqa-test-unlabeled.txt and place it in `data/openqa`.\n",
    "\n",
    "This file contains only questions. The starter code below will help you structure this. It writes a file \"cs224u-openqa-bakeoff-entry.json\" to the current directory. That file should be uploaded as-is. Please do not change its name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "403b0000-5bc0-4657-91e4-5a6e87f2f899",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import tqdm\n",
    "\n",
    "def create_bakeoff_submission(model):\n",
    "    \"\"\"\"\n",
    "    The argument `model` is a `dspy.Module`. The return value of its\n",
    "    `forward` method must have an `answer` attribute.\n",
    "    \"\"\"\n",
    "\n",
    "    filename = os.path.join(\"data\", \"openqa\", \"cs224u-openqa-test-unlabeled.txt\")\n",
    "\n",
    "    # This should become a mapping from questions (str) to response\n",
    "    # dicts from your system.\n",
    "    gens = {}\n",
    "\n",
    "    with open(filename) as f:\n",
    "        questions = f.read().splitlines()\n",
    "\n",
    "    # Here we loop over the questions, run the system `model`, and\n",
    "    # store its `answer` value as the prediction:\n",
    "    for question in tqdm.tqdm(questions):\n",
    "        gens[question] = model(question=question).answer\n",
    "\n",
    "    # Quick tests we advise you to run:\n",
    "    # 1. Make sure `gens` is a dict with the questions as the keys:\n",
    "    assert all(question in gens for q in questions)\n",
    "    # 2. Make sure the values are str:\n",
    "    assert all(isinstance(d, str) for d in gens.values())\n",
    "\n",
    "    # And finally the output file:\n",
    "    with open(\"cs224u-openqa-bakeoff-entry.json\", \"wt\") as f:\n",
    "        json.dump(gens, f, indent=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc0f32a8-547e-4a2c-8283-44adf69657ed",
   "metadata": {},
   "source": [
    "Here's what it looks like to evaluate our first program, `basic_qa_model`, on the bakeoff data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce20e9ae-bb82-4dff-896f-aad7f150177d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "create_bakeoff_submission(basic_qa_model)"
   ]
  }
 ],
 "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.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}