2_training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "4CspUKbHyK7A"
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pmtu1rHfxusY"
   },
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "oZfijJ2kxusc"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import datasets\n",
    "import pandas as pd\n",
    "\n",
    "from src.dataset import add_audio_column, filter_df, prepare_ds, split_df\n",
    "from src.train import end_training, get_model, get_trainer\n",
    "from src.utils import get_csv_name, get_run_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "mFZtkDipxusg"
   },
   "outputs": [],
   "source": [
    "RES_DIR_PATH = \"res\"\n",
    "NOTEBOOK_ENV = \"jupyter\"\n",
    "\n",
    "AUDIOS_DIR_PATH = os.path.join(RES_DIR_PATH, \"mp3_data\")\n",
    "MODELS_DIR_PATH = os.path.join(RES_DIR_PATH, \"models\")\n",
    "DATASETS_DIR_PATH = os.path.join(RES_DIR_PATH, \"datasets\")\n",
    "\n",
    "CSV_PATH = os.path.join(RES_DIR_PATH, \"samples_clustered.csv\")\n",
    "\n",
    "TOP_N_GENRES = 6\n",
    "TOP_N_FEATURES = 9\n",
    "\n",
    "FEATURES_CONFIG_SUBSET = {\"genre\": {\"top_n\": 3, \"samples\": 1000}}\n",
    "FEATURES_CONFIG_GEN = {\"genre\": {\"top_n\": TOP_N_GENRES, \"samples\": None}}\n",
    "FEATURES_CONFIG_CAT = {\"category\": {\"top_n\": TOP_N_FEATURES, \"samples\": None}}\n",
    "\n",
    "VALID_SIZE = 0.1\n",
    "TEST_SIZE = 0.1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAINING_CONFIG = {\n",
    "    \"epochs\": 20,\n",
    "    \"learning_rate\": 5e-5,\n",
    "    \"warmup\": 0.0,\n",
    "    \"train_batch_size\": 8,\n",
    "    \"eval_batch_size\": 16,\n",
    "    \"feature_encoder\": None,\n",
    "    \"freeze_encoder\": None,\n",
    "    \"classifier_layers\": None, \n",
    "    \"classifier_dropout\": None,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_or_load_df(features_config):\n",
    "    filtered_csv_path = get_csv_name(features_config, CSV_PATH)\n",
    "\n",
    "    # If the subset is already in the filesystem, load it directly\n",
    "    if os.path.exists(filtered_csv_path):\n",
    "        print(f\"Loading {filtered_csv_path}\")\n",
    "        df = pd.read_csv(filtered_csv_path)\n",
    "    else:\n",
    "        df = pd.read_csv(CSV_PATH)\n",
    "        # Filter the dataset according to the given configuration and remove rows containing null values\n",
    "        df = filter_df(\n",
    "            df, \n",
    "            remove_nones=True,\n",
    "            features_config=features_config, \n",
    "        )\n",
    "        df.to_csv(filtered_csv_path, index=False)\n",
    "\n",
    "    print(f\"{len(df)} examples in DataFrame\")\n",
    "    # If the split column is not in the dataset, split the dataset into three partisions using \n",
    "    # `TEST_SIZE` and `VALID_SIZE` and save the result\n",
    "\n",
    "    if \"split\" not in df.columns:\n",
    "        df = split_df(df, validation_size=VALID_SIZE, test_size=TEST_SIZE)\n",
    "        df.to_csv(filtered_csv_path, index=False)\n",
    "\n",
    "    print(df.value_counts(\"split\"))\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a function for loading the dataset for the requested model\n",
    "\n",
    "def load_and_prepare_ds(training_config, feature_config, df, clustered=True):\n",
    "    encoded_dataset_path = os.path.join(DATASETS_DIR_PATH, f\"ds-{training_config['feature_encoder']}-full-encoded\")\n",
    "    ds = datasets.load_from_disk(encoded_dataset_path)\n",
    "    ds = add_audio_column(ds, audios_dir_path=AUDIOS_DIR_PATH, training_config={\"feature_encoder\": training_config['feature_encoder']})\n",
    "    return prepare_ds(ds, df, feature_config, clustered=clustered, fixed_mapping=None, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training Variations\n",
    "- Backbones: [Wav2Vec2](https://arxiv.org/abs/2006.11477) and [Whisper](https://cdn.openai.com/papers/whisper.pdf)\n",
    "- Fine-tunining: When disabled, the gradient computatoin in the whole pretrained encoder is disabled\n",
    "- Classification head: Number of layers and hidden dimensions. Each layer is followed by a ReLU activation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "siGNvZU5xusi"
   },
   "source": [
    "## Backbone Implementation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Both models are based on their [Hugging Face Transformers](https://huggingface.co/docs/transformers) implementation.\n",
    "\n",
    "The **Wav2Vec2** classifier extends [Wav2Vec2ForSequenceClassification](https://huggingface.co/docs/transformers/model_doc/wav2vec2#transformers.Wav2Vec2ForSequenceClassification) class, adding the support for a custom classification head.\n",
    "\n",
    "![Wav2Vec2 Architecture](res/report/wav2vec2.png \"Wav2Vec2 Architecture\")\n",
    "\n",
    "For **Whisper**, a classifier implementation didn't exists, so I used the internal [WhisperEncoder](https://huggingface.co/docs/transformers/model_doc/whisper) class and used the latent representation as features for classification.\n",
    "\n",
    "![Whisper Architecture](res/report/whisper.png \"Whisper Architecture\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PxFCxU7Oxusw"
   },
   "source": [
    "# Training (Subset Dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to try the different settings of configuration, I trained the models on genre classification using a subset of the whole dataset, containing 1000 audio files and just the top-3 frequent genres:\n",
    "- World/Ethnic\n",
    "- Rock/Blues\n",
    "- Electronic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading res/samples_clustered_genre3s1000.csv\n",
      "999 examples in DataFrame\n",
      "split\n",
      "train    799\n",
      "test     100\n",
      "valid    100\n",
      "dtype: int64\n"
     ]
    }
   ],
   "source": [
    "# Build the filename indicating the subset of the whole dataset with the specific configurations\n",
    "df = create_or_load_df(FEATURES_CONFIG_SUBSET)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7MWfQIuVxus0"
   },
   "source": [
    "## Wav2Vec2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "wSFCEv3T4Q-m"
   },
   "outputs": [],
   "source": [
    "TRAINING_CONFIG[\"feature_encoder\"] = \"wav2vec2\"\n",
    "TRAINING_CONFIG[\"freeze_encoder\"] = True\n",
    "TRAINING_CONFIG[\"classifier_layers\"] = [256]\n",
    "TRAINING_CONFIG[\"classifier_dropout\"] = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prepared_ds = load_and_prepare_ds(TRAINING_CONFIG, FEATURES_CONFIG_SUBSET, df)\n",
    "\n",
    "prepared_ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "m-AHRUIK2KPL",
    "outputId": "34b6f338-882c-4ebb-8ce0-b4508b66d254"
   },
   "outputs": [],
   "source": [
    "run_name = get_run_name(TRAINING_CONFIG)\n",
    "model = get_model(TRAINING_CONFIG, prepared_ds[\"train\"])\n",
    "\n",
    "trainer = get_trainer(\n",
    "    run_name=run_name,\n",
    "    model=model,\n",
    "    train_ds=prepared_ds[\"train\"],\n",
    "    eval_ds=prepared_ds[\"valid\"],\n",
    "    training_config=TRAINING_CONFIG,\n",
    "    output_dir=\"out\",\n",
    "    debug=False,\n",
    "    env=NOTEBOOK_ENV,\n",
    ")\n",
    "\n",
    "trainer.train()\n",
    "end_training(run_name, trainer, MODELS_DIR_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fine Tuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAINING_CONFIG[\"freeze_encoder\"] = False "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_name = get_run_name(TRAINING_CONFIG)\n",
    "model = get_model(TRAINING_CONFIG, prepared_ds[\"train\"])\n",
    "\n",
    "trainer = get_trainer(\n",
    "    run_name=run_name,\n",
    "    model=model,\n",
    "    train_ds=prepared_ds[\"train\"],\n",
    "    eval_ds=prepared_ds[\"valid\"],\n",
    "    training_config=TRAINING_CONFIG,\n",
    "    output_dir=\"out\",\n",
    "    debug=False,\n",
    "    env=NOTEBOOK_ENV,\n",
    ")\n",
    "\n",
    "trainer.train()\n",
    "end_training(run_name, trainer, MODELS_DIR_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Classification Head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAINING_CONFIG[\"freeze_encoder\"] = False\n",
    "TRAINING_CONFIG[\"classifier_layers\"] = [256, 256]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_name = get_run_name(TRAINING_CONFIG)\n",
    "model = get_model(TRAINING_CONFIG, prepared_ds[\"train\"])\n",
    "\n",
    "trainer = get_trainer(\n",
    "    run_name=run_name,\n",
    "    model=model,\n",
    "    train_ds=prepared_ds[\"train\"],\n",
    "    eval_ds=prepared_ds[\"valid\"],\n",
    "    training_config=TRAINING_CONFIG,\n",
    "    output_dir=\"out\",\n",
    "    debug=False,\n",
    "    env=NOTEBOOK_ENV,\n",
    ")\n",
    "\n",
    "trainer.train()\n",
    "end_training(run_name, trainer, MODELS_DIR_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Wav2Vec2 Loss](res/report/wav2vec2_l.png \"Wav2Vec2 Loss\")\n",
    "![Wav2Vec2 Accuracy](res/report/wav2vec2_a.png \"Wav2Vec2 Accuracy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5Z6ixINFyTA7"
   },
   "source": [
    "## Whisper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAINING_CONFIG[\"feature_encoder\"] = \"whisper\"\n",
    "TRAINING_CONFIG[\"freeze_encoder\"] = True\n",
    "TRAINING_CONFIG[\"classifier_layers\"] = [256]\n",
    "TRAINING_CONFIG[\"classifier_dropout\"] = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prepared_ds = load_and_prepare_ds(TRAINING_CONFIG, FEATURES_CONFIG_SUBSET, df)\n",
    "\n",
    "prepared_ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Frozen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_name = get_run_name(TRAINING_CONFIG)\n",
    "model = get_model(TRAINING_CONFIG, prepared_ds[\"train\"])\n",
    "\n",
    "trainer = get_trainer(\n",
    "    run_name=run_name,\n",
    "    model=model,\n",
    "    train_ds=prepared_ds[\"train\"],\n",
    "    eval_ds=prepared_ds[\"valid\"],\n",
    "    training_config=TRAINING_CONFIG,\n",
    "    output_dir=\"out\",\n",
    "    debug=False,\n",
    "    env=NOTEBOOK_ENV,\n",
    ")\n",
    "\n",
    "trainer.train()\n",
    "end_training(run_name, trainer, MODELS_DIR_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fine-Tuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAINING_CONFIG[\"freeze_encoder\"] = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_name = get_run_name(TRAINING_CONFIG)\n",
    "model = get_model(TRAINING_CONFIG, prepared_ds[\"train\"])\n",
    "\n",
    "trainer = get_trainer(\n",
    "    run_name=run_name,\n",
    "    model=model,\n",
    "    train_ds=prepared_ds[\"train\"],\n",
    "    eval_ds=prepared_ds[\"valid\"],\n",
    "    training_config=TRAINING_CONFIG,\n",
    "    output_dir=\"out\",\n",
    "    debug=False,\n",
    "    env=NOTEBOOK_ENV,\n",
    ")\n",
    "\n",
    "trainer.train()\n",
    "end_training(run_name, trainer, MODELS_DIR_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Classification Head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAINING_CONFIG[\"classifier_layers\"] = [256, 256]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_name = get_run_name(TRAINING_CONFIG)\n",
    "model = get_model(TRAINING_CONFIG, prepared_ds[\"train\"])\n",
    "\n",
    "trainer = get_trainer(\n",
    "    run_name=run_name,\n",
    "    model=model,\n",
    "    train_ds=prepared_ds[\"train\"],\n",
    "    eval_ds=prepared_ds[\"valid\"],\n",
    "    training_config=TRAINING_CONFIG,\n",
    "    output_dir=\"out\",\n",
    "    debug=False,\n",
    "    env=NOTEBOOK_ENV,\n",
    ")\n",
    "\n",
    "trainer.train()\n",
    "end_training(run_name, trainer, MODELS_DIR_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Whisper Loss](res/report/whisper_l.png \"Whisper Loss\")\n",
    "![Whisper Accuracy](res/report/whisper_a.png \"Whisper Accuracy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training (Whole Dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAINING_CONFIG = {\n",
    "    \"epochs\": 3,\n",
    "    \"learning_rate\": 5e-5,\n",
    "    \"warmup\": 0.0,\n",
    "    \"train_batch_size\": 8,\n",
    "    \"eval_batch_size\": 16,\n",
    "    \"feature_encoder\": \"whisper\",\n",
    "    \"freeze_encoder\": False,\n",
    "    \"classifier_layers\": [256], \n",
    "    \"classifier_dropout\": 0.0,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Genre Classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading res/samples_clustered_genre6.csv\n",
      "16932 examples in DataFrame\n",
      "split\n",
      "train    13545\n",
      "test      1694\n",
      "valid     1693\n",
      "dtype: int64\n"
     ]
    },
    {
     "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>path</th>\n",
       "      <th>duration</th>\n",
       "      <th>id</th>\n",
       "      <th>genre</th>\n",
       "      <th>split</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>01 Hip Hop/Abandoned Brass Stabs.mp3</td>\n",
       "      <td>7.262041</td>\n",
       "      <td>01_Hip_Hop_Abandoned_Brass_Stabs</td>\n",
       "      <td>Hip Hop/RnB</td>\n",
       "      <td>test</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>01 Hip Hop/Against Time Keys.mp3</td>\n",
       "      <td>6.948571</td>\n",
       "      <td>01_Hip_Hop_Against_Time_Keys</td>\n",
       "      <td>Hip Hop/RnB</td>\n",
       "      <td>train</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>01 Hip Hop/Against Time Piano.mp3</td>\n",
       "      <td>6.948571</td>\n",
       "      <td>01_Hip_Hop_Against_Time_Piano</td>\n",
       "      <td>Hip Hop/RnB</td>\n",
       "      <td>train</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>01 Hip Hop/Against Time Sax Sample.mp3</td>\n",
       "      <td>6.948571</td>\n",
       "      <td>01_Hip_Hop_Against_Time_Sax_Sample</td>\n",
       "      <td>Hip Hop/RnB</td>\n",
       "      <td>valid</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>01 Hip Hop/Against Time Staccato Strings.mp3</td>\n",
       "      <td>6.948571</td>\n",
       "      <td>01_Hip_Hop_Against_Time_Staccato_Strings</td>\n",
       "      <td>Hip Hop/RnB</td>\n",
       "      <td>train</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                           path  duration  \\\n",
       "0          01 Hip Hop/Abandoned Brass Stabs.mp3  7.262041   \n",
       "1              01 Hip Hop/Against Time Keys.mp3  6.948571   \n",
       "2             01 Hip Hop/Against Time Piano.mp3  6.948571   \n",
       "3        01 Hip Hop/Against Time Sax Sample.mp3  6.948571   \n",
       "4  01 Hip Hop/Against Time Staccato Strings.mp3  6.948571   \n",
       "\n",
       "                                         id        genre  split  \n",
       "0          01_Hip_Hop_Abandoned_Brass_Stabs  Hip Hop/RnB   test  \n",
       "1              01_Hip_Hop_Against_Time_Keys  Hip Hop/RnB  train  \n",
       "2             01_Hip_Hop_Against_Time_Piano  Hip Hop/RnB  train  \n",
       "3        01_Hip_Hop_Against_Time_Sax_Sample  Hip Hop/RnB  valid  \n",
       "4  01_Hip_Hop_Against_Time_Staccato_Strings  Hip Hop/RnB  train  "
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Build the filename indicating the subset of the whole dataset with the specific configurations\n",
    "df = create_or_load_df(FEATURES_CONFIG_GEN)\n",
    "\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prepared_ds = load_and_prepare_ds(TRAINING_CONFIG, FEATURES_CONFIG_GEN, df)\n",
    "\n",
    "prepared_ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_name = get_run_name(TRAINING_CONFIG)\n",
    "model = get_model(TRAINING_CONFIG, prepared_ds[\"train\"])\n",
    "\n",
    "trainer = get_trainer(\n",
    "    run_name=run_name,\n",
    "    model=model,\n",
    "    train_ds=prepared_ds[\"train\"],\n",
    "    eval_ds=prepared_ds[\"valid\"],\n",
    "    training_config=TRAINING_CONFIG,\n",
    "    output_dir=\"out\",\n",
    "    debug=False,\n",
    "    env=NOTEBOOK_ENV,\n",
    ")\n",
    "\n",
    "trainer.train()\n",
    "end_training(run_name, trainer, MODELS_DIR_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Genre Classification Loss](res/report/genre_l.png \"Genre Classification Loss\")\n",
    "![Genre Classification Accuracy](res/report/genre_a.png \"Genre Classification Accuracy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Category Classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build the filename indicating the subset of the whole dataset with the specific configurations\n",
    "df = create_or_load_df(FEATURES_CONFIG_CAT)\n",
    "\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prepared_ds = load_and_prepare_ds(TRAINING_CONFIG, FEATURES_CONFIG_CAT, df)\n",
    "\n",
    "prepared_ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_name = get_run_name(TRAINING_CONFIG)\n",
    "model = get_model(TRAINING_CONFIG, prepared_ds[\"train\"])\n",
    "\n",
    "trainer = get_trainer(\n",
    "    run_name=run_name,\n",
    "    model=model,\n",
    "    train_ds=prepared_ds[\"train\"],\n",
    "    eval_ds=prepared_ds[\"valid\"],\n",
    "    training_config=TRAINING_CONFIG,\n",
    "    output_dir=\"out\",\n",
    "    debug=False,\n",
    "    env=NOTEBOOK_ENV,\n",
    ")\n",
    "\n",
    "trainer.train()\n",
    "end_training(run_name, trainer, MODELS_DIR_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Category Classification Loss](res/report/category_l.png \"Category Classification Loss\")\n",
    "![Category Classification Accuracy](res/report/category_a.png \"Category Classification Accuracy\")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "gpuClass": "standard",
  "interpreter": {
   "hash": "729786e4ac820b34608f27d07ec4fb597776bcc1a9d4d98bff8b4375359b1d17"
  },
  "kernelspec": {
   "display_name": "aii",
   "language": "python",
   "name": "aii"
  },
  "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.9"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "0fee81913e764382bfd0946d1661be8d": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_570e7fe98c364883a87014e220e5baf5",
      "placeholder": "​",
      "style": "IPY_MODEL_5d46f1f40ae24380a6b6d3304d1d1d36",
      "value": "Casting the dataset: 100%"
     }
    },
    "2febbdf5605149f494d9ac510883f262": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "35622cbe24684a7ea5382f1f536bd9ea": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_0fee81913e764382bfd0946d1661be8d",
       "IPY_MODEL_c9e2fff6b66e434e9806d869417f4eff",
       "IPY_MODEL_e3207490c4ba470f8e7cf61b97d84d54"
      ],
      "layout": "IPY_MODEL_2febbdf5605149f494d9ac510883f262"
     }
    },
    "4849127988f64e1e9558a88f220ad1de": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "570e7fe98c364883a87014e220e5baf5": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "5d46f1f40ae24380a6b6d3304d1d1d36": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "76877461d03e4018b8e5c5bc1721d2cb": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "7df7757d71ea46d5a4616f431066aef1": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "a70ff0866d564c96ae19d5d6aa91a727": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "c9e2fff6b66e434e9806d869417f4eff": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_4849127988f64e1e9558a88f220ad1de",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_76877461d03e4018b8e5c5bc1721d2cb",
      "value": 1
     }
    },
    "e3207490c4ba470f8e7cf61b97d84d54": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_a70ff0866d564c96ae19d5d6aa91a727",
      "placeholder": "​",
      "style": "IPY_MODEL_7df7757d71ea46d5a4616f431066aef1",
      "value": " 1/1 [00:12&lt;00:00, 12.42s/ba]"
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}