03b_transformers_words_spacy_wikitext103_audio.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.7/site-packages/sklearn/utils/deprecation.py:143: FutureWarning: The sklearn.datasets.base module is  deprecated in version 0.22 and will be removed in version 0.24. The corresponding classes / functions should instead be imported from sklearn.datasets. Anything that cannot be imported from sklearn.datasets is now part of the private API.\n",
      "  warnings.warn(message, FutureWarning)\n"
     ]
    }
   ],
   "source": [
    "from fastai.vision.all import *\n",
    "from fastai.learner import *\n",
    "from fastai.data.all import *\n",
    "from fastai.callback.tracker import SaveModelCallback\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from pathlib2 import Path\n",
    "import numpy as np\n",
    "import random\n",
    "from torch.nn import MSELoss\n",
    "import logging\n",
    "from six import iteritems\n",
    "from web.datasets.similarity import fetch_MEN, fetch_WS353, fetch_SimLex999\n",
    "from web.embeddings import fetch_GloVe\n",
    "from web.evaluate import evaluate_similarity\n",
    "from web.embedding import Embedding, Vocabulary\n",
    "from gensim.models import Word2Vec\n",
    "from gensim.models import KeyedVectors\n",
    "from python_speech_features.base import mfcc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://s3.amazonaws.com/fast-ai-nlp/wikitext-103.tgz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train = pd.read_csv('data/wikitext-103/train.csv', header=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(29539, 1)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "texts = df_train[0].tolist()\n",
    "texts = texts[:2000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from spacy.tokenizer import Tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import spacy\n",
    "from spacy.attrs import ORTH, LEMMA\n",
    "\n",
    "nlp = spacy.load(\"en_core_web_sm\")\n",
    "\n",
    "nlp.tokenizer.add_special_case('<unk>', [{ORTH: \"<unk>\"}])\n",
    "nlp.tokenizer.add_special_case('<eos>', [{ORTH: \"<eos>\"}])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from joblib import Parallel, delayed\n",
    "from functools import partial\n",
    "from spacy.util import minibatch\n",
    "import spacy\n",
    "\n",
    "nlp = spacy.load(\"en_core_web_sm\")\n",
    "\n",
    "batch_size= 200\n",
    "partitions = minibatch(texts, size=batch_size)\n",
    "\n",
    "executor = Parallel(n_jobs=10, backend=\"multiprocessing\", prefer=\"threads\")\n",
    "\n",
    "def transform_texts(nlp, texts):\n",
    "    for i in range(len(texts)):\n",
    "        texts[i] = [line.strip() for line in texts[i].split('\\n') if len(line) > 100]\n",
    "        texts[i] = ('\\n').join(texts[i])\n",
    "    docs = nlp.tokenizer.pipe(texts)\n",
    "    return list(docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 4s, sys: 3.7 s, total: 1min 8s\n",
      "Wall time: 1min 11s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "do = delayed(partial(transform_texts, nlp))\n",
    "tasks = (do(batch) for batch in partitions)\n",
    "docs_hierarchical = executor(tasks)\n",
    "\n",
    "docs = []\n",
    "for ds in docs_hierarchical: docs += ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 17.1 s, sys: 268 ms, total: 17.4 s\n",
      "Wall time: 17.4 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "words = []\n",
    "\n",
    "for doc in docs:\n",
    "    for token in doc:\n",
    "        if (token.is_alpha or token.is_punct) and '@' not in str(token):\n",
    "            words.append(str(token).lower())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "c = Counter(words)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "tasks = {\n",
    "    \"MEN\": fetch_MEN(),\n",
    "    \"WS353\": fetch_WS353(),\n",
    "    \"SIMLEX999\": fetch_SimLex999()\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "words = []\n",
    "\n",
    "for k, v in tasks.items():\n",
    "    for row in v.X:\n",
    "        for w in row:\n",
    "            words.append(w)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1819"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(set(words))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1584"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(set([w for w, _ in c.most_common(20_000)]).intersection(words))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab = set([w for w, _ in c.most_common(20_000)] + ['<eos>', '<unk>'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# mkdir -p data/audio_wikitext103_20k"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from time import sleep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%time\n",
    "\n",
    "# from google.cloud import texttospeech\n",
    "\n",
    "# client = texttospeech.TextToSpeechClient()\n",
    "\n",
    "# for i, word in enumerate(vocab):\n",
    "#     if i % 300 == 0: sleep(60)\n",
    "\n",
    "#     synthesis_input = texttospeech.SynthesisInput(text=word)\n",
    "\n",
    "#     voice = texttospeech.VoiceSelectionParams(\n",
    "#         language_code=\"en-US\", ssml_gender=texttospeech.SsmlVoiceGender.NEUTRAL\n",
    "#     )\n",
    "\n",
    "#     audio_config = texttospeech.AudioConfig(\n",
    "#         audio_encoding=texttospeech.AudioEncoding.MP3\n",
    "#     )\n",
    "\n",
    "#     response = client.synthesize_speech(\n",
    "#         input=synthesis_input, voice=voice, audio_config=audio_config\n",
    "#     )\n",
    "\n",
    "#     # The response's audio_content is binary.\n",
    "#     with open(f\"data/audio_wikitext103_20k/{word}.mp3\", \"wb\") as out:\n",
    "#         out.write(response.audio_content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import librosa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%time\n",
    "\n",
    "# import warnings\n",
    "# warnings.filterwarnings('ignore')\n",
    "\n",
    "# word2ary = {}\n",
    "\n",
    "# for word in vocab:\n",
    "#     ary, _ = librosa.load(f'data/audio_wikitext103_20k/{word}.mp3', sr=16000)\n",
    "#     word2ary[word] = ary\n",
    "    \n",
    "# pd.to_pickle(word2ary, 'word2ary.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "word2ary = pd.read_pickle('word2ary.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "word2features = {}\n",
    "for k, v in word2ary.items():\n",
    "    word2features[k] = mfcc(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "20002"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(word2features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "lengths = []\n",
    "for k, v in word2features.items():\n",
    "    lengths.append(v.shape[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "193"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "max(lengths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_features(word, pad_to=max(lengths)):\n",
    "    ary = word2features[word][:pad_to]\n",
    "    example = np.zeros((pad_to, 13))\n",
    "    example[:ary.shape[0], :] = ary\n",
    "    return example.astype(np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_mean = -2\n",
    "dataset_std = 10\n",
    "\n",
    "def normalize_data(ary):\n",
    "    return (ary - dataset_mean) / dataset_std"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "word2features_norm = {}\n",
    "for word in vocab:\n",
    "    word2features_norm[word] = normalize_data(prepare_features(word))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = ''\n",
    "\n",
    "for doc in docs:\n",
    "    toks = []\n",
    "    for token in doc:\n",
    "        if (token.is_alpha or token.is_punct) and '@' not in str(token):\n",
    "            if str(token) in vocab:\n",
    "                toks.append(str(token).lower())\n",
    "            else:\n",
    "                toks.append('<unk>')\n",
    "    toks += ['<eos>']\n",
    "    text += ' '.join(toks)\n",
    "    text += ' '"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "32498344"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6224583"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "toks = text.split()\n",
    "len(toks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['<unk>', 'is', 'a', 'fungal', 'genus', 'in', 'the', 'family', '<unk>', ',']"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "toks[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab = list(vocab)\n",
    "# stoi = {s: i for i, s in enumerate(vocab)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "# toks_numeric = [stoi[tok] for tok in toks]\n",
    "# toks_numeric[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "class WordDataset(Dataset):\n",
    "\n",
    "    def __init__(self, data, block_size):\n",
    "        data_size, vocab_size = len(data), len(vocab)\n",
    "        print('data has %d words, %d unique.' % (data_size, vocab_size))\n",
    "        \n",
    "        self.stoi = { word:i for i,word in enumerate(vocab) }\n",
    "        self.itos = { i:word for i,word in enumerate(vocab) }\n",
    "        self.block_size = block_size\n",
    "        self.vocab_size = vocab_size\n",
    "        self.data = data\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.data) - self.block_size\n",
    " \n",
    "    def __getitem__(self, idx):\n",
    "        chunk = self.data[idx:idx + self.block_size + 1]\n",
    "        features = [word2features_norm[s] for s in chunk[:-1]]\n",
    "        dix = [self.stoi[s] for s in chunk]\n",
    "        x = torch.tensor(dix[:-1], dtype=torch.long)\n",
    "        y = torch.tensor(dix[1:], dtype=torch.long)\n",
    "        return np.stack(features), y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LSTM(13, 256, batch_first=True, bidirectional=True)"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "data has 6000000 words, 20002 unique.\n",
      "data has 224583 words, 20002 unique.\n"
     ]
    }
   ],
   "source": [
    "block_size = 128\n",
    "\n",
    "train_dataset = WordDataset(toks[:6_000_000], block_size)\n",
    "valid_dataset = WordDataset(toks[6_000_000:], block_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import logging\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "class GPTConfig:\n",
    "    \"\"\" base GPT config, params common to all GPT versions \"\"\"\n",
    "    embd_pdrop = 0.1\n",
    "    resid_pdrop = 0.1\n",
    "    attn_pdrop = 0.1\n",
    "\n",
    "    def __init__(self, vocab_size, block_size, **kwargs):\n",
    "        self.vocab_size = vocab_size\n",
    "        self.block_size = block_size\n",
    "        for k,v in kwargs.items():\n",
    "            setattr(self, k, v)\n",
    "\n",
    "class GPT1Config(GPTConfig):\n",
    "    \"\"\" GPT-1 like network roughly 125M params \"\"\"\n",
    "    n_layer = 12\n",
    "    n_head = 12\n",
    "    n_embd = 768\n",
    "\n",
    "class CausalSelfAttention(nn.Module):\n",
    "    \"\"\"\n",
    "    A vanilla multi-head masked self-attention layer with a projection at the end.\n",
    "    It is possible to use torch.nn.MultiheadAttention here but I am including an\n",
    "    explicit implementation here to show that there is nothing too scary here.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        assert config.n_embd % config.n_head == 0\n",
    "        # key, query, value projections for all heads\n",
    "        self.key = nn.Linear(config.n_embd, config.n_embd)\n",
    "        self.query = nn.Linear(config.n_embd, config.n_embd)\n",
    "        self.value = nn.Linear(config.n_embd, config.n_embd)\n",
    "        # regularization\n",
    "        self.attn_drop = nn.Dropout(config.attn_pdrop)\n",
    "        self.resid_drop = nn.Dropout(config.resid_pdrop)\n",
    "        # output projection\n",
    "        self.proj = nn.Linear(config.n_embd, config.n_embd)\n",
    "        # causal mask to ensure that attention is only applied to the left in the input sequence\n",
    "        self.register_buffer(\"mask\", torch.tril(torch.ones(config.block_size, config.block_size))\n",
    "                                     .view(1, 1, config.block_size, config.block_size))\n",
    "        self.n_head = config.n_head\n",
    "\n",
    "    def forward(self, x, layer_past=None):\n",
    "        B, T, C = x.size()\n",
    "\n",
    "        # calculate query, key, values for all heads in batch and move head forward to be the batch dim\n",
    "        k = self.key(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n",
    "        q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n",
    "        v = self.value(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n",
    "\n",
    "        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)\n",
    "        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))\n",
    "        att = att.masked_fill(self.mask[:,:,:T,:T] == 0, float('-inf'))\n",
    "        att = F.softmax(att, dim=-1)\n",
    "        att = self.attn_drop(att)\n",
    "        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)\n",
    "        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side\n",
    "\n",
    "        # output projection\n",
    "        y = self.resid_drop(self.proj(y))\n",
    "        return y\n",
    "\n",
    "class Block(nn.Module):\n",
    "    \"\"\" an unassuming Transformer block \"\"\"\n",
    "\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.ln1 = nn.LayerNorm(config.n_embd)\n",
    "        self.ln2 = nn.LayerNorm(config.n_embd)\n",
    "        self.attn = CausalSelfAttention(config)\n",
    "        self.mlp = nn.Sequential(\n",
    "            nn.Linear(config.n_embd, 4 * config.n_embd),\n",
    "            nn.GELU(),\n",
    "            nn.Linear(4 * config.n_embd, config.n_embd),\n",
    "            nn.Dropout(config.resid_pdrop),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.attn(self.ln1(x))\n",
    "        x = x + self.mlp(self.ln2(x))\n",
    "        return x\n",
    "\n",
    "class GPT(nn.Module):\n",
    "    \"\"\"  the full GPT language model, with a context size of block_size \"\"\"\n",
    "\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.encoder= nn.LSTM(\n",
    "            input_size=13,\n",
    "            hidden_size=config.n_embd // 2,\n",
    "            num_layers=1,\n",
    "            batch_first=True,\n",
    "            dropout=0,\n",
    "            bidirectional=True\n",
    "        )\n",
    "        \n",
    "        # input embedding stem\n",
    "#         self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd)\n",
    "        self.pos_emb = nn.Parameter(torch.zeros(1, config.block_size, config.n_embd))\n",
    "        self.drop = nn.Dropout(config.embd_pdrop)\n",
    "        # transformer\n",
    "        self.blocks = nn.Sequential(*[Block(config) for _ in range(config.n_layer)])\n",
    "        # decoder head\n",
    "        self.ln_f = nn.LayerNorm(config.n_embd)\n",
    "        self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False)\n",
    "\n",
    "        self.block_size = config.block_size\n",
    "        self.apply(self._init_weights)\n",
    "\n",
    "        logger.info(\"number of parameters: %e\", sum(p.numel() for p in self.parameters()))\n",
    "\n",
    "    def get_block_size(self):\n",
    "        return self.block_size\n",
    "\n",
    "    def _init_weights(self, module):\n",
    "        if isinstance(module, (nn.Linear, nn.Embedding)):\n",
    "            module.weight.data.normal_(mean=0.0, std=0.02)\n",
    "            if isinstance(module, nn.Linear) and module.bias is not None:\n",
    "                module.bias.data.zero_()\n",
    "        elif isinstance(module, nn.LayerNorm):\n",
    "            module.bias.data.zero_()\n",
    "            module.weight.data.fill_(1.0)\n",
    "\n",
    "    def configure_optimizers(self, train_config):\n",
    "        \"\"\"\n",
    "        This long function is unfortunately doing something very simple and is being very defensive:\n",
    "        We are separating out all parameters of the model into two buckets: those that will experience\n",
    "        weight decay for regularization and those that won't (biases, and layernorm/embedding weights).\n",
    "        We are then returning the PyTorch optimizer object.\n",
    "        \"\"\"\n",
    "\n",
    "        # separate out all parameters to those that will and won't experience regularizing weight decay\n",
    "        decay = set()\n",
    "        no_decay = set()\n",
    "        whitelist_weight_modules = (torch.nn.Linear, )\n",
    "        blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding)\n",
    "        for mn, m in self.named_modules():\n",
    "            for pn, p in m.named_parameters():\n",
    "                fpn = '%s.%s' % (mn, pn) if mn else pn # full param name\n",
    "\n",
    "                if pn.endswith('bias'):\n",
    "                    # all biases will not be decayed\n",
    "                    no_decay.add(fpn)\n",
    "                elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules):\n",
    "                    # weights of whitelist modules will be weight decayed\n",
    "                    decay.add(fpn)\n",
    "                elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules):\n",
    "                    # weights of blacklist modules will NOT be weight decayed\n",
    "                    no_decay.add(fpn)\n",
    "                elif pn.startswith('encoder'):\n",
    "                    no_decay.add(fpn)\n",
    "\n",
    "        # special case the position embedding parameter in the root GPT module as not decayed\n",
    "        no_decay.add('pos_emb')\n",
    "\n",
    "        # validate that we considered every parameter\n",
    "        param_dict = {pn: p for pn, p in self.named_parameters()}\n",
    "        inter_params = decay & no_decay\n",
    "        union_params = decay | no_decay\n",
    "        assert len(inter_params) == 0, \"parameters %s made it into both decay/no_decay sets!\" % (str(inter_params), )\n",
    "        assert len(param_dict.keys() - union_params) == 0, \"parameters %s were not separated into either decay/no_decay set!\" \\\n",
    "                                                    % (str(param_dict.keys() - union_params), )\n",
    "\n",
    "        # create the pytorch optimizer object\n",
    "        optim_groups = [\n",
    "            {\"params\": [param_dict[pn] for pn in sorted(list(decay))], \"weight_decay\": train_config.weight_decay},\n",
    "            {\"params\": [param_dict[pn] for pn in sorted(list(no_decay))], \"weight_decay\": 0.0},\n",
    "        ]\n",
    "        optimizer = torch.optim.AdamW(optim_groups, lr=train_config.learning_rate, betas=train_config.betas)\n",
    "        return optimizer\n",
    "\n",
    "    def forward(self, features, targets=None):        \n",
    "        b, t, _, _ = features.shape\n",
    "        assert t <= self.block_size, \"Cannot forward, model block size is exhausted.\"\n",
    "        \n",
    "        t_embs = []\n",
    "        for i in range(b):\n",
    "            _, (embeddings, _) = self.encoder(features[i])        \n",
    "            t_embs.append(torch.cat((embeddings[-1], embeddings[-2]), 1))\n",
    "\n",
    "        token_embeddings = torch.stack(t_embs)\n",
    "        \n",
    "        # forward the GPT model\n",
    "#         token_embeddings = self.tok_emb(idx) # each index maps to a (learnable) vector\n",
    "        position_embeddings = self.pos_emb[:, :t, :] # each position maps to a (learnable) vector\n",
    "        x = self.drop(token_embeddings + position_embeddings)\n",
    "        x = self.blocks(x)\n",
    "        x = self.ln_f(x)\n",
    "        logits = self.head(x)\n",
    "\n",
    "        # if we are given some desired targets also calculate the loss\n",
    "        loss = None\n",
    "        if targets is not None:\n",
    "            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))\n",
    "\n",
    "        return logits, loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "  \n",
    "\"\"\"\n",
    "Simple training loop; Boilerplate that could apply to any arbitrary neural network,\n",
    "so nothing in this file really has anything to do with GPT specifically.\n",
    "\"\"\"\n",
    "\n",
    "import math\n",
    "import logging\n",
    "\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.optim as optim\n",
    "from torch.optim.lr_scheduler import LambdaLR\n",
    "from torch.utils.data.dataloader import DataLoader\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "class TrainerConfig:\n",
    "    # optimization parameters\n",
    "    max_epochs = 10\n",
    "    batch_size = 64\n",
    "    learning_rate = 3e-4\n",
    "    betas = (0.9, 0.95)\n",
    "    grad_norm_clip = 1.0\n",
    "    weight_decay = 0.1 # only applied on matmul weights\n",
    "    # learning rate decay params: linear warmup followed by cosine decay to 10% of original\n",
    "    lr_decay = False\n",
    "    warmup_tokens = 375e6 # these two numbers come from the GPT-3 paper, but may not be good defaults elsewhere\n",
    "    final_tokens = 260e9 # (at what point we reach 10% of original LR)\n",
    "    # checkpoint settings\n",
    "    ckpt_path = None\n",
    "    num_workers = 0 # for DataLoader\n",
    "\n",
    "    def __init__(self, **kwargs):\n",
    "        for k,v in kwargs.items():\n",
    "            setattr(self, k, v)\n",
    "\n",
    "class Trainer:\n",
    "\n",
    "    def __init__(self, model, train_dataset, test_dataset, config):\n",
    "        self.model = model\n",
    "        self.train_dataset = train_dataset\n",
    "        self.test_dataset = test_dataset\n",
    "        self.config = config\n",
    "\n",
    "        # take over whatever gpus are on the system\n",
    "        self.device = 'cpu'\n",
    "        if torch.cuda.is_available():\n",
    "            self.device = torch.cuda.current_device()\n",
    "            self.model = torch.nn.DataParallel(self.model).to(self.device)\n",
    "\n",
    "    def save_checkpoint(self):\n",
    "        # DataParallel wrappers keep raw model object in .module attribute\n",
    "        raw_model = self.model.module if hasattr(self.model, \"module\") else self.model\n",
    "        logger.info(\"saving %s\", self.config.ckpt_path)\n",
    "        torch.save(raw_model.state_dict(), self.config.ckpt_path)\n",
    "\n",
    "    def train(self):\n",
    "        model, config = self.model, self.config\n",
    "        raw_model = model.module if hasattr(self.model, \"module\") else model\n",
    "        optimizer = raw_model.configure_optimizers(config)\n",
    "\n",
    "        def run_epoch(split):\n",
    "            is_train = split == 'train'\n",
    "            model.train(is_train)\n",
    "            data = self.train_dataset if is_train else self.test_dataset\n",
    "            loader = DataLoader(data, shuffle=True, pin_memory=True,\n",
    "                                batch_size=config.batch_size,\n",
    "                                num_workers=config.num_workers)\n",
    "\n",
    "            losses = []\n",
    "            pbar = tqdm(enumerate(loader), total=len(loader)) if is_train else enumerate(loader)\n",
    "            for it, (x, y) in pbar:\n",
    "                \n",
    "                # place data on the correct device\n",
    "                x = x.to(self.device)\n",
    "                y = y.to(self.device)\n",
    "\n",
    "                # forward the model\n",
    "                with torch.set_grad_enabled(is_train):\n",
    "                    logits, loss = model(x, y)\n",
    "                    loss = loss.mean() # collapse all losses if they are scattered on multiple gpus\n",
    "                    losses.append(loss.item())\n",
    "\n",
    "                if is_train:\n",
    "\n",
    "                    # backprop and update the parameters\n",
    "                    model.zero_grad()\n",
    "                    loss.backward()\n",
    "                    torch.nn.utils.clip_grad_norm_(model.parameters(), config.grad_norm_clip)\n",
    "                    optimizer.step()\n",
    "\n",
    "                    # decay the learning rate based on our progress\n",
    "                    if config.lr_decay:\n",
    "                        self.tokens += (y >= 0).sum() # number of tokens processed this step (i.e. label is not -100)\n",
    "                        if self.tokens < config.warmup_tokens:\n",
    "                            # linear warmup\n",
    "                            lr_mult = float(self.tokens) / float(max(1, config.warmup_tokens))\n",
    "                        else:\n",
    "                            # cosine learning rate decay\n",
    "                            progress = float(self.tokens - config.warmup_tokens) / float(max(1, config.final_tokens - config.warmup_tokens))\n",
    "                            lr_mult = max(0.1, 0.5 * (1.0 + math.cos(math.pi * progress)))\n",
    "                        lr = config.learning_rate * lr_mult\n",
    "                        for param_group in optimizer.param_groups:\n",
    "                            param_group['lr'] = lr\n",
    "                    else:\n",
    "                        lr = config.learning_rate\n",
    "\n",
    "                    # report progress\n",
    "                    pbar.set_description(f\"epoch {epoch+1} iter {it}: train loss {loss.item():.5f}. lr {lr:e}\")\n",
    "\n",
    "            if not is_train:\n",
    "                test_loss = float(np.mean(losses))\n",
    "                logger.info(\"test loss: %f\", test_loss)\n",
    "                return test_loss\n",
    "\n",
    "        best_loss = float('inf')\n",
    "        self.tokens = 0 # counter used for learning rate decay\n",
    "        for epoch in range(config.max_epochs):\n",
    "\n",
    "            run_epoch('train')\n",
    "            if self.test_dataset is not None:\n",
    "                test_loss = run_epoch('test')\n",
    "\n",
    "            # supports early stopping based on the test loss, or just save always if no test set is provided\n",
    "            good_model = self.test_dataset is None or test_loss < best_loss\n",
    "            if self.config.ckpt_path is not None and good_model:\n",
    "                best_loss = test_loss\n",
    "                self.save_checkpoint()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "mconf = GPTConfig(train_dataset.vocab_size, train_dataset.block_size, n_layer=8, n_head=8, n_embd=512)\n",
    "model = GPT(mconf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 1 iter 187495: train loss 2.30185. lr 1.500016e-04: 100%|██████████| 187496/187496 [67:17:22<00:00,  1.29s/it]   \n"
     ]
    }
   ],
   "source": [
    "tconf = TrainerConfig(max_epochs=1, batch_size=32, learning_rate=3e-4,\n",
    "                      lr_decay=True, warmup_tokens=512*20, final_tokens=2*len(train_dataset)*block_size,\n",
    "                      num_workers=8)\n",
    "trainer = Trainer(model, train_dataset, valid_dataset, tconf)\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), 'models/gpt_audio.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "# def top_k_logits(logits, k):\n",
    "#     v, ix = torch.topk(logits, k)\n",
    "#     out = logits.clone()\n",
    "#     out[out < v[:, [-1]]] = -float('Inf')\n",
    "#     return out\n",
    "\n",
    "# @torch.no_grad()\n",
    "# def sample(model, x, steps, temperature=1.0, sample=False, top_k=None):\n",
    "#     \"\"\"\n",
    "#     take a conditioning sequence of indices in x (of shape (b,t)) and predict the next token in\n",
    "#     the sequence, feeding the predictions back into the model each time. Clearly the sampling\n",
    "#     has quadratic complexity unlike an RNN that is only linear, and has a finite context window\n",
    "#     of block_size, unlike an RNN that has an infinite context window.\n",
    "#     \"\"\"\n",
    "#     block_size = model.get_block_size()\n",
    "#     model.eval()\n",
    "#     for k in range(steps):\n",
    "#         x_cond = x if x.size(1) <= block_size else x[:, -block_size:] # crop context if needed\n",
    "#         logits, _ = model(x_cond)\n",
    "#         # pluck the logits at the final step and scale by temperature\n",
    "#         logits = logits[:, -1, :] / temperature\n",
    "#         # optionally crop probabilities to only the top k options\n",
    "#         if top_k is not None:\n",
    "#             logits = top_k_logits(logits, top_k)\n",
    "#         # apply softmax to convert to probabilities\n",
    "#         probs = F.softmax(logits, dim=-1)\n",
    "#         # sample from the distribution or take the most likely\n",
    "#         if sample:\n",
    "#             ix = torch.multinomial(probs, num_samples=1)\n",
    "#         else:\n",
    "#             _, ix = torch.topk(probs, k=1, dim=-1)\n",
    "#         # append to the sequence and continue\n",
    "#         x = torch.cat((x, ix), dim=1)\n",
    "\n",
    "#     return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "# context = [\"the\", \"ship\"]\n",
    "# x = torch.tensor([train_dataset.stoi[s] for s in context], dtype=torch.long)[None,...].to(trainer.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "# y = sample(model, x, 400, temperature=1.0, sample=True, top_k=10)[0]\n",
    "# completion = ' '.join([train_dataset.itos[int(i)] for i in y])\n",
    "# print(completion)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "from six import iteritems\n",
    "from web.datasets.similarity import fetch_MEN, fetch_WS353, fetch_SimLex999\n",
    "from web.embeddings import fetch_GloVe\n",
    "from web.evaluate import evaluate_similarity\n",
    "from web.embedding import Embedding, Vocabulary\n",
    "from gensim.models import Word2Vec\n",
    "from gensim.models import KeyedVectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [],
   "source": [
    "tasks = {\n",
    "    \"MEN\": fetch_MEN(),\n",
    "    \"WS353\": fetch_WS353(),\n",
    "    \"SIMLEX999\": fetch_SimLex999()\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [],
   "source": [
    "speech2vec = KeyedVectors.load_word2vec_format('../speech2vec-pretrained-vectors/speech2vec/50.vec', binary=False) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EmbeddingsDataset(Dataset):\n",
    "    def __len__(self):\n",
    "        return len(vocab)\n",
    "    def __getitem__(self, idx):\n",
    "        return word2features_norm[vocab[idx]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_dl = DataLoader(EmbeddingsDataset(), 128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.52 s, sys: 128 ms, total: 1.64 s\n",
      "Wall time: 1.57 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "vocab_embeddings = []\n",
    "\n",
    "model.train(False);\n",
    "with torch.no_grad():\n",
    "    for batch in pred_dl:\n",
    "        _, (embeddings, _) = model.encoder(batch.cuda())\n",
    "        embeddings = torch.cat((embeddings[-1], embeddings[-2]), 1).cpu().detach()\n",
    "        vocab_embeddings.append(embeddings)\n",
    "        \n",
    "vocab_embeddings = torch.cat(vocab_embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [],
   "source": [
    "our_embeddings = Embedding(\n",
    "    Vocabulary(vocab),\n",
    "    vocab_embeddings\n",
    ")\n",
    "\n",
    "speech2vec = KeyedVectors.load_word2vec_format('../speech2vec-pretrained-vectors/word2vec/50.vec', binary=False) \n",
    "speech2vec_embeddings = Embedding(Vocabulary(list(speech2vec.vocab.keys())), speech2vec.vectors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Missing 883 words. Will replace them with mean vector\n",
      "/opt/conda/lib/python3.7/site-packages/web-0.0.1-py3.7.egg/web/evaluate.py:336: FutureWarning: arrays to stack must be passed as a \"sequence\" type such as list or tuple. Support for non-sequence iterables such as generators is deprecated as of NumPy 1.16 and will raise an error in the future.\n",
      "  A = np.vstack(w.get(word, mean_vector) for word in X[:, 0])\n",
      "/opt/conda/lib/python3.7/site-packages/web-0.0.1-py3.7.egg/web/evaluate.py:337: FutureWarning: arrays to stack must be passed as a \"sequence\" type such as list or tuple. Support for non-sequence iterables such as generators is deprecated as of NumPy 1.16 and will raise an error in the future.\n",
      "  B = np.vstack(w.get(word, mean_vector) for word in X[:, 1])\n",
      "Missing 74 words. Will replace them with mean vector\n",
      "Missing 122 words. Will replace them with mean vector\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Spearman correlation of scores on MEN 0.0226286156446734\n",
      "Spearman correlation of scores on WS353 0.11148925385144111\n",
      "Spearman correlation of scores on SIMLEX999 0.011862104575537205\n"
     ]
    }
   ],
   "source": [
    "for name, data in iteritems(tasks):\n",
    "    print(\"Spearman correlation of scores on {} {}\".format(name, evaluate_similarity(our_embeddings, data.X, data.y)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Missing 392 words. Will replace them with mean vector\n",
      "Missing 61 words. Will replace them with mean vector\n",
      "Missing 24 words. Will replace them with mean vector\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Spearman correlation of scores on MEN 0.6056592803599269\n",
      "Spearman correlation of scores on WS353 0.43349390636024643\n",
      "Spearman correlation of scores on SIMLEX999 0.25938770901422736\n"
     ]
    }
   ],
   "source": [
    "for name, data in iteritems(tasks):\n",
    "    print(\"Spearman correlation of scores on {} {}\".format(name, evaluate_similarity(speech2vec_embeddings, data.X, data.y)))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}