allosaurus-2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "# import os\n",
    "os.environ['TRANSFORMERS_CACHE'] = '/data/cg46773/transformer'\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"   # see issue #152\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"7\" \n",
    "from tqdm.notebook import tqdm\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/data/cg46773/miniconda3/envs/base3.7/lib/python3.7/site-packages/torchaudio/_internal/module_utils.py:87: UserWarning: Failed to import soundfile. 'soundfile' backend is not available.\n",
      "  warnings.warn(\"Failed to import soundfile. 'soundfile' backend is not available.\")\n"
     ]
    }
   ],
   "source": [
    "# import os\n",
    "import pandas as pd\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torchaudio\n",
    "import torchaudio.transforms as transforms\n",
    "import argparse\n",
    "import torch\n",
    "from torch import nn, optim\n",
    "from torch.nn import functional as F\n",
    "from torch.nn.utils.rnn import pad_sequence, pack_padded_sequence, pad_packed_sequence\n",
    "from torch.utils.data.sampler import Sampler, BatchSampler, SubsetRandomSampler\n",
    "import pickle\n",
    "import numpy as np\n",
    "import python_speech_features \n",
    "import scipy.io.wavfile as wav\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ast import literal_eval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument('--batch-size', type=int, default=1, metavar='N',\n",
    "                    help='input batch size for training (default: 10)')\n",
    "parser.add_argument('--epochs', type=int, default=10, metavar='N',\n",
    "                    help='number of epochs to train (default: 10)')\n",
    "parser.add_argument('--no-cuda', action='store_true', default=False,\n",
    "                    help='disables CUDA training')\n",
    "parser.add_argument('--no-mps', action='store_true', default=False,\n",
    "                    help='disables macOS GPU training')\n",
    "parser.add_argument('--seed', type=int, default=1, metavar='S',\n",
    "                    help='random seed (default: 1)')\n",
    "parser.add_argument('--log-interval', type=int, default=300, metavar='N',\n",
    "                    help='how many batches to wait before logging training status')\n",
    "parser.add_argument('--silence', type=int, default=1, choices=[0,1],metavar='N',\n",
    "                    help='If silence==1 then _ char is put in begin/end of phm_seq, if ')\n",
    "parser.add_argument('--hidden_size', type=int, default=10, metavar='N',\n",
    "                    help='hidden size of the LSTM')\n",
    "parser.add_argument('--num_layers', type=int, default=1, metavar='N',\n",
    "                    help='num layers of  LSTM')\n",
    "parser.add_argument('--lr', type=float, default=1e-3,\n",
    "                    help='learning rate')\n",
    "args = parser.parse_args(args=[])\n",
    "args.cuda = not args.no_cuda and torch.cuda.is_available()\n",
    "# args.mps = not args.no_mps and torch.backends.mps.is_available()\n",
    "\n",
    "if args.cuda:\n",
    "    device = torch.device(\"cuda\")\n",
    "elif args.mps:\n",
    "    device = torch.device(\"mps\")\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 1)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "args.num_layers,args.silence"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialize Locale, Phone, Phoneme, Mapping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "allophone_data = pickle.load(open(\"allophone_mappings_new.pkl\",\"rb\"))\n",
    "list(allophone_data.keys())\n",
    "num_universal_phones = len(allophone_data['id to universal_phone'])\n",
    "locale_to_id={locale:idx for idx,locale in enumerate(allophone_data['language_locale to (language-local)id to phoneme'].keys())}\n",
    "id_to_locale={idx:locale for locale,idx in locale_to_id.items()}\n",
    "# maps locale to len of phonemes\n",
    "locale_to_num_phonemes={locale:len(allophone_data['language_locale to (language-local)id to phoneme'][locale]) for locale in locale_to_id}\n",
    "lang_to_phoneme_to_id = allophone_data['language_locale to phoneme to (language-local)id']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "642"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "num_universal_phones"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['id to universal_phone',\n",
       " 'universal_phone to id',\n",
       " 'locale to langcode mapping',\n",
       " 'langcode to locale mapping',\n",
       " 'language_locale to (language-local)id to phoneme',\n",
       " 'language_locale to phoneme to (language-local)id',\n",
       " '(Allophone mappings) language_locale to (language-local)id to universal_phone_ids']"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(allophone_data.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "    This cell initializes the signature matrix\n",
    "'''\n",
    "locale_to_allophone_mappings_id = allophone_data['(Allophone mappings) language_locale to (language-local)id to universal_phone_ids']\n",
    "#locale_to_allophone_mappings_array = {loc:np.zeros((locale_to_num_phonemes[loc], num_universal_phones)) for loc in locale_to_id}\n",
    "signature_matrix = {loc:np.zeros((locale_to_num_phonemes[loc], num_universal_phones)) for loc in locale_to_id}\n",
    "\n",
    "for lang in signature_matrix:\n",
    "    signature_array = signature_matrix[lang]\n",
    "    # print(locale_to_allophone_mappings_id[lang])\n",
    "    for q, phones in locale_to_allophone_mappings_id[lang].items():\n",
    "        for p in phones:\n",
    "            # q (Phoneme) x P (phone)\n",
    "            signature_array[q, p] = 1\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialize Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "#https://pytorch.org/audio/stable/generated/torchaudio.transforms.MFCC.html#torchaudio.transforms.MFCC\n",
    "# def MFCC(filename=\"sample.wav\"):\n",
    "def MFCC(filename=\"output1.wav\"):\n",
    "    waveform, sample_rate = torchaudio.load(filename, normalize=True)\n",
    "    transform = transforms.MFCC(\n",
    "        sample_rate=sample_rate,\n",
    "        n_mfcc=40,\n",
    "        #https://pytorch.org/audio/stable/generated/torchaudio.transforms.MelSpectrogram.html#torchaudio.transforms.MelSpectrogram\n",
    "        melkwargs={\"n_fft\": 400, #\"hop_length\": 160,\n",
    "                   \"n_mels\": 40, \"center\": False}\n",
    "                  # \"win_length\": 0.025},\n",
    "    )\n",
    "    mfcc = transform(waveform)\n",
    "    return mfcc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.signal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "(rate,sig) = wav.read('data/hin/audio/hin-008-001.wav')\n",
    "# mfcc = torch.Tensor(\n",
    "# print(len(sig))\n",
    "signal_resampled=scipy.signal.resample(sig,(len(sig)*16000)//rate)\n",
    "# python_speech_features.mfcc(signal_resampled,16000,numcep=40, nfilt=80)\n",
    "# )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "44100"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MozillaDataset(Dataset):\n",
    "    def __init__(self, args, csv_file, type_,pt_file,audio_dir):\n",
    "        data = pd.read_csv(csv_file,keep_default_na=False, na_values=[''])\n",
    "        data=data[data[\"type\"]==type_].reset_index(drop=True)\n",
    "        self.length = len(data)\n",
    "        audio_paths=audio_dir + '/' + data[\"path\"]\n",
    "        if not data[\"path\"].iloc[0].endswith(\".wav\"):\n",
    "            audio_paths=audio_paths + \".wav\"\n",
    "        self.x_mfcc = []\n",
    "        #JORDI-- torch-mfcc computation\\n\",\n",
    "        # for audio_path in audio_paths:\n",
    "        #     mfcc = MFCC(audio_path)\n",
    "        #     #so that the final dim is: 1 x time x n_mfcc\\n\",\n",
    "        #     mfcc = mfcc.permute(0,2,1)\n",
    "        #     mfcc = torch.squeeze(mfcc)\n",
    "        #     self.x_mfcc.append(mfcc)\n",
    "        \n",
    "        #Chitrank-- via python-speech-features\",\n",
    "        want_to_load_from_pt_file=True\n",
    "        if want_to_load_from_pt_file==True:\n",
    "            self.x_mfcc=torch.load(pt_file)\n",
    "        else:    \n",
    "            for audio_path in tqdm(audio_paths):\n",
    "                (rate,sig) = wav.read(audio_path)\n",
    "                if rate>16000:\n",
    "                    sig=scipy.signal.resample(sig,(len(sig)*16000)//rate)\n",
    "                    rate=16000\n",
    "                mfcc = torch.Tensor(python_speech_features.mfcc(sig,rate,numcep=40, nfilt=80))\n",
    "                # mfcc size is already T x F(=40)\n",
    "                self.x_mfcc.append(mfcc)\n",
    "            torch.save(self.x_mfcc,f=pt_file)\n",
    "        \n",
    "\n",
    "        self.lang = [locale_to_id[i] for i in data[\"langcode\"].values]\n",
    "        self.start_pos_data = [0]\n",
    "        for idx in range(1,len(self.lang)):\n",
    "            if self.lang[idx] != self.lang[idx-1]:\n",
    "                # append when the language changes \n",
    "                self.start_pos_data.append(idx)\n",
    "\n",
    "        if args.silence==1:\n",
    "            phm_seq_raw = [['_']+literal_eval(i)+['_'] for i in data[\"phoneme_sequence\"].values]\n",
    "        else:\n",
    "            phm_seq_raw = [literal_eval(i) for i in data[\"phoneme_sequence\"].values]\n",
    "        phm_seq_raw= [ torch.Tensor([lang_to_phoneme_to_id[id_to_locale[self.lang[idx]]][phm] for phm in phm_seq])\n",
    "                      for idx,phm_seq in enumerate(phm_seq_raw)]\n",
    "        # self.raw = phm_seq_raw\n",
    "\n",
    "        self.phm_seq=phm_seq_raw\n",
    "#         for i in range(self.length):\n",
    "#             sequence_ = []\n",
    "#             lang_ = self.lang[i]\n",
    "#\n",
    "#             # some phonemes occupies two character space, we have to account for that\n",
    "#             j = 0\n",
    "#             special_phm = [\"ɕ\"]\n",
    "#             while j < len(phm_seq_raw[i]):\n",
    "#                 phm = phm_seq_raw[i][j]\n",
    "#                 if phm == \"ɕ\" and lang_ == \"ru\":\n",
    "#                     sequence_.append(lang_to_phoneme_to_id[lang_][\"ɕː\"])\n",
    "#                     j += 2\n",
    "#                 elif j < len(phm_seq_raw[i]) - 1 and phm_seq_raw[i][j+1] not in lang_to_phoneme_to_id[lang_] and phm_seq_raw[i][j+1] not in special_phm:\n",
    "#                     next_phm = phm_seq_raw[i][j+1]\n",
    "#                     new_phm = phm + next_phm\n",
    "#                     sequence_.append(lang_to_phoneme_to_id[lang_][new_phm])\n",
    "#                     j += 2\n",
    "#                 else:\n",
    "#                     sequence_.append(lang_to_phoneme_to_id[lang_][phm])\n",
    "#                     j += 1\n",
    "#             self.phm_seq.append(torch.Tensor(sequence_))\n",
    "        #self.phm_seq = torch.Tensor(self.phm_seq)\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.length\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return {\"x_mfcc\": self.x_mfcc[idx], \"lang\": self.lang[idx], \"phm_seq\": self.phm_seq[idx],}\n",
    "                # \"raw\":self.raw[idx]}\n",
    "\n",
    "mozilla_dataset_train = MozillaDataset(args,csv_file='validated_1000.csv',\n",
    "                                       type_=\"train\",\n",
    "                                       pt_file=\"train_x_mfcc_1000.pt\",\n",
    "                                       audio_dir='mozilla/wav_1000_100_1000/')\n",
    "mozilla_dataset_dev = MozillaDataset(args,csv_file='validated_1000.csv',\n",
    "                                     type_=\"dev\",\n",
    "                                     pt_file=\"dev_x_mfcc_1000.pt\",\n",
    "                                     audio_dir='mozilla/wav_1000_100_1000/')\n",
    "mozilla_dataset_test = MozillaDataset(args,csv_file='validated_1000.csv',\n",
    "                                      type_=\"test\",\n",
    "                                      pt_file=\"test_x_mfcc_1000.pt\",\n",
    "                                      audio_dir='mozilla/wav_1000_100_1000/')\n",
    "ucla_dataset_test = MozillaDataset(args,csv_file='ucla_utterances.csv',\n",
    "                                      type_=\"test\",\n",
    "                                      pt_file=\"ucla_test_x_mfcc_1000.pt\",\n",
    "                                      audio_dir='data/')\n",
    "# print(\"length of dataset: \", len(mozilla_dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collate_fn(batch):\n",
    "    # print(\"printing all the raw phonemes in every batch: \")\n",
    "    # for d in batch:\n",
    "    #     print(d[\"raw\"])\n",
    "    x_mfcc_list = [d['x_mfcc'] for d in batch]\n",
    "    # print(\"shape of first in x_mfcc_list\", x_mfcc_list[0].shape)\n",
    "    # print(\"shape of first in x_mfcc_list\", x_mfcc_list[1].shape)\n",
    "    # print(\"shape of first in x_mfcc_list\", x_mfcc_list[2].shape)\n",
    "    # print(\"shape of first in x_mfcc_list\", x_mfcc_list[3].shape)\n",
    "    x_mfcc_tensor = pad_sequence(x_mfcc_list, batch_first=True, padding_value=0)\n",
    "    #x_mfcc_lengths = sorted([len(x) for x in x_mfcc_list], reverse=True)\n",
    "    x_mfcc_lengths = [len(x) for x in x_mfcc_list]\n",
    "    #print(\"x_mfcc_lengths: \", x_mfcc_lengths)\n",
    "    x_mfcc_tensor = pack_padded_sequence(x_mfcc_tensor, lengths= x_mfcc_lengths, batch_first=True, enforce_sorted = False)\n",
    "    lang_list = [d['lang'] for d in batch]\n",
    "    #lang_list = torch.tensor(lang_list)\n",
    "    phm_seq_list = [d['phm_seq'] for d in batch]\n",
    "    #print(\"phm_seq_list\", phm_seq_list)\n",
    "    phm_seq_tensor = pad_sequence(phm_seq_list, batch_first=True, padding_value=0).to(int)\n",
    "    phm_seq_lengths = [len(x) for x in phm_seq_list]\n",
    "    # print(\"phm_seq_lengths: \", phm_seq_lengths)\n",
    "    #phm_seq_pack_padded_sequence = pack_padded_sequence(phm_seq_tensor, lengths= phm_seq_lengths, batch_first=True, enforce_sorted = False)\n",
    "    # print(\"packed padded sequence: \", x_mfcc_tensor)\n",
    "    return {\n",
    "        \"x_mfcc\": x_mfcc_tensor, # packed padded tensor\n",
    "        \"lang\": lang_list, # python list\n",
    "        \"phm_seq\": phm_seq_tensor, # padded tensor\n",
    "        \"x_mfcc_lengths\": torch.Tensor(x_mfcc_lengths).to(int), # sorted in decreasing order\n",
    "        \"phm_seq_lengths\": torch.Tensor(phm_seq_lengths).to(int) # sorted in decreasing order\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for i in range(len(mozilla_dataset)):\n",
    "#     sample = mozilla_dataset[i]\n",
    "#     print(i, sample['x_mfcc'].shape, sample['lang'], sample[\"phm_seq\"].shape)\n",
    "\n",
    "#     if i == 1:\n",
    "#         break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BucketSampler(Sampler):\n",
    "    '''\n",
    "        creates a bucket sampler where we sample randomly from the same language without replacement\n",
    "    '''\n",
    "    def __init__(self, dataset, batch_size, generator=None) -> None:\n",
    "        self.dataset = dataset\n",
    "        self.batch_size = batch_size\n",
    "        self.generator = generator\n",
    "        start_pos_data = self.dataset.start_pos_data\n",
    "        start_end_indices = []\n",
    "        for i in range(len(start_pos_data) - 1):\n",
    "            start_end_indices.append((start_pos_data[i], start_pos_data[i+1]))\n",
    "        start_end_indices.append((start_pos_data[-1], len(self.dataset)))\n",
    "        ranges  = [range(start, end) for start, end in start_end_indices]\n",
    "        subset_samplers = [SubsetRandomSampler(range_, generator=generator) for range_ in ranges]\n",
    "        self.samplers = [\n",
    "            BatchSampler(subset_sampler, batch_size, drop_last=False) for subset_sampler in subset_samplers\n",
    "        ]\n",
    "        self._len = 0\n",
    "        for sampler in self.samplers:\n",
    "            self._len += len(sampler)\n",
    "        \n",
    "    def __iter__(self):\n",
    "        iterators = [iter(sampler) for sampler in self.samplers]\n",
    "        while iterators:\n",
    "            randint = torch.randint(0, len(iterators),size=(1,), generator=self.generator)[0]\n",
    "            try:\n",
    "                yield next(iterators[randint])\n",
    "            except StopIteration:\n",
    "                iterators.pop(randint)\n",
    "    def __len__(self):\n",
    "        return self._len\n",
    "bucketSampler = BucketSampler(mozilla_dataset_train, batch_size = 4)\n",
    "dataloader = DataLoader(mozilla_dataset_train, batch_sampler = bucketSampler, collate_fn=collate_fn)#, num_workers=4)\n",
    "for i in dataloader:\n",
    "    # print(\"mfcc: \", i[\"x_mfcc\"].data.shape)\n",
    "    # print(\"language: \",i[\"lang\"])\n",
    "    # print(\"phm_seq: \", i[\"phm_seq\"])\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "       1., 1., 1., 1., 1., 1., 1., 1., 1., 0., 1., 0., 0., 0., 1., 0.])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.sum(signature_matrix[\"daa\"],axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Allosaurus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(device(type='cuda'), 10)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device,args.hidden_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_langs=set(mozilla_dataset_train.lang)\n",
    "all_langs=all_langs.union(set(mozilla_dataset_dev.lang))\n",
    "all_langs=all_langs.union(set(mozilla_dataset_test.lang))\n",
    "all_langs=all_langs.union(set(ucla_dataset_test.lang))\n",
    "# all_langs=set(mozilla_dataset_train.langs)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'all_langs' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_25201/719380581.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mall_langs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m: name 'all_langs' is not defined"
     ]
    }
   ],
   "source": [
    "len(all_langs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.3832,  0.4862,  0.7070,  0.2622, -0.3949],\n",
       "        [ 0.3594,  0.3629,  0.1502, -0.4447,  0.4887],\n",
       "        [-0.6370, -0.1855, -0.5306,  0.5091,  1.3228]])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w = torch.empty(3, 5)\n",
    "nn.init.xavier_normal_(w)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "1-- [B x T x 2*H] [2*H x A] [A x P] \n",
    "# 2-- [B x T x H] [P x 2*H]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "ename": "SyntaxError",
     "evalue": "invalid syntax (1509511437.py, line 9)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;36m  File \u001b[0;32m\"/tmp/ipykernel_25201/1509511437.py\"\u001b[0;36m, line \u001b[0;32m9\u001b[0m\n\u001b[0;31m    def\u001b[0m\n\u001b[0m        ^\u001b[0m\n\u001b[0;31mSyntaxError\u001b[0m\u001b[0;31m:\u001b[0m invalid syntax\n"
     ]
    }
   ],
   "source": [
    "class Heirachical(nn.Module):\n",
    "    def __init__(self, args, num_attrs, num_attrtypes, attrs2attrtype):\n",
    "        \"\"\"\n",
    "            attrs2attrtype-- 1-0 matrix of shape [A x AT]\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.allosaurus_layer_matrix[lang_id]=torch.nn.Parameter(torch.Tensor(signature_matrix[locale_]).to(device),requires_grad=True)\n",
    "        if args.original_heir==1:\n",
    "            final_phone_emb_size=2*args.hidden_size\n",
    "            self.attr_embeds=torch.nn.empty(num_attrs, final_phone_emb_size) # It's a (A x 2*H) embeddings matrix \n",
    "        else:\n",
    "            final_phone_emb_size=2*args.hidden_size\n",
    "            attr_emb_size=math.ceil(final_phone_emb_size/num_attrtypes)\n",
    "            self.attr_embeds=torch.nn.empty(num_attrs, attr_emb_size) # It's a (A x [2*H/AT]) embeddings matrix \n",
    "        self.attrs_combination\n",
    "        self.allosaurus_model=Allosaurus()\n",
    "    def forward(self, X_mfcc, lang_id):\n",
    "        \"\"\"\n",
    "            \n",
    "        \"\"\"\n",
    "        packed_padded_sequence_output, (h_n, c_n) = self.encoder(X_mfcc) # returns (B x T x h)\n",
    "        seq_unpacked, lens_unpacked = pad_packed_sequence(packed_padded_sequence_output, batch_first=True)\n",
    "        # the hidden space is times 2 because it is bidirectional\n",
    "        # print(\"shape of padded sequence in hidden space (Batch x Time x (hidden x2 because bi directional)): \", seq_unpacked.shape, flush=True)\n",
    "        # print(\"lens_unpacked: \", lens_unpacked, flush=True)\n",
    "        # returns (B x T x N)\n",
    "        phone_distribution = self.phone_distribution_transf_matrix(seq_unpacked)  \n",
    "        # print(\"shape of phone_distribution (Batch x Time x num_universal_phones): \",phone_distribution.shape ,flush=True)\n",
    "        phoneme_distrib = self.compute_phoneme_distribution(phone_distribution, lang_id)\n",
    "        # print(\"shape of phoneme after amax (Batch x Time x ?): \",phoneme_distrib.shape ,flush=True)\n",
    "        #[compute_phoneme_distribution(allosaurus_layer_matrix[i],langs[i])\n",
    "         #                for i in range(len(langs))] # (return )\n",
    "        return phoneme_distrib\n",
    ".\n",
    "# lang -> phoneme id mapping        \n",
    "        \n",
    "        \n",
    "\n",
    "class Allosaurus(nn.Module):\n",
    "    def __init__(self, args,num_universal_phones, signature_matrix):\n",
    "        '''\n",
    "            signature_matrix: {lang (str): Mi x N}; locale_to_allophone_mappings_array\n",
    "        '''\n",
    "        super().__init__()\n",
    "        self.num_universal_phones = num_universal_phones  # In short N\n",
    "  \n",
    "        self.encoder=Encoder(n_layers=args.num_layers,hidden_dim=args.hidden_size,input_size=40) # In short lstm has (n_lay, h, #mfcc) \n",
    "        self.phone_distribution_transf_matrix=nn.Linear(2 * args.hidden_size, self.num_universal_phones) # It's a (h x N) transformation matrix \n",
    "        self.allosaurus_layer_matrix = {}\n",
    "        self.allosaurus_layer_matrix_orig = {}\n",
    "        for lang_id in all_langs:\n",
    "            # every such matrix is of shape (M_i x N)\n",
    "            locale_=id_to_locale[lang_id]\n",
    "            self.allosaurus_layer_matrix[lang_id]=torch.nn.Parameter(torch.Tensor(signature_matrix[locale_]).to(device),requires_grad=True)\n",
    "            self.allosaurus_layer_matrix_orig[lang_id]=torch.nn.Parameter(torch.Tensor(signature_matrix[locale_]).to(device),requires_grad=False)\n",
    "\n",
    "        self.loss_fn = torch.nn.CTCLoss(blank = 0, reduction='mean', zero_infinity=True)\n",
    "    def compute_phoneme_distribution(self, phone_distrib,lang_id):\n",
    "        # phone_distrib-- ( B x T x N), \n",
    "        # lang-- language index\n",
    "        # retrieve sign_matrix-- (M_i x N)\n",
    "        phoneme_phone_distrib = torch.amax(self.allosaurus_layer_matrix[lang_id]*(phone_distrib.unsqueeze(2)),dim=-1)\n",
    "\n",
    "        # print(\"Final phoneme_phone_distrib matrix(B x T x phonemes):\", phoneme_phone_distrib.shape)\n",
    "        return torch.nn.functional.log_softmax(phoneme_phone_distrib, dim=-1)\n",
    "    \n",
    "    def forward(self, X_mfcc, lang_id):\n",
    "        '''\n",
    "            frames is (B x T x F) matrix where \n",
    "            B is batch_size\n",
    "            T= time_dimension \n",
    "            F- numFeatures in every unit time \n",
    "            Langs is (B) matrix that tell us which lang_index does each batch_elemnet corresponds to.\n",
    "            --  in our special case it is a single value\n",
    "        '''\n",
    "        packed_padded_sequence_output, (h_n, c_n) = self.encoder(X_mfcc) # returns (B x T x h)\n",
    "        seq_unpacked, lens_unpacked = pad_packed_sequence(packed_padded_sequence_output, batch_first=True)\n",
    "        # the hidden space is times 2 because it is bidirectional\n",
    "        # print(\"shape of padded sequence in hidden space (Batch x Time x (hidden x2 because bi directional)): \", seq_unpacked.shape, flush=True)\n",
    "        # print(\"lens_unpacked: \", lens_unpacked, flush=True)\n",
    "        # returns (B x T x N)\n",
    "        phone_distribution = self.phone_distribution_transf_matrix(seq_unpacked)  \n",
    "        # print(\"shape of phone_distribution (Batch x Time x num_universal_phones): \",phone_distribution.shape ,flush=True)\n",
    "        phoneme_distrib = self.compute_phoneme_distribution(phone_distribution, lang_id)\n",
    "        # print(\"shape of phoneme after amax (Batch x Time x ?): \",phoneme_distrib.shape ,flush=True)\n",
    "        #[compute_phoneme_distribution(allosaurus_layer_matrix[i],langs[i])\n",
    "         #                for i in range(len(langs))] # (return )\n",
    "        return phoneme_distrib\n",
    "             \n",
    "        \n",
    "    def loss(self, y_ref, y_scores, src_length, tgt_length, lang_id, alpha=10):\n",
    "        # y_ref is of shape (B x T_o) where T_o= time_stamps in output space. \n",
    "        # y_predicted is of shape [T_inp x M_o_1, T_inp x M_o_2, ..., T_inp x M_o_B] where B is the Batch_size\n",
    "        #   where T_i= time_stamps in input space. Note T_o!=T_i \n",
    "\n",
    "        # ctc_loss= torch.mean([self.loss_fn(y_scores[i], y_ref[i]) for i in range(len(y_ref))])\n",
    "        ctc_loss= torch.mean(self.loss_fn(y_scores.transpose(0,1), y_ref, src_length, tgt_length))\n",
    "        # return ctc_loss+alpha*torch.mean([torc.norm((allosaurus_layer_matrix[i]-signature_matrix[i])) \n",
    "        #                    for i in range(len(languages_in_current_batch))])\n",
    "\n",
    "        weights = self.allosaurus_layer_matrix[lang_id]\n",
    "        weights_orig = self.allosaurus_layer_matrix_orig[lang_id]\n",
    "        a=ctc_loss\n",
    "        b=alpha * torch.square(torch.linalg.matrix_norm((weights-weights_orig)+1e-8))\n",
    "        return a+b\n",
    "        \n",
    "class Encoder(nn.Module):\n",
    "    # def __init__(self, n_layers=6,hidden_dim=1024,input_size=40):\n",
    "    def __init__(self, n_layers=1,hidden_dim=10,input_size=40):\n",
    "        super().__init__()\n",
    "\n",
    "        self.bi_LSTM=torch.nn.LSTM(input_size=input_size,num_layers=n_layers,\n",
    "                         hidden_size=hidden_dim,batch_first=True,\n",
    "                            bidirectional=True)\n",
    "    def forward(self, x):\n",
    "        # x is (B x T x F=40), output (B x T x lstm_hidden_size)\n",
    "        return self.bi_LSTM(x)\n",
    "args.batch_size=64\n",
    "# args.num_layers=6\n",
    "args.num_layers=4\n",
    "# args.hidden_size=1024\n",
    "args.hidden_size=256\n",
    "args.lr=5e-3\n",
    "universal_phone_num = len(allophone_data['id to universal_phone'])\n",
    "model = Allosaurus(args, universal_phone_num, signature_matrix).to(device)\n",
    "# optimizer = optim.Adam(model.parameters(), lr=args.lr)\n",
    "optimizer = optim.SGD(model.parameters(), lr=args.lr)\n",
    "torch.autograd.set_detect_anomaly(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4, 256)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.encoder.bi_LSTM.num_layers,model.encoder.bi_LSTM.hidden_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "class bcolors:\n",
    "    HEADER = '\\033[95m'\n",
    "    OKBLUE = '\\033[94m'\n",
    "    OKCYAN = '\\033[96m'\n",
    "    OKGREEN = '\\033[92m'\n",
    "    WARNING = '\\033[93m'\n",
    "    FAIL = '\\033[91m'\n",
    "    ENDC = '\\033[0m'\n",
    "    BOLD = '\\033[1m'\n",
    "    UNDERLINE = '\\033[4m'\n",
    "\n",
    "class EarlyStopping:\n",
    "    \"\"\"Early stops the training if validation loss doesn't improve after a given patience.\"\"\"\n",
    "    def __init__(self, log_file,patience=7, verbose=False, delta=0, path='allosrs/model/checkpoint.pt', trace_func=print):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            patience (int): How long to wait after last time validation loss improved.\n",
    "                            Default: 7\n",
    "            verbose (bool): If True, prints a message for each validation loss improvement. \n",
    "                            Default: False\n",
    "            delta (float): Minimum change in the monitored quantity to qualify as an improvement.\n",
    "                            Default: 0\n",
    "            path (str): Path for the checkpoint to be saved to.\n",
    "                            Default: 'checkpoint.pt'\n",
    "            trace_func (function): trace print function.\n",
    "                            Default: print            \n",
    "        \"\"\"\n",
    "        self.patience = patience\n",
    "        self.verbose = verbose\n",
    "        self.counter = 0\n",
    "        self.best_score = None\n",
    "        self.early_stop = False\n",
    "        self.val_loss_min = np.Inf\n",
    "        self.delta = delta\n",
    "        self.path = path\n",
    "        self.trace_func = trace_func\n",
    "        self.log_file=log_file\n",
    "    def __call__(self, val_loss, model):\n",
    "\n",
    "        score = -val_loss\n",
    "\n",
    "        if self.best_score is None:\n",
    "            self.best_score = score\n",
    "            self.save_checkpoint(val_loss, model)\n",
    "        elif score < self.best_score + self.delta:\n",
    "            self.counter += 1\n",
    "            with open(self.log_file,mode=\"a\") as wfile:\n",
    "                print(f'EarlyStopping counter: {self.counter} out of {self.patience}\\n',file=wfile)\n",
    "            self.trace_func(f'{bcolors.FAIL}EarlyStopping counter: {self.counter} out of {self.patience}{bcolors.ENDC}\\n')\n",
    "            if self.counter >= self.patience:\n",
    "                self.early_stop = True\n",
    "        else:\n",
    "            self.best_score = score\n",
    "            self.save_checkpoint(val_loss, model)\n",
    "            self.counter = 0\n",
    "\n",
    "    def save_checkpoint(self, val_loss, model):\n",
    "        '''Saves model when validation loss decrease.'''\n",
    "        if self.verbose:\n",
    "            with open(self.log_file,mode=\"a\") as wfile:\n",
    "                print(f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}).  Saving model ...\\n',file=wfile)\n",
    "            \n",
    "            self.trace_func(f'{bcolors.OKGREEN}Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}).  Saving model ...{bcolors.ENDC}\\n')\n",
    "        torch.save(model.state_dict(), self.path)\n",
    "        self.val_loss_min = val_loss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "#TODO: 1. write the validation loss loop while using the new data csv file -- change locale to language\n",
    "#TODO: 2. early stopping\n",
    "#TODO: 3. Phoneme error "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(epoch,dl,log_file):\n",
    "# training loop\n",
    "    model.train()\n",
    "    train_loss = 0\n",
    "    # \"x_mfcc\": x_mfcc_tensor, # packed padded sequence\n",
    "    # \"lang\": lang_list, # python list\n",
    "    # \"phm_seq\": phm_seq_tensor, # packed padded tensor\n",
    "    # \"x_mfcc_lengths\": x_mfcc_lengths, # sorted in decreasing order\n",
    "    # \"phm_seq_lengths\": phm_seq_lengths # sorted in decreasing order\n",
    "    batch_idx=0\n",
    "    predictions=[]\n",
    "    for data in tqdm(dl):\n",
    "        # break\n",
    "        batch_idx+=1\n",
    "        #X_mfcc = data[\"x_mfcc\"].to(device)\n",
    "        X_mfcc = data[\"x_mfcc\"]\n",
    "        lang = data[\"lang\"]\n",
    "        phm_seq = data[\"phm_seq\"]\n",
    "        x_mfcc_lengths = data[\"x_mfcc_lengths\"]\n",
    "        phm_seq_lengths = data[\"phm_seq_lengths\"]\n",
    "        batch_size=len(lang)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        # print(\"the language of this batch is: \", lang[0], flush=True)\n",
    "        # print(\"batch_idx\", batch_idx, flush=True)\n",
    "        # we are assuming one language per batch here\n",
    "        y_predicted = model(X_mfcc.to(device), lang[0])\n",
    "        loss = model.loss(phm_seq.to(device), y_predicted, x_mfcc_lengths.to(device), phm_seq_lengths.to(device), lang[0])\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        train_loss += batch_size*loss.detach().cpu().item()\n",
    "        if batch_idx % args.log_interval == 0:\n",
    "            # print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
    "            #     epoch, batch_idx * len(lang), len(train_loader.dataset),\n",
    "            #     100. * batch_idx / len(train_loader),\n",
    "            #     loss.item() / len(data)))\n",
    "            print('\\t\\tTrain Epoch: {} iter:{} \\tLoss: {:.6f}'.format(\n",
    "                epoch, batch_idx,loss.item() / len(data)))\n",
    "            with open(log_file,mode=\"a\") as wfile:\n",
    "                print('\\t\\tTrain Epoch: {} iter:{} \\tLoss: {:.6f}'.format(\n",
    "                    epoch, batch_idx,loss.item() / len(data)),file=wfile)\n",
    "    if epoch%1==0:\n",
    "        print('Train Epoch: {} ====> Average loss: {:.4f}'.format(\n",
    "          epoch, train_loss / len(dl.dataset)))\n",
    "        with open(log_file,mode=\"a\") as wfile:\n",
    "            print('Train Epoch: {} ====> Average loss: {:.4f}'.format(\n",
    "              epoch, train_loss / len(dl.dataset)),file=wfile)\n",
    "            \n",
    "    return train_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test_loader = DataLoader(mozilla_dataset, batch_sampler = bucketSampler, collate_fn=collate_fn)\n",
    "def evaluate(epoch,dl,log_file):\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    with torch.no_grad():\n",
    "        for data in tqdm(dl):\n",
    "            # break\n",
    "            X_mfcc = data[\"x_mfcc\"]\n",
    "            # print(X_mfcc.data.shape)\n",
    "            lang = data[\"lang\"]\n",
    "            phm_seq = data[\"phm_seq\"]\n",
    "            x_mfcc_lengths = data[\"x_mfcc_lengths\"]\n",
    "            phm_seq_lengths = data[\"phm_seq_lengths\"]\n",
    "            batch_size=len(lang)\n",
    "            y_predicted = model(X_mfcc.to(device), lang[0])\n",
    "            loss = model.loss(phm_seq.to(device), y_predicted, x_mfcc_lengths.to(device), phm_seq_lengths.to(device), lang[0])\n",
    "            test_loss += batch_size*loss.detach().cpu().item()\n",
    "    test_loss /= len(dl.dataset)\n",
    "    with open(log_file,mode=\"a\") as wfile:\n",
    "        print('Validation/Test Epoch {}====> average loss: {:.4f}'.format(epoch,test_loss),file=wfile)\n",
    "    print('Validation/Test Epoch {}====> average loss: {:.4f}'.format(epoch,test_loss))\n",
    "    return test_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MODEL SAVE NAME--- allosrs/model/lstm_nl4_hd256_lr0.005_bs20.pt\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1f8d1060aa3442f18a0348bb4ed1562c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 0====> average loss: 57.8042\n",
      "\u001b[92mValidation loss decreased (inf --> 57.804213).  Saving model ...\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9db5bb4403f44a5cbac3a5f0f18ff32b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 1 iter:300 \tLoss: 0.763395\n",
      "Train Epoch: 1 ====> Average loss: 4.8542\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b2550919954540658e75621d885d147f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 1====> average loss: 3.5306\n",
      "\u001b[92mValidation loss decreased (57.804213 --> 3.530649).  Saving model ...\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e5f20240890d4256af68d6edb5d8349a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 2 iter:300 \tLoss: 0.709110\n",
      "Train Epoch: 2 ====> Average loss: 3.4837\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0a3af72023f848c6b4e56afe2cc76195",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 2====> average loss: 3.4441\n",
      "\u001b[92mValidation loss decreased (3.530649 --> 3.444104).  Saving model ...\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0adbc1990fef4bdc9986a7decddccd10",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 3 iter:300 \tLoss: 0.791222\n",
      "Train Epoch: 3 ====> Average loss: 3.3984\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c74c147040f94042a0627498fd656fc0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 3====> average loss: 3.3203\n",
      "\u001b[92mValidation loss decreased (3.444104 --> 3.320281).  Saving model ...\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "50b488d7dc0c40eb918104655b7c08bc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 4 iter:300 \tLoss: 0.726778\n",
      "Train Epoch: 4 ====> Average loss: 3.3054\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3ef2ae21b547445e921cb1a63105fe03",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 4====> average loss: 3.2379\n",
      "\u001b[92mValidation loss decreased (3.320281 --> 3.237922).  Saving model ...\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "60c7184f8414417fb2df3d8fccea23e0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 5 iter:300 \tLoss: 0.698719\n",
      "Train Epoch: 5 ====> Average loss: 3.2250\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e68fa49733f1468498d5e73ed62bb0ea",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 5====> average loss: 3.1611\n",
      "\u001b[92mValidation loss decreased (3.237922 --> 3.161111).  Saving model ...\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bfc341350b9f4f259c73e09d2fecc8f8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 6 iter:300 \tLoss: 0.761653\n",
      "Train Epoch: 6 ====> Average loss: 3.1674\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "824301a0492640f89e07d464cbb5fbe2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 6====> average loss: 3.1315\n",
      "\u001b[92mValidation loss decreased (3.161111 --> 3.131536).  Saving model ...\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6d84e84a342d4b4088bc6a812006417c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 7 iter:300 \tLoss: 0.757591\n",
      "Train Epoch: 7 ====> Average loss: 3.1309\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0501ca8122954ae9afd1a297a4deca2a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 7====> average loss: 3.0869\n",
      "\u001b[92mValidation loss decreased (3.131536 --> 3.086874).  Saving model ...\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7a8789f8a19d494d8bf8fba0e6d37934",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 8 iter:300 \tLoss: 0.676661\n",
      "Train Epoch: 8 ====> Average loss: 3.1108\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4d9c4c99797040bf9ccdd337a170cb55",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 8====> average loss: 3.0665\n",
      "\u001b[92mValidation loss decreased (3.086874 --> 3.066508).  Saving model ...\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "65b11932a6f64245a225c945bb075f5e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 9 iter:300 \tLoss: 0.575552\n",
      "Train Epoch: 9 ====> Average loss: 3.0994\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c787611fb907463591ae76b2ef6dc5eb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 9====> average loss: 3.0571\n",
      "\u001b[92mValidation loss decreased (3.066508 --> 3.057070).  Saving model ...\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b52d5fc8939b4d30b691df9e96ee0f48",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 10 iter:300 \tLoss: 0.551283\n",
      "Train Epoch: 10 ====> Average loss: 3.0817\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5a9a69a5889f42a0b2c594303ccc4aed",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 10====> average loss: 3.0366\n",
      "\u001b[92mValidation loss decreased (3.057070 --> 3.036558).  Saving model ...\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "121ddcc4560241a2abf1b0ded4e33999",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 11 iter:300 \tLoss: 0.684115\n",
      "Train Epoch: 11 ====> Average loss: 3.0710\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9949ee20d5464ddbbb509ca1d401b263",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 11====> average loss: 3.0479\n",
      "\u001b[91mEarlyStopping counter: 1 out of 7\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6559777cd27047a7bf6fc689b5dbfe4b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 12 iter:300 \tLoss: 0.687688\n",
      "Train Epoch: 12 ====> Average loss: 3.0650\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "57a934290eef4309b34011cc63ecce65",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 12====> average loss: 3.0304\n",
      "\u001b[92mValidation loss decreased (3.036558 --> 3.030394).  Saving model ...\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6ac7c2c322f044cdb0992d1b883b82ad",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 13 iter:300 \tLoss: 0.564369\n",
      "Train Epoch: 13 ====> Average loss: 3.0583\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c268da06875b40749e28a475eb78f085",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 13====> average loss: 3.0531\n",
      "\u001b[91mEarlyStopping counter: 1 out of 7\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "59714d56c6834403b22757787868390b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 14 iter:300 \tLoss: 0.496156\n",
      "Train Epoch: 14 ====> Average loss: 3.0476\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ceb793b67fae49c399123a0dc8fcedd4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 14====> average loss: 3.0246\n",
      "\u001b[92mValidation loss decreased (3.030394 --> 3.024609).  Saving model ...\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "13365abe8f294770b2c56846fd957357",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 15 iter:300 \tLoss: 0.630353\n",
      "Train Epoch: 15 ====> Average loss: 3.0529\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7d26a4e8c8924b5e953b5c0ba3a7d770",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 15====> average loss: 3.0182\n",
      "\u001b[92mValidation loss decreased (3.024609 --> 3.018205).  Saving model ...\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f567652492b04ce6a3f484c8b3287973",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 16 iter:300 \tLoss: 0.512763\n",
      "Train Epoch: 16 ====> Average loss: 3.0451\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ae7d1b3e8a894e98b7aa246f398a739f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 16====> average loss: 3.0264\n",
      "\u001b[91mEarlyStopping counter: 1 out of 7\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "10e4813c5a644572bae6e26797fddbce",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 17 iter:300 \tLoss: 0.663007\n",
      "Train Epoch: 17 ====> Average loss: 3.0487\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "109302ead0554508911a58271d066084",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 17====> average loss: 3.0225\n",
      "\u001b[91mEarlyStopping counter: 2 out of 7\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d7bacd84acb84af1ad0f5695bd2be584",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 18 iter:300 \tLoss: 0.572800\n",
      "Train Epoch: 18 ====> Average loss: 3.0107\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ac876dc817ab4ca09990b4e4115e7382",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 18====> average loss: 2.9761\n",
      "\u001b[92mValidation loss decreased (3.018205 --> 2.976138).  Saving model ...\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b1135d5bc0d544bfbf04470e84506a32",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 19 iter:300 \tLoss: 0.647305\n",
      "Train Epoch: 19 ====> Average loss: 3.0026\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "30882dc348164d2cba4e991f5933a8c1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 19====> average loss: 2.9822\n",
      "\u001b[91mEarlyStopping counter: 1 out of 7\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f21b4489c5c143848c9eae70527e2154",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 20 iter:300 \tLoss: 0.598766\n",
      "Train Epoch: 20 ====> Average loss: 2.9974\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2b897ca6d9d647c3ae2f98128323ce6f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 20====> average loss: 3.0193\n",
      "\u001b[91mEarlyStopping counter: 2 out of 7\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e1b4628d2d5f4fa6983e0586f50b6950",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 21 iter:300 \tLoss: 0.455180\n",
      "Train Epoch: 21 ====> Average loss: 2.9943\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "217f19f9967846b7b844387cc087107c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 21====> average loss: 3.0211\n",
      "\u001b[91mEarlyStopping counter: 3 out of 7\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a485c4ffd931412a89d18552dda671ac",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 22 iter:300 \tLoss: 0.604088\n",
      "Train Epoch: 22 ====> Average loss: 3.0001\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f718e8b296f94507a322aa4bf80af1ee",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 22====> average loss: 2.9829\n",
      "\u001b[91mEarlyStopping counter: 4 out of 7\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fcbde684ed194ca88e654f605f66d418",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 23 iter:300 \tLoss: 0.544923\n",
      "Train Epoch: 23 ====> Average loss: 3.0026\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c42cf1dfdcd94b66b373d7bddb078529",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 23====> average loss: 2.9914\n",
      "\u001b[91mEarlyStopping counter: 5 out of 7\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "21dd245b8dc54a67be4635416c1f3312",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 24 iter:300 \tLoss: 0.508604\n",
      "Train Epoch: 24 ====> Average loss: 2.9938\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6cc888d54c284453b1dd483429327acf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 24====> average loss: 2.9634\n",
      "\u001b[92mValidation loss decreased (2.976138 --> 2.963357).  Saving model ...\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2509cb2a16f942cdac958234e3de3476",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 25 iter:300 \tLoss: 0.430538\n",
      "Train Epoch: 25 ====> Average loss: 2.9900\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "87fe77c06237435ba2f8bc37623b509f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 25====> average loss: 3.0432\n",
      "\u001b[91mEarlyStopping counter: 1 out of 7\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e2289495517641f789aab6f394482bf9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 26 iter:300 \tLoss: 0.698824\n",
      "Train Epoch: 26 ====> Average loss: 3.0011\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f6d1f42ade4d430fa85e259ef816c365",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 26====> average loss: 2.9745\n",
      "\u001b[91mEarlyStopping counter: 2 out of 7\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "70d3a7a7f8724105aa0a5499f3f53672",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 27 iter:300 \tLoss: 0.534325\n",
      "Train Epoch: 27 ====> Average loss: 2.9870\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a4dd41afb0b741a5ae0169f7e4622e76",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 27====> average loss: 2.9987\n",
      "\u001b[91mEarlyStopping counter: 3 out of 7\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ec2104198d7b4658b375bdd2cc0edf52",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 28 iter:300 \tLoss: 0.565815\n",
      "Train Epoch: 28 ====> Average loss: 2.9870\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e18f2f6665a8455990c30e657f4ee73e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 28====> average loss: 2.9622\n",
      "\u001b[92mValidation loss decreased (2.963357 --> 2.962219).  Saving model ...\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3a981e0c23564f129a580686d9041fb5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 29 iter:300 \tLoss: 0.693060\n",
      "Train Epoch: 29 ====> Average loss: 2.9812\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8ae7771a33804a5e9411cb177fd71718",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 29====> average loss: 2.9618\n",
      "\u001b[92mValidation loss decreased (2.962219 --> 2.961824).  Saving model ...\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "32d476e6f987412cb30ab5a7b16e6c72",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 30 iter:300 \tLoss: 0.702064\n",
      "Train Epoch: 30 ====> Average loss: 2.9924\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7ef65914c3b04542a4b2500c285cb1b4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 30====> average loss: 2.9974\n",
      "\u001b[91mEarlyStopping counter: 1 out of 7\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7a36029855b54b7bae61d1e1dea6ce7d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 31 iter:300 \tLoss: 0.667559\n",
      "Train Epoch: 31 ====> Average loss: 2.9894\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "acac713940214f27b28915e62483e6de",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 31====> average loss: 3.0005\n",
      "\u001b[91mEarlyStopping counter: 2 out of 7\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "545578bcaee84466ba72b7b2ab521c5c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 32 iter:300 \tLoss: 0.512055\n",
      "Train Epoch: 32 ====> Average loss: 2.9861\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "87f97d2936bd42f2a325b1d697d292c9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 32====> average loss: 2.9644\n",
      "\u001b[91mEarlyStopping counter: 3 out of 7\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5dc8baa971b14263874d11d78e348309",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 33 iter:300 \tLoss: 0.545099\n",
      "Train Epoch: 33 ====> Average loss: 2.9810\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2840e8c59c6d41f6bbf0b0cd897bf3f6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 33====> average loss: 2.9549\n",
      "\u001b[92mValidation loss decreased (2.961824 --> 2.954854).  Saving model ...\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e615f26fddf440da8563b1618f8944c4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 34 iter:300 \tLoss: 0.677112\n",
      "Train Epoch: 34 ====> Average loss: 2.9809\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2daa1e74447d4ed7b5bf9df9ab77d023",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 34====> average loss: 2.9530\n",
      "\u001b[92mValidation loss decreased (2.954854 --> 2.952971).  Saving model ...\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c49a69e5f61a46658648c6d6e113205a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 35 iter:300 \tLoss: 0.560534\n",
      "Train Epoch: 35 ====> Average loss: 2.9830\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "782f6aebb478474781ae27404faed16c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 35====> average loss: 2.9597\n",
      "\u001b[91mEarlyStopping counter: 1 out of 7\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fd7d6f089ff84643908af15fe909f839",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\tTrain Epoch: 36 iter:300 \tLoss: 0.629809\n",
      "Train Epoch: 36 ====> Average loss: 2.9992\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f5f106b7ee25400fb89491d004f8ded9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation/Test Epoch 36====> average loss: 2.9823\n",
      "\u001b[91mEarlyStopping counter: 2 out of 7\u001b[0m\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d3e23e0c50194631bfb666fd1e0332e9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/450 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "args.batch_size=20\n",
    "args.epochs=1000\n",
    "bucketSampler_train = BucketSampler(mozilla_dataset_train, batch_size = args.batch_size)\n",
    "bucketSampler_dev = BucketSampler(mozilla_dataset_dev, batch_size = 20)\n",
    "# bucketSampler_test = BucketSampler(mozilla_dataset_test, batch_size = 64)\n",
    "train_dataloader = DataLoader(mozilla_dataset_train, batch_sampler = bucketSampler_train, collate_fn=collate_fn)#, num_workers=4)\n",
    "dev_dataloader = DataLoader(mozilla_dataset_dev, batch_sampler = bucketSampler_dev, collate_fn=collate_fn)#, num_workers=4)\n",
    "# test_dataloader = DataLoader(mozilla_dataset_test, batch_sampler = bucketSampler_test, collate_fn=collate_fn)#, num_workers=4)\n",
    "\n",
    "model_filename=f\"allosrs/model/lstm_nl{model.encoder.bi_LSTM.num_layers}_hd{model.encoder.bi_LSTM.hidden_size}_lr{args.lr}_bs{args.batch_size}.pt\"\n",
    "# model_filename=f\"allosrs/model/dummy.pt\"\n",
    "print(\"MODEL SAVE NAME---\",model_filename)\n",
    "\n",
    "log_filename=f\"allosrs/log/lstm_nl{model.encoder.bi_LSTM.num_layers}_hd{model.encoder.bi_LSTM.hidden_size}_lr{args.lr}_bs{args.batch_size}.log\"\n",
    "with open(log_filename,mode=\"w\") as wfile:\n",
    "    print(\"\",end=\"\",file=wfile)\n",
    "\n",
    "\n",
    "es=EarlyStopping(log_file=log_filename,verbose=True,path=model_filename)\n",
    "for epoch in range(1, args.epochs + 1):\n",
    "# for epoch in range(1, 5+ 1):\n",
    "    dev_loss=evaluate(epoch-1,dev_dataloader,log_filename)\n",
    "    # dev_loss=(epoch-5)**2\n",
    "    es(dev_loss,model)\n",
    "    if es.early_stop==True:\n",
    "        break\n",
    "    train_loss=train(epoch,train_dataloader,log_filename)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a=torch.nn.Parameter(torch.zeros(3,3))\n",
    "torch.linalg.matrix_norm(a-a+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set(mozilla_dataset_train.lang)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "signature_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7.0"
  },
  "vscode": {
   "interpreter": {
    "hash": "5e6617ec83aa391bede3392ef2e6ceb28ce1b7bd91d34d663fce0e811022229b"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}