notebooks/medical-text-encrypted-training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Encrypted Training Demo on Medical Text Data using SyferText"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Author:**\n",
    "- Carlos Salgado - [email](mailto:csalgado@uwo.ca) | [GitHub](https://github.com/socd06) | [LinkedIn](www.linkedin.com/in/eng-socd)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem Statement\n",
    "Bob <sup>MD</sup> and Alice <sup>MD</sup> are physicians running their respective medical practices and both have a database of private medical transcriptions. You own a Natural Language Processing (NLP) company and have been contacted by these physicians because both Bob <sup>MD</sup> and Alice <sup>MD</sup> have heard of the high quality of the Machine Learning as a Service (MLaaS) solutions you provide and want you to create a text classifier to help them automatically assign a medical specialty to each new patient text transcription."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Limitations\n",
    "\n",
    "Healthcare data is highly regulated and should be, for most intents and purposes, private. Therefore, if in a medical setting, the Machine Learning model being trained should not actually look at the data. \n",
    "\n",
    "Combining both Bob's and Alice's datasets, you should be able to create a bigger, better dataset that you could use to train your model with higher accuracy, only that you can't because it's all sensitive and private data, which is why you will need [PySyft](https://github.com/OpenMined/pysyft/) and [SyferText](https://github.com/OpenMined/SyferText/) to complete the job at hand."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Downloading the dataset (Skip if Cloned)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "# see repos mentioned above\n",
    "sys.path.append('../scripts')\n",
    "\n",
    "from util import download_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The URL template to all dataset files\n",
    "url_template = 'https://raw.githubusercontent.com/socd06/medical-nlp/master/data/%s'\n",
    "\n",
    "# File names to be downloaded from the using the URL template above\n",
    "files = ['classes.txt','train.csv','test.csv', 'clinical-stopwords.txt', 'vocab.txt']\n",
    "\n",
    "# Construct the list of urls\n",
    "urls = [url_template % file for file in files]\n",
    "\n",
    "# The dataset name and its root folder\n",
    "dataset_name = 'data'\n",
    "root_path = '../data'\n",
    "\n",
    "# Create the dataset folder if it is not already there\n",
    "if not os.path.exists('../data'):\n",
    "    os.mkdir('../data')\n",
    "\n",
    "# Start downloading\n",
    "download_dataset(dataset_name = dataset_name, \n",
    "                 urls = urls, \n",
    "                 root_path = root_path\n",
    "                )\n",
    "print(\"Succesfully downloaded:\",files)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importing libraries\n",
    "Make sure to first install [PySyft](https://github.com/OpenMined/PySyft) and [SyferText](https://github.com/OpenMined/SyferText) before you run this tutorial. \n",
    "Using virtual environments is highly recommended for any PySyft experiment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Falling back to insecure randomness since the required custom op could not be found for the installed version of TensorFlow. Fix this by compiling custom ops. Missing file was '/home/carlos/anaconda3/envs/syft/lib/python3.7/site-packages/tf_encrypted/operations/secure_random/secure_random_module_tf_1.15.2.so'\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/carlos/anaconda3/envs/syft/lib/python3.7/site-packages/tf_encrypted/session.py:24: The name tf.Session is deprecated. Please use tf.compat.v1.Session instead.\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Torch was already hooked... skipping hooking process\n",
      "WARNING:root:Torch was already hooked... skipping hooking process\n",
      "WARNING:root:Torch was already hooked... skipping hooking process\n"
     ]
    }
   ],
   "source": [
    "# PySyft imports\n",
    "import syft as sy\n",
    "from syft.generic.string import String\n",
    "# SyferText imports\n",
    "import syfertext\n",
    "from syfertext.pipeline import SimpleTagger\n",
    "\n",
    "# PyTorch imports\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data import Dataset\n",
    "import torch.optim as optim\n",
    "\n",
    "# Useful imports\n",
    "import numpy as np\n",
    "import csv\n",
    "from sklearn.model_selection import train_test_split\n",
    "import os\n",
    "from pprint import pprint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparing the work environment\n",
    "In this part, we assume each client owns a part of the full dataset and we prepare each worker to perform encrypted training. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Virtual environment\n",
    "\n",
    "If you have not setup [PyGrid](https://github.com/OpenMined/PyGrid/) yet or want to only focus on the training aspect of this demo, run the cell below and skip the next cell.\n",
    "\n",
    "A work environment is simulated with three main actors, a company (us) and two clients owning two private datasets (Bob and Alice) but also a crypto provider which will provide the primitives for Secure Multi-Party Computation (SMPC). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Torch was already hooked... skipping hooking process\n"
     ]
    }
   ],
   "source": [
    "# Create a torch hook for PySyft\n",
    "hook = sy.TorchHook(torch)\n",
    "\n",
    "# Create some PySyft workers\n",
    "me = hook.local_worker # This is the worker representing the deep learning company\n",
    "bob = sy.VirtualWorker(hook, id = 'bob') # Bob owns the first dataset\n",
    "alice = sy.VirtualWorker(hook, id = 'alice') # Alice owns the second dataset\n",
    "\n",
    "crypto_provider = sy.VirtualWorker(hook, id = 'crypto_provider') # provides encryption primitive for SMPC"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Local PyGrid environment\n",
    "\n",
    "A local grid work environment can be initialized using [PyGrid](https://github.com/OpenMined/PyGrid/) with three main actors, a company (us) and two clients owning two private datasets (Bob and Alice) but also a crypto provider which will provide the primitives for Secure Multi-Party Computation (SMPC). **Warning: Highly Experimental**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading dataset locally"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the path to the dataset file\n",
    "dataset_path = '../data/train.csv'\n",
    "\n",
    "# store the dataset as a list of dictionaries\n",
    "# each dictionary has two keys, 'text' and 'label'\n",
    "# the 'text' element is a PySyft String\n",
    "# the 'label' element is an integer with 1 for each surgical specialty and a 0 otherwise\n",
    "dataset_local = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(dataset_path, 'r') as dataset_file:\n",
    "    \n",
    "    # Create a csv reader object\n",
    "    reader = csv.DictReader(dataset_file)\n",
    "    \n",
    "    for elem in reader:\n",
    "        \n",
    "        # Create one entry\n",
    "        # Check if the medical specialty contains 1 (label for surgery) \n",
    "        # otherwise mark it as 0\"\n",
    "        example = dict(text = String(elem['text']),                       \n",
    "                       label = 1 if elem['label'] == '1' else 0\n",
    "                      )\n",
    "        \n",
    "        # add to the local dataset\n",
    "        dataset_local.append(example)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is how an element in the list looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import needed to generate random integer numbers\n",
    "from random import randint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Entry # 942\n",
      "{'label': 0,\n",
      " 'text': \"HISTORY:  ,This 15-day-old female presents to Children's Hospital and transferred from Hospital Emergency Department for further evaluation.  Information is obtained in discussion with the mother and the grandmother in review of previous medical records.  This patient had the onset on the day of presentation of a jelly-like red-brown stool started on Tuesday morning.  Then, the patient was noted to vomit after feeds.  The patient was evaluated at Hospital with further evaluation with laboratory data showing a white blood cell count elevated at 22.2; hemoglobin 14.1; sodium 138; potassium 7.2, possibly hemolyzed; chloride 107; CO2 23; BUN 17; creatinine 1.2; and glucose of 50, which was repeated and found to be stable in that range.  The patient underwent a barium enema, which was read by the radiologist as negative.  The patient was transferred to Children's Hospital for further evaluation after being given doses of ampicillin, cefotaxime, and Rocephin.,PAST MEDICAL HISTORY: , Further, the patient was born in Hospital.  Birth weight was 6 pounds 4 ounces.  There was maternal hypertension.  Mother denies group B strep or herpes.  Otherwise, no past medical history.,IMMUNIZATIONS: , None today.,MEDICATIONS: , Thrush medicine identified as nystatin.,ALLERGIES: , Denied.,PAST SURGICAL HISTORY: , Denied.,SOCIAL HISTORY:  ,Here with mother and grandmother, lives at home.  There is no smoking at home.,FAMILY HISTORY: , None noted exposures.,REVIEW OF SYSTEMS:  ,The patient is fed Enfamil, bottle-fed.  Has had decreased feeding, has had vomiting, has had diarrhea, otherwise negative on the 10 plus systems reviewed.,PHYSICAL EXAMINATION:,VITAL SIGNS/GENERAL:  On physical examination, the initial temperature 97.5, pulse 140, respirations 48 on this 2 kg 15-day-old female who is small, well-developed female, age appropriate.,HEENT: Head is atraumatic and normocephalic with a soft and flat anterior fontanelle.  Pupils are equal, round, and reactive to light.  Grossly conjugate.  Bilateral red reflex appreciated bilaterally.  Clear TMs, nose, and oropharynx.  There is a kind of abundant thrush and white patches on the tongue.,NECK:  Supple, full, painless, and nontender range of motion.,CHEST:  Clear to auscultation, equal, and stable.,HEART:  Regular without rubs or murmurs, and femoral pulses are appreciated bilaterally.,ABDOMEN:  Soft and nontender.  No hepatosplenomegaly or masses.,GENITALIA:  Female genitalia is present on a visual examination.,SKIN:  No significant bruising, lesions, or rash.,EXTREMITIES:  Moves all extremities, and nontender.  No deformity.,NEUROLOGICALLY:  Eyes open, moves all extremities, grossly age appropriate.,MEDICAL DECISION MAKING: , The differential entertained on this patient includes upper respiratory infection, gastroenteritis, urinary tract infection, dehydration, acidosis, and viral syndrome.  The patient is evaluated in the emergency department laboratory data, which shows a white blood cell count of 13.1, hemoglobin 14.0, platelets 267,000, 7 stabs, 68 segs, 15 lymphs, and 9 monos.  Serum electrolytes not normal.  Sodium 138, potassium 5.0, chloride 107, CO2 acidotic at 18, glucose normal at 88, and BUN markedly elevated at 22 as is the creatinine of 1.4.  AST and ALT were elevated as well at 412 and 180 respectively.  A cath urinalysis showing no signs of infection.  Spinal fluid evaluation, please see procedure note below.  White count 0, red count 2060.  Gram stain negative.,PROCEDURE NOTE: , After discussion of the risks, benefits, and indications, and obtaining informed consent with the family and their agreement to proceed, this patient was placed in the left lateral position and using aseptic Betadine preparation, sterile draping, and sterile technique pursued throughout, this patient's L4- L5 interspace was anesthetized with the 1% lidocaine solution following the above sterile preparation, entered with a 22-gauge styletted spinal needle of approximately 0.5 mL clear CSF, they were very slow to obtain.  The fluid was obtained, the needle was removed, and sterile bandage was placed.  The fluid was sent to laboratory for further evaluation (aunt and grandmother) were present throughout the period of time during this procedure and the procedure was tolerated well.  An i-STAT initially obtained showed somewhat of an acidosis with a base excess of -12.  A repeat i-STAT after a bolus of normal saline and a second bolus of normal saline, her maintenance rate of D5 half showed a base excess of -11, which is slowly improving, but not very fast.  Based on the above having this patient consulted to the Hospitalist Service at 2326 hours of request, this patient was consulted to PICU with the plan that the patient need to have continued IV fluids.  Showing signs of dehydration, a third bolus of normal saline was provided, twice maintenance D5 half was continued.  The patient was admitted to the Hospitalist Service for continued IV fluids.  The patient maintains to have clear lungs, has been feeding well here in the department, took virtually a whole small bottle of the appropriate formula.  She has not had any vomiting, is burping.  The patient is admitted for continued close observation and rehydration due to the working diagnoses of gastroenteritis, metabolic acidosis, and dehydration.  Critical care time on this patient is less than 30 minutes, exclusive, otherwise time has been spent evaluating this patient according to this patient's care and admission to the Hospitalist Service.nan\"}\n"
     ]
    }
   ],
   "source": [
    "# Get a random index to verify entry examples\n",
    "random_index = randint(0,len(dataset_local))\n",
    "print(\"Entry #\",random_index)\n",
    "example = dataset_local[random_index]\n",
    "pprint(example)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we verified our examples we can look into our data types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'syft.generic.string.String'>\n",
      "<class 'int'>\n"
     ]
    }
   ],
   "source": [
    "print(type(example['text']))\n",
    "print(type(example['label']))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This transcription text is a PySyft `String` object. The label is an integer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Distributing documents privately\n",
    "\n",
    "We simulate two private datasets owned by two clients (Bob and Alice):\n",
    "\n",
    "1. Load the whole dataset in `mtsamples.csv` locally (the `me` worker). The data will be loaded as a list of dictionaries that has the following format: `[ {'transcription': <transcription text>, 'label': <0 or 1>}, {...}, {...}]`\n",
    "\n",
    "\n",
    "2. Split the dataset into two parts, one for Bob and the other for Alice. Each part will be also split into a training set and a validation set. This will create four lists: `train_bob`, `valid_bob`, `train_alice`, `valid_alice`. Each list has the same format mentioned above.\n",
    "\n",
    "\n",
    "3. Each element in the four lists will be sent to the corresponding worker. This will change the content of the lists as depicted in **Figure 1**. Each list willl hold PySyft pointers to the texts and labels instead of the objects themselves."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div>\n",
    "<br>\n",
    "<img alt = 'medical transcriptions remote datasets' src ='./img/mtsamples_remote.png' style='width:700px'>\n",
    "<div>\n",
    "<div style='width:600px;margin:30px auto 10px auto;text-align:center;'>\n",
    "<strong> Figure 1: </strong> The transcription text and their labels are remotely located in Bob and Alice's remote worker machines, only pointers to them are kept by the local worker (the company's machine).\n",
    "</div>\n",
    "</div>\n",
    "<br>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's split the dataset into two equal parts and send each part to a different worker simulating two remote datasets mentioned before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create two datasets, one for Bob and another one for Alice\n",
    "dataset_bob, dataset_alice = train_test_split(dataset_local[:25000], train_size = 0.5)\n",
    "\n",
    "# Now create a validation set for Bob and another one for Alice\n",
    "train_bob, val_bob = train_test_split(dataset_bob, train_size = 0.9)\n",
    "train_alice, val_alice = train_test_split(dataset_alice, train_size = 0.9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Making the datasets remote:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A function that sends the content of each split to a remote worker\n",
    "def make_remote_dataset(dataset, worker):\n",
    "\n",
    "    # Got through each example in the dataset\n",
    "    for example in dataset:\n",
    "        \n",
    "        # Send each transcription text\n",
    "        example['text'] = example['text'].send(worker)\n",
    "                       \n",
    "        # Send each label as a one-hot-encoded vector\n",
    "        one_hot_label = torch.zeros(2).scatter(0, torch.Tensor([example['label']]).long(), 1)\n",
    "        \n",
    "        # print for debugging purposes\n",
    "        # print(\"mapping\",example['label'],\" to \",one_hot_label)\n",
    "        \n",
    "        # Send the transcription label\n",
    "        example['label'] = one_hot_label.send(worker)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above function transforms the label to a one-hot-encoded format before sending it to a remote worker. Every label corresponds to a 2-digit tensor of binary values (`[1,0]` or `[0,1]`)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can convert the dataset into a remote dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Bob's remote dataset\n",
    "make_remote_dataset(train_bob, bob)\n",
    "make_remote_dataset(val_bob, bob)\n",
    "\n",
    "# Alice's remote dataset\n",
    "make_remote_dataset(train_alice, alice)\n",
    "make_remote_dataset(val_alice, alice)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now Bob's dataset looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'syft.generic.pointers.string_pointer.StringPointer'>\n",
      "(Wrapper)>[PointerTensor | me:82245969193 -> bob:47734380440]\n"
     ]
    }
   ],
   "source": [
    "# Take an element from the dataset\n",
    "example = train_bob[0]\n",
    "\n",
    "print(type(example['text']))\n",
    "print(example['label'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the text type is a PySyft `StringPointer` that points to the real `String` object  located in Bob's machine. The label type is a PySyft `PointerTensor`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can review and see where the label is located."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<VirtualWorker id:bob #objects:4504>\n",
      "<VirtualWorker id:bob #objects:4504>\n"
     ]
    }
   ],
   "source": [
    "print(example['text'].location)\n",
    "print(example['label'].location)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Confirming that the dataset is now remote and also confirming the information in **Figure 1**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the environment and the data are ready for the next step."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a `SyferText` Language object and a pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a Language object with SyferText\n",
    "nlp = syfertext.load('en_core_web_lg', owner = me)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Whenever you create a Language object a pipeline will be created. At initialization, a pipeline only contains a tokenizer. You can see this for yourself using the `pipeline_template` property:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'remote': True, 'name': 'tokenizer'}]"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nlp.pipeline_template"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "syfertext.language.Language"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(nlp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that the tokenizer entry has a propery called `remote` set to `True`. This means we allowed the tokenizer to be sent to a remote worker for the string to be tokenized there.\n",
    "\n",
    "We can add more components to the pipeline by using the `add_pipe` method of the Language class. One component we can add is a `SimpleTagger` object. This is a SyferText object that we can use to set custom attributes to individual tokens. In this tutorial, we will create two taggers: One that tags tokens that are Stop Words and another one that tags each token with their respective class.\n",
    "\n",
    "By tagging we mean setting a custom attribute to a token and assigning it a given value (e.g. An attribute called `is_stop` with `True` and `False` values when evaluating Stop Words. \n",
    "\n",
    "You can refer to **Figure 2** to see how a pipeline is distributed on multiple workers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adding the taggers to the pipeline\n",
    "The `excluded_tokens` dictionary will be used further down, when we create embedding vectors for the transcriptions. This dictionary will enable us to exclude some tokens when we create a document embedding. Such exclusion will be based on the value of the custom attributes we set with the taggers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "use_stop_tagger = True\n",
    "use_vocab_tagger = True\n",
    "\n",
    "# Token with these custom tags\n",
    "# will be excluded from creating\n",
    "# the Doc vector\n",
    "excluded_tokens = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating a  tagger for stop words\n",
    "We will start by creating the Stop Word tagger. First loading the stop word file into a list of words:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the list of stop words\n",
    "with open('../data/clinical-stopwords.txt', 'r') as f:\n",
    "    stop_words = set(f.read().splitlines())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we create the tagger which is an object of the `SimpleTagger` class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a simple tagger object to tag stop words\n",
    "stop_tagger = SimpleTagger(attribute = 'is_stop',\n",
    "                           lookups = stop_words,\n",
    "                           tag = True,\n",
    "                           default_tag = False,\n",
    "                           case_sensitive = False\n",
    "                          )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the `lookups` argument passed was the list of stop words.\n",
    "\n",
    "Every token in the `Doc` object will be given a custom attribute called `is_stop`. Every time a stop word is found, this attribute will be given the value `True` specified by the `tag` argument of the `SimpleTagger` class initializer, otherwise, the `default_tag` will be used (e.g. `False`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "if use_stop_tagger:\n",
    "\n",
    "    # Add the stop word to the pipeline\n",
    "    nlp.add_pipe(name = 'stop tagger',\n",
    "                 component = stop_tagger,\n",
    "                 remote = True\n",
    "                )\n",
    "\n",
    "    # Tokens with 'is_stop' = True are\n",
    "    # not going to be used when creating the \n",
    "    # Doc vector\n",
    "    excluded_tokens['is_stop'] = {True}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating a  tagger for vocab words\n",
    "Likewise we create a SimpleTagger instance that filters out tokens not in our vocabulary file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('../data/vocab.txt', 'r') as f:\n",
    "    vocab_words = set(f.read().splitlines())  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a simple tagger object to tag stop words\n",
    "vocab_tagger = SimpleTagger(attribute = 'is_vocab',\n",
    "                           lookups = vocab_words,\n",
    "                           tag = True,\n",
    "                           default_tag = False,\n",
    "                           case_sensitive = False\n",
    "                          )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "if use_vocab_tagger:\n",
    "\n",
    "    # Add the stop word to the pipeline\n",
    "    nlp.add_pipe(name = 'vocab tagger',\n",
    "                 component = vocab_tagger,\n",
    "                 remote = True\n",
    "                )\n",
    "\n",
    "    # Tokens with 'is_vocab' = False are\n",
    "    # not going to be used when creating the \n",
    "    # Doc vector\n",
    "    excluded_tokens['is_vocab'] = {False}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check out what pipe components are included in the pipeline:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'remote': True, 'name': 'tokenizer'},\n",
       " {'remote': True, 'name': 'stop tagger'},\n",
       " {'remote': True, 'name': 'vocab tagger'}]"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nlp.pipeline_template"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a Dataset class\n",
    "\n",
    "Now that the datasets are remote and ready along with the `Language` object and its pipeline we can create PyTorch loaders to make data batches for training and validation.\n",
    "\n",
    "The batches will be composed of training examples coming from both Bob's and Alice's datasets as if it were only one big dataset.\n",
    "\n",
    "Each example in the batch contains an encrypted version of one transcription's embedding vector and its encrypted label. For this tutorial, the vector will be computed as the average of the transcription's individual token vectors taken from the `en_core_web_lg` language model. Also, tokens with custom tags indicated in `excluded_tokens` won't be taken into account in computing a transcription's vector.\n",
    "\n",
    "From **Figure 2** we can see how the transcription text is remotely preprocessed by SyferText: \n",
    "\n",
    "1. First, the `Language` object `nlp` is used to preprocess one transcription on Bob's or Alice's machine.\n",
    "2. The object `nlp` determines that the real transcription text is actually remote, so it sends a subpipeline containing the required pipeline components we defined to the corresponding worker.\n",
    "3. The subpipeline is run and a `Doc` object is created on the remote worker containing the transcription's individual tokens appropriately tokenized and tagged.\n",
    "4. On the local worker, a `DocPointer` object is created pointing to that `Doc` object.\n",
    "5. By calling `get_encrypted_vector()` on the `DocPointer`, the call is forwarded to `Doc`, which, in turn, computes the `Doc` vector, encrypts it with Secure Multy-Party Computation (SMPC) using PySyft and returns it to the caller at the local worker.\n",
    "6. The PyTorch dataloader takes this encrypted vector and appends it to the training or validation batch.\n",
    "\n",
    "Note that at no moment in the process, the plaintext data of the remote datasets are revealed to the local worker. *Privacy is preserved thanks to SyferText and PySyft!*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div>\n",
    "<br>\n",
    "<img alt =  'SyferText pipeline' src ='./img/mt_pipeline.png' style='width:700px;'>\n",
    "<div>\n",
    "<p style='width:600px;margin:30px auto 10px auto;text-align:center;'>\n",
    "<strong> Figure 2: </strong> A pipeline on the local worker only contains pointers to subpipelines carrying out the actual preprocessing on remote workers.\n",
    "</p>\n",
    "</div>\n",
    "<br>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Take a minute to review the `__getitem__()` method of the custom PyTorch `Dataset` object defined below and see how. Please take a few minutes to check it out below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DatasetMTS(Dataset):\n",
    "    \n",
    "    def __init__(self, sets, share_workers, crypto_provider, nlp):\n",
    "        \"\"\"Initialize the Dataset object\n",
    "        \n",
    "        Args:\n",
    "            sets (list): A list containing all training OR \n",
    "                all validation sets to be used.\n",
    "            share_workers (list): A list of workers that will\n",
    "                be used to hold the SMPC shares.\n",
    "            crypto_provider (worker): A worker that will \n",
    "                provide SMPC primitives for encryption.\n",
    "            nlp: This is SyferText's Language object containing\n",
    "                the preprocessing pipeline.\n",
    "        \"\"\"\n",
    "        self.sets = sets\n",
    "        self.crypto_provider = crypto_provider\n",
    "        self.workers = share_workers\n",
    "    \n",
    "        # Create a single dataset unifying all datasets.\n",
    "        # A property called `self.dataset` is created \n",
    "        # as a result of this call.\n",
    "        self._create_dataset()\n",
    "        \n",
    "        # The language model\n",
    "        self.nlp = nlp\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        \"\"\"In this function, preprocessing with SyferText \n",
    "        of one transcription will be triggered. Encryption will also\n",
    "        be performed and the encrypted vector will be obtained.\n",
    "        The encrypted label will be computed too.\n",
    "        \n",
    "        Args:\n",
    "            index (int): This is an integer received by the \n",
    "                PyTorch DataLoader. It specifies the index of\n",
    "                the example to be fetched. This actually indexes\n",
    "                one example in `self.dataset` which pools over\n",
    "                examples of all the remote datasets.\n",
    "        \"\"\"\n",
    "        \n",
    "        # get the example\n",
    "        example = self.dataset[index]\n",
    "        \n",
    "        # Run the preprocessing pipeline on \n",
    "        # the transcription text and get a DocPointer object\n",
    "        doc_ptr = self.nlp(example['text'])\n",
    "        \n",
    "        # Get the encrypted vector embedding for the document\n",
    "        vector_enc = doc_ptr.get_encrypted_vector(bob, \n",
    "                                                  alice, \n",
    "                                                  crypto_provider = self.crypto_provider,\n",
    "                                                  requires_grad = True,\n",
    "                                                  excluded_tokens = excluded_tokens\n",
    "                                                 )\n",
    "        \n",
    "\n",
    "        # Encrypt the target label\n",
    "        label_enc = example['label'].fix_precision().share(bob, \n",
    "                                                           alice, \n",
    "                                                           crypto_provider = self.crypto_provider,\n",
    "                                                           requires_grad = True\n",
    "                                                          ).get()\n",
    "\n",
    "\n",
    "        return vector_enc, label_enc\n",
    "\n",
    "    \n",
    "    def __len__(self):\n",
    "        \"\"\"Returns the combined size of all of the \n",
    "        remote training/validation sets.\n",
    "        \"\"\"\n",
    "        \n",
    "        # The size of the combined datasets\n",
    "        return len(self.dataset)\n",
    "\n",
    "    def _create_dataset(self):\n",
    "        \"\"\"Create a single list unifying examples from all remote datasets\n",
    "        \"\"\"\n",
    "        \n",
    "        # Initialize the dataset\n",
    "        self.dataset = []\n",
    "      \n",
    "        # populate the dataset list\n",
    "        for dataset in self.sets:\n",
    "            for example in dataset:\n",
    "                self.dataset.append(example)\n",
    "                \n",
    "    @staticmethod\n",
    "    def collate_fn(batch):\n",
    "        \"\"\"The collat_fn method to be used by the\n",
    "        PyTorch data loader.\n",
    "        \"\"\"\n",
    "        \n",
    "        # Unzip the batch\n",
    "        vectors, targets = list(zip(*batch))        \n",
    "            \n",
    "        # concatenate the vectors\n",
    "        vectors = torch.stack(vectors)\n",
    "        \n",
    "        #concatenate the labels\n",
    "        targets = torch.stack(targets)        \n",
    "        \n",
    "        return vectors, targets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now create two such `DatasetMTS` objects, one for training and the other for validation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate a training Dataset object\n",
    "trainset = DatasetMTS(sets = [train_bob,\n",
    "                               train_alice],\n",
    "                       share_workers = [bob, alice],\n",
    "                       crypto_provider = crypto_provider,\n",
    "                       nlp = nlp\n",
    "                      )\n",
    "\n",
    "# Instantiate a validation Dataset object\n",
    "valset = DatasetMTS(sets = [val_bob,\n",
    "                             val_alice],\n",
    "                     share_workers = [bob, alice],\n",
    "                     crypto_provider = crypto_provider,\n",
    "                     nlp = nlp\n",
    "                    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use the `__getitem__` method to obtain the embedding vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Training Vector size is 300\n"
     ]
    }
   ],
   "source": [
    "vec_enc, label_enc = trainset.__getitem__(1)\n",
    "print(f' Training Vector size is {vec_enc.shape[0]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training configuration\n",
    "\n",
    "We will now describe the training hyper-parameters for training and validation and create the PyTorch data loaders:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "#device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set model hyper-parameters constants\n",
    "# batch=32 and LR=0.001 are the original setup\n",
    "# batch=32 and LR=0.0025 also work \n",
    "EMBED_DIM = vec_enc.shape[0]\n",
    "BATCH_SIZE = 128 #16, 64 worked\n",
    "LEARNING_RATE = 0.001 # 0.1, 0.05 worked\n",
    "EPOCHS = 3 # Complete passes in the data\n",
    "NUN_CLASS = 2 # 2 classes since its a binary classifier "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate the DataLoader object for the training set\n",
    "trainloader = DataLoader(trainset, shuffle = True,\n",
    "                         batch_size = BATCH_SIZE, num_workers = 0, \n",
    "                         collate_fn = trainset.collate_fn)\n",
    "\n",
    "\n",
    "# Instantiate the DataLoader object for the validation set\n",
    "valloader = DataLoader(valset, shuffle = True,\n",
    "                       batch_size = BATCH_SIZE, num_workers = 0, \n",
    "                       collate_fn = valset.collate_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating an encrypted classifier model\n",
    "The classifier we will use is a simple fully connected network with 300 input features which is the size of the embedding vectors computed previously by SyferText. The network is a binary classifier with an outputs for surgical specialties and another one for every other specialty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Classifier(torch.nn.Module):\n",
    "    \n",
    "    def __init__(self, in_features, out_features):\n",
    "        super(Classifier, self).__init__()\n",
    "        \n",
    "        self.fc1 = torch.nn.Linear(in_features, 64)\n",
    "        self.fc2 = torch.nn.Linear(64, 32)\n",
    "        self.fc3 = torch.nn.Linear(32, out_features)\n",
    "                \n",
    "    def forward(self, inputs):\n",
    "        x = F.relu(self.fc1(inputs.squeeze(1)))\n",
    "        x = F.relu(self.fc2(x))\n",
    "       \n",
    "        logits = self.fc3(x)\n",
    "        \n",
    "        probs = F.relu(logits)\n",
    "        \n",
    "        return probs, logits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we initialize and encrypt the classifier. The encryption here must use the same workers that hold the share and the same primitives used to encrypt the document vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Classifier(\n",
      "  (fc1): Linear(in_features=300, out_features=64, bias=True)\n",
      "  (fc2): Linear(in_features=64, out_features=32, bias=True)\n",
      "  (fc3): Linear(in_features=32, out_features=2, bias=True)\n",
      ")\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/carlos/anaconda3/envs/syft/lib/python3.7/site-packages/syft/frameworks/torch/tensors/interpreters/additive_shared.py:79: UserWarning: Use dtype instead of field\n",
      "  warnings.warn(\"Use dtype instead of field\")\n",
      "/home/carlos/anaconda3/envs/syft/lib/python3.7/site-packages/syft/frameworks/torch/tensors/interpreters/additive_shared.py:91: UserWarning: Default args selected\n",
      "  warnings.warn(\"Default args selected\")\n"
     ]
    }
   ],
   "source": [
    "# Create the classifer\n",
    "model = Classifier(in_features = EMBED_DIM, out_features = NUN_CLASS)\n",
    "\n",
    "# Apply SMPC encryption\n",
    "model = model.fix_precision().share(bob, alice, \n",
    "                                              crypto_provider = crypto_provider,\n",
    "                                              requires_grad = True\n",
    "                                              )\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last thing to do before training is creating an optimizer. The optimizer doesn't need to be encrypted since it operates separately within each worker holding the classifier and the embeddings' shares. One thing to note is that the optimizer needs to operate on fixed precision numbers to be able to encode shares."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize stochastic gradient descent (SGD) optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.SGD(params = model.parameters(),lr = LEARNING_RATE, momentum=0.3)\n",
    "\n",
    "optimizer = optimizer.fix_precision()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SGD (\n",
      "Parameter Group 0\n",
      "    dampening: 0\n",
      "    lr: FixedPrecisionTensor>tensor(1)\n",
      "    momentum: FixedPrecisionTensor>tensor(300)\n",
      "    nesterov: False\n",
      "    weight_decay: 0\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "print(optimizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Training and Tensorboard\n",
    "We need to create a summary writer in order to view the training and validation curves for loss and accuracy.\n",
    "Then we will be able to run `Tensorboard` and see the information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a summary writer for logging performance with Tensorboard\n",
    "writer = SummaryWriter()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Open a terminal, navigate to the folder containing this notebook, and run:\n",
    "\n",
    "```\n",
    "$ tensorboard --logdir runs/\n",
    "```\n",
    "\n",
    "Then open you favorite web browser and go to `localhost:6006`.\n",
    "\n",
    "You should now be able to see performance curves."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now ready to run the below cell to launch the training.\n",
    "`NLLLoss()` is not yet implemented in PySyft for SMPC mode so we will use Mean Squared Error (MSE) as a training loss even though is not the best choice for a classification task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 \tLoss: 126.72(train)\t|\tAcc: 34.38%(train) 1\n",
      "epoch: 1 \tLoss: 110.43(valid)\t|\tAcc: 25.00%(valid) 1\n",
      "epoch: 1 \tLoss: 110.15(train)\t|\tAcc: 29.69%(train) 2\n",
      "epoch: 1 \tLoss: 70.01(valid)\t|\tAcc: 65.62%(valid) 2\n",
      "epoch: 1 \tLoss: 69.78(train)\t|\tAcc: 66.41%(train) 3\n",
      "epoch: 1 \tLoss: 55.28(valid)\t|\tAcc: 68.75%(valid) 3\n",
      "epoch: 1 \tLoss: 54.78(train)\t|\tAcc: 69.53%(train) 4\n",
      "epoch: 1 \tLoss: 58.21(valid)\t|\tAcc: 65.62%(valid) 4\n",
      "epoch: 1 \tLoss: 58.30(train)\t|\tAcc: 65.62%(train) 5\n",
      "epoch: 1 \tLoss: 52.95(valid)\t|\tAcc: 70.31%(valid) 5\n",
      "epoch: 1 \tLoss: 51.28(train)\t|\tAcc: 71.88%(train) 6\n",
      "epoch: 1 \tLoss: 51.35(valid)\t|\tAcc: 71.88%(valid) 6\n",
      "epoch: 1 \tLoss: 58.55(train)\t|\tAcc: 64.06%(train) 7\n",
      "epoch: 1 \tLoss: 53.36(valid)\t|\tAcc: 69.53%(valid) 7\n",
      "epoch: 1 \tLoss: 57.77(train)\t|\tAcc: 64.06%(train) 8\n",
      "epoch: 1 \tLoss: 50.95(valid)\t|\tAcc: 73.44%(valid) 8\n",
      "epoch: 1 \tLoss: 60.82(train)\t|\tAcc: 60.16%(train) 9\n",
      "epoch: 1 \tLoss: 54.16(valid)\t|\tAcc: 69.53%(valid) 9\n",
      "epoch: 1 \tLoss: 52.32(train)\t|\tAcc: 73.44%(train) 10\n",
      "epoch: 1 \tLoss: 55.87(valid)\t|\tAcc: 65.62%(valid) 10\n",
      "epoch: 1 \tLoss: 52.96(train)\t|\tAcc: 69.53%(train) 11\n",
      "epoch: 1 \tLoss: 47.56(valid)\t|\tAcc: 74.22%(valid) 11\n",
      "epoch: 1 \tLoss: 47.42(train)\t|\tAcc: 74.22%(train) 12\n",
      "epoch: 1 \tLoss: 44.21(valid)\t|\tAcc: 76.56%(valid) 12\n",
      "epoch: 1 \tLoss: 51.01(train)\t|\tAcc: 71.09%(train) 13\n",
      "epoch: 1 \tLoss: 55.31(valid)\t|\tAcc: 66.41%(valid) 13\n",
      "epoch: 1 \tLoss: 47.93(train)\t|\tAcc: 73.44%(train) 14\n",
      "epoch: 1 \tLoss: 46.97(valid)\t|\tAcc: 73.44%(valid) 14\n",
      "epoch: 1 \tLoss: 56.92(train)\t|\tAcc: 64.84%(train) 15\n",
      "epoch: 1 \tLoss: 49.70(valid)\t|\tAcc: 71.09%(valid) 15\n",
      "epoch: 1 \tLoss: 57.12(train)\t|\tAcc: 62.50%(train) 16\n",
      "epoch: 1 \tLoss: 51.90(valid)\t|\tAcc: 69.53%(valid) 16\n",
      "epoch: 1 \tLoss: 53.88(train)\t|\tAcc: 66.41%(train) 17\n",
      "epoch: 1 \tLoss: 55.24(valid)\t|\tAcc: 63.28%(valid) 17\n",
      "epoch: 1 \tLoss: 52.54(train)\t|\tAcc: 68.75%(train) 18\n",
      "epoch: 1 \tLoss: 46.98(valid)\t|\tAcc: 72.66%(valid) 18\n",
      "epoch: 1 \tLoss: 50.98(train)\t|\tAcc: 68.75%(train) 19\n",
      "epoch: 1 \tLoss: 41.29(valid)\t|\tAcc: 77.34%(valid) 19\n",
      "epoch: 1 \tLoss: 56.76(train)\t|\tAcc: 62.50%(train) 20\n",
      "epoch: 1 \tLoss: 44.30(valid)\t|\tAcc: 78.12%(valid) 20\n",
      "epoch: 1 \tLoss: 47.93(train)\t|\tAcc: 71.88%(train) 21\n",
      "epoch: 1 \tLoss: 48.75(valid)\t|\tAcc: 69.53%(valid) 21\n",
      "epoch: 1 \tLoss: 52.59(train)\t|\tAcc: 64.84%(train) 22\n",
      "epoch: 1 \tLoss: 48.28(valid)\t|\tAcc: 68.75%(valid) 22\n",
      "epoch: 1 \tLoss: 49.98(train)\t|\tAcc: 67.19%(train) 23\n",
      "epoch: 1 \tLoss: 45.07(valid)\t|\tAcc: 72.66%(valid) 23\n",
      "epoch: 1 \tLoss: 48.24(train)\t|\tAcc: 67.97%(train) 24\n",
      "epoch: 1 \tLoss: 44.93(valid)\t|\tAcc: 71.09%(valid) 24\n",
      "epoch: 1 \tLoss: 43.34(train)\t|\tAcc: 73.44%(train) 25\n",
      "epoch: 1 \tLoss: 47.41(valid)\t|\tAcc: 68.75%(valid) 25\n",
      "epoch: 1 \tLoss: 51.93(train)\t|\tAcc: 65.62%(train) 26\n",
      "epoch: 1 \tLoss: 44.00(valid)\t|\tAcc: 71.09%(valid) 26\n",
      "epoch: 1 \tLoss: 48.56(train)\t|\tAcc: 64.84%(train) 27\n",
      "epoch: 1 \tLoss: 44.23(valid)\t|\tAcc: 70.31%(valid) 27\n",
      "epoch: 1 \tLoss: 49.11(train)\t|\tAcc: 63.28%(train) 28\n",
      "epoch: 1 \tLoss: 45.68(valid)\t|\tAcc: 68.75%(valid) 28\n",
      "epoch: 1 \tLoss: 50.01(train)\t|\tAcc: 66.41%(train) 29\n",
      "epoch: 1 \tLoss: 41.94(valid)\t|\tAcc: 71.09%(valid) 29\n",
      "epoch: 1 \tLoss: 44.98(train)\t|\tAcc: 69.53%(train) 30\n",
      "epoch: 1 \tLoss: 42.33(valid)\t|\tAcc: 72.66%(valid) 30\n",
      "epoch: 1 \tLoss: 45.15(train)\t|\tAcc: 67.19%(train) 31\n",
      "epoch: 1 \tLoss: 45.28(valid)\t|\tAcc: 66.41%(valid) 31\n",
      "epoch: 1 \tLoss: 27.44(train)\t|\tAcc: 44.53%(train) 32\n",
      "epoch: 1 \tLoss: 42.38(valid)\t|\tAcc: 72.66%(valid) 32\n",
      "epoch: 2 \tLoss: 38.62(train)\t|\tAcc: 76.56%(train) 33\n",
      "epoch: 2 \tLoss: 44.04(valid)\t|\tAcc: 66.41%(valid) 33\n",
      "epoch: 2 \tLoss: 49.57(train)\t|\tAcc: 64.06%(train) 34\n",
      "epoch: 2 \tLoss: 41.53(valid)\t|\tAcc: 76.56%(valid) 34\n",
      "epoch: 2 \tLoss: 46.55(train)\t|\tAcc: 69.53%(train) 35\n",
      "epoch: 2 \tLoss: 42.76(valid)\t|\tAcc: 73.44%(valid) 35\n",
      "epoch: 2 \tLoss: 43.21(train)\t|\tAcc: 79.69%(train) 36\n",
      "epoch: 2 \tLoss: 38.52(valid)\t|\tAcc: 78.12%(valid) 36\n",
      "epoch: 2 \tLoss: 52.69(train)\t|\tAcc: 65.62%(train) 37\n",
      "epoch: 2 \tLoss: 43.92(valid)\t|\tAcc: 77.34%(valid) 37\n",
      "epoch: 2 \tLoss: 43.76(train)\t|\tAcc: 78.12%(train) 38\n",
      "epoch: 2 \tLoss: 38.76(valid)\t|\tAcc: 75.78%(valid) 38\n",
      "epoch: 2 \tLoss: 45.87(train)\t|\tAcc: 70.31%(train) 39\n",
      "epoch: 2 \tLoss: 42.19(valid)\t|\tAcc: 75.78%(valid) 39\n",
      "epoch: 2 \tLoss: 38.02(train)\t|\tAcc: 82.03%(train) 40\n",
      "epoch: 2 \tLoss: 39.79(valid)\t|\tAcc: 71.88%(valid) 40\n",
      "epoch: 2 \tLoss: 43.97(train)\t|\tAcc: 71.09%(train) 41\n",
      "epoch: 2 \tLoss: 40.61(valid)\t|\tAcc: 75.78%(valid) 41\n",
      "epoch: 2 \tLoss: 44.67(train)\t|\tAcc: 69.53%(train) 42\n",
      "epoch: 2 \tLoss: 36.01(valid)\t|\tAcc: 80.47%(valid) 42\n",
      "epoch: 2 \tLoss: 37.85(train)\t|\tAcc: 81.25%(train) 43\n",
      "epoch: 2 \tLoss: 31.40(valid)\t|\tAcc: 79.69%(valid) 43\n",
      "epoch: 2 \tLoss: 41.30(train)\t|\tAcc: 72.66%(train) 44\n",
      "epoch: 2 \tLoss: 42.73(valid)\t|\tAcc: 75.00%(valid) 44\n",
      "epoch: 2 \tLoss: 44.48(train)\t|\tAcc: 70.31%(train) 45\n",
      "epoch: 2 \tLoss: 31.67(valid)\t|\tAcc: 83.59%(valid) 45\n",
      "epoch: 2 \tLoss: 34.28(train)\t|\tAcc: 77.34%(train) 46\n",
      "epoch: 2 \tLoss: 40.18(valid)\t|\tAcc: 75.78%(valid) 46\n",
      "epoch: 2 \tLoss: 36.39(train)\t|\tAcc: 78.12%(train) 47\n",
      "epoch: 2 \tLoss: 38.34(valid)\t|\tAcc: 76.56%(valid) 47\n",
      "epoch: 2 \tLoss: 43.60(train)\t|\tAcc: 73.44%(train) 48\n",
      "epoch: 2 \tLoss: 34.82(valid)\t|\tAcc: 78.91%(valid) 48\n",
      "epoch: 2 \tLoss: 36.67(train)\t|\tAcc: 82.03%(train) 49\n",
      "epoch: 2 \tLoss: 34.09(valid)\t|\tAcc: 76.56%(valid) 49\n",
      "epoch: 2 \tLoss: 37.96(train)\t|\tAcc: 75.78%(train) 50\n",
      "epoch: 2 \tLoss: 35.39(valid)\t|\tAcc: 78.91%(valid) 50\n",
      "epoch: 2 \tLoss: 35.23(train)\t|\tAcc: 83.59%(train) 51\n",
      "epoch: 2 \tLoss: 35.21(valid)\t|\tAcc: 78.12%(valid) 51\n",
      "epoch: 2 \tLoss: 40.82(train)\t|\tAcc: 72.66%(train) 52\n",
      "epoch: 2 \tLoss: 35.76(valid)\t|\tAcc: 75.00%(valid) 52\n",
      "epoch: 2 \tLoss: 35.25(train)\t|\tAcc: 75.78%(train) 53\n",
      "epoch: 2 \tLoss: 30.25(valid)\t|\tAcc: 81.25%(valid) 53\n",
      "epoch: 2 \tLoss: 38.15(train)\t|\tAcc: 75.78%(train) 54\n",
      "epoch: 2 \tLoss: 29.77(valid)\t|\tAcc: 85.94%(valid) 54\n",
      "epoch: 2 \tLoss: 39.91(train)\t|\tAcc: 74.22%(train) 55\n",
      "epoch: 2 \tLoss: 31.22(valid)\t|\tAcc: 86.72%(valid) 55\n",
      "epoch: 2 \tLoss: 35.35(train)\t|\tAcc: 78.91%(train) 56\n",
      "epoch: 2 \tLoss: 33.67(valid)\t|\tAcc: 78.91%(valid) 56\n",
      "epoch: 2 \tLoss: 43.05(train)\t|\tAcc: 71.88%(train) 57\n",
      "epoch: 2 \tLoss: 36.88(valid)\t|\tAcc: 76.56%(valid) 57\n",
      "epoch: 2 \tLoss: 35.58(train)\t|\tAcc: 85.16%(train) 58\n",
      "epoch: 2 \tLoss: 34.51(valid)\t|\tAcc: 79.69%(valid) 58\n",
      "epoch: 2 \tLoss: 34.79(train)\t|\tAcc: 79.69%(train) 59\n",
      "epoch: 2 \tLoss: 39.19(valid)\t|\tAcc: 77.34%(valid) 59\n",
      "epoch: 2 \tLoss: 31.18(train)\t|\tAcc: 82.81%(train) 60\n",
      "epoch: 2 \tLoss: 32.54(valid)\t|\tAcc: 80.47%(valid) 60\n",
      "epoch: 2 \tLoss: 33.41(train)\t|\tAcc: 79.69%(train) 61\n",
      "epoch: 2 \tLoss: 39.83(valid)\t|\tAcc: 78.91%(valid) 61\n",
      "epoch: 2 \tLoss: 32.83(train)\t|\tAcc: 80.47%(train) 62\n",
      "epoch: 2 \tLoss: 42.29(valid)\t|\tAcc: 75.00%(valid) 62\n",
      "epoch: 2 \tLoss: 37.03(train)\t|\tAcc: 80.47%(train) 63\n",
      "epoch: 2 \tLoss: 42.20(valid)\t|\tAcc: 77.34%(valid) 63\n",
      "epoch: 2 \tLoss: 37.74(train)\t|\tAcc: 39.84%(train) 64\n",
      "epoch: 2 \tLoss: 37.84(valid)\t|\tAcc: 78.91%(valid) 64\n",
      "epoch: 3 \tLoss: 42.31(train)\t|\tAcc: 76.56%(train) 65\n",
      "epoch: 3 \tLoss: 37.54(valid)\t|\tAcc: 77.34%(valid) 65\n",
      "epoch: 3 \tLoss: 38.27(train)\t|\tAcc: 74.22%(train) 66\n",
      "epoch: 3 \tLoss: 35.29(valid)\t|\tAcc: 82.03%(valid) 66\n",
      "epoch: 3 \tLoss: 36.27(train)\t|\tAcc: 78.91%(train) 67\n",
      "epoch: 3 \tLoss: 28.40(valid)\t|\tAcc: 85.16%(valid) 67\n",
      "epoch: 3 \tLoss: 43.27(train)\t|\tAcc: 69.53%(train) 68\n",
      "epoch: 3 \tLoss: 43.02(valid)\t|\tAcc: 74.22%(valid) 68\n",
      "epoch: 3 \tLoss: 35.33(train)\t|\tAcc: 78.91%(train) 69\n",
      "epoch: 3 \tLoss: 35.08(valid)\t|\tAcc: 80.47%(valid) 69\n",
      "epoch: 3 \tLoss: 30.93(train)\t|\tAcc: 83.59%(train) 70\n",
      "epoch: 3 \tLoss: 37.12(valid)\t|\tAcc: 80.47%(valid) 70\n",
      "epoch: 3 \tLoss: 59.35(train)\t|\tAcc: 63.28%(train) 71\n",
      "epoch: 3 \tLoss: 54.49(valid)\t|\tAcc: 72.66%(valid) 71\n",
      "epoch: 3 \tLoss: 56.78(train)\t|\tAcc: 70.31%(train) 72\n",
      "epoch: 3 \tLoss: 39.22(valid)\t|\tAcc: 77.34%(valid) 72\n",
      "epoch: 3 \tLoss: 42.92(train)\t|\tAcc: 75.78%(train) 73\n",
      "epoch: 3 \tLoss: 32.24(valid)\t|\tAcc: 82.81%(valid) 73\n",
      "epoch: 3 \tLoss: 37.57(train)\t|\tAcc: 73.44%(train) 74\n",
      "epoch: 3 \tLoss: 32.85(valid)\t|\tAcc: 79.69%(valid) 74\n",
      "epoch: 3 \tLoss: 36.96(train)\t|\tAcc: 75.00%(train) 75\n",
      "epoch: 3 \tLoss: 28.35(valid)\t|\tAcc: 85.16%(valid) 75\n",
      "epoch: 3 \tLoss: 35.07(train)\t|\tAcc: 78.12%(train) 76\n",
      "epoch: 3 \tLoss: 39.21(valid)\t|\tAcc: 75.00%(valid) 76\n",
      "epoch: 3 \tLoss: 38.02(train)\t|\tAcc: 80.47%(train) 77\n",
      "epoch: 3 \tLoss: 40.74(valid)\t|\tAcc: 73.44%(valid) 77\n",
      "epoch: 3 \tLoss: 36.00(train)\t|\tAcc: 82.03%(train) 78\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 3 \tLoss: 37.75(valid)\t|\tAcc: 75.78%(valid) 78\n",
      "epoch: 3 \tLoss: 33.74(train)\t|\tAcc: 82.03%(train) 79\n",
      "epoch: 3 \tLoss: 30.43(valid)\t|\tAcc: 83.59%(valid) 79\n",
      "epoch: 3 \tLoss: 30.11(train)\t|\tAcc: 78.91%(train) 80\n",
      "epoch: 3 \tLoss: 30.86(valid)\t|\tAcc: 79.69%(valid) 80\n",
      "epoch: 3 \tLoss: 35.64(train)\t|\tAcc: 78.91%(train) 81\n",
      "epoch: 3 \tLoss: 39.51(valid)\t|\tAcc: 76.56%(valid) 81\n",
      "epoch: 3 \tLoss: 37.20(train)\t|\tAcc: 75.78%(train) 82\n",
      "epoch: 3 \tLoss: 32.95(valid)\t|\tAcc: 81.25%(valid) 82\n",
      "epoch: 3 \tLoss: 35.23(train)\t|\tAcc: 80.47%(train) 83\n",
      "epoch: 3 \tLoss: 34.29(valid)\t|\tAcc: 78.91%(valid) 83\n",
      "epoch: 3 \tLoss: 35.13(train)\t|\tAcc: 79.69%(train) 84\n",
      "epoch: 3 \tLoss: 27.50(valid)\t|\tAcc: 82.03%(valid) 84\n",
      "epoch: 3 \tLoss: 38.06(train)\t|\tAcc: 79.69%(train) 85\n",
      "epoch: 3 \tLoss: 30.58(valid)\t|\tAcc: 80.47%(valid) 85\n",
      "epoch: 3 \tLoss: 40.42(train)\t|\tAcc: 74.22%(train) 86\n",
      "epoch: 3 \tLoss: 36.16(valid)\t|\tAcc: 77.34%(valid) 86\n",
      "epoch: 3 \tLoss: 33.85(train)\t|\tAcc: 81.25%(train) 87\n",
      "epoch: 3 \tLoss: 30.28(valid)\t|\tAcc: 83.59%(valid) 87\n",
      "epoch: 3 \tLoss: 34.21(train)\t|\tAcc: 80.47%(train) 88\n",
      "epoch: 3 \tLoss: 36.02(valid)\t|\tAcc: 78.12%(valid) 88\n",
      "epoch: 3 \tLoss: 33.50(train)\t|\tAcc: 82.81%(train) 89\n",
      "epoch: 3 \tLoss: 35.06(valid)\t|\tAcc: 79.69%(valid) 89\n",
      "epoch: 3 \tLoss: 34.30(train)\t|\tAcc: 79.69%(train) 90\n",
      "epoch: 3 \tLoss: 31.90(valid)\t|\tAcc: 82.81%(valid) 90\n"
     ]
    }
   ],
   "source": [
    "# save losses for debugging/plotting\n",
    "train_losses = []\n",
    "train_acc = []\n",
    "train_iter = []\n",
    "val_losses = []\n",
    "val_acc = []\n",
    "val_iter = []\n",
    "\n",
    "\n",
    "for epoch in range(EPOCHS):\n",
    "    \n",
    "    for iter, (vectors, targets) in enumerate(trainloader):\n",
    "        \n",
    "        # Set train mode\n",
    "        model.train()\n",
    "\n",
    "        # Zero out previous gradients\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # Predict sentiment probabilities\n",
    "        probs, logits = model(vectors)\n",
    "\n",
    "        # Compute loss and accuracy\n",
    "        loss = ((probs -  targets)**2).sum()\n",
    "\n",
    "\n",
    "        # Get the predicted labels\n",
    "        preds = probs.argmax(dim=1)\n",
    "        targets = targets.argmax(dim=1)\n",
    "        \n",
    "        # Compute the prediction accuracy\n",
    "        accuracy = (preds == targets).sum()\n",
    "        accuracy = accuracy.get().float_precision()\n",
    "        accuracy = 100 * (accuracy / BATCH_SIZE)\n",
    "        \n",
    "        # Backpropagate the loss\n",
    "        loss.backward()\n",
    "\n",
    "        # Update weights\n",
    "        optimizer.step()\n",
    "\n",
    "        # Decrypt the loss for logging\n",
    "        loss = loss.get().float_precision()\n",
    "        \n",
    "        # get iteration number\n",
    "        train_i = 1 + epoch * len(trainloader) + iter \n",
    "        \n",
    "        # append to training losses for plotting\n",
    "        train_losses.append(loss.item())\n",
    "        train_iter.append(train_i)\n",
    "        train_acc.append(accuracy)\n",
    "\n",
    "        # print progress in training    \n",
    "        print(\"epoch:\",epoch+1,f'\\tLoss: {loss:.2f}(train)\\t|\\tAcc: {accuracy:.2f}%(train)', train_i)    \n",
    "        \n",
    "        \n",
    "        # Log to Tensorboard\n",
    "        writer.add_scalar('train/loss', loss, train_i)\n",
    "        writer.add_scalar('train/acc', accuracy, train_i)\n",
    "\n",
    "        # break if over 100 iterations to save time\n",
    "        if train_i>100:\n",
    "            break\n",
    "        \n",
    "        \"\"\" Perform validation on exactly one batch \"\"\"\n",
    "        \n",
    "        # Set validation mode\n",
    "        model.eval()\n",
    "\n",
    "        for vectors, targets in valloader:\n",
    "            \n",
    "            probs, logits = model(vectors)\n",
    "\n",
    "            loss = ((probs -  targets)**2).sum()\n",
    "\n",
    "            preds = probs.argmax(dim=1)\n",
    "            targets = targets.argmax(dim=1)\n",
    "\n",
    "            accuracy = preds.eq(targets).sum()\n",
    "            accuracy = accuracy.get().float_precision()\n",
    "            accuracy = 100 * (accuracy / BATCH_SIZE)\n",
    "\n",
    "            # Decrypt loss for logging/plotting\n",
    "            loss = loss.get().float_precision()\n",
    "            \n",
    "            # get iteration    \n",
    "            val_i = 1 + epoch * len(trainloader) + iter \n",
    "            \n",
    "            # append to validation losses for plotting\n",
    "            val_losses.append(loss.item())\n",
    "            val_iter.append(val_i) \n",
    "            val_acc.append(accuracy)\n",
    "            \n",
    "            # print progress in validation                        \n",
    "            print(\"epoch:\",epoch+1,f'\\tLoss: {loss:.2f}(valid)\\t|\\tAcc: {accuracy:.2f}%(valid)', val_i)\n",
    "            \n",
    "            # Log to tensorboard\n",
    "            writer.add_scalar('val/loss', loss, val_i)\n",
    "            writer.add_scalar('val/acc', accuracy, val_i)\n",
    "            \n",
    "            break\n",
    "\n",
    "            \n",
    "writer.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Results\n",
    "Now that training is finished, you can verify that both Bob and Alice have `SubPipeline` objects on their machines containing the pipeline components defined previously in **Figure 2**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loss\n",
    "![Loss](fig_data/loss.png)\n",
    "\n",
    "### Accuracy\n",
    "![Accuracy](fig_data/accuracy.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# On Alices's machine\n",
    "[alice._objects[id] for id in alice._objects if  isinstance(alice._objects[id], syfertext.SubPipeline)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# On bob's machine\n",
    "[bob._objects[id] for id in bob._objects if  isinstance(bob._objects[id], syfertext.SubPipeline)]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}