notebooks/phish_detection/Phishing_Detection_using_Bert_CLX.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Phishing Detection Using BERT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Authors\n",
    "- Eli Fajardo (NVIDIA)\n",
    "- Gorkem Batmaz (NVIDIA)\n",
    "- Bartley Richardson, PhD (NVIDIA)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Table of Contents \n",
    "* Introduction\n",
    "* List of datasets used\n",
    "* Reading in the datasets\n",
    "* Initialize CLX Phishing Detection and BERT model\n",
    "* Training - CLAIR FRAUDULENT EMAILS dataset\n",
    "* Evaluation of CLAIR Test Set\n",
    "* Training with the the SPAM_ASSASSIN dataset\n",
    "* Evaluation of the SPAM_ASSASSIN Test Set\n",
    "* Training with all three datasets CLAIR+SPAM_ASSASSIN+ENRON\n",
    "* Evaluation of the Test Set of CLAIR+SPAM_ASSASSIN+ENRON Datasets\n",
    "* References"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "Phishing is a method used by fraudsters/hackers to obtain sensitive information from email users by pretending to be from legitimate institutions/people.\n",
    "Various machine learning methods are in use to detect and filter phishing/spam emails. \n",
    "In this notebook, we show how to train a *BERT language model and analyse the performance on various datasets. We have fine-tuned a pre-trained BERT model with a classification layer using HuggingFace library. \n",
    "*BERT stands for Bidirectional Encoder Representations from Transformers. The paper can be found [here.](https://arxiv.org/pdf/1810.04805.pdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Datasets used\n",
    "* [CLAIR-Fraudulent E-mail Corpus](https://www.kaggle.com/rtatman/fraudulent-email-corpus)\n",
    "* [SPAM_ASSASSIN Dataset](https://spamassassin.apache.org/old/publiccorpus/)\n",
    "* [Enron Emails](https://www.cs.cmu.edu/~enron/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Required Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cudf;\n",
    "from cuml.preprocessing.model_selection import train_test_split;\n",
    "from clx.analytics.phishing_detector import PhishingDetector;\n",
    "import s3fs;\n",
    "from os import path;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reading the files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "CLAIR_TSV = \"Phishing_Dataset_Clair_Collection.tsv\"\n",
    "SPAM_TSV = \"spam_assassin_spam_200_20021010.tsv\"\n",
    "EASY_HAM_TSV = \"spam_assassin_easyham_200_20021010.tsv\"\n",
    "HARD_HAM_TSV = \"spam_assassin_hardham_200_20021010.tsv\"\n",
    "ENRON_TSV = \"enron_10000.tsv\"\n",
    "\n",
    "S3_BASE_PATH = \"rapidsai-data/cyber/clx\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clair dataset\n",
    "if not path.exists(CLAIR_TSV):\n",
    "    fs = s3fs.S3FileSystem(anon=True)\n",
    "    fs.get(S3_BASE_PATH + \"/\" + CLAIR_TSV, CLAIR_TSV)\n",
    "    \n",
    "dfclair = cudf.read_csv(CLAIR_TSV, delimiter='\\t', header=None, names=['label', 'email']).dropna()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Phishing emails of the SPAM ASSASSIN dataset\n",
    "if not path.exists(SPAM_TSV):\n",
    "    fs = s3fs.S3FileSystem(anon=True)\n",
    "    fs.get(S3_BASE_PATH + \"/\" + SPAM_TSV, SPAM_TSV)\n",
    " \n",
    "dfspam = cudf.read_csv(SPAM_TSV, delimiter='\\t', header=None, names=['label', 'email'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Benign emails of the SPAM ASSASSIN dataset\n",
    "if not path.exists(EASY_HAM_TSV):\n",
    "    fs = s3fs.S3FileSystem(anon=True)\n",
    "    fs.get(S3_BASE_PATH + \"/\" + EASY_HAM_TSV, EASY_HAM_TSV)\n",
    "    \n",
    "dfeasyham = cudf.read_csv(EASY_HAM_TSV, delimiter='\\t', header=None, names=['label', 'email'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Benign emails of the SPAM ASSASSIN dataset that are easy to be confused with phishing emails\n",
    "if not path.exists(HARD_HAM_TSV):\n",
    "    fs = s3fs.S3FileSystem(anon=True)\n",
    "    fs.get(S3_BASE_PATH + \"/\" + HARD_HAM_TSV, HARD_HAM_TSV)\n",
    "\n",
    "dfhardham = cudf.read_csv(HARD_HAM_TSV, delimiter='\\t', header=None, names=['label', 'email'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Benign Enron emails\n",
    "if not path.exists(ENRON_TSV):\n",
    "    fs = s3fs.S3FileSystem(anon=True)\n",
    "    fs.get(S3_BASE_PATH + \"/\" + ENRON_TSV, ENRON_TSV)\n",
    "\n",
    "dfenron = cudf.read_csv(ENRON_TSV, delimiter='\\t', header=None, names=['label', 'email'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The files contain the first 200 words of each email. The model uses only the first 128 words."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize/Load BERT model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "774d2fa14b24490380fb01beb69abb80",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=433.0, style=ProgressStyle(description_…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "98bd36fa854641c2b775bc2764a9707f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=440473133.0, style=ProgressStyle(descri…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertForSequenceClassification: ['cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias']\n",
      "- This IS expected if you are initializing BertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPretraining model).\n",
      "- This IS NOT expected if you are initializing BertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-uncased and are newly initialized: ['classifier.weight', 'classifier.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "phish_detect = PhishingDetector()\n",
    "phish_detect.init_model()\n",
    "\n",
    "# init_model can also load pre-trained model by passing it model directory path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training - CLAIR FRAUDULENT EMAILS DATASET"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Clair fraudulent e-mails contain criminally deceptive information, usually with the intent of convincing the recipient to give the sender a large amount of money.*\n",
    "The original dataset can be found [here](https://www.kaggle.com/rtatman/fraudulent-email-corpus). \n",
    "\n",
    "Split the dataset into training and test sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(dfclair, 'label', train_size=0.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:   0%|          | 0/1 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train loss: 0.0781611256509879\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch: 100%|██████████| 1/1 [00:59<00:00, 59.21s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Accuracy: 0.9947916666666666\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "phish_detect.train_model(X_train, y_train, epochs=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation of CLAIR Test Set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9974853310980721"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "phish_detect.evaluate_model(X_test, y_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This indicates over 99% of the emails in the test dataset were correctly marked as fraudulent or benign."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training with SPAM_ASSASSIN dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SPAM_ASSASIN public mail corpus consists of three classes. The first category is \"easy ham\" which has benign emails; the second category is \"hard ham\" which has emails which are benign but are closer to spam, and finally, the spam category that includes the fraudulent emails. More info [here](https://spamassassin.apache.org/old/publiccorpus/readme.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Merging the spam assasin dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_assassin = cudf.concat([dfhardham,dfeasyham,dfspam], ignore_index=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Split the dataset into train and test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(df_assassin, 'label', train_size=0.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:   0%|          | 0/1 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train loss: 0.44632136821746826\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch: 100%|██████████| 1/1 [00:16<00:00, 16.62s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Accuracy: 0.862348615916955\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "phish_detect.train_model(X_train, y_train, epochs=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation of the SPAM_ASSASSIN Test Set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.8366111951588502"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "phish_detect.evaluate_model(X_test, y_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "~84% of the spam assassin emails were predicted correctly as spam or benign"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "## Training with CLAIR+SPAM_ASSASSIN datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Merge the two datasets and split as train and test sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_total = cudf.concat([dfhardham,dfeasyham,dfspam,dfclair],ignore_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(df_total, 'label', train_size=0.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:   0%|          | 0/1 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train loss: 0.04908236236842044\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch: 100%|██████████| 1/1 [01:15<00:00, 75.63s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Accuracy: 0.9959415584415584\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "phish_detect.train_model(X_train, y_train, epochs=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation of the Test Set of CLAIR+SPAM_ASSASSIN Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9983585029546946"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "phish_detect.evaluate_model(X_test, y_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When CLAIR and SPAM_ASSASSIN datasets are combined, the test set accuracy is over 99%, which means 99% of the emails in the test set are marked correctly as benign or spam."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training with all three datasets (CLAIR+SPAM_ASSASSIN+ENRON)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Enron Email Dataset was collected and prepared by the CALO Project (A Cognitive Assistant that Learns and Organizes). We use them as benign email class. \n",
    "More info on the ENRON Email Dataset can be found [here](https://www.cs.cmu.edu/~./enron/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Merge all the datasets, split into training and test set and then tokenize the emails"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_total = cudf.concat([dfhardham,dfeasyham,dfspam,dfclair,dfenron],ignore_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(df_total, 'label', train_size=0.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:   0%|          | 0/1 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train loss: 0.01176047741604328\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch: 100%|██████████| 1/1 [02:06<00:00, 126.25s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Accuracy: 0.999261811023622\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "phish_detect.train_model(X_train, y_train, epochs=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation of the Test Set of CLAIR+SPAM_ASSASSIN+ENRON Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9988109393579072"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "phish_detect.evaluate_model(X_test, y_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When all three datasets are combined, over 99% of the emails are labelled correctly. This shows that using a BERT-based phishing detector performs well in identifying the spam emails across these datasets. This notebook is prepared as a POC. Users can experiment with other datasets, increase the coverage and change the number of epochs to fine-tune the results on their datasets. It is also an example of how CLX can be used with huggingface and other libraries to create custom solutions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# References\n",
    "* Radev, D. (2008), CLAIR collection of fraud email, ACL Data and Code Repository, ADCR2008T001, http://aclweb.org/aclwiki\n",
    "* https://www.kaggle.com/rtatman/fraudulent-email-corpus *\n",
    "* https://www.cs.cmu.edu/~./enron/\n",
    "* https://spamassassin.apache.org/old/publiccorpus/readme.html\n",
    "* https://github.com/huggingface/transformers/tree/master/examples#\n",
    "* https://www.depends-on-the-definition.com/named-entity-recognition-with-bert/\n",
    "* https://github.com/ThilinaRajapakse/pytorch-transformers-classification\n",
    "* https://mccormickml.com/2019/07/22/BERT-fine-tuning/"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Copy of BERT Fine-Tuning Sentence Classification.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}