muse-project.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MUSE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-01-09T10:50:15.262984Z",
     "iopub.status.busy": "2025-01-09T10:50:15.262732Z",
     "iopub.status.idle": "2025-01-09T10:50:25.569927Z",
     "shell.execute_reply": "2025-01-09T10:50:25.568795Z",
     "shell.execute_reply.started": "2025-01-09T10:50:15.262959Z"
    },
    "trusted": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/anaconda3/envs/audio/lib/python3.12/site-packages/torchvision/io/image.py:14: UserWarning: Failed to load image Python extension: 'dlopen(/opt/anaconda3/envs/audio/lib/python3.12/site-packages/torchvision/image.so, 0x0006): Library not loaded: @rpath/libjpeg.9.dylib\n",
      "  Referenced from: <367D4265-B20F-34BD-94EB-4F3EE47C385B> /opt/anaconda3/envs/audio/lib/python3.12/site-packages/torchvision/image.so\n",
      "  Reason: tried: '/opt/anaconda3/envs/audio/lib/python3.12/site-packages/torchvision/../../../libjpeg.9.dylib' (no such file), '/opt/anaconda3/envs/audio/lib/python3.12/site-packages/torchvision/../../../libjpeg.9.dylib' (no such file), '/opt/anaconda3/envs/audio/lib/python3.12/lib-dynload/../../libjpeg.9.dylib' (no such file), '/opt/anaconda3/envs/audio/bin/../lib/libjpeg.9.dylib' (no such file)'If you don't plan on using image functionality from `torchvision.io`, you can ignore this warning. Otherwise, there might be something wrong with your environment. Did you have `libjpeg` or `libpng` installed before building `torchvision` from source?\n",
      "  warn(\n",
      "[nltk_data] Downloading package punkt to\n",
      "[nltk_data]     /Users/filiplandin/nltk_data...\n",
      "[nltk_data]   Unzipping tokenizers/punkt.zip.\n",
      "[nltk_data] Downloading package stopwords to\n",
      "[nltk_data]     /Users/filiplandin/nltk_data...\n",
      "[nltk_data]   Unzipping corpora/stopwords.zip.\n"
     ]
    }
   ],
   "source": [
    "\n",
    "import tarfile\n",
    "import os\n",
    "import os\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import librosa\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import Audio\n",
    "import random\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from sklearn.metrics import classification_report, accuracy_score, f1_score\n",
    "import joblib\n",
    "import nltk\n",
    "import torchaudio\n",
    "from torch.nn.utils.rnn import pad_sequence\n",
    "from transformers import AutoModel, AutoProcessor, DistilBertTokenizer\n",
    "\n",
    "# Ensure NLTK data is downloaded\n",
    "nltk.download('punkt')\n",
    "nltk.download('stopwords')\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": [
    "def meld_collate_fn(batch):\n",
    "    # batch is a list of conversation dicts (one per item in dataset)\n",
    "    # We can combine them into a single batch,\n",
    "    # but each conversation may have different # of utterances.\n",
    "\n",
    "    dialog_ids = []\n",
    "    fbank_lists = []\n",
    "    text_lists = []\n",
    "    emotion_lists = []\n",
    "    sentiment_lists = []\n",
    "\n",
    "    for conv in batch:\n",
    "        dialog_ids.append(conv[\"dialog_id\"])\n",
    "\n",
    "        # Convert fbank_list (list of numpy arrays) to tensors and pad\n",
    "        fbank_tensors = [torch.tensor(fbank) for fbank in conv[\"fbank_list\"]]\n",
    "        # Pad along the time dim (T)\n",
    "        fbank_padded = pad_sequence(fbank_tensors, batch_first=True)\n",
    "        fbank_lists.append(fbank_padded)\n",
    "\n",
    "        text_lists.append(conv[\"text_list\"])\n",
    "        emotion_lists.append(torch.tensor(conv[\"emotion_list\"]))\n",
    "        sentiment_lists.append(torch.tensor(conv[\"sentiment_list\"]))\n",
    "\n",
    "    # Return them \"as is\", or do further padding if needed.\n",
    "    return {\n",
    "        \"dialog_ids\": dialog_ids,\n",
    "        \"fbank_lists\": fbank_lists,\n",
    "        \"text_lists\": text_lists,\n",
    "        \"emotion_lists\": emotion_lists,\n",
    "        \"sentiment_lists\": sentiment_lists\n",
    "    }\n",
    "\n",
    "\n",
    "def custom_text_preprocessor(text):\n",
    "    return text"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### fbanks extractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def wav2fbank(filename):\n",
    "    waveform, sr = torchaudio.load(filename)\n",
    "    waveform = waveform - waveform.mean()\n",
    "\n",
    "    try:\n",
    "        fbank = torchaudio.compliance.kaldi.fbank(\n",
    "            waveform,\n",
    "            htk_compat=True,\n",
    "            sample_frequency=sr,\n",
    "            use_energy=False,\n",
    "            window_type='hanning',\n",
    "            num_mel_bins=128,\n",
    "            dither=0.0,\n",
    "            frame_shift=10\n",
    "        )\n",
    "    except:\n",
    "        fbank = torch.zeros([512, 128]) + 0.01\n",
    "        print('there is a loading error')\n",
    "\n",
    "    target_length = 1024\n",
    "    n_frames = fbank.shape[0]\n",
    "\n",
    "    p = target_length - n_frames\n",
    "\n",
    "    # cut and pad\n",
    "    if p > 0:\n",
    "        m = torch.nn.ZeroPad2d((0, 0, 0, p))\n",
    "        fbank = m(fbank)\n",
    "    elif p < 0:\n",
    "        fbank = fbank[0:target_length, :]\n",
    "\n",
    "    return fbank"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: 'data/test/test_sent_emo.csv'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[4], line 3\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# convert audio to fbank from train, dev and test\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m mode \u001b[38;5;129;01min\u001b[39;00m [\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtrain\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mdev\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtest\u001b[39m\u001b[38;5;124m'\u001b[39m]:\n\u001b[0;32m----> 3\u001b[0m     df \u001b[38;5;241m=\u001b[39m \u001b[43mpd\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mread_csv\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43mf\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mdata/\u001b[39;49m\u001b[38;5;132;43;01m{\u001b[39;49;00m\u001b[43mmode\u001b[49m\u001b[38;5;132;43;01m}\u001b[39;49;00m\u001b[38;5;124;43m/\u001b[39;49m\u001b[38;5;132;43;01m{\u001b[39;49;00m\u001b[43mmode\u001b[49m\u001b[38;5;132;43;01m}\u001b[39;49;00m\u001b[38;5;124;43m_sent_emo.csv\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m      4\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m i, row \u001b[38;5;129;01min\u001b[39;00m df\u001b[38;5;241m.\u001b[39miterrows():\n\u001b[1;32m      5\u001b[0m         dia_id \u001b[38;5;241m=\u001b[39m row[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mDialogue_ID\u001b[39m\u001b[38;5;124m'\u001b[39m]\n",
      "File \u001b[0;32m/opt/anaconda3/envs/audio/lib/python3.12/site-packages/pandas/io/parsers/readers.py:1026\u001b[0m, in \u001b[0;36mread_csv\u001b[0;34m(filepath_or_buffer, sep, delimiter, header, names, index_col, usecols, dtype, engine, converters, true_values, false_values, skipinitialspace, skiprows, skipfooter, nrows, na_values, keep_default_na, na_filter, verbose, skip_blank_lines, parse_dates, infer_datetime_format, keep_date_col, date_parser, date_format, dayfirst, cache_dates, iterator, chunksize, compression, thousands, decimal, lineterminator, quotechar, quoting, doublequote, escapechar, comment, encoding, encoding_errors, dialect, on_bad_lines, delim_whitespace, low_memory, memory_map, float_precision, storage_options, dtype_backend)\u001b[0m\n\u001b[1;32m   1013\u001b[0m kwds_defaults \u001b[38;5;241m=\u001b[39m _refine_defaults_read(\n\u001b[1;32m   1014\u001b[0m     dialect,\n\u001b[1;32m   1015\u001b[0m     delimiter,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1022\u001b[0m     dtype_backend\u001b[38;5;241m=\u001b[39mdtype_backend,\n\u001b[1;32m   1023\u001b[0m )\n\u001b[1;32m   1024\u001b[0m kwds\u001b[38;5;241m.\u001b[39mupdate(kwds_defaults)\n\u001b[0;32m-> 1026\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_read\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfilepath_or_buffer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwds\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/anaconda3/envs/audio/lib/python3.12/site-packages/pandas/io/parsers/readers.py:620\u001b[0m, in \u001b[0;36m_read\u001b[0;34m(filepath_or_buffer, kwds)\u001b[0m\n\u001b[1;32m    617\u001b[0m _validate_names(kwds\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mnames\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28;01mNone\u001b[39;00m))\n\u001b[1;32m    619\u001b[0m \u001b[38;5;66;03m# Create the parser.\u001b[39;00m\n\u001b[0;32m--> 620\u001b[0m parser \u001b[38;5;241m=\u001b[39m \u001b[43mTextFileReader\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfilepath_or_buffer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwds\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    622\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m chunksize \u001b[38;5;129;01mor\u001b[39;00m iterator:\n\u001b[1;32m    623\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m parser\n",
      "File \u001b[0;32m/opt/anaconda3/envs/audio/lib/python3.12/site-packages/pandas/io/parsers/readers.py:1620\u001b[0m, in \u001b[0;36mTextFileReader.__init__\u001b[0;34m(self, f, engine, **kwds)\u001b[0m\n\u001b[1;32m   1617\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moptions[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhas_index_names\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwds[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhas_index_names\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n\u001b[1;32m   1619\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandles: IOHandles \u001b[38;5;241m|\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m-> 1620\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_engine \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_make_engine\u001b[49m\u001b[43m(\u001b[49m\u001b[43mf\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mengine\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/anaconda3/envs/audio/lib/python3.12/site-packages/pandas/io/parsers/readers.py:1880\u001b[0m, in \u001b[0;36mTextFileReader._make_engine\u001b[0;34m(self, f, engine)\u001b[0m\n\u001b[1;32m   1878\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mb\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m mode:\n\u001b[1;32m   1879\u001b[0m         mode \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mb\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m-> 1880\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandles \u001b[38;5;241m=\u001b[39m \u001b[43mget_handle\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1881\u001b[0m \u001b[43m    \u001b[49m\u001b[43mf\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1882\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmode\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1883\u001b[0m \u001b[43m    \u001b[49m\u001b[43mencoding\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moptions\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mencoding\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1884\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcompression\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moptions\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mcompression\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1885\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmemory_map\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moptions\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mmemory_map\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1886\u001b[0m \u001b[43m    \u001b[49m\u001b[43mis_text\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mis_text\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1887\u001b[0m \u001b[43m    \u001b[49m\u001b[43merrors\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moptions\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mencoding_errors\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mstrict\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1888\u001b[0m \u001b[43m    \u001b[49m\u001b[43mstorage_options\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moptions\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mstorage_options\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1889\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1890\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandles \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m   1891\u001b[0m f \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandles\u001b[38;5;241m.\u001b[39mhandle\n",
      "File \u001b[0;32m/opt/anaconda3/envs/audio/lib/python3.12/site-packages/pandas/io/common.py:873\u001b[0m, in \u001b[0;36mget_handle\u001b[0;34m(path_or_buf, mode, encoding, compression, memory_map, is_text, errors, storage_options)\u001b[0m\n\u001b[1;32m    868\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(handle, \u001b[38;5;28mstr\u001b[39m):\n\u001b[1;32m    869\u001b[0m     \u001b[38;5;66;03m# Check whether the filename is to be opened in binary mode.\u001b[39;00m\n\u001b[1;32m    870\u001b[0m     \u001b[38;5;66;03m# Binary mode does not support 'encoding' and 'newline'.\u001b[39;00m\n\u001b[1;32m    871\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m ioargs\u001b[38;5;241m.\u001b[39mencoding \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mb\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m ioargs\u001b[38;5;241m.\u001b[39mmode:\n\u001b[1;32m    872\u001b[0m         \u001b[38;5;66;03m# Encoding\u001b[39;00m\n\u001b[0;32m--> 873\u001b[0m         handle \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mopen\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m    874\u001b[0m \u001b[43m            \u001b[49m\u001b[43mhandle\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    875\u001b[0m \u001b[43m            \u001b[49m\u001b[43mioargs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmode\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    876\u001b[0m \u001b[43m            \u001b[49m\u001b[43mencoding\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mioargs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mencoding\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    877\u001b[0m \u001b[43m            \u001b[49m\u001b[43merrors\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43merrors\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    878\u001b[0m \u001b[43m            \u001b[49m\u001b[43mnewline\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m    879\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    880\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    881\u001b[0m         \u001b[38;5;66;03m# Binary mode\u001b[39;00m\n\u001b[1;32m    882\u001b[0m         handle \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mopen\u001b[39m(handle, ioargs\u001b[38;5;241m.\u001b[39mmode)\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: 'data/test/test_sent_emo.csv'"
     ]
    }
   ],
   "source": [
    "# convert audio to fbank from train, dev and test\n",
    "for mode in ['train', 'dev', 'test']:\n",
    "    df = pd.read_csv(f'data/{mode}/{mode}_sent_emo.csv')\n",
    "    for i, row in df.iterrows():\n",
    "        dia_id = row['Dialogue_ID']\n",
    "        utt_id = row['Utterance_ID']\n",
    "        audio_file = f'data/{mode}/audio/dia{dia_id}_utt{utt_id}.wav'\n",
    "        fbank_file = f'data/{mode}/fbank/dia{dia_id}_utt{utt_id}.npy'\n",
    "        \n",
    "        # create the fbank directory if it does not exist\n",
    "        if not os.path.exists(f'data/{mode}/fbank/'):\n",
    "            os.makedirs(f'data/{mode}/fbank/')\n",
    "        \n",
    "        # check if the fbank file already exists\n",
    "        if not os.path.exists(f'data/{mode}/fbank/dia{dia_id}_utt{utt_id}.npy'):\n",
    "            fbank = wav2fbank(audio_file)\n",
    "            fbank = fbank.numpy()\n",
    "            np.save(f'data/{mode}/fbank/dia{dia_id}_utt{utt_id}.npy', fbank)\n",
    "        else:\n",
    "            print(f'{fbank_file} already exists')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataset and dataloader class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-01-07T23:08:42.844157Z",
     "iopub.status.busy": "2025-01-07T23:08:42.842376Z",
     "iopub.status.idle": "2025-01-07T23:08:42.857565Z",
     "shell.execute_reply": "2025-01-07T23:08:42.856552Z",
     "shell.execute_reply.started": "2025-01-07T23:08:42.844115Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "class MELDConversationDataset(Dataset):\n",
    "    def __init__(self, csv_file, root_dir='./data', mode=\"train\", text_processor=None):\n",
    "        \"\"\"\n",
    "        We'll store a list of (dialog_id, [list_of_utterance_dicts]).\n",
    "        Each utterance_dict might contain:\n",
    "          {\n",
    "            \"fbank_path\": str,\n",
    "            \"transcript\": str,\n",
    "            \"emotion\": int,\n",
    "            \"sentiment\": int\n",
    "          }\n",
    "        \"\"\"\n",
    "        df = pd.read_csv(f'{root_dir}/{csv_file}')\n",
    "\n",
    "        # order the df rows according to dialogueID and each dialogue according to utteranceID\n",
    "        df = df.sort_values(by=['Dialogue_ID', 'Utterance_ID'])\n",
    "\n",
    "        self.emotion_class_counts = {0: 0, 1: 0, 2: 0, 3: 0, 4: 0, 5: 0, 6: 0}\n",
    "        self.sentiment_class_counts = {0: 0, 1: 0, 2: 0}\n",
    "        self.max_dialogue_size = 0\n",
    "\n",
    "        self.dialogues = {}  # key: dialogue_id, value: list of utterance dicts\n",
    "        prev_dia_id = None\n",
    "        utt_count = 0\n",
    "\n",
    "        for _, row in df.iterrows():\n",
    "\n",
    "            dia_id = row[\"Dialogue_ID\"]\n",
    "            utt_id = row[\"Utterance_ID\"]\n",
    "\n",
    "            if prev_dia_id == dia_id:\n",
    "                utt_count += 1\n",
    "            else:\n",
    "                if utt_count > self.max_dialogue_size:\n",
    "                    self.max_dialogue_size = utt_count\n",
    "                utt_count = 1\n",
    "\n",
    "            fbank_path = f'data/{mode}/fbank/dia{dia_id}_utt{utt_id}.npy'\n",
    "\n",
    "            emotion = row[\"Emotion\"]\n",
    "            emotion = emotion.lower()\n",
    "            emotion_int = self.emotion_to_int(emotion)\n",
    "            sentiment = row[\"Sentiment\"]\n",
    "            sentiment = sentiment.lower()\n",
    "            sentiment_int = self.sentiment_to_int(sentiment)\n",
    "\n",
    "            self.emotion_class_counts[emotion_int] += 1\n",
    "            self.sentiment_class_counts[sentiment_int] += 1\n",
    "\n",
    "            utter_dict = {\n",
    "                \"fbank_path\": fbank_path,\n",
    "                \"transcript\": row[\"Utterance\"],\n",
    "                \"emotion\": emotion_int,\n",
    "                \"sentiment\": sentiment_int\n",
    "            }\n",
    "\n",
    "            if dia_id not in self.dialogues:\n",
    "                self.dialogues[dia_id] = []\n",
    "            self.dialogues[dia_id].append(utter_dict)\n",
    "\n",
    "            prev_dia_id = dia_id\n",
    "\n",
    "        # Convert to list of (dialog_id, list_of_utterances)\n",
    "        self.dialogues = [(k, sorted(v, key=lambda x: x[\"fbank_path\"]))\n",
    "                          for k, v in self.dialogues.items()]\n",
    "        # The sorting step ensures the utterances are in ascending order of utt_id if needed.\n",
    "\n",
    "    def emotion_to_int(self, str):\n",
    "        str_to_int = {\"neutral\": 0, \"joy\": 1, \"surprise\": 2,\n",
    "                      \"anger\": 3, \"sadness\": 4, \"fear\": 5, \"disgust\": 6}\n",
    "        return str_to_int[str]\n",
    "\n",
    "    def emotion_to_str(self, int):\n",
    "        int_to_str = {0: \"neutral\", 1: \"joy\", 2: \"surprise\",\n",
    "                      3: \"anger\", 4: \"sadness\", 5: \"fear\", 6: \"disgust\"}\n",
    "        return int_to_str[int]\n",
    "\n",
    "    def sentiment_to_int(self, str):\n",
    "        str_to_int = {\"neutral\": 0, \"positive\": 1, \"negative\": 2}\n",
    "        return str_to_int[str]\n",
    "\n",
    "    def sentiment_to_str(self, int):\n",
    "        int_to_str = {0: \"neutral\", 1: \"positive\", 2: \"negative\"}\n",
    "        return int_to_str[int]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.dialogues)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        dialog_id, utterances = self.dialogues[idx]\n",
    "\n",
    "        # For each utterance, load filterbanks, transcript, emotion, sentiment\n",
    "        fbank_list = []\n",
    "        text_list = []\n",
    "        emotion_list = []\n",
    "        sentiment_list = []\n",
    "\n",
    "        for utt in utterances:\n",
    "            # shape e.g. (T, fbank_dim)\n",
    "            fbank = np.load(utt[\"fbank_path\"])\n",
    "            fbank_list.append(fbank)\n",
    "\n",
    "            text_list.append(utt[\"transcript\"])\n",
    "            emotion_list.append(utt[\"emotion\"])         # or mapped to int\n",
    "            sentiment_list.append(utt[\"sentiment\"])     # or mapped to int\n",
    "\n",
    "        return {\n",
    "            \"dialog_id\": dialog_id,\n",
    "            \"fbank_list\": fbank_list,\n",
    "            \"text_list\": text_list,\n",
    "            \"emotion_list\": emotion_list,\n",
    "            \"sentiment_list\": sentiment_list\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-01-07T23:08:42.859145Z",
     "iopub.status.busy": "2025-01-07T23:08:42.858750Z",
     "iopub.status.idle": "2025-01-07T23:08:58.746310Z",
     "shell.execute_reply": "2025-01-07T23:08:58.745266Z",
     "shell.execute_reply.started": "2025-01-07T23:08:42.859108Z"
    },
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data loaded.\n",
      "Data loaders created.\n"
     ]
    }
   ],
   "source": [
    "train_set = MELDConversationDataset(csv_file=\"train_sent_emo.csv\", root_dir=\"data/train\", mode=\"train\")\n",
    "dev_set = MELDConversationDataset(csv_file=\"dev_sent_emo.csv\", root_dir=\"data/dev\", mode=\"dev\")\n",
    "test_set = MELDConversationDataset(csv_file='test_sent_emo.csv', root_dir='data/test', mode='test')\n",
    "\n",
    "print(\"Data loaded.\")\n",
    "\n",
    "train_loader = DataLoader(train_set, batch_size=4, shuffle=True, num_workers=4, collate_fn=meld_collate_fn)\n",
    "dev_loader = DataLoader(dev_set, batch_size=4, shuffle=True, num_workers=4, collate_fn=meld_collate_fn)\n",
    "test_loader = DataLoader(test_set, batch_size=4, shuffle=True, num_workers=4, collate_fn=meld_collate_fn)\n",
    "print(\"Data loaders created.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Miscs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1038 114 169\n"
     ]
    }
   ],
   "source": [
    "print(len(train_set), len(dev_set), len(test_set))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-01-06T16:51:34.622185Z",
     "iopub.status.busy": "2025-01-06T16:51:34.621837Z",
     "iopub.status.idle": "2025-01-06T16:51:34.628399Z",
     "shell.execute_reply": "2025-01-06T16:51:34.627441Z",
     "shell.execute_reply.started": "2025-01-06T16:51:34.622152Z"
    },
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of emotions: 7\n",
      "Number of sentiments: 3\n",
      "Emotion neutral count: 4709\n",
      "Emotion joy count: 1743\n",
      "Emotion surprise count: 1205\n",
      "Emotion anger count: 1109\n",
      "Emotion sadness count: 683\n",
      "Emotion fear count: 268\n",
      "Emotion disgust count: 271\n",
      "Sentiment neutral count: 4709\n",
      "Sentiment positive count: 2334\n",
      "Sentiment negative count: 2945\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "num_emotions = len(train_set.emotion_class_counts)\n",
    "num_sentiments = len(train_set.sentiment_class_counts)\n",
    "\n",
    "print(\"Number of emotions:\", num_emotions)\n",
    "print(\"Number of sentiments:\", num_sentiments)\n",
    "\n",
    "for i in range(7):\n",
    "    print(f\"Emotion {train_set.emotion_to_str(i)} count:\", train_set.emotion_class_counts[i])\n",
    "for i in range(3):\n",
    "    print(f\"Sentiment {train_set.sentiment_to_str(i)} count:\", train_set.sentiment_class_counts[i])\n",
    "    \n",
    "print(\"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total samples: 9988\n",
      "Class weights: tensor([0.0186, 0.0503, 0.0728, 0.0791, 0.1284, 0.3272, 0.3236])\n"
     ]
    }
   ],
   "source": [
    "#get weights for balancing classes\n",
    "class_counts = train_set.emotion_class_counts\n",
    "total_samples = 0\n",
    "\n",
    "for key in class_counts:\n",
    "    total_samples += class_counts[key]\n",
    "    \n",
    "print(\"Total samples:\", total_samples)\n",
    "class_weights = torch.zeros(len(class_counts))\n",
    "\n",
    "for i in range(len(class_counts)):\n",
    "    class_weights[i] = class_counts[i] / total_samples\n",
    "    \n",
    "class_weights = 1 / class_weights  # invert the weights\n",
    "class_weights = class_weights / class_weights.sum()  # normalize the weights\n",
    "class_weights = class_weights.to(device)\n",
    "print(\"Class weights:\", class_weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Max dialogue size: 24\n"
     ]
    }
   ],
   "source": [
    "# PARAMETERS\n",
    "lr = 1e-4\n",
    "criterions = {\n",
    "    'emotion': nn.CrossEntropyLoss(weight=class_weights),\n",
    "    'sentiment': nn.CrossEntropyLoss()\n",
    "}\n",
    "\n",
    "num_epochs = 50\n",
    "\n",
    "#set max utterance size to be the max of all train dev and test sets\n",
    "max_utt = max(train_set.max_dialogue_size, dev_set.max_dialogue_size, test_set.max_dialogue_size)\n",
    "print(\"Max dialogue size:\", max_utt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):\n",
    "    \"\"\"3x3 convolution with padding\"\"\"\n",
    "    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,\n",
    "                     padding=dilation, groups=groups, bias=False, dilation=dilation)\n",
    "\n",
    "\n",
    "def conv1x1(in_planes, out_planes, stride=1):\n",
    "    \"\"\"1x1 convolution\"\"\"\n",
    "    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)\n",
    "\n",
    "\n",
    "class BasicBlock(nn.Module):\n",
    "    expansion = 1\n",
    "\n",
    "    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,\n",
    "                 base_width=64, dilation=1, norm_layer=None, dropout_p=0.5):\n",
    "        super(BasicBlock, self).__init__()\n",
    "        if norm_layer is None:\n",
    "            norm_layer = nn.BatchNorm2d\n",
    "        if groups != 1 or base_width != 64:\n",
    "            raise ValueError(\n",
    "                'BasicBlock only supports groups=1 and base_width=64')\n",
    "        if dilation > 1:\n",
    "            raise NotImplementedError(\n",
    "                \"Dilation > 1 not supported in BasicBlock\")\n",
    "\n",
    "        self.conv1 = conv3x3(inplanes, planes, stride)\n",
    "        self.bn1 = norm_layer(planes)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.dropout = nn.Dropout(p=dropout_p)  # Dropout layer added\n",
    "        self.conv2 = conv3x3(planes, planes)\n",
    "        self.bn2 = norm_layer(planes)\n",
    "        self.downsample = downsample\n",
    "        self.stride = stride\n",
    "\n",
    "    def forward(self, x):\n",
    "        identity = x\n",
    "\n",
    "        out = self.conv1(x)\n",
    "        out = self.bn1(out)\n",
    "        out = self.relu(out)\n",
    "        out = self.dropout(out)  # apply dropout after first ReLU\n",
    "\n",
    "        out = self.conv2(out)\n",
    "        out = self.bn2(out)\n",
    "\n",
    "        if self.downsample is not None:\n",
    "            identity = self.downsample(x)\n",
    "\n",
    "        out += identity\n",
    "        out = self.relu(out)\n",
    "        out = self.dropout(out)  # apply dropout after second ReLU\n",
    "\n",
    "        return out\n",
    "\n",
    "\n",
    "class ResNet(nn.Module):\n",
    "\n",
    "    def __init__(self, block, layers, modality, num_classes=1000, pool='avgpool', zero_init_residual=False,\n",
    "                 groups=1, width_per_group=64, replace_stride_with_dilation=None,\n",
    "                 norm_layer=None):\n",
    "        super(ResNet, self).__init__()\n",
    "        self.modality = modality\n",
    "        self.pool = pool\n",
    "        if norm_layer is None:\n",
    "            norm_layer = nn.BatchNorm2d\n",
    "        self._norm_layer = norm_layer\n",
    "\n",
    "        self.inplanes = 64\n",
    "        self.dilation = 1\n",
    "        if replace_stride_with_dilation is None:\n",
    "            # each element in the tuple indicates if we should replace\n",
    "            # the 2x2 stride with a dilated convolution instead\n",
    "            replace_stride_with_dilation = [False, False, False]\n",
    "        if len(replace_stride_with_dilation) != 3:\n",
    "            raise ValueError(\"replace_stride_with_dilation should be None \"\n",
    "                             \"or a 3-element tuple, got {}\".format(replace_stride_with_dilation))\n",
    "        self.groups = groups\n",
    "        self.base_width = width_per_group\n",
    "        if modality == 'audio':\n",
    "            self.conv1 = nn.Conv2d(1, self.inplanes, kernel_size=7, stride=2, padding=3,\n",
    "                                   bias=False)\n",
    "        elif modality == 'visual':\n",
    "            self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3,\n",
    "                                   bias=False)\n",
    "        else:\n",
    "            raise NotImplementedError(\n",
    "                'Incorrect modality, should be audio or visual but got {}'.format(modality))\n",
    "\n",
    "        self.out_conv = nn.Conv2d(\n",
    "            512 * block.expansion, 768, kernel_size=1, stride=1, bias=False)\n",
    "\n",
    "        self.bn1 = norm_layer(self.inplanes)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n",
    "        self.layer1 = self._make_layer(block, 64, layers[0])\n",
    "        self.layer2 = self._make_layer(block, 128, layers[1], stride=2,\n",
    "                                       dilate=replace_stride_with_dilation[0])\n",
    "        self.layer3 = self._make_layer(block, 256, layers[2], stride=2,\n",
    "                                       dilate=replace_stride_with_dilation[1])\n",
    "        self.layer4 = self._make_layer(block, 512, layers[3], stride=2,\n",
    "                                       dilate=replace_stride_with_dilation[2])\n",
    "\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Conv2d):\n",
    "                nn.init.kaiming_normal_(\n",
    "                    m.weight, mode='fan_out', nonlinearity='relu')\n",
    "            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):\n",
    "                nn.init.normal_(m.weight, mean=1, std=0.02)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "\n",
    "        # Zero-initialize the last BN in each residual branch,\n",
    "        # so that the residual branch starts with zeros, and each residual block behaves like an identity.\n",
    "        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677\n",
    "        if zero_init_residual:\n",
    "            for m in self.modules():\n",
    "                if isinstance(m, Bottleneck):\n",
    "                    nn.init.constant_(m.bn3.weight, 0)\n",
    "                elif isinstance(m, BasicBlock):\n",
    "                    nn.init.constant_(m.bn2.weight, 0)\n",
    "\n",
    "    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):\n",
    "        norm_layer = self._norm_layer\n",
    "        downsample = None\n",
    "        previous_dilation = self.dilation\n",
    "        if dilate:\n",
    "            self.dilation *= stride\n",
    "            stride = 1\n",
    "        if stride != 1 or self.inplanes != planes * block.expansion:\n",
    "            downsample = nn.Sequential(\n",
    "                conv1x1(self.inplanes, planes * block.expansion, stride),\n",
    "                norm_layer(planes * block.expansion),\n",
    "            )\n",
    "\n",
    "        layers = []\n",
    "        layers.append(block(self.inplanes, planes, stride, downsample, self.groups,\n",
    "                            self.base_width, previous_dilation, norm_layer))\n",
    "        self.inplanes = planes * block.expansion\n",
    "        for _ in range(1, blocks):\n",
    "            layers.append(block(self.inplanes, planes, groups=self.groups,\n",
    "                                base_width=self.base_width, dilation=self.dilation,\n",
    "                                norm_layer=norm_layer))\n",
    "\n",
    "        return nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "\n",
    "        if self.modality == 'visual':\n",
    "            (B, C, T, H, W) = x.size()\n",
    "            x = x.permute(0, 2, 1, 3, 4).contiguous()\n",
    "            x = x.view(B * T, C, H, W)\n",
    "\n",
    "        x = self.conv1(x)\n",
    "        x = self.bn1(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.maxpool(x)\n",
    "\n",
    "        x = self.layer1(x)\n",
    "        x = self.layer2(x)\n",
    "        x = self.layer3(x)\n",
    "        x = self.layer4(x)\n",
    "\n",
    "        x = self.out_conv(x)  # Now x has shape [B, 768, H_out, W_out]\n",
    "\n",
    "        out = x\n",
    "\n",
    "        return out\n",
    "\n",
    "\n",
    "class Bottleneck(nn.Module):\n",
    "    expansion = 4\n",
    "\n",
    "    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,\n",
    "                 base_width=64, dilation=1, norm_layer=None):\n",
    "        super(Bottleneck, self).__init__()\n",
    "        if norm_layer is None:\n",
    "            norm_layer = nn.BatchNorm2d\n",
    "        width = int(planes * (base_width / 64.)) * groups\n",
    "        # Both self.conv2 and self.downsample layers downsample the input when stride != 1\n",
    "        self.conv1 = conv1x1(inplanes, width)\n",
    "        self.bn1 = norm_layer(width)\n",
    "        self.conv2 = conv3x3(width, width, stride, groups, dilation)\n",
    "        self.bn2 = norm_layer(width)\n",
    "        self.conv3 = conv1x1(width, planes * self.expansion)\n",
    "        self.bn3 = norm_layer(planes * self.expansion)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.downsample = downsample\n",
    "        self.stride = stride\n",
    "\n",
    "    def forward(self, x):\n",
    "        identity = x\n",
    "\n",
    "        out = self.conv1(x)\n",
    "        out = self.bn1(out)\n",
    "        out = self.relu(out)\n",
    "\n",
    "        out = self.conv2(out)\n",
    "        out = self.bn2(out)\n",
    "        out = self.relu(out)\n",
    "\n",
    "        out = self.conv3(out)\n",
    "        out = self.bn3(out)\n",
    "\n",
    "        if self.downsample is not None:\n",
    "            identity = self.downsample(x)\n",
    "\n",
    "        out += identity\n",
    "        out = self.relu(out)\n",
    "\n",
    "        return out\n",
    "\n",
    "\n",
    "def _resnet(arch, block, layers, modality, progress, **kwargs):\n",
    "    model = ResNet(block, layers, modality, **kwargs)\n",
    "    return model\n",
    "\n",
    "\n",
    "def resnet18(modality, progress=True, **kwargs):\n",
    "    return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], modality, progress,\n",
    "                   **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AudioEncoder(nn.Module):\n",
    "    def __init__(self, model_name=\"facebook/wav2vec2-base\",\n",
    "                 target_sr=16000,\n",
    "                 fine_tune=False,\n",
    "                 unfreeze_last_n=2,  # Number of last layers to unfreeze\n",
    "                 device=device):\n",
    "        super(AudioEncoder, self).__init__()\n",
    "        self.processor = AutoProcessor.from_pretrained(model_name)\n",
    "        self.model = AutoModel.from_pretrained(model_name).to(device)\n",
    "        self.target_sr = target_sr\n",
    "\n",
    "        # Freeze entire model\n",
    "        for param in self.model.parameters():\n",
    "            param.requires_grad = False\n",
    "\n",
    "        if fine_tune:\n",
    "            # Unfreeze only the last `unfreeze_last_n` encoder layers\n",
    "            total_layers = len(self.model.encoder.layers)\n",
    "            for layer_idx in range(total_layers - unfreeze_last_n, total_layers):\n",
    "                for param in self.model.encoder.layers[layer_idx].parameters():\n",
    "                    param.requires_grad = True\n",
    "\n",
    "    def forward(self, waveforms):\n",
    "        \"\"\"\n",
    "        :param waveforms: Tensor of shape [B, T] (already at self.target_sr)\n",
    "        :return: A tensor of shape [B, hidden_dim] (audio embeddings for the batch)\n",
    "        \"\"\"\n",
    "        # Ensure the waveforms are on the correct device\n",
    "        waveforms = waveforms.to(self.model.device)\n",
    "\n",
    "        # Prepare inputs for Wav2Vec2Processor\n",
    "        inputs = self.processor(\n",
    "            waveforms.cpu(),\n",
    "            sampling_rate=self.target_sr,\n",
    "            return_tensors=\"pt\",\n",
    "            padding=True\n",
    "        ).input_values.squeeze((0, 1)).to(self.model.device)\n",
    "\n",
    "        # Forward pass\n",
    "        with torch.no_grad() if not self.training or not any(\n",
    "            p.requires_grad for p in self.model.parameters()\n",
    "        ) else torch.enable_grad():\n",
    "            outputs = self.model(inputs)\n",
    "            hidden_states = outputs.last_hidden_state  # shape [B, T, D]\n",
    "\n",
    "        # Average pooling\n",
    "        audio_emb = hidden_states.mean(dim=1)  # shape [B, D]\n",
    "\n",
    "        return audio_emb\n",
    "\n",
    "\n",
    "class TextEncoder(nn.Module):\n",
    "    \"\"\"\n",
    "    Encodes text using a pretrained BERT model from Hugging Face.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self,\n",
    "                 model_name=\"distilbert-base-uncased\",\n",
    "                 fine_tune=False,\n",
    "                 unfreeze_last_n_layers=2,\n",
    "                 device=device):\n",
    "        super(TextEncoder, self).__init__()\n",
    "        # self.tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "        self.tokenizer = DistilBertTokenizer.from_pretrained(model_name)\n",
    "        self.model = AutoModel.from_pretrained(model_name).to(device)\n",
    "\n",
    "        # If we're not fine-tuning at all, freeze everything.\n",
    "        # Otherwise, freeze everything first, then selectively unfreeze layers.\n",
    "        for param in self.model.parameters():\n",
    "            param.requires_grad = False\n",
    "\n",
    "        if fine_tune and unfreeze_last_n_layers > 0:\n",
    "            # For BERT-base, there are 12 encoder layers: encoder.layer[0] ... encoder.layer[11].\n",
    "            # Unfreeze the last N layers:\n",
    "            total_layers = len(self.model.transformer.layer)\n",
    "            for layer_idx in range(total_layers - unfreeze_last_n_layers, total_layers):\n",
    "                for param in self.model.transformer.layer[layer_idx].parameters():\n",
    "                    param.requires_grad = True\n",
    "\n",
    "            # Optionally unfreeze the pooler layer if you use it\n",
    "            # for param in self.model.pooler.parameters():\n",
    "             #   param.requires_grad = True\n",
    "\n",
    "    def forward(self, text_list):\n",
    "        \"\"\"\n",
    "        :param text_list: A list of strings (or a single string) to encode.\n",
    "        :return: A tensor of shape [batch_size, hidden_dim] with text embeddings\n",
    "        \"\"\"\n",
    "        device = self.model.device\n",
    "\n",
    "        # If a single string is passed, wrap it into a list\n",
    "        if isinstance(text_list, str):\n",
    "            text_list = [text_list]\n",
    "\n",
    "        encodings = self.tokenizer(\n",
    "            text_list,\n",
    "            padding=True,\n",
    "            truncation=True,\n",
    "            return_tensors=\"pt\"\n",
    "        ).to(device)\n",
    "\n",
    "        # If all parameters are frozen (no grad), then no_grad() is fine.\n",
    "        # But if some layers are unfrozen, we want torch.enable_grad()\n",
    "        # so that backprop can proceed for those layers.\n",
    "        use_grad = any(p.requires_grad for p in self.model.parameters())\n",
    "        with torch.enable_grad() if use_grad else torch.no_grad():\n",
    "            outputs = self.model(\n",
    "                input_ids=encodings.input_ids,\n",
    "                attention_mask=encodings.attention_mask\n",
    "            )\n",
    "\n",
    "        # outputs.last_hidden_state -> shape [batch_size, seq_len, hidden_dim]\n",
    "        # Typically we use the [CLS] token embedding as a single representation\n",
    "        cls_emb = outputs.last_hidden_state[:, 0, :]  # shape [B, hidden_dim]\n",
    "        return cls_emb\n",
    "\n",
    "\n",
    "class FusionLSTM(nn.Module):\n",
    "    def __init__(self, input_dim, hidden_dim=256, num_layers=1, bidirectional=False, dropout=0.1):\n",
    "        super(FusionLSTM, self).__init__()\n",
    "        self.lstm = nn.LSTM(\n",
    "            input_size=input_dim,\n",
    "            hidden_size=hidden_dim,\n",
    "            num_layers=num_layers,\n",
    "            batch_first=True,\n",
    "            bidirectional=bidirectional,\n",
    "            dropout=dropout if num_layers > 1 else 0.0\n",
    "        )\n",
    "        self.hidden_dim = hidden_dim * (2 if bidirectional else 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        x: (batch_size, seq_len, input_dim)\n",
    "        returns: (batch_size, seq_len, hidden_dim)\n",
    "        \"\"\"\n",
    "        output, (hn, cn) = self.lstm(x)\n",
    "        return output, (hn, cn)  # output is the sequence of hidden states\n",
    "\n",
    "\n",
    "class MultimodalClassifierWithLSTM(nn.Module):\n",
    "    def __init__(self, audio_encoder, text_encoder, fusion_lstm, hidden_dim=256, num_emotions=7, num_sentiments=3, max_utt=50):\n",
    "        super(MultimodalClassifierWithLSTM, self).__init__()\n",
    "        self.audio_encoder = audio_encoder\n",
    "        self.text_encoder = text_encoder\n",
    "        self.fusion_lstm = fusion_lstm  # e.g., an instance of FusionLSTM\n",
    "        self.emotion_head = nn.Linear(hidden_dim*2, num_emotions)\n",
    "        self.sentiment_head = nn.Linear(hidden_dim*2, num_sentiments)\n",
    "        self.max_utt = max_utt\n",
    "\n",
    "    def forward(self, audio_dialog_utts, text_dialog_utts):\n",
    "\n",
    "        audio_emb = self.audio_encoder(audio_dialog_utts)\n",
    "        audio_emb = F.adaptive_avg_pool2d(audio_emb, 1)\n",
    "        audio_emb = torch.flatten(audio_emb, 1)\n",
    "\n",
    "        text_emb = [self.text_encoder(utt) for utt in text_dialog_utts]\n",
    "        text_emb = torch.stack(text_emb)  # shape (B, hidden_dim)\n",
    "\n",
    "        text_emb = text_emb.squeeze(1)\n",
    "\n",
    "        # print(\"Final Audio emb shape: \", audio_emb.shape)\n",
    "        # print(\"Final Text emb shape: \", text_emb.shape)\n",
    "\n",
    "        # Combine\n",
    "        # (utts, audio_enc_dim + text_enc_dim)\n",
    "        fused_emb = torch.cat([audio_emb, text_emb], dim=-1)\n",
    "\n",
    "        # Reshape back to (B, S, fused_dim) for the LSTM\n",
    "        # (1, utts, audio_enc_dim + text_enc_dim)\n",
    "        fused_emb = fused_emb.unsqueeze(0)\n",
    "\n",
    "        padded_fused_emb = F.pad(\n",
    "            fused_emb, (0, 0, 0, self.max_utt - fused_emb.size(1)))\n",
    "\n",
    "        # Pass through LSTM\n",
    "        lstm_out, (hn, cn) = self.fusion_lstm(padded_fused_emb)\n",
    "\n",
    "        # print(\"LSTM out shape: \", lstm_out.shape)\n",
    "\n",
    "        lstm_out = lstm_out.squeeze(0)\n",
    "\n",
    "        # Classification heads\n",
    "\n",
    "        emotion_logits = self.emotion_head(lstm_out)\n",
    "        sentiment_logits = self.sentiment_head(lstm_out)\n",
    "\n",
    "        return emotion_logits, sentiment_logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MultimodalClassifierWithLSTM(\n",
    "    fusion_lstm=FusionLSTM(input_dim=1536, hidden_dim=256, num_layers=1, bidirectional=True, dropout=0.2),\n",
    "    audio_encoder=resnet18(modality='audio'),\n",
    "    text_encoder=TextEncoder(fine_tune=True, unfreeze_last_n_layers=1),\n",
    "    hidden_dim=256,\n",
    "    num_emotions=num_emotions,\n",
    "    num_sentiments=num_sentiments,\n",
    "    max_utt=max_utt\n",
    ").to(device)\n",
    "\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.0001)\n",
    "experiment_name = 'TEST'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### utils functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from sklearn.metrics import confusion_matrix, classification_report, balanced_accuracy_score, roc_curve, auc\n",
    "from sklearn.preprocessing import label_binarize\n",
    "\n",
    "def compute_metrics(true_labels, predictions):\n",
    "    \"\"\"\n",
    "    Compute classification metrics including accuracy, per-class F1 scores, and weighted average F1 score.\n",
    "    \"\"\"\n",
    "    # Compute overall accuracy\n",
    "    b_accuracy = round(balanced_accuracy_score(true_labels, predictions), 3)   # Compute F1 scores\n",
    "    report = classification_report(\n",
    "        true_labels, predictions, output_dict=True, zero_division=0\n",
    "    )\n",
    "    \n",
    "    per_class_f1 = {label: round(values[\"f1-score\"], 3) for label, values in report.items() if label.isdigit()}\n",
    "    macro_f1 = round(report[\"macro avg\"][\"f1-score\"], 3)\n",
    "\n",
    "    # Compile metrics into a dictionary\n",
    "    metrics = {\n",
    "        \"balanced_acc\": round(b_accuracy, 3),\n",
    "        \"macro_f1\": macro_f1,\n",
    "        \"per_class_f1\": per_class_f1\n",
    "    }\n",
    "\n",
    "    return metrics\n",
    "\n",
    "def task_result_to_table(task_result):\n",
    "    df = pd.DataFrame([\n",
    "        {\n",
    "            'epoch': res['epoch'],\n",
    "            'train_loss': res['train_loss'],\n",
    "            'val_loss': res['val_loss'],\n",
    "            **flatten_metrics(res['train_metrics'], prefix='train'),\n",
    "            **flatten_metrics(res['val_metrics'], prefix='val')\n",
    "        }\n",
    "        for res in task_result\n",
    "    ])\n",
    "    \n",
    "    return df\n",
    "\n",
    "def flatten_metrics(metrics, prefix=''):\n",
    "    \"\"\"\n",
    "    Flatten metrics for easier storage in a tabular format.\n",
    "\n",
    "    :param metrics: Dictionary of metrics (e.g., {'fused': {...}, 'audio': {...}, 'text': {...}})\n",
    "    :param prefix: Prefix for column names (optional, e.g., 'train' or 'val').\n",
    "    :return: Flattened dictionary.\n",
    "    \"\"\"\n",
    "    flattened = {}\n",
    "    for modality, modality_metrics in metrics.items():\n",
    "        for metric_name, value in modality_metrics.items():\n",
    "            flattened[f\"{prefix}_{modality}_{metric_name}\"] = value\n",
    "    return flattened\n",
    "\n",
    "def plot_metrics(results, task='emotions', metric='acc', modality='fused'):\n",
    "    \"\"\"\n",
    "    Plots metrics (e.g., accuracy or F1 score) for a specific task and modality over epochs.\n",
    "    \n",
    "    :param results: Dictionary containing training and validation results.\n",
    "    :param task: Task name ('emotion' or 'sentiment').\n",
    "    :param metric: Metric to plot (e.g., 'acc', 'macro_f1', 'weighted_f1').\n",
    "    :param modality: Modality to plot ('fused', 'audio', 'text').\n",
    "    \"\"\"\n",
    "    train_values = [\n",
    "        epoch_results['train_metrics'][modality][metric] \n",
    "        for epoch_results in results[f'results_{task}']\n",
    "    ]\n",
    "    val_values = [\n",
    "        epoch_results['val_metrics'][modality][metric] \n",
    "        for epoch_results in results[f'results_{task}']\n",
    "    ]\n",
    "    epochs = list(range(1, len(train_values) + 1))\n",
    "\n",
    "    plt.figure(figsize=(10, 6))\n",
    "    plt.plot(epochs, train_values, label=f'Train {metric.capitalize()}', marker='o')\n",
    "    plt.plot(epochs, val_values, label=f'Val {metric.capitalize()}', marker='o')\n",
    "    plt.title(f'{task.capitalize()} {metric.capitalize()} ({modality.capitalize()})')\n",
    "    plt.xlabel('Epochs')\n",
    "    plt.ylabel(metric.capitalize())\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "    plt.show()\n",
    "\n",
    "def analyze_results_per_class(true_labels, predicted_labels, class_names, task_name=\"Sentiment\", mode=\"confusion_matrix\"):\n",
    "    \"\"\"\n",
    "    Analyze results per class with confusion matrix, classification report, or ROC curves.\n",
    "\n",
    "    Args:\n",
    "        true_labels (list or np.ndarray): True labels for the task.\n",
    "        predicted_labels (list or np.ndarray): Predicted labels from the model.\n",
    "        class_names (list): List of class names.\n",
    "        task_name (str): Name of the task (e.g., \"Sentiment\" or \"Emotion\").\n",
    "        mode (str): The type of analysis. Options: \"confusion_matrix\", \"classification_report\", \"roc_curve\".\n",
    "    \"\"\"\n",
    "    os.makedirs('images', exist_ok=True)\n",
    "    \n",
    "    if mode == \"confusion_matrix\":\n",
    "        # Plot confusion matrix\n",
    "        cm = confusion_matrix(true_labels, predicted_labels, labels=range(len(class_names)))\n",
    "        plt.figure(figsize=(8, 6))\n",
    "        sns.heatmap(cm, annot=True, fmt=\"d\", cmap=\"Blues\", xticklabels=class_names, yticklabels=class_names)\n",
    "        plt.title(f\"Confusion Matrix for {task_name}\")\n",
    "        plt.xlabel(\"Predicted\")\n",
    "        plt.ylabel(\"True\")\n",
    "        plt.xticks(rotation=45)\n",
    "        plt.yticks(rotation=45)\n",
    "        plt.savefig(f'images/alternating/{task_name}_confusion_matrix.png')\n",
    "        plt.close()\n",
    "\n",
    "    elif mode == \"classification_report\":\n",
    "        # Print classification report\n",
    "        report = classification_report(true_labels, predicted_labels, target_names=class_names, zero_division=0)\n",
    "        with open(f'images/alternating/{task_name}_classification_report.txt', 'w') as f:\n",
    "            f.write(f\"Classification Report for {task_name}:\\n\\n\")\n",
    "            f.write(report)\n",
    "\n",
    "    elif mode == \"roc_curve\":\n",
    "        # Compute and plot ROC curves\n",
    "        true_binarized = label_binarize(true_labels, classes=range(len(class_names)))\n",
    "        predicted_binarized = label_binarize(predicted_labels, classes=range(len(class_names)))\n",
    "        plt.figure(figsize=(8, 6))\n",
    "        for i, class_name in enumerate(class_names):\n",
    "            fpr, tpr, _ = roc_curve(true_binarized[:, i], predicted_binarized[:, i])\n",
    "            roc_auc = auc(fpr, tpr)\n",
    "            plt.plot(fpr, tpr, label=f\"{class_name} (AUC = {roc_auc:.2f})\")\n",
    "        plt.plot([0, 1], [0, 1], \"k--\")  # Random baseline\n",
    "        plt.title(f\"ROC Curve for {task_name}\")\n",
    "        plt.xlabel(\"False Positive Rate\")\n",
    "        plt.ylabel(\"True Positive Rate\")\n",
    "        plt.legend(loc=\"lower right\")\n",
    "        plt.savefig(f'images/alternating/{task_name}_roc_curve.png')\n",
    "        plt.close()\n",
    "\n",
    "def compute_emotion_class_weights(train_set, device, normalize=True):\n",
    "    \"\"\"\n",
    "    Computes normalized inverse-frequency class weights for emotion labels \n",
    "    from the given training set and moves the result to the specified device.\n",
    "\n",
    "    Args:\n",
    "        train_set: Dataset object that contains samples and a method `get_emotions_dicts()`.\n",
    "        device: The target device (e.g., 'cpu' or 'cuda') to move the tensor.\n",
    "        normalize (bool): Whether to normalize the class weights to sum to 1. Default is True.\n",
    "\n",
    "    Returns:\n",
    "        torch.Tensor: Tensor containing class weights on the specified device.\n",
    "    \"\"\"\n",
    "    # Gather all emotion labels in the train set\n",
    "    emotion_labels = [sample[2] for sample in train_set.samples]\n",
    "\n",
    "    # Count how many samples of each class\n",
    "    unique_classes, counts = np.unique(emotion_labels, return_counts=True)\n",
    "    print(\"Class labels:\", unique_classes)\n",
    "    print(\"Class counts:\", counts)\n",
    "\n",
    "    # Retrieve mapping dictionaries; assume str_to_int maps emotion string to an integer index\n",
    "    _, str_to_int = train_set.get_emotions_dicts()\n",
    "\n",
    "    num_classes = len(str_to_int)\n",
    "    ordered_counts = [0] * num_classes\n",
    "\n",
    "    # Map counts to the correct class index order\n",
    "    for class_label, count in zip(unique_classes, counts):\n",
    "        class_idx = str_to_int[class_label]   # e.g. 'neutral' → 0, 'joy' → 1, etc.\n",
    "        ordered_counts[class_idx] = count\n",
    "\n",
    "    ordered_counts = np.array(ordered_counts)\n",
    "    print(\"Ordered counts:\", ordered_counts)\n",
    "\n",
    "    # Compute inverse frequency, avoiding division by zero\n",
    "    inverse_freq = 1.0 / np.maximum(ordered_counts, 1)\n",
    "\n",
    "    # Create a tensor of class weights\n",
    "    emotions_class_weights = torch.tensor(inverse_freq, dtype=torch.float32)\n",
    "\n",
    "    # Normalize the class weights if requested\n",
    "    if normalize:\n",
    "        emotions_class_weights = emotions_class_weights / emotions_class_weights.sum()\n",
    "\n",
    "    # Move the tensor to the specified device\n",
    "    emotions_class_weights = emotions_class_weights.to(device)\n",
    "\n",
    "    return emotions_class_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_one_epoch(model, dataloader, optimizer, criterions, emotion_reg=0.5, sentiment_reg=0.5, device='cuda'):\n",
    "    model.train()\n",
    "\n",
    "    losses = {'emotion': 0.0, 'sentiment': 0.0}\n",
    "    metrics = {\n",
    "        'emotion': {'fused': [], 'labels': []},\n",
    "        'sentiment': {'fused': [], 'labels': []}\n",
    "    }\n",
    "\n",
    "    loop = tqdm(dataloader, desc=\"Training\", leave=False)\n",
    "    \n",
    "    for batch in loop:\n",
    "        batch_emotion_loss = 0.0\n",
    "        batch_sentiment_loss = 0.0\n",
    "\n",
    "        batch_emotion_preds = []\n",
    "        batch_emotion_labels = []\n",
    "        batch_sentiment_preds = []\n",
    "        batch_sentiment_labels = []\n",
    "\n",
    "        for b_idx in range(len(batch[\"dialog_ids\"])):\n",
    "            fbank_list      = batch[\"fbank_lists\"][b_idx]       \n",
    "            text_list       = batch[\"text_lists\"][b_idx]        \n",
    "            emotion_list    = batch[\"emotion_lists\"][b_idx]\n",
    "            sentiment_list  = batch[\"sentiment_lists\"][b_idx]\n",
    "\n",
    "            # Determine the actual number of utterances before padding\n",
    "            actual_len = len(emotion_list)  # Number of real utterances in this conversation\n",
    "\n",
    "            # Convert lists to tensors\n",
    "            emotion_tensor = torch.as_tensor(emotion_list, dtype=torch.long, device=device)\n",
    "            sentiment_tensor = torch.as_tensor(sentiment_list, dtype=torch.long, device=device)\n",
    "            audio_array = torch.as_tensor(fbank_list, dtype=torch.float, device=device)\n",
    "\n",
    "            audio_array = audio_array.unsqueeze(1)  # adjust dimensions as needed\n",
    "\n",
    "            # Forward pass\n",
    "            emotion_logits, sentiment_logits = model(audio_array, text_list)\n",
    "\n",
    "            # If outputs have shape (1, seq_len, num_classes), squeeze the batch dimension\n",
    "            if len(emotion_logits.shape) == 3:\n",
    "                emotion_logits = emotion_logits.squeeze(0) \n",
    "                sentiment_logits = sentiment_logits.squeeze(0)\n",
    "\n",
    "            # Slice the logits and tensors to ignore padded timesteps\n",
    "            # emotion_logits shape assumed to be (max_utt, num_classes) at this point\n",
    "            emotion_logits = emotion_logits[:actual_len]\n",
    "            sentiment_logits = sentiment_logits[:actual_len]\n",
    "\n",
    "            emotion_tensor = emotion_tensor[:actual_len]\n",
    "            sentiment_tensor = sentiment_tensor[:actual_len]\n",
    "\n",
    "            # Compute Loss\n",
    "            e_loss = criterions['emotion'](emotion_logits, emotion_tensor)\n",
    "            s_loss = criterions['sentiment'](sentiment_logits, sentiment_tensor)\n",
    "\n",
    "            batch_emotion_loss += e_loss\n",
    "            batch_sentiment_loss += s_loss\n",
    "\n",
    "            # Predictions for non-padded time steps\n",
    "            e_preds = torch.argmax(emotion_logits, dim=-1).detach().cpu().numpy()\n",
    "            s_preds = torch.argmax(sentiment_logits, dim=-1).detach().cpu().numpy()\n",
    "\n",
    "            batch_emotion_preds.extend(e_preds)\n",
    "            batch_emotion_labels.extend(emotion_list)\n",
    "\n",
    "            batch_sentiment_preds.extend(s_preds)\n",
    "            batch_sentiment_labels.extend(sentiment_list)\n",
    "\n",
    "        # Average and combine losses over conversations in batch\n",
    "        batch_emotion_loss /= len(batch[\"dialog_ids\"])\n",
    "        batch_sentiment_loss /= len(batch[\"dialog_ids\"])\n",
    "\n",
    "        combined_loss = emotion_reg * batch_emotion_loss + sentiment_reg * batch_sentiment_loss\n",
    "        optimizer.zero_grad()\n",
    "        combined_loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        losses['emotion'] += batch_emotion_loss.item()\n",
    "        losses['sentiment'] += batch_sentiment_loss.item()\n",
    "\n",
    "        metrics['emotion']['fused'].extend(batch_emotion_preds)\n",
    "        metrics['emotion']['labels'].extend(batch_emotion_labels)\n",
    "\n",
    "        metrics['sentiment']['fused'].extend(batch_sentiment_preds)\n",
    "        metrics['sentiment']['labels'].extend(batch_sentiment_labels)\n",
    "\n",
    "    losses['emotion'] /= len(dataloader)\n",
    "    losses['sentiment'] /= len(dataloader)\n",
    "\n",
    "    emotion_metrics = compute_metrics(metrics['emotion']['labels'], metrics['emotion']['fused'])\n",
    "    sentiment_metrics = compute_metrics(metrics['sentiment']['labels'], metrics['sentiment']['fused'])\n",
    "\n",
    "    return losses, {\n",
    "        'emotion': {'fused': emotion_metrics},\n",
    "        'sentiment': {'fused': sentiment_metrics}\n",
    "    }"
   ]
  }
 ],
 "metadata": {
  "kaggle": {
   "accelerator": "gpu",
   "dataSources": [
    {
     "datasetId": 6413518,
     "sourceId": 10356240,
     "sourceType": "datasetVersion"
    },
    {
     "datasetId": 6415331,
     "sourceId": 10358898,
     "sourceType": "datasetVersion"
    },
    {
     "datasetId": 6434748,
     "sourceId": 10386819,
     "sourceType": "datasetVersion"
    }
   ],
   "dockerImageVersionId": 30823,
   "isGpuEnabled": true,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "kernelspec": {
   "display_name": "audio",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}