hw_recogs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e9804d50-4eb0-4f1d-a20d-d6c3dbb0fd65",
   "metadata": {},
   "source": [
    "# Homework and bakeoff: Compositional generalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c401fbca-fc3d-4e35-bb84-15b1ce9b974f",
   "metadata": {},
   "outputs": [],
   "source": [
    "__author__ = \"Christopher Potts and Zhengxuan Wu\"\n",
    "__version__ = \"CS224u, Stanford, Spring 2024\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "117fc266-35fd-4321-beb2-8805073ebcdf",
   "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_recogs.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": "e2111f33-11d7-455d-a1cb-5f0c828559dd",
   "metadata": {},
   "source": [
    "## Overview"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c392958-5e6c-4a57-88b6-dc2880928692",
   "metadata": {},
   "source": [
    "This assignment is about _compositional generalization_. We are going to assess the degree to which our apparently very good models have learned to process and interpret language _systematically_. To do this, we are going to ask them to interpret novel combinations of familiar words and phrases. For humans, these tasks are very easy. For our models, the situation seems to be quite different.\n",
    "\n",
    "The basis for the assignment is the ReCOGS dataset of [Wu, Manning, and Potts 2023](https://arxiv.org/abs/2303.13716). ReCOGS modifies the COGS dataset of [Kim and Linzen 2020](https://aclanthology.org/2020.emnlp-main.731) in a number of ways, with the goal of more directly assessing the interpretive abilities of models.\n",
    "\n",
    "The assignment questions are fairly diverse. Question 1 asks you to conduct a specific analysis of the ReCOGS dataset, and Question 2 follows this up with a corresponding analysis of the errors made by a top-performing ReCOGS model. For Question 3, you try some in-context learning with DSPy. And then we open things up as usual: you can do anything you want for your original system, and you enter that system's predictions into a bakeoff.\n",
    "\n",
    "There is only one rule that we need to enforce throughout this work:\n",
    "\n",
    "__You cannot train your system on any examples from `dataset[\"gen\"]`, nor can the output representations from those examples be included in any prompts used for in-context learning.__\n",
    "\n",
    "The nature of your original system is otherwise unconstrained."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62709719-bfb2-45fe-908f-8ea72caf2470",
   "metadata": {},
   "source": [
    "## Set-up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07497e3f-50cc-4224-b04b-23b7ed5577c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    # This library is our indicator that the required installs\n",
    "    # need to be done.\n",
    "    import datasets\n",
    "except ModuleNotFoundError:\n",
    "    !git clone https://github.com/cgpotts/cs224u/\n",
    "    !pip install -r cs224u/requirements.txt\n",
    "    import sys\n",
    "    sys.path.append(\"cs224u\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b7d665d-000e-4bae-83a8-1f7c78879df5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "from compgen import recogs_exact_match"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5da265b5-a622-4f75-b0df-60cc53f7273e",
   "metadata": {},
   "source": [
    "The default location of the data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23eaea62-edd6-4f41-add5-353b186782ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "SRC_DIRNAME = os.path.join(\"data\", \"recogs\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b11f934-f2c3-4bc3-8c5f-d3dbbd450eb1",
   "metadata": {},
   "source": [
    "The following code should grab the dataset for you; if it fails for any reason, you can manually download it from [this link](https://web.stanford.edu/class/cs224u/data/recogs.tgz) and then put it in `SRC_DIRNAME`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b902be49-6380-46ae-a353-a634b1cab256",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import wget\n",
    "\n",
    "if not os.path.exists(SRC_DIRNAME):\n",
    "    os.makedirs('data', exist_ok=True)\n",
    "    wget.download('https://web.stanford.edu/class/cs224u/data/recogs.tgz', out='data/')\n",
    "    !tar xvf data/recogs.tgz -C data/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5da14956-1813-4372-b747-d7a2dbd10bc1",
   "metadata": {},
   "source": [
    "## Load the COGS and ReCOGS datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49118e3a-a4fc-4a6a-9d88-c0f8371a390c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_split(filename):\n",
    "    return pd.read_csv(\n",
    "        filename,\n",
    "        delimiter=\"\\t\",\n",
    "        names=['input', 'output', 'category'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fda2c30-a788-4a62-806c-f7318797474d",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = {}\n",
    "\n",
    "for splitname in (\"train\", \"dev\", \"gen\"):\n",
    "    dataset[splitname] = load_split(f\"{SRC_DIRNAME}/{splitname}.tsv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e5a50d3-6f83-4c66-bea5-548ccf25ba7a",
   "metadata": {},
   "source": [
    "Here's a look at the dataset. Fundamentally, the task is to map simple English sentences to logical forms. For ReCOGS, you need only predict these forms up to semantic equivalence, which means that we abstract away from the order of the conjuncts and the names of specific variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6900c28-1949-419b-a8e4-9d534f973dae",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset['train'].head(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03961b86-b824-4c5f-8e54-8258b71a6466",
   "metadata": {},
   "source": [
    "The `dataset['gen']` section is divided up into different 21 categories. A category name `X_to_Y` or `only_seen_as_X_as_Y`  means that specific phrases were seen only as `X` in training and will encounter those phrases as `Y` at test time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9eb7e84b-5304-4d4c-a0a5-919ff79a7730",
   "metadata": {},
   "outputs": [],
   "source": [
    "sorted(dataset['gen'].category.unique())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5264a754-0903-40bf-8012-9abd9d7a41dc",
   "metadata": {},
   "source": [
    "## Question 1: Proper names and their semantic roles"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "989aaeec-4de8-4d7f-a90f-334117b4ea04",
   "metadata": {},
   "source": [
    "A number of the COGS/ReCOGS generalization categories assess models on their ability to handle proper names appearing in novel positions at test time. For example, in the `obj_to_subj_proper` category, models encounter proper names that appeared in the train set only in grammatical object position (e.g., _see Sandy_), and then they are asked to make predictions about cases where those names are grammatical subjects (_Sandy left_). These changes have systematic effects on the grammatical roles that the meanings of these names play semantically. In particular, subjects are likely to play `agent` roles and objects are likely to play `theme` roles."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4acb7cb6-5e43-4697-89eb-c16c592209b1",
   "metadata": {},
   "source": [
    "### Task 1: Pattern-based analysis function [1 point]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "552d0f81-cd75-4895-9414-1ada5c9c9f28",
   "metadata": {},
   "source": [
    "Write a function that scans ReCOGS logical forms to determine what role proper names play. The following are the core steps:\n",
    "\n",
    "1. Identify proper names. All and only proper names begin with capital letters in these LFs, and proper names consist only of ascii letters. The format is, informally, `Name ( d+ )`, as in `Sandy ( 47 )`.\n",
    "\n",
    "2. Identify role expressions. The pattern is always `role ( d+ , d+ )`, as in `agent ( 1 , 47 )`. Here, the first variable is for the associated event, and the second is the role argument. The possible roles are `agent`, `theme`, and `recipient`. (The dataset includes other roles, but these involve events, not people.)\n",
    "\n",
    "3. Determine which of the above are linked in the sense that the variable names are the same. A given name can link to multiple role expressions (or none at all), and LFs can contain multiple names and multiple role expressions.\n",
    "\n",
    "To do the above, you just need to complete the function `get_propername_role`. The test should clear up any ambiguity and help you iterate to a solution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3a19570-3a3c-4843-a9de-f48f2b746dfe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "def get_propername_role(s):\n",
    "    \"\"\"Extract from `s` all the pairs `(name, role)` determined by\n",
    "    binding relationships. There can be multiple tokens of the same\n",
    "    name with different variables, as in \"Kim ( 1 )\" and \"Kim ( 47 )\",\n",
    "    and there can be instances in which a single name with variable\n",
    "    like \"Kim ( 1 )\" binds into multiple role expressions like\n",
    "    \"agent ( 4 , 1 )\" and \"theme ( 6 , 1 )\". Your function should\n",
    "    cover all these cases.\n",
    "\n",
    "    We've suggested a particular program design to get you started,\n",
    "    but you are free to do something different and perhaps cleverer\n",
    "    if you wish!\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    s: str\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    set of tuples `(name, role)` where `name` and `role` are str\n",
    "    \"\"\"\n",
    "    # Step 1: Define a regex for \"name ( var )\" expressions:\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "    # Step 2: Define a regex for \"role ( var , var )\" expressions:\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "    # Step 3: Use `findall` with both of your regexs:\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "    # Step 4: Loop overall combinations of matches from your regexs\n",
    "    # to build `data` as a set of pairs `(name, role)`:\n",
    "    data = set()\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "    # Step 5: Return `data`:\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60f3e98e-78f2-4096-b5df-db38700997dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_get_propername_role(func):\n",
    "    examples = [\n",
    "        # Standard case:\n",
    "        (\n",
    "            \"Bella ( 7 ) ; smile ( 4 ) AND agent ( 4 , 7 )\",\n",
    "            {(\"Bella\", \"agent\")}\n",
    "        ),\n",
    "        # No binding:\n",
    "        (\n",
    "            \"Riley ( 37 ) ; theme ( 4 , 7 )\",\n",
    "            set()\n",
    "        ),\n",
    "        # Two tokens of the same name referring to different entities:\n",
    "        (\n",
    "            \"Riley ( 37 ) ; Riley ( 4 ) ; theme ( 1 , 37 ) AND agent ( 1 , 4 )\",\n",
    "            {(\"Riley\", \"theme\"), (\"Riley\", \"agent\")},\n",
    "        ),\n",
    "        # Two names:\n",
    "        (\n",
    "            \"Riley ( 4 ) ; Emma ( 243 ) ; recipient ( 6 , 4 ) AND agent ( 6 , 243 )\",\n",
    "            {(\"Riley\", \"recipient\"), (\"Emma\", \"agent\")},\n",
    "        ),\n",
    "        # One name binding into multiple role expressions:\n",
    "        (\n",
    "            \"Riley ( 4 ) ; agent ( 6 , 4 ) AND theme ( 6 , 4 )\",\n",
    "            {(\"Riley\", \"theme\"), (\"Riley\", \"agent\")},\n",
    "        ),\n",
    "        # Nothing to match:\n",
    "        (\n",
    "            \"no proper names\",\n",
    "            set()\n",
    "        )\n",
    "    ]\n",
    "    errcount = 0\n",
    "    for ex, expected in examples:\n",
    "        result = func(ex)\n",
    "        if expected != result:\n",
    "            errcount += 1\n",
    "            print(f\"Error for `{func.__name__}`:\"\n",
    "                  f\"\\n\\tInput: {ex}\"\n",
    "                  f\"\\n\\tExpected: {expected}\"\n",
    "                  f\"\\n\\tGot: {result}\")\n",
    "    if errcount == 0:\n",
    "        print(f\"No errors detected for `{func.__name__}`\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01c1d893-6e10-4913-8852-b0df7da1eb67",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_get_propername_role(get_propername_role)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7908627-e0ac-4f97-bfd9-55a83ba70d3a",
   "metadata": {},
   "source": [
    "### Task 2: Finding challenging names [1 point]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab578426-4f74-4e93-879c-e40d8ab92fc7",
   "metadata": {},
   "source": [
    "You can now use your code to find the names that will be the most challenging because their train/gen roles are disjoint. To do this, you just need to complete the function `find_name_roles`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "728fe9a2-35a7-47d3-98b9-56a13362f5c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "\n",
    "def find_name_roles(split_df, colname=\"output\"):\n",
    "    \"\"\"Create a map from names to dicts mapping roles to counts: the\n",
    "    number of time the name appears with role in `split_df`:\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    split_df : pd.DataFrame\n",
    "        Needs to have a column called `colname`.\n",
    "    colname: str\n",
    "        Column to target with `get_propername_role`. Default: \"output\".\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    `defaultdict` mapping names to roles to counts\n",
    "    \"\"\"\n",
    "    # This is a convenient way to create a multidimensional count dict:\n",
    "    # You can access it out of the box as `all_roles[key1][key2] += 1`.\n",
    "    all_roles = defaultdict(lambda : defaultdict(int))\n",
    "\n",
    "    # Apply `get_propername_role` to every value in the target column\n",
    "    # and aggregate the results into `all_roles`:\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "    # Return `all_roles`:\n",
    "    return all_roles"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57613848-3acd-4114-b5f0-49374d7e4c5b",
   "metadata": {},
   "source": [
    "A quick test:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72997def-1245-4c25-98f6-c7a601e35593",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_find_name_roles(func):\n",
    "    df = pd.DataFrame({\n",
    "        \"tester\": [\n",
    "            \"Bella ( 7 ) ; agent ( 4 , 7 )\",\n",
    "            \"Bella ( 7 ) ; agent ( 4 , 7 )\",\n",
    "            \"Riley ( 37 ) ; agent ( 4 , 37 )\",\n",
    "            \"Riley ( 3 ) ; theme ( 4 , 3 )\",\n",
    "            \"Emma ( 37 ) ; theme ( 4 , 7 )\"\n",
    "        ]})\n",
    "    expected = {\n",
    "        \"Bella\": {\"agent\": 2},\n",
    "        \"Riley\": {\"agent\": 1, \"theme\": 1}\n",
    "    }\n",
    "    result = func(df, colname=\"tester\")\n",
    "    if result != expected:\n",
    "        print(f\"Error for `{func.__name__}`:\"\n",
    "              f\"\\n\\tExpected:{expected}\"\n",
    "              f\"\\n\\tGot: {result}\")\n",
    "    else:\n",
    "        print(f\"No errors found for `{func.__name__}`\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb1e75e4-d5f0-4181-85d9-73b9d8b1db8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_find_name_roles(find_name_roles)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6532c822-d02f-4930-b119-d167ee797384",
   "metadata": {},
   "source": [
    "Once the test passes, this analysis should be informative:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb0c22b2-845e-41b8-a958-e5cd9628bf7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_roles = find_name_roles(dataset['train'])\n",
    "\n",
    "sorted(train_roles.items(), key=lambda x: len(x[1]))[: 3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3dc45e2a-78ef-46fd-8119-e054c58b148a",
   "metadata": {},
   "outputs": [],
   "source": [
    "gen_roles = find_name_roles(dataset[\"gen\"])\n",
    "\n",
    "sorted(gen_roles.items(), key=lambda x: len(x[1]))[: 3]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9338e56-1171-48f0-ab8a-0c556d2dfdbf",
   "metadata": {},
   "source": [
    "We will return to these troublemakers in a bit."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6517a076-72b5-4837-8898-fef1f8339657",
   "metadata": {},
   "source": [
    "## Pretrained ReCOGS models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0076aa3-9eea-475d-aea4-b7653a865120",
   "metadata": {},
   "source": [
    "We launch now into an extended interlude before Question 2. For Question 2, you will work with a ReCOGS model that we trained for you. This interlude presents the code needed to work with this model. We are exposing these details to you in case you want to use this code to train or fine-tune your own models for your original system."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80c3e96e-dfc8-4e75-9b24-51e7f33a781e",
   "metadata": {},
   "source": [
    "### Tokenization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43392af3-cc12-449e-9301-983a3d554737",
   "metadata": {},
   "source": [
    "Here is a function for creating Hugging Face `PreTrainedTokenizerFast` tokenizers based on a provided vocab file. It pretty much just splits on whitespace and adds special tokens. Chris originally planned to have writing this be a homework question, but it turned out to be very difficult and confusing for him to write, so he decided to just present it to you in the hope that it helps you with similar tasks in the future."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de55c95c-ef3d-453f-9336-97e9e93e31c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tokenizers import Tokenizer\n",
    "from tokenizers.models import WordLevel\n",
    "from tokenizers.pre_tokenizers import WhitespaceSplit\n",
    "from tokenizers.processors import TemplateProcessing\n",
    "from transformers import PreTrainedTokenizerFast\n",
    "\n",
    "\n",
    "def get_tokenizer(vocab_filename):\n",
    "    with open(vocab_filename) as f:\n",
    "        vocab = f.read().splitlines()\n",
    "    vocab_size = len(vocab)\n",
    "    vocab = dict(zip(vocab, list(range(vocab_size))))\n",
    "    tok = Tokenizer(WordLevel(vocab, unk_token='[UNK]'))\n",
    "    # This definitely needs to be done here and in the construction of\n",
    "    # `PreTrainedTokenizerFast`. Don't be tempted to \"clean this up\"!\n",
    "    tok.add_special_tokens([\"[BOS]\", \"[UNK]\", \"[PAD]\", \"[EOS]\"])\n",
    "    tok.pre_tokenizer = WhitespaceSplit()\n",
    "    tok.post_processor = TemplateProcessing(\n",
    "        single=f\"[BOS]:0 $A:0 [EOS]:0\",\n",
    "        special_tokens=[\n",
    "            (\"[BOS]\", tok.token_to_id(\"[BOS]\")),\n",
    "            (\"[EOS]\", tok.token_to_id(\"[EOS]\"))])\n",
    "    return PreTrainedTokenizerFast(\n",
    "        tokenizer_object=tok,\n",
    "        bos_token=\"[BOS]\",\n",
    "        unk_token=\"[UNK]\",\n",
    "        pad_token=\"[PAD]\",\n",
    "        eos_token=\"[EOS]\",\n",
    "        # This vital; otherwise any periods will have their leading\n",
    "        # spaces removed, which is wrong for COGS/ReCOGS.\n",
    "        clean_up_tokenization_spaces=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a541a47b-a3ed-421c-bb9e-59f3b5f7f2af",
   "metadata": {},
   "source": [
    "We will have separate tokens for the encoder and the decoder:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa80d0a2-4599-4212-9bdb-a6d06263a0b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "enc_tokenizer = get_tokenizer(os.path.join(SRC_DIRNAME, \"src_vocab.txt\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fa0fc58-7b49-4e19-98c9-369eb4b3ce64",
   "metadata": {},
   "outputs": [],
   "source": [
    "enc_tokenizer.tokenize(\n",
    "    \"A sailor was helped\", \n",
    "    add_special_tokens=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb21e534-1d80-4642-983b-96fa415fdaba",
   "metadata": {},
   "outputs": [],
   "source": [
    "dec_tokenizer = get_tokenizer(os.path.join(SRC_DIRNAME, \"tgt_vocab.txt\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "474ceec8-6378-45e3-ac52-659d442f259d",
   "metadata": {},
   "outputs": [],
   "source": [
    "dec_tokenizer.tokenize(\n",
    "    \"sailor ( 53 ) ; help ( 7 ) AND theme ( 7 , 53 )\", \n",
    "    add_special_tokens=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25600c63-2c1f-47e7-82ed-a9f5eacb73b7",
   "metadata": {},
   "source": [
    "### Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c1ede76-83b6-4573-8067-f50dfc97e5dd",
   "metadata": {},
   "source": [
    "Next is a dataset utility. Chris was originally going to have you write this yourselves, since it is useful to know how to write these utilities, and the task is really just to use our tokenizers appropriately. However, since `collate_fn` has to be a static method with fixed arguments, we can't easily pass in these tokenizers to it! As a result, we have to do all the tokenization at once ahead of time and then redo all the masking work for each batch. So Chris did this for you in the hope that this will be useful to you in the future."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2857b954-c426-4f47-a9c3-10761203883d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "class RecogsDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, enc_tokenizer, dec_tokenizer, X, y=None):\n",
    "        self.X = [enc_tokenizer.encode(s) for s in X]\n",
    "        self.y = y\n",
    "        if y is not None:\n",
    "            self.y = [dec_tokenizer.encode(s) for s in y]\n",
    "\n",
    "    @staticmethod\n",
    "    def collate_fn(batch):\n",
    "        \"\"\"Unfortunately, we can't pass the tokenizer in as an argument\n",
    "        to this method, since it is a static method, so we need to do\n",
    "        the work of creating the necessary attention masks.\"\"\"\n",
    "        def get_pad_and_mask(vals):\n",
    "            lens = [len(i) for i in vals]\n",
    "            maxlen = max(lens)\n",
    "            pad = []\n",
    "            mask = []\n",
    "            for ex, length in zip(vals, lens):\n",
    "                diff = maxlen - length\n",
    "                pad.append(ex + ([0] * diff))\n",
    "                mask.append(([1] * length) + ([0] * diff))\n",
    "            return torch.tensor(pad), torch.tensor(mask)\n",
    "        batch_elements = list(zip(*batch))\n",
    "        X = batch_elements[0]\n",
    "        X_pad, X_mask = get_pad_and_mask(X)\n",
    "        if len(batch_elements) == 1:\n",
    "            return X_pad, X_mask\n",
    "        else:\n",
    "            y = batch_elements[1]\n",
    "            y_pad, y_mask = get_pad_and_mask(y)\n",
    "            # Repeat `y_pad` because our optimizer expects to find\n",
    "            # labels in final position. These will not be used because\n",
    "            # Hugging Face will calculate the loss for us.\n",
    "            return X_pad, X_mask, y_pad, y_mask, y_pad\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.X)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        if self.y is None:\n",
    "            return (self.X[idx],)\n",
    "        else:\n",
    "            return (self.X[idx], self.y[idx])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3021faf-dfa6-4d89-9a39-5b4546ac2008",
   "metadata": {},
   "source": [
    "The following just illustrate how to work with the above utility:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3586efad-f12d-4811-8345-e7b18fc46ae9",
   "metadata": {},
   "outputs": [],
   "source": [
    "ex_dataset = RecogsDataset(\n",
    "    enc_tokenizer,\n",
    "    dec_tokenizer,\n",
    "    dataset['train'].input.head(20),\n",
    "    y=dataset['train'].output.head(20))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b80b9ddc-cc02-4b0c-863d-8453fc5bf450",
   "metadata": {},
   "outputs": [],
   "source": [
    "ex_dataloader = torch.utils.data.DataLoader(\n",
    "    ex_dataset,\n",
    "    batch_size=2,\n",
    "    shuffle=True,\n",
    "    pin_memory=True,\n",
    "    collate_fn=ex_dataset.collate_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a08b88f4-abc9-4921-8bf6-78c9c4157198",
   "metadata": {},
   "outputs": [],
   "source": [
    "ex_batch = iter(ex_dataloader)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ef750fa-794c-4fe0-a27f-330377521db3",
   "metadata": {},
   "source": [
    "This will show you batches. Since `batch_size=2` for `dataloader`, this will be a tuple where each element has two lists. The structure is determined by `collate_fn` in `RecogsDataset`: \n",
    "\n",
    "`X_pad, X_mask, y_pad, y_mask, y_pad`\n",
    "\n",
    "where `y_pad` is repeated in the final position to meet the interface specifications of `torch_base_model.py`, in case you decide to train models yourself. (See details below; Hugging Face calculates the loss itself, which is ultimately nice but a bit non-standard.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "016cb2a5-869a-4860-af8a-c94e7586128e",
   "metadata": {},
   "outputs": [],
   "source": [
    "next(ex_batch)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf807690-a8dd-4f8c-8053-943c60e4e295",
   "metadata": {},
   "source": [
    "### Model basics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0349a372-8a06-4a04-b787-2568dceb87f9",
   "metadata": {},
   "source": [
    "Now we come to the model itself. We will first load it and explore it a bit, and then we will define a nice classifier interface for it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76471758-07c0-42b7-b31e-997dc2cd9bd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import EncoderDecoderModel\n",
    "\n",
    "encdec = EncoderDecoderModel.from_pretrained(f\"ReCOGS/ReCOGS-model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a470ef9b-32ed-43f1-85f7-17da743a40b1",
   "metadata": {},
   "source": [
    "A single illustrative example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "666df81f-beb3-4a30-b202-1193908b78b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "ex_inputs = enc_tokenizer.batch_encode_plus(\n",
    "    [\"A rose was helped by a dog .\"], \n",
    "    return_tensors='pt')\n",
    "\n",
    "ex_outputs = dec_tokenizer.batch_encode_plus(\n",
    "    ['rose ( 53 ) ; dog ( 38 ) ; help ( 7 ) AND theme ( 7 , 53 ) AND agent ( 7 , 38 )'], \n",
    "    return_tensors='pt')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0eeca365-29b8-42bb-a696-cc9bcb6dec7d",
   "metadata": {},
   "source": [
    "Here is the forward method. For training, it is vital to have `labels=` here so that the model return a loss value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f64f13fb-b812-445c-8a31-e76233757a4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "ex_rep = encdec(\n",
    "    ex_inputs['input_ids'],\n",
    "    ex_inputs['attention_mask'],\n",
    "    ex_outputs['input_ids'],\n",
    "    labels=ex_outputs['attention_mask'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12266de7-a99d-4ec8-9d46-bdb62485495a",
   "metadata": {},
   "outputs": [],
   "source": [
    "ex_rep.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d934b5a8-dfcb-4355-8355-406ee8d2d164",
   "metadata": {},
   "source": [
    "And here is how we will do generation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71339418-1d43-4e3a-a9a1-a4b55be343c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "ex_gen = encdec.generate(\n",
    "    ex_inputs['input_ids'],\n",
    "    attention_mask=ex_inputs['attention_mask'],\n",
    "    max_new_tokens=512,\n",
    "    eos_token_id=encdec.config.eos_token_id)\n",
    "\n",
    "ex_gen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91da0447-d244-4fd0-8a6a-1c5bec561d69",
   "metadata": {},
   "outputs": [],
   "source": [
    "ex_pred = dec_tokenizer.batch_decode(\n",
    "    ex_gen, \n",
    "    skip_special_tokens=False, \n",
    "    # Out tokenizer have this set already, but I am nervous:\n",
    "    clean_up_tokenization_spaces=False)\n",
    "\n",
    "ex_pred"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac10a038-99b9-467c-a86d-ed82e7bd8fdc",
   "metadata": {},
   "source": [
    "### Model interface"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d70b3db7-b1f5-4fab-84b0-4b96d7ced611",
   "metadata": {},
   "source": [
    "Okay, finally, the main interface. If you do not plan to train your own models using our code, then you can treat `RecogsModel` as an interface and not worry about these details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b17643dd-35bc-4e38-99e4-a61d47d3ed10",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_model_base import TorchModelBase\n",
    "import torch.nn as nn\n",
    "from transformers import EncoderDecoderModel"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d599c05a-1be6-447c-a931-2cadd1112758",
   "metadata": {},
   "source": [
    "As I mentioned above, Hugging Face `EncoderDecoderModel` instances will calculate a loss internally if you provide them with `labels`. Normally, one's optimization loop would need to do this manually. In order to rely on Hugging Face and still use the trainer in `torch_model_base.py`, we define this simple loss that just takes in model outputs and labels and returns `outputs.loss`. The labels argument is present for compatibility; it was already used internally to get the value of `outputs.loss` and so can be ignored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f69d742e-58c2-4993-b6d8-6f18b3d157f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RecogsLoss(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.reduction = \"mean\"\n",
    "\n",
    "    def forward(self, outputs, labels):\n",
    "        \"\"\"`labels` is ignored, as it was already used to assign a\n",
    "        value of `outputs.loss`, and that value is all we need.\"\"\"\n",
    "        return outputs.loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd27db42-46d1-436e-8870-70fdd38e1219",
   "metadata": {},
   "source": [
    "Here is a basic `nn.Module`. Its sole purpose is to organize the examples created by our `RecogsDataset` and feed them to the trained `EncoderDecoderModel`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49dd0d13-41fd-4c9a-b7de-90bd2a7a1840",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RecogsModule(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.encdec = EncoderDecoderModel.from_pretrained(\n",
    "            f\"ReCOGS/ReCOGS-model\")\n",
    "\n",
    "    def forward(self, X_pad, X_mask, y_pad, y_mask, labels=None):\n",
    "        outputs = self.encdec(\n",
    "            input_ids=X_pad, \n",
    "            attention_mask=X_mask,\n",
    "            decoder_attention_mask=y_mask,\n",
    "            labels=y_pad)\n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f996ee6c-ef7c-4f51-b596-a1d69ecc59a2",
   "metadata": {},
   "source": [
    "And, at last, our interface. The keyword parameter `initialize=True` is the default because we are initially going to use this just for making predictions, and so we need the instance to establish all its parameters when we initialize it as opposed to waiting to do that when we call `fit` (which we may never do)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7105b19c-a8b7-44c7-9028-f62936a578b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RecogsModel(TorchModelBase):\n",
    "    def __init__(self, *args,\n",
    "            initialize=True,\n",
    "            enc_vocab_filename=f\"{SRC_DIRNAME}/src_vocab.txt\",\n",
    "            dec_vocab_filename=f\"{SRC_DIRNAME}/tgt_vocab.txt\",\n",
    "            **kwargs):\n",
    "        self.enc_vocab_filename = enc_vocab_filename\n",
    "        self.dec_vocab_filename = dec_vocab_filename\n",
    "        self.enc_tokenizer = get_tokenizer(self.enc_vocab_filename)\n",
    "        self.dec_tokenizer = get_tokenizer(self.dec_vocab_filename)\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.loss = RecogsLoss()\n",
    "        if initialize:\n",
    "            self.initialize()\n",
    "\n",
    "    def build_graph(self):\n",
    "        return RecogsModule()\n",
    "\n",
    "    def build_dataset(self, X, y=None):\n",
    "        return RecogsDataset(\n",
    "            self.enc_tokenizer, self.dec_tokenizer, X, y=y)\n",
    "\n",
    "    def predict(self, X, device=None):\n",
    "        device = self.device if device is None else torch.device(device)\n",
    "        dataset = self.build_dataset(X)\n",
    "        dataloader = self._build_dataloader(dataset, shuffle=False)\n",
    "        self.model.to(device)\n",
    "        self.model.eval()\n",
    "        preds = []\n",
    "        with torch.no_grad():\n",
    "            for batch in dataloader:\n",
    "                X_pad, X_mask = [x.to(device) for x in batch]\n",
    "                outputs = self.model.encdec.generate(\n",
    "                    X_pad,\n",
    "                    attention_mask=X_mask,\n",
    "                    max_new_tokens=512,\n",
    "                    eos_token_id=self.model.encdec.config.eos_token_id)\n",
    "                results = self.dec_tokenizer.batch_decode(\n",
    "                    outputs, \n",
    "                    skip_special_tokens=True,\n",
    "                    clean_up_tokenization_spaces=False)\n",
    "                preds += results\n",
    "        return preds\n",
    "\n",
    "    def score(self, X, y, device=None):\n",
    "        # An overall accuracy score:\n",
    "        preds = self.predict(X, device=device)\n",
    "        vals = [int(recogs_exact_match(gold, pred)) for gold, pred in zip(y, preds)]\n",
    "        return sum(vals) / len(vals)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95d30490-b7ce-4efe-85f7-e101e3481454",
   "metadata": {},
   "outputs": [],
   "source": [
    "recogs_model = RecogsModel()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5f41201-b28b-4352-9bcf-b62b9d977af9",
   "metadata": {},
   "source": [
    "Predictions for our first to train cases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87b99739-eb29-4876-9248-dd72eff1e926",
   "metadata": {},
   "outputs": [],
   "source": [
    "recogs_model.predict(dataset['dev'].input[: 2])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da86e9b3-4281-4c74-92e5-9288f202e881",
   "metadata": {},
   "source": [
    "## Question 2: Exploring predictions [2 points]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06170626-679b-4710-a4a2-f444c104f44b",
   "metadata": {},
   "source": [
    "Now that we are set up to use the model, we can move to Question 2. There is just one final preliminary: for ReCOGs, we want to come as close as possible to assessing systems purely on semantic criteria, as opposed to assessing their ability to predict arbitrary features of logical forms. In particular, we want predictions to be independent of the particular choice of variable names and independent of the order of conjuncts. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa31f1aa-26a0-41a6-a01b-0a4fe85e6367",
   "metadata": {},
   "source": [
    "### ReCOGS assessment function"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e9451d8-2e9d-4a3d-943f-0d06e747b12f",
   "metadata": {},
   "source": [
    "The function `recogs_exact_match` does this. It's a complex function, and so you can ignore its precise implementation details. Here are some illustrative examples to give you a feel for it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08c49a60-082e-42d0-ac75-0a606fd5a588",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The precise names of bound variables do not matter:\n",
    "\n",
    "recogs_exact_match(\n",
    "    \"dog ( 4 ) AND happy ( 4 )\", \n",
    "    \"dog ( 7 ) AND happy ( 7 ) \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "657bc28c-505b-4247-bcbe-dc49718dca2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The order of conjuncts does not matter:\n",
    "\n",
    "recogs_exact_match(\n",
    "    \"dog ( 4 ) AND happy ( 4 )\", \n",
    "    \"happy ( 7 ) AND dog ( 7 )\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a49eb573-2153-4f63-aad1-cb5ac8f2df62",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Consistency of variable names does matter:\n",
    "\n",
    "recogs_exact_match(\n",
    "    \"dog ( 4 ) AND happy ( 4 )\", \n",
    "    \"dog ( 4 ) AND happy ( 7 )\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39be2e2c-d1d1-4f44-8ca0-2ee0f23cf3ee",
   "metadata": {},
   "source": [
    "### Task"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4b83476-d018-4b24-9126-38c3577b68de",
   "metadata": {},
   "source": [
    "Your task is to write a utility function to see how well a model does on a specific generalization category in the generalization dataset. The metric is accuracy according to `recogs_exact_match`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "378d7f44-a317-433e-befa-301cdcd16978",
   "metadata": {},
   "outputs": [],
   "source": [
    "def category_assess(gen_df, model, category):\n",
    "    \"\"\"Assess `model` against the `category` examples in `gen_df`.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    gen_df: pd.DataFrame\n",
    "        Should be `dataset[\"gen\"]`\n",
    "    model: A `RecogsModel instance\n",
    "    category: str\n",
    "        A string from `gen_df.category`\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    `pd.DataFrame` limited to `category` examples and with columns\n",
    "    \"prediction\" and \"correct\" added by this function\n",
    "    \"\"\"\n",
    "    # This line is done for you because of how important it is to\n",
    "    # operate on a copy of the dataframe rather than the original!\n",
    "    cat_df = gen_df[gen_df.category == category].copy()\n",
    "\n",
    "    # Step 1: Add a column called \"prediction\" to `cat_df`. This should\n",
    "    # give the predicted LFs:\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "    # Step 2: Add a column \"correct\" that says whether the prediction\n",
    "    # and the gold output are the same. Must use `recogs_exact_match`.\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "    # Step 3: Return the `pd.DataFrame` `cat_df`:\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14331e44-6eb6-4c87-9930-28d6b0f1506f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_category_assess(func):\n",
    "    testmod = RecogsModel()\n",
    "    samp_df = dataset['gen'].head(150)\n",
    "    examples = [\n",
    "        (\"active_to_passive\", 0.80),\n",
    "        (\"unacc_to_transitive\", 0.86),\n",
    "        (\"obj_to_subj_proper\", 0.78)\n",
    "    ]\n",
    "    result_df = func(samp_df, testmod, \"active_to_passive\")\n",
    "    if not isinstance(result_df, pd.DataFrame):\n",
    "        print(f\"Error `{func.__name__}`: \"\n",
    "              \"Return value should be a `pd.DataFrame`\")\n",
    "        return\n",
    "    errcount = 0\n",
    "    for colname in (\"input\", \"output\", \"category\", \"prediction\", \"correct\"):\n",
    "        if colname not in result_df.columns:\n",
    "            errcount += 1\n",
    "            print(f\"Error `{func.__name__}`: column '{colname}' is missing\")\n",
    "    if errcount != 0:\n",
    "        return\n",
    "    expected_len = 5\n",
    "    result_len = result_df.shape[0]\n",
    "    if not result_df.shape[0] == expected_len:\n",
    "        print(f\"Error `{func.__name__}`: \"\n",
    "              f\"Expected {expected_len} results, got {result_len}.\")\n",
    "        return\n",
    "    errcount = 0\n",
    "    for cat, expected in examples:\n",
    "        result_df = func(samp_df, testmod, cat)\n",
    "        result = result_df.correct.sum() / result_df.shape[0]\n",
    "        result = round(result, 2)\n",
    "        if result != expected:\n",
    "            errcount += 1\n",
    "            print(f\"Error `{func.__name__}` with category {cat}: \"\n",
    "                  f\"Expected acc {expected}, got {result}\")\n",
    "    if errcount == 0:\n",
    "        print(f\"No errors for `{func.__name__}`\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3a178d8-b45b-411c-8995-3e6869322565",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_category_assess(category_assess)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3fcde66-9b8c-4ca6-947b-2bb8a98a3671",
   "metadata": {},
   "source": [
    "Question 1 above might lead you to expect that our model will struggle with examples in which proper names appear with totally unfamiliar roles. For that question, you wrote `get_propername_role` to get `(name, role)` pairs from examples and `find_name_roles` to do analyses with that function. We can now run that same analysis on our errors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67a5305d-4144-45a2-91c4-75e64021fa38",
   "metadata": {},
   "outputs": [],
   "source": [
    "gen_df = dataset['gen']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ede22c2d-6fe3-4dbe-ab0d-ababd9d48d7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Depending on your computer, this could take a while. On a relatively\n",
    "# new Apple laptop, it took about 3 minutes. Colab will be much more\n",
    "# variable in the time it takes, depending on what kind of instance\n",
    "# you are running.\n",
    "\n",
    "pred_df = category_assess(gen_df, recogs_model, \"obj_to_subj_proper\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6dfda72-8d80-4137-9f34-ec39d3e2de95",
   "metadata": {},
   "source": [
    "Extract the errors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "974d1a73-a08b-478f-b463-9c079836755b",
   "metadata": {},
   "outputs": [],
   "source": [
    "err_df = pred_df[pred_df.correct == False]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d38e4376-ce08-417c-abcd-d414cbc8d2c4",
   "metadata": {},
   "source": [
    "Use `find_name_roles` to get the role distribution in the error set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df67cec8-3e19-4141-93db-0b5e6c5380d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "err_roles = find_name_roles(err_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56f4fbcb-e9ad-4789-9c8c-f0d66af74426",
   "metadata": {},
   "outputs": [],
   "source": [
    "sorted(err_roles.items(), key=lambda x: len(x[1]))[: 3]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5ab221a-f50d-42af-a4ec-c59fcd7eef31",
   "metadata": {},
   "source": [
    "It's our old friend Charlie – in training, always a theme; in the generalization tests, always an agent."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7a19b17-f58a-4a40-a07a-71f362dd33d1",
   "metadata": {},
   "source": [
    "## Question 3: In-context learning with DSPy [2 points]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9489b4a-849d-4b00-835a-a5a6cc6f61d5",
   "metadata": {},
   "source": [
    "For this question, we are going to switch gears, from using our trained ReCOGS model to seeing whether we can get traction on this problem using only in-context learning. This question is meant to be very straightforward – our sole goal is to get you to the point where you have a working DSPy program that you can build on."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7e9a648-1535-4f31-945f-1488a5883a70",
   "metadata": {},
   "source": [
    "### Set-up"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99ae6c14-4aab-41c3-9d12-177c5de38a09",
   "metadata": {},
   "source": [
    "Standard set-up for DSPy, but we don't need a retriever:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4148343-7182-4e75-be98-f40481e8f4d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "import openai\n",
    "import os\n",
    "import dspy\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "root_path = '.'\n",
    "\n",
    "os.environ[\"DSP_NOTEBOOK_CACHEDIR\"] = os.path.join(root_path, 'cache')\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": "31b8932c-2916-4d35-b05f-533713414c0d",
   "metadata": {},
   "source": [
    "Our language model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d1d1d00-4589-45d3-86f3-d6c1b81e8eb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "lm = dspy.OpenAI(model='gpt-3.5-turbo', api_key=openai_key)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7338f9bd-6f4d-4131-bd35-efadd6e2a72e",
   "metadata": {},
   "source": [
    "DSPy settings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18e2cbc2-60bd-4893-8de0-5ef2ad3950f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "dspy.settings.configure(lm=lm)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "840dbd4b-7652-4bb8-bafb-c38bd752c10e",
   "metadata": {},
   "source": [
    "### Train examples in DSPy format"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80ceb9c5-183e-45f0-9f1b-3783b334053c",
   "metadata": {},
   "source": [
    "This will convert the train set into a list of `dspy.Example` instances to use for demonstrations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a145a809-6cc7-4711-aa94-850fd1383ea1",
   "metadata": {},
   "outputs": [],
   "source": [
    "dspy_recogs_train = [\n",
    "    dspy.Example(\n",
    "        sentence=row['input'], logical_form=row['output']\n",
    "    ).with_inputs(\"sentence\")\n",
    "    for _, row in dataset['train'].iterrows()]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ac35f2a-9b07-4b0c-bab4-1e4e022f5b03",
   "metadata": {},
   "source": [
    "### Basic signature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7860e77a-5799-4975-9b6a-50373aeb743f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RecogsSignature(dspy.Signature):\n",
    "    __doc__ = \"\"\"Translate English sentences into logical form.\"\"\"\n",
    "\n",
    "    sentence = dspy.InputField()\n",
    "    logical_form = dspy.OutputField()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d86488e1-25a8-4ae6-b34c-629685c35732",
   "metadata": {},
   "source": [
    "### Task 1: Basic Module"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6909ee5b-d473-4cf4-84ee-14b39b42226f",
   "metadata": {},
   "source": [
    "Your task is to complete `RecogsDspy` so that it uses `RecogsSignature`. This is simply meant to be a foundation for developing more complex systems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a5d99e5-1dca-44fb-b849-1f2c209f166d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class RecogsDspy(dspy.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        # Create an attribute called `self.predict`:\n",
    "        ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "    def forward(self, sentence):\n",
    "        \"\"\"Use the prediction later defined in the `init` method,\n",
    "        and return a `dspy.Prediction` instance with attribute\n",
    "        `logical_form`.\n",
    "        \"\"\"\n",
    "        pass\n",
    "        ##### YOUR CODE HERE\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88b771f7-8a89-4c29-ad3e-9c9a6e01dfa9",
   "metadata": {},
   "source": [
    "The following just tests that your `RecogsDspy` has the basic behavior we need for subsequent evaluations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5304dde4-c238-49cc-b2e2-bbc5b47f3497",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def test_RecogsDspy(classname):\n",
    "    model = classname()\n",
    "    pred = model(sentence='A rose was helped by a dog .')\n",
    "    if not hasattr(pred, \"logical_form\"):\n",
    "        print(f\"Error for {classname.__name__}: \"\n",
    "              \"predictions don't have an attribute `logical_form`.\")\n",
    "        return None\n",
    "    print(f\"No errors detected for `{classname.__name__}`\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4899af1-1de1-435c-82e4-97eb8920aa59",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "test_RecogsDspy(RecogsDspy)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b81f4a0d-c917-4944-93ec-44e7b7ac77c2",
   "metadata": {},
   "source": [
    "Basic usage; wihthout few-shot examples, the model cannot guess the required format:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb4a8329-ecb1-4080-8e50-faa933cd5662",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "recogs_dspy_model = RecogsDspy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6696a9e9-7401-4278-847f-0f018f0dd9ae",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dspy_recogs_train[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "624096e7-f943-4487-b0d7-cf0d8e9797ee",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "recogs_dspy_model(sentence=dspy_recogs_train[0].sentence)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa596d1b-61ed-4c5d-b3af-dceb1c0861f8",
   "metadata": {},
   "source": [
    "### Task 2: Few-shot usage with teleprompting"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8c27587-4c06-453d-9579-6429a762df8e",
   "metadata": {},
   "source": [
    "Your task is just to use `LabeledFewShot` to add few-shot demonstrations when using `RecogsDspy` or a model like it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11a45b47-e2ab-4a5d-a0b3-b5507c4d4f10",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dspy.teleprompt import LabeledFewShot\n",
    "\n",
    "def fewshot_recogs_dspy(model):\n",
    "    pass\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c77496c-dc04-4b09-85ac-b3113ba3dca5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def test_fewshot_recogs_dspy(func):\n",
    "    model = RecogsDspy()\n",
    "    compiled = func(model)\n",
    "    state = compiled.dump_state()\n",
    "    if not state['predict']['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": "45a0d78a-b67f-4896-96b1-b79f4d57772e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "test_fewshot_recogs_dspy(fewshot_recogs_dspy)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97cf7623-501f-400f-8db3-1da59674d16b",
   "metadata": {},
   "source": [
    "The demonstrations help the model make predictions in the correct general format:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4666f6fa-7c63-401e-9a54-42f93ae252cb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "compiled_recogs_dspy_model = fewshot_recogs_dspy(recogs_dspy_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c194f250-c6d9-4b9f-8a29-55910470d7ef",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "compiled_recogs_dspy_model(sentence=dspy_recogs_train[0].sentence)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0eaa937-771e-4202-88a7-c6e708a3d3ad",
   "metadata": {},
   "source": [
    "### Optional assessment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8ca366b-aa51-4d81-b400-c6c9500eb482",
   "metadata": {},
   "source": [
    "Here we sample 10 dev cases for a small evaluation. If you adapt this code, remember to use `recogs_exact_match` so that you aren't unfairly penalized for conjunct order or varible name differences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a20ae08-86ef-40da-a963-69a7333647ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "ssamp = dataset['dev'].sample(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "268a2271-8a65-4355-8885-6f2bc214ff91",
   "metadata": {},
   "outputs": [],
   "source": [
    "ssamp['prediction'] = ssamp.input.apply(\n",
    "    lambda x: compiled_recogs_dspy_model(sentence=x).logical_form)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9415ce6-a4e2-4b03-81dd-5ad85e0de78f",
   "metadata": {},
   "outputs": [],
   "source": [
    "ssamp['correct'] = ssamp.apply(\n",
    "    lambda row: recogs_exact_match(row['output'], row['prediction']), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa9e743e-5941-47dd-8953-cc3b302a314a",
   "metadata": {},
   "outputs": [],
   "source": [
    "ssamp['correct'].sum() / ssamp.shape[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5294287-5808-4eb7-ac2d-19bd0b660bd2",
   "metadata": {},
   "source": [
    "A random example to see what's going on:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "909fa032-f559-4971-94a8-16c7eba054d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "ssamp.sample(1).to_dict(orient='records')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5faf84b5-af29-4a55-8f83-6f71fdb1c054",
   "metadata": {},
   "source": [
    "## Question 4: Original System [3 points]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cfa81ab-c4a1-4f4e-9d50-d223f81c90a6",
   "metadata": {},
   "source": [
    "For your original system, you can do anything at all. The only constraint (repeated from above):\n",
    "\n",
    "__You cannot train your system on any examples from `dataset[\"gen\"]`, nor can the output representations from those examples be included in any prompts used for in-context learning.__\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": "b9b50392-e869-4315-bb2f-09b76b8000ae",
   "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",
    "##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# STOP COMMENT: Please do not remove this comment."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fcb4dab-c0a9-4b1c-a556-148d66b2a77a",
   "metadata": {},
   "source": [
    "Here are some potential paths – just a few of many options, though!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e20adf07-b7f3-4e6f-b172-73cfe781b7a4",
   "metadata": {},
   "source": [
    "### Option: DSPy program\n",
    "\n",
    "This could build on Question 3 very directly. All we have tried ourselves so far is the simple approach from that question."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da281ffa-ecf4-4598-bd1f-484153378a18",
   "metadata": {},
   "source": [
    "### Option: Further training of our model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8eb6de7-bab8-4f38-9567-49ea8d47dda7",
   "metadata": {},
   "source": [
    "This is very easy to do. For example, here we do some training on the first 10 dev examples, and we've exposed some keyword arguments that may be of interest:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76cf1b21-8059-4bc2-9cf8-a23b910e8b7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "recogs_ff = RecogsModel(\n",
    "    batch_size=5,\n",
    "    gradient_accumulation_steps=20,\n",
    "    max_iter=100, \n",
    "    early_stopping=True,\n",
    "    n_iter_no_change=10,\n",
    "    optimizer_class=torch.optim.Adam,\n",
    "    eta=0.00001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bd269d1-d835-4241-a5b6-0463883a42c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = recogs_ff.fit(dataset['dev'].input[: 40], dataset['dev'].output[: 40])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58af5ee4-d7e9-4700-8700-0021f11a1e1a",
   "metadata": {},
   "source": [
    "For this, you will want to pay a lot of attention to the optimization-related parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b64ec5e-56dd-4174-96ea-b9574a39ea58",
   "metadata": {},
   "source": [
    "### Option: Using a pretrained model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d71f245a-c474-4723-a1c8-ea4666c22ce5",
   "metadata": {},
   "source": [
    "The code used for Question 2 should make this very easy. For example, the following is the start of a complete solution using T5:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d75cf2f9-bedc-48e6-a367-747008002eed",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM\n",
    "\n",
    "class T5RecogsModule(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.encdec = AutoModelForSeq2SeqLM.from_pretrained(\"t5-small\")\n",
    "\n",
    "    def forward(self, X_pad, X_mask, y_pad, y_mask, labels=None):\n",
    "        outputs = self.encdec(\n",
    "            input_ids=X_pad, \n",
    "            attention_mask=X_mask,\n",
    "            decoder_attention_mask=y_mask,\n",
    "            labels=y_pad)\n",
    "        return outputs\n",
    "\n",
    "class T5RecogsModel(RecogsModel):\n",
    "    def __init__(self, *args, initialize=True, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.enc_tokenizer = AutoTokenizer.from_pretrained(\"t5-small\")\n",
    "        self.dec_tokenizer = self.enc_tokenizer\n",
    "\n",
    "    def build_graph(self):\n",
    "        return T5RecogsModule()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "613ca475-f788-441a-a32e-f3feed6aa88b",
   "metadata": {},
   "source": [
    "This will make predictions, but they will be pretty totally disconnected from our task right now:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eded1da8-e9ba-4cf8-a5ef-6d6c7c41ccf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "t5mod = T5RecogsModel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9074288-99ea-41d8-8430-7a6887957953",
   "metadata": {},
   "outputs": [],
   "source": [
    "t5_exs = dataset['dev'].input[: 2]\n",
    "\n",
    "t5_exs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a683db7-38de-49b0-b2d2-92cebeb1daa0",
   "metadata": {},
   "outputs": [],
   "source": [
    "t5mod.predict(t5_exs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e21f911-69d3-4fa1-b3eb-6c73b494a15c",
   "metadata": {},
   "source": [
    "This model needs to be fine-tuned on ReCOGS, which you can do with its `fit` method. In that case, you will want to pay a lot of attention to the optimization-related parameters to `TorchModelBase`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c15c129-8ada-4114-9c59-42ebc8025b59",
   "metadata": {},
   "source": [
    "### Option: Training a seq2seq model from scratch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "168e69f3-a5dd-4288-8d6c-ae652dac6365",
   "metadata": {},
   "source": [
    "The above code for T5 is easily adapted to use a randomly initialized model. The config files used to train our core model are `encoder_config.json` and `decoder_config.json` in `SRC_DIRNAME`. These might be a good starting point in terms of parameters and other set-up details."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5f6ee76-5119-4e89-ac12-42d2933bb2b0",
   "metadata": {},
   "source": [
    "### There are lots more options!\n",
    "\n",
    "Maybe a symbolic solver? A learned semantic parser? Tree-structured neural network?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "521ebabe-bfa7-4559-b7a6-e677b9241815",
   "metadata": {},
   "source": [
    "## Question 5: Bakeoff entry [1 point]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cab31552-ed3c-41cc-b187-d45ec408a3ba",
   "metadata": {},
   "source": [
    "Here we read in the bakeoff dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8f943f3-de8c-4f0c-924e-8c705be013b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "bakeoff_df = pd.read_csv(\n",
    "    os.path.join(SRC_DIRNAME, \"cs224u-recogs-test-unlabeled.tsv\"), \n",
    "    sep=\"\\t\", index_col=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "588b8d2c-169e-4fd5-a4d6-004d70b0b0d2",
   "metadata": {},
   "source": [
    "For the bakeoff entry, you should add a column \"prediction\" containing your predicted LFs and then use the following command to write the file to disk:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9eb3120b-a736-464a-a387-9a3bdc9bdcfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "bakeoff_df.to_csv(\"cs224u-recogs-bakeoff-entry.tsv\", sep=\"\\t\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "943bac05-b45d-4f0b-9ad8-e52b83b9760f",
   "metadata": {},
   "source": [
    "Here is what the first couple of lines of the file should look like:\n",
    "\n",
    "```\n",
    "\tinput\tcategory\tprediction\n",
    "0\tA cake was blessed by the wolf .\tactive_to_passive\tPREDICTED LF\n",
    "1\tA melon was blessed by a boy .\tactive_to_passive\tPREDICTED LF\n",
    "```\n",
    "\n",
    "where `PREDICTED LF` is what you predicted. Here is a quick test you can run locally to ensure that the autograder won't fail:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd5eec46-2508-4818-a9ff-1989f99bfdfb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_bakeoff_file(filename=\"cs224u-recogs-bakeoff-entry.tsv\"):\n",
    "    ref_filename = os.path.join(SRC_DIRNAME, \"cs224u-recogs-test-unlabeled.tsv\")\n",
    "    ref_df = pd.read_csv(ref_filename, sep=\"\\t\", index_col=0)\n",
    "\n",
    "    entry_df = pd.read_csv(filename, sep=\"\\t\", index_col=0)\n",
    "\n",
    "    errcount = 0\n",
    "\n",
    "    # Check expected columns:\n",
    "    expected_cols = [\"input\", \"category\", \"prediction\"]\n",
    "    for col in expected_cols:\n",
    "        if col not in entry_df.columns:\n",
    "            errcount += 1\n",
    "            print(f\"Missing column: {col}\")\n",
    "    if errcount > 0:\n",
    "        return\n",
    "\n",
    "    # Use the \"category\" column as a check that the rows have not\n",
    "    # been shuffled:\n",
    "    if not entry_df.category.equals(ref_df.category):\n",
    "        errcount += 1\n",
    "        print(\"Rows do not seem to be aligned with reference file. \"\n",
    "              \"Might they have gotten shuffled?\")\n",
    "\n",
    "    # Check that the predictions all have type str:\n",
    "    for line_num, x in enumerate(entry_df.prediction, start=1):\n",
    "        if not isinstance(x, str):\n",
    "            errcount += 1\n",
    "            print(f\"Prediction on line {line_num} is not a str: {x}\")\n",
    "\n",
    "    if errcount == 0:\n",
    "        print(\"Bakeoff file seems to be in good shape!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26d37fe9-6014-4fac-b9d6-5552a01ce72c",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_bakeoff_file(\"cs224u-recogs-bakeoff-entry.tsv\")"
   ]
  }
 ],
 "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.18"
  },
  "toc-showtags": false
 },
 "nbformat": 4,
 "nbformat_minor": 5
}