HomomorphicEncryption.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Homomorphic Encryption Example\n",
    "\n",
    "## This demo is a copy from the following github repo:\n",
    "https://github.com/anujdutt9/Encrypted-Machine-Learning/blob/master/Homomorphic-Encryption-Demo.ipynb\n",
    "+ Might need to install library: !pip install phe\n",
    "\n",
    "+Logistic Regression for Spam/Not Spam e-mail Classification.\n",
    "\n",
    "For this problem we have two users:\n",
    "\n",
    "USER-1\n",
    "\n",
    "USER-2\n",
    "\n",
    "AI Inc. makes a Machine Learning model that is trained on some email data for classification between Spam/Not Spam. Now, they want to take that model, encrypt it and send to USER-1 and USER-2 who will train the model on their data, fully Homomorphically Encrypted, and send the trained, a bit better model back to AI Inc.\n",
    "\n",
    "In this process, AI Inc. get a better trained model every time without even looking at USER-1 or USER-2 data. This way AI Inc. can serve the customers better with a smart Machine Learning model and the USER has complete control of his/her data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import Dependencies\n",
    "import phe as paillier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create Public and Private Keys\n",
    "key_length = 1024\n",
    "pub_key, private_key = paillier.generate_paillier_keypair(n_length=key_length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<PaillierPublicKey 7a26580b2d>"
      ]
     },
     "execution_count": 99,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pub_key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<PaillierPrivateKey for <PaillierPublicKey 7a26580b2d>>"
      ]
     },
     "execution_count": 100,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "private_key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a:  10\n",
      "Encrypted a:  <phe.paillier.EncryptedNumber object at 0x7f5715d8a450>\n",
      "Encrypted a Public Key:  <PaillierPublicKey 7a26580b2d>\n"
     ]
    }
   ],
   "source": [
    "# Encrypt an operation using Public Key\n",
    "a = 10\n",
    "print(\"a: \",a)\n",
    "\n",
    "encrypted_a = pub_key.encrypt(a)\n",
    "print(\"Encrypted a: \",encrypted_a)\n",
    "\n",
    "print(\"Encrypted a Public Key: \", encrypted_a.public_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "b:  5\n",
      "Encrypted b:  <phe.paillier.EncryptedNumber object at 0x7f5715d8a950>\n",
      "Encrypted b Public Key:  <PaillierPublicKey 7a26580b2d>\n"
     ]
    }
   ],
   "source": [
    "# Encrypt another variable\n",
    "b = 5\n",
    "print(\"b: \", b)\n",
    "\n",
    "encrypted_b = pub_key.encrypt(b)\n",
    "print(\"Encrypted b: \", encrypted_b)\n",
    "\n",
    "print(\"Encrypted b Public Key: \",encrypted_b.public_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "c:  15\n"
     ]
    }
   ],
   "source": [
    "# Do an operation on Encrypted Variables\n",
    "c = a + b\n",
    "print(\"c: \", c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "d:  50\n"
     ]
    }
   ],
   "source": [
    "d = a * b\n",
    "print(\"d: \",d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Encrypted e:  <phe.paillier.EncryptedNumber object at 0x7f5715d99410>\n"
     ]
    }
   ],
   "source": [
    "e = a - b\n",
    "\n",
    "encrypted_e = pub_key.encrypt(e)\n",
    "print(\"Encrypted e: \", encrypted_e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Decrypt the Encrypted Data\n",
    "decrypted_e = private_key.decrypt(encrypted_e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Decrypted e:  5\n"
     ]
    }
   ],
   "source": [
    "print(\"Decrypted e: \", decrypted_e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Import Dependencies\n",
    "\n",
    "import time\n",
    "import os.path\n",
    "from zipfile import ZipFile\n",
    "from urllib.request import urlopen\n",
    "from contextlib import contextmanager\n",
    "\n",
    "import numpy as np\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.feature_extraction.text import CountVectorizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "#download dataset. this only needs to be run the first time.\n",
    "\n",
    "url = [\n",
    "    'https://cloudstor.aarnet.edu.au/plus/index.php/s/RpHZ57z2E3BTiSQ/download',\n",
    "    'https://cloudstor.aarnet.edu.au/plus/index.php/s/QVD4Xk5Cz3UVYLp/download'\n",
    "]\n",
    "\n",
    "\n",
    "def download_data():\n",
    "    \"\"\"Download two sets of Enron1 spam/ham e-mails if they are not here\n",
    "    We will use the first as trainset and the second as testset.\n",
    "    Return the path prefix to us to load the data from disk.\"\"\"\n",
    "\n",
    "    n_datasets = 2\n",
    "    for d in range(1, n_datasets + 1):\n",
    "        if not os.path.isdir('./enron%d' % d):\n",
    "\n",
    "            URL = url[d-1]\n",
    "            print(\"Downloading %d/%d: %s\" % (d, n_datasets, URL))\n",
    "            folderzip = './enron%d.zip' % d\n",
    "\n",
    "            with urlopen(URL) as remotedata:\n",
    "                with open(folderzip, 'wb') as z:\n",
    "                    z.write(remotedata.read())\n",
    "\n",
    "            with ZipFile(folderzip) as z:\n",
    "                z.extractall()\n",
    "            os.remove(folderzip)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading 1/2: https://cloudstor.aarnet.edu.au/plus/index.php/s/RpHZ57z2E3BTiSQ/download\n",
      "Downloading 2/2: https://cloudstor.aarnet.edu.au/plus/index.php/s/QVD4Xk5Cz3UVYLp/download\n"
     ]
    }
   ],
   "source": [
    "#only run if we need to get data\n",
    "\n",
    "download_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data Preprocessing\n",
    "def preprocess_data():\n",
    "    \"\"\"\n",
    "    Load the email dataset and Represent them as bag-of-words.\n",
    "    Shuffle and split train/test.\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"Importing dataset...\")\n",
    "    path = '/home/jupyter/enron1/ham/'\n",
    "    ham1 = [open(path + f, 'r', errors='replace').read().strip(r\"\\n\")\n",
    "            for f in os.listdir(path) if os.path.isfile(path + f)]\n",
    "    path = '/home/jupyter/enron1/spam/'\n",
    "    spam1 = [open(path + f, 'r', errors='replace').read().strip(r\"\\n\")\n",
    "             for f in os.listdir(path) if os.path.isfile(path + f)]\n",
    "    path = '/home/jupyter/enron2/ham/'\n",
    "    ham2 = [open(path + f, 'r', errors='replace').read().strip(r\"\\n\")\n",
    "            for f in os.listdir(path) if os.path.isfile(path + f)]\n",
    "    path = '/home/jupyter/enron2/spam/'\n",
    "    spam2 = [open(path + f, 'r', errors='replace').read().strip(r\"\\n\")\n",
    "             for f in os.listdir(path) if os.path.isfile(path + f)]\n",
    "\n",
    "    # Merge and create labels\n",
    "    emails = ham1 + spam1 + ham2 + spam2\n",
    "    y = np.array([-1] * len(ham1) + [1] * len(spam1) +\n",
    "                 [-1] * len(ham2) + [1] * len(spam2))\n",
    "\n",
    "    # Words count, keep only frequent words\n",
    "    # Minimum Document Word Frequency: 0.001\n",
    "    count_vect = CountVectorizer(decode_error='replace', stop_words='english', min_df=0.001)\n",
    "    X = count_vect.fit_transform(emails)\n",
    "\n",
    "    print('Vocabulary size: %d' % X.shape[1])\n",
    "\n",
    "    # Shuffle\n",
    "    perm = np.random.permutation(X.shape[0])\n",
    "    X, y = X[perm, :], y[perm]\n",
    "\n",
    "    # Split train and test\n",
    "    split = 500\n",
    "    X_train, X_test = X[-split:, :], X[:-split, :]\n",
    "    y_train, y_test = y[-split:], y[:-split]\n",
    "\n",
    "    print(\"Labels in trainset are {:.2f} spam : {:.2f} ham\".format(\n",
    "        np.mean(y_train == 1), np.mean(y_train == -1)))\n",
    "\n",
    "    return X_train, y_train, X_test, y_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [],
   "source": [
    "@contextmanager\n",
    "def timer():\n",
    "    \"\"\"Helper for measuring runtime\"\"\"\n",
    "\n",
    "    time0 = time.perf_counter()\n",
    "    yield\n",
    "    print('[elapsed time: %.2f s]' % (time.perf_counter() - time0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AI_Inc:\n",
    "    \"\"\"\n",
    "    AI Inc. Trains a Logistic Regression model on plaintext data, encrypts the model for remote use by USER-1 and USER-2,\n",
    "    decrypts encrypted scores using the paillier private key.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        self.model = LogisticRegression()\n",
    "\n",
    "    # Generate Public and Private Key Pairs\n",
    "    # Public Key is used to Encrypt the Data, Private Key to Decrypt\n",
    "    def generate_paillier_keypair(self, n_length):\n",
    "        self.pubkey, self.privkey = paillier.generate_paillier_keypair(n_length=n_length)\n",
    "\n",
    "    # Train the Model\n",
    "    def fit(self, X, y):\n",
    "        self.model = self.model.fit(X, y)\n",
    "\n",
    "    # Make Predictions for Email \"Spam/Not Spam\"\n",
    "    def predict(self, X):\n",
    "        return self.model.predict(X)\n",
    "\n",
    "    # Encypt the Coefficients for the Logistic Regression Equation\n",
    "    # Weights can tell about the data, so Encrypt them\n",
    "    # Equation: y = mX + b\n",
    "    def encrypt_weights(self):\n",
    "        coef = self.model.coef_[0, :]\n",
    "        encrypted_weights = [self.pubkey.encrypt(coef[i])\n",
    "                             for i in range(coef.shape[0])]\n",
    "        encrypted_intercept = self.pubkey.encrypt(self.model.intercept_[0])\n",
    "        return encrypted_weights, encrypted_intercept\n",
    "\n",
    "    # Decrypt the Scores for the Model\n",
    "    def decrypt_scores(self, encrypted_scores):\n",
    "        return [self.privkey.decrypt(s) for s in encrypted_scores]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now the USER-1 gets a trained model from AI Inc. and trains on its own data all using Homomorphic Encryption.\n",
    "class User_1:\n",
    "    \"\"\"\n",
    "    USER-1/USER-2 are given the encrypted model trained by AI Inc. and the public key.\n",
    "\n",
    "    Scores local plaintext data with the encrypted model, but cannot decrypt\n",
    "    the scores without the private key held by AI Inc..\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, pubkey):\n",
    "        self.pubkey = pubkey\n",
    "\n",
    "    # Set Initial Values of Coefficients\n",
    "    def set_weights(self, weights, intercept):\n",
    "        self.weights = weights\n",
    "        self.intercept = intercept\n",
    "\n",
    "    # Compute the Prediction Scores for the Model all while being totally Encrypted.\n",
    "    def encrypted_score(self, x):\n",
    "        \"\"\"Compute the score of `x` by multiplying with the encrypted model,\n",
    "        which is a vector of `paillier.EncryptedNumber`\"\"\"\n",
    "        score = self.intercept\n",
    "        _, idx = x.nonzero()\n",
    "        for i in idx:\n",
    "            score += x[0, i] * self.weights[i]\n",
    "        return score\n",
    "\n",
    "    # Get the Evaluation Scores for the Model\n",
    "    def encrypted_evaluate(self, X):\n",
    "        return [self.encrypted_score(X[i, :]) for i in range(X.shape[0])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Importing dataset...\n",
      "Vocabulary size: 7997\n",
      "Labels in trainset are 0.27 spam : 0.73 ham\n"
     ]
    }
   ],
   "source": [
    "# Get the Preprocessed Split Data\n",
    "X_train, y_train, X_test, y_test = preprocess_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AI Inc.: Generating Paillier Public Private Keypair\n"
     ]
    }
   ],
   "source": [
    "# Now firstly the AI Inc. Generates the Public and Private Keys\n",
    "print(\"AI Inc.: Generating Paillier Public Private Keypair\")\n",
    "ai_inc = AI_Inc()\n",
    "# NOTE: using smaller keys sizes wouldn't be cryptographically safe\n",
    "ai_inc.generate_paillier_keypair(n_length=1024)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AI Inc.: Training Initial Spam Classifier\n",
      "[elapsed time: 0.07 s]\n"
     ]
    }
   ],
   "source": [
    "print(\"AI Inc.: Training Initial Spam Classifier\")\n",
    "with timer() as t:\n",
    "    ai_inc.fit(X_train, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AI Inc.'s Classification on Test Data, what it would expect the performance to be on USER-1/2's data...\n",
      "[elapsed time: 0.00 s]\n",
      "Error 0.040\n"
     ]
    }
   ],
   "source": [
    "print(\"AI Inc.'s Classification on Test Data, what it would expect the performance to be on USER-1/2's data...\")\n",
    "with timer() as t:\n",
    "    error = np.mean(ai_inc.predict(X_test) != y_test)\n",
    "print(\"Error {:.3f}\".format(error))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AI Inc.: Encrypting Trained Classifier before sending to USER-1/2\n",
      "[elapsed time: 15.56 s]\n"
     ]
    }
   ],
   "source": [
    "print(\"AI Inc.: Encrypting Trained Classifier before sending to USER-1/2\")\n",
    "with timer() as t:\n",
    "    encrypted_weights, encrypted_intercept = ai_inc.encrypt_weights()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A Few Encrypted Weights:  [<phe.paillier.EncryptedNumber object at 0x7f5715551790>, <phe.paillier.EncryptedNumber object at 0x7f5715551150>, <phe.paillier.EncryptedNumber object at 0x7f5715551250>, <phe.paillier.EncryptedNumber object at 0x7f5716209b50>, <phe.paillier.EncryptedNumber object at 0x7f5714fc0050>]\n",
      "Encrypted Intercept:  <phe.paillier.EncryptedNumber object at 0x7f571542e410>\n"
     ]
    }
   ],
   "source": [
    "# Confirming the Weights are Encrypted\n",
    "print(\"A Few Encrypted Weights: \", encrypted_weights[0:5]) #Not showing full list\n",
    "print(\"Encrypted Intercept: \", encrypted_intercept)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[<phe.paillier.EncryptedNumber object at 0x7f5714fc0a90>]\n",
      "<phe.paillier.EncryptedNumber object at 0x7f5714fc0a90>\n"
     ]
    }
   ],
   "source": [
    "#TC trying to complete. User makes prediction on unecrypted data using encrypted model.  \n",
    "user1 = User_1(pubkey=ai_inc.pubkey)\n",
    "user1.set_weights(encrypted_weights, encrypted_intercept)\n",
    "evaluateValue = user1.encrypted_evaluate(X_test[0])\n",
    "\n",
    "print(evaluateValue)\n",
    "print(evaluateValue[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-16.413734225260374"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# AI company decrypts score\n",
    "ai_inc.privkey.decrypt(evaluateValue[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "environment": {
   "name": "pytorch-gpu.1-4.m48",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/pytorch-gpu.1-4:m48"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}