src/toxicity_type_detection.ipynb

{
 "cells": [
  {
   "attachments": {},
   "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",
    "- [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 best model](#Training-best-model)\n",
    "- [Documentation](#Documentation)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "In this section, we will import the necessary libraries and set up the environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dotenv import load_dotenv\n",
    "load_dotenv(\"../.env\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "NotebookArguments(num_train_epochs=30, early_stopping_patience=2, 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\"}', sagemaker_image_uri='215993976552.dkr.ecr.us-east-1.amazonaws.com/sagemaker-transformers:1.12.0-gpu-py38')"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "from ml.arguments import NotebookArguments\n",
    "from ml.utils import remove_checkpoints\n",
    "\n",
    "params = NotebookArguments(\n",
    "    mlflow_experiment_name=\"toxicity-type-detection\",\n",
    "    mlflow_tags={\n",
    "        \"project\": \"ToChiquinho\",\n",
    "        \"dataset\": \"OLID-BR\",\n",
    "        \"model_type\": \"bert\",\n",
    "        \"problem_type\": \"multi_label_classification\"\n",
    "    },\n",
    "    sagemaker_execution_role_arn=os.environ.get(\"SAGEMAKER_EXECUTION_ROLE_ARN\"),\n",
    "    sagemaker_tuning_job_name = \"pytorch-training-221118-2303\",\n",
    "    aws_profile_name=os.environ.get(\"AWS_PROFILE\")\n",
    ")\n",
    "\n",
    "params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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 = f\"ToChiquinho/{params.mlflow_experiment_name}\"\n",
    "\n",
    "if params.sagemaker_execution_role_arn is None:\n",
    "    params.sagemaker_execution_role_arn = sagemaker.get_execution_role(sagemaker_session)"
   ]
  },
  {
   "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": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration dougtrajano--olid-br-f83aad8215e23434\n",
      "Found cached dataset parquet (C:/Users/trajano/.cache/huggingface/datasets/dougtrajano___parquet/dougtrajano--olid-br-f83aad8215e23434/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ad09a62a00c840a7bec6c286166af3f6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\"dougtrajano/olid-br\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at C:\\Users\\trajano\\.cache\\huggingface\\datasets\\dougtrajano___parquet\\dougtrajano--olid-br-f83aad8215e23434\\0.0.0\\2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec\\cache-33e930f77f9076e5.arrow\n",
      "Loading cached processed dataset at C:\\Users\\trajano\\.cache\\huggingface\\datasets\\dougtrajano___parquet\\dougtrajano--olid-br-f83aad8215e23434\\0.0.0\\2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec\\cache-2aedd0dd5e9a3144.arrow\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f6026236696e4573831b914e6e40f936",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cc6c75ad0fd249b5837d8b0fe9d0cfce",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    test: Dataset({\n",
       "        features: ['text', 'health', 'ideology', 'insult', 'lgbtqphobia', 'other_lifestyle', 'physical_aspects', 'profanity_obscene', 'racism', 'sexism', 'xenophobia'],\n",
       "        num_rows: 1438\n",
       "    })\n",
       "    train: Dataset({\n",
       "        features: ['text', 'health', 'ideology', 'insult', 'lgbtqphobia', 'other_lifestyle', 'physical_aspects', 'profanity_obscene', 'racism', 'sexism', 'xenophobia'],\n",
       "        num_rows: 3417\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['text', 'health', 'ideology', 'insult', 'lgbtqphobia', 'other_lifestyle', 'physical_aspects', 'profanity_obscene', 'racism', 'sexism', 'xenophobia'],\n",
       "        num_rows: 855\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dd4f02ee54ea4d0f895304f7e0f762e4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Flattening the indices:   0%|          | 0/2 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "61c926e379cd413f89710cb371c6310c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Flattening the indices:   0%|          | 0/4 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6dc210ad216c49a5a2dea736a929ea0f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Flattening the indices:   0%|          | 0/1 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "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": 4,
   "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": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "shutil.rmtree(\"data\")"
   ]
  },
  {
   "attachments": {},
   "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 `ml` 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": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MLFlow run ID: 481954e8af32432e83b2af1151d781a0\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}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the estimator\n",
    "\n",
    "We will use the `sagemaker.huggingface.HuggingFace` class to define our estimator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "from sagemaker.huggingface import HuggingFace\n",
    "\n",
    "checkpoint_s3_uri = f\"s3://{bucket_name}/{prefix}/checkpoints\"\n",
    "\n",
    "instance_type = \"ml.g4dn.xlarge\" # 4 vCPUs, 16 GB RAM, 1 x NVIDIA T4 16GB GPU - $ 0.736 per hour\n",
    "# instance_type = \"ml.g4dn.2xlarge\" # 8 vCPUs, 32 GB RAM, 1 x NVIDIA T4 16GB GPU - $ 0.94 per hour\n",
    "# instance_type = \"ml.g5.xlarge\" # 4 vCPUs, 16 GB RAM, 1 x NVIDIA A10G 24GB GPU - $ 1.408 per hour\n",
    "# instance_type = \"ml.g5.2xlarge\" # 8 vCPUs, 32 GB RAM, 1 x NVIDIA A10G 24GB GPU - $ 1.515 per hour\n",
    "# instance_type = \"ml.g5.4xlarge\" # 16 vCPUs, 64 GB RAM, 2 x NVIDIA A10G 24GB GPU - $ 2.03 per hour\n",
    "# instance_type = \"ml.g5.8xlarge\" # 32 vCPUs, 128 GB RAM, 4 x NVIDIA A10G 24GB GPU - $ 3.06 per hour\n",
    "\n",
    "estimator = HuggingFace(\n",
    "    entry_point=\"train.py\",\n",
    "    source_dir=\"ml\",\n",
    "    base_job_name=params.mlflow_experiment_name,\n",
    "    container_log_level=logging.DEBUG,\n",
    "    role=params.sagemaker_execution_role_arn,\n",
    "    sagemaker_session=sagemaker_session,\n",
    "    py_version=\"py38\",\n",
    "    pytorch_version=\"1.10.2\",\n",
    "    transformers_version=\"4.17.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_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\": 100,\n",
    "        # \"max_val_samples\": 100,\n",
    "        # \"max_test_samples\": 100,\n",
    "        \"num_train_epochs\": params.num_train_epochs,\n",
    "        \"early_stopping_patience\": params.early_stopping_patience,\n",
    "        \"batch_size\": params.batch_size,\n",
    "        \"eval_dataset\": \"validation\",\n",
    "        \"seed\": params.seed\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlflow.log_params(\n",
    "    {\n",
    "        \"instance_type\": estimator.instance_type,\n",
    "        \"instance_count\": estimator.instance_count,\n",
    "        \"early_stopping_patience\": params.early_stopping_patience\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": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:sagemaker.image_uris:image_uri is not presented, retrieving image_uri based on instance_type, framework etc.\n",
      "INFO:sagemaker:Creating training-job with name: pytorch-training-2022-12-29-14-41-19-455\n"
     ]
    }
   ],
   "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": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.tuner import (\n",
    "    IntegerParameter,\n",
    "    CategoricalParameter,\n",
    "    ContinuousParameter,\n",
    "    HyperparameterTuner,\n",
    ")\n",
    "\n",
    "estimator._hyperparameters.pop(\"max_train_samples\", None)\n",
    "estimator._hyperparameters.pop(\"max_val_samples\", None)\n",
    "estimator._hyperparameters.pop(\"max_test_samples\", None)\n",
    "\n",
    "tuner = HyperparameterTuner(\n",
    "    estimator,\n",
    "    max_jobs=18,\n",
    "    max_parallel_jobs=3,\n",
    "    objective_type=\"Maximize\",\n",
    "    objective_metric_name=\"eval_f1\",\n",
    "    metric_definitions=[\n",
    "        {\n",
    "            \"Name\": \"eval_f1\",\n",
    "            \"Regex\": \"eval_f1_weighted: ([0-9\\\\.]+)\"\n",
    "        }\n",
    "    ],\n",
    "    hyperparameter_ranges={\n",
    "        \"learning_rate\": ContinuousParameter(1e-5, 1e-3),\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),\n",
    "        \"optim\": CategoricalParameter(\n",
    "            [\n",
    "                \"adamw_hf\",\n",
    "                \"adamw_torch\",\n",
    "                \"adamw_apex_fused\",\n",
    "                \"adafactor\"\n",
    "            ]\n",
    "        )\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tuner.fit(inputs, wait=False)\n",
    "\n",
    "params.sagemaker_tuning_job_name = tuner.latest_tuning_job.name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SageMaker tuning job name: pytorch-training-221118-2303\n"
     ]
    }
   ],
   "source": [
    "print(f\"SageMaker tuning job name: {params.sagemaker_tuning_job_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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>17</th>\n",
       "      <td>pytorch-training-221118-2303-001-2c25be0c</td>\n",
       "      <td>0.759114</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>pytorch-training-221118-2303-006-49d87b23</td>\n",
       "      <td>0.758660</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>pytorch-training-221118-2303-007-1182e099</td>\n",
       "      <td>0.725817</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>pytorch-training-221118-2303-008-10c5fe8b</td>\n",
       "      <td>0.719377</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>pytorch-training-221118-2303-005-f1adc05a</td>\n",
       "      <td>0.718727</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>pytorch-training-221118-2303-004-1fdab0a9</td>\n",
       "      <td>0.714234</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>pytorch-training-221118-2303-002-c40b0d07</td>\n",
       "      <td>0.709980</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>pytorch-training-221118-2303-009-5743f9ce</td>\n",
       "      <td>0.709527</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>pytorch-training-221118-2303-012-c98d6e1c</td>\n",
       "      <td>0.691050</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>pytorch-training-221118-2303-011-cd4e57b8</td>\n",
       "      <td>0.675431</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>pytorch-training-221118-2303-003-ee3d2e41</td>\n",
       "      <td>0.675431</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>pytorch-training-221118-2303-014-677d0d92</td>\n",
       "      <td>0.671177</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>pytorch-training-221118-2303-015-f432d514</td>\n",
       "      <td>0.671177</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>pytorch-training-221118-2303-018-c70bac75</td>\n",
       "      <td>0.665301</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>pytorch-training-221118-2303-010-7c1a0ed4</td>\n",
       "      <td>0.546034</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>pytorch-training-221118-2303-017-ec8b49e7</td>\n",
       "      <td>0.544258</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>pytorch-training-221118-2303-013-8440bbfd</td>\n",
       "      <td>0.544258</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>pytorch-training-221118-2303-016-f1f0f0a2</td>\n",
       "      <td>0.544258</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                              TrainingJobName  FinalObjectiveValue  \\\n",
       "17  pytorch-training-221118-2303-001-2c25be0c             0.759114   \n",
       "12  pytorch-training-221118-2303-006-49d87b23             0.758660   \n",
       "11  pytorch-training-221118-2303-007-1182e099             0.725817   \n",
       "10  pytorch-training-221118-2303-008-10c5fe8b             0.719377   \n",
       "13  pytorch-training-221118-2303-005-f1adc05a             0.718727   \n",
       "14  pytorch-training-221118-2303-004-1fdab0a9             0.714234   \n",
       "16  pytorch-training-221118-2303-002-c40b0d07             0.709980   \n",
       "9   pytorch-training-221118-2303-009-5743f9ce             0.709527   \n",
       "6   pytorch-training-221118-2303-012-c98d6e1c             0.691050   \n",
       "7   pytorch-training-221118-2303-011-cd4e57b8             0.675431   \n",
       "15  pytorch-training-221118-2303-003-ee3d2e41             0.675431   \n",
       "4   pytorch-training-221118-2303-014-677d0d92             0.671177   \n",
       "3   pytorch-training-221118-2303-015-f432d514             0.671177   \n",
       "0   pytorch-training-221118-2303-018-c70bac75             0.665301   \n",
       "8   pytorch-training-221118-2303-010-7c1a0ed4             0.546034   \n",
       "1   pytorch-training-221118-2303-017-ec8b49e7             0.544258   \n",
       "5   pytorch-training-221118-2303-013-8440bbfd             0.544258   \n",
       "2   pytorch-training-221118-2303-016-f1f0f0a2             0.544258   \n",
       "\n",
       "   TrainingJobStatus  \n",
       "17         Completed  \n",
       "12         Completed  \n",
       "11         Completed  \n",
       "10         Completed  \n",
       "13         Completed  \n",
       "14         Completed  \n",
       "16         Completed  \n",
       "9          Completed  \n",
       "6          Completed  \n",
       "7          Completed  \n",
       "15         Completed  \n",
       "4          Completed  \n",
       "3          Completed  \n",
       "0          Completed  \n",
       "8          Completed  \n",
       "1          Completed  \n",
       "5          Completed  \n",
       "2          Completed  "
      ]
     },
     "execution_count": 9,
     "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[\"optim\"] = tuner_metrics[\"optim\"].apply(lambda x: x.strip('\"'))\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": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'adam_beta1': 0.9339215524915885,\n",
       " 'adam_beta2': 0.9916979096990963,\n",
       " 'adam_epsilon': 3.4435900142455904e-07,\n",
       " 'learning_rate': 7.044186985160909e-05,\n",
       " 'optim': 'adamw_apex_fused',\n",
       " 'weight_decay': 0.02426675806866223,\n",
       " 'TrainingJobName': 'pytorch-training-221118-2303-001-2c25be0c',\n",
       " 'TrainingJobStatus': 'Completed',\n",
       " 'FinalObjectiveValue': 0.7591144442558289,\n",
       " 'TrainingStartTime': Timestamp('2022-11-18 23:04:37-0300', tz='tzlocal()'),\n",
       " 'TrainingEndTime': Timestamp('2022-11-18 23:51:48-0300', tz='tzlocal()'),\n",
       " 'TrainingElapsedTimeSeconds': 2831.0}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "best_job = tuner_metrics.iloc[0]\n",
    "best_job.to_dict()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training best model\n",
    "\n",
    "We will train the model using the best hyperparameters found.\n",
    "\n",
    "We will concatenate the training and validation datasets to train the model with more data and evaluate it using the test dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:sagemaker.image_uris:image_uri is not presented, retrieving image_uri based on instance_type, framework etc.\n",
      "INFO:sagemaker:Creating training-job with name: toxicity-type-detection-2023-02-26-21-42-23-182\n"
     ]
    }
   ],
   "source": [
    "estimator.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_FLATTEN_PARAMS\": \"True\",\n",
    "    \"HF_MLFLOW_LOG_ARTIFACTS\": \"True\",\n",
    "    \"HUGGINGFACE_HUB_TOKEN\": params.huggingface_hub_token\n",
    "}\n",
    "\n",
    "estimator._hyperparameters = {\n",
    "    \"push_to_hub\": \"True\",\n",
    "    \"hub_model_id\": f\"dougtrajano/{params.mlflow_experiment_name}\",\n",
    "    \"num_train_epochs\": params.num_train_epochs,\n",
    "    \"early_stopping_patience\": params.early_stopping_patience,\n",
    "    \"batch_size\": params.batch_size,\n",
    "    \"seed\": params.seed,\n",
    "    \"concat_validation_set\": \"True\",\n",
    "    \"eval_dataset\": \"test\",\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",
    "    \"optim\": best_job[\"optim\"]\n",
    "}\n",
    "\n",
    "estimator.fit(inputs, wait=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:botocore.credentials:Found credentials in shared credentials file: ~/.aws/credentials\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Deleted 85 checkpoints from sagemaker-us-east-1-215993976552/ToChiquinho/toxicity-type-detection/checkpoints.\n"
     ]
    }
   ],
   "source": [
    "remove_checkpoints(\n",
    "    bucket_name=bucket_name,\n",
    "    checkpoint_prefix=f\"{prefix}/checkpoints\",\n",
    "    aws_profile_name=params.aws_profile_name\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
}