examples/text_classification/tc_bert_azureml.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Copyright (c) Microsoft Corporation. All rights reserved.*\n",
    "\n",
    "*Licensed under the MIT License.*\n",
    "\n",
    "# Text Classification of MultiNLI Sentences using BERT with Azure ML Pipelines"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Introduction\n",
    "\n",
    "In this notebook, we fine-tune and evaluate a pretrained [BERT](https://arxiv.org/abs/1810.04805) model on a subset of the [MultiNLI](https://www.nyu.edu/projects/bowman/multinli/) dataset using [AzureML](https://azure.microsoft.com/en-us/services/machine-learning-service/) Pipelines.\n",
    "\n",
    "We use a [distributed sequence classifier](../../utils_nlp/bert/sequence_classification_distributed.py) that wraps [Hugging Face's PyTorch implementation](https://github.com/huggingface/pytorch-pretrained-BERT) of Google's [BERT](https://github.com/google-research/bert).\n",
    "\n",
    "The notebooks acts as a template to,\n",
    "1. Process a massive dataset in parallel by dividing the dataset into chunks using [DASK](https://dask.org/) .\n",
    "2. Perform distributed training on AzureML compute on these processed chunks.\n",
    "\n",
    "We create an [AzureML Pipeline](https://docs.microsoft.com/en-us/azure/machine-learning/service/concept-ml-pipelines) for the two steps mentioned above. With this pipeline, the notebook can be scheduled regularly to fine tune BERT with new data and get a model which can be further deployed on [Azure Container Instance](https://docs.microsoft.com/en-us/azure/container-service/).\n",
    "\n",
    "AzureML Pipeline define reusable machine learning workflows that can be used as a template for your machine learning scenarios. Pipelines allow you to optimize your workflow and spend time on machine learning rather than infrastructure. If you are new to the concept of pipelines, [this would be a good place to get started](https://github.com/Azure/MachineLearningNotebooks/tree/master/how-to-use-azureml/machine-learning-pipelines/intro-to-pipelines).\n",
    "\n",
    "**Note: To learn how to do pre-training on your own, please reference the [AzureML-BERT repo](https://github.com/microsoft/AzureML-BERT) created by Microsoft.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "System version: 3.6.8 |Anaconda, Inc.| (default, Feb 21 2019, 18:30:04) [MSC v.1916 64 bit (AMD64)]\n",
      "Azure ML SDK Version: 1.0.48\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "sys.path.append(\"../../\")\n",
    "import os\n",
    "import json\n",
    "import random\n",
    "import shutil\n",
    "import pandas as pd\n",
    "\n",
    "from utils_nlp.azureml import azureml_utils\n",
    "from utils_nlp.dataset.multinli import get_generator\n",
    "\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "import azureml.core\n",
    "from azureml.core import Datastore, Experiment,  get_run\n",
    "from azureml.core.conda_dependencies import CondaDependencies\n",
    "from azureml.core.runconfig import RunConfiguration\n",
    "from azureml.core.compute import ComputeTarget,  AmlCompute\n",
    "from azureml.exceptions import ComputeTargetException\n",
    "from azureml.data.data_reference import DataReference\n",
    "from azureml.pipeline.steps import PythonScriptStep\n",
    "from azureml.pipeline.core import Pipeline, PipelineData\n",
    "from azureml.widgets import RunDetails\n",
    "from azureml.train.dnn import PyTorch\n",
    "from azureml.core.runconfig import MpiConfiguration\n",
    "from azureml.pipeline.steps import EstimatorStep\n",
    "\n",
    "print(\"System version: {}\".format(sys.version))\n",
    "print(\"Azure ML SDK Version:\", azureml.core.VERSION)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define a few variables before we get started, these variables define the folder where the data would reside, the batch size and the number of epochs we are training for. \n",
    "We also define the variables for AzureML workspace, which you can use to create a new workspace. You can ignore these variables if you have `config.json` in `.azureml` directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "LABEL_COL = \"genre\"\n",
    "DATA_FOLDER = \"../../data/temp\"\n",
    "TRAIN_FOLDER = \"../../data/temp/train\"\n",
    "TEST_FOLDER = \"../../data/temp/test\"\n",
    "ENCODED_LABEL_COL = \"label\"\n",
    "NUM_PARTITIONS = None\n",
    "LABELS = ['telephone', 'government', 'travel', 'slate', 'fiction']\n",
    "PROJECT_FOLDER = \"../../\"\n",
    "NODE_COUNT = 4\n",
    "\n",
    "config_path = (\n",
    "    \"./.azureml\"\n",
    ")  # Path to the directory containing config.json with azureml credentials\n",
    "\n",
    "# Azure resources\n",
    "subscription_id = \"YOUR_SUBSCRIPTION_ID\"\n",
    "resource_group = \"YOUR_RESOURCE_GROUP_NAME\"  \n",
    "workspace_name = \"YOUR_WORKSPACE_NAME\"  \n",
    "workspace_region = \"YOUR_WORKSPACE_REGION\" #Possible values eastus, eastus2 and so on.\n",
    "cluster_name = \"pipelines-tc-12\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example we will use AzureML pipelines to execute training pipelines. Each preprocessing step is included as a step in the pipeline. For a more detailed walkthrough of what pipelines are with a getting started guidelines check this [notebook](https://github.com/Azure/MachineLearningNotebooks/blob/master/how-to-use-azureml/machine-learning-pipelines/intro-to-pipelines/aml-pipelines-getting-started.ipynb). We start by doing some AzureML related setup below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 0.1 Initialize a workspace\n",
    "\n",
    "The following cell looks to set up the connection to your [Azure Machine Learning service Workspace](https://docs.microsoft.com/azure/machine-learning/service/concept-azure-machine-learning-architecture#workspace). You can choose to connect to an existing workspace or create a new one. \n",
    "\n",
    "**To access an existing workspace:**\n",
    "1. If you have a `config.json` file, you do not need to provide the workspace information; you will only need to update the `config_path` variable that is defined above which contains the file.\n",
    "2. Otherwise, you will need to supply the following:\n",
    "    * The name of your workspace\n",
    "    * Your subscription id\n",
    "    * The resource group name\n",
    "\n",
    "**To create a new workspace:**\n",
    "\n",
    "Set the following information:\n",
    "* A name for your workspace\n",
    "* Your subscription id\n",
    "* The resource group name\n",
    "* [Azure region](https://azure.microsoft.com/en-us/global-infrastructure/regions/) to create the workspace in, such as `eastus2`. \n",
    "\n",
    "This will automatically create a new resource group for you in the region provided if a resource group with the name given does not already exist. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "ws = azureml_utils.get_or_create_workspace(\n",
    "    config_path=config_path,\n",
    "    subscription_id=subscription_id,\n",
    "    resource_group=resource_group,\n",
    "    workspace_name=workspace_name,\n",
    "    workspace_region=workspace_region,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 0.2 Create a compute target\n",
    "We create and attach a [compute target](https://docs.microsoft.com/azure/machine-learning/service/concept-azure-machine-learning-architecture#compute-target) for training the model. Here we use the AzureML-managed compute target ([AmlCompute](https://docs.microsoft.com/azure/machine-learning/service/how-to-set-up-training-targets#amlcompute)) as our remote training compute resource. Our cluster autoscales from 0 to 8 `STANDARD_NC12` GPU nodes.\n",
    "\n",
    "Creating and configuring the AmlCompute cluster takes approximately 5 minutes the first time around. Once a cluster with the given configuration is created, it does not need to be created again.\n",
    "\n",
    "As with other Azure services, there are limits on certain resources (e.g. AmlCompute) associated with the Azure Machine Learning service. Read more about the default limits and how to request more quota [here](https://docs.microsoft.com/en-us/azure/machine-learning/service/how-to-manage-quotas)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found existing compute target.\n",
      "{'currentNodeCount': 0, 'targetNodeCount': 0, 'nodeStateCounts': {'preparingNodeCount': 0, 'runningNodeCount': 0, 'idleNodeCount': 0, 'unusableNodeCount': 0, 'leavingNodeCount': 0, 'preemptedNodeCount': 0}, 'allocationState': 'Steady', 'allocationStateTransitionTime': '2019-08-11T08:53:18.284000+00:00', 'errors': None, 'creationTime': '2019-07-25T04:16:20.598768+00:00', 'modifiedTime': '2019-08-05T06:40:12.292030+00:00', 'provisioningState': 'Succeeded', 'provisioningStateTransitionTime': None, 'scaleSettings': {'minNodeCount': 0, 'maxNodeCount': 10, 'nodeIdleTimeBeforeScaleDown': 'PT120S'}, 'vmPriority': 'Dedicated', 'vmSize': 'STANDARD_NC12'}\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    compute_target = ComputeTarget(workspace=ws, name=cluster_name)\n",
    "    print(\"Found existing compute target.\")\n",
    "except ComputeTargetException:\n",
    "    print(\"Creating a new compute target...\")\n",
    "    compute_config = AmlCompute.provisioning_configuration(\n",
    "        vm_size=\"STANDARD_NC12\", max_nodes=8\n",
    "    )\n",
    "\n",
    "    # create the cluster\n",
    "    compute_target = ComputeTarget.create(ws, cluster_name, compute_config)\n",
    "\n",
    "    compute_target.wait_for_completion(show_output=True)\n",
    "\n",
    "# use get_status() to get a detailed status for the current AmlCompute.\n",
    "print(compute_target.get_status().serialize())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Preprocessing\n",
    "\n",
    "The pipeline is defined by a series of steps, the first being a PythonScriptStep which utilizes [DASK](https://dask.org/) to load dataframes in partitions allowing us to load and preprocess different sets of data in parallel."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 Read Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_batches = get_generator(DATA_FOLDER, \"train\", num_batches=NUM_PARTITIONS, batch_size=10e6)\n",
    "test_batches = get_generator(DATA_FOLDER, \"dev_matched\", num_batches=NUM_PARTITIONS, batch_size=10e6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 Preprocess and Tokenize\n",
    "\n",
    "In the classification task, we use the first sentence only as the text input, and the corresponding genre as the label. Select the examples corresponding to one of the entailment labels (*neutral* in this case) to avoid duplicate rows, as the sentences are not unique, whereas the sentence pairs are.\n",
    "\n",
    "Once filtered, we encode the labels. To do this, fit a label encoder with the known labels in a MNLI dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "os.makedirs(TRAIN_FOLDER, exist_ok=True)\n",
    "os.makedirs(TEST_FOLDER, exist_ok=True)\n",
    "\n",
    "labels = LABELS\n",
    "label_encoder = LabelEncoder()\n",
    "label_encoder.fit(labels)\n",
    "\n",
    "num_train_partitions = 0\n",
    "for batch in train_batches:\n",
    "    batch = batch[batch[\"gold_label\"]==\"neutral\"]\n",
    "    batch[ENCODED_LABEL_COL] = label_encoder.transform(batch[LABEL_COL])\n",
    "    batch.to_csv(TRAIN_FOLDER+\"/batch{}.csv\".format(str(num_train_partitions)))\n",
    "    num_train_partitions += 1\n",
    "    \n",
    "num_test_partitions = 0\n",
    "for batch in test_batches:\n",
    "    batch = batch[batch[\"gold_label\"]==\"neutral\"]\n",
    "    batch[ENCODED_LABEL_COL] = label_encoder.transform(batch[LABEL_COL])\n",
    "    batch.to_csv(TEST_FOLDER+\"/batch{}.csv\".format(str(num_test_partitions)))\n",
    "    num_test_partitions += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once we have the partitions of data ready they are uploaded to the datastore."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "$AZUREML_DATAREFERENCE_9609849b541244d396d06017b5729edb"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds = ws.get_default_datastore()\n",
    "ds.upload(src_dir=TRAIN_FOLDER, target_path=\"mnli_data/train\", overwrite=True, show_progress=False)\n",
    "ds.upload(src_dir=TEST_FOLDER, target_path=\"mnli_data/test\", overwrite=True, show_progress=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "shutil.rmtree(TRAIN_FOLDER)\n",
    "shutil.rmtree(TEST_FOLDER)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now parallely operate on each batch to tokenize the data and preprocess the tokens. To do this, we create a PythonScript step below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing preprocess.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile preprocess.py\n",
    "# Copyright (c) Microsoft Corporation. All rights reserved.\n",
    "# Licensed under the MIT License.\n",
    "import argparse\n",
    "import logging\n",
    "import os\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "from utils_nlp.models.bert.common import Language, Tokenizer\n",
    "\n",
    "LABEL_COL = \"genre\"\n",
    "TEXT_COL = \"sentence1\"\n",
    "LANGUAGE = Language.ENGLISH\n",
    "TO_LOWER = True\n",
    "MAX_LEN = 150\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "\n",
    "def tokenize(df):\n",
    "    \"\"\"Tokenize the text documents and convert them to lists of tokens using the BERT tokenizer.\n",
    "    Args:\n",
    "        df(pd.Dataframe): Dataframe with training or test samples\n",
    "\n",
    "    Returns:\n",
    "\n",
    "        list: List of lists of tokens for train set.\n",
    "\n",
    "    \"\"\"\n",
    "    tokenizer = Tokenizer(\n",
    "        LANGUAGE, to_lower=TO_LOWER)\n",
    "    tokens = tokenizer.tokenize(list(df[TEXT_COL]))\n",
    "\n",
    "    return tokens\n",
    "\n",
    "\n",
    "def preprocess(tokens):\n",
    "    \"\"\" Preprocess method that does the following,\n",
    "            Convert the tokens into token indices corresponding to the BERT tokenizer's vocabulary\n",
    "            Add the special tokens [CLS] and [SEP] to mark the beginning and end of a sentence\n",
    "            Pad or truncate the token lists to the specified max length\n",
    "            Return mask lists that indicate paddings' positions\n",
    "            Return token type id lists that indicate which sentence the tokens belong to (not needed\n",
    "            for one-sequence classification)\n",
    "\n",
    "    Args:\n",
    "        tokens(pd.Dataframe): Dataframe with tokens for train set.\n",
    "\n",
    "    Returns:\n",
    "        list: List of lists of tokens for train or test set with special tokens added.\n",
    "        list: Input mask.\n",
    "    \"\"\"\n",
    "    tokenizer = Tokenizer(\n",
    "        LANGUAGE, to_lower=TO_LOWER)\n",
    "    tokens, mask, _ = tokenizer.preprocess_classification_tokens(\n",
    "        tokens, MAX_LEN\n",
    "    )\n",
    "\n",
    "    return tokens, mask\n",
    "\n",
    "\n",
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument(\"--input_data\", type=str, help=\"input data\")\n",
    "parser.add_argument(\"--output_data\", type=str, help=\"Path to the output file.\")\n",
    "\n",
    "args = parser.parse_args()\n",
    "input_data = args.input_data\n",
    "output_data = args.output_data\n",
    "output_dir = os.path.dirname(os.path.abspath(output_data))\n",
    "\n",
    "if output_dir is not None:\n",
    "    os.makedirs(output_dir, exist_ok=True)\n",
    "    logger.info(\"%s created\" % output_dir)\n",
    "\n",
    "df = pd.read_csv(args.input_data)\n",
    "tokens_array = tokenize(df)\n",
    "tokens_array, mask_array = preprocess(tokens_array)\n",
    "\n",
    "df['tokens'] = tokens_array\n",
    "df['mask'] = mask_array\n",
    "\n",
    "# Filter columns\n",
    "cols = ['tokens', 'mask', 'label']\n",
    "df = df[cols]\n",
    "df.to_csv(output_data, header=False, index=False)\n",
    "logger.info(\"Completed\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'../../utils_nlp/models/bert/preprocess.py'"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "preprocess_file = os.path.join(PROJECT_FOLDER,'utils_nlp/models/bert/preprocess.py')\n",
    "shutil.move('preprocess.py',preprocess_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a conda environment for the steps below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "conda_dependencies = CondaDependencies.create(\n",
    "    conda_packages=[\n",
    "        \"numpy\",\n",
    "        \"scikit-learn\",\n",
    "        \"pandas\",\n",
    "    ],\n",
    "    pip_packages=[\"azureml-sdk==1.0.43.*\", \n",
    "                  \"torch==1.1\", \n",
    "                  \"tqdm==4.31.1\",\n",
    "                 \"pytorch-pretrained-bert>=0.6\"],\n",
    "    python_version=\"3.6.8\",\n",
    ")\n",
    "run_config = RunConfiguration(conda_dependencies=conda_dependencies)\n",
    "run_config.environment.docker.enabled = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then create the list of steps that use the preprocess.py created above. We use the output of these steps as input to training in the next section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "processed_train_files = []\n",
    "processed_test_files = []\n",
    "ds = ws.get_default_datastore()\n",
    "\n",
    "for i in range(num_train_partitions):\n",
    "        input_data = DataReference(datastore=ds, \n",
    "                                   data_reference_name='train_batch_{}'.format(str(i)), \n",
    "                                   path_on_datastore='mnli_data/train/batch{}.csv'.format(str(i)),\n",
    "                                   overwrite=False)\n",
    "\n",
    "        output_data = PipelineData(name=\"train{}\".format(str(i)), datastore=ds,\n",
    "                       output_path_on_compute='mnli_data/processed_train/batch{}.csv'.format(str(i)))\n",
    "\n",
    "        step = PythonScriptStep(\n",
    "            name='preprocess_step_train_{}'.format(str(i)),\n",
    "            arguments=[\"--input_data\", input_data, \"--output_data\", output_data],\n",
    "            script_name= 'utils_nlp/models/bert/preprocess.py',\n",
    "            inputs=[input_data],\n",
    "            outputs=[output_data],\n",
    "            source_directory=PROJECT_FOLDER,\n",
    "            compute_target=compute_target,\n",
    "            runconfig=run_config,\n",
    "            allow_reuse=False,\n",
    "        )\n",
    "        \n",
    "        processed_train_files.append(output_data)         \n",
    "            \n",
    "for i in range(num_test_partitions):\n",
    "            input_data = DataReference(datastore=ds, \n",
    "                                       data_reference_name='test_batch_{}'.format(str(i)), \n",
    "                                       path_on_datastore='mnli_data/test/batch{}.csv'.format(str(i)),\n",
    "                                       overwrite=False)\n",
    "        \n",
    "            output_data = PipelineData(name=\"test{}\".format(str(i)), datastore=ds,\n",
    "                        output_path_on_compute='mnli_data/processed_test/batch{}.csv'.format(str(i)))\n",
    "            \n",
    "            step = PythonScriptStep(\n",
    "                name='preprocess_step_test_{}'.format(str(i)),\n",
    "                arguments=[\"--input_data\", input_data, \"--output_data\", output_data],\n",
    "                script_name= 'utils_nlp/models/bert/preprocess.py',\n",
    "                inputs=[input_data],\n",
    "                outputs=[output_data],\n",
    "                source_directory=PROJECT_FOLDER,\n",
    "                compute_target=compute_target,\n",
    "                runconfig=run_config,\n",
    "                allow_reuse=False,\n",
    "            )\n",
    "            \n",
    "            processed_test_files.append(output_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Train and Score"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the data is processed and available on datastore, we  train the classifier using the training examples. This involves fine-tuning the BERT Transformer and learning a linear classification layer on top of that. After training is complete we score the performance of the model on the test dataset\n",
    "\n",
    "The training is distributed and is done AzureML's capability to support distributed using MPI with horovod. \n",
    "\n",
    "**Please note** that training requires a GPU enabled cluster in AzureML Compute. We suggest using NC12. If you would like to change the GPU configuration, please changes `NUM_GPUS` variable accordingly.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Setup training script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing train.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile train.py\n",
    "# Copyright (c) Microsoft Corporation. All rights reserved.\n",
    "# Licensed under the MIT License.\n",
    "\n",
    "import argparse\n",
    "import json\n",
    "import logging\n",
    "import os\n",
    "import torch\n",
    "\n",
    "from sklearn.metrics import classification_report\n",
    "\n",
    "from utils_nlp.common.timer import Timer\n",
    "from utils_nlp.models.bert.common import Language, get_dataset_multiple_files\n",
    "from utils_nlp.models.bert.sequence_classification_distributed import (\n",
    "    BERTSequenceClassifier,\n",
    ")\n",
    "\n",
    "BATCH_SIZE = 32\n",
    "NUM_GPUS = 2\n",
    "NUM_EPOCHS = 1\n",
    "LABELS = [\"telephone\", \"government\", \"travel\", \"slate\", \"fiction\"]\n",
    "OUTPUT_DIR = \"./outputs/\"\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument(\n",
    "    \"--train_files\",\n",
    "    nargs=\"+\",\n",
    "    default=[],\n",
    "    help=\"List of file paths to all the files in train dataset.\",\n",
    ")\n",
    "\n",
    "parser.add_argument(\n",
    "    \"--test_files\",\n",
    "    nargs=\"+\",\n",
    "    default=[],\n",
    "    help=\"List of file paths to all the files in test dataset.\",\n",
    ")\n",
    "\n",
    "args = parser.parse_args()\n",
    "train_files = [file.strip() for file in args.train_files]\n",
    "test_files = [file.strip() for file in args.test_files]\n",
    "\n",
    "# Handle square brackets from train list\n",
    "train_files[0] = train_files[0][1:]\n",
    "train_files[len(train_files) - 1] = train_files[len(train_files) - 1][:-1]\n",
    "train_dataset = get_dataset_multiple_files(train_files)\n",
    "\n",
    "# Handle square brackets from test list\n",
    "test_files[0] = test_files[0][1:]\n",
    "test_files[len(test_files) - 1] = test_files[len(test_files) - 1][:-1]\n",
    "test_dataset = get_dataset_multiple_files(test_files)\n",
    "\n",
    "# Train\n",
    "classifier = BERTSequenceClassifier(\n",
    "    language=Language.ENGLISH, num_labels=len(LABELS), use_distributed=True\n",
    ")\n",
    "\n",
    "# Create data loaders.\n",
    "kwargs = (\n",
    "    {\"num_workers\": NUM_GPUS, \"pin_memory\": True} if torch.cuda.is_available() else {}\n",
    ")\n",
    "train_data_loader = classifier.create_data_loader(\n",
    "    train_dataset, batch_size=BATCH_SIZE, **kwargs\n",
    ")\n",
    "test_data_loader = classifier.create_data_loader(\n",
    "    test_dataset, batch_size=BATCH_SIZE, mode=\"test\", **kwargs\n",
    ")\n",
    "\n",
    "# Create optimizer\n",
    "num_examples = len(train_dataset)\n",
    "num_batches = int(num_examples / BATCH_SIZE)\n",
    "num_train_optimization_steps = num_batches * NUM_EPOCHS\n",
    "optimizer = classifier.create_optimizer(num_train_optimization_steps)\n",
    "\n",
    "with Timer() as t:\n",
    "    for epoch in range(1, NUM_EPOCHS + 1):\n",
    "        train_data_loader.sampler.set_epoch(epoch)\n",
    "        classifier.fit(\n",
    "            train_data_loader,\n",
    "            epoch=epoch,\n",
    "            bert_optimizer=optimizer,\n",
    "            num_gpus=NUM_GPUS,\n",
    "            num_epochs=NUM_EPOCHS,\n",
    "        )\n",
    "\n",
    "# Predict\n",
    "preds, labels_test = classifier.predict(test_data_loader, num_gpus=NUM_GPUS)\n",
    "\n",
    "# Evaluate\n",
    "results = classification_report(\n",
    "    labels_test, preds, target_names=LABELS, output_dict=True\n",
    ")\n",
    "\n",
    "# Write out results.\n",
    "classifier.save_model()\n",
    "result_file = os.path.join(OUTPUT_DIR, \"results.json\")\n",
    "with open(result_file, \"w+\") as fp:\n",
    "    json.dump(results, fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'../../utils_nlp/models/bert/train.py'"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_file = os.path.join(PROJECT_FOLDER,'utils_nlp/models/bert/train.py')\n",
    "shutil.move('train.py',train_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Create a Pytorch Estimator\n",
    "\n",
    "We create a Pytorch Estimator using AzureML SDK and additonally define an EstimatorStep to run it on AzureML pipelines.\n",
    "\n",
    "The Azure ML SDK's PyTorch Estimator allows us to submit PyTorch training jobs for both single-node and distributed runs. For more information on the PyTorch estimator, refer [here](https://docs.microsoft.com/azure/machine-learning/service/how-to-train-pytorch).\n",
    "\n",
    "This Estimator specifies that the training script will run on 4 nodes, with 2 worker per node. In order to execute a distributed run using GPU, we must define `use_gpu` and `distributed_backend` to use MPI/Horovod. PyTorch, Horovod, and other necessary dependencies are installed automatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING - framework_version is not specified, defaulting to version 1.1.\n"
     ]
    }
   ],
   "source": [
    "estimator = PyTorch(source_directory=PROJECT_FOLDER,\n",
    "                    compute_target=compute_target,\n",
    "                    entry_script='utils_nlp/models/bert/train.py',\n",
    "                    node_count= NODE_COUNT,\n",
    "                    distributed_training= MpiConfiguration(),\n",
    "                    use_gpu=True,\n",
    "                    conda_packages=['scikit-learn=0.20.3', 'numpy>=1.16.0', 'pandas'],\n",
    "                    pip_packages=[\"tqdm==4.31.1\",\"pytorch-pretrained-bert>=0.6\"]\n",
    "                   )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = processed_train_files + processed_test_files\n",
    "\n",
    "est_step = EstimatorStep(name=\"Estimator-Train\", \n",
    "                         estimator=estimator, \n",
    "                         estimator_entry_script_arguments=[\n",
    "                             '--train_files',  str(processed_train_files),\n",
    "                             '--test_files', str(processed_test_files)],\n",
    "                         inputs = inputs,\n",
    "                         runconfig_pipeline_params=None, \n",
    "                         compute_target=compute_target)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Submit the pipeline\n",
    "\n",
    "The model is fine tuned on AML Compute and takes **45 minutes** to train. The total time to run the pipeline will be around **1h 30 minutes** if you use the default value `max_epoch=1`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "pipeline = Pipeline(workspace=ws, steps=[est_step])\n",
    "experiment = Experiment(ws, 'NLP-TC-BERT-distributed')\n",
    "pipeline_run = experiment.submit(pipeline)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "060659321062486694c0acbb0184eeed",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "_PipelineWidget(widget_settings={'childWidgetDisplay': 'popup', 'send_telemetry': False, 'log_level': 'INFO', …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "RunDetails(pipeline_run).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "#If you would like to cancel the job for any reasons uncomment the code below.\n",
    "#pipeline_run.cancel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#wait for the run to complete before continuing in the notebook\n",
    "pipeline_run.wait_for_completion()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.4 Download and analyze results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading file outputs/results.json to ./outputs\\results.json...\n",
      "Downloading file outputs/bert-large-uncased to ./outputs\\bert-large-uncased...\n",
      "Downloading file outputs/bert_config.json to ./outputs\\bert_config.json...\n"
     ]
    }
   ],
   "source": [
    "step_run = pipeline_run.find_step_run(\"Estimator-Train\")[0]\n",
    "file_names = ['outputs/results.json', 'outputs/bert-large-uncased', 'outputs/bert_config.json' ]\n",
    "azureml_utils.get_output_files(step_run, './outputs', file_names=file_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              f1-score  precision    recall  support\n",
      "telephone     0.904130   0.843191  0.974563    629.0\n",
      "government    0.955857   0.972366  0.939900    599.0\n",
      "travel        0.839966   0.935849  0.761905    651.0\n",
      "slate         0.986411   0.974724  0.998382    618.0\n",
      "fiction       0.938871   0.918712  0.959936    624.0\n",
      "micro avg     0.925344   0.925344  0.925344   3121.0\n",
      "macro avg     0.925047   0.928968  0.926937   3121.0\n",
      "weighted avg  0.923913   0.928455  0.925344   3121.0\n"
     ]
    }
   ],
   "source": [
    "with open('outputs/results.json', 'r') as handle:\n",
    "    parsed = json.load(handle)\n",
    "    print(pd.DataFrame.from_dict(parsed).transpose())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the above chart we can notice the performance of the model trained on a distributed setup in AzureML Compute. From our comparison to fine tuning the same model on MNLI dataset on a `STANDARD_NC12` machine [here](tc_mnli_bert.ipynb) we notice a gain of 20% in the model training time with no drop in performance for AzureML Compute. We present the comparison of weight avg of the metrics along with the training time below,\n",
    "\n",
    "| Training Setup | F1-Score | Precision | Recall | Training Time |\n",
    "| --- | --- | --- | --- | --- |\n",
    "|Standard NC12 | 0.93 |0.93 |0.93 | 58 min |\n",
    "|AzureML Compute*|0.934| 0.934 | 0.934| 46 min |\n",
    "\n",
    "* AzureML Compute - The setup used 4 nodes with `STANDARD_NC12` machines.\n",
    "\n",
    "We also observe common tradeoffs associated with distributed training. We make use of [Horovod](https://github.com/horovod/horovod), a distributed training tool for many popular deep learning frameworks that enables parallelization of work across the nodes in the cluster. Distributed training decreases the time it takes for the model to converge in theory, but the model may also take more time in communicating with each node. Note that the communication time will eventually become negligible when training on larger and larger datasets, but being aware of this tradeoff is helpful for choosing the node configuration when training on smaller datasets. We expect the gains of using AzureML to increase with increased dataset size."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally clean up any intermediate files we created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.remove(train_file)\n",
    "os.remove(preprocess_file)"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python nlp_cpu",
   "language": "python",
   "name": "ame"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}