03aaa_transformers_words_spacy_wikitext103_larger_vocab.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 npkj\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": 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 56.4 s, sys: 3.42 s, total: 59.9 s\n",
      "Wall time: 1min 2s\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 15.1 s, sys: 268 ms, total: 15.4 s\n",
      "Wall time: 15.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": [
       "1775"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(set([w for w, _ in c.most_common(60_000)]).intersection(words))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab = set([w for w, _ in c.most_common(60_000)] + ['<eos>', '<unk>'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "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": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "32855178"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6224583"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "toks = text.split()\n",
    "len(toks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['<unk>', 'is', 'a', 'fungal', 'genus', 'in', 'the', 'family', '<unk>', ',']"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "toks[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "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, 60002 unique.\n",
      "data has 224583 words, 60002 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": 28,
   "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": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 1 iter 93747: train loss 1.68898. lr 7.500079e-05: 100%|██████████| 93748/93748 [7:46:47<00:00,  3.35it/s]  \n",
      "epoch 2 iter 93747: train loss 1.34694. lr 1.500000e-05: 100%|██████████| 93748/93748 [7:48:15<00:00,  3.34it/s]  \n",
      "epoch 3 iter 93747: train loss 1.48858. lr 7.500079e-05: 100%|██████████| 93748/93748 [7:48:10<00:00,  3.34it/s]  \n",
      "epoch 4 iter 93747: train loss 1.61718. lr 1.500000e-04: 100%|██████████| 93748/93748 [7:48:31<00:00,  3.33it/s]  \n"
     ]
    }
   ],
   "source": [
    "tconf = TrainerConfig(max_epochs=4, batch_size=64, learning_rate=1.5e-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": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), 'models/gpt_words.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "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": 32,
   "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": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "the ship to a rendezvous point . <unk> the two aircraft arms fire occurred in the hangar deck and the hangar fire killed one man and wounding others in the hangar . <unk> <unk> , the ship was able to withstand fire from all seven men . <unk> ship was also fitted with a pair of millimetre ( in ) <unk> pounder guns and twelve millimetre ( in ) anti aircraft guns . <unk> <unk> threat , <unk> had to carry out and cover after each of the ship main deck landing sequence . <unk> torpedo six inch ( mm ) <unk> torpedo tubes were initially fitted , but this changed in to a more powerful gave . <unk> of the ship sisters were removed , with the exception of <unk> <unk> . <unk> a result , her guns were removed , and <unk> was returned to service in the <unk> . <unk> the end of the war , four of the four new class were sold to ship breakers . <unk> following month , the two new <unk> class monitors were rebuilt at the outbreak of <unk> <unk> <unk> . <unk> the war , <unk> was assigned to the <unk> <unk> in the <unk> , and as a training ship , the ship was decommissioned in <unk> . <unk> was ultimately sold for scrap in and broken up for scrap in – . <unk> two guns in the <unk> class pre dreadnought battleships were intended to be retained as a target or out of service , a multi ship armament was such as the water ship <unk> <unk> , but plans called for better repair them . <unk> more triple mounts were added in early to improve the hull distance , and the ship catch was not removed . <unk> ship was protected with an anti aircraft armament used in <unk> and operated with eleven single mm anti aircraft guns . <unk> the war assessments , the ship crew was equipped with a pair of millimeter ( in ) torpedo tubes , and the other three built by the <unk> states was reduced to six to eighteen . <unk> being propelled with the torpedo carried by two and a second stint in the <unk> . <unk> she conducted coast bombardment and anti submarine gunnery training missions . <unk> was not until <unk> , when she was transferred again to <unk> .\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": 34,
   "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": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "tasks = {\n",
    "    \"MEN\": fetch_MEN(),\n",
    "    \"WS353\": fetch_WS353(),\n",
    "    \"SIMLEX999\": fetch_SimLex999()\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "speech2vec = KeyedVectors.load_word2vec_format('../speech2vec-pretrained-vectors/speech2vec/50.vec', binary=False) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "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": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Missing 131 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 28 words. Will replace them with mean vector\n",
      "Missing 5 words. Will replace them with mean vector\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Spearman correlation of scores on MEN 0.39577794507394065\n",
      "Spearman correlation of scores on WS353 0.4357339372469612\n",
      "Spearman correlation of scores on SIMLEX999 0.2039707080228247\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": 39,
   "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
}