wav2vec2_fine_tuning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LKJ_ajH5W8yO"
   },
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "executionInfo": {
     "elapsed": 27623,
     "status": "ok",
     "timestamp": 1738348170382,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "wQS82V80W8yQ"
   },
   "outputs": [],
   "source": [
    "from dataset_generator import DatasetGenerator\n",
    "from utils import extract_all_chars, save_dict_as_json\n",
    "from data_preprocessor import Preprocessor\n",
    "from data_augmentation import AudioAugmentation\n",
    "from data_collator import DataCollatorCTCWithPadding\n",
    "\n",
    "import os\n",
    "import torch\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from torch.utils.data import Dataset\n",
    "import evaluate\n",
    "from transformers import (\n",
    "    Wav2Vec2CTCTokenizer,\n",
    "    Wav2Vec2FeatureExtractor,\n",
    "    Wav2Vec2Processor,\n",
    "    Wav2Vec2ForCTC,\n",
    "    TrainingArguments,\n",
    "    Trainer\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "executionInfo": {
     "elapsed": 2,
     "status": "ok",
     "timestamp": 1738348170382,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "0z93WBNcW8yR"
   },
   "outputs": [],
   "source": [
    "AUDIO_DIR = 'dataset'\n",
    "DATA_PATH = 'data.csv'\n",
    "\n",
    "word_character_map = {\n",
    "    'iskljuci': 'isključi',\n",
    "    'ukljuci': 'uključi'\n",
    "}\n",
    "\n",
    "TORCH_DATASETS_DIR = 'torch_datasets'\n",
    "\n",
    "MODEL_NAME = \"wav2vec2-finetuned-voice-commands\"\n",
    "MODELS_DIR = 'models'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "u-nc292GW8yS"
   },
   "source": [
    "# Gather from folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 847,
     "status": "ok",
     "timestamp": 1738346366174,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "ws3DMAUjW8yS",
    "outputId": "774a8f26-1db2-4024-f2f9-d82e61e1ff2f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset saved to data.csv\n"
     ]
    }
   ],
   "source": [
    "dg = DatasetGenerator(word_character_map)\n",
    "\n",
    "\n",
    "dg.generate(input_dir=AUDIO_DIR, output_file=DATA_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 206
    },
    "executionInfo": {
     "elapsed": 3,
     "status": "ok",
     "timestamp": 1738346366599,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "ZyaadFVuW8yS",
    "outputId": "49468b2b-4450-4133-8c34-3f6e7777a1e7"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "summary": "{\n  \"name\": \"df\",\n  \"rows\": 120,\n  \"fields\": [\n    {\n      \"column\": \"audio_filepath\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 120,\n        \"samples\": [\n          \"dataset/iskljuci-78-22-1.wav\",\n          \"dataset/iskljuci-37-21-2.wav\",\n          \"dataset/zvuk-38-21-2.wav\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"text\",\n      \"properties\": {\n        \"dtype\": \"category\",\n        \"num_unique_values\": 7,\n        \"samples\": [\n          \"zatvori\",\n          \"zvuk\",\n          \"uklju\\u010di\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    }\n  ]\n}",
       "type": "dataframe",
       "variable_name": "df"
      },
      "text/html": [
       "\n",
       "  <div id=\"df-664ac358-3253-4a53-bebd-7aa159fef70f\" class=\"colab-df-container\">\n",
       "    <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>audio_filepath</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>dataset/zatvori-38-21-1.wav</td>\n",
       "      <td>zatvori</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>dataset/zvuk-38-21-1.wav</td>\n",
       "      <td>zvuk</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>dataset/zvuk-38-21-3.wav</td>\n",
       "      <td>zvuk</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>dataset/zvuk-19-21-2.wav</td>\n",
       "      <td>zvuk</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>dataset/zvuk-38-21-2.wav</td>\n",
       "      <td>zvuk</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "    <div class=\"colab-df-buttons\">\n",
       "\n",
       "  <div class=\"colab-df-container\">\n",
       "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-664ac358-3253-4a53-bebd-7aa159fef70f')\"\n",
       "            title=\"Convert this dataframe to an interactive table.\"\n",
       "            style=\"display:none;\">\n",
       "\n",
       "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
       "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
       "  </svg>\n",
       "    </button>\n",
       "\n",
       "  <style>\n",
       "    .colab-df-container {\n",
       "      display:flex;\n",
       "      gap: 12px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert {\n",
       "      background-color: #E8F0FE;\n",
       "      border: none;\n",
       "      border-radius: 50%;\n",
       "      cursor: pointer;\n",
       "      display: none;\n",
       "      fill: #1967D2;\n",
       "      height: 32px;\n",
       "      padding: 0 0 0 0;\n",
       "      width: 32px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert:hover {\n",
       "      background-color: #E2EBFA;\n",
       "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "      fill: #174EA6;\n",
       "    }\n",
       "\n",
       "    .colab-df-buttons div {\n",
       "      margin-bottom: 4px;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert {\n",
       "      background-color: #3B4455;\n",
       "      fill: #D2E3FC;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert:hover {\n",
       "      background-color: #434B5C;\n",
       "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
       "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
       "      fill: #FFFFFF;\n",
       "    }\n",
       "  </style>\n",
       "\n",
       "    <script>\n",
       "      const buttonEl =\n",
       "        document.querySelector('#df-664ac358-3253-4a53-bebd-7aa159fef70f button.colab-df-convert');\n",
       "      buttonEl.style.display =\n",
       "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "\n",
       "      async function convertToInteractive(key) {\n",
       "        const element = document.querySelector('#df-664ac358-3253-4a53-bebd-7aa159fef70f');\n",
       "        const dataTable =\n",
       "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
       "                                                    [key], {});\n",
       "        if (!dataTable) return;\n",
       "\n",
       "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
       "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
       "          + ' to learn more about interactive tables.';\n",
       "        element.innerHTML = '';\n",
       "        dataTable['output_type'] = 'display_data';\n",
       "        await google.colab.output.renderOutput(dataTable, element);\n",
       "        const docLink = document.createElement('div');\n",
       "        docLink.innerHTML = docLinkHtml;\n",
       "        element.appendChild(docLink);\n",
       "      }\n",
       "    </script>\n",
       "  </div>\n",
       "\n",
       "\n",
       "<div id=\"df-ecc0024a-9267-4781-91c2-62176cc31483\">\n",
       "  <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-ecc0024a-9267-4781-91c2-62176cc31483')\"\n",
       "            title=\"Suggest charts\"\n",
       "            style=\"display:none;\">\n",
       "\n",
       "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
       "     width=\"24px\">\n",
       "    <g>\n",
       "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
       "    </g>\n",
       "</svg>\n",
       "  </button>\n",
       "\n",
       "<style>\n",
       "  .colab-df-quickchart {\n",
       "      --bg-color: #E8F0FE;\n",
       "      --fill-color: #1967D2;\n",
       "      --hover-bg-color: #E2EBFA;\n",
       "      --hover-fill-color: #174EA6;\n",
       "      --disabled-fill-color: #AAA;\n",
       "      --disabled-bg-color: #DDD;\n",
       "  }\n",
       "\n",
       "  [theme=dark] .colab-df-quickchart {\n",
       "      --bg-color: #3B4455;\n",
       "      --fill-color: #D2E3FC;\n",
       "      --hover-bg-color: #434B5C;\n",
       "      --hover-fill-color: #FFFFFF;\n",
       "      --disabled-bg-color: #3B4455;\n",
       "      --disabled-fill-color: #666;\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart {\n",
       "    background-color: var(--bg-color);\n",
       "    border: none;\n",
       "    border-radius: 50%;\n",
       "    cursor: pointer;\n",
       "    display: none;\n",
       "    fill: var(--fill-color);\n",
       "    height: 32px;\n",
       "    padding: 0;\n",
       "    width: 32px;\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart:hover {\n",
       "    background-color: var(--hover-bg-color);\n",
       "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "    fill: var(--button-hover-fill-color);\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart-complete:disabled,\n",
       "  .colab-df-quickchart-complete:disabled:hover {\n",
       "    background-color: var(--disabled-bg-color);\n",
       "    fill: var(--disabled-fill-color);\n",
       "    box-shadow: none;\n",
       "  }\n",
       "\n",
       "  .colab-df-spinner {\n",
       "    border: 2px solid var(--fill-color);\n",
       "    border-color: transparent;\n",
       "    border-bottom-color: var(--fill-color);\n",
       "    animation:\n",
       "      spin 1s steps(1) infinite;\n",
       "  }\n",
       "\n",
       "  @keyframes spin {\n",
       "    0% {\n",
       "      border-color: transparent;\n",
       "      border-bottom-color: var(--fill-color);\n",
       "      border-left-color: var(--fill-color);\n",
       "    }\n",
       "    20% {\n",
       "      border-color: transparent;\n",
       "      border-left-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "    }\n",
       "    30% {\n",
       "      border-color: transparent;\n",
       "      border-left-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "      border-right-color: var(--fill-color);\n",
       "    }\n",
       "    40% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "    }\n",
       "    60% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "    }\n",
       "    80% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "      border-bottom-color: var(--fill-color);\n",
       "    }\n",
       "    90% {\n",
       "      border-color: transparent;\n",
       "      border-bottom-color: var(--fill-color);\n",
       "    }\n",
       "  }\n",
       "</style>\n",
       "\n",
       "  <script>\n",
       "    async function quickchart(key) {\n",
       "      const quickchartButtonEl =\n",
       "        document.querySelector('#' + key + ' button');\n",
       "      quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
       "      quickchartButtonEl.classList.add('colab-df-spinner');\n",
       "      try {\n",
       "        const charts = await google.colab.kernel.invokeFunction(\n",
       "            'suggestCharts', [key], {});\n",
       "      } catch (error) {\n",
       "        console.error('Error during call to suggestCharts:', error);\n",
       "      }\n",
       "      quickchartButtonEl.classList.remove('colab-df-spinner');\n",
       "      quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
       "    }\n",
       "    (() => {\n",
       "      let quickchartButtonEl =\n",
       "        document.querySelector('#df-ecc0024a-9267-4781-91c2-62176cc31483 button');\n",
       "      quickchartButtonEl.style.display =\n",
       "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "    })();\n",
       "  </script>\n",
       "</div>\n",
       "    </div>\n",
       "  </div>\n"
      ],
      "text/plain": [
       "                audio_filepath     text\n",
       "0  dataset/zatvori-38-21-1.wav  zatvori\n",
       "1     dataset/zvuk-38-21-1.wav     zvuk\n",
       "2     dataset/zvuk-38-21-3.wav     zvuk\n",
       "3     dataset/zvuk-19-21-2.wav     zvuk\n",
       "4     dataset/zvuk-38-21-2.wav     zvuk"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.read_csv(DATA_PATH)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ey7xpv9TW8yT"
   },
   "source": [
    "# Create vocabulary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1ctZxHsUW8yT"
   },
   "outputs": [],
   "source": [
    "VOCAB_PATH = 'vocab.json'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 5,
     "status": "ok",
     "timestamp": 1738341855645,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "I24lDpzHW8yT",
    "outputId": "a7f98470-5406-4581-eb39-33ac38d413ea"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'i': 0,\n",
       " 'a': 1,\n",
       " 'e': 2,\n",
       " 'k': 3,\n",
       " 'č': 4,\n",
       " 'j': 5,\n",
       " 'u': 6,\n",
       " 'z': 7,\n",
       " 'l': 8,\n",
       " 'o': 9,\n",
       " 'v': 10,\n",
       " 's': 11,\n",
       " 'r': 12,\n",
       " 't': 13,\n",
       " '|': 14,\n",
       " '[UNK]': 15,\n",
       " '[PAD]': 16}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "words = df['text'].unique()\n",
    "\n",
    "vocab_list = extract_all_chars(words)\n",
    "\n",
    "vocab_list.extend(['|', '[UNK]', '[PAD]'])\n",
    "vocab_dict = {v: k for k, v in enumerate(vocab_list)}\n",
    "vocab_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8i0WozqFW8yT"
   },
   "outputs": [],
   "source": [
    "save_dict_as_json(VOCAB_PATH, vocab_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VxZA0VmcW8yU"
   },
   "source": [
    "# Loading the tokenizer, feature extractor and processor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "executionInfo": {
     "elapsed": 933,
     "status": "ok",
     "timestamp": 1738346315903,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "dyF8lCtaW8yU"
   },
   "outputs": [],
   "source": [
    "tokenizer = Wav2Vec2CTCTokenizer(\"./vocab.json\", unk_token=\"[UNK]\", pad_token=\"[PAD]\", word_delimiter_token=\"|\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "executionInfo": {
     "elapsed": 2,
     "status": "ok",
     "timestamp": 1738346315903,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "svxuKiVDW8yU"
   },
   "outputs": [],
   "source": [
    "feature_extractor = Wav2Vec2FeatureExtractor(feature_size=1,\n",
    "                                             sampling_rate=16000,\n",
    "                                             padding_value=0.0,\n",
    "                                             do_normalize=True,\n",
    "                                             return_attention_mask=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "executionInfo": {
     "elapsed": 1,
     "status": "ok",
     "timestamp": 1738346315903,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "zHNsnS0AW8yU"
   },
   "outputs": [],
   "source": [
    "processor = Wav2Vec2Processor(feature_extractor=feature_extractor, tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "W66sktZPW8yU"
   },
   "source": [
    "# Splitting, Preprocessing and Augmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "executionInfo": {
     "elapsed": 368,
     "status": "ok",
     "timestamp": 1738346349059,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "7cCS-tTZW8yU"
   },
   "outputs": [],
   "source": [
    "aug = AudioAugmentation(min_noise=0, max_noise=.005, time_stretch_rate=.9, pitch_shift_n_steps=2)\n",
    "\n",
    "train_preprocessor = Preprocessor(processor=processor, sr=16000, audio_augmentation=aug, augment_count=2)\n",
    "val_preprocessor = Preprocessor(processor=processor, sr=16000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "i66RcypZx0EW"
   },
   "source": [
    "We need to make sure that the same speaker's audios cannot be in both train and validation dataset. I will also remove my audios from the dataset so that I can test it on my voice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 366,
     "status": "ok",
     "timestamp": 1738346420888,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "SFbk06YAmNOL",
    "outputId": "9fa95fb1-f617-4a99-95ea-85064aedc607"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Set:\n",
      "                  audio_filepath     text\n",
      "3       dataset/zvuk-19-21-2.wav     zvuk\n",
      "6       dataset/zvuk-19-21-3.wav     zvuk\n",
      "8    dataset/zatvori-19-21-3.wav  zatvori\n",
      "9       dataset/zvuk-19-21-1.wav     zvuk\n",
      "11    dataset/otvori-19-21-1.wav   otvori\n",
      "..                           ...      ...\n",
      "114  dataset/zatvori-78-22-1.wav  zatvori\n",
      "115   dataset/zatvori-148-22.wav  zatvori\n",
      "116     dataset/zvuk-37-21-2.wav     zvuk\n",
      "117     dataset/zvuk-37-21-1.wav     zvuk\n",
      "118     dataset/zvuk-37-21-3.wav     zvuk\n",
      "\n",
      "[85 rows x 2 columns]\n",
      "\n",
      "Validation Set:\n",
      "                   audio_filepath      text\n",
      "46       dataset/zvuk-63-21-1.wav      zvuk\n",
      "55   dataset/iskljuci-63-21-1.wav  isključi\n",
      "57       dataset/zvuk-89-22-1.wav      zvuk\n",
      "63   dataset/iskljuci-89-22-1.wav  isključi\n",
      "70     dataset/otvori-63-21-1.wav    otvori\n",
      "78     dataset/otvori-89-22-1.wav    otvori\n",
      "85     dataset/svetlo-89-22-1.wav    svetlo\n",
      "86     dataset/svetlo-63-21-1.wav    svetlo\n",
      "89    dataset/ukljuci-89-22-1.wav   uključi\n",
      "98      dataset/vrata-89-22-1.wav     vrata\n",
      "102     dataset/vrata-63-21-1.wav     vrata\n",
      "103   dataset/ukljuci-63-21-1.wav   uključi\n",
      "110   dataset/zatvori-63-21-1.wav   zatvori\n",
      "119   dataset/zatvori-89-22-1.wav   zatvori\n"
     ]
    }
   ],
   "source": [
    "def extract_parts(row):\n",
    "    # Safely remove 'text' and split remaining path\n",
    "    audio_path = row['audio_filepath'].split('/',1)[1]\n",
    "    audio_path = audio_path.split('-',1)[1]\n",
    "    audio_path = audio_path.replace('.wav', '')\n",
    "    parts = audio_path.split('-', 2)  # Split into at most 3 parts (speaker, year, version)\n",
    "    return pd.Series(parts + [None] * (3 - len(parts)))  # Pad with None if version is missing\n",
    "\n",
    "df[['speaker', 'year', 'version']] = df.apply(extract_parts, axis=1)\n",
    "\n",
    "df['version'].fillna(1)\n",
    "df[['speaker', 'year']] = df[['speaker', 'year']].astype(int)  # Convert to integer\n",
    "# Remove files with speaker 38 and year 21 - This is my voice, I want to test it later on my voice\n",
    "df = df[~((df['speaker'] == 38) & (df['year'] == 21))]\n",
    "\n",
    "# Group by speaker-year\n",
    "groups = list(df.groupby(['speaker', 'year']).groups.keys())\n",
    "np.random.seed(42)\n",
    "np.random.shuffle(groups)  # Shuffle the groups\n",
    "\n",
    "# Split into train (80%) and validation (20%)\n",
    "split_idx = int(0.8 * len(groups))\n",
    "train_groups = set(groups[:split_idx])\n",
    "val_groups = set(groups[split_idx:])\n",
    "\n",
    "# Assign to train and validation sets\n",
    "train_df = df[df.set_index(['speaker', 'year']).index.isin(train_groups)]\n",
    "val_df = df[df.set_index(['speaker', 'year']).index.isin(val_groups)]\n",
    "\n",
    "# Drop extra columns if needed\n",
    "train_df = train_df[['audio_filepath', 'text']]\n",
    "val_df = val_df[['audio_filepath', 'text']]\n",
    "\n",
    "print(\"Train Set:\")\n",
    "print(train_df)\n",
    "print(\"\\nValidation Set:\")\n",
    "print(val_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 206
    },
    "executionInfo": {
     "elapsed": 18680,
     "status": "ok",
     "timestamp": 1738346449732,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "SToic1xwW8yV",
    "outputId": "9ca9384a-2e0b-4e5f-be6c-b115120e31dd"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "summary": "{\n  \"name\": \"train_df\",\n  \"rows\": 255,\n  \"fields\": [\n    {\n      \"column\": \"input_values\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 255,\n        \"samples\": [\n          \"tensor([-0.1769,  0.0629,  0.1565,  ..., -0.0818, -0.0528, -0.1062])\",\n          \"tensor([0.2387, 0.3726, 0.3026,  ..., 0.2616, 0.2065, 0.2265])\",\n          \"tensor([ 0.0806,  0.0333,  0.5164,  ...,  0.0149, -0.0210, -0.1682])\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"labels\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 255,\n        \"samples\": [\n          \"tensor([ 7,  1, 13, 10,  9, 12,  0])\",\n          \"tensor([ 7,  1, 13, 10,  9, 12,  0])\",\n          \"tensor([ 0, 11,  3,  8,  5,  6,  4,  0])\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    }\n  ]\n}",
       "type": "dataframe",
       "variable_name": "train_df"
      },
      "text/html": [
       "\n",
       "  <div id=\"df-a6400ab7-8cb2-480b-8309-fa83b2588fa3\" class=\"colab-df-container\">\n",
       "    <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>input_values</th>\n",
       "      <th>labels</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>[tensor(0.1722), tensor(0.2742), tensor(0.3092...</td>\n",
       "      <td>[tensor(7), tensor(10), tensor(6), tensor(3)]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>[tensor(0.4646), tensor(0.7110), tensor(0.1649...</td>\n",
       "      <td>[tensor(7), tensor(10), tensor(6), tensor(3)]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>[tensor(-0.0084), tensor(0.7243), tensor(0.259...</td>\n",
       "      <td>[tensor(7), tensor(10), tensor(6), tensor(3)]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>[tensor(-0.0727), tensor(-0.9619), tensor(-0.0...</td>\n",
       "      <td>[tensor(7), tensor(10), tensor(6), tensor(3)]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>[tensor(-0.2181), tensor(-0.8109), tensor(1.01...</td>\n",
       "      <td>[tensor(7), tensor(10), tensor(6), tensor(3)]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "    <div class=\"colab-df-buttons\">\n",
       "\n",
       "  <div class=\"colab-df-container\">\n",
       "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-a6400ab7-8cb2-480b-8309-fa83b2588fa3')\"\n",
       "            title=\"Convert this dataframe to an interactive table.\"\n",
       "            style=\"display:none;\">\n",
       "\n",
       "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
       "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
       "  </svg>\n",
       "    </button>\n",
       "\n",
       "  <style>\n",
       "    .colab-df-container {\n",
       "      display:flex;\n",
       "      gap: 12px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert {\n",
       "      background-color: #E8F0FE;\n",
       "      border: none;\n",
       "      border-radius: 50%;\n",
       "      cursor: pointer;\n",
       "      display: none;\n",
       "      fill: #1967D2;\n",
       "      height: 32px;\n",
       "      padding: 0 0 0 0;\n",
       "      width: 32px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert:hover {\n",
       "      background-color: #E2EBFA;\n",
       "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "      fill: #174EA6;\n",
       "    }\n",
       "\n",
       "    .colab-df-buttons div {\n",
       "      margin-bottom: 4px;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert {\n",
       "      background-color: #3B4455;\n",
       "      fill: #D2E3FC;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert:hover {\n",
       "      background-color: #434B5C;\n",
       "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
       "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
       "      fill: #FFFFFF;\n",
       "    }\n",
       "  </style>\n",
       "\n",
       "    <script>\n",
       "      const buttonEl =\n",
       "        document.querySelector('#df-a6400ab7-8cb2-480b-8309-fa83b2588fa3 button.colab-df-convert');\n",
       "      buttonEl.style.display =\n",
       "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "\n",
       "      async function convertToInteractive(key) {\n",
       "        const element = document.querySelector('#df-a6400ab7-8cb2-480b-8309-fa83b2588fa3');\n",
       "        const dataTable =\n",
       "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
       "                                                    [key], {});\n",
       "        if (!dataTable) return;\n",
       "\n",
       "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
       "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
       "          + ' to learn more about interactive tables.';\n",
       "        element.innerHTML = '';\n",
       "        dataTable['output_type'] = 'display_data';\n",
       "        await google.colab.output.renderOutput(dataTable, element);\n",
       "        const docLink = document.createElement('div');\n",
       "        docLink.innerHTML = docLinkHtml;\n",
       "        element.appendChild(docLink);\n",
       "      }\n",
       "    </script>\n",
       "  </div>\n",
       "\n",
       "\n",
       "<div id=\"df-82a92394-3723-4204-87b8-a8c42f8ac6c4\">\n",
       "  <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-82a92394-3723-4204-87b8-a8c42f8ac6c4')\"\n",
       "            title=\"Suggest charts\"\n",
       "            style=\"display:none;\">\n",
       "\n",
       "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
       "     width=\"24px\">\n",
       "    <g>\n",
       "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
       "    </g>\n",
       "</svg>\n",
       "  </button>\n",
       "\n",
       "<style>\n",
       "  .colab-df-quickchart {\n",
       "      --bg-color: #E8F0FE;\n",
       "      --fill-color: #1967D2;\n",
       "      --hover-bg-color: #E2EBFA;\n",
       "      --hover-fill-color: #174EA6;\n",
       "      --disabled-fill-color: #AAA;\n",
       "      --disabled-bg-color: #DDD;\n",
       "  }\n",
       "\n",
       "  [theme=dark] .colab-df-quickchart {\n",
       "      --bg-color: #3B4455;\n",
       "      --fill-color: #D2E3FC;\n",
       "      --hover-bg-color: #434B5C;\n",
       "      --hover-fill-color: #FFFFFF;\n",
       "      --disabled-bg-color: #3B4455;\n",
       "      --disabled-fill-color: #666;\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart {\n",
       "    background-color: var(--bg-color);\n",
       "    border: none;\n",
       "    border-radius: 50%;\n",
       "    cursor: pointer;\n",
       "    display: none;\n",
       "    fill: var(--fill-color);\n",
       "    height: 32px;\n",
       "    padding: 0;\n",
       "    width: 32px;\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart:hover {\n",
       "    background-color: var(--hover-bg-color);\n",
       "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "    fill: var(--button-hover-fill-color);\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart-complete:disabled,\n",
       "  .colab-df-quickchart-complete:disabled:hover {\n",
       "    background-color: var(--disabled-bg-color);\n",
       "    fill: var(--disabled-fill-color);\n",
       "    box-shadow: none;\n",
       "  }\n",
       "\n",
       "  .colab-df-spinner {\n",
       "    border: 2px solid var(--fill-color);\n",
       "    border-color: transparent;\n",
       "    border-bottom-color: var(--fill-color);\n",
       "    animation:\n",
       "      spin 1s steps(1) infinite;\n",
       "  }\n",
       "\n",
       "  @keyframes spin {\n",
       "    0% {\n",
       "      border-color: transparent;\n",
       "      border-bottom-color: var(--fill-color);\n",
       "      border-left-color: var(--fill-color);\n",
       "    }\n",
       "    20% {\n",
       "      border-color: transparent;\n",
       "      border-left-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "    }\n",
       "    30% {\n",
       "      border-color: transparent;\n",
       "      border-left-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "      border-right-color: var(--fill-color);\n",
       "    }\n",
       "    40% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "    }\n",
       "    60% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "    }\n",
       "    80% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "      border-bottom-color: var(--fill-color);\n",
       "    }\n",
       "    90% {\n",
       "      border-color: transparent;\n",
       "      border-bottom-color: var(--fill-color);\n",
       "    }\n",
       "  }\n",
       "</style>\n",
       "\n",
       "  <script>\n",
       "    async function quickchart(key) {\n",
       "      const quickchartButtonEl =\n",
       "        document.querySelector('#' + key + ' button');\n",
       "      quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
       "      quickchartButtonEl.classList.add('colab-df-spinner');\n",
       "      try {\n",
       "        const charts = await google.colab.kernel.invokeFunction(\n",
       "            'suggestCharts', [key], {});\n",
       "      } catch (error) {\n",
       "        console.error('Error during call to suggestCharts:', error);\n",
       "      }\n",
       "      quickchartButtonEl.classList.remove('colab-df-spinner');\n",
       "      quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
       "    }\n",
       "    (() => {\n",
       "      let quickchartButtonEl =\n",
       "        document.querySelector('#df-82a92394-3723-4204-87b8-a8c42f8ac6c4 button');\n",
       "      quickchartButtonEl.style.display =\n",
       "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "    })();\n",
       "  </script>\n",
       "</div>\n",
       "    </div>\n",
       "  </div>\n"
      ],
      "text/plain": [
       "                                        input_values  \\\n",
       "0  [tensor(0.1722), tensor(0.2742), tensor(0.3092...   \n",
       "1  [tensor(0.4646), tensor(0.7110), tensor(0.1649...   \n",
       "2  [tensor(-0.0084), tensor(0.7243), tensor(0.259...   \n",
       "3  [tensor(-0.0727), tensor(-0.9619), tensor(-0.0...   \n",
       "4  [tensor(-0.2181), tensor(-0.8109), tensor(1.01...   \n",
       "\n",
       "                                          labels  \n",
       "0  [tensor(7), tensor(10), tensor(6), tensor(3)]  \n",
       "1  [tensor(7), tensor(10), tensor(6), tensor(3)]  \n",
       "2  [tensor(7), tensor(10), tensor(6), tensor(3)]  \n",
       "3  [tensor(7), tensor(10), tensor(6), tensor(3)]  \n",
       "4  [tensor(7), tensor(10), tensor(6), tensor(3)]  "
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Preprocess data\n",
    "preprocessed_train_data = []\n",
    "preprocessed_val_data = []\n",
    "\n",
    "# Preprocess training data\n",
    "for _, row in train_df.iterrows():\n",
    "    preprocessed_train_data.extend(train_preprocessor.preprocess(row))\n",
    "\n",
    "# Preprocess validation data\n",
    "for _, row in val_df.iterrows():\n",
    "    preprocessed_val_data.extend(val_preprocessor.preprocess(row))\n",
    "\n",
    "# Create new dataframes for the preprocessed data\n",
    "train_df = pd.DataFrame(preprocessed_train_data)\n",
    "val_df = pd.DataFrame(preprocessed_val_data)\n",
    "\n",
    "train_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JtEnPlIsW8yV"
   },
   "source": [
    "# Generate PyTorch dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "executionInfo": {
     "elapsed": 497,
     "status": "ok",
     "timestamp": 1738346467135,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "QXSYuaUjW8yV"
   },
   "outputs": [],
   "source": [
    "# Create a PyTorch Dataset class\n",
    "class AudioDataset(Dataset):\n",
    "    def __init__(self, dataframe):\n",
    "        self.data = dataframe\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return {\n",
    "            \"input_values\": self.data.iloc[idx][\"input_values\"],\n",
    "            \"labels\": self.data.iloc[idx][\"labels\"],\n",
    "        }\n",
    "\n",
    "# Prepare datasets\n",
    "train_dataset = AudioDataset(train_df)\n",
    "val_dataset = AudioDataset(val_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gP44O4uKW8yV"
   },
   "source": [
    "## Save the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "executionInfo": {
     "elapsed": 3543,
     "status": "ok",
     "timestamp": 1738346471105,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "W_n655bhW8yV"
   },
   "outputs": [],
   "source": [
    "if not os.path.exists(TORCH_DATASETS_DIR):\n",
    "    os.mkdir(TORCH_DATASETS_DIR)\n",
    "\n",
    "torch.save(train_dataset, os.path.join(TORCH_DATASETS_DIR, 'train.pt'))\n",
    "torch.save(val_dataset, os.path.join(TORCH_DATASETS_DIR, 'val.pt'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jC4dX-voW8yV"
   },
   "source": [
    "# Loading the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HF_TOKEN = os.environ.get('HF_TOKEN')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 1342,
     "status": "ok",
     "timestamp": 1738346320341,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "-C_RVikKW8yV",
    "outputId": "daae6c9e-76b0-4b7a-c34b-48f3d498ab7f"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of Wav2Vec2ForCTC were not initialized from the model checkpoint at facebook/wav2vec2-large-xlsr-53 and are newly initialized: ['lm_head.bias', 'lm_head.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "model = Wav2Vec2ForCTC.from_pretrained(\n",
    "    \"facebook/wav2vec2-large-xlsr-53\",\n",
    "    attention_dropout=0.2,\n",
    "    hidden_dropout=0.2,\n",
    "    feat_proj_dropout=0.05,\n",
    "    mask_time_prob=0.04,\n",
    "    layerdrop=0.15,\n",
    "    ctc_loss_reduction=\"mean\",\n",
    "    pad_token_id=processor.tokenizer.pad_token_id,\n",
    "    vocab_size=len(processor.tokenizer),\n",
    "    token = HF_TOKEN\n",
    ")\n",
    "\n",
    "model.config.vocab_size = len(processor.tokenizer)\n",
    "\n",
    "model.freeze_feature_encoder()\n",
    "model.gradient_checkpointing_enable()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "c1Isq6w6W8yV"
   },
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "executionInfo": {
     "elapsed": 404,
     "status": "ok",
     "timestamp": 1738346324037,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "UTLY_FWFW8yV"
   },
   "outputs": [],
   "source": [
    "data_collator = DataCollatorCTCWithPadding(processor=processor, padding=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "executionInfo": {
     "elapsed": 503,
     "status": "ok",
     "timestamp": 1738346325888,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "oh7PMBk8W8yW"
   },
   "outputs": [],
   "source": [
    "if not os.path.exists(MODELS_DIR):\n",
    "    os.mkdir(MODELS_DIR)\n",
    "\n",
    "model_out_dir = os.path.join(MODELS_DIR, MODEL_NAME)\n",
    "if not os.path.exists(model_out_dir):\n",
    "    os.mkdir(model_out_dir)\n",
    "\n",
    "logs_dir = os.path.join(MODELS_DIR, MODEL_NAME, 'logs')\n",
    "if not os.path.exists(logs_dir):\n",
    "    os.mkdir(logs_dir)\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=model_out_dir,                               # Directory to save model checkpoints\n",
    "    overwrite_output_dir=True,\n",
    "    eval_strategy=\"steps\",                                  # Evaluate every N steps\n",
    "    per_device_train_batch_size=8,                          # Batch size for training\n",
    "    per_device_eval_batch_size=8,                           # Batch size for evaluation\n",
    "    gradient_accumulation_steps=2,                          # Gradient accumulation\n",
    "    learning_rate=3e-4,                                     # Learning rate\n",
    "    warmup_steps=500,                                       # Warmup steps for LR scheduler\n",
    "    num_train_epochs=100,                                   # Number of epochs\n",
    "    logging_dir=logs_dir,                                   # Directory for logging\n",
    "    logging_steps=10,                                       # Log every N steps\n",
    "    save_steps=50,                                          # Save checkpoint every N steps\n",
    "    eval_steps=50,\n",
    "    save_total_limit=2,                                     # Only keep the last 2 checkpoints\n",
    "    fp16=True,                                              # Use mixed precision\n",
    "    dataloader_num_workers=2,                               # Number of workers for DataLoader\n",
    "    load_best_model_at_end=True,                            # Load the best model at the end\n",
    "    metric_for_best_model=\"wer\",                            # Metric to determine best model\n",
    "    greater_is_better=False,                                # Smaller WER is better\n",
    "    save_safetensors=False,\n",
    "    seed=42,                                                # Random seed for reproducibility\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "executionInfo": {
     "elapsed": 3576,
     "status": "ok",
     "timestamp": 1738346332172,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "6RUeiukaW8yW"
   },
   "outputs": [],
   "source": [
    "# Define WER metric\n",
    "wer_metric = evaluate.load(\"wer\")\n",
    "\n",
    "def compute_metrics(pred):\n",
    "    pred_logits = pred.predictions\n",
    "    pred_ids = torch.argmax(torch.tensor(pred_logits), dim=-1)\n",
    "\n",
    "    # Decode predictions and labels\n",
    "    pred_str = processor.batch_decode(pred_ids, skip_special_tokens=True)\n",
    "    label_str = processor.batch_decode(pred.label_ids, skip_special_tokens=True)\n",
    "\n",
    "    # Compute WER\n",
    "    wer = wer_metric.compute(predictions=pred_str, references=label_str)\n",
    "    return {\"wer\": wer}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "executionInfo": {
     "elapsed": 9209,
     "status": "ok",
     "timestamp": 1738346484088,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "XDp3vV9QW8yW"
   },
   "outputs": [],
   "source": [
    "# Initialize Trainer\n",
    "trainer = Trainer(\n",
    "    model=model,                               # Wav2Vec2 model\n",
    "    data_collator=data_collator,               # Data Collator\n",
    "    args=training_args,                        # Training arguments\n",
    "    train_dataset=train_dataset,               # Training dataset\n",
    "    eval_dataset=val_dataset,                  # Validation dataset\n",
    "    processing_class=processor,                # Processor\n",
    "    compute_metrics=compute_metrics,           # WER metric\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "executionInfo": {
     "elapsed": 1586832,
     "status": "error",
     "timestamp": 1738343829458,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "06xXPN-7W8yW",
    "outputId": "3a1dde73-47cb-4734-f139-7c0a5a650318"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m The `run_name` is currently set to the same value as `TrainingArguments.output_dir`. If this was not intended, please specify a different run name by setting the `TrainingArguments.run_name` parameter.\n"
     ]
    },
    {
     "data": {
      "application/javascript": "\n        window._wandbApiKey = new Promise((resolve, reject) => {\n            function loadScript(url) {\n            return new Promise(function(resolve, reject) {\n                let newScript = document.createElement(\"script\");\n                newScript.onerror = reject;\n                newScript.onload = resolve;\n                document.body.appendChild(newScript);\n                newScript.src = url;\n            });\n            }\n            loadScript(\"https://cdn.jsdelivr.net/npm/postmate/build/postmate.min.js\").then(() => {\n            const iframe = document.createElement('iframe')\n            iframe.style.cssText = \"width:0;height:0;border:none\"\n            document.body.appendChild(iframe)\n            const handshake = new Postmate({\n                container: iframe,\n                url: 'https://wandb.ai/authorize'\n            });\n            const timeout = setTimeout(() => reject(\"Couldn't auto authenticate\"), 5000)\n            handshake.then(function(child) {\n                child.on('authorize', data => {\n                    clearTimeout(timeout)\n                    resolve(data)\n                });\n            });\n            })\n        });\n    ",
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m If you're specifying your api key in code, ensure this code is not shared publicly.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Consider setting the WANDB_API_KEY environment variable, or running `wandb login` from the command line.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Appending key for api.wandb.ai to your netrc file: /root/.netrc\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mmihailo-radovic11\u001b[0m (\u001b[33mmihailo-radovic11-student\u001b[0m) to \u001b[32mhttps://api.wandb.ai\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Using wandb-core as the SDK backend.  Please refer to https://wandb.me/wandb-core for more information.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.19.5"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/content/drive/MyDrive/voice_command_recognition/wandb/run-20250131_164726-pd5zsjc1</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/mihailo-radovic11-student/huggingface/runs/pd5zsjc1' target=\"_blank\">models/wav2vec2-finetuned-voice-commands</a></strong> to <a href='https://wandb.ai/mihailo-radovic11-student/huggingface' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/developer-guide' target=\"_blank\">docs</a>)<br>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/mihailo-radovic11-student/huggingface' target=\"_blank\">https://wandb.ai/mihailo-radovic11-student/huggingface</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/mihailo-radovic11-student/huggingface/runs/pd5zsjc1' target=\"_blank\">https://wandb.ai/mihailo-radovic11-student/huggingface/runs/pd5zsjc1</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='129' max='1600' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [ 129/1600 03:05 < 35:52, 0.68 it/s, Epoch 8/100]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Wer</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>50</td>\n",
       "      <td>82.420400</td>\n",
       "      <td>36.380486</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>100</td>\n",
       "      <td>40.076100</td>\n",
       "      <td>11.197175</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='1251' max='1600' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [1251/1600 28:53 < 08:04, 0.72 it/s, Epoch 78.12/100]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Wer</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>50</td>\n",
       "      <td>82.420400</td>\n",
       "      <td>36.380486</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>100</td>\n",
       "      <td>40.076100</td>\n",
       "      <td>11.197175</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>150</td>\n",
       "      <td>7.223500</td>\n",
       "      <td>3.168533</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>200</td>\n",
       "      <td>6.062800</td>\n",
       "      <td>2.939375</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>250</td>\n",
       "      <td>5.949000</td>\n",
       "      <td>2.916493</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>300</td>\n",
       "      <td>5.839600</td>\n",
       "      <td>2.906181</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>350</td>\n",
       "      <td>5.808700</td>\n",
       "      <td>2.914468</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>400</td>\n",
       "      <td>5.738100</td>\n",
       "      <td>2.841084</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>450</td>\n",
       "      <td>5.093800</td>\n",
       "      <td>2.332966</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>500</td>\n",
       "      <td>4.631900</td>\n",
       "      <td>2.114949</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>550</td>\n",
       "      <td>3.266400</td>\n",
       "      <td>1.013510</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>600</td>\n",
       "      <td>2.014100</td>\n",
       "      <td>0.432637</td>\n",
       "      <td>0.785714</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>650</td>\n",
       "      <td>1.223800</td>\n",
       "      <td>0.241662</td>\n",
       "      <td>0.428571</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>700</td>\n",
       "      <td>1.674000</td>\n",
       "      <td>0.194477</td>\n",
       "      <td>0.571429</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>750</td>\n",
       "      <td>1.003000</td>\n",
       "      <td>0.172041</td>\n",
       "      <td>0.285714</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>800</td>\n",
       "      <td>0.669000</td>\n",
       "      <td>0.063172</td>\n",
       "      <td>0.142857</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>850</td>\n",
       "      <td>0.609000</td>\n",
       "      <td>0.019306</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>900</td>\n",
       "      <td>1.194900</td>\n",
       "      <td>0.011274</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>950</td>\n",
       "      <td>0.401200</td>\n",
       "      <td>0.025712</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1000</td>\n",
       "      <td>0.552300</td>\n",
       "      <td>0.005187</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1050</td>\n",
       "      <td>0.494600</td>\n",
       "      <td>0.004375</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1100</td>\n",
       "      <td>0.367300</td>\n",
       "      <td>0.003725</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1150</td>\n",
       "      <td>0.369200</td>\n",
       "      <td>0.003601</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1200</td>\n",
       "      <td>0.293100</td>\n",
       "      <td>0.002365</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1250</td>\n",
       "      <td>0.237300</td>\n",
       "      <td>0.001879</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "RuntimeError",
     "evalue": "[enforce fail at inline_container.cc:603] . unexpected pos 2304414144 vs 2304414032",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m/usr/local/lib/python3.11/dist-packages/torch/serialization.py\u001b[0m in \u001b[0;36msave\u001b[0;34m(obj, f, pickle_module, pickle_protocol, _use_new_zipfile_serialization, _disable_byteorder_record)\u001b[0m\n\u001b[1;32m    849\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0m_open_zipfile_writer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mopened_zipfile\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 850\u001b[0;31m             _save(\n\u001b[0m\u001b[1;32m    851\u001b[0m                 \u001b[0mobj\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.11/dist-packages/torch/serialization.py\u001b[0m in \u001b[0;36m_save\u001b[0;34m(obj, zip_file, pickle_module, pickle_protocol, _disable_byteorder_record)\u001b[0m\n\u001b[1;32m   1113\u001b[0m             \u001b[0;31m# Now that it is on the CPU we can directly copy it into the zip file\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1114\u001b[0;31m             \u001b[0mzip_file\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwrite_record\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstorage\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_bytes\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1115\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: [enforce fail at inline_container.cc:778] . PytorchStreamWriter failed writing file data/65: file write failed",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-31-635967aca760>\u001b[0m in \u001b[0;36m<cell line: 0>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m# Start training\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m/usr/local/lib/python3.11/dist-packages/transformers/trainer.py\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)\u001b[0m\n\u001b[1;32m   2162\u001b[0m                 \u001b[0mhf_hub_utils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0menable_progress_bars\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2163\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2164\u001b[0;31m             return inner_training_loop(\n\u001b[0m\u001b[1;32m   2165\u001b[0m                 \u001b[0margs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2166\u001b[0m                 \u001b[0mresume_from_checkpoint\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mresume_from_checkpoint\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.11/dist-packages/transformers/trainer.py\u001b[0m in \u001b[0;36m_inner_training_loop\u001b[0;34m(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)\u001b[0m\n\u001b[1;32m   2589\u001b[0m                         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstate\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mepoch\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mepoch\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mstep\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;36m1\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0msteps_skipped\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0msteps_in_epoch\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2590\u001b[0m                         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcontrol\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcallback_handler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mon_step_end\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstate\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcontrol\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2591\u001b[0;31m                         self._maybe_log_save_evaluate(\n\u001b[0m\u001b[1;32m   2592\u001b[0m                             \u001b[0mtr_loss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_norm\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrial\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mepoch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mignore_keys_for_eval\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstart_time\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2593\u001b[0m                         )\n",
      "\u001b[0;32m/usr/local/lib/python3.11/dist-packages/transformers/trainer.py\u001b[0m in \u001b[0;36m_maybe_log_save_evaluate\u001b[0;34m(self, tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval, start_time)\u001b[0m\n\u001b[1;32m   3054\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3055\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcontrol\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshould_save\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3056\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_save_checkpoint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrial\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   3057\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcontrol\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcallback_handler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mon_save\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstate\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcontrol\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3058\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.11/dist-packages/transformers/trainer.py\u001b[0m in \u001b[0;36m_save_checkpoint\u001b[0;34m(self, model, trial)\u001b[0m\n\u001b[1;32m   3190\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msave_only_model\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3191\u001b[0m             \u001b[0;31m# Save optimizer and scheduler\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3192\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_save_optimizer_and_scheduler\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput_dir\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   3193\u001b[0m             \u001b[0;31m# Save RNG state\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3194\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_save_rng_state\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput_dir\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.11/dist-packages/transformers/trainer.py\u001b[0m in \u001b[0;36m_save_optimizer_and_scheduler\u001b[0;34m(self, output_dir)\u001b[0m\n\u001b[1;32m   3311\u001b[0m         \u001b[0;32melif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshould_save\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3312\u001b[0m             \u001b[0;31m# deepspeed.save_checkpoint above saves model/optim/sched\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3313\u001b[0;31m             \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msave\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstate_dict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput_dir\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mOPTIMIZER_NAME\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   3314\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3315\u001b[0m         \u001b[0;31m# Save SCHEDULER & SCALER\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.11/dist-packages/torch/serialization.py\u001b[0m in \u001b[0;36msave\u001b[0;34m(obj, f, pickle_module, pickle_protocol, _use_new_zipfile_serialization, _disable_byteorder_record)\u001b[0m\n\u001b[1;32m    847\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    848\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0m_use_new_zipfile_serialization\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 849\u001b[0;31m         \u001b[0;32mwith\u001b[0m \u001b[0m_open_zipfile_writer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mopened_zipfile\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    850\u001b[0m             _save(\n\u001b[1;32m    851\u001b[0m                 \u001b[0mobj\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.11/dist-packages/torch/serialization.py\u001b[0m in \u001b[0;36m__exit__\u001b[0;34m(self, *args)\u001b[0m\n\u001b[1;32m    688\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    689\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__exit__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 690\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfile_like\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwrite_end_of_file\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    691\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfile_stream\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    692\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfile_stream\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclose\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: [enforce fail at inline_container.cc:603] . unexpected pos 2304414144 vs 2304414032"
     ]
    }
   ],
   "source": [
    "# Start training\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 854,
     "referenced_widgets": [
      "34099aca857e4a18a68207004b2f24a9",
      "796ffb7b80b24575aab46137ca8fe457",
      "e0b0ffef05de4183b954553ecd80c402",
      "07d26ed10bf0476f94d6d4e385936557",
      "b3e80eaaa8024dc9b6aa0e490539f39f",
      "a849ef312b01417ca19182f15cf9f7bb",
      "c5d480a18ff6461da9955554857be736",
      "3e9252b1f92246d1b8f7b5a80f60ff66"
     ]
    },
    "executionInfo": {
     "elapsed": 336423,
     "status": "ok",
     "timestamp": 1738347228472,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "quanl1WjXoCh",
    "outputId": "cbe23512-e09b-4515-c884-23bd2d4f2802"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.11/dist-packages/transformers/trainer.py:3420: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  torch.load(os.path.join(checkpoint, OPTIMIZER_NAME), map_location=map_location)\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m The `run_name` is currently set to the same value as `TrainingArguments.output_dir`. If this was not intended, please specify a different run name by setting the `TrainingArguments.run_name` parameter.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mmihailo-radovic11\u001b[0m (\u001b[33mmihailo-radovic11-student\u001b[0m) to \u001b[32mhttps://api.wandb.ai\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Using wandb-core as the SDK backend.  Please refer to https://wandb.me/wandb-core for more information.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "34099aca857e4a18a68207004b2f24a9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='Waiting for wandb.init()...\\r'), FloatProgress(value=0.011112468488893986, max=1.0…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.19.5"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/content/drive/MyDrive/voice_command_recognition/wandb/run-20250131_180301-zhucvkip</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/mihailo-radovic11-student/huggingface/runs/zhucvkip' target=\"_blank\">models/wav2vec2-finetuned-voice-commands</a></strong> to <a href='https://wandb.ai/mihailo-radovic11-student/huggingface' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/developer-guide' target=\"_blank\">docs</a>)<br>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/mihailo-radovic11-student/huggingface' target=\"_blank\">https://wandb.ai/mihailo-radovic11-student/huggingface</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/mihailo-radovic11-student/huggingface/runs/zhucvkip' target=\"_blank\">https://wandb.ai/mihailo-radovic11-student/huggingface/runs/zhucvkip</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.11/dist-packages/transformers/trainer.py:3083: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  checkpoint_rng_state = torch.load(rng_file)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='1401' max='1600' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [1401/1600 04:49 < 04:49, 0.69 it/s, Epoch 87.50/100]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Wer</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1250</td>\n",
       "      <td>0.253400</td>\n",
       "      <td>0.002069</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1300</td>\n",
       "      <td>0.344300</td>\n",
       "      <td>0.001716</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1350</td>\n",
       "      <td>0.266600</td>\n",
       "      <td>0.002186</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1400</td>\n",
       "      <td>0.156500</td>\n",
       "      <td>0.001543</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='1600' max='1600' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [1600/1600 10:36, Epoch 100/100]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Wer</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1250</td>\n",
       "      <td>0.253400</td>\n",
       "      <td>0.002069</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1300</td>\n",
       "      <td>0.344300</td>\n",
       "      <td>0.001716</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1350</td>\n",
       "      <td>0.266600</td>\n",
       "      <td>0.002186</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1400</td>\n",
       "      <td>0.156500</td>\n",
       "      <td>0.001543</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1450</td>\n",
       "      <td>0.275100</td>\n",
       "      <td>0.002841</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1500</td>\n",
       "      <td>0.331500</td>\n",
       "      <td>0.001397</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1550</td>\n",
       "      <td>0.168100</td>\n",
       "      <td>0.001240</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1600</td>\n",
       "      <td>0.243100</td>\n",
       "      <td>0.001211</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=1600, training_loss=0.06577926870435476, metrics={'train_runtime': 646.6687, 'train_samples_per_second': 39.433, 'train_steps_per_second': 2.474, 'total_flos': 2.692254420058742e+18, 'train_loss': 0.06577926870435476, 'epoch': 100.0})"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train(resume_from_checkpoint=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 74
    },
    "executionInfo": {
     "elapsed": 403,
     "status": "ok",
     "timestamp": 1738347228872,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "sNe8YgLrw3ET",
    "outputId": "feafe71e-0829-437e-bb54-9ab38ee7ad8d"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='2' max='2' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [2/2 00:00]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.019305603578686714, 'eval_wer': 0.0, 'eval_runtime': 0.4752, 'eval_samples_per_second': 29.46, 'eval_steps_per_second': 4.209, 'epoch': 100.0}\n"
     ]
    }
   ],
   "source": [
    "final_metrics = trainer.evaluate(val_dataset)\n",
    "print(final_metrics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 16574,
     "status": "ok",
     "timestamp": 1738347246862,
     "user": {
      "displayName": "Mihailo Radovic",
      "userId": "01884334755596507994"
     },
     "user_tz": -60
    },
    "id": "Rwn2dDY6z5UN",
    "outputId": "862130e9-25be-4f7a-ba1c-ff6d33bfed3a"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.save_pretrained(f\"{MODELS_DIR}/best_model\")\n",
    "processor.save_pretrained(f\"{MODELS_DIR}/best_model_processor\")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "u-nc292GW8yS",
    "Ey7xpv9TW8yT",
    "VxZA0VmcW8yU"
   ],
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.11.5"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "07d26ed10bf0476f94d6d4e385936557": {
     "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
     }
    },
    "34099aca857e4a18a68207004b2f24a9": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "VBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "VBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "VBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_796ffb7b80b24575aab46137ca8fe457",
       "IPY_MODEL_e0b0ffef05de4183b954553ecd80c402"
      ],
      "layout": "IPY_MODEL_07d26ed10bf0476f94d6d4e385936557"
     }
    },
    "3e9252b1f92246d1b8f7b5a80f60ff66": {
     "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": ""
     }
    },
    "796ffb7b80b24575aab46137ca8fe457": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "LabelModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "LabelModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "LabelView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_b3e80eaaa8024dc9b6aa0e490539f39f",
      "placeholder": "​",
      "style": "IPY_MODEL_a849ef312b01417ca19182f15cf9f7bb",
      "value": "Waiting for wandb.init()...\r"
     }
    },
    "a849ef312b01417ca19182f15cf9f7bb": {
     "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": ""
     }
    },
    "b3e80eaaa8024dc9b6aa0e490539f39f": {
     "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
     }
    },
    "c5d480a18ff6461da9955554857be736": {
     "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
     }
    },
    "e0b0ffef05de4183b954553ecd80c402": {
     "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": "",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_c5d480a18ff6461da9955554857be736",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_3e9252b1f92246d1b8f7b5a80f60ff66",
      "value": 1
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}