speaker_prediction/train_speaker_prediction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Speaker Identification using Whisper\n",
    "* Train a classification model to identify if the speaker in a audio segment is lex or not\n",
    "\n",
    "### Data:\n",
    "* Expecting training audio clips in data/audio_dataset. Download from [here](https://drive.google.com/file/d/1SF0j1UmMxpwFNeY1wkj3R20pRB7L0a4t/view?usp=share_link)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "is cuda available: True\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from pathlib import Path\n",
    "import os\n",
    "import sys\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "\n",
    "print('is cuda available:', torch.cuda.is_available())\n",
    "\n",
    "# Add whisper repo to path to import\n",
    "repo_dir = Path(os.getcwd()).parents[0]/'whisper'\n",
    "sys.path.append(str(repo_dir))\n",
    "import whisper\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Whisper model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# if this model is too slow, try the other smaller models such as small.en and base.en\n",
    "model = whisper.load_model(\"medium.en\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Labelled dataset\n",
    "\n",
    "* Contains start and end defining segment of the audio clip\n",
    "* audio_name is the name of the podcast. The podcasts audio files should be present in data/lex_podcasts\n",
    "* `is_lex` is the label. 1 if speaker from start time to end time of the audio clip is lex. 0 if not\n",
    "\n",
    "#### Data augmentation:\n",
    "* All samples with audio_idx < 0 are heuristically labelled based on keyword, not hand labelled (for eg. \"the following is a conv\" in the start of a podcast is always lex\"). Choose to select a subsample of this in training, since all audio clips will share a very similar spectral pattern"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>start</th>\n",
       "      <th>end</th>\n",
       "      <th>text</th>\n",
       "      <th>fname</th>\n",
       "      <th>audio_name</th>\n",
       "      <th>audio_idx</th>\n",
       "      <th>is_lex</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>02:49:11.280</td>\n",
       "      <td>02:49:15.120</td>\n",
       "      <td>And, you know, some people also ask, are you ...</td>\n",
       "      <td>episode_215</td>\n",
       "      <td>Wojciech Zaremba： OpenAI Codex, GPT-3, Robotic...</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>02:20:14.140</td>\n",
       "      <td>02:20:17.260</td>\n",
       "      <td>I still do that often.</td>\n",
       "      <td>episode_215</td>\n",
       "      <td>Wojciech Zaremba： OpenAI Codex, GPT-3, Robotic...</td>\n",
       "      <td>1</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>00:19:15.360</td>\n",
       "      <td>00:19:17.320</td>\n",
       "      <td>things that you put into context of GPT.</td>\n",
       "      <td>episode_215</td>\n",
       "      <td>Wojciech Zaremba： OpenAI Codex, GPT-3, Robotic...</td>\n",
       "      <td>2</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>02:45:11.760</td>\n",
       "      <td>02:45:16.000</td>\n",
       "      <td>and that also gives, you know, huge perspecti...</td>\n",
       "      <td>episode_215</td>\n",
       "      <td>Wojciech Zaremba： OpenAI Codex, GPT-3, Robotic...</td>\n",
       "      <td>3</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>01:33:44.600</td>\n",
       "      <td>01:33:49.160</td>\n",
       "      <td>You, it's often the way how it works is you o...</td>\n",
       "      <td>episode_215</td>\n",
       "      <td>Wojciech Zaremba： OpenAI Codex, GPT-3, Robotic...</td>\n",
       "      <td>4</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          start           end  \\\n",
       "0  02:49:11.280  02:49:15.120   \n",
       "1  02:20:14.140  02:20:17.260   \n",
       "2  00:19:15.360  00:19:17.320   \n",
       "3  02:45:11.760  02:45:16.000   \n",
       "4  01:33:44.600  01:33:49.160   \n",
       "\n",
       "                                                text        fname  \\\n",
       "0   And, you know, some people also ask, are you ...  episode_215   \n",
       "1                             I still do that often.  episode_215   \n",
       "2           things that you put into context of GPT.  episode_215   \n",
       "3   and that also gives, you know, huge perspecti...  episode_215   \n",
       "4   You, it's often the way how it works is you o...  episode_215   \n",
       "\n",
       "                                          audio_name  audio_idx  is_lex  \n",
       "0  Wojciech Zaremba： OpenAI Codex, GPT-3, Robotic...          0     0.0  \n",
       "1  Wojciech Zaremba： OpenAI Codex, GPT-3, Robotic...          1     1.0  \n",
       "2  Wojciech Zaremba： OpenAI Codex, GPT-3, Robotic...          2     0.0  \n",
       "3  Wojciech Zaremba： OpenAI Codex, GPT-3, Robotic...          3     0.0  \n",
       "4  Wojciech Zaremba： OpenAI Codex, GPT-3, Robotic...          4     0.0  "
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "audio_dataset_dir = Path('data/audio_dataset')\n",
    "labelled_path = 'data/labelled_dataset.csv'\n",
    "df = pd.read_csv(labelled_path)\n",
    "\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Num audio clips: 698\n",
      "Num unique podcasts: 68\n",
      "\n",
      "\n",
      "Postcasts containing the most tagged clips\n",
      "Elon Musk： Neuralink, AI, Autopilot, and the Pale Blue Dot ｜ Lex Fridman Podcast #49                   65\n",
      "Ray Dalio： Principles, the Economic Machine, AI & the Arc of Life ｜ Lex Fridman Podcast #54            21\n",
      "Judea Pearl： Causal Reasoning, Counterfactuals, and the Path to AGI ｜ Lex Fridman Podcast #56          21\n",
      "Dmitry Korkin： Computational Biology of Coronavirus ｜ Lex Fridman Podcast #90                          21\n",
      "Jeremy Howard： fast.ai Deep Learning Courses and Research ｜ Lex Fridman Podcast #35                    21\n",
      "Cumrun Vafa： String Theory ｜ Lex Fridman Podcast #204                                                  21\n",
      "Po-Shen Loh： Mathematics, Math Olympiad, Combinatorics & Contact Tracing ｜ Lex Fridman Podcast #183    20\n",
      "Jim Keller： Moore's Law, Microprocessors, and First Principles ｜ Lex Fridman Podcast #70               20\n",
      "Name: audio_name, dtype: int64\n",
      "\n",
      "\n",
      "Postcasts containing the least tagged clips\n",
      "Zach Bitter： Ultramarathon Running ｜ Lex Fridman Podcast #205                        1\n",
      "Whitney Cummings： Comedy, Robotics, Neurology, and Love ｜ Lex Fridman Podcast #55    1\n",
      "Name: audio_name, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "podcasts = list(df['audio_name'].unique())\n",
    "print('Num audio clips:', len(df))\n",
    "print(f'Num unique podcasts: {len(podcasts)}')\n",
    "\n",
    "print(f'\\n\\nPostcasts containing the most tagged clips\\n{df[\"audio_name\"].value_counts().head(8)}')\n",
    "print(f'\\n\\nPostcasts containing the least tagged clips\\n{df[\"audio_name\"].value_counts().tail(2)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create features for classifer\n",
    "* Each hidden state is of shape (batch_size, 1500, hidden_size). The 1500 is the hidden state across 1500 time periods. We need to summarize features across the 3 time windows and create a feature vector of size (batch_size, hidden_size)\n",
    "* This what get_feature_vector functions do: \n",
    "    * `get_feature_vector3`: Performs the best. Creates mean and std features across time window. Concatenate mean and std across time. Creates feature of shape (batch_size, hidden_size + hidden_size).\n",
    "    * `get_feature_vector2`: Creates mean features across time windows. (batch_size, hidden_size)\n",
    "    * `get_feature_vector1`: Calculates mean features across 3 time windows of 1500 (1-500, 500-1000, 1000-1500). Creates feature shape of (batch_size, hidden_size * 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_feature_vector1(batch):\n",
    "    \"\"\"\n",
    "    Get features to train classifier.  Input: batch with 1 sample (bs, timewindows, hidden_size)\n",
    "    \n",
    "    Concatate mean features across 3 time windows \n",
    "    \n",
    "    \"\"\"\n",
    "    out = [batch[0, :500, :].mean(1).flatten(), batch[0, 500:1000, :].mean(1).flatten(), batch[0, 1000:, :].mean(1).flatten()]\n",
    "    X = torch.cat(out, dim=-1)\n",
    "        \n",
    "    return X[None, :]\n",
    "\n",
    "def get_feature_vector2(batch):\n",
    "    \"\"\"\n",
    "    Get features to train classifier.  Input: batch with 1 sample (bs, timewindows, hidden_size)\n",
    "    \n",
    "    Mean of features across entire timewindow of a clip\n",
    "    \n",
    "    \"\"\"\n",
    "    X = batch[0, :, :].mean(0)\n",
    "    return X[None, :]\n",
    "\n",
    "def get_feature_vector3(batch):\n",
    "    \"\"\"\n",
    "    Get features to train classifier. Input: batch with 1 sample (bs, timewindows, hidden_size)\n",
    "    \n",
    "    Mean and std of features across entire timewindow of a clip\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "\n",
    "    out = batch[0, :, :].mean(0)\n",
    "    out2 = batch[0, :, :].std(0)\n",
    "    X = torch.cat([out, out2], dim=-1)\n",
    "        \n",
    "    return X[None, :]\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get Whisper Embeddings\n",
    "* Requires GPU to finish fast. If it takes too long, consider using a smaller Whisper model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 698/698 [04:16<00:00,  2.72it/s]\n"
     ]
    }
   ],
   "source": [
    "if torch.cuda.is_available():\n",
    "    model = model.cuda()\n",
    "    \n",
    "if not os.path.exists(audio_dataset_dir):\n",
    "    if not audio_dataset_dir.exists():\n",
    "        raise ValueError('Expecting audio clips data in ', audio_dataset_dir)\n",
    "        \n",
    "filenames = df['audio_idx'].apply(lambda x: str(x) + '.mp3')\n",
    "audio_paths  = [audio_dataset_dir/filename for filename in filenames]\n",
    "\n",
    "idx_to_path = {idx: path for idx, path in enumerate(audio_paths)}\n",
    "\n",
    "hidden_l1 = []\n",
    "hidden_l2 = []\n",
    "hidden_last = []\n",
    "hidden_middle = []\n",
    "\n",
    "metadata_outputs = []\n",
    "\n",
    "\n",
    "model = model.eval()\n",
    "\n",
    "for audio_path in tqdm(audio_paths, total=len(audio_paths), disable=False):\n",
    "    \n",
    "    assert os.path.exists(audio_path)\n",
    "\n",
    "    # load audio\n",
    "    audio = whisper.load_audio(str(audio_path))\n",
    "    audio = whisper.pad_or_trim(audio)\n",
    "    \n",
    "    # create mel spectogram input for whisper encoder\n",
    "    mel = whisper.log_mel_spectrogram(audio).to(model.device)\n",
    "    mel = mel[None, :, :]\n",
    "    \n",
    "    # forward pass through encoder\n",
    "    with torch.no_grad():\n",
    "        _ = model.embed_audio(mel)\n",
    "    del _\n",
    "    \n",
    "    # get various hidden states of enocder\n",
    "    hidden_l1.append(get_feature_vector3(model.encoder.encoder_out1.cpu()))\n",
    "    hidden_l2.append(get_feature_vector3(model.encoder.encoder_out2.cpu()))\n",
    "    hidden_last.append(get_feature_vector3(model.encoder.encoder_out_last.cpu()))\n",
    "    hidden_middle.append(get_feature_vector3(model.encoder.encoder_out_middle.cpu()))\n",
    "    \n",
    "    metadata_outputs.append({\"metadata\": None, \"audio_path\": audio_path.name})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>start</th>\n",
       "      <th>end</th>\n",
       "      <th>text</th>\n",
       "      <th>fname</th>\n",
       "      <th>audio_name</th>\n",
       "      <th>audio_idx</th>\n",
       "      <th>is_lex</th>\n",
       "      <th>metadata</th>\n",
       "      <th>audio_path</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>02:49:11.280</td>\n",
       "      <td>02:49:15.120</td>\n",
       "      <td>And, you know, some people also ask, are you ...</td>\n",
       "      <td>episode_215</td>\n",
       "      <td>Wojciech Zaremba： OpenAI Codex, GPT-3, Robotic...</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>None</td>\n",
       "      <td>0.mp3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>02:20:14.140</td>\n",
       "      <td>02:20:17.260</td>\n",
       "      <td>I still do that often.</td>\n",
       "      <td>episode_215</td>\n",
       "      <td>Wojciech Zaremba： OpenAI Codex, GPT-3, Robotic...</td>\n",
       "      <td>1</td>\n",
       "      <td>1.0</td>\n",
       "      <td>None</td>\n",
       "      <td>1.mp3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>00:19:15.360</td>\n",
       "      <td>00:19:17.320</td>\n",
       "      <td>things that you put into context of GPT.</td>\n",
       "      <td>episode_215</td>\n",
       "      <td>Wojciech Zaremba： OpenAI Codex, GPT-3, Robotic...</td>\n",
       "      <td>2</td>\n",
       "      <td>0.0</td>\n",
       "      <td>None</td>\n",
       "      <td>2.mp3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>02:45:11.760</td>\n",
       "      <td>02:45:16.000</td>\n",
       "      <td>and that also gives, you know, huge perspecti...</td>\n",
       "      <td>episode_215</td>\n",
       "      <td>Wojciech Zaremba： OpenAI Codex, GPT-3, Robotic...</td>\n",
       "      <td>3</td>\n",
       "      <td>0.0</td>\n",
       "      <td>None</td>\n",
       "      <td>3.mp3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>01:33:44.600</td>\n",
       "      <td>01:33:49.160</td>\n",
       "      <td>You, it's often the way how it works is you o...</td>\n",
       "      <td>episode_215</td>\n",
       "      <td>Wojciech Zaremba： OpenAI Codex, GPT-3, Robotic...</td>\n",
       "      <td>4</td>\n",
       "      <td>0.0</td>\n",
       "      <td>None</td>\n",
       "      <td>4.mp3</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          start           end  \\\n",
       "0  02:49:11.280  02:49:15.120   \n",
       "1  02:20:14.140  02:20:17.260   \n",
       "2  00:19:15.360  00:19:17.320   \n",
       "3  02:45:11.760  02:45:16.000   \n",
       "4  01:33:44.600  01:33:49.160   \n",
       "\n",
       "                                                text        fname  \\\n",
       "0   And, you know, some people also ask, are you ...  episode_215   \n",
       "1                             I still do that often.  episode_215   \n",
       "2           things that you put into context of GPT.  episode_215   \n",
       "3   and that also gives, you know, huge perspecti...  episode_215   \n",
       "4   You, it's often the way how it works is you o...  episode_215   \n",
       "\n",
       "                                          audio_name  audio_idx  is_lex  \\\n",
       "0  Wojciech Zaremba： OpenAI Codex, GPT-3, Robotic...          0     0.0   \n",
       "1  Wojciech Zaremba： OpenAI Codex, GPT-3, Robotic...          1     1.0   \n",
       "2  Wojciech Zaremba： OpenAI Codex, GPT-3, Robotic...          2     0.0   \n",
       "3  Wojciech Zaremba： OpenAI Codex, GPT-3, Robotic...          3     0.0   \n",
       "4  Wojciech Zaremba： OpenAI Codex, GPT-3, Robotic...          4     0.0   \n",
       "\n",
       "  metadata audio_path  \n",
       "0     None      0.mp3  \n",
       "1     None      1.mp3  \n",
       "2     None      2.mp3  \n",
       "3     None      3.mp3  \n",
       "4     None      4.mp3  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df2 = pd.concat([df, pd.DataFrame(metadata_outputs)], axis=1)\n",
    "df2.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train test split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Split found with ratio: {0.0: 0.5815217391304348, 1.0: 0.41847826086956524}: seed: 63\n",
      "Split found with ratio: {0.0: 0.6, 1.0: 0.4}: seed: 650\n",
      "Split found with ratio: {0.0: 0.5869565217391305, 1.0: 0.41304347826086957}: seed: 932\n",
      "Split found with ratio: {0.0: 0.5934065934065934, 1.0: 0.4065934065934066}: seed: 56\n",
      "Split found with ratio: {0.0: 0.5766423357664233, 1.0: 0.4233576642335766}: seed: 173\n"
     ]
    }
   ],
   "source": [
    "from collections import Counter\n",
    "import numpy as np \n",
    "\n",
    "def cv_split(df, seed=42):\n",
    "    \"\"\"\n",
    "    Split on podcast\n",
    "    \"\"\"\n",
    "    \n",
    "    import random\n",
    "    \n",
    "    randgen = random.Random(seed)\n",
    "        \n",
    "    subdf = df[df['audio_idx'] >= 0]\n",
    "    \n",
    "    # negative audio_idxs are not manually labelled, heuristically labelled. \n",
    "    # Near 100% are accurate, but speech pattern is same so only use few samples\n",
    "    augmented_df = df[df['audio_idx'] < -1].sample(50)\n",
    "    \n",
    "    # split based on speakerid \n",
    "    sources = list(subdf['audio_name'].unique())\n",
    "    \n",
    "    test_split = randgen.sample(sources, len(sources) // 4)\n",
    "    train_split = list(set(sources).difference(test_split))\n",
    "\n",
    "    test_df = df[df['audio_name'].isin(test_split)]\n",
    "    train_df = df[df['audio_name'].isin(train_split)]\n",
    "    \n",
    "    # add heuristic data to labelled data\n",
    "    train_df = pd.concat([train_df, augmented_df], axis=0)\n",
    "                               \n",
    "    \n",
    "    return train_df, test_df\n",
    "    \n",
    "def generate_splits(df, num_splits=5):\n",
    "    \"\"\"\n",
    "    Splitting on speaker ID randomly leads to very high class imbalance in test set - in some podcasts lex tags are very few. \n",
    "    The strategy is to keep splitting on random seeds until a split of 40%-60% is reached\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    cvs = []\n",
    "    seen_seeds = set()\n",
    "    for split in range(num_splits):\n",
    "        seeds = np.random.randint(low=0, high=1000, size=(50,))\n",
    "        for seed in seeds: \n",
    "            # cant use same seed again \n",
    "            if seed in seen_seeds:\n",
    "                continue\n",
    "            train_df, test_df = cv_split(df, seed=seed)\n",
    "            counts = test_df['is_lex'].value_counts()\n",
    "            counts = counts/counts.sum()\n",
    "            if (counts.loc[1.0] >= 0.40) and (counts.loc[1.0] <= 0.60):\n",
    "                print(f'Split found with ratio: {counts.to_dict()}: seed: {seed}')\n",
    "                break \n",
    "                \n",
    "        cvs.append((train_df, test_df))\n",
    "        seen_seeds.add(seed)\n",
    "        \n",
    "    return cvs \n",
    "\n",
    "        \n",
    "splits = generate_splits(df, num_splits=5)        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import f1_score, precision_score, recall_score, accuracy_score, classification_report\n",
    "from sklearn.exceptions import ConvergenceWarning\n",
    "import warnings \n",
    "from sklearn import preprocessing\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.pipeline import make_pipeline, Pipeline\n",
    "from sklearn import svm\n",
    "\n",
    "with warnings.catch_warnings():\n",
    "    warnings.filterwarnings(\"ignore\", category=ConvergenceWarning)\n",
    "\n",
    "def get_metrics(true, pred):\n",
    "    f1, recall, precision, accuracy = f1_score(true, pred), recall_score(true, pred), precision_score(true, pred), accuracy_score(true, pred)\n",
    "    \n",
    "    return {'f1': f1, 'recall': recall, 'precision': precision, 'accuracy': accuracy}\n",
    "\n",
    "def train_eval(X, splits):\n",
    "    \"\"\"\n",
    "    Train model\n",
    "    \n",
    "    Inputs: \n",
    "        - X: All features X \n",
    "        - splits: List of tuples of (train_df, test_df)\n",
    "    \n",
    "    \"\"\"\n",
    "    test_metrics = []\n",
    "    train_metrics = []\n",
    "    fold_preds = []\n",
    "    for train_df, test_df in splits:\n",
    "        X_train, y_train = X[list(train_df.index), :], train_df['is_lex']\n",
    "        X_test, y_test = X[list(test_df.index), :], test_df['is_lex']\n",
    "\n",
    "        scalar = preprocessing.StandardScaler()\n",
    "        \n",
    "        # overfits fast with logistic regression\n",
    "#         clf = LogisticRegression(random_state=0, max_iter=15, C=.7)\n",
    "        clf = svm.SVC(kernel='rbf', C=.7)\n",
    "\n",
    "        pipeline = Pipeline([('transformer', scalar), ('estimator', clf)])\n",
    "\n",
    "        pipeline.fit(X_train, y_train)\n",
    "        pred_train = pipeline.predict(X_train)\n",
    "        pred_test = pipeline.predict(X_test)\n",
    "\n",
    "        m1 = get_metrics(y_train, pred_train)\n",
    "        m1['num_positive_samples'] = (y_train==1).sum()\n",
    "        m1['num_negative_samples'] = (y_train==0).sum()\n",
    "        train_metrics.append(m1)\n",
    "        \n",
    "        m2 = get_metrics(y_test, pred_test)\n",
    "        m2['num_positive_samples'] = (y_test==1).sum()\n",
    "        m2['num_negative_samples'] = (y_test==0).sum()\n",
    "        test_metrics.append(m2)\n",
    "        \n",
    "        fold_df = test_df.copy()\n",
    "        fold_df['preds'] = pred_test\n",
    "        \n",
    "        fold_preds.append(fold_df)\n",
    "        \n",
    "        \n",
    "        \n",
    "\n",
    "    train_metrics = pd.DataFrame(train_metrics)\n",
    "    test_metrics = pd.DataFrame(test_metrics)\n",
    "\n",
    "    display('Test stats ', test_metrics.describe().loc[[ 'mean','std']])\n",
    "    display('Train stats ', train_metrics.describe().loc[[ 'mean','std']])\n",
    "    \n",
    "    return train_metrics, test_metrics, fold_preds\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train and Predict\n",
    "* Train classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "----Metrics of output of first_hidden encoder block output----\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Test stats '"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>f1</th>\n",
       "      <th>recall</th>\n",
       "      <th>precision</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>num_positive_samples</th>\n",
       "      <th>num_negative_samples</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>0.741579</td>\n",
       "      <td>0.733944</td>\n",
       "      <td>0.763987</td>\n",
       "      <td>0.786568</td>\n",
       "      <td>67.800000</td>\n",
       "      <td>96.600000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.059654</td>\n",
       "      <td>0.077747</td>\n",
       "      <td>0.111614</td>\n",
       "      <td>0.067345</td>\n",
       "      <td>10.917875</td>\n",
       "      <td>15.175638</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            f1    recall  precision  accuracy  num_positive_samples  \\\n",
       "mean  0.741579  0.733944   0.763987  0.786568             67.800000   \n",
       "std   0.059654  0.077747   0.111614  0.067345             10.917875   \n",
       "\n",
       "      num_negative_samples  \n",
       "mean             96.600000  \n",
       "std              15.175638  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "'Train stats '"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>f1</th>\n",
       "      <th>recall</th>\n",
       "      <th>precision</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>num_positive_samples</th>\n",
       "      <th>num_negative_samples</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>0.902497</td>\n",
       "      <td>0.861970</td>\n",
       "      <td>0.947235</td>\n",
       "      <td>0.932789</td>\n",
       "      <td>196.200000</td>\n",
       "      <td>346.400000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.014132</td>\n",
       "      <td>0.020646</td>\n",
       "      <td>0.012351</td>\n",
       "      <td>0.008928</td>\n",
       "      <td>10.917875</td>\n",
       "      <td>15.175638</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            f1    recall  precision  accuracy  num_positive_samples  \\\n",
       "mean  0.902497  0.861970   0.947235  0.932789            196.200000   \n",
       "std   0.014132  0.020646   0.012351  0.008928             10.917875   \n",
       "\n",
       "      num_negative_samples  \n",
       "mean            346.400000  \n",
       "std              15.175638  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "----Metrics of output of second_hidden encoder block output----\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Test stats '"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>f1</th>\n",
       "      <th>recall</th>\n",
       "      <th>precision</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>num_positive_samples</th>\n",
       "      <th>num_negative_samples</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>0.794997</td>\n",
       "      <td>0.795201</td>\n",
       "      <td>0.813036</td>\n",
       "      <td>0.828643</td>\n",
       "      <td>67.800000</td>\n",
       "      <td>96.600000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.050030</td>\n",
       "      <td>0.088273</td>\n",
       "      <td>0.115795</td>\n",
       "      <td>0.058580</td>\n",
       "      <td>10.917875</td>\n",
       "      <td>15.175638</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            f1    recall  precision  accuracy  num_positive_samples  \\\n",
       "mean  0.794997  0.795201   0.813036  0.828643             67.800000   \n",
       "std   0.050030  0.088273   0.115795  0.058580             10.917875   \n",
       "\n",
       "      num_negative_samples  \n",
       "mean             96.600000  \n",
       "std              15.175638  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "'Train stats '"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>f1</th>\n",
       "      <th>recall</th>\n",
       "      <th>precision</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>num_positive_samples</th>\n",
       "      <th>num_negative_samples</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>0.920326</td>\n",
       "      <td>0.884621</td>\n",
       "      <td>0.959119</td>\n",
       "      <td>0.944697</td>\n",
       "      <td>196.200000</td>\n",
       "      <td>346.400000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.009157</td>\n",
       "      <td>0.014203</td>\n",
       "      <td>0.005749</td>\n",
       "      <td>0.005722</td>\n",
       "      <td>10.917875</td>\n",
       "      <td>15.175638</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            f1    recall  precision  accuracy  num_positive_samples  \\\n",
       "mean  0.920326  0.884621   0.959119  0.944697            196.200000   \n",
       "std   0.009157  0.014203   0.005749  0.005722             10.917875   \n",
       "\n",
       "      num_negative_samples  \n",
       "mean            346.400000  \n",
       "std              15.175638  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "----Metrics of output of middle_hidden encoder block output----\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Test stats '"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>f1</th>\n",
       "      <th>recall</th>\n",
       "      <th>precision</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>num_positive_samples</th>\n",
       "      <th>num_negative_samples</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>0.930347</td>\n",
       "      <td>0.931942</td>\n",
       "      <td>0.930254</td>\n",
       "      <td>0.942578</td>\n",
       "      <td>67.800000</td>\n",
       "      <td>96.600000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.034269</td>\n",
       "      <td>0.038175</td>\n",
       "      <td>0.051011</td>\n",
       "      <td>0.028605</td>\n",
       "      <td>10.917875</td>\n",
       "      <td>15.175638</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            f1    recall  precision  accuracy  num_positive_samples  \\\n",
       "mean  0.930347  0.931942   0.930254  0.942578             67.800000   \n",
       "std   0.034269  0.038175   0.051011  0.028605             10.917875   \n",
       "\n",
       "      num_negative_samples  \n",
       "mean             96.600000  \n",
       "std              15.175638  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "'Train stats '"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>f1</th>\n",
       "      <th>recall</th>\n",
       "      <th>precision</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>num_positive_samples</th>\n",
       "      <th>num_negative_samples</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>0.987889</td>\n",
       "      <td>0.986776</td>\n",
       "      <td>0.989027</td>\n",
       "      <td>0.991244</td>\n",
       "      <td>196.200000</td>\n",
       "      <td>346.400000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.003545</td>\n",
       "      <td>0.004278</td>\n",
       "      <td>0.005904</td>\n",
       "      <td>0.002605</td>\n",
       "      <td>10.917875</td>\n",
       "      <td>15.175638</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            f1    recall  precision  accuracy  num_positive_samples  \\\n",
       "mean  0.987889  0.986776   0.989027  0.991244            196.200000   \n",
       "std   0.003545  0.004278   0.005904  0.002605             10.917875   \n",
       "\n",
       "      num_negative_samples  \n",
       "mean            346.400000  \n",
       "std              15.175638  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "----Metrics of output of last_hidden encoder block output----\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Test stats '"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>f1</th>\n",
       "      <th>recall</th>\n",
       "      <th>precision</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>num_positive_samples</th>\n",
       "      <th>num_negative_samples</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>0.907284</td>\n",
       "      <td>0.890827</td>\n",
       "      <td>0.926739</td>\n",
       "      <td>0.924785</td>\n",
       "      <td>67.800000</td>\n",
       "      <td>96.600000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.018847</td>\n",
       "      <td>0.028368</td>\n",
       "      <td>0.051824</td>\n",
       "      <td>0.016446</td>\n",
       "      <td>10.917875</td>\n",
       "      <td>15.175638</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            f1    recall  precision  accuracy  num_positive_samples  \\\n",
       "mean  0.907284  0.890827   0.926739  0.924785             67.800000   \n",
       "std   0.018847  0.028368   0.051824  0.016446             10.917875   \n",
       "\n",
       "      num_negative_samples  \n",
       "mean             96.600000  \n",
       "std              15.175638  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "'Train stats '"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>f1</th>\n",
       "      <th>recall</th>\n",
       "      <th>precision</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>num_positive_samples</th>\n",
       "      <th>num_negative_samples</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>0.980991</td>\n",
       "      <td>0.969397</td>\n",
       "      <td>0.992885</td>\n",
       "      <td>0.986433</td>\n",
       "      <td>196.200000</td>\n",
       "      <td>346.400000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.006044</td>\n",
       "      <td>0.007926</td>\n",
       "      <td>0.005553</td>\n",
       "      <td>0.004272</td>\n",
       "      <td>10.917875</td>\n",
       "      <td>15.175638</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            f1    recall  precision  accuracy  num_positive_samples  \\\n",
       "mean  0.980991  0.969397   0.992885  0.986433            196.200000   \n",
       "std   0.006044  0.007926   0.005553  0.004272             10.917875   \n",
       "\n",
       "      num_negative_samples  \n",
       "mean            346.400000  \n",
       "std              15.175638  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# utilize different hidden states for training and evaluate\n",
    "for X, hidden_name in zip([hidden_l1, hidden_l2, hidden_middle, hidden_last], ['first_hidden', 'second_hidden', 'middle_hidden', 'last_hidden']):          \n",
    "    X_use = torch.cat(X, dim=0)\n",
    "    \n",
    "    print(f'\\n----Metrics of output of {hidden_name} encoder block output----\\n')\n",
    "    train_metrics, test_metrics, pred_dfs = train_eval(X_use, splits)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Explore predictions\n",
    "* Get predictions of one fold"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Test stats '"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>f1</th>\n",
       "      <th>recall</th>\n",
       "      <th>precision</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>num_positive_samples</th>\n",
       "      <th>num_negative_samples</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>0.930347</td>\n",
       "      <td>0.931942</td>\n",
       "      <td>0.930254</td>\n",
       "      <td>0.942578</td>\n",
       "      <td>67.800000</td>\n",
       "      <td>96.600000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.034269</td>\n",
       "      <td>0.038175</td>\n",
       "      <td>0.051011</td>\n",
       "      <td>0.028605</td>\n",
       "      <td>10.917875</td>\n",
       "      <td>15.175638</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            f1    recall  precision  accuracy  num_positive_samples  \\\n",
       "mean  0.930347  0.931942   0.930254  0.942578             67.800000   \n",
       "std   0.034269  0.038175   0.051011  0.028605             10.917875   \n",
       "\n",
       "      num_negative_samples  \n",
       "mean             96.600000  \n",
       "std              15.175638  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "'Train stats '"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>f1</th>\n",
       "      <th>recall</th>\n",
       "      <th>precision</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>num_positive_samples</th>\n",
       "      <th>num_negative_samples</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>0.987889</td>\n",
       "      <td>0.986776</td>\n",
       "      <td>0.989027</td>\n",
       "      <td>0.991244</td>\n",
       "      <td>196.200000</td>\n",
       "      <td>346.400000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.003545</td>\n",
       "      <td>0.004278</td>\n",
       "      <td>0.005904</td>\n",
       "      <td>0.002605</td>\n",
       "      <td>10.917875</td>\n",
       "      <td>15.175638</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            f1    recall  precision  accuracy  num_positive_samples  \\\n",
       "mean  0.987889  0.986776   0.989027  0.991244            196.200000   \n",
       "std   0.003545  0.004278   0.005904  0.002605             10.917875   \n",
       "\n",
       "      num_negative_samples  \n",
       "mean            346.400000  \n",
       "std              15.175638  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "X_use = torch.cat(hidden_middle, dim=0)\n",
    "train_metrics, test_metrics, pred_dfs = train_eval(X_use, splits)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Metrics '"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "f1                        0.962025\n",
       "recall                    0.987013\n",
       "precision                 0.938272\n",
       "accuracy                  0.967391\n",
       "num_positive_samples     77.000000\n",
       "num_negative_samples    107.000000\n",
       "Name: 0, dtype: float64"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "'Preds'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>start</th>\n",
       "      <th>end</th>\n",
       "      <th>text</th>\n",
       "      <th>fname</th>\n",
       "      <th>audio_name</th>\n",
       "      <th>audio_idx</th>\n",
       "      <th>is_lex</th>\n",
       "      <th>preds</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>39</th>\n",
       "      <td>00:13:00.920</td>\n",
       "      <td>00:13:04.400</td>\n",
       "      <td>of what the genetics is like and the real,</td>\n",
       "      <td>episode_076</td>\n",
       "      <td>John Hopfield： Physics View of the Mind and Ne...</td>\n",
       "      <td>41</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>40</th>\n",
       "      <td>00:40:42.800</td>\n",
       "      <td>00:40:45.280</td>\n",
       "      <td>After all, all the things your brain does,</td>\n",
       "      <td>episode_076</td>\n",
       "      <td>John Hopfield： Physics View of the Mind and Ne...</td>\n",
       "      <td>42</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>41</th>\n",
       "      <td>01:08:23.320</td>\n",
       "      <td>01:08:26.320</td>\n",
       "      <td>us weird descendants of apes.</td>\n",
       "      <td>episode_076</td>\n",
       "      <td>John Hopfield： Physics View of the Mind and Ne...</td>\n",
       "      <td>43</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>42</th>\n",
       "      <td>00:12:15.400</td>\n",
       "      <td>00:12:20.400</td>\n",
       "      <td>And since I can't see the evolutionary proces...</td>\n",
       "      <td>episode_076</td>\n",
       "      <td>John Hopfield： Physics View of the Mind and Ne...</td>\n",
       "      <td>44</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>43</th>\n",
       "      <td>01:00:06.200</td>\n",
       "      <td>01:00:08.440</td>\n",
       "      <td>in the near or maybe even far future</td>\n",
       "      <td>episode_076</td>\n",
       "      <td>John Hopfield： Physics View of the Mind and Ne...</td>\n",
       "      <td>45</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           start           end  \\\n",
       "39  00:13:00.920  00:13:04.400   \n",
       "40  00:40:42.800  00:40:45.280   \n",
       "41  01:08:23.320  01:08:26.320   \n",
       "42  00:12:15.400  00:12:20.400   \n",
       "43  01:00:06.200  01:00:08.440   \n",
       "\n",
       "                                                 text        fname  \\\n",
       "39         of what the genetics is like and the real,  episode_076   \n",
       "40         After all, all the things your brain does,  episode_076   \n",
       "41                      us weird descendants of apes.  episode_076   \n",
       "42   And since I can't see the evolutionary proces...  episode_076   \n",
       "43               in the near or maybe even far future  episode_076   \n",
       "\n",
       "                                           audio_name  audio_idx  is_lex  \\\n",
       "39  John Hopfield： Physics View of the Mind and Ne...         41     0.0   \n",
       "40  John Hopfield： Physics View of the Mind and Ne...         42     0.0   \n",
       "41  John Hopfield： Physics View of the Mind and Ne...         43     1.0   \n",
       "42  John Hopfield： Physics View of the Mind and Ne...         44     0.0   \n",
       "43  John Hopfield： Physics View of the Mind and Ne...         45     1.0   \n",
       "\n",
       "    preds  \n",
       "39    0.0  \n",
       "40    0.0  \n",
       "41    1.0  \n",
       "42    0.0  \n",
       "43    1.0  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# predictions for one fold\n",
    "fold_idx = 0\n",
    "foldk_preds = pred_dfs[fold_idx]\n",
    "\n",
    "display('Metrics ', test_metrics.loc[fold_idx])\n",
    "display('Preds', foldk_preds.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Listen to audio_idx audio clip and explore the predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>start</th>\n",
       "      <th>end</th>\n",
       "      <th>text</th>\n",
       "      <th>fname</th>\n",
       "      <th>audio_name</th>\n",
       "      <th>audio_idx</th>\n",
       "      <th>is_lex</th>\n",
       "      <th>preds</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>195</th>\n",
       "      <td>00:04:00.080</td>\n",
       "      <td>00:04:01.080</td>\n",
       "      <td>anything like that.</td>\n",
       "      <td>episode_035</td>\n",
       "      <td>Jeremy Howard： fast.ai Deep Learning Courses a...</td>\n",
       "      <td>198</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>421</th>\n",
       "      <td>00:14:49.900</td>\n",
       "      <td>00:14:51.500</td>\n",
       "      <td>from the start and from the end.</td>\n",
       "      <td>episode_120</td>\n",
       "      <td>François Chollet： Measures of Intelligence ｜ L...</td>\n",
       "      <td>432</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>55</th>\n",
       "      <td>00:36:51.640</td>\n",
       "      <td>00:36:53.240</td>\n",
       "      <td>on a form of the Boltzmann machine</td>\n",
       "      <td>episode_076</td>\n",
       "      <td>John Hopfield： Physics View of the Mind and Ne...</td>\n",
       "      <td>57</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>629</th>\n",
       "      <td>00:14:12.400</td>\n",
       "      <td>00:14:18.720</td>\n",
       "      <td>will be on the machine side. This is just, th...</td>\n",
       "      <td>episode_49</td>\n",
       "      <td>Elon Musk： Neuralink, AI, Autopilot, and the P...</td>\n",
       "      <td>646</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>583</th>\n",
       "      <td>00:00:00.000</td>\n",
       "      <td>00:00:07.280</td>\n",
       "      <td>The following is a conversation with Elon Mus...</td>\n",
       "      <td>episode_49</td>\n",
       "      <td>Elon Musk： Neuralink, AI, Autopilot, and the P...</td>\n",
       "      <td>600</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>275</th>\n",
       "      <td>01:20:23.740</td>\n",
       "      <td>01:20:24.940</td>\n",
       "      <td>Who cares about birthdays?</td>\n",
       "      <td>episode_189</td>\n",
       "      <td>David Sinclair： Extending the Human Lifespan B...</td>\n",
       "      <td>281</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>168</th>\n",
       "      <td>00:05:02.960</td>\n",
       "      <td>00:05:03.760</td>\n",
       "      <td>even love to do.</td>\n",
       "      <td>episode_098</td>\n",
       "      <td>Kate Darling： Social Robotics ｜ Lex Fridman Po...</td>\n",
       "      <td>171</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>621</th>\n",
       "      <td>00:11:22.000</td>\n",
       "      <td>00:11:28.720</td>\n",
       "      <td>So you can see the consequences of if you fir...</td>\n",
       "      <td>episode_49</td>\n",
       "      <td>Elon Musk： Neuralink, AI, Autopilot, and the P...</td>\n",
       "      <td>638</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>630</th>\n",
       "      <td>00:14:44.720</td>\n",
       "      <td>00:14:51.040</td>\n",
       "      <td>It's not the cortex that's steering the monke...</td>\n",
       "      <td>episode_49</td>\n",
       "      <td>Elon Musk： Neuralink, AI, Autopilot, and the P...</td>\n",
       "      <td>647</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>585</th>\n",
       "      <td>00:00:13.120</td>\n",
       "      <td>00:00:20.720</td>\n",
       "      <td>sequel of all time, Godfather Part 2. As many...</td>\n",
       "      <td>episode_49</td>\n",
       "      <td>Elon Musk： Neuralink, AI, Autopilot, and the P...</td>\n",
       "      <td>602</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>608</th>\n",
       "      <td>00:06:29.040</td>\n",
       "      <td>00:06:35.920</td>\n",
       "      <td>in hope of finding an algorithmic or maybe a ...</td>\n",
       "      <td>episode_49</td>\n",
       "      <td>Elon Musk： Neuralink, AI, Autopilot, and the P...</td>\n",
       "      <td>625</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>588</th>\n",
       "      <td>00:00:39.600</td>\n",
       "      <td>00:00:44.880</td>\n",
       "      <td>most of us don't question the way things are ...</td>\n",
       "      <td>episode_49</td>\n",
       "      <td>Elon Musk： Neuralink, AI, Autopilot, and the P...</td>\n",
       "      <td>605</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>196</th>\n",
       "      <td>00:50:49.200</td>\n",
       "      <td>00:50:51.080</td>\n",
       "      <td>Which is very tough, so you have to.</td>\n",
       "      <td>episode_035</td>\n",
       "      <td>Jeremy Howard： fast.ai Deep Learning Courses a...</td>\n",
       "      <td>199</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>292</th>\n",
       "      <td>01:31:47.700</td>\n",
       "      <td>01:31:49.740</td>\n",
       "      <td>we can recreate that,</td>\n",
       "      <td>episode_189</td>\n",
       "      <td>David Sinclair： Extending the Human Lifespan B...</td>\n",
       "      <td>299</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>611</th>\n",
       "      <td>00:06:49.280</td>\n",
       "      <td>00:06:54.640</td>\n",
       "      <td>most interesting, as a civilization that we s...</td>\n",
       "      <td>episode_49</td>\n",
       "      <td>Elon Musk： Neuralink, AI, Autopilot, and the P...</td>\n",
       "      <td>628</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            start           end  \\\n",
       "195  00:04:00.080  00:04:01.080   \n",
       "421  00:14:49.900  00:14:51.500   \n",
       "55   00:36:51.640  00:36:53.240   \n",
       "629  00:14:12.400  00:14:18.720   \n",
       "583  00:00:00.000  00:00:07.280   \n",
       "275  01:20:23.740  01:20:24.940   \n",
       "168  00:05:02.960  00:05:03.760   \n",
       "621  00:11:22.000  00:11:28.720   \n",
       "630  00:14:44.720  00:14:51.040   \n",
       "585  00:00:13.120  00:00:20.720   \n",
       "608  00:06:29.040  00:06:35.920   \n",
       "588  00:00:39.600  00:00:44.880   \n",
       "196  00:50:49.200  00:50:51.080   \n",
       "292  01:31:47.700  01:31:49.740   \n",
       "611  00:06:49.280  00:06:54.640   \n",
       "\n",
       "                                                  text        fname  \\\n",
       "195                                anything like that.  episode_035   \n",
       "421                   from the start and from the end.  episode_120   \n",
       "55                  on a form of the Boltzmann machine  episode_076   \n",
       "629   will be on the machine side. This is just, th...   episode_49   \n",
       "583   The following is a conversation with Elon Mus...   episode_49   \n",
       "275                         Who cares about birthdays?  episode_189   \n",
       "168                                   even love to do.  episode_098   \n",
       "621   So you can see the consequences of if you fir...   episode_49   \n",
       "630   It's not the cortex that's steering the monke...   episode_49   \n",
       "585   sequel of all time, Godfather Part 2. As many...   episode_49   \n",
       "608   in hope of finding an algorithmic or maybe a ...   episode_49   \n",
       "588   most of us don't question the way things are ...   episode_49   \n",
       "196               Which is very tough, so you have to.  episode_035   \n",
       "292                              we can recreate that,  episode_189   \n",
       "611   most interesting, as a civilization that we s...   episode_49   \n",
       "\n",
       "                                            audio_name  audio_idx  is_lex  \\\n",
       "195  Jeremy Howard： fast.ai Deep Learning Courses a...        198     0.0   \n",
       "421  François Chollet： Measures of Intelligence ｜ L...        432     1.0   \n",
       "55   John Hopfield： Physics View of the Mind and Ne...         57     0.0   \n",
       "629  Elon Musk： Neuralink, AI, Autopilot, and the P...        646     0.0   \n",
       "583  Elon Musk： Neuralink, AI, Autopilot, and the P...        600     1.0   \n",
       "275  David Sinclair： Extending the Human Lifespan B...        281     0.0   \n",
       "168  Kate Darling： Social Robotics ｜ Lex Fridman Po...        171     1.0   \n",
       "621  Elon Musk： Neuralink, AI, Autopilot, and the P...        638     0.0   \n",
       "630  Elon Musk： Neuralink, AI, Autopilot, and the P...        647     0.0   \n",
       "585  Elon Musk： Neuralink, AI, Autopilot, and the P...        602     1.0   \n",
       "608  Elon Musk： Neuralink, AI, Autopilot, and the P...        625     1.0   \n",
       "588  Elon Musk： Neuralink, AI, Autopilot, and the P...        605     1.0   \n",
       "196  Jeremy Howard： fast.ai Deep Learning Courses a...        199     1.0   \n",
       "292  David Sinclair： Extending the Human Lifespan B...        299     0.0   \n",
       "611  Elon Musk： Neuralink, AI, Autopilot, and the P...        628     1.0   \n",
       "\n",
       "     preds  \n",
       "195    0.0  \n",
       "421    1.0  \n",
       "55     0.0  \n",
       "629    0.0  \n",
       "583    1.0  \n",
       "275    0.0  \n",
       "168    1.0  \n",
       "621    0.0  \n",
       "630    0.0  \n",
       "585    1.0  \n",
       "608    1.0  \n",
       "588    1.0  \n",
       "196    1.0  \n",
       "292    0.0  \n",
       "611    1.0  "
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "foldk_preds.sample(15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "whisper",
   "language": "python",
   "name": "whisper"
  },
  "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.9.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}