src/toxic_spans_detection.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Toxic Spans 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 toxic spans in a Portuguese text. We used the [OLID-BR](https://dougtrajano.github.io/olid-br/) dataset as the training data.\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=5, batch_size=8, validation_split=0.2, seed=1993, mlflow_experiment_name='toxic-spans-detection', mlflow_tags='{\"project\": \"ToChiquinho\", \"dataset\": \"OLID-BR\", \"model_type\": \"bert\", \"problem_type\": \"span_categorization\"}')"
      ]
     },
     "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",
    "    num_train_epochs=30,\n",
    "    early_stopping_patience=5,\n",
    "    mlflow_tracking_uri=os.environ.get(\"MLFLOW_TRACKING_URI\"),\n",
    "    mlflow_experiment_name=\"toxic-spans-detection\",\n",
    "    mlflow_tracking_username=os.environ.get(\"MLFLOW_TRACKING_USERNAME\"),\n",
    "    mlflow_tracking_password=os.environ.get(\"MLFLOW_TRACKING_PASSWORD\"),\n",
    "    mlflow_tags={\n",
    "        \"project\": \"ToChiquinho\",\n",
    "        \"dataset\": \"OLID-BR\",\n",
    "        \"model_type\": \"bert\",\n",
    "        \"problem_type\": \"span_categorization\"\n",
    "    },\n",
    "    sagemaker_execution_role_arn=os.environ.get(\"SAGEMAKER_EXECUTION_ROLE_ARN\"),\n",
    "    sagemaker_tuning_job_name = \"pytorch-training-221218-1712\",\n",
    "    aws_profile_name=os.environ.get(\"AWS_PROFILE\")\n",
    ")\n",
    "\n",
    "params"
   ]
  },
  {
   "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": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:datasets.builder:Using custom data configuration dougtrajano--olid-br-f83aad8215e23434\n",
      "WARNING:datasets.builder: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": "da3793b9acec4744a2acc33211901456",
       "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": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:datasets.fingerprint:Parameter 'function'=<function prepare_dataset.<locals>.<lambda> at 0x000001E670255240> of the transform datasets.arrow_dataset.Dataset.filter@2.0.1 couldn't be hashed properly, a random hash was used instead. Make sure your transforms and parameters are serializable with pickle or dill for the dataset fingerprinting and caching to work. If you reuse this transform, the caching mechanism will consider it to be different from the previous calls and recompute everything. This warning is only showed once. Subsequent hashing failures won't be showed.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "25d14162df1e47558517ef7491f0a526",
       "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": "44ab60d16dda4c61b5f9a3d36635b0d0",
       "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', 'toxic_spans'],\n",
       "        num_rows: 1438\n",
       "    })\n",
       "    train: Dataset({\n",
       "        features: ['text', 'toxic_spans'],\n",
       "        num_rows: 3433\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['text', 'toxic_spans'],\n",
       "        num_rows: 859\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import ast\n",
    "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 offensive examples with toxic spans\n",
    "    dataset = dataset.filter(\n",
    "        lambda example: example[\"is_offensive\"] == \"OFF\"\n",
    "        and example.get(\"toxic_spans\") is not None\n",
    "    )\n",
    "    \n",
    "    # Filter only \"text\" and \"toxic_spans\" columns\n",
    "    dataset = dataset.remove_columns(\n",
    "        [\n",
    "            col for col in dataset[\"train\"].column_names if col not in [\"text\", \"toxic_spans\"]\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": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5f998dbf568045d9948896e76cbb9e47",
       "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": "c0b8689aca1d4aeeb43c68d039a7eae5",
       "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": "76c7aa394eef484e8e179b263cbe17aa",
       "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/toxic-spans-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/toxic-spans-detection/data\"\n",
    "\n",
    "print(\"input spec (in this case, just an S3 path): {}\".format(inputs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "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: d37904ee5cad40e1bc2a6159e951a98a\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",
    "    },\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",
    "        \"model_name\": \"pt_core_news_lg\",\n",
    "        \"learning_rate\": 0.001,\n",
    "        \"num_train_epochs\": params.num_train_epochs,\n",
    "        \"early_stopping_patience\": params.early_stopping_patience,\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": 8,
   "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-19-32-02-356\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": 10,
   "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: ([0-9\\\\.]+)\"\n",
    "        }\n",
    "    ],\n",
    "    hyperparameter_ranges={\n",
    "        \"learning_rate\": ContinuousParameter(1e-4, 1e-2),\n",
    "        \"dropout\": CategoricalParameter([0.0, 0.1, 0.2, 0.3, 0.4, 0.5]),\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([\"adam\", \"radam\"])\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": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SageMaker tuning job name: pytorch-training-221218-1712\n"
     ]
    }
   ],
   "source": [
    "print(f\"SageMaker tuning job name: {params.sagemaker_tuning_job_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "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>8</th>\n",
       "      <td>pytorch-training-221218-1712-010-e1b19c80</td>\n",
       "      <td>0.570670</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>pytorch-training-221218-1712-005-96b28f81</td>\n",
       "      <td>0.567621</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>pytorch-training-221218-1712-009-8c6afd13</td>\n",
       "      <td>0.544218</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>pytorch-training-221218-1712-002-92a59f2e</td>\n",
       "      <td>0.539523</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>pytorch-training-221218-1712-007-29bfbac0</td>\n",
       "      <td>0.513867</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>pytorch-training-221218-1712-003-f23e8359</td>\n",
       "      <td>0.505005</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>pytorch-training-221218-1712-015-4a7eb2fb</td>\n",
       "      <td>0.491720</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>pytorch-training-221218-1712-008-5d6d9bb0</td>\n",
       "      <td>0.473096</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>pytorch-training-221218-1712-017-838075ab</td>\n",
       "      <td>0.469928</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>pytorch-training-221218-1712-016-bc1a09a1</td>\n",
       "      <td>0.461928</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>pytorch-training-221218-1712-001-7d88e9c0</td>\n",
       "      <td>0.446936</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>pytorch-training-221218-1712-014-6e1da9c1</td>\n",
       "      <td>0.417162</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>pytorch-training-221218-1712-011-872d0c31</td>\n",
       "      <td>0.088475</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>pytorch-training-221218-1712-006-973d27e2</td>\n",
       "      <td>0.088475</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>pytorch-training-221218-1712-012-93cd5dcc</td>\n",
       "      <td>0.088475</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>pytorch-training-221218-1712-004-6ab20369</td>\n",
       "      <td>0.088475</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>pytorch-training-221218-1712-013-97b44e95</td>\n",
       "      <td>0.088475</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>pytorch-training-221218-1712-018-279689bd</td>\n",
       "      <td>0.088475</td>\n",
       "      <td>Completed</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                              TrainingJobName  FinalObjectiveValue  \\\n",
       "8   pytorch-training-221218-1712-010-e1b19c80             0.570670   \n",
       "13  pytorch-training-221218-1712-005-96b28f81             0.567621   \n",
       "9   pytorch-training-221218-1712-009-8c6afd13             0.544218   \n",
       "16  pytorch-training-221218-1712-002-92a59f2e             0.539523   \n",
       "11  pytorch-training-221218-1712-007-29bfbac0             0.513867   \n",
       "15  pytorch-training-221218-1712-003-f23e8359             0.505005   \n",
       "3   pytorch-training-221218-1712-015-4a7eb2fb             0.491720   \n",
       "10  pytorch-training-221218-1712-008-5d6d9bb0             0.473096   \n",
       "1   pytorch-training-221218-1712-017-838075ab             0.469928   \n",
       "2   pytorch-training-221218-1712-016-bc1a09a1             0.461928   \n",
       "17  pytorch-training-221218-1712-001-7d88e9c0             0.446936   \n",
       "4   pytorch-training-221218-1712-014-6e1da9c1             0.417162   \n",
       "7   pytorch-training-221218-1712-011-872d0c31             0.088475   \n",
       "12  pytorch-training-221218-1712-006-973d27e2             0.088475   \n",
       "6   pytorch-training-221218-1712-012-93cd5dcc             0.088475   \n",
       "14  pytorch-training-221218-1712-004-6ab20369             0.088475   \n",
       "5   pytorch-training-221218-1712-013-97b44e95             0.088475   \n",
       "0   pytorch-training-221218-1712-018-279689bd             0.088475   \n",
       "\n",
       "   TrainingJobStatus  \n",
       "8          Completed  \n",
       "13         Completed  \n",
       "9          Completed  \n",
       "16         Completed  \n",
       "11         Completed  \n",
       "15         Completed  \n",
       "3          Completed  \n",
       "10         Completed  \n",
       "1          Completed  \n",
       "2          Completed  \n",
       "17         Completed  \n",
       "4          Completed  \n",
       "7          Completed  \n",
       "12         Completed  \n",
       "6          Completed  \n",
       "14         Completed  \n",
       "5          Completed  \n",
       "0          Completed  "
      ]
     },
     "execution_count": 19,
     "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[\"dropout\"] = tuner_metrics[\"dropout\"].apply(lambda x: float(x.strip('\"')))\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": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'adam_beta1': 0.9978242993498763,\n",
       " 'adam_beta2': 0.9988901284249041,\n",
       " 'adam_epsilon': 3.12576102525027e-08,\n",
       " 'dropout': 0.3,\n",
       " 'learning_rate': 0.00038798590315954165,\n",
       " 'optim': 'adam',\n",
       " 'weight_decay': 0.1,\n",
       " 'TrainingJobName': 'pytorch-training-221218-1712-010-e1b19c80',\n",
       " 'TrainingJobStatus': 'Completed',\n",
       " 'FinalObjectiveValue': 0.5706698298454285,\n",
       " 'TrainingStartTime': Timestamp('2022-12-18 18:03:54-0300', tz='tzlocal()'),\n",
       " 'TrainingEndTime': Timestamp('2022-12-18 18:24:26-0300', tz='tzlocal()'),\n",
       " 'TrainingElapsedTimeSeconds': 1232.0}"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "best_job = tuner_metrics.iloc[0].copy()\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": 21,
   "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-19-56-31-774\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",
    "}\n",
    "\n",
    "estimator._hyperparameters = {\n",
    "    \"model_name\": \"pt_core_news_lg\",\n",
    "    \"adam_beta1\": best_job[\"adam_beta1\"],\n",
    "    \"adam_beta2\": best_job[\"adam_beta2\"],\n",
    "    \"adam_epsilon\": best_job[\"adam_epsilon\"],\n",
    "    \"dropout\": best_job[\"dropout\"],\n",
    "    \"optim\": best_job[\"optim\"],\n",
    "    \"weight_decay\": best_job[\"weight_decay\"],\n",
    "    \"learning_rate\": best_job[\"learning_rate\"],\n",
    "    \"num_train_epochs\": params.num_train_epochs,\n",
    "    \"early_stopping_patience\": params.early_stopping_patience,\n",
    "    \"concat_validation_set\": \"True\",\n",
    "    \"eval_dataset\": \"test\",\n",
    "    \"seed\": params.seed\n",
    "}\n",
    "\n",
    "estimator.fit(inputs, wait=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "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 181 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",
    ")\n",
    "\n",
    "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
}