src/toxicity_type_detection.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Toxicity Type Detection\n",
    "\n",
    "![SageMaker](https://img.shields.io/badge/SageMaker-%23FF9900.svg?style=for-the-badge&logo=amazon-aws&logoColor=white)\n",
    "\n",
    "This notebook is a part of [ToChiquinho](https://dougtrajano.github.io/ToChiquinho/) project, which trains a model to detect toxicity types in a Portuguese text using the [OLID-BR](https://dougtrajano.github.io/olid-br/) dataset.\n",
    "\n",
    "The model is trained using [Amazon SageMaker](https://aws.amazon.com/sagemaker/).\n",
    "\n",
    "- [Parameters](#Parameters)\n",
    "- [Setup](#Setup)\n",
    "- [Prepare the data](#Prepare-the-data)\n",
    "  - [Uploading the data to S3](#Uploading-the-data-to-S3)\n",
    "- [Training process](#Training-process)\n",
    "  - [Define the estimator](#Define-the-estimator)\n",
    "  - [Hyperparameter tuning](#Hyperparameter-tuning)\n",
    "  - [Training model with best hyperparameters](#Training-model-with-best-hyperparameters)\n",
    "- [Documentation](#Documentation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameters\n",
    "\n",
    "In this section, we define the parameters used in the notebook.\n",
    "\n",
    "The `Parameters` class is used to store the parameters and it has the description of each parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from dataclasses import dataclass, field\n",
    "from typing import Any, Dict, Optional\n",
    "\n",
    "@dataclass\n",
    "class Parameters:\n",
    "    num_train_epochs: Optional[int] = field(\n",
    "        default=5,\n",
    "        metadata={\n",
    "            \"help\": \"Number of training epochs\",\n",
    "        }\n",
    "    )\n",
    "\n",
    "    num_train_epochs_per_child: Optional[int] = field(\n",
    "        default=1,\n",
    "        metadata={\n",
    "            \"help\": \"The number of epochs to train the model. An epoch is an iteration over the entire training set.\"\n",
    "        }\n",
    "    )\n",
    "\n",
    "    batch_size: Optional[int] = field(\n",
    "        default=8,\n",
    "        metadata={\n",
    "            \"help\": \"Batch size for training and evaluation.\"\n",
    "        }\n",
    "    )\n",
    "    \n",
    "    validation_split: Optional[float] = field(\n",
    "        default=0.2,\n",
    "        metadata={\n",
    "            \"help\": \"The percentage of the training set to use as validation set.\"\n",
    "        }\n",
    "    )\n",
    "    \n",
    "    seed: Optional[int] = field(\n",
    "        default=1993,\n",
    "        metadata={\n",
    "            \"help\": \"The seed to use for random number generation.\"\n",
    "        }\n",
    "    )\n",
    "\n",
    "    mlflow_tracking_uri: Optional[str] = field(\n",
    "        default=os.environ.get(\"MLFLOW_TRACKING_URI\"),\n",
    "        repr=False,\n",
    "        metadata={\n",
    "            \"help\": \"The URI of the MLFlow tracking server.\"\n",
    "        }\n",
    "    )\n",
    "\n",
    "    mlflow_experiment_name: Optional[str] = field(\n",
    "        default=os.environ.get(\"MLFLOW_EXPERIMENT_NAME\", \"Default\"),\n",
    "        metadata={\n",
    "            \"help\": \"The name of the MLFlow experiment.\"\n",
    "        }\n",
    "    )\n",
    "\n",
    "    mlflow_tags: Optional[Dict[str, Any]] = field(\n",
    "        default=None,\n",
    "        metadata={\n",
    "            \"help\": \"The tags to use for the MLFlow run.\"\n",
    "        }\n",
    "    )\n",
    "    \n",
    "    mlflow_tracking_username: Optional[str] = field(\n",
    "        default=None,\n",
    "        repr=False,\n",
    "        metadata={\n",
    "            \"help\": \"The username to use to authenticate with the MLFlow tracking server.\"\n",
    "        }\n",
    "    )\n",
    "\n",
    "    mlflow_tracking_password: Optional[str] = field(\n",
    "        default=None,\n",
    "        repr=False,\n",
    "        metadata={\n",
    "            \"help\": \"The password to use to authenticate with the MLFlow tracking server.\"\n",
    "        }\n",
    "    )\n",
    "\n",
    "    mlflow_run_id: Optional[str] = field(\n",
    "        default=os.environ.get(\"MLFLOW_RUN_ID\"),\n",
    "        repr=False,\n",
    "        metadata={\n",
    "            \"help\": \"The ID of the MLFlow run.\"\n",
    "        }\n",
    "    )\n",
    "\n",
    "    sagemaker_tuning_job_name: Optional[str] = field(\n",
    "        default=None,\n",
    "        repr=False,\n",
    "        metadata={\n",
    "            \"help\": \"The name of the SageMaker hyperparameter tuning job.\"\n",
    "        }\n",
    "    )\n",
    "\n",
    "    sagemaker_execution_role_arn: Optional[str] = field(\n",
    "        default=None,\n",
    "        repr=False,\n",
    "        metadata={\n",
    "            \"help\": \"The ARN of the SageMaker execution role.\"\n",
    "        }\n",
    "    )\n",
    "\n",
    "    aws_profile_name: Optional[str] = field(\n",
    "        default=\"default\",\n",
    "        repr=False,\n",
    "        metadata={\n",
    "            \"help\": \"The name of the AWS profile to use.\"\n",
    "        }\n",
    "    )\n",
    "\n",
    "    def __post_init__(self):\n",
    "        if self.mlflow_tracking_uri is not None:\n",
    "            os.environ[\"MLFLOW_TRACKING_URI\"] = self.mlflow_tracking_uri\n",
    "        if self.mlflow_experiment_name is not None:\n",
    "            os.environ[\"MLFLOW_EXPERIMENT_NAME\"] = self.mlflow_experiment_name\n",
    "        if self.mlflow_tracking_username is not None:\n",
    "            os.environ[\"MLFLOW_TRACKING_USERNAME\"] = self.mlflow_tracking_username\n",
    "        if self.mlflow_tracking_password is not None:\n",
    "            os.environ[\"MLFLOW_TRACKING_PASSWORD\"] = self.mlflow_tracking_password        \n",
    "        if self.mlflow_run_id is not None:\n",
    "            os.environ[\"MLFLOW_RUN_ID\"] = self.mlflow_run_id\n",
    "        \n",
    "        if isinstance(self.mlflow_tags, dict):\n",
    "            self.mlflow_tags = json.dumps(self.mlflow_tags)\n",
    "            os.environ[\"MLFLOW_TAGS\"] = self.mlflow_tags\n",
    "        elif self.mlflow_tags is not None:\n",
    "            raise ValueError(\"The mlflow_tags parameter must be a dictionary.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the next cell, you can change the parameters to fit your needs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameters(num_train_epochs=5, num_train_epochs_per_child=1, batch_size=8, validation_split=0.2, seed=1993, mlflow_experiment_name='toxicity-type-detection', mlflow_tags='{\"project\": \"ToChiquinho\", \"dataset\": \"OLID-BR\", \"model_type\": \"bert\", \"problem_type\": \"multi_label_classification\"}')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "params = Parameters()\n",
    "\n",
    "params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import sagemaker\n",
    "\n",
    "sagemaker_session = sagemaker.Session(\n",
    "    boto_session=boto3.Session(profile_name=params.aws_profile_name)\n",
    ")\n",
    "\n",
    "bucket_name = sagemaker_session.default_bucket()\n",
    "prefix = \"ToChiquinho/toxicity-type-detection\"\n",
    "\n",
    "if params.sagemaker_execution_role_arn is None:\n",
    "    params.sagemaker_execution_role_arn = sagemaker.get_execution_role(sagemaker_session)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def remove_checkpoints(\n",
    "        bucket_name: str,\n",
    "        checkpoint_prefix: str,\n",
    "        aws_profile_name: str = \"default\"):\n",
    "    \"\"\"\n",
    "    Remove all checkpoints from the specified S3 bucket and prefix.\n",
    "\n",
    "    Args:\n",
    "    - bucket_name: The name of the S3 bucket.\n",
    "    - checkpoint_prefix: The prefix of the checkpoints to remove.\n",
    "    - aws_profile_name: The name of the AWS profile to use.\n",
    "    \"\"\"\n",
    "    session = boto3.Session(profile_name=aws_profile_name)\n",
    "    s3 = session.resource(\"s3\")\n",
    "    bucket = s3.Bucket(bucket_name)\n",
    "    response = bucket.objects.filter(Prefix=checkpoint_prefix).delete()\n",
    "\n",
    "    if len(response) == 0:\n",
    "        print(\"No checkpoints found.\")\n",
    "    elif response[0][\"ResponseMetadata\"][\"HTTPStatusCode\"] == 200:\n",
    "        count = len(response[0][\"Deleted\"])\n",
    "        print(f\"Deleted {count} checkpoints.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare the data\n",
    "\n",
    "In this section, we will prepare the data to be used in the training process.\n",
    "\n",
    "We will download OLID-BR dataset from [HuggingFace Datasets](https://huggingface.co/datasets/olidbr), process it and upload it to S3 to be used in the training process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\"dougtrajano/olid-br\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "from typing import Union\n",
    "\n",
    "def prepare_dataset(\n",
    "    dataset: Union[datasets.Dataset, datasets.DatasetDict],\n",
    "    test_size: float = 0.2,\n",
    "    seed: int = 42\n",
    ") -> Union[datasets.Dataset, datasets.DatasetDict]:\n",
    "\n",
    "    # Filter only rows with is_offensive = \"OFF\"\n",
    "    dataset = dataset.filter(lambda example: example[\"is_offensive\"] == \"OFF\")\n",
    "\n",
    "    # Filter only offensive comments with at least one toxicity label\n",
    "    labels = [\n",
    "        \"health\",\n",
    "        \"ideology\",\n",
    "        \"insult\",\n",
    "        \"lgbtqphobia\",\n",
    "        \"other_lifestyle\",\n",
    "        \"physical_aspects\",\n",
    "        \"profanity_obscene\",\n",
    "        \"racism\",\n",
    "        \"sexism\",\n",
    "        \"xenophobia\"\n",
    "    ]\n",
    "\n",
    "    dataset = dataset.filter(\n",
    "        lambda example: any([example[label] == True for label in labels])\n",
    "    )\n",
    "\n",
    "    # Keep only toxicity labels columns and text column\n",
    "    dataset = dataset.remove_columns(\n",
    "        [\n",
    "            col for col in dataset[\"train\"].column_names if col not in labels + [\"text\"]\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    train_dataset = dataset[\"train\"].train_test_split(\n",
    "        test_size=test_size,\n",
    "        shuffle=True,\n",
    "        seed=seed\n",
    "    )\n",
    "\n",
    "    dataset[\"train\"] = train_dataset[\"train\"]\n",
    "    dataset[\"validation\"] = train_dataset[\"test\"]\n",
    "\n",
    "    return dataset\n",
    "\n",
    "\n",
    "dataset = prepare_dataset(\n",
    "    dataset,\n",
    "    test_size=params.validation_split,\n",
    "    seed=params.seed\n",
    ")\n",
    "\n",
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset.save_to_disk(\"data\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Uploading the data to S3\n",
    "\n",
    "We are going to use the `sagemaker.Session.upload_data` function to upload our datasets to an S3 location.\n",
    "\n",
    "The return value inputs identifies the location -- we will use later when we start the training job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input spec (in this case, just an S3 path): s3://sagemaker-us-east-1-215993976552/ToChiquinho/toxicity-type-detection/data\n"
     ]
    }
   ],
   "source": [
    "# inputs = sagemaker_session.upload_data(\n",
    "#     path=\"data\",\n",
    "#     bucket=bucket_name,\n",
    "#     key_prefix=f\"{prefix}/data\"\n",
    "# )\n",
    "\n",
    "inputs = \"s3://sagemaker-us-east-1-215993976552/ToChiquinho/toxicity-type-detection/data\"\n",
    "\n",
    "print(\"input spec (in this case, just an S3 path): {}\".format(inputs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "shutil.rmtree(\"data\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training session\n",
    "\n",
    "In this section, we will run the training process.\n",
    "\n",
    "To use Amazon SageMaker to run Docker containers, we need to provide a Python script for the container to run. In our case, all the code is in the `modeling` folder, including the `train.py` script.\n",
    "\n",
    "We will start doing a hyperparameter tuning process to find the best hyperparameters for our model.\n",
    "\n",
    "Then, we will train the model using the best hyperparameters found."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MLFlow run ID: 5bbfc5522d8040a3bc29113eb053f98a\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import mlflow\n",
    "\n",
    "mlflow.start_run()\n",
    "\n",
    "print(f\"MLFlow run ID: {mlflow.active_run().info.run_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the estimator\n",
    "\n",
    "We will use the `sagemaker.pytorch.PyTorch` class to define our estimator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.pytorch import PyTorch\n",
    "\n",
    "checkpoint_s3_uri = f\"s3://{bucket_name}/{prefix}/checkpoints\"\n",
    "\n",
    "# instance_type = \"ml.m5.4xlarge\" # 16 vCPUs, 64 GB RAM, no GPU - $ 0.922 per hour\n",
    "# instance_type = \"ml.m4.4xlarge\" # 16 vCPUs, 64 GB RAM, no GPU - $ 0.966 per hour\n",
    "# instance_type = \"ml.c4.8xlarge\" # 36 vCPUs, 60 GB RAM, no GPU - $ 1.909 per hour\n",
    "# instance_type = \"ml.c5.9xlarge\" # 36 vCPUs, 72 GB RAM, no GPU - $ 1.836 per hour\n",
    "\n",
    "instance_type = \"ml.g4dn.2xlarge\" # 8 vCPUs, 32 GB RAM, 1 x NVIDIA T4 GPU - $ 0.94 per hour\n",
    "# instance_type = \"ml.p3.2xlarge\" # 8 vCPUs, 61 GB RAM, 1 x NVIDIA Tesla V100 GPU - $ 3.825 per hour\n",
    "\n",
    "estimator = PyTorch(\n",
    "    entry_point=\"train.py\",\n",
    "    source_dir=\"modeling\",\n",
    "    role=params.sagemaker_execution_role_arn,\n",
    "    sagemaker_session=sagemaker_session,\n",
    "    py_version=\"py38\",\n",
    "    framework_version=\"1.12.0\",\n",
    "    instance_count=1,\n",
    "    instance_type=instance_type,\n",
    "    use_spot_instances=True,\n",
    "    max_wait=10800,\n",
    "    max_run=10800,\n",
    "    checkpoint_s3_uri=checkpoint_s3_uri,\n",
    "    checkpoint_local_path=\"/opt/ml/checkpoints\",\n",
    "    environment={\n",
    "        \"MLFLOW_TRACKING_URI\": params.mlflow_tracking_uri,\n",
    "        \"MLFLOW_EXPERIMENT_NAME\": params.mlflow_experiment_name,\n",
    "        \"MLFLOW_TRACKING_USERNAME\": params.mlflow_tracking_username,\n",
    "        \"MLFLOW_TRACKING_PASSWORD\": params.mlflow_tracking_password,\n",
    "        \"MLFLOW_TAGS\": params.mlflow_tags,\n",
    "        \"MLFLOW_RUN_ID\": mlflow.active_run().info.run_id,\n",
    "        \"MLFLOW_NESTED_RUN\": \"True\",\n",
    "        \"MLFLOW_FLATTEN_PARAMS\": \"True\",\n",
    "        \"WANDB_DISABLED\": \"True\"\n",
    "    },\n",
    "    hyperparameters={\n",
    "        ## If you want to test the code, uncomment the following lines to use smaller datasets\n",
    "        # \"max_train_samples\": 50,\n",
    "        # \"max_val_samples\": 50,\n",
    "        # \"max_test_samples\": 50,\n",
    "        \"num_train_epochs\": params.num_train_epochs_per_child,\n",
    "        \"batch_size\": params.batch_size,\n",
    "        \"seed\": params.seed\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To test our training job before hyperparameter tuning, we will run it with a small number of samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator.fit(inputs, wait=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hyperparameter Tuning\n",
    "\n",
    "We will use the `sagemaker.tuner.HyperparameterTuner` class to run a hyperparameter tuning process.\n",
    "\n",
    "We use MLflow to track the training process, so we can analyze the results through the MLflow UI."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Workaround for boto/boto3/issues/3488 issue\n",
    "\n",
    "Due to the issue [Estimator.environment not using in SageMaker.Client.create_hyper_parameter_tuning_job() · Issue #3488 · boto/boto3](https://github.com/boto/boto3/issues/3488), we need to include our environment variables in the `hyperparameters` parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "for k, v in estimator.environment.items():\n",
    "    if k != \"MLFLOW_TAGS\":\n",
    "        estimator._hyperparameters[k] = v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.tuner import (\n",
    "    IntegerParameter,\n",
    "    CategoricalParameter,\n",
    "    ContinuousParameter,\n",
    "    HyperparameterTuner,\n",
    ")\n",
    "\n",
    "tuner = HyperparameterTuner(\n",
    "    estimator,\n",
    "    max_jobs=18,\n",
    "    max_parallel_jobs=3,\n",
    "    objective_type=\"Maximize\",\n",
    "    objective_metric_name=\"test_f1\",\n",
    "    metric_definitions=[\n",
    "        {\n",
    "            \"Name\": \"test_f1\",\n",
    "            \"Regex\": \"Test F1-score: ([0-9\\\\.]+)\"\n",
    "        }\n",
    "    ],\n",
    "    hyperparameter_ranges={\n",
    "        \"learning_rate\": ContinuousParameter(1e-5, 1e-3, scaling_type=\"Logarithmic\"),\n",
    "        \"weight_decay\": ContinuousParameter(0.0, 0.1),\n",
    "        \"adam_beta1\": ContinuousParameter(0.8, 0.999),\n",
    "        \"adam_beta2\": ContinuousParameter(0.8, 0.999),\n",
    "        \"adam_epsilon\": ContinuousParameter(1e-8, 1e-6, scaling_type=\"Logarithmic\")\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SageMaker tuning job name: pytorch-training-221110-0849\n"
     ]
    }
   ],
   "source": [
    "tuner.fit(inputs, wait=False)\n",
    "\n",
    "params.sagemaker_tuning_job_name = tuner.latest_tuning_job.name\n",
    "\n",
    "print(f\"SageMaker tuning job name: {params.sagemaker_tuning_job_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlflow.log_params({\n",
    "    \"instance_type\": estimator.instance_type,\n",
    "    \"instance_count\": estimator.instance_count,\n",
    "    \"num_train_epochs_per_child\": params.num_train_epochs_per_child,\n",
    "    \"batch_size\": params.batch_size,\n",
    "    \"seed\": params.seed\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>TrainingJobName</th>\n",
       "      <th>FinalObjectiveValue</th>\n",
       "      <th>TrainingJobStatus</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>pytorch-training-221110-0849-011-b73293fa</td>\n",
       "      <td>0.761191</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>pytorch-training-221110-0849-012-94a0142f</td>\n",
       "      <td>0.750948</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>pytorch-training-221110-0849-013-c71b29a3</td>\n",
       "      <td>0.743556</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>pytorch-training-221110-0849-006-3955426c</td>\n",
       "      <td>0.737553</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>pytorch-training-221110-0849-004-a1003b73</td>\n",
       "      <td>0.694309</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>pytorch-training-221110-0849-010-39dcc340</td>\n",
       "      <td>0.694248</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>pytorch-training-221110-0849-003-2b58641f</td>\n",
       "      <td>0.694248</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>pytorch-training-221110-0849-002-404d00c0</td>\n",
       "      <td>0.694248</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>pytorch-training-221110-0849-001-286e660c</td>\n",
       "      <td>0.694248</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>pytorch-training-221110-0849-005-dfe707c1</td>\n",
       "      <td>0.672951</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>pytorch-training-221110-0849-009-e2f206a5</td>\n",
       "      <td>0.651342</td>\n",
       "      <td>Stopped</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>pytorch-training-221110-0849-008-81429e04</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Failed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>pytorch-training-221110-0849-007-00eb664d</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Stopped</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                              TrainingJobName  FinalObjectiveValue  \\\n",
       "2   pytorch-training-221110-0849-011-b73293fa             0.761191   \n",
       "1   pytorch-training-221110-0849-012-94a0142f             0.750948   \n",
       "0   pytorch-training-221110-0849-013-c71b29a3             0.743556   \n",
       "7   pytorch-training-221110-0849-006-3955426c             0.737553   \n",
       "9   pytorch-training-221110-0849-004-a1003b73             0.694309   \n",
       "3   pytorch-training-221110-0849-010-39dcc340             0.694248   \n",
       "10  pytorch-training-221110-0849-003-2b58641f             0.694248   \n",
       "11  pytorch-training-221110-0849-002-404d00c0             0.694248   \n",
       "12  pytorch-training-221110-0849-001-286e660c             0.694248   \n",
       "8   pytorch-training-221110-0849-005-dfe707c1             0.672951   \n",
       "4   pytorch-training-221110-0849-009-e2f206a5             0.651342   \n",
       "5   pytorch-training-221110-0849-008-81429e04                  NaN   \n",
       "6   pytorch-training-221110-0849-007-00eb664d                  NaN   \n",
       "\n",
       "   TrainingJobStatus  \n",
       "2          Completed  \n",
       "1          Completed  \n",
       "0          Completed  \n",
       "7          Completed  \n",
       "9          Completed  \n",
       "3          Completed  \n",
       "10         Completed  \n",
       "11         Completed  \n",
       "12         Completed  \n",
       "8          Completed  \n",
       "4            Stopped  \n",
       "5             Failed  \n",
       "6            Stopped  "
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "tuner_metrics: pd.DataFrame = sagemaker.HyperparameterTuningJobAnalytics(\n",
    "    hyperparameter_tuning_job_name=params.sagemaker_tuning_job_name,\n",
    "    sagemaker_session=sagemaker_session\n",
    ").dataframe()\n",
    "\n",
    "tuner_metrics.sort_values(\"FinalObjectiveValue\", ascending=False, inplace=True)\n",
    "tuner_metrics[[\"TrainingJobName\", \"FinalObjectiveValue\", \"TrainingJobStatus\"]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we can sort the results by the `FinalObjectiveValue` metric and see the best hyperparameters found."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'adam_beta1': 0.9057509584600282,\n",
       " 'adam_beta2': 0.9548132227438636,\n",
       " 'adam_epsilon': 1e-08,\n",
       " 'learning_rate': 1e-05,\n",
       " 'weight_decay': 0.027187699510002696,\n",
       " 'TrainingJobName': 'pytorch-training-221110-0849-011-b73293fa',\n",
       " 'TrainingJobStatus': 'Completed',\n",
       " 'FinalObjectiveValue': 0.7611912488937378,\n",
       " 'TrainingStartTime': Timestamp('2022-11-10 14:44:24-0300', tz='tzlocal()'),\n",
       " 'TrainingEndTime': Timestamp('2022-11-10 16:01:50-0300', tz='tzlocal()'),\n",
       " 'TrainingElapsedTimeSeconds': 4646.0}"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "best_job = tuner_metrics.iloc[0]\n",
    "best_job.to_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training model with best hyperparameters\n",
    "\n",
    "After analyzing the results of the hyperparameter tuning process, we can train the model using the best hyperparameters found."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No checkpoints found.\n"
     ]
    }
   ],
   "source": [
    "remove_checkpoints(\n",
    "    bucket_name=bucket_name,\n",
    "    checkpoint_prefix=f\"{prefix}/checkpoints\",\n",
    "    aws_profile_name=params.aws_profile_name\n",
    ")\n",
    "\n",
    "estimator = PyTorch(\n",
    "    entry_point=\"train.py\",\n",
    "    source_dir=\"modeling\",\n",
    "    role=params.sagemaker_execution_role_arn,\n",
    "    sagemaker_session=sagemaker_session,\n",
    "    py_version=\"py38\",\n",
    "    framework_version=\"1.12.0\",\n",
    "    instance_count=1,\n",
    "    instance_type=instance_type,\n",
    "    use_spot_instances=True,\n",
    "    max_wait=21600,\n",
    "    max_run=21600,\n",
    "    checkpoint_s3_uri=checkpoint_s3_uri,\n",
    "    checkpoint_local_path=\"/opt/ml/checkpoints\",\n",
    "    environment={\n",
    "        \"MLFLOW_TRACKING_URI\": params.mlflow_tracking_uri,\n",
    "        \"MLFLOW_EXPERIMENT_NAME\": params.mlflow_experiment_name,\n",
    "        \"MLFLOW_TRACKING_USERNAME\": params.mlflow_tracking_username,\n",
    "        \"MLFLOW_TRACKING_PASSWORD\": params.mlflow_tracking_password,\n",
    "        \"MLFLOW_TAGS\": params.mlflow_tags,\n",
    "        \"MLFLOW_RUN_ID\": mlflow.active_run().info.run_id,\n",
    "        \"MLFLOW_NESTED_RUN\": \"True\",\n",
    "        \"MLFLOW_FLATTEN_PARAMS\": \"True\",\n",
    "        \"HF_MLFLOW_LOG_ARTIFACTS\": \"True\",\n",
    "        \"WANDB_DISABLED\": \"True\"\n",
    "    },\n",
    "    hyperparameters={\n",
    "        \"num_train_epochs\": params.num_train_epochs,\n",
    "        \"batch_size\": params.batch_size,\n",
    "        \"adam_beta1\": best_job[\"adam_beta1\"],\n",
    "        \"adam_beta2\": best_job[\"adam_beta2\"],\n",
    "        \"adam_epsilon\": best_job[\"adam_epsilon\"],\n",
    "        \"learning_rate\": best_job[\"learning_rate\"],\n",
    "        \"weight_decay\": best_job[\"weight_decay\"],\n",
    "        \"seed\": params.seed\n",
    "    },\n",
    ")\n",
    "\n",
    "estimator.fit(inputs, wait=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlflow.log_params(\n",
    "    {\n",
    "        \"best_adam_beta1\": best_job[\"adam_beta1\"],\n",
    "        \"best_adam_beta2\": best_job[\"adam_beta2\"],\n",
    "        \"best_adam_epsilon\": best_job[\"adam_epsilon\"],\n",
    "        \"best_learning_rate\": best_job[\"learning_rate\"],\n",
    "        \"best_weight_decay\": best_job[\"weight_decay\"]\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlflow.end_run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Documentation\n",
    "\n",
    "- [Estimators — sagemaker documentation](https://sagemaker.readthedocs.io/en/stable/api/training/estimators.html)\n",
    "- [HyperparameterTuner — sagemaker documentation](https://sagemaker.readthedocs.io/en/stable/api/training/tuner.html)\n",
    "- [Configure and Launch a Hyperparameter Tuning Job - Amazon SageMaker](https://docs.aws.amazon.com/sagemaker/latest/dg/automatic-model-tuning-ex-tuning-job.html)\n",
    "- [Managed Spot Training in Amazon SageMaker - Amazon SageMaker](https://docs.aws.amazon.com/sagemaker/latest/dg/model-managed-spot-training.html)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.8 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "369f2c481f4da34e4445cda3fffd2e751bd1c4d706f27375911949ba6bb62e1c"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}