03a_transformers_words_spacy_wikitext103.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"
   ]
  },
  {
   "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": 54,
   "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": 55,
   "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": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 5s, sys: 4.02 s, total: 1min 9s\n",
      "Wall time: 1min 12s\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": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 16.4 s, sys: 208 ms, total: 16.6 s\n",
      "Wall time: 16.6 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": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "c = Counter(words)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "tasks = {\n",
    "    \"MEN\": fetch_MEN(),\n",
    "    \"WS353\": fetch_WS353(),\n",
    "    \"SIMLEX999\": fetch_SimLex999()\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "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": 62,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1819"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(set(words))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1584"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(set([w for w, _ in c.most_common(20_000)]).intersection(words))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab = set([w for w, _ in c.most_common(20_000)] + ['<eos>', '<unk>'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "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": 66,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "32498344"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6224583"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "toks = text.split()\n",
    "len(toks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['<unk>', 'is', 'a', 'fungal', 'genus', 'in', 'the', 'family', '<unk>', ',']"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "toks[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab = list(vocab)\n",
    "# stoi = {s: i for i, s in enumerate(vocab)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "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": 24,
   "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",
    "        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 x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "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": 26,
   "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",
    "        # 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",
    "\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, idx, targets=None):\n",
    "        b, t = idx.size()\n",
    "        assert t <= self.block_size, \"Cannot forward, model block size is exhausted.\"\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": 27,
   "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": 31,
   "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": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 1 iter 44423: train loss 1.41769. lr 1.623031e-04:  95%|█████████▍| 44424/46874 [4:31:17<15:00,  2.72it/s]  IOPub message rate exceeded.\n",
      "The notebook 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",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n",
      "epoch 2 iter 2479: train loss 1.39965. lr 1.375497e-04:   5%|▌         | 2479/46874 [15:11<4:32:57,  2.71it/s]IOPub message rate exceeded.\n",
      "The notebook 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",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n",
      "epoch 2 iter 7380: train loss 1.34147. lr 1.132767e-04:  16%|█▌        | 7380/46874 [45:12<4:02:10,  2.72it/s]IOPub message rate exceeded.\n",
      "The notebook 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",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n",
      "epoch 2 iter 12300: train loss 1.27336. lr 8.990451e-05:  26%|██▌       | 12300/46874 [1:15:20<3:31:47,  2.72it/s]IOPub message rate exceeded.\n",
      "The notebook 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",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n",
      "epoch 2 iter 17227: train loss 1.20002. lr 6.813277e-05:  37%|███▋      | 17228/46874 [1:45:30<3:00:59,  2.73it/s]IOPub message rate exceeded.\n",
      "The notebook 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",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n",
      "epoch 2 iter 21905: train loss 1.14360. lr 4.951354e-05:  47%|████▋     | 21905/46874 [2:14:10<2:32:59,  2.72it/s]IOPub message rate exceeded.\n",
      "The notebook 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",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n",
      "epoch 2 iter 26837: train loss 1.13330. lr 3.256022e-05:  57%|█████▋    | 26838/46874 [2:44:23<2:02:25,  2.73it/s]IOPub message rate exceeded.\n",
      "The notebook 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",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n",
      "epoch 2 iter 31709: train loss 1.10016. lr 3.000000e-05:  68%|██████▊   | 31709/46874 [3:14:14<1:32:53,  2.72it/s]IOPub message rate exceeded.\n",
      "The notebook 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",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n",
      "epoch 2 iter 36305: train loss 1.09231. lr 3.000000e-05:  77%|███████▋  | 36306/46874 [3:42:24<1:04:44,  2.72it/s]IOPub message rate exceeded.\n",
      "The notebook 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",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n",
      "epoch 2 iter 41172: train loss 1.08893. lr 3.000000e-05:  88%|████████▊ | 41172/46874 [4:12:13<34:53,  2.72it/s]IOPub message rate exceeded.\n",
      "The notebook 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",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n",
      "epoch 2 iter 46054: train loss 1.08740. lr 3.000000e-05:  98%|█████████▊| 46054/46874 [4:42:07<05:00,  2.73it/s]IOPub message rate exceeded.\n",
      "The notebook 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",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "tconf = TrainerConfig(max_epochs=2, batch_size=128, 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": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), 'models/gpt_words.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "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": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "the ship to be named <unk> <unk> after the ship was laid down . <unk> she was the first <unk> ship to be named for <unk> <unk> , the <unk> <unk> . <unk> was launched on <unk> and completed on <unk> . <unk> served with the fleet until she was transferred to <unk> <unk> <unk> in response to support the convoys carried there . <unk> became increasingly involved in the <unk> of <unk> , while the rest of the <unk> <unk> <unk> <unk> was under consideration to succeed the <unk> in training . <unk> cruiser operated the <unk> <unk> <unk> , <unk> , for operation into the waters of the <unk> <unk> of the <unk> <unk> . <unk> fleet returned to <unk> and was employed as a target ship in the <unk> <unk> <unk> ( the <unk> <unk> and <unk> <unk> ) with this mission . <unk> <unk> entered <unk> <unk> <unk> <unk> in <unk> , and was decommissioned on <unk> , entering the <unk> <unk> <unk> <unk> on <unk> . <unk> arrived at <unk> on <unk> for a fourth time , which charted in the top in both the <unk> and <unk> <unk> <unk> , with the assistance of the <unk> <unk> . <unk> <unk> , the <unk> embarked for the <unk> <unk> <unk> in the <unk> of the <unk> . <unk> put to action on <unk> and arrived at <unk> , <unk> on <unk> and <unk> , where she joined the <unk> . <unk> the three week trial , <unk> was chosen to direct the birth of <unk> <unk> . <unk> was reinstated in <unk> and the charges were dropped . <unk> promised that , if a second trial would be held as a <unk> if all charges were dropped , those who had burned him to death were dropped . <unk> lived <unk> at <unk> . <unk> <unk> ( – ) was a <unk> lawyer from the state of <unk> , and he served most of the <unk> of any crew that worked with him . <unk> was when <unk> was four years old , her father paid a monthly <unk> service fee at a time when he felt that his value was owed more than the great amount of <unk> . <unk> , in early he became frustrated at the extent by the new publisher and took responsibility for the conspiracy . <unk> <unk> was sentenced to two\n"
     ]
    }
   ],
   "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": 45,
   "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": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "tasks = {\n",
    "    \"MEN\": fetch_MEN(),\n",
    "    \"WS353\": fetch_WS353(),\n",
    "    \"SIMLEX999\": fetch_SimLex999()\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "37622"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "speech2vec = KeyedVectors.load_word2vec_format('../speech2vec-pretrained-vectors/speech2vec/50.vec', binary=False) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "our_embeddings = Embedding(\n",
    "    Vocabulary(vocab),\n",
    "    model.tok_emb.weight.detach().cpu().numpy()\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": 71,
   "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.3592662509635285\n",
      "Spearman correlation of scores on WS353 0.44876905406056566\n",
      "Spearman correlation of scores on SIMLEX999 0.23353966528046177\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": 72,
   "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
}