hyperparameter_tuning/huggingface_multiclass_text_classification_20_newsgroups/hpo_huggingface_text_classification_20_newsgroups.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Text Classification with Amazon SageMaker HuggingFace and Hyperparameter Tuning\n",
    "\n",
    "Automatic model tuning, also known as hyperparameter tuning, finds the best version of a model by running many jobs that test a range of hyperparameters on your dataset. You choose the tunable hyperparameters, a range of values for each, and an objective metric. You choose the objective metric from the metrics that the algorithm computes. Automatic model tuning searches the hyperparameters chosen to find the combination of values that result in the model that optimizes the objective metric.\n",
    "\n",
    "\n",
    "## Introduction\n",
    "\n",
    "Text Classification can be used to solve various use-cases like sentiment analysis, spam detection, hashtag prediction etc. \n",
    "\n",
    "\n",
    "This notebook demonstrates the use of the [HuggingFace Transformers library](https://huggingface.co/transformers/) together with a custom Amazon sagemaker-sdk extension to fine-tune a pre-trained transformer on multi class text classification. In particular, the pre-trained model will be fine-tuned using the [20 Newsgroups dataset](http://qwone.com/~jason/20Newsgroups/). To get started, we need to set up the environment with a few prerequisite steps, for permissions, configurations, and so on."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install Python packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "!{sys.executable} -m pip install \"scikit_learn==0.20.0\" \"sagemaker>=2.48.0\" \"transformers==4.6.1\" \"datasets[s3]==1.6.2\" \"nltk==3.4.4\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you run this notebook in SageMaker Studio, you need to make sure `ipywidgets` is installed and restart the kernel, so please uncomment the code in the next cell, and run it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%capture\n",
    "# import IPython\n",
    "# import sys\n",
    "\n",
    "# !{sys.executable} -m pip install ipywidgets\n",
    "# IPython.Application.instance().kernel.do_shutdown(True)  # has to restart kernel so changes are used"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Let's start by specifying:\n",
    "\n",
    "- The S3 bucket and prefix that you want to use for training and model data. This should be within the same region as the Notebook Instance, training, and hosting. If you don't specify a bucket, SageMaker SDK will create a default bucket following a pre-defined naming convention in the same region. \n",
    "- The IAM role ARN used to give SageMaker access to your data. It can be fetched using the **get_execution_role** method from sagemaker python SDK."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "isConfigCell": true
   },
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "from sagemaker import get_execution_role\n",
    "import json\n",
    "import boto3\n",
    "import pandas as pd\n",
    "import re\n",
    "import string\n",
    "from sklearn.model_selection import train_test_split\n",
    "import sagemaker.huggingface\n",
    "\n",
    "\n",
    "sess = sagemaker.Session()\n",
    "\n",
    "role = get_execution_role()\n",
    "print(\n",
    "    role\n",
    ")  # This is the role that SageMaker would use to leverage AWS resources (S3, CloudWatch) on your behalf\n",
    "\n",
    "bucket = sess.default_bucket()  # Replace with your own bucket name if needed\n",
    "print(bucket)\n",
    "s3_prefix = \"huggingface/20_newsgroups\"  # Replace with the prefix under which you want to store the data if needed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data Preparation\n",
    "\n",
    "Now we'll download a dataset from the web on which we want to train the text classification model.\n",
    "\n",
    "In this example, let us train the text classification model on the [20 Newsgroups dataset](http://qwone.com/~jason/20Newsgroups/). The 20 Newsgroups dataset consists of 20000 messages taken from 20 Usenet newsgroups."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "\n",
    "data_dir = \"20_newsgroups_bulk\"\n",
    "if os.path.exists(data_dir):  # cleanup existing data folder\n",
    "    shutil.rmtree(data_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!aws s3 cp s3://sagemaker-sample-files/datasets/text/20_newsgroups/20_newsgroups_bulk.tar.gz ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tar xzf 20_newsgroups_bulk.tar.gz\n",
    "!ls 20_newsgroups_bulk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_list = [os.path.join(data_dir, f) for f in os.listdir(data_dir)]\n",
    "print(\"Number of files:\", len(file_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "documents_count = 0\n",
    "for file in file_list:\n",
    "    df = pd.read_csv(file, header=None, names=[\"text\"])\n",
    "    documents_count = documents_count + df.shape[0]\n",
    "print(\"Number of documents:\", documents_count)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's inspect the dataset files and analyze the categories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "categories_list = [f.split(\"/\")[1] for f in file_list]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "categories_list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the dataset consists of 20 topics, each in different file."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us inspect the dataset to get some understanding about how the data and the label is provided in the dataset. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(\"./20_newsgroups_bulk/rec.motorcycles\", header=None, names=[\"text\"])\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"text\"][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(\"./20_newsgroups_bulk/comp.sys.mac.hardware\", header=None, names=[\"text\"])\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"text\"][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see from the above, there is a single file for each class in the dataset. Each record is just a plain text paragraphs with header, body, footer and quotes. We will need to process them into a suitable data format."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Preprocessing\n",
    "We need to preprocess the dataset to remove the header, footer, quotes, leading/trailing whitespace, extra spaces, tabs, and HTML tags/markups. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Download the `nltk` tokenizer and other libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nltk\n",
    "from nltk.tokenize import word_tokenize\n",
    "\n",
    "nltk.download(\"punkt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.datasets.twenty_newsgroups import (\n",
    "    strip_newsgroup_header,\n",
    "    strip_newsgroup_quoting,\n",
    "    strip_newsgroup_footer,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This following function will remove the header, footer and quotes (of earlier messages in each text)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def strip_newsgroup_item(item):\n",
    "    item = strip_newsgroup_header(item)\n",
    "    item = strip_newsgroup_quoting(item)\n",
    "    item = strip_newsgroup_footer(item)\n",
    "    return item"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following function will take care of removing leading/trailing whitespace, extra spaces, tabs, and HTML tags/markups."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_text(texts):\n",
    "    final_text_list = []\n",
    "    for text in texts:\n",
    "\n",
    "        # Check if the sentence is a missing value\n",
    "        if isinstance(text, str) == False:\n",
    "            text = \"\"\n",
    "\n",
    "        filtered_sentence = []\n",
    "\n",
    "        # Lowercase\n",
    "        text = text.lower()\n",
    "\n",
    "        # Remove leading/trailing whitespace, extra space, tabs, and HTML tags/markups\n",
    "        text = text.strip()\n",
    "        text = re.sub(\"\\[.*?\\]\", \"\", text)\n",
    "        text = re.sub(\"https?://\\S+|www\\.\\S+\", \"\", text)\n",
    "        text = re.sub(\"<.*?>+\", \"\", text)\n",
    "        text = re.sub(\"[%s]\" % re.escape(string.punctuation), \"\", text)\n",
    "        text = re.sub(\"\\n\", \"\", text)\n",
    "        text = re.sub(\"\\w*\\d\\w*\", \"\", text)\n",
    "\n",
    "        for w in word_tokenize(text):\n",
    "            # We are applying some custom filtering here, feel free to try different things\n",
    "            # Check if it is not numeric\n",
    "            if not w.isnumeric():\n",
    "                filtered_sentence.append(w)\n",
    "        final_string = \" \".join(filtered_sentence)  # final string of cleaned words\n",
    "\n",
    "        final_text_list.append(final_string)\n",
    "\n",
    "    return final_text_list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we will read each of the `20_newsgroups` dataset files, call `strip_newsgroup_item` and `process_text` functions we defined earlier, and then aggregate all data into one dataframe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_categories_df = pd.DataFrame()\n",
    "\n",
    "for file in file_list:\n",
    "    print(f\"Processing {file}\")\n",
    "    label = file.split(\"/\")[1]\n",
    "    df = pd.read_csv(file, header=None, names=[\"text\"])\n",
    "    df[\"text\"] = df[\"text\"].apply(strip_newsgroup_item)\n",
    "    df[\"text\"] = process_text(df[\"text\"].tolist())\n",
    "    df[\"label\"] = label\n",
    "    all_categories_df = all_categories_df.append(df, ignore_index=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's inspect how many categories there are in our dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_categories_df[\"label\"].value_counts()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In our dataset there are 20 categories which is too much, so we will combine the sub-categories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# replace to politics\n",
    "all_categories_df[\"label\"].replace(\n",
    "    {\n",
    "        \"talk.politics.misc\": \"politics\",\n",
    "        \"talk.politics.guns\": \"politics\",\n",
    "        \"talk.politics.mideast\": \"politics\",\n",
    "    },\n",
    "    inplace=True,\n",
    ")\n",
    "\n",
    "# replace to recreational\n",
    "all_categories_df[\"label\"].replace(\n",
    "    {\n",
    "        \"rec.sport.hockey\": \"recreational\",\n",
    "        \"rec.sport.baseball\": \"recreational\",\n",
    "        \"rec.autos\": \"recreational\",\n",
    "        \"rec.motorcycles\": \"recreational\",\n",
    "    },\n",
    "    inplace=True,\n",
    ")\n",
    "\n",
    "# replace to religion\n",
    "all_categories_df[\"label\"].replace(\n",
    "    {\n",
    "        \"soc.religion.christian\": \"religion\",\n",
    "        \"talk.religion.misc\": \"religion\",\n",
    "        \"alt.atheism\": \"religion\",\n",
    "    },\n",
    "    inplace=True,\n",
    ")\n",
    "\n",
    "# replace to computer\n",
    "all_categories_df[\"label\"].replace(\n",
    "    {\n",
    "        \"comp.windows.x\": \"computer\",\n",
    "        \"comp.sys.ibm.pc.hardware\": \"computer\",\n",
    "        \"comp.os.ms-windows.misc\": \"computer\",\n",
    "        \"comp.graphics\": \"computer\",\n",
    "        \"comp.sys.mac.hardware\": \"computer\",\n",
    "    },\n",
    "    inplace=True,\n",
    ")\n",
    "# replace to sales\n",
    "all_categories_df[\"label\"].replace({\"misc.forsale\": \"sales\"}, inplace=True)\n",
    "\n",
    "# replace to science\n",
    "all_categories_df[\"label\"].replace(\n",
    "    {\n",
    "        \"sci.crypt\": \"science\",\n",
    "        \"sci.electronics\": \"science\",\n",
    "        \"sci.med\": \"science\",\n",
    "        \"sci.space\": \"science\",\n",
    "    },\n",
    "    inplace=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we are left with 6 categories, which is much better."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_categories_df[\"label\"].value_counts()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's calculate number of words for each row."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_categories_df[\"word_count\"] = all_categories_df[\"text\"].apply(lambda x: len(str(x).split()))\n",
    "all_categories_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's get basic statistics about the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_categories_df[\"word_count\"].describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the mean value is around 159 words. However, there are outliers, such as a text with 11351 words. This can make it harder for the model to result in good performance. We will take care to drop those rows."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's drop empty rows first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "no_text = all_categories_df[all_categories_df[\"word_count\"] == 0]\n",
    "print(len(no_text))\n",
    "\n",
    "# drop these rows\n",
    "all_categories_df.drop(no_text.index, inplace=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's drop the rows that are longer than 256 words, as it is a length close to the mean value of the word count. This is done to make it easy for the model to train without outliers. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "long_text = all_categories_df[all_categories_df[\"word_count\"] > 256]\n",
    "print(len(long_text))\n",
    "\n",
    "# drop these rows\n",
    "all_categories_df.drop(long_text.index, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_categories_df[\"label\"].value_counts()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's get basic statistics about the dataset after our outliers fixes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_categories_df[\"word_count\"].describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This looks much more balanced."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we drop the `word_count` columns as we will not need it anymore."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_categories_df.drop(columns=\"word_count\", axis=1, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_categories_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's convert categorical label to integer number, in order to prepare the dataset for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "categories = all_categories_df[\"label\"].unique().tolist()\n",
    "categories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "categories.index(\"recreational\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_categories_df[\"label\"] = all_categories_df[\"label\"].apply(lambda x: categories.index(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_categories_df[\"label\"].value_counts()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We partition the dataset into 80% training and 20% validation set and save to `csv` files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df, test_df = train_test_split(all_categories_df, test_size=0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df.to_csv(\"train.csv\", index=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_df.to_csv(\"test.csv\", index=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's inspect the label distribution in the training dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df[\"label\"].value_counts()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's inspect the label distribution in the test dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_df[\"label\"].value_counts()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tokenization "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A tokenizer is in charge of preparing the inputs for a model. The library contains tokenizers for all the models. Most of the tokenizers are available in two flavors: a full python implementation and a “Fast” implementation based on the Rust library [tokenizers](https://github.com/huggingface/tokenizers). The “Fast” implementations allows:\n",
    "\n",
    " - A significant speed-up in particular when doing batched tokenization.\n",
    " - Additional methods to map between the original string (character and words) and the token space (e.g. getting the index of the token comprising a given character or the span of characters corresponding to a given token). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "# tokenizer used in preprocessing\n",
    "tokenizer_name = \"distilbert-base-uncased\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# download tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load train and test datasets\n",
    "\n",
    "Let's create a [Dataset](https://huggingface.co/docs/datasets/loading_datasets.html) from our local `csv` files for training and test we saved earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = load_dataset(\"csv\", data_files={\"train\": \"train.csv\", \"test\": \"test.csv\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset[\"train\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset[\"train\"][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset[\"test\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset[\"test\"][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# tokenizer helper function\n",
    "def tokenize(batch):\n",
    "    return tokenizer(batch[\"text\"], padding=\"max_length\", truncation=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = dataset[\"train\"]\n",
    "test_dataset = dataset[\"test\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tokenize train and test datasets\n",
    "\n",
    "Let's tokenize the train dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = train_dataset.map(tokenize, batched=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's tokenize the test dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset = test_dataset.map(tokenize, batched=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set format for PyTorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = train_dataset.rename_column(\"label\", \"labels\")\n",
    "train_dataset.set_format(\"torch\", columns=[\"input_ids\", \"attention_mask\", \"labels\"])\n",
    "test_dataset = test_dataset.rename_column(\"label\", \"labels\")\n",
    "test_dataset.set_format(\"torch\", columns=[\"input_ids\", \"attention_mask\", \"labels\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Uploading data to `sagemaker_session_bucket`\n",
    "\n",
    "After we processed the datasets, we are going to upload it to S3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import botocore\n",
    "from datasets.filesystems import S3FileSystem\n",
    "\n",
    "s3 = S3FileSystem()\n",
    "\n",
    "# save train_dataset to s3\n",
    "training_input_path = f\"s3://{sess.default_bucket()}/{s3_prefix}/train\"\n",
    "train_dataset.save_to_disk(training_input_path, fs=s3)\n",
    "\n",
    "# save test_dataset to s3\n",
    "test_input_path = f\"s3://{sess.default_bucket()}/{s3_prefix}/test\"\n",
    "test_dataset.save_to_disk(test_input_path, fs=s3)\n",
    "\n",
    "print(training_input_path)\n",
    "print(test_input_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up hyperparameter tuning job\n",
    "Now that we are done with all the setup that is needed, we are ready to train our HuggingFace model. To begin, let us create a `HuggingFace` estimator object. This estimator will launch the training job."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training the HuggingFace model for supervised text classification\n",
    "\n",
    "In order to create a sagemaker training job we need a `HuggingFace` Estimator. The Estimator handles end-to-end Amazon SageMaker training and deployment tasks. In an Estimator we define, which fine-tuning script should be used as `entry_point`, which `instance_type` should be used, which `hyperparameters` are passed in .....\n",
    "\n",
    "\n",
    "\n",
    "```python\n",
    "huggingface_estimator = HuggingFace(entry_point='train.py',\n",
    "                            source_dir='./code',\n",
    "                            instance_type='ml.p3.2xlarge',\n",
    "                            instance_count=1,\n",
    "                            volume_size=256,\n",
    "                            role=role,\n",
    "                            transformers_version='4.6',\n",
    "                            pytorch_version='1.7',\n",
    "                            py_version='py36',\n",
    "                            hyperparameters = {'epochs': 1,\n",
    "                                               'model_name':'distilbert-base-uncased',\n",
    "                                               'num_labels': 6\n",
    "                                              })\n",
    "```\n",
    "\n",
    "When we create a SageMaker training job, SageMaker takes care of starting and managing all the required ec2 instances for us with the `huggingface` container, uploads the provided fine-tuning script `train.py` and downloads the data from our `sagemaker_session_bucket` into the container at `/opt/ml/input/data`. Then, it starts the training job by running. \n",
    "\n",
    "```python\n",
    "/opt/conda/bin/python train.py --epochs 1 --model_name distilbert-base-uncased --num_labels 6\n",
    "```\n",
    "\n",
    "The `hyperparameters` you define in the `HuggingFace` estimator are passed in as named arguments. \n",
    "\n",
    "SageMaker is providing useful properties about the training environment through various environment variables, including the following:\n",
    "\n",
    "* `SM_MODEL_DIR`: A string that represents the path where the training job writes the model artifacts to. After training, artifacts in this directory are uploaded to S3 for model hosting.\n",
    "\n",
    "* `SM_NUM_GPUS`: An integer representing the number of GPUs available to the host.\n",
    "\n",
    "* `SM_CHANNEL_XXXX:` A string that represents the path to the directory that contains the input data for the specified channel. For example, if you specify two input channels in the HuggingFace estimator’s fit call, named `train` and `test`, the environment variables `SM_CHANNEL_TRAIN` and `SM_CHANNEL_TEST` are set.\n",
    "\n",
    "\n",
    "To run your training job locally you can define `instance_type='local'` or `instance_type='local-gpu'` for `gpu` usage. _Note: this does not work within SageMaker Studio_\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We create a metric_definition dictionary that contains regex-based definitions that will be used to parse the job logs and extract metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metric_definitions = [\n",
    "    {\"Name\": \"loss\", \"Regex\": \"'loss': ([0-9]+(.|e\\-)[0-9]+),?\"},\n",
    "    {\"Name\": \"learning_rate\", \"Regex\": \"'learning_rate': ([0-9]+(.|e\\-)[0-9]+),?\"},\n",
    "    {\"Name\": \"eval_loss\", \"Regex\": \"'eval_loss': ([0-9]+(.|e\\-)[0-9]+),?\"},\n",
    "    {\"Name\": \"eval_accuracy\", \"Regex\": \"'eval_accuracy': ([0-9]+(.|e\\-)[0-9]+),?\"},\n",
    "    {\"Name\": \"eval_f1\", \"Regex\": \"'eval_f1': ([0-9]+(.|e\\-)[0-9]+),?\"},\n",
    "    {\"Name\": \"eval_precision\", \"Regex\": \"'eval_precision': ([0-9]+(.|e\\-)[0-9]+),?\"},\n",
    "    {\"Name\": \"eval_recall\", \"Regex\": \"'eval_recall': ([0-9]+(.|e\\-)[0-9]+),?\"},\n",
    "    {\"Name\": \"eval_runtime\", \"Regex\": \"'eval_runtime': ([0-9]+(.|e\\-)[0-9]+),?\"},\n",
    "    {\n",
    "        \"Name\": \"eval_samples_per_second\",\n",
    "        \"Regex\": \"'eval_samples_per_second': ([0-9]+(.|e\\-)[0-9]+),?\",\n",
    "    },\n",
    "    {\"Name\": \"epoch\", \"Regex\": \"'epoch': ([0-9]+(.|e\\-)[0-9]+),?\"},\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.huggingface import HuggingFace\n",
    "\n",
    "# hyperparameters, which are passed into the training job\n",
    "hyperparameters = {\"epochs\": 1, \"model_name\": \"distilbert-base-uncased\", \"num_labels\": 6}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's define the SageMaker `HuggingFace` estimator with resource configurations and hyperparameters to train Text Classification on 20 Newsgroups dataset, running on a `p3.2xlarge` instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "huggingface_estimator = HuggingFace(\n",
    "    entry_point=\"train.py\",\n",
    "    source_dir=\"./code\",\n",
    "    instance_type=\"ml.p3.2xlarge\",\n",
    "    instance_count=1,\n",
    "    volume_size=256,\n",
    "    role=role,\n",
    "    transformers_version=\"4.6\",\n",
    "    pytorch_version=\"1.7\",\n",
    "    py_version=\"py36\",\n",
    "    hyperparameters=hyperparameters,\n",
    "    metric_definitions=metric_definitions,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once we've defined our estimator we can specify the hyperparameters we'd like to tune and their possible values.  We have three different types of hyperparameters.\n",
    "- Categorical parameters need to take one value from a discrete set.  We define this by passing the list of possible values to `CategoricalParameter(list)`\n",
    "- Continuous parameters can take any real number value between the minimum and maximum value, defined by `ContinuousParameter(min, max)`\n",
    "- Integer parameters can take any integer value between the minimum and maximum value, defined by `IntegerParameter(min, max)`\n",
    "\n",
    "*Note, if possible, it's almost always best to specify a value as the least restrictive type.  For example, tuning learning rate as a continuous value between 0.01 and 0.2 is likely to yield a better result than tuning as a categorical parameter with values 0.01, 0.1, 0.15, or 0.2.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.tuner import (\n",
    "    IntegerParameter,\n",
    "    CategoricalParameter,\n",
    "    ContinuousParameter,\n",
    "    HyperparameterTuner,\n",
    ")\n",
    "\n",
    "hyperparameter_ranges = {\n",
    "    \"train_batch_size\": IntegerParameter(8, 32),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we'll specify the objective metric that we'd like to tune and its definition, which includes the regular expression (Regex) needed to extract that metric from the CloudWatch logs of the training job.  If you bring your own algorithm, your algorithm emits metrics by itself. In that case, you'll need to add a `MetricDefinition` object here to define the format of those metrics through regex, so that SageMaker knows how to extract those metrics from your CloudWatch logs.\n",
    "\n",
    "In this case, we elected to monitor `eval_accuracy` as you can see below. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "objective_metric_name = \"eval_accuracy\"\n",
    "objective_type = \"Maximize\"\n",
    "hpo_metric_definitions = [\n",
    "    {\"Name\": \"eval_accuracy\", \"Regex\": \"'eval_accuracy': ([0-9]+(.|e\\-)[0-9]+),?\"}\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we'll create a `HyperparameterTuner` object, to which we pass:\n",
    "- The `HuggingFace` estimator we created above\n",
    "- Our hyperparameter ranges\n",
    "- Objective metric name and definition\n",
    "- Tuning resource configurations such as Number of training jobs to run in total and how many training jobs can be run in parallel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tuner = HyperparameterTuner(\n",
    "    huggingface_estimator,\n",
    "    objective_metric_name,\n",
    "    hyperparameter_ranges,\n",
    "    hpo_metric_definitions,\n",
    "    max_jobs=6,\n",
    "    max_parallel_jobs=3,\n",
    "    objective_type=objective_type,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Launch hyperparameter tuning job\n",
    "Now we can launch a hyperparameter tuning job by calling *fit()* function. After the hyperparameter tuning job is created, we can go to SageMaker console to track the progress of the hyperparameter tuning job until it is completed.\n",
    "\n",
    "This should take around 28 minutes to complete."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "tuner.fit({\"train\": training_input_path, \"test\": test_input_path}, logs=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analyze Results of a Hyperparameter Tuning job\n",
    "\n",
    "Once you have completed a tuning job, (or even while the job is still running) you can use the code below to analyze the results to understand how each hyperparameter effects the quality of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm_client = boto3.Session().client(\"sagemaker\")\n",
    "\n",
    "tuning_job_name = tuner.latest_tuning_job.name\n",
    "tuning_job_name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Track hyperparameter tuning job progress\n",
    "After you launch a tuning job, you can see its progress by calling `describe_tuning_job` API. The output from describe-tuning-job is a JSON object that contains information about the current state of the tuning job. You can call `list_training_jobs_for_tuning_job` to see a detailed list of the training jobs that the tuning job launched."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tuning_job_result = sm_client.describe_hyper_parameter_tuning_job(\n",
    "    HyperParameterTuningJobName=tuning_job_name\n",
    ")\n",
    "\n",
    "status = tuning_job_result[\"HyperParameterTuningJobStatus\"]\n",
    "if status != \"Completed\":\n",
    "    print(\"Reminder: the tuning job has not been completed.\")\n",
    "\n",
    "job_count = tuning_job_result[\"TrainingJobStatusCounters\"][\"Completed\"]\n",
    "print(\"%d training jobs have completed\" % job_count)\n",
    "\n",
    "is_minimize = (\n",
    "    tuning_job_result[\"HyperParameterTuningJobConfig\"][\"HyperParameterTuningJobObjective\"][\"Type\"]\n",
    "    != \"Maximize\"\n",
    ")\n",
    "objective_name = tuning_job_result[\"HyperParameterTuningJobConfig\"][\n",
    "    \"HyperParameterTuningJobObjective\"\n",
    "][\"MetricName\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "if tuning_job_result.get(\"BestTrainingJob\", None):\n",
    "    print(\"Best model found so far:\")\n",
    "    pprint(tuning_job_result[\"BestTrainingJob\"])\n",
    "else:\n",
    "    print(\"No training jobs have reported results yet.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fetch all results as `DataFrame`\n",
    "We can list hyperparameters and objective metrics of all training jobs and pick up the training job with the best objective metric."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "tuner_analytics = sagemaker.HyperparameterTuningJobAnalytics(tuning_job_name)\n",
    "\n",
    "full_df = tuner_analytics.dataframe()\n",
    "\n",
    "if len(full_df) > 0:\n",
    "    df = full_df[full_df[\"FinalObjectiveValue\"] > -float(\"inf\")]\n",
    "    if len(df) > 0:\n",
    "        df = df.sort_values(\"FinalObjectiveValue\", ascending=is_minimize)\n",
    "        print(\"Number of training jobs with valid objective: %d\" % len(df))\n",
    "        print({\"lowest\": min(df[\"FinalObjectiveValue\"]), \"highest\": max(df[\"FinalObjectiveValue\"])})\n",
    "        pd.set_option(\"display.max_colwidth\", -1)  # Don't truncate TrainingJobName\n",
    "    else:\n",
    "        print(\"No training jobs have reported valid results yet.\")\n",
    "\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deploy the best trained model\n",
    "Once the training is done, we can deploy the trained model as an Amazon SageMaker real-time hosted endpoint. This will allow us to make predictions (or inference) from the model. Note that we don't have to host on the same type of instance that we used to train, because usually for inference, less compute power is needed than for training, and in addition, instance endpoints will be up and running for long, it's advisable to choose a cheaper instance for inference.\n",
    "\n",
    "- `ml.p3.2xlarge` - deliver high performance compute in the cloud with up to 8 NVIDIA® V100 Tensor Core GPUs and up to 100 `Gbps` of networking throughput for machine learning and HPC applications.\n",
    "- `ml.g4dn.xlarge` - the industry’s most cost-effective and versatile GPU instances for deploying machine learning models such as image classification, object detection, and speech recognition, and for graphics-intensive applications such as remote graphics workstations, game streaming, and graphics rendering."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor = tuner.deploy(1, \"ml.g4dn.xlarge\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we use the returned predictor object to call the endpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_sentence(sentence):\n",
    "    result = predictor.predict({\"inputs\": sentence})\n",
    "    index = int(result[0][\"label\"].split(\"LABEL_\")[1])\n",
    "    print(categories[index])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sentences = [\n",
    "    \"The modem is an internal AT/(E)ISA 8-bit card (just a little longer than a half-card).\",\n",
    "    \"In the cage I usually wave to bikers.  They usually don't wave back.  My wife thinks it's strange but I don't care.\",\n",
    "    \"Voyager has the unusual luck to be on a stable trajectory out of the solar system.\",\n",
    "]\n",
    "\n",
    "# using the same processing logic that we used during data preparation for training\n",
    "processed_sentences = process_text(sentences)\n",
    "\n",
    "for sentence in processed_sentences:\n",
    "    predict_sentence(sentence)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Clean up\n",
    "Endpoints should be deleted when no longer in use, since (per the [SageMaker pricing page](https://aws.amazon.com/sagemaker/pricing/)) they're billed by time deployed.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor.delete_endpoint()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "instance_type": "ml.t3.medium",
  "kernelspec": {
   "display_name": "Python 3 (PyTorch 1.6 Python 3.6 CPU Optimized)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-east-1:081325390199:image/pytorch-1.6-cpu-py36-ubuntu16.04-v1"
  },
  "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.13"
  },
  "notice": "Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.  Licensed under the Apache License, Version 2.0 (the \"License\"). You may not use this file except in compliance with the License. A copy of the License is located at http://aws.amazon.com/apache2.0/ or in the \"license\" file accompanying this file. This file is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."
 },
 "nbformat": 4,
 "nbformat_minor": 4
}