hw_openqa.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cf85bea0-68bf-4405-96ec-37579b2e9587",
   "metadata": {},
   "source": [
    "# Homework and bakeoff: Few-shot OpenQA with DSP"
   ]
  },
  {
   "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, Spring 2023\""
   ]
  },
  {
   "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",
    "[![Open in SageMaker Studio Lab](https://studiolab.sagemaker.aws/studiolab.svg)](https://studiolab.sagemaker.aws/import/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 [DSP programming library](https://github.com/stanfordnlp/dsp) 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 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",
    "As usual, this notebook sets up the task and provides starter code. We will be relying on the DSP library, which allows us to define retrieval-augmented in-context learning systems in code. We first provide two fully implemented examples:\n",
    "\n",
    "* _Few-shot OpenQA_: The given input is a question and the goal is to provide an answer. Some _demonstration_ Q/A pairs are sampled from a train set (in our case, SQuAD).\n",
    "\n",
    "* _Few-shot QA with context_: The given input is a question with an associated evidence passage, and the goal is to provide an answer. The _demonstrations_ are now Q/A pairs with associated gold evidence passages. These are sampled from a train set (in our case, SQuAD).\n",
    "\n",
    "The above examples are followed by some assignment questions aimed at helping you to think creatively about the problem. The first of these defines a core system for our target task:\n",
    "\n",
    "* _Few-shot OpenQA with context_: This is like _few-shot QA with context_ except the passages are now retrieved from a large search index using ColBERT. \n",
    "\n",
    "The second question illustrates how to use the powerful DSP `annotate` function to improve the set of demonstrations used by the system.\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 an Eleuther model on a heavy-duty cluster computer, or you can pay with time by using an Eleuther model on a more modest computer.  __For now, though, the Cohere models are free to use, so they should be your first choice; see [setup.ipynb](setup.ipynb) if you don't have an account__."
   ]
  },
  {
   "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 = 'dsp'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a04cb488-cd40-4f9d-b884-8ff83b012042",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cohere\n",
    "from datasets import load_dataset\n",
    "import openai\n",
    "import os\n",
    "import dsp"
   ]
  },
  {
   "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",
    "openai_key = os.getenv('OPENAI_API_KEY')  # or replace with your API key (optional)\n",
    "\n",
    "cohere_key = os.getenv('COHERE_API_KEY')  # or replace with your API key (optional)\n",
    "\n",
    "colbert_server = 'http://index.contextual.ai:8893/api/search'"
   ]
  },
  {
   "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c118b014-e13f-433d-ad60-074636c7e738",
   "metadata": {},
   "outputs": [],
   "source": [
    "lm = dsp.GPT3(model='text-davinci-001', api_key=openai_key)\n",
    "\n",
    "# Options for Cohere: command-medium-nightly, command-xlarge-nightly\n",
    "#lm = dsp.Cohere(model='command-xlarge-nightly', api_key=cohere_key)\n",
    "\n",
    "rm = dsp.ColBERTv2(url=colbert_server)\n",
    "\n",
    "dsp.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": [
    "# [d[\"root\"] for d in openai.Model.list()[\"data\"]]"
   ]
  },
  {
   "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. It is also useful that it has gold passages supporting the standard QA formulation, so we can see how well our LM performs with an \"oracle\" retriever that always retrieves the gold passage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5eaf2fd0-d060-4100-8702-f7311efd6129",
   "metadata": {},
   "outputs": [],
   "source": [
    "squad = load_dataset(\"squad\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36965402-e3da-4531-b7e9-4b12cebcdf30",
   "metadata": {},
   "source": [
    "The following utility just reads a SQuAD split in as a list of `SquadExample` 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 SquadExample named tuples with attributes\n",
    "    id, title, context, question, answers\n",
    "\n",
    "    \"\"\"\n",
    "    data = zip(*[squad[split][field] for field in squad[split].features])\n",
    "    return [dsp.Example(id=eid, title=title, context=context, question=q, answer=a['text']) \n",
    "            for eid, title, context, q, a in data]"
   ]
  },
  {
   "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": "3ecce47a-ecd2-4143-8c52-81700c060b9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_exs = sorted(squad_dev, key=lambda x: hash(x.id))[: 200]"
   ]
  },
  {
   "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 protocols are the standard ones for SQuAD and related tasks: exact match of the answer (EM) and token-level F1. We'll reply primarily on DSP for these evaluation utilities; the following is a light modification of `dsp.evaluation.utils.evaluateAnswer`, which is itself built evaluation code from [apple/ml-qrecc](https://github.com/apple/ml-qrecc/blob/main/utils/evaluate_qa.py) repository. It performs very basic string normalization before doing the core comparisons."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e8bf612-5195-4f4a-b72a-bf9faf128142",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dsp.utils import EM, F1\n",
    "import tqdm\n",
    "import pandas as pd\n",
    "\n",
    "def evaluateAnswer(fn, dev):\n",
    "    \"\"\"Evaluate a DSP program on `dev`.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    fn : DSP system\n",
    "    def : list of `dsp.Example` instances\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    dict with keys \"df\", \"em\", \"f1\" storung assessment data\n",
    "    \"\"\"\n",
    "    data = []\n",
    "    for example in tqdm.tqdm(dev):\n",
    "        prediction = fn(example)\n",
    "        d = dict(example)\n",
    "        pred = prediction.answer\n",
    "        d['prediction'] = pred\n",
    "        d['em'] = EM(pred, example.answer)\n",
    "        d['f1'] = F1(pred, example.answer)\n",
    "        data.append(d)\n",
    "    df = pd.DataFrame(data)\n",
    "    em = round(100.0 * df['em'].sum() / len(dev), 1)\n",
    "    df['em'] = df['em'].apply(lambda x: '✔️' if x else '❌')\n",
    "    f1 = df['f1'].mean()\n",
    "    return {'df': df, 'em': em, 'f1': f1}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28265d01-890d-4f04-b518-da0e8a1cb235",
   "metadata": {},
   "source": [
    "## DSP 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(\"Which U.S. states border no U.S. states?\")"
   ]
  },
  {
   "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": [
    "### Prompt templates"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f60d49d4-f690-496a-a0ab-341096d8c1fb",
   "metadata": {},
   "source": [
    "In DSP, the more usual way to call the LM is to define a prompt template. Here we define a generic QA prompt template:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f6f50d6-e296-412f-a959-62e776b6f27c",
   "metadata": {},
   "outputs": [],
   "source": [
    "Question = dsp.Type(\n",
    "    prefix=\"Question:\", \n",
    "    desc=\"${the question to be answered}\")\n",
    "\n",
    "Answer = dsp.Type(\n",
    "    prefix=\"Answer:\", \n",
    "    desc=\"${a short factoid answer, often between 1 and 5 words}\", \n",
    "    format=dsp.format_answers)\n",
    "\n",
    "qa_template = dsp.Template(\n",
    "    instructions=\"Answer questions with short factoid answers.\", \n",
    "    question=Question(), \n",
    "    answer=Answer())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cfd383f-aa80-44ab-a381-f53e3b8e6631",
   "metadata": {},
   "source": [
    "And here is a self-contained example that uses our question and template to create a prompt:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b846dd8b-f8d6-44c6-9079-342bc3f54969",
   "metadata": {},
   "outputs": [],
   "source": [
    "states_ex = dsp.Example(\n",
    "    question=\"Which U.S. states border no U.S. states?\",\n",
    "    demos=dsp.sample(squad_train, k=2))\n",
    "\n",
    "print(qa_template(states_ex))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb58e275-99a9-45dc-8ae9-fcbdc191f66b",
   "metadata": {},
   "source": [
    "### Prompt-based generation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e0be5e0-7e7f-4e81-bcc0-4ce7a882eecd",
   "metadata": {},
   "source": [
    "We can how put the above pieces together to call the model with our constructed prompt:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c7d277c-a0aa-46c7-b0c3-1c7264efe75e",
   "metadata": {},
   "outputs": [],
   "source": [
    "states_ex, states_compl = dsp.generate(qa_template)(states_ex, stage='basics')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "413b75c4-8706-4936-a26b-b45c254b3e13",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(states_compl.answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8460291-5c58-44bf-9d7e-02176ee6fc43",
   "metadata": {},
   "source": [
    "And here's precisely what the model saw and did:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "820774e4-339f-4263-9493-9f91e7743a17",
   "metadata": {},
   "outputs": [],
   "source": [
    "lm.inspect_history(n=1)"
   ]
  },
  {
   "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": "code",
   "execution_count": null,
   "id": "2e395070-4c6f-4194-81ff-331e32eecd1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "states_ex.question"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5dfd114-96cf-468a-bac3-d3d39d6f3ca6",
   "metadata": {},
   "source": [
    "The basic `dsp.retrieve` method returns only passages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "891fc391-c177-4da7-9332-ab20cdba3c0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "passages = dsp.retrieve(states_ex.question, k=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abdac37b-b5fe-421c-826f-4fd699cb2e36",
   "metadata": {},
   "outputs": [],
   "source": [
    "passages"
   ]
  },
  {
   "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(states_ex.question, k=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db2017ee-1375-4251-a24f-7f792852ffac",
   "metadata": {},
   "source": [
    "## Few-shot OpenQA"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa1dc2c2-8d8b-4cd2-a40f-ee023932359a",
   "metadata": {},
   "source": [
    "With the above pieces in place, we can define our first DSP system. This one does few-shot OpenQA with no context passages. In essense, our prompts contain\n",
    "\n",
    "1. A sequences of Q/A demonstrations (no context passages).\n",
    "2. The target question (no context passage).\n",
    "\n",
    "Here is the full system; note the use of the decorator `@dsp.transformation` – this will ensure that no `example` instances are modified when the program is used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df25de1c-ee5c-42ec-b1f0-a7194ba6063a",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dsp.transformation\n",
    "def few_shot_openqa(example, train=squad_train, k=2): \n",
    "    example.demos = dsp.sample(train, k=k)\n",
    "    example, completions = dsp.generate(qa_template)(example, stage='qa')\n",
    "    return completions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "952dd9d4-ba1f-4689-a548-4e216610f268",
   "metadata": {},
   "source": [
    "There are really just two steps here. Let's go through them individually. Our example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b15bb28d-d46e-4a26-ab01-42c73ef953ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "ex = squad_dev[0].copy()\n",
    "\n",
    "ex"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed8407fb-a556-45f0-9c2a-92d08fbba5f7",
   "metadata": {},
   "source": [
    "We add some demonstrations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cb1e73b-7368-4b95-bf63-d4ef05f1ec2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "ex.demos = dsp.sample(squad_train, k=2)\n",
    "\n",
    "ex"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ae66b9f-ff8d-4952-bc59-f2da0d949089",
   "metadata": {},
   "source": [
    "And then we call the LM using `qa_template`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc564cf6-d1d2-4a11-b0f2-58bfce3e34ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "ex, ex_compl = dsp.generate(qa_template)(ex, stage='qa')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a000eec9-a63b-41fc-9641-60110996f231",
   "metadata": {},
   "source": [
    "Here, `ex_compl` is a `Completions` instance. We will typically use only the `answer` attribute:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba4656bd-2f21-47b0-85f7-91720076be6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(ex_compl.answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "297c0a19-84c4-4342-ac54-4750a2e70710",
   "metadata": {},
   "source": [
    "And, as a final check, we can see precisely what the LM saw:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57cb7634-4f80-4f8a-a855-f96c696717ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "lm.inspect_history(n=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a4b36c9-5549-4e67-8f07-c640229e69cc",
   "metadata": {},
   "source": [
    "## Few-shot QA with context"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4888840f-27d4-4f09-bcdc-6a95946fa8c5",
   "metadata": {},
   "source": [
    "The above system makes no use of evidence passages. As a first step toward bringing in such passages, we define a regular few-shot QA system. For this system, prompts contain:\n",
    "\n",
    "1. A sequences of Q/A demonstrations, each with a gold context passage.\n",
    "2. The target question with a gold context passage.\n",
    "\n",
    "This kind of system is very demanding in terms of data, since we need to have gold evidence passages for every Q/A pair used for demonstations and the Q that is our target. Datasets like SQuAD support this, but it's a rare situation in the world. (Our next system will address this by dropping the need for gold passages)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d451db1-a2c3-4c4f-982f-3a00e58c3ff6",
   "metadata": {},
   "source": [
    "### Template with context"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67513718-09e7-4386-ab78-aabac3b6e9ef",
   "metadata": {},
   "source": [
    "The first step toward defining this system is a new prompt template that includes context:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47cd46e0-ca91-4fe3-bd2c-e4005a16669d",
   "metadata": {},
   "outputs": [],
   "source": [
    "Context = dsp.Type(\n",
    "    prefix=\"Context:\\n\",\n",
    "    desc=\"${sources that may contain relevant content}\",\n",
    "    format=dsp.passages2text)\n",
    "\n",
    "qa_template_with_passages = dsp.Template(\n",
    "    instructions=qa_template.instructions,\n",
    "    context=Context(), \n",
    "    question=Question(), \n",
    "    answer=Answer())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a24aefd-0aa9-4b82-b8c3-42a06247bd14",
   "metadata": {},
   "source": [
    "Here's what this does for a SQUaD example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c0b3571-16e6-48a0-9225-e92430ed0082",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(qa_template_with_passages(ex))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78c0a313-ba07-4b1d-9da4-a3fbf11321ca",
   "metadata": {},
   "source": [
    "### The system"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "987a7406-dc46-4d6a-af19-198c11214211",
   "metadata": {},
   "source": [
    "And here is the full system; the code is identical to `few_shot_openqa` except we now use `qa_template_with_passages`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90323a84-b5bb-4bb0-a20e-ee89d85d09e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dsp.transformation\n",
    "def few_shot_qa_with_context(example, train=squad_train, k=3):\n",
    "    example.demos = dsp.sample(train, k=k)\n",
    "    generator = dsp.generate(qa_template_with_passages)\n",
    "    example, completions = generator(example, stage='qa')\n",
    "    return completions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4871b0f5-f736-4d85-bfbe-199be78ff4eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(few_shot_qa_with_context(squad_dev[0]).answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f6d5962-0993-4cac-9fc0-8ed18f608e11",
   "metadata": {},
   "outputs": [],
   "source": [
    "lm.inspect_history(n=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad912442-d489-48a7-90b3-eb85d5318771",
   "metadata": {},
   "source": [
    "## Dev evaluations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5204ad57-6d89-44f3-bde4-d788784a9fff",
   "metadata": {},
   "source": [
    "This quick section shows some full evaluations using `evaluateAnswer` (see [Evaluation](#Evaluation) above). Depending on which model you're using, these evaluations could be expensive, so you might want to run them only sparingly. Here I am running them on just 25 dev examples to further avoid cost run-ups."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5465c228-9ada-4c70-a694-669b9629ee75",
   "metadata": {},
   "outputs": [],
   "source": [
    "tiny_dev = dev_exs[: 25]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28ea0d8d-e51a-4b06-b7c7-aed2fcadb74d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# few_shot_openqa_results = evaluateAnswer(few_shot_openqa, tiny_dev)\n",
    "#\n",
    "# print(few_shot_openqa_results['em'])\n",
    "# print(few_shot_openqa_results['f1'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39992dc4-c066-453e-8e1a-247823b41865",
   "metadata": {},
   "source": [
    "You can also see the full set of results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47d06cbf-1a36-4492-8b2a-7ffdb1e324e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# few_shot_openqa_results['df'].head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fcb9157-bfde-4990-9379-3194b2cd4705",
   "metadata": {},
   "outputs": [],
   "source": [
    "# few_shot_qa_results = evaluateAnswer(few_shot_qa_with_context, tiny_dev)\n",
    "#\n",
    "# print(few_shot_qa_results['em'])\n",
    "# print(few_shot_qa_results['f1'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eeb28556-d901-4b6b-8a7d-93040203294d",
   "metadata": {},
   "source": [
    "## Question 1: Few-shot OpenQA with context [3 points]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f481081-3e16-4ccd-b379-c7e0c3011286",
   "metadata": {},
   "source": [
    "Your task here is to define a first instance of our target system: Few-shot OpenQA with context passages. To do this, you simply complete `few_shot_openqa_with_context`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ca394a8-be1a-4dfb-b9f6-b40f6b1fb146",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dsp.transformation\n",
    "def few_shot_openqa_with_context(example, train=squad_train, k=3):\n",
    "    pass\n",
    "    # Sample `k` demonstrations from `train`:\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "    # For each demonstration, retrieve one passage and add it\n",
    "    # as the `context` attribute` so we can use our template\n",
    "    # `qa_template_with_passages`:\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "    # Add the list of demonstrations to `example` as the `demos` attribute:\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "    # Retrieve a context passage for `example` itself and add it\n",
    "    # as the `context` attribute:\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "    # Use `dsp.generate` to call the model on `example` using\n",
    "    # `qa_template_with_passages`:\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "    # Return the Completions instance returned by `dsp.generate`:\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "271fd165-fd7f-4943-8af4-a83f0d976453",
   "metadata": {},
   "source": [
    "A quick test you can use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1568f224-e6bb-443c-ae75-dd55a062d9fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_few_shot_openqa_with_context(func):\n",
    "    ex = dsp.Example(question=\"Q0\", context=\"C0\", answer=[\"A0\"])\n",
    "    train = [\n",
    "        dsp.Example(question=\"Q1\", context=None, answer=[\"A1\"]),\n",
    "        dsp.Example(question=\"Q2\", context=None, answer=[\"A2\"]),\n",
    "        dsp.Example(question=\"Q3\", context=None, answer=[\"A3\"])]\n",
    "    compl = func(ex, train=train, k=2)\n",
    "    errcount = 0\n",
    "    # Check the LM was used as expected:\n",
    "    if len(compl.data) != 1:\n",
    "        errcount += 1\n",
    "        print(f\"Error for `{func.__name__}`: Unexpected LM output.\")\n",
    "    data = compl.data[0]\n",
    "    # Check that the right number of demos was used:\n",
    "    demos = data['demos']\n",
    "    if len(demos) > 2:\n",
    "        errcount += 1\n",
    "        print(f\"Error for `{func.__name__}`: \"\n",
    "              f\"Unexpected demo count: {len(demos)}\")\n",
    "    # Check that context passages were included in the prompt:\n",
    "    fields = compl.template.fields\n",
    "    if not any(f.name == 'Context:' for f in fields):\n",
    "        errcount += 1\n",
    "        print(f\"Error for `{func.__name__}`: \"\n",
    "              f\"No context passages in the prompt.\")\n",
    "    # Check that the context passages were retrieved:\n",
    "    if data['context'] == \"C0\":\n",
    "        errcount += 1\n",
    "        print(f\"Error for `{func.__name__}`: \"\n",
    "              f\"No context passage retrieved for the target.\")\n",
    "    for d in demos:\n",
    "        if d['context'] is None:\n",
    "            errcount += 1\n",
    "            print(f\"Error for `{func.__name__}`: \"\n",
    "                  f\"No context passage retrieved for demo {d}.\")\n",
    "    if errcount == 0:\n",
    "        print(f\"No errors found for `{func.__name__}`\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8efe5f07-df96-4d95-8334-8c4966b66269",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_few_shot_openqa_with_context(few_shot_openqa_with_context)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2448dfa0-e945-44b6-a487-70e49675cf61",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(few_shot_openqa_with_context(dev_exs[0]).answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9855ec2-84cc-4a2f-ad47-502c26e44dd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "lm.inspect_history(n=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddccccd5-78f9-47f6-b1d1-7f1171a12b51",
   "metadata": {},
   "source": [
    "Here's an optional evaluation of the system using `tiny_dev`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11f2f036-f18f-4be2-983f-283c07d3e019",
   "metadata": {},
   "outputs": [],
   "source": [
    "# few_shot_openqa_with_context_results = evaluateAnswer(\n",
    "#     few_shot_openqa_with_context, tiny_dev)\n",
    "#\n",
    "# few_shot_openqa_with_context_results['f1']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9971ffd0-a563-4d56-a896-0735a63ae92f",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Question 2: Using annotate"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "054b0fbd-b21b-4b6f-a24c-54235ad1833a",
   "metadata": {},
   "source": [
    "This question is designed to give you some experience with DSP's powerful `annotate` method. You can think of this as a generic tool for defining general aspects of your prompt. Here we will use it to filter the set of demonstrations we use.\n",
    "\n",
    "The overall idea here is that the demonstrations we sample might vary in quality in ways that could impact model performance. For example, if we want to try to push the model to provide extractive answers as in classical QA – answers that are substrings of the evidence passage – then it works against our interests to include demonstrations where the model is unabel to do this.\n",
    "\n",
    "We will do this in two parts to facilitate testing."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d268d575-7cfc-4fb6-9d12-87fcc14c59c1",
   "metadata": {},
   "source": [
    "### Task 1: Filtering demonstrations 1 [2 points]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae208ecb-1569-4c4a-b105-828084f458ef",
   "metadata": {},
   "source": [
    "This is the heart of the question: complete `filter_demos` so that, given a demonstration `d` and a list of demonstrations `demos`, it keeps `d` if and only if\n",
    "\n",
    "1. The passage retrieved for `d` contrains `d.answer`, and\n",
    "2. The model's generation for `d` based on `qa_template_with_passages` contains `d.answer`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d90e4d8-9ecb-4273-93f0-5d0550ec8da1",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dsp.transformation\n",
    "def filter_demos(d):\n",
    "\n",
    "    # Retrieve a passage for `d.question` and make sure that it\n",
    "    # contains `d.answer`. Use `dsp.passage_match` for this!\n",
    "    # return None if there is no match.\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "    # Sample `k=3` demonstrations to help the model assess this\n",
    "    # potential demonstration:\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "    # Generate an answer based on `qa_template_with_passages`\n",
    "    # and use `dsp.answer_match` to check that the predicted answer\n",
    "    # contains `d.answer`. If it does not, return None.\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "    # Return d, if you got this far:\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9674c760-e507-4b60-9f17-5fc91500a68e",
   "metadata": {},
   "source": [
    "Here's a test; this is not an ideal unit test because we don't know which LM you will be using, but it should clarify our intentions and help you with debugging."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66812d32-cc0b-4736-8483-4bb3812f1e9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_filter_demos(func):\n",
    "    # This example should be filtered at the retrieval step, since\n",
    "    # 👽 is not in the index:\n",
    "    ex1 = dsp.Example(\n",
    "        question=\"Who is 👽?\", context=\"C0\", answer=[\"👽\"])\n",
    "    result1 = func(ex1)\n",
    "    errcount = 0\n",
    "    if result1 is not None:\n",
    "        errcount += 1\n",
    "        print(f\"Error for `{func.__name__}`: Expected {None}, got {result1}\")\n",
    "    # This example should not be filtered given our tester LM:\n",
    "    ex2 = dsp.Example(\n",
    "        question=\"Who is Beyoncé?\", context=\"C0\", answer=[\"Beyoncé\"])\n",
    "    # This example should be filtered given our tester LM:\n",
    "    ex3 = dsp.Example(\n",
    "        question=\"Who is Beyoncé?\", context=\"C0\", answer=[\"NO MATCH\"])\n",
    "    class TestLM:\n",
    "        def __init__(self, **kwargs):\n",
    "            self.kwargs = kwargs\n",
    "            self.history = []\n",
    "\n",
    "        def __call__(self, prompt, **kwargs):\n",
    "            answer = [\"Beyoncé\"]\n",
    "            return answer\n",
    "    dsp.settings.configure(lm=TestLM(), rm=rm)\n",
    "    try:\n",
    "        result2 = func(ex2)\n",
    "        if result2 is None:\n",
    "            errcount += 1\n",
    "            print(f\"Error for `{func.__name__}`: \"\n",
    "                  f\"Expected example not to be filtered by `answer_match`.\")\n",
    "        result3 = func(ex3)\n",
    "        if result3 is not None:\n",
    "            errcount += 1\n",
    "            print(f\"Error for `{func.__name__}`: \"\n",
    "                  f\"Expected example to be filtered by `answer_match`.\")\n",
    "    except:\n",
    "        raise\n",
    "    finally:\n",
    "        # Restore the actual model:\n",
    "        dsp.settings.configure(lm=lm, rm=rm)\n",
    "    if errcount == 0:\n",
    "        print(f\"No errors detected for `{func.__name__}`\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef762c5a-68cf-422a-bec4-aa369f73aa87",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_filter_demos(filter_demos)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b897f65-b966-499a-87f1-dd7f8625d069",
   "metadata": {},
   "source": [
    "### Task 2: Full filtering program [1 point]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c08168b-64fe-4965-94f5-c3dc4675dc8a",
   "metadata": {},
   "source": [
    "The task is to complete `few_shot_openqa_with_context_and_demo_filtering` as a few-shot OpenQA system like the one from Question 1, but using the filtering mechanism defined by `filter_demos`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2d3601f-1f61-4f6a-b8d9-09ce515f19c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dsp.transformation\n",
    "def few_shot_openqa_with_context_and_demo_filtering(example, train=squad_train, k=3):\n",
    "\n",
    "    # Sample 20 demonstrations:\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "    # Filter the demonstrations using `annotate` and `filter_demos`.\n",
    "    # The user's `k` should be used to specify the maximum number of\n",
    "    # demonstrations kept at this stage.\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "    # Add the list of filtered demonstrations as a the `demos`\n",
    "    # attribute of `example`:\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "    # Retrieve a context passage for `example.question` and add it\n",
    "    # as the `context` attribute for the example:\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "    # Generate a prediction using `qa_template_with_passages` as\n",
    "    # we did before:\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "    # Return the generated `Completions` instance:\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa4db956-79a8-4092-a553-91f0d79f3084",
   "metadata": {},
   "source": [
    "Our previous test should suffice to help with debugging this program:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fc93545-cd13-49a3-840c-385d29bb397f",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_few_shot_openqa_with_context(\n",
    "    few_shot_openqa_with_context_and_demo_filtering)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c55dbeb-031d-4429-8669-17162d2ea05b",
   "metadata": {},
   "source": [
    "Quiick example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26f0964f-136a-4e85-9944-4472dcd5cbc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(few_shot_openqa_with_context_and_demo_filtering(dev_exs[0]).answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea4224d4-cedc-4f58-a0b6-20b8f22d0317",
   "metadata": {},
   "outputs": [],
   "source": [
    "lm.inspect_history(n=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5ae2262-3eea-487b-8e58-96f310892a19",
   "metadata": {},
   "source": [
    "Here is code for an optional initial evaluation with `tiny_dev`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9c06c67-dd96-46f2-9fac-653d562fc385",
   "metadata": {},
   "outputs": [],
   "source": [
    "# filtering_results = evaluateAnswer(\n",
    "#     few_shot_openqa_with_context_and_demo_filtering, tiny_dev)\n",
    "\n",
    "# filtering_results['f1']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19655d70-007c-4a41-9f3b-e20df9c5169b",
   "metadata": {},
   "source": [
    "## Question 3: 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 so far sampled randomly from the SQuaD train set to create few-shot prompts. One might instead sample passages that have some connection to the target question. See `dsp.knn`, for example.\n",
    "\n",
    "* There are a lot of parameters to our LMs that we have so far ignored. Exploring different values might lead to better results. The `temperature` parameter is highly impactful for our task.\n",
    "\n",
    "* We have so far made no use of the scores from the LM or the RM.\n",
    "\n",
    "* We have so far made no use of DSP's functionality for self-consistency. See the DPS intro notebook for examples.\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",
    "\n",
    "# STOP COMMENT: Please do not remove this comment."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5b39c60-7494-46a6-b450-42b7e9fe3aad",
   "metadata": {},
   "source": [
    "## Question 4: 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": [
    "if not os.path.exists(os.path.join(\"data\", \"openqa\", \"cs224u-openqa-test-unlabeled.txt\")):\n",
    "    !mkdir -p data/openqa\n",
    "    !wget https://web.stanford.edu/class/cs224u/data/cs224u-openqa-test-unlabeled.txt -P 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": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "def create_bakeoff_submission(fn):\n",
    "    \"\"\"\"\n",
    "    The argument `fn` is a DSP program with the same signature as the \n",
    "    ones we wrote above: `dsp.Example` to `dsp.Completions`.\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",
    "    questions = [dsp.Example(question=q) for q in questions]\n",
    "\n",
    "    # `questions` is the list of `dsp.Example` instances you need to \n",
    "    # evaluate your system on. \n",
    "    #\n",
    "    # Here we loop over the questions, run the system `fn`, and\n",
    "    # store its `answer` value as the prediction:\n",
    "    for question in tqdm.tqdm(questions):\n",
    "        gens[question.question] = fn(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(q.question in gens for q in questions)\n",
    "    # 2. Make sure the values are dicts and have the key we will use:\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, `few_shot_openqa`, on the bakeoff data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce20e9ae-bb82-4dff-896f-aad7f150177d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create_bakeoff_submission(few_shot_openqa_with_context)"
   ]
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}