ria_rnn.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "ria_rnn.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "authorship_tag": "ABX9TyNW0Jui5ay5vp/cQSPIGRBV",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/DmitriyValetov/nlp_course_project/blob/master/ria_rnn.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FO83gr3P1Cqc",
        "colab_type": "text"
      },
      "source": [
        "# Dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QgutxGKrctFP",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from google.colab import drive\n",
        "import os\n",
        "\n",
        "drive.mount('/content/drive')\n",
        "root = '/content/drive/My Drive/'\n",
        "data_fn = 'stop_lem_norm_sents_ria.json.gz'\n",
        "vocab_fn = 'stop_lem_norm_sents_ria.json_vocab.json'\n",
        "enc_fn = 'stop_lem_norm_sents_ria.json_enc.json.gz'\n",
        "data_path = os.path.join(root, data_fn)\n",
        "vocab_path = os.path.join(root, vocab_fn)\n",
        "enc_path = os.path.join(root, enc_fn)\n",
        "print(f'Check {data_path}: {os.path.exists(data_path)}')\n",
        "print(f'Check {vocab_path}: {os.path.exists(vocab_path)}')\n",
        "print(f'Check {enc_path}: {os.path.exists(enc_path)}')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dRlGjfpXdU88",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "from torch.utils.data import Dataset, DataLoader, random_split\n",
        "import os\n",
        "import json\n",
        "import gzip\n",
        "from pprint import pprint\n",
        "from tqdm.notebook import tqdm\n",
        "import numpy as np\n",
        "import gc\n",
        "import shutil\n",
        "\n",
        "class RiaDataset(Dataset):\n",
        "  def __init__(self, data=None, max_vocab=10000, max_samples=500,\n",
        "               vocab=None, enc_data=None, \n",
        "               special_chars=('<pad>', '<unk>', '<sos>', '<eos>')):\n",
        "    super(RiaDataset).__init__()\n",
        "    # Vocabulary dict {word: count)\n",
        "    if vocab is None:  # Create vocabualry\n",
        "      print('Creating vocabulary')\n",
        "      if data is None:\n",
        "        raise ValueError('Requires raw data path for vocabulary creating')\n",
        "      vocab = {}\n",
        "      cnt, bad_ids = 0, []\n",
        "      root, ext = os.path.splitext(data)\n",
        "      with gzip.open(data, 'rb') if ext == '.gz' else open(data) as f:\n",
        "        try:\n",
        "          for _ in tqdm(range(max_samples), desc='loading samples'):\n",
        "            n = json.loads(next(f))\n",
        "            if not len(n['text']) == 0 and not len(n['title']) == 0:\n",
        "              for s in n['text']:\n",
        "                for w in s.split():\n",
        "                  vocab[w] = vocab.setdefault(w, 0) + 1\n",
        "              for s in n['title']:\n",
        "                for w in s.split():\n",
        "                  vocab[w] = vocab.setdefault(w, 0) + 1\n",
        "            else:\n",
        "              bad_ids.append(cnt)\n",
        "            cnt += 1\n",
        "        except StopIteration as e:\n",
        "          print(f'max_samples {max_samples} > len dataset {cnt}')\n",
        "      print(f'bad texts: {len(bad_ids)} {bad_ids}')\n",
        "      vocab_path = root + '_vocab.json'\n",
        "      print(f'Saving vocabulary to {vocab_path}')\n",
        "      with open(vocab_path, 'w') as f:\n",
        "        json.dump(vocab, f)\n",
        "    elif isinstance(vocab, str):  # Load vocabualry\n",
        "      print(f'Loading vocabulary from {vocab}')\n",
        "      with open(vocab) as f:\n",
        "        vocab = json.load(f)\n",
        "    # else use as it is   \n",
        "    print('Full vocabulary')\n",
        "    print(f'unique: {len(vocab)}, total: {sum(vocab.values())}')\n",
        "    self.itos = list(special_chars)\n",
        "    top_words = [(w, c) for w, c in sorted(vocab.items(), key=lambda x: x[1],\n",
        "                                            reverse=True)][:max_vocab]\n",
        "    self.itos += [w for w, _ in top_words]\n",
        "    self.stoi = {x: i for i, x in enumerate(self.itos)}\n",
        "    assert len(self.itos) == len(self.stoi)\n",
        "    print('Reduced vocabulary')\n",
        "    print(f'unique: {len(self.stoi)}, total: {sum(c for _, c in top_words)} ~ \\\n",
        "    {sum(c for _, c in top_words)/sum(vocab.values())*100:6.3f}% of full')\n",
        "    # Samples [list of np.arrays (text), list of np.arrays (title)]\n",
        "    self.samples = [] \n",
        "    if enc_data is None:  # Encode samples\n",
        "      print('Encoding samples')\n",
        "      if data is None:\n",
        "        raise ValueError('Requires raw data path for samples encoding')\n",
        "      cnt, bad_ids = 0, []\n",
        "      root, ext = os.path.splitext(data)\n",
        "      with gzip.open(data, 'rb') if ext == '.gz' else open(data) as f:\n",
        "        try:\n",
        "          for _ in tqdm(range(max_samples), desc='loading samples'):\n",
        "            n = json.loads(next(f))\n",
        "            if not len(n['text']) == 0 and not len(n['title']) == 0:        \n",
        "              x1 = [self.encode(x) for x in n['text']]\n",
        "              x2 = [self.encode(x) for x in n['title']]\n",
        "              self.samples.append([x1, x2])\n",
        "            else:\n",
        "              bad_ids.append(cnt)\n",
        "            cnt += 1\n",
        "        except StopIteration as e:\n",
        "            print(f'max_samples {max_samples} > len dataset {cnt}')\n",
        "      print(f'bad texts: {len(bad_ids)} {bad_ids}')\n",
        "      enc_path = root + '_enc.json'\n",
        "      print(f'Saving encoded samples to {enc_path}')\n",
        "      with open(enc_path, 'w') as f:\n",
        "        for s in tqdm(self.samples):\n",
        "          json_str = json.dumps(s, default=lambda x: x.tolist())\n",
        "          f.write(json_str + '\\n')\n",
        "      comp_enc_path = enc_path + '.gz'\n",
        "      print(f'Compressing {enc_path} to {comp_enc_path}')\n",
        "      with gzip.open(comp_enc_path, 'wb') as gz_file:\n",
        "        with open(enc_path, 'rb') as json_file:\n",
        "          shutil.copyfileobj(json_file, gz_file)\n",
        "      print(f'Deleting uncompressed samples: {enc_path}')\n",
        "      os.remove(enc_path)\n",
        "    elif isinstance(enc_data, str):  # Load encodings\n",
        "      print(f'Loading encoded samples from {enc_data}')\n",
        "      _, ext = os.path.splitext(enc_data)\n",
        "      with gzip.open(enc_data, 'rb') if ext == '.gz' else open(enc_data) as f:\n",
        "        try:\n",
        "          for _ in tqdm(range(max_samples), desc='loading samples'):\n",
        "            s = json.loads(next(f))\n",
        "            x1 = [np.array(x) for x in s[0]]\n",
        "            x2 = [np.array(x) for x in s[1]]\n",
        "            self.samples.append([x1, x2])\n",
        "        except StopIteration as e:\n",
        "          print(f'max_samples {max_samples} > len dataset {len(self.samples)}')\n",
        "    else:\n",
        "      self.samples = enc_data\n",
        "    print(f'samples: {len(self.samples)}')\n",
        "\n",
        "  def __len__(self):\n",
        "    return len(self.samples)\n",
        "\n",
        "  def __getitem__(self, i):\n",
        "    return self.samples[i]\n",
        "\n",
        "  def encode(self, s):\n",
        "    return np.array([self.stoi['<sos>']] \\\n",
        "                  + [self.stoi.get(x, self.stoi['<unk>']) for x in s.split()] \\\n",
        "                  + [self.stoi['<eos>']])\n",
        "  def decode(self, s):\n",
        "    return [self.itos[x] for x in s]\n",
        "\n",
        "class Collate():\n",
        "  def __init__(self, n_x1=1, padding_value=0):\n",
        "     self.n_x1 = n_x1\n",
        "     self.padding_value = padding_value\n",
        "\n",
        "  def __call__(self, batch):\n",
        "    bx1, bx2 = [], []\n",
        "    for b in batch:\n",
        "      x1s, x2s = b\n",
        "      # skip first sentence (usually date and place)\n",
        "      ei = min(len(x1s), 1 + self.n_x1)\n",
        "      si = 1 if ei > 1 else 0\n",
        "      # print(self.n_x1, si, ei, len(x1s))\n",
        "      for x1 in x1s[si:ei]:\n",
        "        for x2 in x2s:\n",
        "          bx1.append(torch.as_tensor(x1))\n",
        "          bx2.append(torch.as_tensor(x2))\n",
        "    bx1 = torch.nn.utils.rnn.pad_sequence(bx1, batch_first=True, \n",
        "                                          padding_value=0)\n",
        "    bx2 = torch.nn.utils.rnn.pad_sequence(bx2, batch_first=True,\n",
        "                                          padding_value=0)\n",
        "    batch = [bx1, bx2]\n",
        "    return batch\n",
        "\n",
        "# Create vocabulary and encodings  # max_samples=1003869 max_vocab=50000\n",
        "# ds = RiaDataset(data=data_path, max_vocab=50000, max_samples=2003869, \n",
        "#                 vocab=None, enc_data=None)\n",
        "# Load vocabulary and encodings\n",
        "ds = RiaDataset(data=None, max_vocab=50000, max_samples=1003869, \n",
        "                vocab=vocab_path, enc_data=enc_path) \n",
        "# unique words: 601956, total words: 198199252\n",
        "# vocabulary: 50004, words in vocabulary: 194385367\n",
        "\n",
        "torch.manual_seed(0)\n",
        "torch.cuda.manual_seed(0)\n",
        "# dl = DataLoader(ds, batch_size=1, num_workers=1, shuffle=False, \n",
        "#                 drop_last=False, collate_fn=Collate())\n",
        "# for s in tqdm(dl, desc='dataset'):\n",
        "#   bx1, bx2 = s\n",
        "#   assert bx1.size()[0] == bx2.size()[0]\n",
        "#   print(bx1.size(), bx2.size())\n",
        "#   for x1, x2 in zip(bx1, bx2):\n",
        "#     print(ds.decode(x1), ds.decode(x2))\n",
        "#   print([ds.decode(x) for x in bx1])\n",
        "#   print([ds.decode(x) for x in bx2])\n",
        "train_len = int(0.7*len(ds))\n",
        "test_len = int(0.2*len(ds))\n",
        "val_len = len(ds) - train_len - test_len\n",
        "lens = [train_len, test_len, val_len]\n",
        "print(lens, sum(lens), len(ds))\n",
        "train_ds, test_ds, val_ds = random_split(ds, lens)\n",
        "train_dl = DataLoader(train_ds, batch_size=2, num_workers=1, \n",
        "                      shuffle=True, drop_last=False, \n",
        "                      collate_fn=Collate(n_x1=1))\n",
        "test_dl = DataLoader(test_ds, batch_size=3, num_workers=1, \n",
        "                     shuffle=False, drop_last=False, \n",
        "                     collate_fn=Collate(n_x1=1))\n",
        "val_dl = DataLoader(val_ds, batch_size=3, num_workers=1, \n",
        "                    shuffle=False, drop_last=False, \n",
        "                    collate_fn=Collate(n_x1=1))\n",
        "# for s in tqdm(train_dl, desc='train'):\n",
        "#   bx1, bx2 = s\n",
        "#   assert bx1.size()[0] == bx2.size()[0]\n",
        "#   # print(bx1.size(), bx2.size())\n",
        "# for s in tqdm(test_dl, desc='test'):\n",
        "#   bx1, bx2 = s\n",
        "#   assert bx1.size()[0] == bx2.size()[0]\n",
        "#   # print(bx1.size(), bx2.size())\n",
        "# for s in tqdm(val_dl, desc='val'):\n",
        "#   bx1, bx2 = s\n",
        "#   assert bx1.size()[0] == bx2.size()[0]\n",
        "#   # print(bx1.size(), bx2.size())\n",
        "gc.collect()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eYgtjOABt2rJ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# %matplotlib inline\n",
        "# import matplotlib.pyplot as plt\n",
        "# import seaborn as sns\n",
        "# import pandas as pd\n",
        "\n",
        "# max_vocab = 1000\n",
        "# with open(vocab_path) as f:\n",
        "#   vocab = json.load(f)\n",
        "# top_words = dict([(w, c) for w, c in sorted(vocab.items(), key=lambda x: x[1],\n",
        "#                                             reverse=True)][:max_vocab])\n",
        "# # # sns.distplot(list(map(len, vocab.keys())), kde=False)\n",
        "# # # plt.figure()\n",
        "# # # sns.distplot(list(map(len, top_words.keys())), kde=True)\n",
        "# # # plt.figure()\n",
        "# sns.distplot(list(vocab.values()), kde=False, hist_kws={'log': True})\n",
        "# sns.distplot(list(top_words.values()), kde=False, hist_kws={'log': True})\n",
        "# ax = sns.distplot(list(vocab.values()), kde=False, hist=True, \n",
        "#              hist_kws={'log': True}, norm_hist=False)\n",
        "# # ax.set_xscale('log')\n",
        "# ax = sns.distplot(list(top_words.values()), kde=False, hist=True, \n",
        "#              hist_kws={'log': True}, norm_hist=False)\n",
        "# # ax.set_xscale('log')\n",
        "# plt.figure()\n",
        "# plt.title('Sentences lengths')\n",
        "# sns.kdeplot([len(y) for x in ds.samples for y in x[0]], \n",
        "#             label='texts', shade=True, bw=2)\n",
        "# sns.kdeplot([len(y) for x in ds.samples for y in x[1]], \n",
        "#             label='titles', shade=True, bw=2)\n",
        "# plt.legend()\n",
        "# plt.figure()\n",
        "# sns.jointplot(x='titles lengths', y='texts lengths', data=pd.DataFrame({\n",
        "#     \"texts lengths\": [sum(map(len, x[0])) for x in ds.samples],\n",
        "#     \"titles lengths\": [sum(map(len, x[1])) for x in ds.samples]}\n",
        "# ), kind='kde', n_levels=10)\n",
        "# sns.jointplot(x='titles 1 sentense length', y='texts 1 sentense length',\n",
        "#               data=pd.DataFrame({\n",
        "#     \"texts 1 sentense length\": [len(x[0][0]) for x in ds.samples],\n",
        "#     \"titles 1 sentense length\": [len(x[1][0]) for x in ds.samples]}\n",
        "# ), kind='kde', n_levels=10)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bwYBLWZRXQqS",
        "colab_type": "text"
      },
      "source": [
        "# Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AfI1NnKX9c1o",
        "colab_type": "code",
        "cellView": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence\n",
        "\n",
        "class EncoderRNN(nn.Module):\n",
        "  def __init__(self, num_embeddings, embedding_dim, padding_idx,\n",
        "               rnn_type, hidden_size, num_layers=1, rnn_dropout=0,\n",
        "               bidirectional=False, dropout=0, pack=False):\n",
        "    super(EncoderRNN, self).__init__()\n",
        "    rnn_map = {'RNN': nn.RNN, 'GRU': nn.GRU, 'LSTM': nn.LSTM}\n",
        "    self.embedding = nn.Embedding(num_embeddings=num_embeddings, \n",
        "                                  embedding_dim=embedding_dim,\n",
        "                                  padding_idx=padding_idx)\n",
        "    self.rnn = rnn_map[rnn_type](input_size=embedding_dim, \n",
        "                                 hidden_size=hidden_size,\n",
        "                                 num_layers=num_layers,\n",
        "                                 batch_first=True,\n",
        "                                 dropout=rnn_dropout,\n",
        "                                 bidirectional=bidirectional)\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "    self.pack = pack\n",
        "\n",
        "  def forward(self, x1):  # [B, L]\n",
        "    x = self.embedding(x1)  # [B, L] -> [B, L, E]\n",
        "    if self.pack:  # [B, L, E] -> Packed [B, L, E]\n",
        "      ls = torch.sum(x1 != 0, dim=1)\n",
        "      x = pack_padded_sequence(x, ls, batch_first=True, enforce_sorted=False)\n",
        "    # [B, L, E] -> (Packed) [B, L, ND*H], ([NL*ND, B, H], ([NL*ND, B, H]))\n",
        "    if isinstance(self.rnn, nn.LSTM):\n",
        "      ht, (hn, cn) = self.rnn(x)\n",
        "      hn, cn = self.dropout(hn), self.dropout(cn)\n",
        "      return ht, (hn, cn)\n",
        "    else:\n",
        "      ht, hn = self.rnn(x)\n",
        "      hn = self.dropout(hn)\n",
        "      return ht, hn\n",
        "\n",
        "\n",
        "class AttentionDecoderRNN(nn.Module):\n",
        "  def __init__(self, num_embeddings, embedding_dim, padding_idx,\n",
        "               rnn_type, hidden_size, num_layers=1, rnn_dropout=0,\n",
        "               bidirectional=False, dropout=0, out_hidden=0,\n",
        "               attn_type='soft_dot', pack=False):\n",
        "    super(AttentionDecoderRNN, self).__init__()\n",
        "    rnn_map = {'RNN': nn.RNN, 'GRU': nn.GRU, 'LSTM': nn.LSTM}\n",
        "    attn_types = ['dot', 'cos', 'dist', 'soft_dot', 'soft_cos', 'soft_dist', 'none']\n",
        "    self.attn_type = attn_type\n",
        "    self.embedding = nn.Embedding(num_embeddings=num_embeddings, \n",
        "                                  embedding_dim=embedding_dim,\n",
        "                                  padding_idx=padding_idx)\n",
        "    self.rnn = rnn_map[rnn_type](input_size=embedding_dim, \n",
        "                                 hidden_size=hidden_size,\n",
        "                                 num_layers=num_layers,\n",
        "                                 batch_first=True,\n",
        "                                 dropout=rnn_dropout,\n",
        "                                 bidirectional=bidirectional)\n",
        "    out_input = hidden_size if self.attn_type == 'none' else 2*hidden_size\n",
        "    if out_hidden > 0:\n",
        "      self.out_hidden = nn.Linear(out_input, out_hidden)\n",
        "      self.out = nn.Linear(out_hidden, num_embeddings)\n",
        "    else:\n",
        "      self.out_hidden = None\n",
        "      self.out = nn.Linear(out_input, num_embeddings)\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "    self.softmax = nn.LogSoftmax(dim=2)\n",
        "    self.pack = pack\n",
        "\n",
        "  def forward(self, x2, h1, x1):  # [B, L], [B, L, ND*H], ([NL*ND, B, H], [NL*ND, B, H])\n",
        "    x = self.embedding(x2)  # [B, L] -> [B, L, E]\n",
        "    if self.pack:  # [B, L, E] -> Packed [B, L, E]\n",
        "      lengths = torch.sum(x2 != 0, dim=1)\n",
        "      x = pack_padded_sequence(x, lengths, batch_first=True, enforce_sorted=False)\n",
        "    ht1, hns1 = h1  # [B, L, ND*H], ([NL*ND, B, H], [NL*ND, B, H])\n",
        "    # [B, L, D], ([NL*ND, B, H], [NL*ND, B, H]) -> [B, L, ND*H], ([NL*ND, B, H], [NL*ND, B, H])\n",
        "    ht2, hns2 = self.rnn(x, hns1)\n",
        "    if self.pack:  # Packed [B, L, E] -> [B, L, E]\n",
        "      ht1, _ = pad_packed_sequence(ht1, batch_first=True)\n",
        "      ht2, _ = pad_packed_sequence(ht2, batch_first=True)\n",
        "    if self.attn_type == 'none':\n",
        "      if self.rnn.bidirectional:\n",
        "        B2, L2, NDH2 = ht2.size()\n",
        "        # [B, L, ND*H] -> [B, L, ND, H] -> [B, L, H]\n",
        "        ht2 = ht2.view(B2, L2, 2, int(NDH2/2)).mean(2)\n",
        "      x = self.dropout(ht2)\n",
        "      if self.out_hidden is not None:\n",
        "        x = self.out_hidden(x)\n",
        "      x = self.out(x)\n",
        "      x = self.softmax(x)\n",
        "      return x, None\n",
        "    else:\n",
        "      if self.rnn.bidirectional:\n",
        "        B1, L1, NDH1 = ht1.size()\n",
        "        B2, L2, NDH2 = ht2.size()  # B1 == B2, NDH1 == NDH2\n",
        "        # [B, L, ND*H] -> [B, L, ND, H] -> [B, L, H]\n",
        "        ht1 = ht1.view(B1, L1, 2, int(NDH1/2)).mean(2)\n",
        "        ht2 = ht2.view(B2, L2, 2, int(NDH2/2)).mean(2)\n",
        "      # [L2, H], [L1, H] -> [L2, L1]\n",
        "      # mask where x1 and x2 token is <PAD>\n",
        "      # pad_mask = torch.einsum('bi,bj->bij', x2, x1) == 0\n",
        "      # mask where only x1 token is <PAD>\n",
        "      pad_mask = torch.einsum('bi,bj->bij', torch.ones_like(x2), x1) == 0\n",
        "      # [B, L2, H], [B, L1, H] -> [B, L2, L1]  # attention\n",
        "      # [B, L2, L1], [B, L1, H] -> [B, L2, H]  # weighted h1\n",
        "      if self.attn_type == 'dot':\n",
        "        attn = torch.einsum('bih,bjh->bij', ht2, ht1)  # dot product\n",
        "        attn[pad_mask] = 0\n",
        "      elif self.attn_type == 'soft_dot':\n",
        "        attn = torch.einsum('bih,bjh->bij', ht2, ht1)  # dot product\n",
        "        attn[pad_mask] = float('-inf')\n",
        "        attn = F.softmax(attn, 2)\n",
        "      elif self.attn_type == 'dist':\n",
        "        ht1, ht2 = ht1.contiguous(), ht2.contiguous()\n",
        "        attn = torch.cdist(ht2, ht1)  # euclidian distance\n",
        "        attn = torch.masked_fill(attn, pad_mask, float('inf'))\n",
        "        attn = F.threshold(attn, threshold=1e-6, value=1e-6)  # short dist\n",
        "        attn = 1/attn  # inverse\n",
        "        attn = F.normalize(attn, p=1, dim=2)  # to [0, 1]\n",
        "      elif self.attn_type == 'soft_dist':\n",
        "        ht1, ht2 = ht1.contiguous(), ht2.contiguous()\n",
        "        attn = torch.cdist(ht2, ht1)  # euclidian distance\n",
        "        attn = torch.masked_fill(attn, pad_mask, float('inf'))\n",
        "        attn = F.softmin(attn, 2)\n",
        "      elif self.attn_type == 'cos':\n",
        "        ht1n = F.normalize(ht1, p=2, dim=2)  # normalize to length = 1\n",
        "        ht2n = F.normalize(ht2, p=2, dim=2)  # normalize to length = 1\n",
        "        attn = torch.einsum('bih,bjh->bij', ht2n, ht1n)  # dot product\n",
        "        attn[pad_mask] = 0\n",
        "      elif self.attn_type == 'soft_cos':\n",
        "        ht1n = F.normalize(ht1, p=2, dim=2)  # normalize to length = 1\n",
        "        ht2n = F.normalize(ht2, p=2, dim=2)  # normalize to length = 1\n",
        "        attn = torch.einsum('bih,bjh->bij', ht2n, ht1n)  # dot product\n",
        "        attn[pad_mask] = float('-inf')\n",
        "        attn = F.softmax(attn, 2)\n",
        "      hw1 = torch.einsum('bij,bjh->bih', attn, ht1)  # weighted h1\n",
        "      ha = torch.cat((ht2, hw1), 2)  # [B, L2, H], [B, L2, H] -> [B, L2, H+H]\n",
        "      x = self.dropout(ha)  # [B, L2, H+H] -> [B, L2, H+H]\n",
        "      if self.out_hidden is not None:\n",
        "        x = self.out_hidden(x)   # [B, L2, H+H] -> [B, L2, OH]\n",
        "      x = self.out(x)  # [B, L2, H+H] or [B, L2, OH] -> [B, L2, D2]\n",
        "      y2 = self.softmax(x)  # [B, L2, D2] -> [B, L2, D2]\n",
        "      return y2, attn\n",
        "\n",
        "\n",
        "class EncoderDecoder(nn.Module):\n",
        "  def __init__(self, \n",
        "               enc_num_embeddings, enc_embedding_dim, enc_padding_idx,\n",
        "               dec_num_embeddings, dec_embedding_dim, dec_padding_idx,\n",
        "               rnn_type, hidden_size, num_layers=1, out_hidden=0,\n",
        "               enc_rnn_dropout=0, dec_rnn_dropout=0,\n",
        "               bidirectional=False, enc_dropout=0, dec_dropout=0, \n",
        "               attn_type='dot', pack=False):\n",
        "    super(EncoderDecoder, self).__init__()\n",
        "    self.encoder = EncoderRNN(num_embeddings=enc_num_embeddings, \n",
        "                              embedding_dim=enc_embedding_dim, \n",
        "                              padding_idx=enc_padding_idx,\n",
        "                              hidden_size=hidden_size, \n",
        "                              rnn_type=rnn_type, \n",
        "                              bidirectional=bidirectional,\n",
        "                              num_layers=num_layers,\n",
        "                              dropout=enc_dropout,\n",
        "                              rnn_dropout=enc_rnn_dropout,\n",
        "                              pack=pack)\n",
        "    self.decoder = AttentionDecoderRNN(num_embeddings=dec_num_embeddings, \n",
        "                                       embedding_dim=dec_embedding_dim, \n",
        "                                       padding_idx=dec_padding_idx,\n",
        "                                       hidden_size=hidden_size, \n",
        "                                       rnn_type=rnn_type, \n",
        "                                       bidirectional=bidirectional,\n",
        "                                       num_layers=num_layers,\n",
        "                                       dropout=dec_dropout,\n",
        "                                       out_hidden=out_hidden,\n",
        "                                       rnn_dropout=dec_rnn_dropout,\n",
        "                                       attn_type=attn_type,\n",
        "                                       pack=pack)\n",
        "\n",
        "  def forward(self, x1, x2):\n",
        "    #  [B, L1] -> [B, L1, ND*H], ([NL*ND, B, H], [NL*ND, B, H])\n",
        "    h1 = self.encoder(x1)\n",
        "    # [B, L2], ([B, L1, ND*H], ([NL*ND, B, H], [NL*ND, B, H])) -> [B, L2, E2]\n",
        "    y2 = self.decoder(x2, h1, x1)\n",
        "    return y2\n",
        "\n",
        "\n",
        "def external_attn(x1, x2, attn_dict):\n",
        "    attn_dict = {1: [1], 2: [2], 4: [4, 3]}\n",
        "    attn = []\n",
        "    B1, L1 = x1.size()\n",
        "    B2, L2 = x2.size()\n",
        "    for i in range(B1):\n",
        "      b = []\n",
        "      for j in range(L2):\n",
        "        l2 = []\n",
        "        # x2t = x2[i, j].item()\n",
        "        x2t = x2[i, j + 1].item() if j + 2 < L2 else 0 # decoder shift\n",
        "        for k in range(L1):\n",
        "          x1t = x1[i, k].item()\n",
        "          # x1t = x1[i, k + 1].item() if k + 1 < L1 else 0 # encoder shift\n",
        "          x1ts = self.attn_dict.get(x2t, [])\n",
        "          if x1t in x1ts:\n",
        "            l2.append(1.)\n",
        "          else:\n",
        "            l2.append(0.)\n",
        "        b.append(l2)\n",
        "      attn.append(b)\n",
        "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "    attn = torch.tensor(attn).to(device)\n",
        "    attn = F.normalize(attn, p=1, dim=2)  # to [0, 1]\n",
        "    return attn\n",
        "\n",
        "from IPython.display import Image\n",
        "# Image(make_dot(loss).render('loss', format='png'))\n",
        "from tqdm.notebook import tqdm\n",
        "from torch import optim\n",
        "import numpy as np\n",
        "\n",
        "def plot_attention(a, x1, x2, shift=True, mask=True, \n",
        "                   suptitle='Attention', figsize=None, \n",
        "                   decoder_x1=None, decoder_x2=None, tight=False,\n",
        "                   labels=True):\n",
        "  # %matplotlib inline\n",
        "  import matplotlib.pyplot as plt\n",
        "  from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
        "\n",
        "  b = a.shape[0]\n",
        "  fig, axs = plt.subplots(1, b, figsize=figsize)\n",
        "  if not isinstance(axs, np.ndarray):  # if batch size == 1\n",
        "    axs = [axs]\n",
        "  fig.suptitle(suptitle)\n",
        "  for i in range(b):\n",
        "    ax = axs[i]\n",
        "    if shift:\n",
        "      ba, bx1, bx2 = a[i,:-1,1:], x1[i,1:], x2[i,1:]\n",
        "    else:\n",
        "      ba, bx1, bx2 = a[i], x1[i], x2[i]\n",
        "    if mask:\n",
        "      mask_x1 = np.flatnonzero(bx1)\n",
        "      mask_x2 = np.flatnonzero(bx2)\n",
        "      ba, bx1, bx2 =  ba[mask_x2,:][:,mask_x1], bx1[mask_x1], bx2[mask_x2]\n",
        "    # ax.set_title(f'{i+1}', y=-0.2)\n",
        "    im = ax.imshow(ba, cmap='gray')\n",
        "    ax.set_xticks(np.arange(len(bx1)))\n",
        "    ax.set_yticks(np.arange(len(bx2)))\n",
        "    if labels:\n",
        "      ax.set_xlabel('x1')\n",
        "      ax.set_ylabel('x2')\n",
        "    ax.xaxis.tick_top()\n",
        "    ax.xaxis.set_label_position('top')\n",
        "    if decoder_x1 is not None:\n",
        "      bx1 = decoder_x1(bx1)\n",
        "    if decoder_x2 is not None:\n",
        "      bx2 = decoder_x2(bx2)\n",
        "    ax.set_xticklabels(bx1, rotation=90)  # rotation=90\n",
        "    ax.set_yticklabels(bx2)\n",
        "    fig.colorbar(im, ax=ax, fraction=0.05, pad=0.05)\n",
        "  if tight:\n",
        "    plt.tight_layout()\n",
        "\n",
        "\n",
        "# torch.backends.cudnn.enabled=False\n",
        "# torch.backends.cudnn.deterministic=True\n",
        "# torch.autograd.set_detect_anomaly(True)\n",
        "\n",
        "# enc_dec_config = {\n",
        "#   'enc_num_embeddings': 6,\n",
        "#   'enc_embedding_dim': 2,\n",
        "#   'enc_padding_idx': 0,\n",
        "#   'dec_num_embeddings': 6,\n",
        "#   'dec_embedding_dim': 2,\n",
        "#   'dec_padding_idx': 0,\n",
        "#   'rnn_type': 'RNN',\n",
        "#   'hidden_size': 2,\n",
        "#   'num_layers': 1,\n",
        "#   'bidirectional': False,\n",
        "#   'enc_rnn_dropout': 0,\n",
        "#   'dec_rnn_dropout': 0,\n",
        "#   'enc_dropout': 0,\n",
        "#   'dec_dropout': 0,\n",
        "#   'attn_type': 'soft_dist',\n",
        "#   'out_hidden': 16,\n",
        "#   'pack': True\n",
        "# }\n",
        "# # <pad>, <unk>, <go>, <eos>\n",
        "# device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "# seed = 0\n",
        "# torch.manual_seed(seed)\n",
        "# torch.cuda.manual_seed(seed)\n",
        "# model = EncoderDecoder(**enc_dec_config).to(device)\n",
        "# x1 = torch.tensor([[2, 1, 4, 3, 0], [2, 1, 4, 1, 3]]).to(device)\n",
        "# x2 = torch.tensor([[2, 4, 1, 5, 3, 0], [2, 4, 5, 1, 1, 3]]).to(device)\n",
        "# opt = optim.Rprop(model.parameters(), lr=0.01)\n",
        "# loss_fn = torch.nn.NLLLoss(ignore_index=0)  # 0 is <PAD>\n",
        "# pbar = tqdm(range(200))\n",
        "# for i in pbar:\n",
        "#   opt.zero_grad()\n",
        "#   t = x2[:,1:].flatten(start_dim=0)\n",
        "#   y2, attn = model(x1, x2)\n",
        "#   p = y2[:,:-1,:].flatten(start_dim=0, end_dim=1)\n",
        "#   loss = loss_fn(p, t)\n",
        "#   loss.backward()\n",
        "#   opt.step()\n",
        "#   with torch.no_grad():\n",
        "#     idx = torch.nonzero(t).view(-1)\n",
        "#     acc = torch.sum(torch.argmax(p, 1)[idx] == t[idx]).float()/idx.size()[0]\n",
        "#     ps = torch.argmax(y2[:,:-1,:], 2).detach().cpu().numpy()\n",
        "#     ts = x2[:,1:].detach().cpu().numpy()\n",
        "#     for sp, st in zip(ps, ts):\n",
        "#       t_set = set(st) - {0, 1, 2, 3}\n",
        "#       eos = sp[sp == 3][0] if 3 in sp else len(sp)\n",
        "#       p_eos = sp[:eos]\n",
        "#       p_set = set(p_eos) - {0, 1, 2, 3}\n",
        "#       i_set = t_set.intersection(p_set)\n",
        "#       pre = len(i_set)/len(p_set) if len(p_set) != 0 else 0\n",
        "#       rec = len(i_set)/len(t_set) if len(t_set) != 0 else 0\n",
        "#       f1 = 2*pre*rec/(pre + rec) if pre + rec != 0 else 0\n",
        "#       # print(st, sp, p_eos)\n",
        "#       # print(t_set, p_set)\n",
        "#       # print(f'precision: {pre}, recall: {rec}, f1: {f1}')\n",
        "#     if i % 25 == 0:\n",
        "#       # print(attn)\n",
        "#       # print(p)\n",
        "#       # print(torch.argmax(p, 1))\n",
        "#       # print(t)\n",
        "#       if attn is not None:\n",
        "#          plot_attention(attn.detach().cpu().numpy(), \n",
        "#                         x1.detach().cpu().numpy(), \n",
        "#                         x2.detach().cpu().numpy(),\n",
        "#                         labels=True, tight=False,\n",
        "#                         shift=False, mask=False, figsize=(10, 5))\n",
        "#   pbar.set_description(f'loss: {loss:.3f}, acc: {acc:.3f}')\n",
        "#   if acc == 1:\n",
        "#     break\n",
        "# # Image(make_dot(attn).render('attn', format='png'))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AZPHNG9iXXou",
        "colab_type": "text"
      },
      "source": [
        "# Train"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bDVQqjgHhR9O",
        "colab_type": "text"
      },
      "source": [
        "## GPU/CPU stats"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PUcqiqInpBGs",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install gputil  # for monitoring\n",
        "!pip install psutil  # for monitoring\n",
        "!pip install humanize  # for monitoring\n",
        "\n",
        "!pip install optuna  # for train\n",
        "\n",
        "!pip install rouge_score  # for metrics"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "icz-Yt2dikV-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!ln -sf /opt/bin/nvidia-smi /usr/bin/nvidia-smi\n",
        "import subprocess\n",
        "print(subprocess.getstatusoutput('nvidia-smi')[1])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Jb4ephjLpFFY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Import packages\n",
        "import os, sys, humanize, psutil, GPUtil\n",
        "\n",
        "# Define function\n",
        "def mem_report(verbose=True):\n",
        "  h = humanize.naturalsize\n",
        "  cmf = psutil.virtual_memory().available\n",
        "  cmt = psutil.virtual_memory().total\n",
        "  if verbose: \n",
        "    print(f\"CPU mem Free: {h(cmf)} / {h(cmt)}\")\n",
        "  gmfs, gmts, gmus = [], [], []\n",
        "  GPUs = GPUtil.getGPUs()\n",
        "  for i, gpu in enumerate(GPUs):\n",
        "    gmfs.append(gpu.memoryFree)\n",
        "    gmts.append(gpu.memoryTotal)\n",
        "    gmus.append(gpu.memoryUtil)\n",
        "    if verbose:\n",
        "      print(f'GPU {i} mem free: {h(gmfs[-1]*1000000)} / {h(gmts[-1]*1000000)} util: {int(gmus[-1]*100)}%')\n",
        "  return cmf, gmfs, gmts, gmus\n",
        "mem_report()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LksNjL0ohmNR",
        "colab_type": "text"
      },
      "source": [
        "## Beam"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KcU-0oSABot9",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%matplotlib inline\n",
        "import matplotlib.pyplot as plt\n",
        "import seaborn as sns\n",
        "\n",
        "reduction_map = {\n",
        "  'sum': torch.sum,\n",
        "  'mean': torch.mean,\n",
        "  'none': None\n",
        "}\n",
        "\n",
        "class Forced():\n",
        "  def __init__(self, **kwargs):\n",
        "    pass\n",
        "\n",
        "  def __call__(self, model, bx1, bx2):\n",
        "    by2p, attn = model(bx1, bx2)\n",
        "    y2 = torch.argmax(by2p[:,:-1,:], axis=2)\n",
        "    return y2, attn  \n",
        "\n",
        "class Greedy():\n",
        "  def __init__(self, sos=2, eos=3, max_len=30, **kwargs):\n",
        "    self.sos, self.eos, self.max_len = sos, eos, max_len\n",
        "    self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "\n",
        "  def __call__(self, model, bx1, bx2=None):\n",
        "    batch_size = bx1.size(0)  # input batch_size\n",
        "    bx2 = torch.full((batch_size, 1), self.sos, dtype=bx1.dtype).to(self.device)  # batch with <sos>\n",
        "    # stop when predictions len > max_len or all have <eos> token\n",
        "    while bx2.size(1) - 1 < self.max_len and not torch.all(torch.any(bx2 == self.eos, axis=1)):\n",
        "      by2p, attn = model(bx1, bx2)\n",
        "      next_bx2 = torch.argmax(by2p[:,-1,:], axis=1, keepdim=True)\n",
        "      bx2 = torch.cat((bx2, next_bx2), 1)\n",
        "    return bx2[:,1:], attn  \n",
        "\n",
        "class Beam():\n",
        "  def __init__(self, sos=2, eos=3, max_len=30, beam_width=2, beam_depth=2,\n",
        "               depth_reduction='sum', beam_reduction='none', \n",
        "               batch_reduction='none', max_input_size=30000, \n",
        "               first_beam_width=10, verbose=0,\n",
        "               plot_dist=False, plot_size=100, **kwargs):\n",
        "    self.sos, self.eos, self.max_len = sos, eos, max_len\n",
        "    self.beam_width, self.beam_depth = beam_width, beam_depth\n",
        "    self.depth_reduction = reduction_map[depth_reduction]\n",
        "    self.beam_reduction = reduction_map[beam_reduction]\n",
        "    self.batch_reduction = reduction_map[batch_reduction]\n",
        "    self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "    self.max_input_size = max_input_size\n",
        "    self.first_beam_width = first_beam_width\n",
        "    self.plot_size = plot_size\n",
        "    self.plot_dist = plot_dist\n",
        "    self.verbose = verbose\n",
        "\n",
        "  def __call__(self, model, bx1, bx2=None):\n",
        "    batch_size = bx1.size(0)  # input batch_size\n",
        "    device = torch\n",
        "    bx2 = torch.full((batch_size, 1), self.sos, dtype=bx1.dtype).to(self.device)  # batch with <sos>\n",
        "    # stop when len decoder output > max_len or all decoder outputs have <eos> token\n",
        "    while bx2.size(1) - 1 < self.max_len and not torch.all(torch.any(bx2 == self.eos, axis=1)):\n",
        "      beam_scores = torch.empty((batch_size, 0)).to(self.device)  # scores for each beam\n",
        "      bx1t = bx1 if batch_size == 1 else bx1.clone()\n",
        "      for i in range(self.beam_depth):\n",
        "        eos_mask = torch.any(bx2 == self.eos, axis=1)  # beams with <eos>\n",
        "        cur_batch_size = torch.sum(~eos_mask)  # beams without <eos>\n",
        "        cur_input_size = bx1t.numel() + bx2.numel()  # B*L1 + B*L2\n",
        "        cur_beam_width = self.beam_width if i > 0 else self.first_beam_width\n",
        "        if self.verbose:\n",
        "          print(f'{i} {bx1t.size()}, {bx2.size()}, {cur_batch_size}, {cur_input_size}')\n",
        "        if bx2.size(1) - 1 >= self.max_len or cur_batch_size == 0:\n",
        "          break\n",
        "        if cur_input_size > self.max_input_size:  # split into parts\n",
        "          n_parts = cur_input_size // self.max_input_size + 1\n",
        "          bx1t_parts = torch.chunk(bx1t[~eos_mask], n_parts)\n",
        "          bx2_parts = torch.chunk(bx2[~eos_mask], n_parts)\n",
        "          masked_by2p = torch.empty((0, cur_beam_width)).to(self.device)\n",
        "          masked_bx2 = torch.empty((0, cur_beam_width), dtype=bx2.dtype).to(self.device)\n",
        "          if self.plot_dist:\n",
        "            fig, axs = plt.subplots(ncols=2)\n",
        "          for j, (bx1t_p, bx2_p) in enumerate(zip(bx1t_parts, bx2_parts)):\n",
        "            if self.verbose:\n",
        "              print(f'{j+1}/{n_parts} {bx1t_p.size()}, {bx2_p.size()}, {bx1t_p.numel() + bx2_p.numel()}')\n",
        "            by2p_p, attn = model(bx1t_p, bx2_p)\n",
        "            masked_by2p_p, masked_bx2_p = torch.topk(by2p_p[:,-1,:], \n",
        "                                                     cur_beam_width)\n",
        "            masked_by2p = torch.cat((masked_by2p, masked_by2p_p))\n",
        "            masked_bx2 = torch.cat((masked_bx2, masked_bx2_p))\n",
        "            if self.plot_dist:\n",
        "              sns.distplot(torch.topk(by2p_p[:,-1,:], self.plot_size)[0], \n",
        "                            hist=True, kde=False, label=f'{j+1}/{n_parts}',\n",
        "                            ax=axs[1])\n",
        "              for k in range(by2p_p.size(0)):\n",
        "                sns.distplot(torch.topk(by2p_p[k,-1,:], self.plot_size)[0], \n",
        "                             hist=True, kde=False, ax=axs[0])\n",
        "          if self.plot_dist:\n",
        "            fig.suptitle(f'step {bx2.size(1)} top {self.plot_size} scores')\n",
        "            axs[0].set_title(f'by beam')\n",
        "            axs[1].set_title(f'by part')\n",
        "            plt.legend()\n",
        "        else:\n",
        "          by2p, attn = model(bx1t[~eos_mask], bx2[~eos_mask])  # predict\n",
        "          if i > 0 and self.beam_reduction is not None:  # beam reduction\n",
        "            if i == 1:\n",
        "              prev_prev_batch_size = batch_size\n",
        "            else:\n",
        "              prev_prev_batch_size = batch_size*self.first_beam_width*self.beam_width**(i-1)\n",
        "            by2p = by2p.view(prev_prev_batch_size, cur_beam_width, by2p.size(1), -1)  # split predictions into batches\n",
        "            by2p = self.beam_reduction(by2p, axis=1) # cumulative beams predictions\n",
        "            by2p = by2p.repeat(cur_beam_width, 1, 1)  # return to prev batch size\n",
        "          if self.batch_reduction is not None:  # batch reduction\n",
        "            if i > 0:\n",
        "              prev_batch_size = batch_size*self.first_beam_width*self.beam_width**(i-1)\n",
        "            else:\n",
        "              prev_batch_size = batch_size\n",
        "            by2p = self.batch_reduction(by2p, axis=0) # cumulative batch predictions\n",
        "            masked_by2p, masked_bx2 = torch.topk(by2p[-1,:], cur_beam_width)  # beams to top k last predictions\n",
        "            masked_bx2 = masked_bx2.repeat(prev_batch_size, 1)  # return to prev batch size\n",
        "            masked_by2p = masked_by2p.repeat(prev_batch_size, 1)  # return to prev batch size\n",
        "          else:\n",
        "            if self.plot_dist:\n",
        "              fig, axs = plt.subplots(ncols=2)\n",
        "              # sns.distplot(by2p[:,-1,:], hist=True, kde=False, label=f'{bx2.size(1)}')\n",
        "              sns.distplot(torch.topk(by2p[:,-1,:], self.plot_size)[0], \n",
        "                           hist=True, kde=False, ax=axs[1])\n",
        "              for k in range(by2p.size(0)):\n",
        "                sns.distplot(torch.topk(by2p[k,-1,:], self.plot_size)[0], \n",
        "                             hist=True, kde=False, ax=axs[0])\n",
        "              fig.suptitle(f'step {bx2.size(1)} top {self.plot_size} scores')\n",
        "              axs[0].set_title(f'by beam')\n",
        "              axs[1].set_title(f'cumulative')\n",
        "            masked_by2p, masked_bx2 = torch.topk(by2p[:,-1,:], cur_beam_width)  # beams to top k last predictions\n",
        "        next_by2p = torch.full((bx2.size(0), cur_beam_width), \n",
        "                               float('-inf'), dtype=by2p.dtype).to(self.device)\n",
        "        next_bx2 = torch.full((bx2.size(0), cur_beam_width), \n",
        "                              0, dtype=bx2.dtype).to(self.device)\n",
        "        next_by2p[~eos_mask] = masked_by2p  # from beams without <eos>\n",
        "        next_bx2[~eos_mask] = masked_bx2  # from beams without <eos>\n",
        "        # update bx2 and scores\n",
        "        new_batch_size = batch_size*self.first_beam_width*self.beam_width**i\n",
        "        next_bx2 = next_bx2.view(new_batch_size, 1)  # new beams\n",
        "        next_by2p = next_by2p.view(new_batch_size, 1)  # new scores\n",
        "        beam_scores = torch.repeat_interleave(beam_scores, cur_beam_width, 0)  # increase batch for new scores\n",
        "        bx2 = torch.repeat_interleave(bx2, cur_beam_width, 0)  # increase batch for new beams\n",
        "        bx2 = torch.cat((bx2, next_bx2), 1)  # add beams\n",
        "        beam_scores = torch.cat((beam_scores, next_by2p), 1)  # add beams scores\n",
        "        if batch_size == 1:\n",
        "          bx1t = bx1.expand(new_batch_size, -1)  # increase batch for new beams\n",
        "        else:\n",
        "          bx1t = torch.repeat_interleave(bx1t, cur_beam_width, 0)  # increase batch for new beams\n",
        "      if self.verbose:\n",
        "        print(f'{bx1t.size()}, {bx2.size()}, {bx1t.numel() + bx2.numel()}')\n",
        "      # print(bx2)\n",
        "      # print(beam_scores)\n",
        "      if self.depth_reduction == torch.sum:\n",
        "        pad_mask = beam_scores == float('-inf')\n",
        "        # print(pad_mask)\n",
        "        beam_scores[pad_mask] = 0\n",
        "        # print(beam_scores)\n",
        "        beam_scores = self.depth_reduction(beam_scores, axis=1) # cumulative beams scores\n",
        "      elif self.depth_reduction == torch.mean:\n",
        "        pad_mask = beam_scores == float('-inf')\n",
        "        # print(pad_mask)\n",
        "        beam_scores[pad_mask] = 0\n",
        "        # print(beam_scores)\n",
        "        beam_scores = self.depth_reduction(beam_scores, axis=1) # cumulative beams scores\n",
        "        # print(beam_scores)\n",
        "        beam_lens = torch.sum(~pad_mask, axis=1)\n",
        "        # print(beam_lens)\n",
        "        beam_scores = beam_scores / beam_lens\n",
        "      else:\n",
        "        self.depth_reduction(beam_scores, axis=1)\n",
        "      if self.plot_dist:\n",
        "        plt.figure()\n",
        "        sns.distplot(torch.topk(beam_scores, top_beams)[0], hist=True, kde=False)\n",
        "        plt.title(f'step {bx2.size(1)} top {top_beams} beams scores')\n",
        "      # print(beam_scores)\n",
        "      if self.batch_reduction is None:\n",
        "        beam_scores = beam_scores.view(batch_size, -1)  # split scores into batches\n",
        "        best_beams = torch.argmax(beam_scores, axis=1, keepdim=True)  # best beams\n",
        "        bx2 = bx2.view(batch_size, int(bx2.size(0)/batch_size), -1)  # split beams into batches\n",
        "        # XXX its fucking magic... (return to input batch_size)\n",
        "        best_beams = best_beams.unsqueeze(2).expand(best_beams.size(0), \n",
        "                                                    best_beams.size(1),\n",
        "                                                    bx2.size(2))\n",
        "        bx2 = torch.gather(bx2, 1, best_beams)\n",
        "        bx2 = bx2.view(batch_size, -1)\n",
        "      else:\n",
        "        best_beams = torch.argmax(beam_scores, axis=0, keepdim=True)  # best beams\n",
        "        bx2 = bx2[best_beams]  # best beam of all batches\n",
        "        bx2 = bx2.repeat(batch_size, 1)  # return to input batch size\n",
        "    return bx2[:,1:], attn\n",
        "\n",
        "class BeamBkw():\n",
        "  def __init__(self, sos=2, eos=3, max_len=30, beam_width=2, beam_depth=2,\n",
        "               depth_reduction='sum', beam_reduction='none', \n",
        "               batch_reduction='none', bkw_beam_width=2, **kwargs):\n",
        "    self.sos, self.eos, self.max_len = sos, eos, max_len\n",
        "    self.beam_width, self.beam_depth = beam_width, beam_depth\n",
        "    self.bkw_beam_width = bkw_beam_width\n",
        "    self.depth_reduction = reduction_map[depth_reduction]\n",
        "    self.beam_reduction = reduction_map[beam_reduction]\n",
        "    self.batch_reduction = reduction_map[batch_reduction]\n",
        "    self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "    if self.bkw_beam_width > self.beam_width**self.beam_depth:\n",
        "      print(f'Warning! Backward beam width {self.bkw_beam_width} greater\\\n",
        "      than beam_width**beam_depth {self.beam_width**self.beam_depth}') \n",
        "      print(f'setting backword beam width to {self.beam_width**self.beam_depth}')\n",
        "      self.bkw_beam_width = self.beam_width**self.beam_depth\n",
        "\n",
        "  def __call__(self, model, bx1, bx2=None):\n",
        "    batch_size = bx1.size(0)  # input batch_size\n",
        "    bx2 = torch.full((batch_size, 1), self.sos, dtype=bx1.dtype).to(self.device)  # batch with <sos>\n",
        "    # stop when len decoder output > max_len or all decoder outputs have <eos> token\n",
        "    while bx2.size(1) - 1 < self.max_len and not torch.all(torch.any(bx2 == self.eos, axis=1)):\n",
        "      fwd_scores = torch.empty((batch_size, 0)).to(self.device)  # scores for each beam\n",
        "      fwd_bx2 = bx2.clone()\n",
        "      bx1t = bx1.clone()\n",
        "      for i in range(self.beam_depth):\n",
        "        new_batch_size = batch_size*self.beam_width**(i+1)\n",
        "        by2p, attn = model(bx1t, fwd_bx2)  # predict\n",
        "        if i > 0 and self.beam_reduction is not None:  # beam reduction\n",
        "          prev_prev_batch_size = batch_size*self.beam_width**(i-1)\n",
        "          by2p = by2p.view(prev_prev_batch_size, self.beam_width, by2p.size(1), -1)  # split predictions into batches\n",
        "          by2p = self.beam_reduction(by2p, axis=1) # cumulative beams predictions\n",
        "          by2p = by2p.repeat(self.beam_width, 1, 1)  # return to prev batch size\n",
        "        if self.batch_reduction is not None:\n",
        "          prev_batch_size = batch_size*self.beam_width**(i)\n",
        "          by2p = self.batch_reduction(by2p, axis=0) # cumulative batch predictions\n",
        "          next_by2p, next_bx2 = torch.topk(by2p[-1,:], self.beam_width)  # beams to top k last predictions\n",
        "          next_bx2 = next_bx2.repeat(prev_batch_size, 1)  # return to prev batch size\n",
        "          next_by2p = next_by2p.repeat(prev_batch_size, 1)  # return to prev batch size        \n",
        "        else:\n",
        "          next_by2p, next_bx2 = torch.topk(by2p[:,-1,:], self.beam_width)  # beams to top k last predictions        \n",
        "        next_bx2 = next_bx2.view(new_batch_size, 1)  # new beams\n",
        "        next_by2p = next_by2p.view(new_batch_size, 1)  # new scores\n",
        "        fwd_scores = torch.repeat_interleave(fwd_scores, self.beam_width, 0)  # increase batch for new scores\n",
        "        fwd_bx2 = torch.repeat_interleave(fwd_bx2, self.beam_width, 0)  # increase batch for new beams\n",
        "        fwd_bx2 = torch.cat((fwd_bx2, next_bx2), 1)  # add beams\n",
        "        fwd_scores = torch.cat((fwd_scores, next_by2p), 1)  # add beams scores\n",
        "        bx1t = torch.repeat_interleave(bx1t, self.beam_width, 0)  # increase batch for new beams\n",
        "      if self.batch_reduction is None:\n",
        "        fwd_scores = fwd_scores.view(batch_size, self.beam_width**self.beam_depth, -1)  # split scores into batches\n",
        "        fwd_bx2 = fwd_bx2.view(batch_size, self.beam_width**self.beam_depth, -1)  # split beams into batches\n",
        "        bkw_bx2 = torch.empty((batch_size, self.beam_width**self.beam_depth, 0), dtype=bx1.dtype)  # scores for each beam\n",
        "        bkw_scores = torch.empty((batch_size, self.beam_width**self.beam_depth, 0))  # scores for each beam\n",
        "        mask = torch.full((fwd_bx2.size(0), fwd_bx2.size(1), 1), False, dtype=torch.bool)\n",
        "        n_beams_per_batch = self.beam_width**self.beam_depth\n",
        "        for i in range(fwd_scores.size(2)):\n",
        "          old_n_beams_per_batch = n_beams_per_batch*self.bkw_beam_width**(i)\n",
        "          new_n_beams_per_batch = n_beams_per_batch*self.bkw_beam_width**(i+1)\n",
        "          cur_values = fwd_bx2[:,:,-i-1:-i if i > 0 else None]\n",
        "          cur_values = torch.repeat_interleave(cur_values, self.bkw_beam_width**i, 1)\n",
        "          cur_scores = fwd_scores[:,:,-i-1:-i if i > 0 else None]\n",
        "          cur_scores = torch.repeat_interleave(cur_scores, self.bkw_beam_width**i, 1)\n",
        "          cur_scores[mask] = float('-inf')  # mask from prev step\n",
        "          best_scores, best_beams = torch.topk(cur_scores, self.bkw_beam_width, 1)  # best beams\n",
        "          best_values = torch.gather(cur_values, 1, best_beams).view(batch_size, -1)\n",
        "          best_values = best_values.view(best_values.size(0), self.bkw_beam_width, -1)\n",
        "          best_values = torch.repeat_interleave(best_values, old_n_beams_per_batch, 1)\n",
        "          cur_values = cur_values.repeat(1, self.bkw_beam_width, 1)\n",
        "          mask = cur_values != best_values\n",
        "          bkw_bx2 = bkw_bx2.repeat(1, self.bkw_beam_width, 1)  # increase batch for new beams\n",
        "          bkw_bx2 = torch.cat((best_values, bkw_bx2), 2)\n",
        "          best_scores = torch.repeat_interleave(best_scores, old_n_beams_per_batch, 1)\n",
        "          bkw_scores = bkw_scores.repeat(1, self.bkw_beam_width, 1)  # increase batch for new beams\n",
        "          bkw_scores = torch.cat((best_scores, bkw_scores), 2)\n",
        "        beam_scores = self.depth_reduction(bkw_scores, axis=2) # cumulative beams scores\n",
        "        best_beams = torch.argmax(beam_scores, axis=1, keepdim=True)  # best beams\n",
        "        best_beams = best_beams.unsqueeze(2).expand(best_beams.size(0), best_beams.size(1), bkw_bx2.size(2))\n",
        "        bkw_bx2 = torch.gather(bkw_bx2, 1, best_beams) # best beam of all batches\n",
        "        bkw_bx2 = bkw_bx2.view(batch_size, -1)\n",
        "        bx2 = torch.cat((bx2, bkw_bx2), 1)  # <sos> + bkw_bx2\n",
        "      else:  # batch reduction\n",
        "        bkw_bx2 = torch.empty((new_batch_size, 0), dtype=bx1.dtype)  # scores for each beam\n",
        "        bkw_scores = torch.empty((new_batch_size, 0))  # scores for each beam\n",
        "        mask = torch.full((new_batch_size, 1), False, dtype=torch.bool)\n",
        "        fwd_batch_size = new_batch_size\n",
        "        for i in range(fwd_scores.size(1)):\n",
        "          prev_batch_size = fwd_batch_size*self.bkw_beam_width**(i)\n",
        "          new_batch_size = fwd_batch_size*self.bkw_beam_width**(i+1)\n",
        "          cur_values = fwd_bx2[:,-i-1:-i if i > 0 else None]\n",
        "          cur_values = torch.repeat_interleave(cur_values, self.bkw_beam_width**i, 0)\n",
        "          cur_scores = fwd_scores[:,-i-1:-i if i > 0 else None]\n",
        "          cur_scores = torch.repeat_interleave(cur_scores, self.bkw_beam_width**i, 0)\n",
        "          cur_scores[mask] = float('-inf')  # mask from prev step\n",
        "          best_scores, best_beams = torch.topk(cur_scores, self.bkw_beam_width, 0)  # best beams\n",
        "          best_values = cur_values[best_beams]\n",
        "          mask = cur_values != best_values\n",
        "          mask = mask.view(new_batch_size, -1)\n",
        "          best_values = best_values.view(self.bkw_beam_width, -1)\n",
        "          best_values = torch.repeat_interleave(best_values, prev_batch_size, 0)\n",
        "          bkw_bx2 = bkw_bx2.repeat(self.bkw_beam_width, 1)  # increase batch for new beams\n",
        "          bkw_bx2 = torch.cat((best_values, bkw_bx2), 1)\n",
        "          best_scores = torch.repeat_interleave(best_scores, prev_batch_size, 0)\n",
        "          bkw_scores = bkw_scores.repeat(self.bkw_beam_width, 1)  # increase batch for new beams\n",
        "          bkw_scores = torch.cat((best_scores, bkw_scores), 1)\n",
        "        # bkw batch reduction\n",
        "        bkw_scores = self.depth_reduction(bkw_scores, axis=1) # cumulative beams scores\n",
        "        best_beams = torch.argmax(bkw_scores, axis=0, keepdim=True)  # best beams\n",
        "        bkw_bx2 = bkw_bx2[best_beams]  # best beam of all batches\n",
        "        bkw_bx2 = bkw_bx2.repeat(batch_size, 1)  # return to input batch size\n",
        "        bx2 = torch.cat((bx2, bkw_bx2), 1)  # <sos> + bkw_bx2\n",
        "    return bx2[:,1:], attn"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I3dpA2Fe3BME",
        "colab_type": "text"
      },
      "source": [
        "## Metric"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tCyCP83RUbXi",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from rouge_score import rouge_scorer\n",
        "from pprint import pprint\n",
        "\n",
        "def metrics(bx2, by2, eos=3):\n",
        "  scorer = rouge_scorer.RougeScorer(['rouge1', 'rouge2', 'rougeL'])\n",
        "  scores = {}\n",
        "  ps = by2.detach().cpu().numpy()\n",
        "  ts = bx2.detach().cpu().numpy()\n",
        "  for sp, st in zip(ps, ts):\n",
        "    # target to set\n",
        "    st_eos_i = np.where(st==eos)[0][0] if eos in st else len(st)\n",
        "    st_eos = st[:st_eos_i]\n",
        "    t_set = set(st_eos) - {0, 1, 2, 3}\n",
        "    # prediction to set\n",
        "    sp_eos_i = np.where(sp==eos)[0][0] if eos in sp else len(sp)\n",
        "    sp_eos = sp[:sp_eos_i]\n",
        "    p_set = set(sp_eos) - {0, 1, 2, 3}\n",
        "    # metrics\n",
        "    i_set = t_set.intersection(p_set)\n",
        "    pre = len(i_set)/len(p_set) if len(p_set) != 0 else 0\n",
        "    rec = len(i_set)/len(t_set) if len(t_set) != 0 else 0\n",
        "    f1 = 2*pre*rec/(pre + rec) if pre + rec != 0 else 0\n",
        "    rl = len(sp_eos)/len(st_eos) if len(st_eos) != 0 else 0\n",
        "    # to string\n",
        "    t_str = ' '.join(map(str, st_eos))\n",
        "    p_str = ' '.join(map(str, sp_eos))\n",
        "    rs = scorer.score(t_str, p_str)\n",
        "    # print(f'st: {st}, sp: {sp}')\n",
        "    # print(f'eos in st?: {eos in st}, sp_eos_i: {st_eos_i}')\n",
        "    # print(f'eos in sp?: {eos in sp}, sp_eos_i: {sp_eos_i}')\n",
        "    # print(f'st_eos: {st_eos}, sp_eos: {sp_eos}')\n",
        "    # print(f't_set: {t_set}, p_set: {p_set}')\n",
        "    # print(f't_str: \"{t_str}\", p_str: \"{p_str}\"')\n",
        "    # pprint(rs)\n",
        "    # print(f'precision: {pre}, recall: {rec}, f1: {f1}. rl: {rl}\\n')\n",
        "    scores.setdefault('set_precision', []).append(pre)\n",
        "    scores.setdefault('set_recall', []).append(rec)\n",
        "    scores.setdefault('set_f1', []).append(f1)\n",
        "    scores.setdefault('relen', []).append(rl)\n",
        "    for k, v in rs.items():\n",
        "      scores.setdefault(f'{k}_precision', []).append(v.precision)\n",
        "      scores.setdefault(f'{k}_recall', []).append(v.recall)\n",
        "      scores.setdefault(f'{k}_f1', []).append(v.fmeasure)\n",
        "  return scores\n",
        "\n",
        "x2 = torch.tensor([[1, 5, 3, 4, 4, 3, 3, 3, 3], \n",
        "                   [5, 3, 0, 4, 3, 0, 0, 0, 0]])\n",
        "y2 = torch.tensor([[3, 1, 4, 3, 4, 6, 7, 4], \n",
        "                   [0, 5, 4, 3, 3, 0, 0, 3]])\n",
        "metrics(x2, y2)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zhWq-kdb3Ep7",
        "colab_type": "text"
      },
      "source": [
        "## Run"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rZe5hRHIWcZD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import humanize\n",
        "import psutil\n",
        "import GPUtil\n",
        "\n",
        "def run(train, device, model, opt, loss_fn, dl, ds, print_triples=1, \n",
        "        plot_attn=False, inference=None, metric=False, pin_memory=False,\n",
        "        max_samples=None):\n",
        "  scores = {}\n",
        "  pbar = tqdm(dl, total=max_samples)\n",
        "  batch_cnt = 0\n",
        "  if train:\n",
        "    model.train()\n",
        "    for x1, x2 in pbar:\n",
        "      opt.zero_grad()\n",
        "      x1 = x1.to(device, non_blocking=pin_memory)\n",
        "      x2 = x2.to(device, non_blocking=pin_memory)\n",
        "      t = x2[:,1:].flatten(0)\n",
        "      y2p, attn = model(x1, x2)\n",
        "      p = y2p[:,:-1,:].flatten(0, 1)\n",
        "      loss = loss_fn(p, t)\n",
        "      loss.backward()\n",
        "      opt.step()\n",
        "      scores.setdefault('loss', []).extend(loss.item())\n",
        "      y2 = None\n",
        "      if metric:\n",
        "        y2 = torch.argmax(y2p[:,:-1,:], 2)\n",
        "        # idx = torch.nonzero(t).view(-1)\n",
        "        # acc = torch.sum(y2.view(-1)[idx] == t[idx]).float()/idx.size()[0]\n",
        "        # accs.append(acc.item())\n",
        "        bs = metrics(x2[:,1:], y2)  # batch scores\n",
        "        r1f = sum(bs['rouge1_f1'])/len(bs['rouge1_f1'])\n",
        "        r2f = sum(bs['rouge2_f1'])/len(bs['rouge2_f1'])\n",
        "        rlf = sum(bs['rougeL_f1'])/len(bs['rougeL_f1'])\n",
        "        sf = sum(bs['set_f1'])/len(bs['set_f1'])\n",
        "        rl = sum(bs['relen'])/len(bs['relen'])\n",
        "        for score, values in bs.items():\n",
        "          scores.setdefault(score, []).extend(values)\n",
        "        cmt = psutil.virtual_memory().total\n",
        "        cma = psutil.virtual_memory().available\n",
        "        cmu = (cmt-cma)/cmt\n",
        "        gmu = GPUtil.getGPUs()[0].memoryUtil if len(GPUtil.getGPUs()) > 0 else 0\n",
        "        pbar.set_description(f'train {int(cmu*100)}%/{int(gmu*100)}% \\\n",
        "          l: {loss.item():.3}, \\\n",
        "          sf: {int(sf*100)}, r1f: {int(r1f*100)}, r2f: {int(r2f*100)}, \\\n",
        "          rlf: {int(rlf*100)}, rl: {int(rl*100)}')\n",
        "      else:\n",
        "        pbar.set_description(f'loss: {loss.item():.3}')\n",
        "      batch_cnt += 1\n",
        "      if max_samples is not None:\n",
        "        if batch_cnt == max_samples:\n",
        "          break\n",
        "  else:\n",
        "    model.eval()\n",
        "    with torch.no_grad():\n",
        "      for x1, x2 in pbar:\n",
        "        x1 = x1.to(device, non_blocking=pin_memory)\n",
        "        x2 = x2.to(device, non_blocking=pin_memory)\n",
        "        if inference is not None:\n",
        "          y2, attn = inference(model, x1, x2)\n",
        "        else:  # forced with loss (e.g. for training)\n",
        "          t = x2[:,1:].flatten(0)\n",
        "          y2p, attn = model(x1, x2)\n",
        "          p = y2p[:,:-1,:].flatten(0, 1)\n",
        "          loss = loss_fn(p, t)\n",
        "          y2 = torch.argmax(y2p[:,:-1,:], 2)\n",
        "          scores.setdefault('loss', []).extend(loss.item())\n",
        "          pbar.set_description(f'loss: {loss.item():.3}')\n",
        "        if metric:\n",
        "          bs = metrics(x2[:,1:], y2)  # batch scores\n",
        "          for score, values in bs.items():\n",
        "            scores.setdefault(score, []).extend(values)\n",
        "          # r1f = sum(bs['rouge1_f1'])/len(bs['rouge1_f1'])\n",
        "          # r2f = sum(bs['rouge2_f1'])/len(bs['rouge2_f1'])\n",
        "          # rlf = sum(bs['rougeL_f1'])/len(bs['rougeL_f1'])\n",
        "          # sf = sum(bs['set_f1'])/len(bs['set_f1'])\n",
        "          # rl = sum(bs['relen'])/len(bs['relen'])\n",
        "          # cmt = psutil.virtual_memory().total\n",
        "          # cma = psutil.virtual_memory().available\n",
        "          # cmu = (cmt-cma)/cmt\n",
        "          # gmu = GPUtil.getGPUs()[0].memoryUtil if len(GPUtil.getGPUs()) > 0 else 0\n",
        "          # pbar.set_description(f'eval {int(cmu*100)}%/{int(gmu*100)}% \\\n",
        "          #   sf: {int(sf*100)}, r1f: {int(r1f*100)}, r2f: {int(r2f*100)}, \\\n",
        "          #   rlf: {int(rlf*100)}, rl: {int(rl*100)}')\n",
        "          # pbar.set_description(f'gmu: {int(gmu*100)}%')\n",
        "        batch_cnt += 1\n",
        "        if max_samples is not None:\n",
        "          if batch_cnt == max_samples:\n",
        "            break\n",
        "  if print_triples:\n",
        "    if y2 is None:\n",
        "      y2 = torch.argmax(y2p[:,:-1,:], 2)\n",
        "    dy2 = map(ds.decode, y2.detach().cpu().numpy())\n",
        "    dx2 = map(ds.decode, x2[:,1:].detach().cpu().numpy())\n",
        "    dx1 = map(ds.decode, x1[:,1:].detach().cpu().numpy())\n",
        "    triples = list(zip(dx1, dx2, dy2))[:print_triples]\n",
        "    for triple in triples:\n",
        "      print('')\n",
        "      print(triple[0])\n",
        "      print(triple[1])\n",
        "      print(triple[2])\n",
        "  if plot_attn:\n",
        "    plot_attention(attn.detach().cpu().numpy()[:1], \n",
        "                   x1.detach().cpu().numpy()[:1], \n",
        "                   x2.detach().cpu().numpy()[:1],\n",
        "                   decoder_x1 = ds.decode,\n",
        "                   decoder_x2 = ds.decode,\n",
        "                   figsize=(10, 10),\n",
        "                   shift=True,\n",
        "                   mask=True,\n",
        "                   suptitle=None)\n",
        "  return scores"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hziAG-KBqOO8",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def tensors_mem():\n",
        "  for obj in gc.get_objects():\n",
        "    try:\n",
        "      if isinstance(obj, torch.Tensor):\n",
        "      # if hasattr(obj, 'data'):\n",
        "      #   if torch.is_tensor(obj.data):\n",
        "        print(type(obj), obj.size(), obj.dtype, obj.device, obj.layout, obj.memory_format)\n",
        "    except:\n",
        "      pass\n",
        "tensors_mem()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ohVsoShwhWws",
        "colab_type": "text"
      },
      "source": [
        "## Main"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-OElaFmqpOwK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "import optuna\n",
        "from tqdm.notebook import tqdm\n",
        "from pprint import pprint\n",
        "import random\n",
        "import numpy as np\n",
        "from shutil import copyfile\n",
        "import GPUtil\n",
        "from torch import optim\n",
        "from torchsummary import summary\n",
        "\n",
        "def objective(trial):\n",
        "  device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "  trial.set_user_attr('device', str(device))\n",
        "  if device.type == \"cuda\":\n",
        "    print(f'gmu: {\"% \".join(str(int(x.memoryUtil*100)) for x in GPUtil.getGPUs())}%')\n",
        "    torch.cuda.empty_cache()\n",
        "    torch.cuda.ipc_collect()\n",
        "    print(torch.cuda.memory_stats(device=device))\n",
        "    print(torch.cuda.memory_summary(device=device))\n",
        "    print(f'gmu: {\"% \".join(str(int(x.memoryUtil*100)) for x in GPUtil.getGPUs())}%')\n",
        "  opt_fn_map = {\n",
        "    'SGD': optim.SGD,\n",
        "    'Adam': optim.Adam,\n",
        "    'Adagrad': optim.Adagrad,\n",
        "    'ASGD': optim.ASGD,\n",
        "    'Adamax': optim.Adamax,\n",
        "    'SparseAdam': optim.SparseAdam,\n",
        "    'AdamW': optim.AdamW,\n",
        "    'Adadelta': optim.Adadelta,\n",
        "    'LBFGS': optim.LBFGS,\n",
        "    'RMSprop': optim.RMSprop,\n",
        "    'Rprop': optim.Rprop\n",
        "  }\n",
        "  eval_fn_map = {\n",
        "      'Beam': Beam,\n",
        "      'BeamBkw': BeamBkw,\n",
        "      'Greedy': Greedy,\n",
        "      'Forced': Forced\n",
        "  }\n",
        "  # fn = 'model.pth'\n",
        "  fn = f'model_{trial.study.study_name}_{trial.number}.pth'\n",
        "  print(fn)\n",
        "  # checkpoint = None\n",
        "  checkpoint = torch.load('/content/drive/My Drive/model_ria_base.pth',\n",
        "                          map_location=device)\n",
        "  train = False  # ?\n",
        "  # train_config = checkpoint['train_config']\n",
        "  train_config = {\n",
        "    'seed': trial.suggest_int('seed', 0, 0),\n",
        "    'n_epoches': trial.suggest_int('n_epoches', 1, 1),\n",
        "    'lr': trial.suggest_loguniform('lr', 1e-3, 1e-3),\n",
        "      # 'opt_fn': trial.suggest_categorical('opt_fn', ['SGD', 'Adam', 'Adagrad', \n",
        "      #                                           'ASGD', 'Adamax',\n",
        "      #                                           'AdamW', 'Adadelta',\n",
        "      #                                           'RMSprop', 'Rprop']),\n",
        "    'opt_fn': trial.suggest_categorical('opt_fn', ['Adam']),\n",
        "    'batch_size': trial.suggest_int('batch_size', 1, 1),\n",
        "    'weight_decay': None,\n",
        "    'pin_memory': True,\n",
        "    # 'weight_decay': trial.suggest_loguniform('weight_decay', 1e-6, 1e-6),\n",
        "    'n_x1': trial.suggest_int('n_x1', 1, 1)\n",
        "  }\n",
        "  eval_config = {\n",
        "      # 'eval_fn': trial.suggest_categorical('eval_fn', ['Beam', 'BeamBkw']),\n",
        "      # 'eval_fn': trial.suggest_categorical('eval_fn', ['Beam']),\n",
        "      'eval_fn': 'Beam',\n",
        "      'max_samples': 100}\n",
        "  if eval_config['eval_fn'] in ['Beam', 'BeamBkw']:\n",
        "    eval_config['beam_width'] = trial.suggest_int('beam_width', 2, 3)\n",
        "    eval_config['beam_depth'] = trial.suggest_int('beam_depth', 8, 8)\n",
        "    # eval_config['beam_reduction'] = trial.suggest_categorical('beam_reduction', ['sum', 'mean', 'none']),\n",
        "    eval_config['beam_reduction'] = 'none'\n",
        "    eval_config['max_input_size'] = 20000  # BATCH*(LEN1+LEN2) ~ BATCH*(10+30)\n",
        "    eval_config['batch_reduction'] = 'none'\n",
        "    eval_config['max_len'] = 14\n",
        "    eval_config['plot_size'] = 10\n",
        "    eval_config['plot_dist'] = False\n",
        "    eval_config['verbose'] = False\n",
        "    # eval_config['first_beam_width'] = eval_config['beam_width']\n",
        "    eval_config['first_beam_width'] = trial.suggest_int('first_beam_width', 10, 10)\n",
        "    # eval_config['depth_reduction'] = trial.suggest_categorical('depth_reduction', ['sum', 'mean'])\n",
        "    eval_config['depth_reduction'] = 'mean'\n",
        "    if eval_config['eval_fn'] == 'BeamBkw':\n",
        "      eval_config['bkw_beam_width'] = trial.suggest_int('bkw_beam_width', 2, 2)\n",
        "  if checkpoint is None:\n",
        "    model_config = {\n",
        "      'enc_num_embeddings': len(ds.stoi),\n",
        "      'enc_embedding_dim': trial.suggest_int('enc_embedding_dim', 300, 300),\n",
        "      'enc_padding_idx': ds.stoi['<pad>'],\n",
        "      'dec_num_embeddings': len(ds.stoi),\n",
        "      'dec_embedding_dim': trial.suggest_int('dec_embedding_dim', 300, 300),\n",
        "      'dec_padding_idx': ds.stoi['<pad>'],\n",
        "      'rnn_type': trial.suggest_categorical('rnn_type', ['RNN']),\n",
        "      # 'rnn_type': trial.suggest_categorical('rnn_type', ['RNN']),\n",
        "      'hidden_size': trial.suggest_int('hidden_size', 300, 300),\n",
        "      'num_layers': trial.suggest_int('num_layers', 1, 1),\n",
        "      'bidirectional': trial.suggest_categorical('bidirectional', [False]),\n",
        "      # 'bidirectional': trial.suggest_categorical('bidirectional', [True]),\n",
        "      'enc_dropout': trial.suggest_uniform('enc_dropout', 0.0, 0.0),\n",
        "      'dec_dropout': trial.suggest_uniform('dec_dropout', 0.0, 0.0),\n",
        "      'enc_rnn_dropout': trial.suggest_uniform('enc_rnn_dropout', 0.0, 0.0),\n",
        "      'dec_rnn_dropout': trial.suggest_uniform('dec_rnn_dropout', 0.0, 0.0),\n",
        "      'attn_type': trial.suggest_categorical('attn_type', ['soft_dist']),\n",
        "      # 'attn_type': trial.suggest_categorical('attn_type', ['dot', 'dist', 'soft_dot', 'soft_dist', 'cos', 'soft_cos', 'none']),\n",
        "      'out_hidden': trial.suggest_int('out_hidden', 600, 600),\n",
        "      'pack': trial.suggest_categorical('pack', [True]),\n",
        "    }\n",
        "  else:\n",
        "    model_config = checkpoint['model_config'] \n",
        "  # pprint(model_config)   \n",
        "  # pprint(train_config)\n",
        "  # pprint(eval_config)\n",
        "  for k, v in model_config.items():\n",
        "    trial.set_user_attr('model_' + k, v)\n",
        "  for k, v in train_config.items():\n",
        "    trial.set_user_attr('train_' + k, v)\n",
        "  for k, v in eval_config.items():\n",
        "    trial.set_user_attr('eval_' + k, v)\n",
        "  seed = train_config['seed']\n",
        "  torch.backends.cudnn.benchmark = False\n",
        "  torch.backends.cudnn.deterministic = True\n",
        "  torch.manual_seed(seed)\n",
        "  torch.cuda.manual_seed(seed)\n",
        "  random.seed(seed)\n",
        "  np.random.seed(seed)\n",
        "  opt_fn = train_config['opt_fn']\n",
        "  lr = train_config['lr']\n",
        "  n_epoches = train_config['n_epoches']\n",
        "  batch_size = train_config['batch_size']\n",
        "  weight_decay = train_config['weight_decay']\n",
        "  pin_memory = train_config['pin_memory']\n",
        "  model = EncoderDecoder(**model_config).to(device)\n",
        "  if checkpoint is not None:\n",
        "    model.load_state_dict(checkpoint['model_state_dict'])\n",
        "  pprint(model)\n",
        "  n_params = sum(x.numel() for x in model.parameters() if x.requires_grad)\n",
        "  trial.set_user_attr('n_params', n_params)\n",
        "  if opt_fn not in ['Rprop'] and weight_decay is not None:\n",
        "    opt = opt_fn_map[opt_fn](model.parameters(), lr=lr, \n",
        "                             weight_decay=weight_decay)\n",
        "  else:\n",
        "    opt = opt_fn_map[opt_fn](model.parameters(), lr=lr)\n",
        "  loss_fn = torch.nn.NLLLoss(ignore_index=0)\n",
        "  train_len = int(0.7*len(ds))\n",
        "  test_len = int(0.2*len(ds))\n",
        "  val_len = len(ds) - train_len - test_len\n",
        "  lens = [train_len, test_len, val_len]\n",
        "  trial.set_user_attr('n_samples', len(ds))\n",
        "  trial.set_user_attr('n_samples_train', train_len)\n",
        "  trial.set_user_attr('n_samples_test', test_len)\n",
        "  trial.set_user_attr('n_samples_val', val_len)\n",
        "  train_ds, test_ds, val_ds = random_split(ds, lens)\n",
        "  train_dl = DataLoader(train_ds, batch_size=batch_size, num_workers=1, \n",
        "                        shuffle=True, drop_last=False,\n",
        "                        collate_fn=Collate(train_config['n_x1']),\n",
        "                        pin_memory=pin_memory)\n",
        "  test_dl = DataLoader(test_ds, batch_size=batch_size, num_workers=1, \n",
        "                       shuffle=False, drop_last=False,\n",
        "                       collate_fn=Collate(train_config['n_x1']),\n",
        "                       pin_memory=pin_memory)\n",
        "  val_dl = DataLoader(val_ds, batch_size=batch_size, num_workers=1, \n",
        "                      shuffle=False, drop_last=False,\n",
        "                      collate_fn=Collate(train_config['n_x1']),\n",
        "                      pin_memory=pin_memory)\n",
        "  pprint(trial.user_attrs)\n",
        "  pprint(trial.params)\n",
        "  if train:\n",
        "    pbar_epoch = tqdm(range(n_epoches))\n",
        "    for i in tqdm(pbar_epoch):\n",
        "      train_scores = run(True, device, model, opt, loss_fn, train_dl, ds, \n",
        "                          print_triples=3, \n",
        "                          metric=False,\n",
        "                          pin_memory=pin_memory,\n",
        "                          plot_attn=True if model_config['attn_type'] != 'none' else False)\n",
        "      train_loss = sum(train_scores['loss'])/len(train_scores['loss'])\n",
        "      trial.report(train_loss, step=i+1)\n",
        "      # validation\n",
        "      # eval_fn = eval_fn_map[eval_config['eval_fn']](**eval_config)\n",
        "      val_scores = run(False, device, model, opt, loss_fn, val_dl, ds,  \n",
        "                       print_triples=3, \n",
        "                       inference=None,\n",
        "                       metric=False,\n",
        "                       pin_memory=pin_memory,\n",
        "                       plot_attn=False)\n",
        "      val_loss = sum(val_scores['loss'])/len(val_scores['loss'])\n",
        "      # print(f'{type(eval_fn).__name__} p: {val_pre:.3f}, r: {val_rec:.3f}, f1: {val_f1:.3f}, rl: {val_rl:.3f}')\n",
        "      pbar_epoch.set_description(f'epoch: {i+1}, tloss: {train_loss:.3}, vloss: {val_loss:.3}')\n",
        "      # pbar_epoch.set_description(f'l: {loss:.2f}/{val_loss:.2f} \\\n",
        "      # a: {int(acc*100)}/{int(val_acc*100)} p: {int(pre*100)}/{int(val_pre*100)} \\\n",
        "      # r: {int(rec*100)}/{int(val_rec*100)} f1: {int(f1*100)}/{int(val_f1*100)}\\\n",
        "      # rl: {int(rl*100)}/{int(val_rl*100)}')\n",
        "      # torch.save({'epoch': i,\n",
        "      #             'model_state_dict': model.state_dict(),\n",
        "      #             'optimizer_state_dict': opt.state_dict(),\n",
        "      #             'loss': loss,\n",
        "      #             'model_config': model_config,\n",
        "      #             'train_config': train_config,\n",
        "      #             'eval_config': eval_config},\n",
        "      #              fn)\n",
        "      # try:\n",
        "      #   copyfile(fn, '/content/drive/My Drive/' + fn)\n",
        "      # except Exception as e:\n",
        "      #   print(e)\n",
        "  # test\n",
        "  eval_fn = eval_fn_map[eval_config['eval_fn']](**eval_config)\n",
        "  test_scores = run(False, device, model, opt, loss_fn, test_dl, ds, \n",
        "                    print_triples=1, \n",
        "                    inference=eval_fn,\n",
        "                    metric=True,\n",
        "                    pin_memory=pin_memory,\n",
        "                    max_samples=eval_config['max_samples'],\n",
        "                    plot_attn=False)\n",
        "  scores = test_scores\n",
        "  test_scores = {}\n",
        "  for score, values in scores.items():\n",
        "    test_scores[score] = sum(values)/len(values)\n",
        "    trial.set_user_attr(score, test_scores[score])\n",
        "  pprint(test_scores)\n",
        "  return loss if train else test_scores['rouge1_f1']\n",
        "\n",
        "import os\n",
        "fn = 'optuna_20.db'\n",
        "if not os.path.exists(fn):\n",
        "  copyfile('/content/drive/My Drive/' + fn, fn)\n",
        "study = optuna.create_study(study_name='2_vs_3_100', \n",
        "                            direction='maximize', \n",
        "                            storage=f'sqlite:///{fn}', \n",
        "                            sampler=optuna.samplers.RandomSampler(0),  # optuna.samplers.GridSampler()\n",
        "                            load_if_exists=True)\n",
        "\n",
        "def dump_optuna(study, frozen_trial):\n",
        "  fn = 'optuna_20.db'\n",
        "  print(f'saving optuna {fn}')\n",
        "  try:\n",
        "    copyfile(fn, '/content/drive/My Drive/' + fn)\n",
        "  except Exception as e:\n",
        "    print(e)\n",
        "  print('done')\n",
        "\n",
        "study.optimize(objective, n_trials=2, callbacks=[dump_optuna])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gIc7J6j5CZM7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import gc\n",
        "gc.collect()\n",
        "# try:\n",
        "#   copyfile(fn, '/content/drive/My Drive/' + fn)\n",
        "# except Exception as e:\n",
        "#   print(e)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UlcXPvqiyHbQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Greedy 200k\n",
        "# {'relen': 0.9892607721833844,\n",
        "#  'rouge1_f1': 0.38423491185005065,\n",
        "#  'rouge1_precision': 0.4026889198348889,\n",
        "#  'rouge1_recall': 0.37827062666600875,\n",
        "#  'rouge2_f1': 0.16623936163781577,\n",
        "#  'rouge2_precision': 0.17458141088353102,\n",
        "#  'rouge2_recall': 0.16407297971294255,\n",
        "#  'rougeL_f1': 0.3567234511838257,\n",
        "#  'rougeL_precision': 0.3739773768993677,\n",
        "#  'rougeL_recall': 0.35121314215560523,\n",
        "#  'set_f1': 0.39064669071317015,\n",
        "#  'set_precision': 0.4199063773953252,\n",
        "#  'set_recall': 0.3771313731307635}\n",
        "\n",
        "# try:\n",
        "# except KeyboardInterrupt:\n",
        "#   print(e)\n",
        "# except Exception as e:\n",
        "#   print(e)\n",
        "# finally:\n",
        "#   torch.save({'epoch': i,\n",
        "#               'model_state_dict': model.state_dict(),\n",
        "#               'optimizer_state_dict': opt.state_dict(),\n",
        "#               'loss': loss,\n",
        "#               'model_config': model_config,\n",
        "#               'train_config': train_config},\n",
        "#               'model_exception.pth')\n",
        "#   try:\n",
        "#     copyfile('model.pth', '/content/drive/My Drive/' + 'model.pth')\n",
        "#   except Exception as e:\n",
        "#     print(e)\n",
        "# raise KeyboardInterrupt"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BJernh1FIvU2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from pprint import pprint\n",
        "pprint(study.best_params)\n",
        "pprint(study.best_value)\n",
        "pprint(study.best_trial)\n",
        "pprint(study.direction)\n",
        "%reload_ext google.colab.data_table\n",
        "df = study.trials_dataframe()\n",
        "df.params_eval_fn.fillna(df.user_attrs_eval_eval_fn, inplace=True)\n",
        "df"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZgmudFU2IykS",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# optuna.visualization.plot_contour(study, params=['beam_width', 'beam_depth'])\n",
        "# optuna.visualization.plot_optimization_history(study)\n",
        "optuna.visualization.plot_slice(study)\n",
        "# optuna.visualization.plot_parallel_coordinate(study)  # BUG\n",
        "# optuna.visualization.plot_intermediate_values(study)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_cgvmWMJHCV_",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%matplotlib inline\n",
        "import matplotlib.pyplot as plt\n",
        "import seaborn as sns\n",
        "import pandas as pd\n",
        "\n",
        "# max_vocab = 1000\n",
        "# with open(vocab_path) as f:\n",
        "#   vocab = json.load(f)\n",
        "# top_words = dict([(w, c) for w, c in sorted(vocab.items(), key=lambda x: x[1],\n",
        "#                                             reverse=True)][:max_vocab])\n",
        "# # sns.distplot(list(map(len, vocab.keys())), kde=False)\n",
        "# # plt.figure()\n",
        "# # sns.distplot(list(map(len, top_words.keys())), kde=True)\n",
        "# # plt.figure()\n",
        "# sns.distplot(list(vocab.values()), kde=False, hist_kws={'log': True})\n",
        "# sns.distplot(list(top_words.values()), kde=False, hist_kws={'log': True})\n",
        "# ax = sns.distplot(list(vocab.values()), kde=False, hist=True, \n",
        "#              hist_kws={'log': True}, norm_hist=False)\n",
        "# # ax.set_xscale('log')\n",
        "# ax = sns.distplot(list(top_words.values()), kde=False, hist=True, \n",
        "#              hist_kws={'log': True}, norm_hist=False)\n",
        "# # ax.set_xscale('log')\n",
        "# plt.figure()\n",
        "# plt.title('Sentences lengths')\n",
        "# sns.kdeplot([len(y) for x in ds.samples for y in x[0]], \n",
        "#             label='texts', shade=True, bw=2)\n",
        "# sns.kdeplot([len(y) for x in ds.samples for y in x[1]], \n",
        "#             label='titles', shade=True, bw=2)\n",
        "# plt.legend()\n",
        "# plt.figure()\n",
        "# sns.jointplot(x='titles lengths', y='texts lengths', data=pd.DataFrame({\n",
        "#     \"texts lengths\": [sum(map(len, x[0])) for x in ds.samples],\n",
        "#     \"titles lengths\": [sum(map(len, x[1])) for x in ds.samples]}\n",
        "# ), kind='kde', n_levels=10)\n",
        "# sns.jointplot(x='user_attrs_rouge1_recall', y='user_attrs_rouge1_precision',\n",
        "#               data=df, kind='kde', n_levels=10) \n",
        "# hue=\"sex\"\n",
        "# sns.catplot(x=\"params_beam_width\", y=\"params_beam_depth\", \n",
        "#             kind=\"violin\", hue=\"params_eval_fn\", split=True, bw=1., data=df)\n",
        "# sns.violinplot(x=df.params_beam_depth, y=df.user_attrs_rouge1_f1);\n",
        "\n",
        "# sns.catplot(x=\"params_beam_width\", y=\"user_attrs_rougeL_f1\", \n",
        "#             # hue=\"smoker\",\n",
        "#             col=\"params_eval_fn\", aspect=.6,\n",
        "#             kind=\"swarm\", data=df);\n",
        "# sns.pairplot(df, hue=\"value\", vars=[\"params_beam_width\", \"params_beam_depth\"], kind=\"kde\")\n",
        "\n",
        "# g = sns.PairGrid(df, vars=[\"params_beam_width\", \"params_beam_depth\",\n",
        "#                            \"user_attrs_rougeL_f1\", \"user_attrs_rouge2_f1\",\n",
        "#                            \"user_attrs_rouge1_f1\"])\n",
        "# g.map_upper(plt.scatter)\n",
        "# g.map_lower(sns.kdeplot)\n",
        "# g.map_diag(sns.kdeplot, lw=3, legend=False);\n",
        "\n",
        "# sns.violinplot(x=df.params_beam_depth, y=df.user_attrs_rouge1_f1);\n",
        "\n",
        "# sns.catplot(x=\"params_beam_width\", \n",
        "#             y=\"params_beam_depth\", \n",
        "#             hue=\"value\",\n",
        "#             col=\"params_eval_fn\", \n",
        "#             aspect=.618,\n",
        "#             kind=\"swarm\", palette='Blues', data=df);\n",
        "# sns.catplot(x=\"params_beam_width\", \n",
        "#             y=\"user_attrs_rouge2_f1\", \n",
        "#             hue=\"params_eval_fn\",\n",
        "#             col=\"params_beam_depth\", \n",
        "#             aspect=.618,\n",
        "#             kind=\"swarm\", data=df);\n",
        "\n",
        "# sns.lmplot(x=\"params_beam_width\", y=\"user_attrs_rouge2_f1\", \n",
        "#           #  hue=\"smoker\", \n",
        "#            col=\"params_beam_depth\", row=\"params_eval_fn\", data=df);\n",
        "\n",
        "# df.params_eval_fn.fillna(df.user_attrs_eval_params_eval_fn, inplace=True)\n",
        "\n",
        "sns.pairplot(df, x_vars=[\n",
        "                        #  \"params_beam_width\", \n",
        "                         \"params_beam_depth\",\n",
        "                        #  \"params_first_beam_width\",\n",
        "                         \"user_attrs_eval_input_limit\"\n",
        "                         ], \n",
        "             y_vars=[\"user_attrs_rouge1_f1\", \n",
        "                     \"user_attrs_rouge2_f1\",\n",
        "                     \"user_attrs_rougeL_f1\",\n",
        "                     \"user_attrs_set_f1\"],\n",
        "             hue=\"user_attrs_eval_max_samples\", \n",
        "             height=5, aspect=1., kind=\"reg\");\n",
        "\n",
        "# sns.jointplot(x='params_beam_width', y='params_beam_depth',\n",
        "#               data=df, kind='kde', n_levels=10) "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NbikIK--_CV3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}