bert_tpu_tweet_model.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Copy of Text Classification with BERT, TPUs and Engineered Features - cleaned",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/solidgose/solidgose_experiments/blob/master/bert_tpu_tweet_model.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "metadata": {
        "id": "r7cyhQ7-ASqp",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# Text Classification with BERT, TPUs, and engineered features on Google Colab\n",
        "## Overview and Objective\n",
        "This Colab notebook demonstrates the ability to conduct fast text classification using the Tensor Processing Units (TPUs) provided through Google Colab.\n",
        "\n",
        "Using two publicly-downloadable text databases: the fivethirtyeight.com's list of Internet Research Agency's (IRA) associated tweets during the 2016 General Election in the United States, and George Washington University's database of tweets identified with the 2016 election, a BERT model will be fine-tuned and trained to determine if a tweet from these dataset is associated with the IRA (or not).\n",
        "\n",
        "This notebook uses the BERT model checkpoints, not the TensorFlow Hub-hosted BERT models. This is because the TensorFlow Hub model uses a tf.Placeholder/feed_dict that reduces performance on TPUs and is incompatible for conducting inference. This notebook instead uses tf.Dataset with TFRecords.\n",
        "\n",
        "For now, this notebook only uses the English BERT models and English examples. More languages may be added later.\n",
        "\n",
        "## Disclaimer and License\n",
        "This notebook is intended for demonstration and research purposes only, and is not intended to be used in any production or actual predictive capacity. Also, Google Colab is intended for research; any computationally intensive forks of this code should probably be operated on a paid GCS instance. \n",
        "\n",
        "Where code has been copied or altered from other sources, the code will be commented to indicate the original source of the code. \n",
        "\n",
        "The twitter data downloaded by this notebook contains expletives, sexual references, and other strong language. URLs contained in tweets may link to web pages with malware, pornography, or other nefarious content. (Or, a link to a perfectly normal web page!) These links are not accessed or followed in this notebook. Nevertheless, proceed with caution! \n",
        "\n",
        "Code utilized in this notebook is adopted from the Apache License 2.0, and thus this notebook falls under the same license.\n",
        "\n",
        "This notebook is not intended to serve in any production or actual capacity - this is for experimentation and demonstration purposes only.\n",
        "\n",
        "## Prerequisites\n",
        "**Google Colab**: This notebook is designed to run on Google Colab; in Google Colab make sure your Runtime -> Change Runtime Session is set to *Python 3* and that your hardware accelerator is *TPU*. Otherwise, this code won't run. This notebook has not been tested in other environments.\n",
        "\n",
        "**Google Cloud Storage**: A Google Cloud Compute storage account is required to operate this notebook. The account is required because TPUs can only access data from a Google Storage Container (gs://) and not a local file system. Information on establishing a Google Cloud account is here: (URL Goes Here). Typical use of this notebook will result in very minimal costs (< USD 1.00).\n",
        "\n",
        "**Twitter Developer Account**: While the fivethirtyeight.com tweets are available for download without a license, the GWU tweets must be downloaded from Twitter's API. This requires the use of a Twitter Developer Account - the \"free\" tier of this service is sufficient to run this demonstration. The public and private \"consumer\" and \"app\" keys are needed to pull this data. More information on this is available here: (info here)\n",
        "\n",
        "**Questions or Feedback**: Post an issue to https://github.com/solidgose\n"
      ]
    },
    {
      "metadata": {
        "id": "LzRYXhEQMq50",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "import pickle\n",
        "import pprint\n",
        "import json\n",
        "import re\n",
        "import pandas as pd\n",
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "from datetime import datetime\n",
        "from tqdm import tqdm\n",
        "from os import listdir\n",
        "from os.path import isfile, join, exists\n",
        "import os\n",
        "import requests\n",
        "import sys\n",
        "from glob import glob\n",
        "from itertools import chain\n",
        "import time\n",
        "import math\n",
        "import numpy as np\n",
        "import urllib.request\n",
        "!pip install twitter\n",
        "from twitter import *\n",
        "!pip install namegenerator\n",
        "import namegenerator\n",
        "!pip install bert-tensorflow\n",
        "import bert\n",
        "import ipywidgets as widgets\n",
        "from bert import run_classifier\n",
        "from bert import optimization\n",
        "from bert import tokenization\n",
        "from bert import modeling\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "bcdfRTF4uPR7",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# About Settings and Defaults\n",
        "\n",
        "\n",
        "### GCS Bucket\n",
        "A GS bucket is a \"gs://bucket\" (where bucket is your bucket) location. If the bucket is not present under your GCS project ID, a bucket will be created with this code.\n",
        "\n",
        "### tf_checkpoint_root\n",
        "This is the root location where new fine-tuned and trained BERT model checkpoints will be located. \n",
        "\n",
        "### data_loc\n",
        "Where the source data for this notebook will be downloaded.\n",
        "\n",
        "### data_cache\n",
        "Location for temporary files created for use by the TPU and other procedures.\n",
        "\n",
        "### project_id\n",
        "This is your Google Compute Services project-id. A GCS billable project with a project ID is required to utilize this workbook.\n",
        "\n",
        "### Twitter API IDs\n",
        "These IDs are required for the *twitter* pypi package to download the tweets for this exercise. A Twitter developer account is needed to complete this information. \n",
        "\n",
        "### What happens when you run this code block\n",
        "1. The global constants needed to set file locations for the notebook are created.\n",
        "\n",
        "2. This Google Colab instance will log into your Google Cloud Compute API, confirm your identity, and display an authentication token needed to re-enter back into this notebook. The token is needed for the notebook to access your Google Cloud Storage."
      ]
    },
    {
      "metadata": {
        "id": "ng7XJ6vEyckX",
        "colab_type": "code",
        "cellView": "both",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#@title Settings and defaults\n",
        "gcs_bucket = '<bucket name here>' #@param {type:\"string\"}\n",
        "tf_checkpoint_root = 'tf_checkpoint_root' #@param {type:\"string\"}\n",
        "#tf_hub_cache = 'tf_hub_cache' #@param {type:\"string\"}\n",
        "data_loc = 'data_loc' #@param {type:\"string\"}\n",
        "data_cache = 'data_cache' #@param {type:\"string\"}\n",
        "project_id = '<input project id>' #@param {type:\"string\"}\n",
        "consumer_twitter_api = '<input consumer twitter api>' #@param {type:\"string\"} \n",
        "consumer_twitter_api_private = '<input consumer twitter api private>' #@param {type:\"string\"} \n",
        "access_api = '<input access api>' #@param {type:\"string\"}\n",
        "access_api_private = '<input access api private>' #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "BUCKET_LOC = f'gs://{gcs_bucket}/'\n",
        "TF_CHECKPOINT_ROOT = f'gs://{gcs_bucket}/{tf_checkpoint_root}'\n",
        "DATA_DIR = f'gs://{gcs_bucket}/{data_loc}'\n",
        "#GS_TF_HUB_CACHE = f'gs://{gcs_bucket}/{tf_hub_cache}'\n",
        "DATA_CACHE = f'gs://{gcs_bucket}/{data_cache}'\n",
        "CONSUMER_TWITTER_API = consumer_twitter_api\n",
        "CONSUMER_TWITTER_API_PRIVATE = consumer_twitter_api_private\n",
        "ACCESS_TWITTER_API = access_api\n",
        "ACCESS_TWITTER_API_PRIVATE = access_api_private\n",
        "\n",
        "# Now credentials are set for all future sessions on this TPU.\n",
        "!gcloud config set project {project_id}\n",
        "\n",
        "from google.colab import auth\n",
        "auth.authenticate_user()\n",
        "\n",
        "assert 'COLAB_TPU_ADDR' in os.environ, 'ERROR: Not connected to a TPU runtime; please see the first cell in this notebook for instructions!'\n",
        "TPU_ADDRESS = 'grpc://' + os.environ['COLAB_TPU_ADDR']\n",
        "print('TPU address is', TPU_ADDRESS)\n",
        "with tf.Session(TPU_ADDRESS) as session:\n",
        "  print('TPU devices:')\n",
        "  pprint.pprint(session.list_devices())\n",
        "\n",
        "  # Upload credentials to TPU.\n",
        "  with open('/content/adc.json', 'r') as f:\n",
        "    auth_info = json.load(f)\n",
        "  tf.contrib.cloud.configure_gcs(session, credentials=auth_info)\n",
        "\n",
        "print(\"All defaults successfully established\")\n",
        "\n",
        "label_list = [0,1]\n",
        "use_TPU = True\n",
        "random_seed = 42"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "91XXBpT-NmjE",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Create any buckets or bucket folders needed for code execution.\n",
        "\n",
        "bucket_status = !gsutil ls {BUCKET_LOC}\n",
        "if \"AccessDeniedException\" in bucket_status:\n",
        "    !gsutil mb -p {project_id} {BUCKET_LOC}\n",
        "    print(f'Created {BUCKET_LOC}')\n",
        "else:\n",
        "    print(f'{BUCKET_LOC} already exists.')\n",
        "    \n",
        "if not tf.io.gfile.exists(TF_CHECKPOINT_ROOT):\n",
        "    print(f'Created {TF_CHECKPOINT_ROOT}')\n",
        "    tf.gfile.MakeDirs(TF_CHECKPOINT_ROOT)\n",
        "else:\n",
        "    print(f'{TF_CHECKPOINT_ROOT} already exists.')\n",
        "        \n",
        "if not tf.io.gfile.exists(DATA_DIR):\n",
        "    print(f'Created {DATA_DIR}')\n",
        "    tf.gfile.MakeDirs(DATA_DIR)\n",
        "else:\n",
        "    print(f'{DATA_DIR} already exists.')\n",
        "          \n",
        "if not tf.io.gfile.exists(DATA_CACHE):\n",
        "    print(f'Created {DATA_CACHE}')\n",
        "    tf.gfile.MakeDirs(DATA_CACHE)\n",
        "else:\n",
        "    print(f'{DATA_CACHE} already exists.')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "jwyuLHvS6kBM",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# As of Feb 2019, these are the BERT models available. This list may change as new models are introduced.\n",
        "\n",
        "bert_model_urls = [\"https://storage.googleapis.com/bert_models/2018_10_18/uncased_L-12_H-768_A-12.zip\",\n",
        "                   \"https://storage.googleapis.com/bert_models/2018_10_18/uncased_L-24_H-1024_A-16.zip\",\n",
        "                   \"https://storage.googleapis.com/bert_models/2018_10_18/cased_L-12_H-768_A-12.zip\",\n",
        "                   \"https://storage.googleapis.com/bert_models/2018_10_18/cased_L-24_H-1024_A-16.zip\"]\n",
        "\n",
        "# If there are no sub-folders located in your tensorflow checkpoint root,\n",
        "# this code will create one root to start with\n",
        "\n",
        "list_of_tf_roots = tf.io.gfile.listdir(TF_CHECKPOINT_ROOT)\n",
        "    \n",
        "if len(list_of_tf_roots) == 0:\n",
        "    print(\"There are no models present in your tf checkpoint root. \\n Create a new model using this name or a name of your choice:\")\n",
        "    cur_name = namegenerator.gen()\n",
        "    text_input = input(f'Enter a model name, or <ENTER> for {cur_name}:')\n",
        "    if text_input is '':\n",
        "        text_input=cur_name\n",
        "    print(f'Creating folder for {text_input}')\n",
        "    tf.gfile.MakeDirs(TF_CHECKPOINT_ROOT + '/' + text_input)\n",
        "    print(\"Folder created\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "F0SIz33XVLBs",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "The following code uses ipywidgets for dynamic selection of constants. (Google Colab does not support dynamic forms, unforuntately.) Running the code will allow you to select your BERT model and location for your model, for the code that will subsequently execute.\n",
        "\n",
        "If you don't know which model to select, try 768/cased first."
      ]
    },
    {
      "metadata": {
        "id": "Hx1lj2_zSwCg",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "\n",
        "print(\"Select a bert model to use\")\n",
        "bert_model_choice = widgets.Dropdown(options=bert_model_urls, value=\"https://storage.googleapis.com/bert_models/2018_10_18/cased_L-12_H-768_A-12.zip\")\n",
        "bert_model_choice\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "3CR3a4zTGr0R",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "print(\"Select a location for the tensorflow model.\")\n",
        "tf_checkpoint_sub_dir = widgets.Dropdown(options=tf.io.gfile.listdir(TF_CHECKPOINT_ROOT), value=tf.io.gfile.listdir(TF_CHECKPOINT_ROOT)[0])\n",
        "tf_checkpoint_sub_dir"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "HlUBCoZaGz2d",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "print(f'Do you want to download the BERT model into {tf_checkpoint_sub_dir.value}?')\n",
        "should_download = widgets.Dropdown(options=['Yes','No'], value='Yes')\n",
        "should_download"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "ztT4dr8uImiU",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def create_and_populate_checkpoint_folder(subfolder_name, bert_model_name_with_zip):\n",
        "    try:\n",
        "        \n",
        "        # Check to see if there is a trained model in this folder, in this case,\n",
        "        # halt and ask the user what to do.\n",
        "        \n",
        "        if tf.gfile.Exists(TF_CHECKPOINT_ROOT + '/' + subfolder_name +  'checkpoint'):\n",
        "            print(\"There appears to be an existing checkpoint in this sub-folder.\")\n",
        "            print(\"Do you want to wipe this subfolder and insert a new BERT step 0 checkpoint?\")\n",
        "            answer = input(\"Type \\'yes\\' and <Enter> if you want to wipe. Type any other character to not wipe.'\")\n",
        "            if answer == \"yes\":\n",
        "                !gsutil rm {TF_CHECKPOINT_ROOT + '/' + subfolder_name + '*'}\n",
        "            else:\n",
        "                print(\"Ending function early.\")\n",
        "                return True\n",
        "\n",
        "        \n",
        "        bert_file_name = bert_model_name_with_zip.split('/')[-1]\n",
        "        print(f'Retrieving {bert_file_name}')\n",
        "\n",
        "        # Check to see if this file has already been downloaded. If so,\n",
        "        # skip the download step\n",
        "\n",
        "        if not exists(bert_model_choice.value):\n",
        "            urllib.request.urlretrieve(bert_model_choice.value, bert_file_name)\n",
        "            print(f'{bert_file_name} successfully downloaded')\n",
        "        else:\n",
        "            print(f'{bert_file_name} already downloaded')\n",
        "        # Re-unzip the bert model, just in case something funky happened\n",
        "        !unzip -o {bert_file_name}\n",
        "        !cd {bert_file_name.split('.')[0]};gsutil cp * {TF_CHECKPOINT_ROOT + '/' + subfolder_name}\n",
        "\n",
        "        return True\n",
        "    except Exception as e:\n",
        "              print(e)\n",
        "              return False"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "4dQ_QttKX67l",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# If you chose to download the bert model, the following function will\n",
        "# download and unzip the model in the designated location\n",
        "\n",
        "if should_download.value == \"Yes\":\n",
        "    create_and_populate_checkpoint_folder(subfolder_name = tf_checkpoint_sub_dir.value,\n",
        "                                         bert_model_name_with_zip = bert_model_choice.value)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "YJLmJqsCbPJS",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Now, we download the archive of tweets identified by fivethirtyeight.com as sourced from Russia's Internet Research Agency (IRA) during the 2016 election. This archive of tweets is served as a github repositry, which will be cloned and copied to the GCS bucket. \n",
        "\n",
        "Downloading the IRA tweets for the first time will take 5-10 minutes; the status may not update as the files are downloaded."
      ]
    },
    {
      "metadata": {
        "id": "DQPD4OCSy9eL",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "print(f'fivethirtyeight.com \\'s tweets will be downloaded to {DATA_DIR}')\n",
        "if tf.gfile.Exists(DATA_DIR +'/ru_tweets.h5'):\n",
        "    print(\"Tweets already downloaded to GCS; mirroring copy onto local instance.\")\n",
        "    if not os.path.exists('ru_tweets.h5'):\n",
        "        !gsutil cp {DATA_DIR}/ru_tweets.h5 ru_tweets.h5 \n",
        "    ru_tweets = pd.read_hdf('ru_tweets.h5',key='ru_tweets')\n",
        "    print('Tweets loaded into local instance and accessible as \\'ru_tweets\\'')\n",
        "else:\n",
        "    print(\"Cloning fivethirtyeight.com repo. This may take 5-10 minutes.\")\n",
        "    !git clone https://github.com/fivethirtyeight/russian-troll-tweets.git\n",
        "    print(\"Tweet list successfully downloaded\")\n",
        "    cur_path = '/content/russian-troll-tweets/'\n",
        "    ru_tweets = pd.concat([pd.read_csv(cur_path + f) if os.path.splitext(f)[1] == '.csv' else None for f in tqdm(listdir(cur_path))])\n",
        "    # save as an hd5\n",
        "    print(\"Caching tweets on GCS\")\n",
        "    print(\"Saving as h5\")\n",
        "    ru_tweets.to_hdf('ru_tweets.h5',key='ru_tweets')\n",
        "    print(\"copying to gcs bucket\")\n",
        "    !gsutil cp ru_tweets.h5 {DATA_DIR}/ru_tweets.h5\n",
        "    print(\"Operation completed.\")\n",
        "    "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "JB7nB_O1zHud",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Next, we download a list of tweet ids, collected by George Washington University (GWU) and hosted on dataverse.harvard.edu. These tweet ids were collected by GWU for purposes of collecting all tweets relating to the 2016 election. Unlike the fivethirtyeight.com tweets that include content, GWU provides only the ID and not the actual content. The Twitter Developer keys will be used to download the tweet content directly from Twitter.\n",
        "\n",
        "A collab instance does not have enough RAM/Disk Space to hold a large quantity of downloaded tweets. This code takes randomly samples 1/50 of each file."
      ]
    },
    {
      "metadata": {
        "id": "Ua-s3hk50Ltw",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "print(\"Checking to see if we already have a cached version of this data.\")\n",
        "if tf.gfile.Exists(DATA_DIR +'/filter_list_pd.h5'):\n",
        "    print(\"Data already downloaded and cached. Loadng from GCS and copying to local file system.\")\n",
        "    if not os.path.exists('filter_list_pd.h5'):\n",
        "        !gsutil cp {DATA_DIR}/filter_list_pd.h5 filter_list_pd.h5  \n",
        "    filter_list_pd = pd.read_hdf('filter_list_pd.h5',key='filter_list')\n",
        "    print('tweet indices loaded')\n",
        "else:    \n",
        "    # Note: This code has been written in long-form, so that new different samples,\n",
        "    # files, etc., can be input. Feel free to experiment as you wish.\n",
        "    \n",
        "    print(\"Downloading twitter IDs using pandas...\")\n",
        "    election_filter1_pd = pd.read_csv('https://dataverse.harvard.edu/api/access/datafile/:persistentId?persistentId=doi:10.7910/DVN/PDI7IN/KAONZQ', header=None, sep = \" \") \n",
        "    print(\"Election filter 1 downloaded\")\n",
        "    election_filter2_pd = pd.read_csv('https://dataverse.harvard.edu/api/access/datafile/:persistentId?persistentId=doi:10.7910/DVN/PDI7IN/WMEGGS', header=None, sep = \" \")\n",
        "    print(\"Election filter 2 downloaded\")\n",
        "    election_filter3_pd = pd.read_csv('https://dataverse.harvard.edu/api/access/datafile/:persistentId?persistentId=doi:10.7910/DVN/PDI7IN/LGVDZB', header=None, sep = \" \")\n",
        "    print(\"Election filter 3 downloaded\")\n",
        "    election_filter4_pd = pd.read_csv('https://dataverse.harvard.edu/api/access/datafile/:persistentId?persistentId=doi:10.7910/DVN/PDI7IN/UQ7JRQ', header=None, sep = \" \")\n",
        "    print(\"Election filter 4 downloaded\")\n",
        "    election_filter5_pd = pd.read_csv('https://dataverse.harvard.edu/api/access/datafile/:persistentId?persistentId=doi:10.7910/DVN/PDI7IN/XQYUAY', header=None, sep = \" \")\n",
        "    print(\"Election filter 5 downloaded\")\n",
        "    election_filter6_pd = pd.read_csv('https://dataverse.harvard.edu/api/access/datafile/:persistentId?persistentId=doi:10.7910/DVN/PDI7IN/ZJALBN', header=None, sep = \" \")\n",
        "    print(\"Election filter 6 downloaded\")\n",
        "    # now, we are only going to keep 1 million from each, since these files are too large.\n",
        "\n",
        "    filter_list_pd = pd.concat([election_filter1_pd.sample(frac=(1/50), random_state=random_seed), \n",
        "                                election_filter2_pd.sample(frac=(1/50), random_state=random_seed),\n",
        "                                election_filter3_pd.sample(frac=(1/50), random_state=random_seed), \n",
        "                                election_filter4_pd.sample(frac=(1/50), random_state=random_seed), \n",
        "                                election_filter5_pd.sample(frac=(1/50), random_state=random_seed), \n",
        "                                election_filter6_pd.sample(frac=(1/50), random_state=random_seed)]).reset_index()\n",
        "    filter_list_pd.columns = ['index','tweet_id'] \n",
        "    filter_list_len = len(filter_list_pd)\n",
        "    print(f'Loaded {filter_list_len} columns')\n",
        "    filter_list_pd.info()\n",
        "    filter_list_pd.to_hdf('filter_list_pd.h5',key='filter_list')\n",
        "    print(\"copying to gcs bucket\")\n",
        "    !gsutil cp filter_list_pd.h5 {DATA_DIR}/filter_list_pd.h5\n",
        "    print('Freeing up RAM...removing dataframes of original downloads')\n",
        "    del election_filter1_pd\n",
        "    del election_filter2_pd\n",
        "    del election_filter3_pd\n",
        "    del election_filter4_pd\n",
        "    del election_filter5_pd\n",
        "    del election_filter6_pd\n",
        "    print(\"Twitter ID collection completed.\")\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "TAfVBtsx0nIX",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "The following operations are conducted to confirm our requests to the rate limits of Twitter's API, which allows 100 IDs per request, and 900 requests per 15 minutes. \n",
        "\n",
        "Of course, the more data you download, the larger the training set, and hypothetically, the greater likelihood that a random downloaded sample of tweets will be representative of the universe of potential tweets. (In other words, the more data, the less likely an overfit model will be fit.)"
      ]
    },
    {
      "metadata": {
        "id": "p57MVmpWGM5B",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# divide the current filtered tweet list into buckets of 100 tweets\n",
        "MAX_IDs_PER_REQUEST = 100\n",
        "number_of_shards = math.ceil(len(filter_list_pd)/MAX_IDs_PER_REQUEST)\n",
        "filter_lst_pd_only_indices = filter_list_pd[['tweet_id']]\n",
        "filter_list_pd_sharded = np.array_split(filter_list_pd,number_of_shards)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "09g52g4pO_VF",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# The data provided in the original files are integers.\n",
        "# These integers need to be converted to strings, for purposes\n",
        "# of calling the Twitter API\n",
        "list_to_string = [list(map(str,f.tweet_id.tolist())) for f in filter_list_pd_sharded]\n",
        "string_list_to_concat = [\",\".join(f) for f in list_to_string]\n",
        "print('String conversion completed.')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "AOGSrqfvhbUZ",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "MAX_REQUESTS_PER_15_MIN = 900\n",
        "# now, divide string_list_to_concat into chunks of 900 shards, for submission to the twitter api\n",
        "number_of_shards = math.ceil(len(string_list_to_concat)/MAX_REQUESTS_PER_15_MIN)\n",
        "print(f'Created {number_of_shards} number_of_shards')\n",
        "string_list_to_concat_sharded = np.array_split(string_list_to_concat,number_of_shards)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "X8E7ZVNBDjrD",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Now, we instantiate the pypi Twitter API wrapper with the API keys collected.\n",
        "try:\n",
        "    t = Twitter(auth=OAuth(ACCESS_TWITTER_API, ACCESS_TWITTER_API_PRIVATE, CONSUMER_TWITTER_API, CONSUMER_TWITTER_API_PRIVATE))\n",
        "except:\n",
        "    except Exception as e:\n",
        "        print(\"Sorry! You may not be using valid credentials. Please check the following error message and try again.\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "J3wXCt4O1k39",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#@title Select start shard and number to collect. (15 min per shard)\n",
        "start_shard = 0 #@param {type:\"integer\"}\n",
        "num_of_shards = 3 #@param {type:\"integer\"}\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "GGVNkO-834ch",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def countdown(t):\n",
        "    while t > 0:\n",
        "        sys.stdout.write('\\rCountdown to next API pull : {}s'.format(t))\n",
        "        t -= 1\n",
        "        sys.stdout.flush()\n",
        "        time.sleep(1)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "zGNFLJvHltO2",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def save_tweetdata_to_gs(in_json, i):\n",
        "    pickle.dump(in_json, open(f'tweet100-{i}.pickle','wb'))\n",
        "    !gsutil cp tweet100-{i}.pickle {DATA_DIR}/tweet100-{i}.pickle "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "WGlv4ZgdVh4u",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Downloads the selected tweets, every 15 minutes\n",
        "# Note: This will overwrite tweets with the same index number\n",
        "\n",
        "for i in range(start_shard, start_shard + num_of_shards):\n",
        "    ret_tweets = [t.statuses.lookup(_id=f) for f in tqdm(string_list_to_concat_sharded[i])]\n",
        "    save_tweetdata_to_gs(ret_tweets, i)\n",
        "    countdown(910)\n",
        "    # Wait 15 min / 900 seconds until the next pull. Add 10 secs, just in case."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "XN4KCDCO_geY",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Find all the available tweets saved on the gs bucket,\n",
        "# copy them onto the local file system, and load them\n",
        "# into RAM.\n",
        "\n",
        "# This will overwrite any tweets still on the local file system.\n",
        "# If you have already downloaded tweets using the above code,\n",
        "# Then just run this code to restore what you have already downloaded.\n",
        "\n",
        "!gsutil cp {DATA_DIR}/tweet100* .\n",
        "\n",
        "# double-check that they have been loaded\n",
        "\n",
        "!ls tweet100*\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "AG62QXc2_SUl",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Convert the downloaded tweets to pandas data frames\n",
        "\n",
        "cur_tweets_files = glob('tweet100*')\n",
        "print(cur_tweets_files)\n",
        "\n",
        "tweet_dicts = [pickle.load(open(f,'rb')) for f in tqdm(cur_tweets_files)]\n",
        "num_of_tweet_dicts = len(tweet_dicts)\n",
        "print(f'{num_of_tweet_dicts} tweet files loaded')\n",
        "\n",
        "# Unnest the dictionaries"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "BUfeQIJLX9JX",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def flatten_tweet_json(in_json):\n",
        "    \"\"\"\n",
        "    Input is a json file from the twitter API.\n",
        "    Returns a python dictionary that unnests the\n",
        "    user-related content from the input json\n",
        "    \"\"\"\n",
        "    \n",
        "    \n",
        "    ret_dict = {}\n",
        "    cur_keys = in_json.keys()\n",
        "    for k in cur_keys:\n",
        "        if not type(in_json[k]) is dict:\n",
        "            ret_dict.update({k:in_json[k]})\n",
        "        else:\n",
        "            for kk in in_json[k].keys():\n",
        "                ret_dict.update({k + '.' + kk:in_json[k][kk]})\n",
        "    return(ret_dict)\n",
        "    "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "7EGYFUu7nmIc",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Convert the tweet_dicts to a data frame, replacing our tweet_dicts to conserve memory\n",
        "tweet_dicts = pd.DataFrame(list(map(flatten_tweet_json, tqdm(tweet_dicts))))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "XoAsQpNewkJV",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# The above steps are a bit time intensive. The following code\n",
        "# allows you to backup and restore the data generated thus far.\n",
        "\n",
        "def dump_prepared_tweets():\n",
        "    pickle.dump(tweet_dicts,open('tweet_dicts_prepped.pickle','wb'))\n",
        "    pickle.dump(ru_tweets, open('ru_tweets_prepped.pickle','wb'))\n",
        "    !gsutil cp tweet_dicts_prepped.pickle {DATA_CACHE}/tweet_dicts_prepped.pickle \n",
        "    !gsutil cp ru_tweets_prepped.pickle {DATA_CACHE}/ru_tweets_prepped.pickle\n",
        "    print('Prepped tweet dump and copy to GCS completed')\n",
        "    return True\n",
        "\n",
        "\n",
        "def load_prepared_tweets():\n",
        "    if not os.path.exists('tweet_dicts_prepped.pickle'):\n",
        "        print(\"Copying tweet_dicts_prepped.pickle from gcs\")\n",
        "        !gsutil cp {DATA_CACHE}/tweet_dicts_prepped.pickle .\n",
        "    if not os.path.exists('ru_tweets_prepped.pickle'):\n",
        "        !gsutil cp {DATA_CACHE}/ru_tweets_prepped.pickle .\n",
        "        print(\"Copying ru_tweets_prepped.pickle from gc\")\n",
        "    tweet_dicts = pickle.load(open('tweet_dicts_prepped.pickle','rb'))\n",
        "    ru_tweets = pickle.load(open('ru_tweets_prepped.pickle','rb'))    \n",
        "    print(\"Prepared tweets loaded\")\n",
        "    return tweet_dicts, ru_tweets"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "mWXSxcVwxN5B",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Uncomment these lines to dump and/or load the prepared tweets\n",
        "\n",
        "# dump_prepared_tweets\n",
        "tweet_dicts, ru_tweets = load_prepared_tweets()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "gCxgZIS66WwT",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "The tweet content collected using the GWU list of twitter IDs may have tweets identified by fivethirtyeight.com as being authorited by IRA. Ideally, a training set should have an accurate a ground truth as possible: the list of GWU tweets should ideally not have any tweets from IRA.\n",
        "\n",
        "For immediate purposes, we'll check the author name and text. Howeer, this step could probably be a bit more comprehensie."
      ]
    },
    {
      "metadata": {
        "id": "Q40YzlMv5-VG",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Collect a unique list of tweets written by the IRA, by screen name\n",
        "\n",
        "ru_authors_list = [x.lower() for x in list(set(ru_tweets.author.tolist()))]\n",
        "print(f'Number of unique screen names in the IRA dataset: {len(ru_authors_list)}')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "EFc_7Nw36qH-",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Conduct filtering\n",
        "filtered_tweets_screen_names = [x.lower() for x in tweet_dicts['user.screen_name'].tolist()]\n",
        "non_unique_screen_names_len = len(filtered_tweets_screen_names)\n",
        "unique_screen_names_len = len(set(filtered_tweets_screen_names))\n",
        "unique_screen_names = set(filtered_tweets_screen_names)\n",
        "print(f'Found {non_unique_screen_names_len} non-unique screen names, and {len(unique_screen_names)} unique screen names in filtered tweets')\n",
        "\n",
        "overlapping_screen_names = set(unique_screen_names) & set(ru_authors_list)\n",
        "print(f'There are {len(overlapping_screen_names)} overlapping screen names')\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "rXHSiW7X9OQ9",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Now, check for exact text content of tweets between ru and filtered.\n",
        "# It is likely that a russian tweet is in our list of filtered tweets\n",
        "\n",
        "overlapping_text = set(tweet_dicts['text'].tolist()) & set(ru_tweets['content'].tolist())\n",
        "\n",
        "# Remove the overlapping tweets from our filtered list\n",
        "\n",
        "gwu_tweets = tweet_dicts[~tweet_dicts['text'].isin(ru_tweets['content'])]\n",
        "\n",
        "# The IRA tweet database has non-English tweets\n",
        "# For now, let's limit the IRA database to English tweets\n",
        "\n",
        "ru_tweets = ru_tweets.loc[ru_tweets.language == 'English']"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "A09dRLvu1aZi",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "\"\"\"\n",
        "Scrub the tweets of urls, hash tags, emojis, RTs, and other twitter handles.\n",
        "BERT only analyzes sentence-like text.\n",
        "\n",
        "These other features could be used as engineered features...\n",
        "\"\"\" \n",
        "\n",
        "\n",
        "def clean_tweet_text(in_text):\n",
        "    \n",
        "    try:\n",
        "        \n",
        "        # remove URLS\n",
        "        # regex from https://regexr.com/36fcc\n",
        "        url_re = re.compile(r'(http|ftp|https)://([\\w+?\\.\\w+])+([a-zA-Z0-9\\~\\!\\@\\#\\$\\%\\^\\&\\*\\(\\)_\\-\\=\\+\\\\\\/\\?\\.\\:\\;\\'\\,]*)?')\n",
        "        in_text = url_re.sub(\"\",in_text)\n",
        "\n",
        "        # remove hashtags\n",
        "        hashtag_re = re.compile(r'(#\\w+)')\n",
        "        in_text = hashtag_re.sub(\"\",in_text)\n",
        "\n",
        "        # remove screen name references\n",
        "        screen_name_re = re.compile(r'(@\\w+)')\n",
        "        in_text = screen_name_re.sub(\"\", in_text)\n",
        "\n",
        "        # remove RTs\n",
        "        in_text = in_text.replace(\"RT\",\"\")\n",
        "\n",
        "        # strip text of extra spaces , and keep one space between each word\n",
        "        in_text = \" \".join(in_text.split())\n",
        "\n",
        "        # strip text of emjoi\n",
        "        # https://stackoverflow.com/questions/51217909/removing-all-emojis-from-text\n",
        "        emoji_re = re.compile(\"[\"\n",
        "            u\"\\U0001F600-\\U0001F64F\"  # emoticons\n",
        "            u\"\\U0001F300-\\U0001F5FF\"  # symbols & pictographs\n",
        "            u\"\\U0001F680-\\U0001F6FF\"  # transport & map symbols\n",
        "            u\"\\U0001F1E0-\\U0001F1FF\"  # flags (iOS)\n",
        "            u\"\\U0001F1F2-\\U0001F1F4\"  # Macau flag\n",
        "            u\"\\U0001F1E6-\\U0001F1FF\"  # flags\n",
        "            u\"\\U0001F600-\\U0001F64F\"\n",
        "            u\"\\U00002702-\\U000027B0\"\n",
        "            u\"\\U000024C2-\\U0001F251\"\n",
        "            u\"\\U0001f926-\\U0001f937\"\n",
        "            u\"\\U0001F1F2\"\n",
        "            u\"\\U0001F1F4\"\n",
        "            u\"\\U0001F620\"\n",
        "            u\"\\u200d\"\n",
        "            u\"\\u2640-\\u2642\"\n",
        "            \"]+\", flags=re.UNICODE)\n",
        "\n",
        "        in_text = emoji_re.sub(\"\",in_text)\n",
        "    \n",
        "    except:\n",
        "        in_text = \"\"\n",
        "    return in_text\n",
        "\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "fXG5v-rSD0oy",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "\n",
        "# create a new column with the filtered text\n",
        "tqdm.pandas()\n",
        "ru_tweets['filtered_text'] = ru_tweets['content'].progress_apply(clean_tweet_text).tolist()\n",
        "gwu_tweets['filtered_text'] = gwu_tweets['text'].progress_apply(clean_tweet_text).tolist()\n",
        "\n",
        "# remove any tweets less than 10 characters\n",
        "\n",
        "ru_tweets= ru_tweets.loc[ru_tweets['filtered_text'].progress_apply(lambda x: len(x) >= 10)]\n",
        "gwu_tweets= gwu_tweets.loc[gwu_tweets['filtered_text'].progress_apply(lambda x: len(x) >= 10)]\n",
        "\n",
        "\n",
        "# Let's print a sample\n",
        "\n",
        "print(ru_tweets.sample(10)['filtered_text'].tolist())\n",
        "print(gwu_tweets.sample(10)['filtered_text'].tolist())"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "6qOyVQrr9m1H",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Define the characteristics on which we will fit a model. The settings indicated are recommended, and will work for one iteration of the twitter API pull. If you have downloaded more tweets, then the sample size may be increased."
      ]
    },
    {
      "metadata": {
        "id": "vzbzX-XY8f71",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#@title Selecting sample size and train/test split\n",
        "sample_size = 125000 #@param {type:\"integer\"}\n",
        "training_pct = 0.8 #@param {type:\"slider\", min:0, max:1, step:0.05}\n",
        "proportion_ira_tweets_to_gwu_tweets = 0.5 #@param {type:\"slider\", min:0, max:1, step:0.05}\n",
        "\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "ikSDMoRoAQXv",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# create a new dataframe with our annotated set\n",
        "# start with 100k total examples, divided 50/50\n",
        "\n",
        "tqdm.pandas()\n",
        "random_seed = 42\n",
        "number_of_ira_tweets = int(sample_size * .5)\n",
        "number_of_gwu_tweets = int(sample_size * (1-.5))\n",
        "train_set_size = int(sample_size * training_pct)\n",
        "test_set_size = int(sample_size * (1-training_pct))\n",
        "\n",
        "labeled_set = pd.DataFrame({'text':ru_tweets['filtered_text'].sample(number_of_ira_tweets, random_state=random_seed).tolist() + gwu_tweets['filtered_text'].sample(number_of_gwu_tweets, random_state=random_seed).tolist(), \n",
        "                            'label':[1 for x in range(0,number_of_ira_tweets)] + [0 for x in range(0,number_of_gwu_tweets)]})\n",
        "\n",
        "\n",
        "# re-shuffle\n",
        "\n",
        "labeled_set = labeled_set.sample(frac=1, random_state=random_seed)\n",
        "train_labeled_set = labeled_set[:train_set_size]\n",
        "test_labeled_set = labeled_set[train_set_size:]\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "yijnsrR4Fz_j",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "labeled_set.info()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Rh83vp2CGDNV",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Importing our BERT-specific functions\n",
        "\n",
        "import bert\n",
        "from bert import run_classifier\n",
        "from bert import optimization\n",
        "from bert import tokenization\n",
        "from bert import modeling"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "iQK9Vqli6uBR",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Adopted from https://github.com/google-research/bert/blob/master/predicting_movie_reviews_with_bert_on_tf_hub.ipynb\n",
        "\n",
        "DATA_COLUMN = 'text'\n",
        "LABEL_COLUMN = 'label'\n",
        "\n",
        "train_InputExamples = train_labeled_set.apply(lambda x: bert.run_classifier.InputExample(guid=None, # Globally unique ID for bookkeeping, unused in this example\n",
        "                                                                   text_a = x[DATA_COLUMN], \n",
        "                                                                   text_b = None, \n",
        "                                                                   label = x[LABEL_COLUMN]), axis = 1)\n",
        "\n",
        "test_InputExamples = test_labeled_set.apply(lambda x: bert.run_classifier.InputExample(guid=None, \n",
        "                                                                   text_a = x[DATA_COLUMN], \n",
        "                                                                   text_b = None, \n",
        "                                                                   label = x[LABEL_COLUMN]), axis = 1)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "H0JujQ0bGNmM",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def create_tokenizer():\n",
        "  \n",
        "  \"\"\"\n",
        "  Reviews the type of BERT module selected earlier, \n",
        "  and create the appropriate tokenizer.\n",
        "  \n",
        "  A tokenizer is required for inputing text into BERT\n",
        "  \"\"\"\n",
        "    \n",
        "  do_lower_case = False\n",
        "  if 'uncased' in bert_model_choice.value:\n",
        "        cased_option = False\n",
        "  else:\n",
        "        cased_option = True\n",
        "        \n",
        "  vocab_file = TF_CHECKPOINT_ROOT + '/' + tf_checkpoint_sub_dir.value + 'vocab.txt'\n",
        "  \n",
        "  with tf.Graph().as_default():      \n",
        "      return bert.tokenization.FullTokenizer(\n",
        "          vocab_file=vocab_file, do_lower_case=do_lower_case)\n",
        "\n",
        "tokenizer = create_tokenizer()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "GSrgFznBGkOa",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# We'll set sequences to be at most 256 tokens long\n",
        "# Note, this is not 256 words, but 256 tokens\n",
        "# obviously, for tweets, this should be sufficient\n",
        "MAX_SEQ_LENGTH = 256\n",
        "\n",
        "# Convert our train and test features to InputFeatures that BERT understands.\n",
        "# Save the results as a TFRecord, which will be copied to Google Cloud Storage,\n",
        "# so, that the TPU can read the TFRecord directly\n",
        "\n",
        "train_features = bert.run_classifier.file_based_convert_examples_to_features(train_InputExamples, label_list, MAX_SEQ_LENGTH, tokenizer,'train_features.TFRecord')\n",
        "test_features = bert.run_classifier.file_based_convert_examples_to_features(test_InputExamples, label_list, MAX_SEQ_LENGTH, tokenizer,'test_features.TFRecord')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "6OJ6p6NuZSSP",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Copy the generated TFRecords to the data_cache, so that the TPU may directly access\n",
        "\n",
        "!gsutil cp train_features.TFRecord {DATA_CACHE}/train_features.TFRecord\n",
        "!gsutil cp test_features.TFRecord {DATA_CACHE}/test_features.TFRecord\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Cgb5YFo1G1e4",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Adopted from https://github.com/google-research/bert/blob/ffbda2a1aafe530525212d13194cc84d92ed0313/run_classifier.py#L574\n",
        "\n",
        "def create_model(bert_config, is_training, input_ids, input_mask, segment_ids,\n",
        "                 labels, num_labels, use_one_hot_embeddings):\n",
        "  \"\"\"Creates a classification model.\"\"\"\n",
        "  model = modeling.BertModel(\n",
        "      config=bert_config,\n",
        "      is_training=is_training,\n",
        "      input_ids=input_ids,\n",
        "      input_mask=input_mask,\n",
        "      token_type_ids=segment_ids,\n",
        "      use_one_hot_embeddings=use_one_hot_embeddings)\n",
        "\n",
        "  # Since we are classifying sentences, we are using pooled, not\n",
        "  # token output.\n",
        "  output_layer = model.get_pooled_output()\n",
        "\n",
        "  hidden_size = output_layer.shape[-1].value\n",
        "\n",
        "  output_weights = tf.get_variable(\n",
        "      \"output_weights\", [num_labels, hidden_size],\n",
        "      initializer=tf.truncated_normal_initializer(stddev=0.02))\n",
        "\n",
        "  output_bias = tf.get_variable(\n",
        "      \"output_bias\", [num_labels], initializer=tf.zeros_initializer())\n",
        "\n",
        "  with tf.variable_scope(\"loss\"):\n",
        "    if is_training:\n",
        "      # I.e., 0.1 dropout\n",
        "      output_layer = tf.nn.dropout(output_layer, keep_prob=0.9)\n",
        "\n",
        "    logits = tf.matmul(output_layer, output_weights, transpose_b=True)\n",
        "    logits = tf.nn.bias_add(logits, output_bias)\n",
        "    probabilities = tf.nn.softmax(logits, axis=-1)\n",
        "    log_probs = tf.nn.log_softmax(logits, axis=-1)\n",
        "\n",
        "    one_hot_labels = tf.one_hot(labels, depth=num_labels, dtype=tf.float32)\n",
        "\n",
        "    per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs, axis=-1)\n",
        "    loss = tf.reduce_mean(per_example_loss)\n",
        "\n",
        "    return (loss, per_example_loss, logits, probabilities)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "2L5iz0kHLxqI",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Adopted from https://github.com/google-research/bert/blob/ffbda2a1aafe530525212d13194cc84d92ed0313/run_classifier.py#L619\n",
        "\n",
        "def model_fn_builder(bert_config, num_labels, init_checkpoint, learning_rate,\n",
        "                     num_train_steps, num_warmup_steps, use_tpu,\n",
        "                     use_one_hot_embeddings):\n",
        "  \"\"\"Returns `model_fn` closure for TPUEstimator.\"\"\"\n",
        "\n",
        "  def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument\n",
        "    \"\"\"The `model_fn` for TPUEstimator.\"\"\"\n",
        "\n",
        "    tf.logging.info(\"*** Features ***\")\n",
        "    for name in sorted(features.keys()):\n",
        "      tf.logging.info(\"  name = %s, shape = %s\" % (name, features[name].shape))\n",
        "\n",
        "    input_ids = features[\"input_ids\"]\n",
        "    input_mask = features[\"input_mask\"]\n",
        "    segment_ids = features[\"segment_ids\"]\n",
        "    label_ids = features[\"label_ids\"]\n",
        "\n",
        "    is_training = (mode == tf.estimator.ModeKeys.TRAIN)\n",
        "\n",
        "    (total_loss, per_example_loss, logits, probabilities) = create_model(\n",
        "        bert_config, is_training, input_ids, input_mask, segment_ids, label_ids,\n",
        "        num_labels, use_one_hot_embeddings)\n",
        "\n",
        "    tvars = tf.trainable_variables()\n",
        "    initialized_variable_names = {}\n",
        "    scaffold_fn = None\n",
        "    if init_checkpoint:\n",
        "      (assignment_map, initialized_variable_names\n",
        "      ) = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint)\n",
        "      if use_tpu:\n",
        "\n",
        "        def tpu_scaffold():\n",
        "          tf.train.init_from_checkpoint(init_checkpoint, assignment_map)\n",
        "          return tf.train.Scaffold()\n",
        "\n",
        "        scaffold_fn = tpu_scaffold\n",
        "      else:\n",
        "        tf.train.init_from_checkpoint(init_checkpoint, assignment_map)\n",
        "\n",
        "    tf.logging.info(\"**** Trainable Variables ****\")\n",
        "    for var in tvars:\n",
        "      init_string = \"\"\n",
        "      if var.name in initialized_variable_names:\n",
        "        init_string = \", *INIT_FROM_CKPT*\"\n",
        "      tf.logging.info(\"  name = %s, shape = %s%s\", var.name, var.shape,\n",
        "                      init_string)\n",
        "\n",
        "    output_spec = None\n",
        "    if mode == tf.estimator.ModeKeys.TRAIN:\n",
        "\n",
        "      train_op = optimization.create_optimizer(\n",
        "          total_loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu)\n",
        "\n",
        "      output_spec = tf.contrib.tpu.TPUEstimatorSpec(\n",
        "          mode=mode,\n",
        "          loss=total_loss,\n",
        "          train_op=train_op,\n",
        "          scaffold_fn=scaffold_fn)\n",
        "    elif mode == tf.estimator.ModeKeys.EVAL:\n",
        "\n",
        "      def metric_fn(per_example_loss, label_ids, logits):\n",
        "        predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)\n",
        "        accuracy = tf.metrics.accuracy(label_ids, predictions)\n",
        "        loss = tf.metrics.mean(values=per_example_loss)\n",
        "        f1_score = tf.contrib.metrics.f1_score(label_ids,predictions)\n",
        "        auc = tf.metrics.auc(label_ids,predictions)\n",
        "        recall = tf.metrics.recall(label_ids,predictions)\n",
        "        precision = tf.metrics.precision(label_ids,predictions) \n",
        "        true_pos = tf.metrics.true_positives(label_ids,predictions)\n",
        "        true_neg = tf.metrics.true_negatives(label_ids,predictions)   \n",
        "        false_pos = tf.metrics.false_positives(label_ids,predictions)  \n",
        "        false_neg = tf.metrics.false_negatives(label_ids,predictions)\n",
        "\n",
        "        return {\n",
        "            \"eval_accuracy\": accuracy,\n",
        "            \"eval_loss\": loss,\n",
        "            \"eval_f1_score\": f1_score,\n",
        "            \"eval_auc\": auc,\n",
        "            \"eval_recall\": recall,\n",
        "            \"eval_precision\": precision,\n",
        "            \"eval_true_pos\": true_pos,\n",
        "            \"eval_true_neg\": true_neg,\n",
        "            \"eval_false_pos\": false_pos,\n",
        "            \"eval_false_neg\": false_neg,\n",
        "        }\n",
        "\n",
        "      eval_metrics = (metric_fn,\n",
        "                      [per_example_loss, label_ids, logits])\n",
        "      output_spec = tf.contrib.tpu.TPUEstimatorSpec(\n",
        "          mode=mode,\n",
        "          loss=total_loss,\n",
        "          eval_metrics=eval_metrics,\n",
        "          scaffold_fn=scaffold_fn)\n",
        "    else:\n",
        "      output_spec = tf.contrib.tpu.TPUEstimatorSpec(\n",
        "          mode=mode,\n",
        "          predictions={\"probabilities\": probabilities},\n",
        "          scaffold_fn=scaffold_fn)\n",
        "    return output_spec\n",
        "\n",
        "  return model_fn"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "ENVKDSijHPtw",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# These default settings are recommended. Feel feel to change to experiment.\n",
        "\n",
        "TRAIN_BATCH_SIZE = 32\n",
        "EVAL_BATCH_SIZE = 32\n",
        "PREDICT_BATCH_SIZE = 32\n",
        "LEARNING_RATE = 2e-5\n",
        "NUM_TRAIN_EPOCHS = 6.0\n",
        "# Warmup is a period of time where hte learning rate \n",
        "# is small and gradually increases--usually helps training.\n",
        "WARMUP_PROPORTION = 0.1\n",
        "# Model configs\n",
        "SAVE_CHECKPOINTS_STEPS = 5000\n",
        "SAVE_SUMMARY_STEPS = 2500\n",
        "NUM_TPU_CORES = 8\n",
        "OUTPUT_DIR = TF_CHECKPOINT_ROOT + '/' + tf_checkpoint_sub_dir.value\n",
        "ITERATIONS_PER_LOOP = 1000\n",
        "MAX_SEQ_LENGTH= 256"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "fwW1YTouKRMx",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Adopted from https://github.com/google-research/bert/blob/bee6030e31e42a9394ac567da170a89a98d2062f/run_classifier_with_tfhub.py#L180\n",
        "\n",
        "# Create an object to represent the TPU cluster\n",
        "tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(TPU_ADDRESS)\n",
        "\n",
        "# Create our configuration file\n",
        "run_config = tf.contrib.tpu.RunConfig(\n",
        "    cluster=tpu_cluster_resolver,\n",
        "    model_dir=OUTPUT_DIR,\n",
        "    save_checkpoints_steps=SAVE_CHECKPOINTS_STEPS,\n",
        "    tpu_config=tf.contrib.tpu.TPUConfig(\n",
        "        iterations_per_loop=ITERATIONS_PER_LOOP,\n",
        "        num_shards=NUM_TPU_CORES,\n",
        "        per_host_input_for_training=tf.contrib.tpu.InputPipelineConfig.PER_HOST_V2))\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "RGYIeEHkHUhq",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Compute # train and warmup steps from batch size\n",
        "num_train_steps = int(len(train_InputExamples) / TRAIN_BATCH_SIZE * NUM_TRAIN_EPOCHS)\n",
        "num_warmup_steps = int(num_train_steps * WARMUP_PROPORTION)\n",
        "print(f'Number of training steps is {num_train_steps}, and number of warmup steps is {num_warmup_steps}')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "zFdHRrvQeWv5",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Adopted from https://github.com/google-research/bert/blob/bee6030e31e42a9394ac567da170a89a98d2062f/run_classifier_with_tfhub.py#L89\n",
        "# Generate a model from our model builder factory  \n",
        "    \n",
        "model_fn = model_fn_builder(\n",
        "    bert_config=  modeling.BertConfig.from_json_file(TF_CHECKPOINT_ROOT + '/' + tf_checkpoint_sub_dir.value + 'bert_config.json'),\n",
        "    num_labels=len(label_list),\n",
        "    init_checkpoint=TF_CHECKPOINT_ROOT + '/' + tf_checkpoint_sub_dir.value + 'bert_model.ckpt',\n",
        "    learning_rate=2e-5,\n",
        "    num_train_steps=num_train_steps,\n",
        "    num_warmup_steps=num_warmup_steps,\n",
        "    use_tpu=True,\n",
        "    use_one_hot_embeddings=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "gj807Lj3Hh5u",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Create the estimator\n",
        "    \n",
        "estimator = tf.contrib.tpu.TPUEstimator(\n",
        "use_tpu=True,\n",
        "model_fn=model_fn,\n",
        "config=run_config,\n",
        "train_batch_size=TRAIN_BATCH_SIZE,\n",
        "eval_batch_size=EVAL_BATCH_SIZE,\n",
        "predict_batch_size=PREDICT_BATCH_SIZE)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "ew3qKpJUID7N",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Create an input function for training, using the TFRecord we generated\n",
        "\n",
        "train_input_fn = bert.run_classifier.file_based_input_fn_builder(\n",
        "        input_file=DATA_CACHE + '/train_features.TFRecord',\n",
        "        seq_length=MAX_SEQ_LENGTH,\n",
        "        is_training=True,\n",
        "        drop_remainder=True)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "xvx7CzFFHlj4",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "print(f'Beginning Training!')\n",
        "current_time = datetime.now()\n",
        "estimator.train(input_fn=train_input_fn, max_steps=num_train_steps)\n",
        "print(\"Training took time \", datetime.now() - current_time)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "C0j2ZOoLkfAr",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Training is done! Let's see how accurate our model is:"
      ]
    },
    {
      "metadata": {
        "id": "4ijigDbSkf1Z",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Now, let's evaluate to see how we've done.\n",
        "\n",
        "test_input_fn = bert.run_classifier.file_based_input_fn_builder(\n",
        "        input_file=DATA_CACHE + '/test_features.TFRecord',\n",
        "        seq_length=MAX_SEQ_LENGTH,\n",
        "        is_training=False,\n",
        "        drop_remainder=True)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "wn1LiUrNlu_G",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# The calculation for evaluation steps is separate from the calculation for training steps\n",
        "eval_steps = int(len(test_InputExamples) / EVAL_BATCH_SIZE)\n",
        "print(eval_steps)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "STgsMgAwlBJe",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "results_dict = estimator.evaluate(input_fn=test_input_fn, steps=eval_steps)\n",
        "print(results_dict)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "idK62s6npWrP",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "4smaY4AQSuBO",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "If everything went well, you should have received at least 90 percent accuracy, recall, and precision!"
      ]
    },
    {
      "metadata": {
        "id": "Km1txbQ4nLe-",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Let's make some predictions - for this to work, we have to create a TF Record...\n",
        "# Adopted from https://github.com/google-research/bert/blob/0a0ea64a3ac1f43ed27d75278b9578708f9febcf/predicting_movie_reviews_with_bert_on_tf_hub.ipynb#L1090\n",
        "\n",
        "def getPrediction(in_sentences):\n",
        "  ret = []\n",
        "  try:\n",
        "        \n",
        "      labels = [0, 1]\n",
        "      input_examples = [run_classifier.InputExample(guid=\"\", text_a = x, text_b = None, label = 0) for x in in_sentences] # here, \"\" is just a dummy label\n",
        "      input_features = run_classifier.file_based_convert_examples_to_features(input_examples, label_list, MAX_SEQ_LENGTH, tokenizer,'predict_features.TFRecord')\n",
        "      !gsutil cp predict_features.TFRecord {DATA_CACHE}/predict_features.TFRecord\n",
        "      predict_input_fn = run_classifier.file_based_input_fn_builder(input_file=DATA_CACHE + '/predict_features.TFRecord', seq_length=MAX_SEQ_LENGTH, is_training=False, drop_remainder=True)\n",
        "      predictions = estimator.predict(predict_input_fn)\n",
        "      #ret = [(sentence, prediction['probabilities'], labels[prediction['labels']]) for sentence, prediction in zip(in_sentences, predictions)]\n",
        "      for p in predictions:\n",
        "            ret.append(p)\n",
        "  except IndexError:\n",
        "      return(ret)\n",
        "  return(ret)  "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "ujBO-hv9p9eV",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Now, try your own predictions! With a batch size set at 32, for TPU prediction, make sure you are sending predictions that match the batch size."
      ]
    },
    {
      "metadata": {
        "id": "Qtt1-tN-pVkV",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Yg3CzjB-nonx",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "pred_sentences = [\n",
        "  \"Sentence 1\",\n",
        "  \"Sentence 2\",\n",
        "  \"Sentence 3\",\n",
        "  \"Sentence 4\"\n",
        "]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "yVLAKMG7nyl9",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Matching batch size\n",
        "pred_sentences = pred_sentences * 8\n",
        "preds = getPrediction(pred_sentences)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "8CclkvPtuVz-",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "pred_probs = [np.argmax(p['probabilities']) for p in preds]\n",
        "final_predictions = list(zip(pred_sentences, pred_probs))\n",
        "print(final_predictions)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "TtT3ueYl13eP",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "WG8AtiTXEbkp",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Adding Features to the BERT Classification Model\n",
        "In the example above, text classification was accomplished through the BERT model analyzing only the filtered text of the tweet. With 125,000 features on an 80/20 train-test split,94 percent recall and 93 percent precision were observed in the test set. That's quite good!\n",
        "\n",
        "Perhaps, we can add some aditional features using the tweet metadata to further improve the scores? This will require some feature engineering.\n",
        "\n",
        "For this exercise, three features will be added: 1) a feature-scaled count of the number of followers, a feature-scaled count of the number of screen names followed by the respective tweeter, and third, the ratio between the number of followers and followees. \n",
        "\n"
      ]
    },
    {
      "metadata": {
        "id": "Vsd8qwTHj7oY",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#@title (Keep or change) Selecting sample size and train/test split\n",
        "sample_size = 125000 #@param {type:\"integer\"}\n",
        "training_pct = 0.8 #@param {type:\"slider\", min:0, max:1, step:0.05}\n",
        "proportion_ira_tweets_to_gwu_tweets = 0.5 #@param {type:\"slider\", min:0, max:1, step:0.05}\n",
        "\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "4PQjL2zEmNbN",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "sKSZfUzgFdPO",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "\"\"\"\n",
        "\n",
        "# Let's go back to our original dictionary of tweets\n",
        "\n",
        "# The relevant features for each original dataset are:\n",
        "\n",
        "ru_tweets['followers']\n",
        "ru_tweets['following']\n",
        "\n",
        "gwu_tweets['user.followers_count']\n",
        "gwu_tweets['users.friends_count']\n",
        "\n",
        "\"\"\"\n",
        "\n",
        "tqdm.pandas()\n",
        "random_seed = 42\n",
        "number_of_ira_tweets = int(sample_size * .5)\n",
        "number_of_gwu_tweets = int(sample_size * (1-.5))\n",
        "train_set_size = int(sample_size * training_pct)\n",
        "test_set_size = int(sample_size * (1-training_pct))\n",
        "\n",
        "extra_features = pd.DataFrame({'text':ru_tweets['filtered_text'].sample(number_of_ira_tweets, random_state=random_seed).tolist() + \n",
        "                                           gwu_tweets['filtered_text'].sample(number_of_gwu_tweets, random_state=random_seed).tolist(), \n",
        "                            'label':[1 for x in range(0,number_of_ira_tweets)] + [0 for x in range(0,number_of_gwu_tweets)],\n",
        "                                          'followers':ru_tweets['followers'].sample(number_of_ira_tweets, random_state=random_seed).tolist() + \n",
        "                                           gwu_tweets['user.followers_count'].sample(number_of_gwu_tweets, random_state=random_seed).tolist(),\n",
        "                                          'following': ru_tweets['followers'].sample(number_of_ira_tweets, random_state=random_seed).tolist() + \n",
        "                                           gwu_tweets['user.friends_count'].sample(number_of_gwu_tweets, random_state=random_seed).tolist()})\n",
        "\n",
        "\n",
        "# re-shuffle\n",
        "\n",
        "extra_features = extra_features.sample(frac=1, random_state=random_seed)\n",
        "\n",
        "# let's use some feature scaling\n",
        "\n",
        "followers_stats={\n",
        "'max_values':max(extra_features.followers.tolist()),\n",
        "'min_values':min(extra_features.followers.tolist()),\n",
        "'avg_values':np.mean(extra_features.followers.tolist()),\n",
        "'median_values':np.median(extra_features.followers.tolist())\n",
        "}\n",
        "print(followers_stats)\n",
        "\n",
        "# Yikes, this distribution is heavily skewed towards zero!\n",
        "\n",
        "following_stats={\n",
        "'max_values':max(extra_features.following.tolist()),\n",
        "'min_values':min(extra_features.following.tolist()),\n",
        "'avg_values':np.mean(extra_features.following.tolist()),\n",
        "'median_values':np.median(extra_features.following.tolist())\n",
        "}\n",
        "\n",
        "# This distribution does not look as skewed.\n",
        "\n",
        "# let's conduct feature scaling based on the median.\n",
        "\n",
        "def feature_scale_median(x,x_median,x_min,x_max):\n",
        "    \n",
        "    return((x-x_median)/(x_max-x_min))\n",
        "\n",
        "extra_features['followers_scaled'] = extra_features['followers'].apply(feature_scale_median, \n",
        "                                                                       x_median=followers_stats['median_values'],\n",
        "                                                                      x_min=followers_stats['min_values'],\n",
        "                                                                      x_max=followers_stats['max_values'])\n",
        "\n",
        "\n",
        "extra_features['following_scaled'] = extra_features['following'].apply(feature_scale_median, \n",
        "                                                                       x_median=following_stats['median_values'],\n",
        "                                                                      x_min=following_stats['min_values'],\n",
        "                                                                      x_max=following_stats['max_values'])\n",
        "\n",
        "\n",
        "extra_features['following_follower_ratio'] = extra_features['following'] / extra_features['followers']\n",
        "\n",
        "# replace any non-numbers or infinite values with zero\n",
        "\n",
        "extra_features['following_follower_ratio'] = extra_features['following_follower_ratio'].apply(lambda x: 0 if (np.isnan(x) or np.isinf(x)) else x)\n",
        "extra_features['following_follower_ratio'].describe()\n",
        "\n",
        "# convert our three extra features into a single list\n",
        "\n",
        "extra_features_list = list(zip(extra_features['followers_scaled'].tolist(),\n",
        "                              extra_features['following_scaled'].tolist(),\n",
        "                              extra_features['following_follower_ratio'].tolist()))\n",
        "\n",
        "# convert from a list of tuples to a list of lists\n",
        "extra_features_list = list(zip(extra_features['followers_scaled'].tolist(),\n",
        "                              extra_features['following_scaled'].tolist(),\n",
        "                              extra_features['following_follower_ratio'].tolist()))\n",
        "\n",
        "\n",
        "extra_features_list = [[a,b,c] for a,b,c in extra_features_list]\n",
        "extra_features['extra_features'] = extra_features_list                           \n",
        "\n",
        "\n",
        "train_extra_features = extra_features[:train_set_size]\n",
        "test_extra_features = extra_features[train_set_size:]\n",
        "\n",
        "# setting a global variable for extra features. This isn't the best practice,\n",
        "# but for now...\n",
        "\n",
        "num_of_extra_features = 3\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "7AhxAq1AzfaJ",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "UH78pEO0WBA-",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "alTyOOgIoddD",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Since we are creating a generating a new model, and because we do not want to wipe the previous model, we need to create a new folder in your tf_checkpoint_root with the new model. Made sure that you've selected a BERT model in the previous cells - otherwise, this function will fail."
      ]
    },
    {
      "metadata": {
        "id": "Aojgk2eFTeCi",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "print(\"Select a location for the tensorflow model. Run the function two cells down to create a new folder.\")\n",
        "extra_features_model_folder = widgets.Dropdown(options=tf.io.gfile.listdir(TF_CHECKPOINT_ROOT), value=tf.io.gfile.listdir(TF_CHECKPOINT_ROOT)[0])\n",
        "extra_features_model_folder"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "LPuigaDkTkmi",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# assign the var\n",
        "extra_features_model_folder = str(extra_features_model_folder.value[:-1])\n",
        "print(f'Working from {extra_features_model_folder} as {type(extra_features_model_folder)}')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "rI99xXLf6uy3",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "## This code wil create a new sub-folder and populate a folder with a BERT model\n",
        "## **Only run if you want to create a new folder. Otherwise, select a model from above.\n",
        "import urllib\n",
        "cur_name = namegenerator.gen()\n",
        "text_input = input(f'Enter a model name, or <ENTER> for {cur_name}:')\n",
        "if text_input is '':\n",
        "    text_input=cur_name\n",
        "print(f'Creating folder for {text_input}')\n",
        "tf.gfile.MakeDirs(TF_CHECKPOINT_ROOT + '/' + text_input)\n",
        "extra_features_model_folder = text_input\n",
        "bert_file_name = bert_model_choice.value.split('/')[-1]\n",
        "print(bert_file_name)\n",
        "urllib.request.urlretrieve(bert_model_choice.value, bert_file_name )\n",
        "print(\"downloaded\")\n",
        "!ls\n",
        "!unzip -o {bert_file_name}\n",
        "!cd {bert_file_name.split('.')[0]};gsutil cp * {TF_CHECKPOINT_ROOT + '/' + extra_features_model_folder}\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "XfvO04qESocP",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Now, we're going to lightly edit the original code from the BERT repo\n",
        "# to add our extra features\n",
        "\n",
        "# From https://github.com/google-research/bert/blob/ffbda2a1aafe530525212d13194cc84d92ed0313/run_classifier.py#L161\n",
        "\n",
        "\n",
        "class InputExampleExtraFeatures(object):\n",
        "  \"\"\"A single training/test example for simple sequence classification.\"\"\"\n",
        "\n",
        "  def __init__(self, guid, text_a, text_b=None, label=None,extra_features=None):\n",
        "    \"\"\"Constructs a InputExample.\n",
        "    Args:\n",
        "      guid: Unique id for the example.\n",
        "      text_a: string. The untokenized text of the first sequence. For single\n",
        "        sequence tasks, only this sequence must be specified.\n",
        "      text_b: (Optional) string. The untokenized text of the second sequence.\n",
        "        Only must be specified for sequence pair tasks.\n",
        "      label: (Optional) string. The label of the example. This should be\n",
        "        specified for train and dev examples, but not for test examples.\n",
        "      extra_features: 1-D numpy array of extra features\n",
        "    \"\"\"\n",
        "    self.guid = guid\n",
        "    self.text_a = text_a\n",
        "    self.text_b = text_b\n",
        "    self.label = label\n",
        "    self.extra_features = extra_features\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "EHxnJV2vsXgm",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# From https://github.com/google-research/bert/blob/ffbda2a1aafe530525212d13194cc84d92ed0313/run_classifier.py#L479\n",
        "\n",
        "def file_based_convert_examples_to_features_with_extra_features(\n",
        "    examples, label_list, max_seq_length, tokenizer, output_file):\n",
        "  \"\"\"Convert a set of `InputExample`s to a TFRecord file.\"\"\"\n",
        "\n",
        "  writer = tf.python_io.TFRecordWriter(output_file)\n",
        "\n",
        "  for (ex_index, example) in enumerate(examples):\n",
        "    if ex_index % 10000 == 0:\n",
        "      tf.logging.info(\"Writing example %d of %d\" % (ex_index, len(examples)))\n",
        "\n",
        "    feature = convert_single_example_extra_features(ex_index, example, label_list,\n",
        "                                     max_seq_length, tokenizer)\n",
        "\n",
        "    def create_int_feature(values):\n",
        "      f = tf.train.Feature(int64_list=tf.train.Int64List(value=list(values)))\n",
        "      return f\n",
        "    \n",
        "    def create_float_feature(values):\n",
        "      f = tf.train.Feature(float_list=tf.train.FloatList(value=list(values)))\n",
        "      return f\n",
        "    \n",
        "\n",
        "    features = collections.OrderedDict()\n",
        "    features[\"input_ids\"] = create_int_feature(feature.input_ids)\n",
        "    features[\"input_mask\"] = create_int_feature(feature.input_mask)\n",
        "    features[\"segment_ids\"] = create_int_feature(feature.segment_ids)\n",
        "    features[\"label_ids\"] = create_int_feature([feature.label_id])\n",
        "\n",
        "    # Adding extra feature here\n",
        "    features[\"extra_features\"] = create_float_feature(feature.extra_features)\n",
        "    features[\"is_real_example\"] = create_int_feature(\n",
        "        [int(feature.is_real_example)])\n",
        "\n",
        "    tf_example = tf.train.Example(features=tf.train.Features(feature=features))\n",
        "    writer.write(tf_example.SerializeToString())\n",
        "  writer.close()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "BuxfcKYxygL4",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# From https://github.com/google-research/bert/blob/master/predicting_movie_reviews_with_bert_on_tf_hub.ipynb\n",
        "\n",
        "DATA_COLUMN = 'text'\n",
        "LABEL_COLUMN = 'label'\n",
        "\n",
        "train_InputExamples_extra_features = train_extra_features.apply(lambda x: InputExampleExtraFeatures(guid=None, # Globally unique ID for bookkeeping, unused in this example\n",
        "                                                                   text_a = x[DATA_COLUMN], \n",
        "                                                                   text_b = None, \n",
        "                                                                   label = x[LABEL_COLUMN], # adding extra features\n",
        "                                                                   extra_features = x['extra_features']), axis = 1)\n",
        "\n",
        "test_InputExamples_extra_features = test_extra_features.apply(lambda x: InputExampleExtraFeatures(guid=None, \n",
        "                                                                   text_a = x[DATA_COLUMN], \n",
        "                                                                   text_b = None, \n",
        "                                                                   label = x[LABEL_COLUMN], # adding extra features\n",
        "                                                                   extra_features = x['extra_features']), axis = 1)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "3rjl4HWF0gr5",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# From https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/run_classifier.py#L161\n",
        "\n",
        "import collections\n",
        "\n",
        "class PaddingInputExample_ExtraFeatures(object):\n",
        "    pass\n",
        "\n",
        "class InputFeatures_ExtraFeatures(object):\n",
        "  \"\"\"A single set of features of data.\"\"\"\n",
        "\n",
        "  def __init__(self,\n",
        "               input_ids,\n",
        "               input_mask,\n",
        "               segment_ids,\n",
        "               label_id,\n",
        "               extra_features,\n",
        "               is_real_example=True):\n",
        "    self.input_ids = input_ids\n",
        "    self.input_mask = input_mask\n",
        "    self.segment_ids = segment_ids\n",
        "    self.label_id = label_id\n",
        "    self.extra_features = extra_features\n",
        "    self.is_real_example = is_real_example\n",
        "\n",
        "def convert_single_example_extra_features(ex_index, example, label_list, max_seq_length,\n",
        "                           tokenizer):\n",
        "  \"\"\"Converts a single `InputExample` into a single `InputFeatures`.\"\"\"\n",
        "\n",
        "  if isinstance(example, PaddingInputExample_ExtraFeatures):\n",
        "    return InputFeatures(\n",
        "        input_ids=[0] * max_seq_length,\n",
        "        input_mask=[0] * max_seq_length,\n",
        "        segment_ids=[0] * max_seq_length,\n",
        "        label_id=0,\n",
        "        # if this is padding, insert 3 zeros for extra features\n",
        "        extra_features=[0] * num_of_extra_features,\n",
        "        is_real_example=False)\n",
        "\n",
        "  label_map = {}\n",
        "  for (i, label) in enumerate(label_list):\n",
        "    label_map[label] = i\n",
        "\n",
        "  tokens_a = tokenizer.tokenize(example.text_a)\n",
        "  tokens_b = None\n",
        "  if example.text_b:\n",
        "    tokens_b = tokenizer.tokenize(example.text_b)\n",
        "\n",
        "  if tokens_b:\n",
        "    # Modifies `tokens_a` and `tokens_b` in place so that the total\n",
        "    # length is less than the specified length.\n",
        "    # Account for [CLS], [SEP], [SEP] with \"- 3\"\n",
        "    _truncate_seq_pair(tokens_a, tokens_b, max_seq_length - 3)\n",
        "  else:\n",
        "    # Account for [CLS] and [SEP] with \"- 2\"\n",
        "    if len(tokens_a) > max_seq_length - 2:\n",
        "      tokens_a = tokens_a[0:(max_seq_length - 2)]\n",
        "\n",
        "  tokens = []\n",
        "  segment_ids = []\n",
        "  tokens.append(\"[CLS]\")\n",
        "  segment_ids.append(0)\n",
        "  for token in tokens_a:\n",
        "    tokens.append(token)\n",
        "    segment_ids.append(0)\n",
        "  tokens.append(\"[SEP]\")\n",
        "  segment_ids.append(0)\n",
        "\n",
        "  if tokens_b:\n",
        "    for token in tokens_b:\n",
        "      tokens.append(token)\n",
        "      segment_ids.append(1)\n",
        "    tokens.append(\"[SEP]\")\n",
        "    segment_ids.append(1)\n",
        "\n",
        "  input_ids = tokenizer.convert_tokens_to_ids(tokens)\n",
        "\n",
        "  # The mask has 1 for real tokens and 0 for padding tokens. Only real\n",
        "  # tokens are attended to.\n",
        "  input_mask = [1] * len(input_ids)\n",
        "\n",
        "  # Zero-pad up to the sequence length.\n",
        "  while len(input_ids) < max_seq_length:\n",
        "    input_ids.append(0)\n",
        "    input_mask.append(0)\n",
        "    segment_ids.append(0)\n",
        "\n",
        "  assert len(input_ids) == max_seq_length\n",
        "  assert len(input_mask) == max_seq_length\n",
        "  assert len(segment_ids) == max_seq_length\n",
        "\n",
        "  label_id = label_map[example.label]\n",
        "  if ex_index < 5:\n",
        "    tf.logging.info(\"*** Example ***\")\n",
        "    tf.logging.info(\"guid: %s\" % (example.guid))\n",
        "    tf.logging.info(\"tokens: %s\" % \" \".join(\n",
        "        [tokenization.printable_text(x) for x in tokens]))\n",
        "    tf.logging.info(\"input_ids: %s\" % \" \".join([str(x) for x in input_ids]))\n",
        "    tf.logging.info(\"input_mask: %s\" % \" \".join([str(x) for x in input_mask]))\n",
        "    tf.logging.info(\"segment_ids: %s\" % \" \".join([str(x) for x in segment_ids]))\n",
        "    tf.logging.info(\"label: %s (id = %d)\" % (example.label, label_id))\n",
        "    tf.logging.info(\"extra_features: %s\" % \" \".join([str(x) for x in example.extra_features]))\n",
        "    \n",
        "\n",
        "  feature = InputFeatures_ExtraFeatures(\n",
        "      input_ids=input_ids,\n",
        "      input_mask=input_mask,\n",
        "      segment_ids=segment_ids,\n",
        "      label_id=label_id,\n",
        "      extra_features=example.extra_features,\n",
        "      is_real_example=True)\n",
        "  return feature\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "aFJSToHkzn-d",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "MAX_SEQ_LENGTH = 256\n",
        "label_list = [0,1]\n",
        "train_features = file_based_convert_examples_to_features_with_extra_features(train_InputExamples_extra_features, label_list, MAX_SEQ_LENGTH, tokenizer,'train_extra_features.TFRecord')\n",
        "test_features = file_based_convert_examples_to_features_with_extra_features(test_InputExamples_extra_features, label_list, MAX_SEQ_LENGTH, tokenizer,'test_extra_features.TFRecord')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "cVmq5mzJATn2",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "\n",
        "!gsutil cp train_extra_features.TFRecord {DATA_CACHE}/train_extra_features.TFRecord\n",
        "!gsutil cp test_extra_features.TFRecord {DATA_CACHE}/test_extra_features.TFRecord\n",
        "# Uncomment these lines if you want to copy from data_cache to your local\n",
        "# \n",
        "#!gsutil cp {DATA_CACHE}/train_extra_features.TFRecord train_extra_features.TFRecord \n",
        "#!gsutil cp {DATA_CACHE}/test_extra_features.TFRecord test_extra_features.TFRecord\n",
        "#"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "JcRL9h7kdLz8",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "The original create_model function is altered here to add the extra features. In this case, the extra features are concatenated to the pooled output vector, and the weights for this vector are then exposed as trainable.\n",
        "\n",
        "The same optimizer for the BERT model is then utilized for the entire set of trainable variables - including the variables that were just created. In this case, the model is being both fine-tuned and trained with these three extra variables. \n",
        "\n",
        "An alternative approach could have also included fine-tuning the model and the extracting the fine-tuned vector before the binary classification layer, and then running a separate model on the concatenated output features + extra features.\n",
        "\n",
        "Additional experimentation is needed to determine the best approach. However, for immediate purposes, editing the existing create_model code to accomodate the new extra features, and then allowing the model to train on this amended model was the most straightforward approach."
      ]
    },
    {
      "metadata": {
        "id": "F6k-MT922q6w",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Adopted from https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/run_classifier.py#L574\n",
        "\n",
        "def create_model_extra_features(bert_config, is_training, input_ids, input_mask, segment_ids,extra_features,\n",
        "                 labels, num_labels, use_one_hot_embeddings):\n",
        "  \"\"\"Creates a classification model.\"\"\"\n",
        "  model = modeling.BertModel(\n",
        "      config=bert_config,\n",
        "      is_training=is_training,\n",
        "      input_ids=input_ids,\n",
        "      input_mask=input_mask,\n",
        "      token_type_ids=segment_ids,\n",
        "      use_one_hot_embeddings=use_one_hot_embeddings)\n",
        "\n",
        "\n",
        "  output_layer = model.get_pooled_output()\n",
        "  # Here, we make alterations to add the extra features\n",
        "  output_layer_extra_features = tf.concat([output_layer,tf.convert_to_tensor(extra_features, dtype=tf.float32)],axis=1)  \n",
        "    \n",
        "  hidden_size = output_layer_extra_features.shape[-1].value\n",
        "\n",
        "  output_weights = tf.get_variable(\n",
        "      \"output_weights\", [num_labels, hidden_size],\n",
        "      initializer=tf.truncated_normal_initializer(stddev=0.02))\n",
        "\n",
        "  output_bias = tf.get_variable(\n",
        "      \"output_bias\", [num_labels], initializer=tf.zeros_initializer())\n",
        "\n",
        "  with tf.variable_scope(\"loss\"):\n",
        "    if is_training:\n",
        "      # I.e., 0.1 dropout\n",
        "      output_layer_extra_features = tf.nn.dropout(output_layer_extra_features, keep_prob=0.9)\n",
        "\n",
        "    logits = tf.matmul(output_layer_extra_features, output_weights, transpose_b=True)\n",
        "    logits = tf.nn.bias_add(logits, output_bias)\n",
        "    probabilities = tf.nn.softmax(logits, axis=-1)\n",
        "    log_probs = tf.nn.log_softmax(logits, axis=-1)\n",
        "\n",
        "    one_hot_labels = tf.one_hot(labels, depth=num_labels, dtype=tf.float32)\n",
        "\n",
        "    per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs, axis=-1)\n",
        "    loss = tf.reduce_mean(per_example_loss)\n",
        "\n",
        "    return (loss, per_example_loss, logits, probabilities)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "s8o1z24C3wdj",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#Adopted from https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/run_classifier.py#L509\n",
        "\n",
        "def file_based_input_fn_builder_extra_features(input_file, seq_length, is_training,\n",
        "                                drop_remainder):\n",
        "  \"\"\"Creates an `input_fn` closure to be passed to TPUEstimator.\"\"\"\n",
        "\n",
        "  name_to_features = {\n",
        "      \"input_ids\": tf.FixedLenFeature([seq_length], tf.int64),\n",
        "      \"input_mask\": tf.FixedLenFeature([seq_length], tf.int64),\n",
        "      \"segment_ids\": tf.FixedLenFeature([seq_length], tf.int64),\n",
        "      \"label_ids\": tf.FixedLenFeature([], tf.int64),\n",
        "      \"extra_features\": tf.FixedLenFeature([num_of_extra_features], tf.float32), #Adding extra features\n",
        "      \"is_real_example\": tf.FixedLenFeature([], tf.int64),\n",
        "  }\n",
        "\n",
        "  def _decode_record(record, name_to_features):\n",
        "    \"\"\"Decodes a record to a TensorFlow example.\"\"\"\n",
        "    example = tf.parse_single_example(record, name_to_features)\n",
        "\n",
        "    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.\n",
        "    # So cast all int64 to int32.\n",
        "    for name in list(example.keys()):\n",
        "      t = example[name]\n",
        "      if t.dtype == tf.int64:\n",
        "        t = tf.to_int32(t)\n",
        "      example[name] = t\n",
        "    \n",
        "    return example\n",
        "\n",
        "  def input_fn(params):\n",
        "    \"\"\"The actual input function.\"\"\"\n",
        "    batch_size = params[\"batch_size\"]\n",
        "\n",
        "    # For training, we want a lot of parallel reading and shuffling.\n",
        "    # For eval, we want no shuffling and parallel reading doesn't matter.\n",
        "    d = tf.data.TFRecordDataset(input_file)\n",
        "    if is_training:\n",
        "      d = d.repeat()\n",
        "      d = d.shuffle(buffer_size=100)\n",
        "\n",
        "    d = d.apply(\n",
        "        tf.contrib.data.map_and_batch(\n",
        "            lambda record: _decode_record(record, name_to_features),\n",
        "            batch_size=batch_size,\n",
        "            drop_remainder=drop_remainder))\n",
        "\n",
        "    return d\n",
        "\n",
        "  return input_fn\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "JtkeHbML48BK",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Adopted from https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/run_classifier.py#L619\n",
        "\n",
        "def model_fn_builder_extra_features(bert_config, num_labels, init_checkpoint, learning_rate,\n",
        "                     num_train_steps, num_warmup_steps, use_tpu,\n",
        "                     use_one_hot_embeddings):\n",
        "  \"\"\"Returns `model_fn` closure for TPUEstimator.\"\"\"\n",
        "\n",
        "  def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument\n",
        "    \"\"\"The `model_fn` for TPUEstimator.\"\"\"\n",
        "\n",
        "    tf.logging.info(\"*** Features ***\")\n",
        "    for name in sorted(features.keys()):\n",
        "      tf.logging.info(\"  name = %s, shape = %s\" % (name, features[name].shape))\n",
        "\n",
        "    input_ids = features[\"input_ids\"]\n",
        "    input_mask = features[\"input_mask\"]\n",
        "    segment_ids = features[\"segment_ids\"]\n",
        "    label_ids = features[\"label_ids\"]\n",
        "    extra_features = features[\"extra_features\"] # Adding extra features\n",
        "\n",
        "    is_training = (mode == tf.estimator.ModeKeys.TRAIN)\n",
        "\n",
        "    (total_loss, per_example_loss, logits, probabilities) = create_model_extra_features(\n",
        "        bert_config, is_training, input_ids, input_mask, segment_ids, extra_features, label_ids,\n",
        "        num_labels, use_one_hot_embeddings)\n",
        "\n",
        "    tvars = tf.trainable_variables()\n",
        "    initialized_variable_names = {}\n",
        "    scaffold_fn = None\n",
        "    if init_checkpoint:\n",
        "      (assignment_map, initialized_variable_names\n",
        "      ) = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint)\n",
        "      if use_tpu:\n",
        "\n",
        "        def tpu_scaffold():\n",
        "          tf.train.init_from_checkpoint(init_checkpoint, assignment_map)\n",
        "          return tf.train.Scaffold()\n",
        "\n",
        "        scaffold_fn = tpu_scaffold\n",
        "      else:\n",
        "        tf.train.init_from_checkpoint(init_checkpoint, assignment_map)\n",
        "\n",
        "    tf.logging.info(\"**** Trainable Variables ****\")\n",
        "    for var in tvars:\n",
        "      init_string = \"\"\n",
        "      if var.name in initialized_variable_names:\n",
        "        init_string = \", *INIT_FROM_CKPT*\"\n",
        "      tf.logging.info(\"  name = %s, shape = %s%s\", var.name, var.shape,\n",
        "                      init_string)\n",
        "\n",
        "    output_spec = None\n",
        "    if mode == tf.estimator.ModeKeys.TRAIN:\n",
        "\n",
        "      train_op = optimization.create_optimizer(\n",
        "          total_loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu)\n",
        "\n",
        "      output_spec = tf.contrib.tpu.TPUEstimatorSpec(\n",
        "          mode=mode,\n",
        "          loss=total_loss,\n",
        "          train_op=train_op,\n",
        "          scaffold_fn=scaffold_fn)\n",
        "    elif mode == tf.estimator.ModeKeys.EVAL:\n",
        "\n",
        "      def metric_fn(per_example_loss, label_ids, logits):\n",
        "        predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)\n",
        "        accuracy = tf.metrics.accuracy(label_ids, predictions)\n",
        "        loss = tf.metrics.mean(values=per_example_loss)\n",
        "        f1_score = tf.contrib.metrics.f1_score(label_ids,predictions)\n",
        "        auc = tf.metrics.auc(label_ids,predictions)\n",
        "        recall = tf.metrics.recall(label_ids,predictions)\n",
        "        precision = tf.metrics.precision(label_ids,predictions) \n",
        "        true_pos = tf.metrics.true_positives(label_ids,predictions)\n",
        "        true_neg = tf.metrics.true_negatives(label_ids,predictions)   \n",
        "        false_pos = tf.metrics.false_positives(label_ids,predictions)  \n",
        "        false_neg = tf.metrics.false_negatives(label_ids,predictions)\n",
        "\n",
        "        return {\n",
        "            \"eval_accuracy\": accuracy,\n",
        "            \"eval_loss\": loss,\n",
        "            \"eval_f1_score\": f1_score,\n",
        "            \"eval_auc\": auc,\n",
        "            \"eval_recall\": recall,\n",
        "            \"eval_precision\": precision,\n",
        "            \"eval_true_pos\": true_pos,\n",
        "            \"eval_true_neg\": true_neg,\n",
        "            \"eval_false_pos\": false_pos,\n",
        "            \"eval_false_neg\": false_neg,\n",
        "        }\n",
        "\n",
        "      eval_metrics = (metric_fn,\n",
        "                      [per_example_loss, label_ids, logits])\n",
        "      output_spec = tf.contrib.tpu.TPUEstimatorSpec(\n",
        "          mode=mode,\n",
        "          loss=total_loss,\n",
        "          eval_metrics=eval_metrics,\n",
        "          scaffold_fn=scaffold_fn)\n",
        "    else:\n",
        "      output_spec = tf.contrib.tpu.TPUEstimatorSpec(\n",
        "          mode=mode,\n",
        "          predictions={\"probabilities\": probabilities},\n",
        "          scaffold_fn=scaffold_fn)\n",
        "    return output_spec\n",
        "\n",
        "  return model_fn"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "i3Zt0GIOm1kj",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "EOI5Y20R5vQp",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Compute # train and warmup steps from batch size\n",
        "num_train_steps = int(len(train_InputExamples_extra_features) / TRAIN_BATCH_SIZE * 6)\n",
        "num_warmup_steps = int(num_train_steps * WARMUP_PROPORTION)\n",
        "print(f'Number of training steps is {num_train_steps}, and number of warmup steps is {num_warmup_steps}')\n",
        "\n",
        "model_fn_extra_features = model_fn_builder_extra_features(\n",
        "  bert_config=  modeling.BertConfig.from_json_file(TF_CHECKPOINT_ROOT + '/' + extra_features_model_folder + '/bert_config.json'),\n",
        "  num_labels=len(label_list),\n",
        "  init_checkpoint=TF_CHECKPOINT_ROOT + '/' + extra_features_model_folder + '/bert_model.ckpt',\n",
        "  learning_rate=2e-5,\n",
        "  num_train_steps=num_train_steps,\n",
        "  num_warmup_steps=num_warmup_steps,\n",
        "  use_tpu=True,\n",
        "  use_one_hot_embeddings=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "TReragyO8kfC",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "run_config_extra_features = tf.contrib.tpu.RunConfig(\n",
        "    cluster=tpu_cluster_resolver,\n",
        "    model_dir=TF_CHECKPOINT_ROOT + '/' + extra_features_model_folder,\n",
        "    save_checkpoints_steps=SAVE_CHECKPOINTS_STEPS,\n",
        "    tpu_config=tf.contrib.tpu.TPUConfig(\n",
        "        iterations_per_loop=ITERATIONS_PER_LOOP,\n",
        "        num_shards=NUM_TPU_CORES,\n",
        "        per_host_input_for_training=tf.contrib.tpu.InputPipelineConfig.PER_HOST_V2))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "lm7cTQYD51R1",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "\n",
        "estimator_extra_features = tf.contrib.tpu.TPUEstimator(\n",
        "use_tpu=True,\n",
        "model_fn=model_fn_extra_features,\n",
        "config=run_config_extra_features,\n",
        "train_batch_size=TRAIN_BATCH_SIZE,\n",
        "eval_batch_size=EVAL_BATCH_SIZE,\n",
        "predict_batch_size=PREDICT_BATCH_SIZE)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "O18icNAS6dP3",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Create an input function for training. drop_remainder = True for using TPUs.\n",
        "\n",
        "train_input_fn_extra_features = file_based_input_fn_builder_extra_features(\n",
        "        input_file=DATA_CACHE + '/train_extra_features.TFRecord',\n",
        "        seq_length=MAX_SEQ_LENGTH,\n",
        "        is_training=True,\n",
        "        drop_remainder=True)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "sSt8Q5rE6JBk",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "print(f'Beginning Training!')\n",
        "current_time = datetime.now()\n",
        "estimator_extra_features.train(input_fn=train_input_fn_extra_features, max_steps=num_train_steps)\n",
        "print(\"Training took time \", datetime.now() - current_time)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "PVym7_48rijM",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Finally, create a test input function for evaluation\n",
        "\n",
        "test_input_fn_extra_features = file_based_input_fn_builder_extra_features(\n",
        "        input_file=DATA_CACHE + '/test_extra_features.TFRecord',\n",
        "        seq_length=MAX_SEQ_LENGTH,\n",
        "        is_training=False,\n",
        "        drop_remainder=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "lOG_wSC3rnhe",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "eval_steps = int(len(test_InputExamples_extra_features) / EVAL_BATCH_SIZE)\n",
        "print(eval_steps)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "3-eJWBs_rrFY",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "results_dict_extra_features = estimator_extra_features.evaluate(input_fn=test_input_fn_extra_features, steps=eval_steps)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "1cOponatKlIX",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "print(results_dict_extra_features)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "3RqOurHaauF5",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## All done!\n",
        "\n",
        "*What's next?*\n",
        "\n",
        "- This notebook is obviously a bit unwieldy, so converting the code into class modules would be a useful exercise.\n",
        "\n",
        "- The three added features did not seem to have much effect on precision or recall. Perhaps new features could be evaluated - or, perhaps the BERT language embeddings provide enough information to fit a model.\n",
        "\n",
        "- More robust testing on out-of-sample observations is needed to determine if the model has overfit the training set. The test set evaluation appears to indicate that overfitting is not a concern; however, a more robust test is needed.\n",
        "\n",
        "- The limitations of the Colab environment limit the ability to ingest and train on a larger training and test set. A model could be created outside the Colab environment on a much larger training/test set, and then wrapped as a RESTful API for real-time predictions.\n",
        "\n",
        "- Who knows!\n"
      ]
    },
    {
      "metadata": {
        "id": "zMgaDtrcvI62",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}