causal_reasoner.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "signal-tucson",
   "metadata": {},
   "source": [
    "# ROCK: Reasoning About Commonsense Causality"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "nasty-denmark",
   "metadata": {},
   "source": [
    "Use this notebook for performing CCR using ROCK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "aboriginal-giant",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function, absolute_import, division\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib widget\n",
    "\n",
    "import sys, os, json, time, datetime, logging, warnings, multiprocessing, itertools\n",
    "import tqdm, json, sqlite3, ast\n",
    "from pathlib import Path\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch\n",
    "import nltk, spacy\n",
    "import transformers, allennlp\n",
    "from transformers import (AutoTokenizer, AutoModelForMaskedLM,\n",
    "                          RobertaModel,RobertaForMaskedLM, \n",
    "                          RobertaTokenizer, GPT2LMHeadModel, GPT2Tokenizer)\n",
    "import allennlp_models\n",
    "import allennlp_models.pretrained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "behind-acoustic",
   "metadata": {},
   "outputs": [],
   "source": [
    "import src\n",
    "import src.pipeline\n",
    "import src.utils as utils\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "norwegian-amount",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "print(torch.cuda.is_available())\n",
    "TORCH_DEV = torch.device(f'cuda:0') if torch.cuda.is_available() \\\n",
    "                                    else torch.device(\"cpu\")\n",
    "\n",
    "logging.getLogger('allennlp.common.params').disabled = True \n",
    "logging.getLogger('allennlp.nn.initializers').disabled = True \n",
    "logging.getLogger('allennlp.modules.token_embedders.embedding').setLevel(logging.INFO) \n",
    "logging.getLogger('urllib3.connectionpool').disabled = True \n",
    "logging.getLogger().setLevel(logging.CRITICAL)\n",
    "warnings.filterwarnings('ignore')\n",
    "logging.disable(sys.maxsize)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "hungry-consensus",
   "metadata": {},
   "outputs": [],
   "source": [
    "def console_log(msg, end='\\n'):\n",
    "    os.write(1, ('[LOG/{}]'.format(multiprocessing.current_process().name)+msg+end).encode('utf-8'))\n",
    "\n",
    "\n",
    "def col_print(*args, cw=12, sep='|'):\n",
    "    print(f\" {sep} \".join(('{'+f\":<{cw}\"+'}').format(s) for s in args))\n",
    "    \n",
    "def set_seed(seed):\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    if torch.cuda.is_available():\n",
    "        torch.cuda.manual_seed_all(seed)\n",
    "    \n",
    "def set_ts_seed():\n",
    "    set_seed(int(str(datetime.datetime.now().timestamp()).replace('.', '')) % (2 ** 31))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "racial-respondent",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_PATH = Path(\"./exp_data\")\n",
    "MODEL_PATH = Path(\"./models\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "thick-injection",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set_seed(hsh('random_string') % (2 ** 31))\n",
    "set_ts_seed()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "binary-tribune",
   "metadata": {},
   "source": [
    "## Load Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "minimal-assets",
   "metadata": {},
   "outputs": [],
   "source": [
    "copa_dev = pd.read_json(DATA_PATH / \"copa_dev.json\", lines=True, orient='records').set_index('idx')\n",
    "copa_test = pd.read_csv(DATA_PATH / \"copa_test.json\")\n",
    "glt_d1 glt_d1 = pd.read_csv(DATA_PATH / \"glucose_d1_probs.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dated-monaco",
   "metadata": {},
   "source": [
    "# ROCK Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "inappropriate-savannah",
   "metadata": {},
   "outputs": [],
   "source": [
    "spacy_model = spacy.load('en_core_web_md')\n",
    "allensrl = src.pipeline.AllenSRLWrapper(allennlp_models.pretrained.load_predictor(\"structured-prediction-srl-bert\", cuda_device=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cardiac-martial",
   "metadata": {},
   "source": [
    "## Temporal Predictor"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "received-mailing",
   "metadata": {},
   "source": [
    "As a bare minimum, a customized temporal predictor needs to overwrite the `predict` method.\n",
    "Below is the implmentation we used based on mask language modeling.\n",
    "\n",
    "```python\n",
    "class TempPredictor:\n",
    "    def __init__(self, model, tokenizer, device, spacy_model=\"en_core_web_sm\"):\n",
    "        self._model = model\n",
    "        self._model.to(device)\n",
    "        self._model.eval()\n",
    "        self._tokenizer = tokenizer\n",
    "        self._mtoken = self._tokenizer.mask_token\n",
    "        self.unmasker = transformers.pipeline(\"fill-mask\", model=self._model, tokenizer=self._tokenizer, device=0)\n",
    "        try:\n",
    "            self._spacy = spacy.load(spacy_model)\n",
    "        except Exception as e:\n",
    "            self._spacy = spacy.load(\"en_core_web_sm\")\n",
    "            print(f\"Failed to load spacy model {spacy_model}, use default 'en_core_web_sm'\\n{e}\")\n",
    "\n",
    "\n",
    "    def predict(self, e1, e2, top_k=5):\n",
    "        \"\"\"\n",
    "        returns\n",
    "        \"\"\"\n",
    "        txt = self._remove_punct(e1) + \" \" + self._mtoken + \" \" + self._sent_lowercase(e2)\n",
    "        return self.unmasker(txt, top_k=top_k)\n",
    "\n",
    "\n",
    "    def get_temp(self, e1, e2, top_k=5, crop=1):\n",
    "        inst1 = self.predict(e1, e2, top_k)\n",
    "        inst2 = self.predict(e2, e1, top_k)\n",
    "\n",
    "        # e1 before e2\n",
    "        b1 = self._extract_token_prob(inst1, \"before\", crop=crop)\n",
    "        b2 = self._extract_token_prob(inst2, \"after\", crop=crop)\n",
    "\n",
    "        # e1 after e2\n",
    "        a1 = self._extract_token_prob(inst1, \"after\", crop=crop)\n",
    "        a2 = self._extract_token_prob(inst2, \"before\", crop=crop)\n",
    "\n",
    "        return (b1+b2)/2, (a1+a2)/2\n",
    "\n",
    "    def __call__(self, *args, **kwargs):\n",
    "        return self.get_temp(*args, **kwargs)\n",
    "    \n",
    "    # other methods omitted\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sitting-probability",
   "metadata": {},
   "source": [
    "**NB** Fine-tuned RoBERTa model checkpoint can be downloaded using [this anonymous Dropbox link](https://www.dropbox.com/s/9egrzn1ny3oq2qa/roberta_ft.tar.gz?dl=0) (1.29GB)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "unable-probability",
   "metadata": {},
   "outputs": [],
   "source": [
    "tp_roberta_ft = src.pipeline.TempPredictor(\n",
    "    model=RobertaForMaskedLM.from_pretrained(MODEL_PATH/\"roberta_ft\"),\n",
    "    tokenizer=RobertaTokenizer.from_pretrained(\"roberta-base\"),\n",
    "    device=TORCH_DEV\n",
    ")\n",
    "\n",
    "tp_roberta_base = src.pipeline.TempPredictor(\n",
    "    model=RobertaForMaskedLM.from_pretrained(\"roberta-base\"),\n",
    "    tokenizer=RobertaTokenizer.from_pretrained(\"roberta-base\"),\n",
    "    device=TORCH_DEV\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "affiliated-extent",
   "metadata": {},
   "source": [
    "##### Sanity Check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "encouraging-husband",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Premise: I doubted the salesman's pitch.\n",
      "C1: I turned his offer down.\n",
      "C2: He persuaded me to buy the product.\n",
      "Question: effect\tCorrect choice: Choice 1\n",
      "\n",
      "============== PREMISE <---> CHOICE 1 (expect before > after)\n",
      "I doubted the salesman's pitch. <---> I turned his offer down.\n",
      "Base model:\tbefore: 0.000\tafter: 0.000\n",
      "FT model:\tbefore: 0.486\tafter: 0.514\n",
      "\n",
      "============== PREMISE <---> CHOICE 2\n",
      "I doubted the salesman's pitch. <---> He persuaded me to buy the product.\n",
      "Base model:\tbefore: 0.015\tafter: 0.000\n",
      "FT model:\tbefore: 0.488\tafter: 0.511\n"
     ]
    }
   ],
   "source": [
    "utils.test_copa_run(copa_dev.iloc[5], tp_roberta_base, tp_roberta_ft, top_k=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "square-subscriber",
   "metadata": {},
   "source": [
    "## Event Sampler"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "announced-robinson",
   "metadata": {},
   "source": [
    "The event sampler subclasses `EventGenerator`.\n",
    "As a bare minimum, a custom implmentation should provide the `__call__` method.\n",
    "\n",
    "```python\n",
    "class EventGenerator:\n",
    "    def __init__(self, model, tokenizer, spacy_model, device):\n",
    "        self.model = model.to(device)\n",
    "        self.tokenizer = tokenizer\n",
    "        self.tokenizer.pad_token = self.tokenizer.eos_token\n",
    "        self.device = device\n",
    "\n",
    "    def __call__(self, prompt, max_length=30, **kwargs):\n",
    "        # pass\n",
    "```\n",
    "\n",
    "Below is our wrapper for GPT-J that is used in our paper:\n",
    "\n",
    "\n",
    "```python\n",
    "\n",
    "class GPTJGenerator:\n",
    "    def __init__(self, model, tokenizer, device=None):\n",
    "\n",
    "        self.model = model\n",
    "        \n",
    "        if device is not None:\n",
    "            self.model = self.model.to(device)\n",
    "\n",
    "        self.tokenizer = tokenizer\n",
    "        self.tokenizer.pad_token = self.tokenizer.eos_token\n",
    "        self.device = device\n",
    "\n",
    "    def __call__(self, prompt, **kwargs):\n",
    "        output_id = self.model.generate(self.tokenizer(prompt, return_tensors=\"pt\", padding=True).input_ids, **kwargs)\n",
    "        return self.tokenizer.batch_decode(output_id)\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "macro-meditation",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "gptj_tokenizer = AutoTokenizer.from_pretrained(\"EleutherAI/gpt-j-6B\")\n",
    "gptj_model = AutoModelForCausalLM.from_pretrained(\"EleutherAI/gpt-j-6B\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "about-colon",
   "metadata": {},
   "outputs": [],
   "source": [
    "gpt_generator = src.pipeline.GPTJGenerator(model=gptj_model, tokenizer=gptj_tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "optimum-lewis",
   "metadata": {},
   "outputs": [],
   "source": [
    "gen_kwargs = dict(max_length=30,\n",
    "                  do_sample=True,\n",
    "                  temperature=0.9,\n",
    "                  num_return_sequences=10,\n",
    "                 )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "checked-graphics",
   "metadata": {},
   "outputs": [],
   "source": [
    "gen_sents = gpt_generator(\"The man turned on the faucet.\", **gen_kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "three-johns",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The man turned on the faucet. His face was a mask of concentration, and his hands were steady as he washed the car.\n",
      "\n",
      "\n",
      "The man turned on the faucet. It was a small plastic one, nothing fancy. His hands shook, but he managed to turn the water\n",
      "The man turned on the faucet. He drank from it until the bottle was empty.\n",
      "\n",
      "After a moment he pulled the cap off the\n",
      "The man turned on the faucet. Water began splashing into the sink. He was washing the dishes, his expression blank, his eyes dead\n",
      "The man turned on the faucet.\n",
      "\n",
      "Bruno had to take a step back. The water made a sound. The man looked\n",
      "The man turned on the faucet. The water gushed out. A loud splash, and water went everywhere, flowing down the man's body\n",
      "The man turned on the faucet. He watched it fill with cold water. He heard it gurgle into the sink. He watched it\n",
      "The man turned on the faucet. He scrubbed down the sink and the tub, washed the toilet bowl, then flushed. In the meantime\n",
      "The man turned on the faucet. A strong current of water hit him in the face. He wiped it away and spat. The sink was\n",
      "The man turned on the faucet. He took off his tie and loosened his collar. No, not his collar; his shirt. He\n"
     ]
    }
   ],
   "source": [
    "print(\"\\n\".join(gen_sents))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "rural-london",
   "metadata": {},
   "source": [
    "## Interventions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "respected-samuel",
   "metadata": {},
   "source": [
    "As a bare minimum, the intervention generator should\n",
    "implement `__call__` method that takes a prompt and additional\n",
    "kwargs as arguments and return a list of interventions.\n",
    "\n",
    "```python\n",
    "class InterventionGenerator:\n",
    "    def __init__(self, **kwargs):\n",
    "        pass\n",
    "    \n",
    "    def __call__(self, prompt, **kwargs:)\n",
    "        pass\n",
    "\n",
    "```\n",
    "\n",
    "\n",
    "We use PolyJuice in our implementation based on their implementation [0].\n",
    "\n",
    "\n",
    "[0] https://github.com/tongshuangwu/polyjuice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "martial-legend",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "cf_gen = src.pipeline.PJGenerator(srl_processor=allensrl)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "fifth-munich",
   "metadata": {},
   "outputs": [],
   "source": [
    "cf_texts = cf_gen(\"The man turned off the faucet.\",\n",
    "      ctrl_codes=[\n",
    "          \"resemantic\", \n",
    "          \"negation\", \n",
    "          \"lexical\",\n",
    "          \"quantifier\"\n",
    "          \"insert\",\n",
    "          \"restructure\",\n",
    "          \"shuffle\",\n",
    "          \"delete\"\n",
    "                 ]\n",
    "                 )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "integrated-anime",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The woman turned off the faucet.\n",
      "The man turned off the faucet.\n",
      "The man poured water from the watering can into the pitcher until the watering can was empty. off the faucet.\n",
      "The man replaced the bell brand off the faucet.\n",
      "The man lit the cigarette off the faucet.\n",
      "The man turned off the water main.\n",
      "The man turned off the dishwasher.\n",
      "The man turned off the fan and replaced it with a fan outside of his house.\n"
     ]
    }
   ],
   "source": [
    "print(\"\\n\".join(cf_texts['resemantic']))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "scheduled-subscriber",
   "metadata": {},
   "source": [
    "## Processing Datasets for Evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "genuine-incident",
   "metadata": {},
   "source": [
    "### Construct DataFrames for processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "disciplinary-interface",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen_copa_proc_df(copa):\n",
    "    return pd.DataFrame(list(itertools.chain.from_iterable(\n",
    "        [[s[0], 'premise', s[1]['premise']]] if s[1]['question'] == 'effect'\n",
    "        else [[s[0], 'choice1', s[1]['choice1']], [s[0], 'choice2', s[1]['choice2']]]\n",
    "        for s in copa.iterrows()\n",
    "    )), columns=['index', 'name', 'text'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "blank-identification",
   "metadata": {},
   "outputs": [],
   "source": [
    "copa_proc = gen_copa_proc_df(copa_dev)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "rolled-persian",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>index</th>\n",
       "      <th>name</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>premise</td>\n",
       "      <td>The man turned on the faucet.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>premise</td>\n",
       "      <td>The girl found a bug in her cereal.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2</td>\n",
       "      <td>premise</td>\n",
       "      <td>The woman retired.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>premise</td>\n",
       "      <td>I wanted to conserve energy.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>4</td>\n",
       "      <td>choice1</td>\n",
       "      <td>The cook froze it.</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   index     name                                 text\n",
       "0      0  premise        The man turned on the faucet.\n",
       "1      1  premise  The girl found a bug in her cereal.\n",
       "2      2  premise                   The woman retired.\n",
       "3      3  premise         I wanted to conserve energy.\n",
       "4      4  choice1                   The cook froze it."
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "copa_proc.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "complimentary-dietary",
   "metadata": {},
   "source": [
    "We can apply the components row by row, but it is more efficient to let\n",
    "each component batch process the data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "every-surface",
   "metadata": {},
   "source": [
    "#### Sample Covariates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "structured-radar",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_cov(df, model, tokenizer):\n",
    "    output_ids = []\n",
    "    for s in df.iterrows():\n",
    "        prompt = f\"{s[1]['text']} Before that, \"\n",
    "\n",
    "        gen_tokens = model.generate(tokenizer(prompt,\n",
    "                          return_tensors=\"pt\", padding=True).input_ids, \n",
    "                    do_sample=True,\n",
    "                    temperature=0.9,\n",
    "                    max_length=30,\n",
    "                    num_return_sequences=100,\n",
    "            )\n",
    "        output_ids.append(gen_tokens)\n",
    "    return [tokenizer.batch_decode(tks) for tks in output_ids]\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "vertical-player",
   "metadata": {},
   "outputs": [],
   "source": [
    "copa_proc['covariates'] = sample_cov(copa_proc, gptj_model, gptj_tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "actual-semiconductor",
   "metadata": {},
   "source": [
    "#### Generating Interventions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "overall-brisbane",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_interventions(self, s, cf_gen, **kwargs):\n",
    "    interventions = self.cf_gen(s, gen_kwargs=kwargs)\n",
    "    intvers = list(itertools.chain(*[ints for _, ints in interventions.items()]))\n",
    "    self.last_gen['interventions'] = intvers\n",
    "    return intvers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "abandoned-charge",
   "metadata": {},
   "outputs": [],
   "source": [
    "copa_proc['interventions'] = copa_proc.apply(lambda s : get_interventions(s, cf_gen, \n",
    "                                ctrl_codes=[\n",
    "                                      \"resemantic\", \n",
    "                                      \"negation\", \n",
    "                                      \"lexical\",\n",
    "                                      \"quantifier\"\n",
    "                                      \"insert\",\n",
    "                                      \"restructure\",\n",
    "                                      \"shuffle\",\n",
    "                                      \"delete\"\n",
    "                                ], axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "green-bread",
   "metadata": {},
   "source": [
    "### Obtain Temporal Probabilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "paperback-video",
   "metadata": {},
   "outputs": [],
   "source": [
    "# use `utils.glt_get_probs`\n",
    "# if working on glucose-d1\n",
    "copa_proc = copa_proc.apply(lambda s : utils.copa_get_probs(s, model=tp_roberta_ft, \n",
    "                                                            top_k=5, spacy_model=spacy_model), \n",
    "                            axis=1)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "respiratory-setting",
   "metadata": {},
   "source": [
    "#### Add a few columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "moral-pressure",
   "metadata": {},
   "outputs": [],
   "source": [
    "def postproc_copa(df):\n",
    "    df['label_idx'] = df['name'].apply(\n",
    "        lambda s: -1 if s == 'premise' else int(s[-1])-1\n",
    "    )\n",
    "\n",
    "    df['outcome'] = df.apply(lambda s:\n",
    "        None if s['label_idx'] == -1 else copa_test.iloc[s['index']]['premise'], axis=1)\n",
    "\n",
    "\n",
    "    tmp_df = df[df['label_idx']==-1].copy()\n",
    "    tmp_df['label_idx'] = 1\n",
    "    df.loc[df['label_idx']==-1, 'label_idx']=0\n",
    "    df = pd.concat([df, tmp_df])\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "isolated-celebration",
   "metadata": {},
   "outputs": [],
   "source": [
    "copa_proc = postproc_copa(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "strange-spank",
   "metadata": {},
   "source": [
    "#### Save Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "supposed-flooring",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>index</th>\n",
       "      <th>name</th>\n",
       "      <th>text</th>\n",
       "      <th>covariates</th>\n",
       "      <th>interventions</th>\n",
       "      <th>outcome</th>\n",
       "      <th>label_idx</th>\n",
       "      <th>p_xd</th>\n",
       "      <th>p_dy</th>\n",
       "      <th>p_xy</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>premise</td>\n",
       "      <td>The man turned on the faucet.</td>\n",
       "      <td>[\"The man turned on the faucet. Before that, h...</td>\n",
       "      <td>['The man chose to turned on the faucet.', 'At...</td>\n",
       "      <td>The toilet filled with water.</td>\n",
       "      <td>0</td>\n",
       "      <td>[[(0.210379958152771, 0.30398909747600555, 0.0...</td>\n",
       "      <td>[(0.45566828548908234, 0.4977114349603653), (0...</td>\n",
       "      <td>[(0.04539947956800461, 0.03326283395290375), (...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0</td>\n",
       "      <td>premise</td>\n",
       "      <td>The man turned on the faucet.</td>\n",
       "      <td>[\"The man turned on the faucet. Before that, h...</td>\n",
       "      <td>['The man chose to turned on the faucet.', 'At...</td>\n",
       "      <td>Water flowed from the spout.</td>\n",
       "      <td>1</td>\n",
       "      <td>[[(0.210379958152771, 0.30398909747600555, 0.0...</td>\n",
       "      <td>[(0.5247508734464645, 0.4717450588941574), (0....</td>\n",
       "      <td>[(0.3355148509144783, 0.25840914994478226), (0...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>premise</td>\n",
       "      <td>The girl found a bug in her cereal.</td>\n",
       "      <td>[\"The girl found a bug in her cereal. Before t...</td>\n",
       "      <td>['While digging for well water Monica stayed o...</td>\n",
       "      <td>She poured milk in the bowl.</td>\n",
       "      <td>0</td>\n",
       "      <td>[[(0.5899830460548401, 0.39629100263118744, 0....</td>\n",
       "      <td>[(0.47232694923877716, 0.5270598828792572), (0...</td>\n",
       "      <td>[(0.5464454889297485, 0.40834908187389374), (0...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>premise</td>\n",
       "      <td>The girl found a bug in her cereal.</td>\n",
       "      <td>[\"The girl found a bug in her cereal. Before t...</td>\n",
       "      <td>['While digging for well water Monica stayed o...</td>\n",
       "      <td>She lost her appetite.</td>\n",
       "      <td>1</td>\n",
       "      <td>[[(0.5899830460548401, 0.39629100263118744, 0....</td>\n",
       "      <td>[(0.4848470687866211, 0.5141351819038391), (0....</td>\n",
       "      <td>[(0.6019861996173859, 0.38734038174152374), (0...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2</td>\n",
       "      <td>premise</td>\n",
       "      <td>The woman retired.</td>\n",
       "      <td>[\"The woman retired. Before that, she'd writte...</td>\n",
       "      <td>['Suddenly the woman retired.', 'The author ra...</td>\n",
       "      <td>She received her pension.</td>\n",
       "      <td>0</td>\n",
       "      <td>[[(0.3462740257382393, 0.30608220398426056, 0....</td>\n",
       "      <td>[(0.4986240118741989, 0.49911610782146454), (0...</td>\n",
       "      <td>[(0.5325719714164734, 0.45344893634319305), (0...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   index     name                                 text  \\\n",
       "0      0  premise        The man turned on the faucet.   \n",
       "1      0  premise        The man turned on the faucet.   \n",
       "2      1  premise  The girl found a bug in her cereal.   \n",
       "3      1  premise  The girl found a bug in her cereal.   \n",
       "4      2  premise                   The woman retired.   \n",
       "\n",
       "                                          covariates  \\\n",
       "0  [\"The man turned on the faucet. Before that, h...   \n",
       "1  [\"The man turned on the faucet. Before that, h...   \n",
       "2  [\"The girl found a bug in her cereal. Before t...   \n",
       "3  [\"The girl found a bug in her cereal. Before t...   \n",
       "4  [\"The woman retired. Before that, she'd writte...   \n",
       "\n",
       "                                       interventions  \\\n",
       "0  ['The man chose to turned on the faucet.', 'At...   \n",
       "1  ['The man chose to turned on the faucet.', 'At...   \n",
       "2  ['While digging for well water Monica stayed o...   \n",
       "3  ['While digging for well water Monica stayed o...   \n",
       "4  ['Suddenly the woman retired.', 'The author ra...   \n",
       "\n",
       "                         outcome  label_idx  \\\n",
       "0  The toilet filled with water.          0   \n",
       "1   Water flowed from the spout.          1   \n",
       "2   She poured milk in the bowl.          0   \n",
       "3         She lost her appetite.          1   \n",
       "4      She received her pension.          0   \n",
       "\n",
       "                                                p_xd  \\\n",
       "0  [[(0.210379958152771, 0.30398909747600555, 0.0...   \n",
       "1  [[(0.210379958152771, 0.30398909747600555, 0.0...   \n",
       "2  [[(0.5899830460548401, 0.39629100263118744, 0....   \n",
       "3  [[(0.5899830460548401, 0.39629100263118744, 0....   \n",
       "4  [[(0.3462740257382393, 0.30608220398426056, 0....   \n",
       "\n",
       "                                                p_dy  \\\n",
       "0  [(0.45566828548908234, 0.4977114349603653), (0...   \n",
       "1  [(0.5247508734464645, 0.4717450588941574), (0....   \n",
       "2  [(0.47232694923877716, 0.5270598828792572), (0...   \n",
       "3  [(0.4848470687866211, 0.5141351819038391), (0....   \n",
       "4  [(0.4986240118741989, 0.49911610782146454), (0...   \n",
       "\n",
       "                                                p_xy  \n",
       "0  [(0.04539947956800461, 0.03326283395290375), (...  \n",
       "1  [(0.3355148509144783, 0.25840914994478226), (0...  \n",
       "2  [(0.5464454889297485, 0.40834908187389374), (0...  \n",
       "3  [(0.6019861996173859, 0.38734038174152374), (0...  \n",
       "4  [(0.5325719714164734, 0.45344893634319305), (0...  "
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "copa_proc.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "visible-flavor",
   "metadata": {},
   "outputs": [],
   "source": [
    "copa_proc.to_csv(DATA_PATH/\"copa_dev_probs.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "million-strap",
   "metadata": {},
   "source": [
    "## Next Steps\n",
    "\n",
    "Please see `result_presentation.ipynb` notebook for evaulation."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py3.8-torch",
   "language": "python",
   "name": "py3.8-torch"
  },
  "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}