data_generation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Generation \n",
    "\n",
    "In this notebook, we generate our dataset from scratch. \n",
    "\n",
    "We follow the steps in the figure. \n",
    "\n",
    "![title](figures/data_generation.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Expert-defined ground-truth Bayesian network "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.data import DiseaseModelGT\n",
    "\n",
    "data_model = DiseaseModelGT()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------+-----+\n",
      "| season(cold) | 0.4 |\n",
      "+--------------+-----+\n",
      "| season(warm) | 0.6 |\n",
      "+--------------+-----+\n",
      "+-----------+--------------+--------------+\n",
      "| season    | season(cold) | season(warm) |\n",
      "+-----------+--------------+--------------+\n",
      "| pneu(yes) | 0.015        | 0.005        |\n",
      "+-----------+--------------+--------------+\n",
      "| pneu(no)  | 0.985        | 0.995        |\n",
      "+-----------+--------------+--------------+\n",
      "+----------+--------------+--------------+\n",
      "| season   | season(cold) | season(warm) |\n",
      "+----------+--------------+--------------+\n",
      "| inf(yes) | 0.5          | 0.05         |\n",
      "+----------+--------------+--------------+\n",
      "| inf(no)  | 0.5          | 0.95         |\n",
      "+----------+--------------+--------------+\n",
      "+-----------+-----------+----------+\n",
      "| pneu      | pneu(yes) | pneu(no) |\n",
      "+-----------+-----------+----------+\n",
      "| dysp(yes) | 0.3       | 0.15     |\n",
      "+-----------+-----------+----------+\n",
      "| dysp(no)  | 0.7       | 0.85     |\n",
      "+-----------+-----------+----------+\n",
      "+------------+----------+---------+\n",
      "| inf        | inf(yes) | inf(no) |\n",
      "+------------+----------+---------+\n",
      "| nasal(yes) | 0.7      | 0.2     |\n",
      "+------------+----------+---------+\n",
      "| nasal(no)  | 0.3      | 0.8     |\n",
      "+------------+----------+---------+\n",
      "+------------+-----------+-----------+----------+----------+\n",
      "| pneu       | pneu(yes) | pneu(yes) | pneu(no) | pneu(no) |\n",
      "+------------+-----------+-----------+----------+----------+\n",
      "| inf        | inf(yes)  | inf(no)   | inf(yes) | inf(no)  |\n",
      "+------------+-----------+-----------+----------+----------+\n",
      "| cough(yes) | 0.9       | 0.9       | 0.8      | 0.05     |\n",
      "+------------+-----------+-----------+----------+----------+\n",
      "| cough(no)  | 0.1       | 0.1       | 0.2      | 0.95     |\n",
      "+------------+-----------+-----------+----------+----------+\n",
      "+-------------+-----------+-----------+----------+----------+\n",
      "| pneu        | pneu(yes) | pneu(yes) | pneu(no) | pneu(no) |\n",
      "+-------------+-----------+-----------+----------+----------+\n",
      "| inf         | inf(yes)  | inf(no)   | inf(yes) | inf(no)  |\n",
      "+-------------+-----------+-----------+----------+----------+\n",
      "| fever(none) | 0.05      | 0.1       | 0.85     | 0.8      |\n",
      "+-------------+-----------+-----------+----------+----------+\n",
      "| fever(low)  | 0.15      | 0.1       | 0.14     | 0.15     |\n",
      "+-------------+-----------+-----------+----------+----------+\n",
      "| fever(high) | 0.8       | 0.8       | 0.01     | 0.05     |\n",
      "+-------------+-----------+-----------+----------+----------+\n",
      "+-----------+-----------+-----------+----------+----------+\n",
      "| pneu      | pneu(yes) | pneu(yes) | pneu(no) | pneu(no) |\n",
      "+-----------+-----------+-----------+----------+----------+\n",
      "| inf       | inf(yes)  | inf(no)   | inf(yes) | inf(no)  |\n",
      "+-----------+-----------+-----------+----------+----------+\n",
      "| pain(yes) | 0.3       | 0.3       | 0.1      | 0.05     |\n",
      "+-----------+-----------+-----------+----------+----------+\n",
      "| pain(no)  | 0.7       | 0.7       | 0.9      | 0.95     |\n",
      "+-----------+-----------+-----------+----------+----------+\n"
     ]
    }
   ],
   "source": [
    "print(data_model.model.get_cpds(\"season\"))\n",
    "print(data_model.model.get_cpds(\"pneu\"))\n",
    "print(data_model.model.get_cpds(\"inf\"))\n",
    "print(data_model.model.get_cpds(\"dysp\"))\n",
    "print(data_model.model.get_cpds(\"nasal\"))\n",
    "print(data_model.model.get_cpds(\"cough\"))\n",
    "print(data_model.model.get_cpds(\"fever\"))\n",
    "print(data_model.model.get_cpds(\"pain\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Sample tabular variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating for node: pain: 100%|██████████| 8/8 [00:00<00:00, 129.99it/s]\n"
     ]
    }
   ],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "N = 5000\n",
    "seed = 547\n",
    "\n",
    "samples = data_model.sample(N, seed=seed)\n",
    "train, test = train_test_split(samples, train_size=0.8, random_state=seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Prompt large language model to create clinical notes\n",
    "\n",
    "This part is not fully reproducible, since ChatGPT will provide different responses every time it's prompted. However, we still include a demonstration of how we obtained the text descriptions for clarity. \n",
    "\n",
    "Only 2/3 of the training set needs an associated text, so first we randomly split off this portion of the training samples. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "seed = 547\n",
    "train_text, _ = train_test_split(train, train_size=2667, random_state=seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We generate all the symptom combinations and the number of texts we need to generate for each. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_count_df(df):\n",
    "    prompt_data = df[[\"dysp\", \"cough\", \"fever\", \"pain\", \"nasal\"]]\n",
    "    prompt_data = prompt_data.rename(columns={\"dysp\": \"dyspnea\", \"pain\": \"chest pain / pain attributed to airways\", \"nasal\": \"sneezing / blocked nose\"})\n",
    "    column_names = prompt_data.columns.tolist()\n",
    "    prompt_data[\"combo\"] = prompt_data.agg(\",\".join, axis=1)\n",
    "    combo_counts = prompt_data[\"combo\"].value_counts().reset_index()\n",
    "    combo_counts[column_names] = combo_counts[\"index\"].str.split(\",\", expand=True)\n",
    "    count_df = combo_counts.drop(columns=[\"index\"])\n",
    "    count_df = count_df.rename(columns={\"combo\": \"count\"})\n",
    "    return count_df\n",
    "\n",
    "count_combinations = create_count_df(train_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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>count</th>\n",
       "      <th>dyspnea</th>\n",
       "      <th>cough</th>\n",
       "      <th>fever</th>\n",
       "      <th>chest pain / pain attributed to airways</th>\n",
       "      <th>sneezing / blocked nose</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1032</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>none</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>317</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>none</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>229</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>none</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>186</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>low</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>150</td>\n",
       "      <td>yes</td>\n",
       "      <td>no</td>\n",
       "      <td>none</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   count  ... sneezing / blocked nose\n",
       "0   1032  ...                      no\n",
       "1    317  ...                     yes\n",
       "2    229  ...                     yes\n",
       "3    186  ...                      no\n",
       "4    150  ...                      no\n",
       "\n",
       "[5 rows x 6 columns]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "count_combinations.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To generate our data, we used 20 example notes on real patient encounters from a GP's clinical practice. For privacy reasons, we cannot publish this data. To illustrate our prompting method, we create some fake examples that stand in for the real patient encounter notes. Every note is annotated with the symptoms experienced by this patient. Not all symptoms are necessary described in the text. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "dummy_example_notes = [\n",
    "                       {\"dyspnea\": \"no\", \"cough\": \"no\", \"fever\": \"none\", \"chest pain / pain attributed to airways\": \"no\", \"sneezing / blocked nose\": \"no\",\n",
    "                       \"text\": \"The patient presents with shoulder pain. They have trouble raising their right arm all the way. This has been going on for a week.\"},\n",
    "                       {\"dyspnea\": \"no\", \"cough\": \"no\", \"fever\": \"none\", \"chest pain / pain attributed to airways\": \"no\", \"sneezing / blocked nose\": \"no\",\n",
    "                        \"text\": \"Patient presents with long-term fatigue that is not solved by a good night's sleep. They also show a loss of apetite\"},\n",
    "                       {\"dyspnea\": \"no\", \"cough\": \"yes\", \"fever\": \"none\", \"chest pain / pain attributed to airways\": \"no\", \"sneezing / blocked nose\": \"yes\",\n",
    "                        \"text\": \"Young child presents with a heavy cough with phlegm. He shows no signs of fever, though has been sneezing for the past couple of days.\"},\n",
    "                       {\"dyspnea\": \"yes\", \"cough\": \"yes\", \"fever\": \"none\", \"chest pain / pain attributed to airways\": \"yes\", \"sneezing / blocked nose\": \"no\",\n",
    "                        \"text\": \"Patient experiences trouble breathing which intensifies during exercise. Sometimes they have to pause when walking up the stairs. \\\n",
    "                        They also worry about their painful scraping cough.\"},\n",
    "                       {\"dyspnea\": \"no\", \"cough\": \"no\", \"fever\": \"none\", \"chest pain / pain attributed to airways\": \"yes\", \"sneezing / blocked nose\": \"yes\",\n",
    "                        \"text\": \"Young man feels a stabbing pain in the chest when taking deep breaths. He has been experiencing cold symptoms, such as a stuffy nose, for 2 days now.\"},\n",
    "                       {\"dyspnea\": \"no\", \"cough\": \"no\", \"fever\": \"high\", \"chest pain / pain attributed to airways\": \"no\", \"sneezing / blocked nose\": \"no\",\n",
    "                        \"text\": \"Newborn is sufferring from reduced appetite and high fever of up to 40°C. Mother is growing increasingly worried.\"}\n",
    "                       ]\n",
    "df_examples = pd.DataFrame(dummy_example_notes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "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>dyspnea</th>\n",
       "      <th>cough</th>\n",
       "      <th>fever</th>\n",
       "      <th>chest pain / pain attributed to airways</th>\n",
       "      <th>sneezing / blocked nose</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>none</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>The patient presents with shoulder pain. They ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>none</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>Patient presents with long-term fatigue that i...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>none</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>Young child presents with a heavy cough with p...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>yes</td>\n",
       "      <td>yes</td>\n",
       "      <td>none</td>\n",
       "      <td>yes</td>\n",
       "      <td>no</td>\n",
       "      <td>Patient experiences trouble breathing which in...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>none</td>\n",
       "      <td>yes</td>\n",
       "      <td>yes</td>\n",
       "      <td>Young man feels a stabbing pain in the chest w...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  dyspnea cough fever chest pain / pain attributed to airways  \\\n",
       "0      no    no  none                                      no   \n",
       "1      no    no  none                                      no   \n",
       "2      no   yes  none                                      no   \n",
       "3     yes   yes  none                                     yes   \n",
       "4      no    no  none                                     yes   \n",
       "\n",
       "  sneezing / blocked nose                                               text  \n",
       "0                      no  The patient presents with shoulder pain. They ...  \n",
       "1                      no  Patient presents with long-term fatigue that i...  \n",
       "2                     yes  Young child presents with a heavy cough with p...  \n",
       "3                      no  Patient experiences trouble breathing which in...  \n",
       "4                     yes  Young man feels a stabbing pain in the chest w...  "
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_examples.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The most commonly occurring symptom combination is the one where all symptoms are absent from the patient (all features are \"no\"/\"none\"). We first generate clinical text notes for all *other* combinations, then return to the most occurring one. \n",
    "\n",
    "We use the following two templates, of which we provide an example in the OpenAI playground: \n",
    "- `prompt_known_combo`: The symptom combination is available in the list of examples: https://platform.openai.com/playground/p/poSdvoy9dipYIwVPXepRrUyL?model=gpt-3.5-turbo&mode=chat \n",
    "- `prompt_unknown_combo`: The symptom combination is not available in the list of examples: https://platform.openai.com/playground/p/6KOm6pP6DXmMwDxUJWdGGHP1?model=gpt-3.5-turbo&mode=chat \n",
    "\n",
    "The code for creating these templates for an arbitrary symptom combination is provided in utils/prompt.py. \n",
    "\n",
    "Below, we iterate over all symptom combinations, except the first one, and generate the necessary number of clinical texts. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "import random \n",
    "\n",
    "from utils.prompt import prompt_known_combo, prompt_unknown_combo, get_random_examples\n",
    "\n",
    "openai.api_key = \"fill-in-your-openai-key\"\n",
    "openai.organization = \"fill-in-your-organization\"\n",
    "\n",
    "count_combinations[\"text\"] = \"\"\n",
    "max_req = 5\n",
    "temp = 1\n",
    "freq_pen = 0.5\n",
    "n_rand_ex = 2\n",
    "debug = False\n",
    "random.seed(547)\n",
    "\n",
    "# iterate over all possible symptom combinations\n",
    "for i in range(1, len(count_combinations)): # skip first one (all symptoms not present, gets special prompt)\n",
    "\n",
    "    # request set of symptoms\n",
    "    req_combo = count_combinations.iloc[i]\n",
    "    n = req_combo[\"count\"]\n",
    "    req_symptoms = req_combo.drop(labels=[\"count\", \"text\"]).to_dict()\n",
    "\n",
    "    # get symptom combination from examples \n",
    "    sel_examples = df_examples\n",
    "    for symp, val in req_symptoms.items(): \n",
    "        sel_examples = sel_examples[sel_examples[symp] == val]\n",
    "    examples = sel_examples[\"text\"].to_list()\n",
    "    \n",
    "    # if example exists, prompt known combo \n",
    "    if len(examples) > 0:\n",
    "        response = []\n",
    "        if n > max_req:\n",
    "            mask_neg = False # mask out the negative symptoms in the prompt half of the time\n",
    "            for _ in range(n//max_req): # process request in groups of max_req\n",
    "                messages, resp, full_resp = prompt_known_combo(req_symptoms, examples, max_req, mask_neg, temp, freq_pen, debug=debug) \n",
    "                response += resp\n",
    "                mask_neg = True if mask_neg == False else False\n",
    "            n_remaining = n%max_req\n",
    "            if n_remaining > 0:\n",
    "                messages, resp, full_resp = prompt_known_combo(req_symptoms, examples, n_remaining, mask_neg, temp, freq_pen, debug=debug) \n",
    "                response += resp\n",
    "        else: \n",
    "            mask_neg = False # don't mask out negative symptoms when only using one prompt\n",
    "            messages, resp, full_resp = prompt_known_combo(req_symptoms, examples, n, mask_neg, temp, freq_pen, debug=debug)\n",
    "            response += resp\n",
    "\n",
    "        count_combinations.at[i, \"text\"] = response\n",
    "    \n",
    "    # if no example exists, prompt unknown combo\n",
    "    else: \n",
    "\n",
    "        response = []\n",
    "        if n > max_req:\n",
    "            \n",
    "            mask_neg = False # mask out the negative symptoms in the prompt half of the time\n",
    "            for _ in range(n//max_req): # process request in groups of max_req\n",
    "                examples = get_random_examples(df_examples, n_rand_ex)\n",
    "                messages, resp, full_resp = prompt_unknown_combo(req_symptoms, examples, max_req, mask_neg, temp, freq_pen, debug=debug) \n",
    "                response += resp\n",
    "                mask_neg = True if mask_neg == False else False\n",
    "            n_remaining = n%max_req\n",
    "            if n_remaining > 0:\n",
    "                examples = get_random_examples(df_examples, n_rand_ex)\n",
    "                messages, resp, full_resp = prompt_unknown_combo(req_symptoms, examples, n_remaining, mask_neg, temp, freq_pen, debug=debug) \n",
    "                response += resp\n",
    "        else: \n",
    "            mask_neg = False # don't mask out negative symptoms when only using one prompt\n",
    "            examples = get_random_examples(df_examples, n_rand_ex)\n",
    "            messages, resp, full_resp = prompt_unknown_combo(req_symptoms, examples, n, mask_neg, temp, freq_pen, debug=debug)\n",
    "            response += resp\n",
    "\n",
    "        count_combinations.at[i, \"text\"] = response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Afterwards, we checked whether the appropriate number of texts were generated for each symptom combination and checked for obvious mistakes. \n",
    "\n",
    "Since we need a large number of examples for the symptom combination where none of the symptoms are observed, we use a variety of prompting strategies to get sufficient variation in the generated clinical notes. The following templates are used as prompting strategies. We mention the amount of clinical notes we obtain this way for both the train and test set. \n",
    "1. `prompt_unrelated_strat13`: Prompt with 2 random in-context examples from the set of clinical note examples where all symptoms are \"no\": https://platform.openai.com/playground/p/Zw9y4EZ8RRfGaZTgBCPu5DPT?model=gpt-3.5-turbo&mode=chat (400 in train set, 151 in test set)\n",
    "2. `prompt_unknown_combo` (from before): Prompt with 2 random out-of-context examples from the full set of clinical note examples, where symptom combination can be anything: https://platform.openai.com/playground/p/jYJeFfzXVZ5akndHGkB6cmPk?model=gpt-3.5-turbo&mode=chat (50 in train set, 19 in test set)\n",
    "3. `prompt_unrelated_strat13`: Prompt with 2 random in-context examples from the set of clinical note examples where all symptoms are \"no\", but don't mention in the prompt that these symptoms are absent. Afterwards, we filter for mentions of our undesired symptoms manually: https://platform.openai.com/playground/p/Ikszv18Zbqr162kF0iSErbCT?model=gpt-3.5-turbo&mode=chat (107 in train set, 39 in test set)\n",
    "4. `prompt_unrelated_strat45`: Prompt without any examples, without mentioning the symptoms that should not be present. Afterwards, we filter for mentions of our undesired symptoms manually: https://platform.openai.com/playground/p/8KaVpr7chxLMHyeKz6Y2mMVE?model=gpt-3.5-turbo&mode=chat (111 in train set, 43 in test set)\n",
    "5. `prompt_unrelated_strat45`: Prompt without any examples, but mention the symptoms that should not be present: https://platform.openai.com/playground/p/eXDqpwMkqcvUA8wWiW1H06et?model=gpt-3.5-turbo&mode=chat (364 in train set, 136 in test set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.prompt import prompt_unrelated_strat13\n",
    "\n",
    "# Strategy 1\n",
    "temp = 1 \n",
    "freq_pen = 0.5\n",
    "mention_sympt = True\n",
    "random.seed(547)\n",
    "n_req = 10\n",
    "n_ex = 2\n",
    "df_examples_no = df_examples.iloc[0:4]\n",
    "responses = []\n",
    "full_resp = []\n",
    "\n",
    "for i in range(40):\n",
    "    messages, resp_split, response = prompt_unrelated_strat13(n_req, df_examples_no, n_ex, mention_sympt, temp, freq_pen, debug=False)\n",
    "    full_resp.append(response)\n",
    "    responses += resp_split\n",
    "    print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Strategy 2\n",
    "\n",
    "n_rand_ex = 2\n",
    "n_req = 10\n",
    "mask_neg = False\n",
    "temp = 1\n",
    "freq_pen = 0.5\n",
    "\n",
    "req_combo = count_combinations.iloc[0]\n",
    "req_symptoms = req_combo.drop(labels=[\"count\", \"text\"]).to_dict()\n",
    "\n",
    "responses = []\n",
    "full_resp = []\n",
    "\n",
    "for i in range(5):\n",
    "    examples = get_random_examples(df_examples, n_rand_ex)\n",
    "    messages, resp_split, response = prompt_unknown_combo(req_symptoms, examples, n_req, mask_neg, temp, freq_pen, debug=False)\n",
    "    full_resp.append(response)\n",
    "    responses += resp_split\n",
    "    print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Strategy 3\n",
    "temp = 1 \n",
    "freq_pen = 0.5\n",
    "mention_sympt = False\n",
    "random.seed(547)\n",
    "n_req = 10\n",
    "n_ex = 2\n",
    "df_examples_no = df_examples.iloc[0:4]\n",
    "responses = []\n",
    "full_resp = []\n",
    "\n",
    "for i in range(11):\n",
    "    messages, resp_split, response = prompt_unrelated_strat13(n_req, df_examples_no, n_ex, mention_sympt, temp, freq_pen, debug=False)\n",
    "    full_resp.append(response)\n",
    "    responses += resp_split\n",
    "    print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.prompt import prompt_unrelated_strat45\n",
    "\n",
    "# Strategy 4\n",
    "temp = 1 \n",
    "freq_pen = 0.5\n",
    "mention_sympt = False\n",
    "responses = []\n",
    "full_resp = []\n",
    "for i in range(12):\n",
    "    messages, resp_split, response = prompt_unrelated_strat45(10, temp, freq_pen, mention_sympt, debug=False)\n",
    "    full_resp.append(response)\n",
    "    responses += resp_split\n",
    "    print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Strategy 5\n",
    "temp = 1 \n",
    "freq_pen = 0.5\n",
    "mention_sympt = True\n",
    "responses = []\n",
    "full_resp = []\n",
    "for i in range(35):\n",
    "    messages, resp_split, response = prompt_unrelated_strat45(10, temp, freq_pen, mention_sympt, debug=False)\n",
    "    full_resp.append(response)\n",
    "    responses += resp_split\n",
    "    print(i)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After generating the text according to the strategies outlined above, we end up with a dataframe where every symptom combination is associated with a list of n possible text descriptions, with n being the number of times this combination occurs in the training set. \n",
    "\n",
    "We then added these clinical notes to the 2/3 of the training set in need of text. We repeat the same process for the test set. \n",
    "\n",
    "The resulting dataframes are loaded in below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle \n",
    "with open(\"data/suppl/train_2667_text.p\", \"rb\") as file: \n",
    "    train_text = pickle.load(file)\n",
    "with open(\"data/suppl/train_1333_notext.p\", \"rb\") as file: \n",
    "    train_no_text = pickle.load(file)\n",
    "with open(\"data/suppl/test_1000_text.p\", \"rb\") as file: \n",
    "    test = pickle.load(file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "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>season</th>\n",
       "      <th>pneu</th>\n",
       "      <th>inf</th>\n",
       "      <th>dysp</th>\n",
       "      <th>cough</th>\n",
       "      <th>fever</th>\n",
       "      <th>pain</th>\n",
       "      <th>nasal</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>warm</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>high</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>Patient presenting with a high fever (39.4°C) ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>warm</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>high</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>High-grade fever (39.6°C) persisting for two d...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>warm</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>high</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>Patient complaining of a high fever (38.7°C) t...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>warm</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>high</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>Fever of 39°C reported since yesterday with as...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>warm</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>high</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>Significant fever (40°C) for the past 24 hours...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  season pneu inf  ... pain nasal                                               text\n",
       "0   warm   no  no  ...   no    no  Patient presenting with a high fever (39.4°C) ...\n",
       "1   warm   no  no  ...   no    no  High-grade fever (39.6°C) persisting for two d...\n",
       "2   warm   no  no  ...   no    no  Patient complaining of a high fever (38.7°C) t...\n",
       "3   warm   no  no  ...   no    no  Fever of 39°C reported since yesterday with as...\n",
       "4   warm   no  no  ...   no    no  Significant fever (40°C) for the past 24 hours...\n",
       "\n",
       "[5 rows x 9 columns]"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_text.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "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>season</th>\n",
       "      <th>pneu</th>\n",
       "      <th>inf</th>\n",
       "      <th>dysp</th>\n",
       "      <th>cough</th>\n",
       "      <th>fever</th>\n",
       "      <th>pain</th>\n",
       "      <th>nasal</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>4504</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>none</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2491</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>none</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>825</th>\n",
       "      <td>warm</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>low</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2766</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>low</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2681</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>yes</td>\n",
       "      <td>yes</td>\n",
       "      <td>none</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     season pneu  inf dysp cough fever pain nasal text\n",
       "4504   cold   no  yes   no   yes  none   no   yes     \n",
       "2491   cold   no  yes   no   yes  none   no   yes     \n",
       "825    warm   no   no   no    no   low   no    no     \n",
       "2766   cold   no   no   no    no   low   no    no     \n",
       "2681   cold   no  yes  yes   yes  none   no   yes     "
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_no_text.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2667, 1333)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(train_text), len(train_no_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "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>season</th>\n",
       "      <th>pneu</th>\n",
       "      <th>inf</th>\n",
       "      <th>dysp</th>\n",
       "      <th>cough</th>\n",
       "      <th>fever</th>\n",
       "      <th>pain</th>\n",
       "      <th>nasal</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>high</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>Patient presents with high-grade fever of 39.5...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>warm</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>high</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>Chief complaint of persistent high fever (up t...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>warm</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>high</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>Febrile patient with history of high-grade fev...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>high</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>Patient reports ongoing high fever up to 38.9°...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>warm</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>high</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>Examined patient exhibiting high fever spikes ...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  season pneu inf  ... pain nasal                                               text\n",
       "0   cold   no  no  ...   no    no  Patient presents with high-grade fever of 39.5...\n",
       "1   warm   no  no  ...   no    no  Chief complaint of persistent high fever (up t...\n",
       "2   warm   no  no  ...   no    no  Febrile patient with history of high-grade fev...\n",
       "3   cold   no  no  ...   no    no  Patient reports ongoing high fever up to 38.9°...\n",
       "4   warm   no  no  ...   no    no  Examined patient exhibiting high fever spikes ...\n",
       "\n",
       "[5 rows x 9 columns]"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1000"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Make data partially observed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train set\n",
    "\n",
    "For the training set, we have some additional steps to do. We already masked out the text in 1/3 of the training data (step 4c) by not generating a clinical note for those samples. \n",
    "\n",
    "For the remaining 2/3 of the data, the symptoms are masked in 1/3 of the samples (step 4b)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "seed = 547\n",
    "train_obs, train_no_obs = train_test_split(train_text, train_size=1333, random_state=seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1333, 1334)"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(train_obs), len(train_no_obs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "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>season</th>\n",
       "      <th>pneu</th>\n",
       "      <th>inf</th>\n",
       "      <th>dysp</th>\n",
       "      <th>cough</th>\n",
       "      <th>fever</th>\n",
       "      <th>pain</th>\n",
       "      <th>nasal</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>281</th>\n",
       "      <td>warm</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Patient complains of low-grade fever for the p...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1466</th>\n",
       "      <td>warm</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Patient complaining of nasal congestion for th...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>543</th>\n",
       "      <td>warm</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Patient seeking advice on smoking cessation, r...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2322</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Patient presents with mild dyspnea, mainly wit...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2113</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Complains of coughing fits, along with constan...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     season pneu  ... nasal                                               text\n",
       "281    warm   no  ...   NaN  Patient complains of low-grade fever for the p...\n",
       "1466   warm   no  ...   NaN  Patient complaining of nasal congestion for th...\n",
       "543    warm   no  ...   NaN  Patient seeking advice on smoking cessation, r...\n",
       "2322   cold   no  ...   NaN  Patient presents with mild dyspnea, mainly wit...\n",
       "2113   cold   no  ...   NaN  Complains of coughing fits, along with constan...\n",
       "\n",
       "[5 rows x 9 columns]"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "train_no_obs.loc[:, [\"dysp\", \"cough\", \"pain\", \"fever\", \"nasal\"]] = np.nan\n",
    "train_no_obs.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now combine the three parts, making one big training set of 4000 samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "train_full = pd.concat([train_no_text, train_no_obs, train_obs])\n",
    "train_full = train_full.reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "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>season</th>\n",
       "      <th>pneu</th>\n",
       "      <th>inf</th>\n",
       "      <th>dysp</th>\n",
       "      <th>cough</th>\n",
       "      <th>fever</th>\n",
       "      <th>pain</th>\n",
       "      <th>nasal</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>none</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>none</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>warm</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>low</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>low</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>yes</td>\n",
       "      <td>yes</td>\n",
       "      <td>none</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  season pneu  inf dysp cough fever pain nasal text\n",
       "0   cold   no  yes   no   yes  none   no   yes     \n",
       "1   cold   no  yes   no   yes  none   no   yes     \n",
       "2   warm   no   no   no    no   low   no    no     \n",
       "3   cold   no   no   no    no   low   no    no     \n",
       "4   cold   no  yes  yes   yes  none   no   yes     "
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_full.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we remove the symptoms \"fever\" and \"pain\" from the training set (step 4a)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_full_masked = train_full[[\"season\", \"pneu\", \"inf\", \"dysp\", \"cough\", \"nasal\", \"text\"]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the training set found in \"data/train_4000_final.p\". We also save the version where the symptoms \"fever\" and \"pain\" are still observed, in \"data/train_4000_final_obs.p\", so we can use it to train the BN++ baseline. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "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>season</th>\n",
       "      <th>pneu</th>\n",
       "      <th>inf</th>\n",
       "      <th>dysp</th>\n",
       "      <th>cough</th>\n",
       "      <th>nasal</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>yes</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>yes</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>warm</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>yes</td>\n",
       "      <td>yes</td>\n",
       "      <td>yes</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  season pneu  inf dysp cough nasal text\n",
       "0   cold   no  yes   no   yes   yes     \n",
       "1   cold   no  yes   no   yes   yes     \n",
       "2   warm   no   no   no    no    no     \n",
       "3   cold   no   no   no    no    no     \n",
       "4   cold   no  yes  yes   yes   yes     "
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_full_masked.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "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>season</th>\n",
       "      <th>pneu</th>\n",
       "      <th>inf</th>\n",
       "      <th>dysp</th>\n",
       "      <th>cough</th>\n",
       "      <th>fever</th>\n",
       "      <th>pain</th>\n",
       "      <th>nasal</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>none</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>none</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>warm</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>low</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>low</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td>yes</td>\n",
       "      <td>yes</td>\n",
       "      <td>none</td>\n",
       "      <td>no</td>\n",
       "      <td>yes</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  season pneu  inf dysp cough fever pain nasal text\n",
       "0   cold   no  yes   no   yes  none   no   yes     \n",
       "1   cold   no  yes   no   yes  none   no   yes     \n",
       "2   warm   no   no   no    no   low   no    no     \n",
       "3   cold   no   no   no    no   low   no    no     \n",
       "4   cold   no  yes  yes   yes  none   no   yes     "
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_full.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "with open(f\"data/train_4000_final.p\", \"wb\") as file: \n",
    "    pickle.dump(train_full_masked, file)\n",
    "with open(f\"data/train_4000_final_fever_pain.p\", \"wb\") as file: \n",
    "    pickle.dump(train_full, file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test set\n",
    "\n",
    "For the test, we only remove the symptoms \"fever\" and \"pain\" (step 4a). We don't do steps 4b and 4c. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_masked = test[[\"season\", \"pneu\", \"inf\", \"dysp\", \"cough\", \"nasal\", \"text\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "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>season</th>\n",
       "      <th>pneu</th>\n",
       "      <th>inf</th>\n",
       "      <th>dysp</th>\n",
       "      <th>cough</th>\n",
       "      <th>nasal</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>Patient presents with high-grade fever of 39.5...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>warm</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>Chief complaint of persistent high fever (up t...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>warm</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>Febrile patient with history of high-grade fev...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>cold</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>Patient reports ongoing high fever up to 38.9°...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>warm</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>no</td>\n",
       "      <td>Examined patient exhibiting high fever spikes ...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  season pneu  ... nasal                                               text\n",
       "0   cold   no  ...    no  Patient presents with high-grade fever of 39.5...\n",
       "1   warm   no  ...    no  Chief complaint of persistent high fever (up t...\n",
       "2   warm   no  ...    no  Febrile patient with history of high-grade fev...\n",
       "3   cold   no  ...    no  Patient reports ongoing high fever up to 38.9°...\n",
       "4   warm   no  ...    no  Examined patient exhibiting high fever spikes ...\n",
       "\n",
       "[5 rows x 7 columns]"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_masked.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, we save both this final version of the test set, and the version where fever and pain are observed, which will be used to evaluate the BN++ model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "with open(f\"data/test_1000_final.p\", \"wb\") as file: \n",
    "    pickle.dump(test_masked, file)\n",
    "with open(f\"data/test_1000_final_fever_pain.p\", \"wb\") as file: \n",
    "    pickle.dump(test, file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add BioLORD embeddings\n",
    "\n",
    "We use the SentenceTransformer class from Huggingface to generate BioLORD embeddings for the clinical notes in both our train and test set. We calculate these upfront and store them in the train and test dataframes for easier use in our models. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "with open(f\"data/train_4000_final.p\", \"rb\") as file: \n",
    "    train_df = pickle.load(file)\n",
    "with open(f\"data/test_1000_final.p\", \"rb\") as file: \n",
    "    test_df = pickle.load(file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\Paloma\\anaconda3\\envs\\transformers\\Lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "model = SentenceTransformer(\"FremyCompany/BioLORD-2023\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode_text_BioLORD(row, model):\n",
    "    if row.name % 25 == 0: # print index to show progress\n",
    "        print(row.name)\n",
    "    \n",
    "    sent = row[\"text\"]\n",
    "\n",
    "    sent_emb = model.encode(sent)\n",
    "\n",
    "    return sent_emb "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df[\"BioLORD emb\"] = train_df.apply(encode_text_BioLORD, args=(model,), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_df[\"BioLORD emb\"] = test_df.apply(encode_text_BioLORD, args=(model,), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "with open(f\"data/test_1000_final.p\", \"wb\") as file: \n",
    "    pickle.dump(test_df, file)\n",
    "with open(f\"data/train_4000_final.p\", \"wb\") as file: \n",
    "    pickle.dump(train_df, file)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "bn-text",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}