main_transformer_headline_generation_simple.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "main_transformer_headline_generation_simple.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "BB8HXRc_chEG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "config = {\n",
        "    'num_epochs': 20,\n",
        "    'num_tokens': 2000,\n",
        "    # 'vocab_size': 10000+4,\n",
        "    'd_model': 512,\n",
        "    'num_encoder_layers': 3, # 6\n",
        "    'num_decoder_layers': 3, # 6\n",
        "    'dim_feedforward': 2048,\n",
        "    'nhead': 4, # 8,\n",
        "    'pos_dropout': 0.2,\n",
        "    'trans_dropout': 0.2,\n",
        "    'n_warmup_steps': 4000,\n",
        "    'batch_size': 64,\n",
        "}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WRpxX6rCdGVM",
        "colab_type": "code",
        "outputId": "d8f4fdd7-71df-4709-8ddd-9876fe129476",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "root = \"drive/My Drive/Colab Notebooks/huawei/project/generations/\""
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-KEJUKEFM8OR",
        "colab_type": "code",
        "outputId": "f34272c6-990b-46a9-b6af-578d304a69ac",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 108
        }
      },
      "source": [
        "!pip install pymorphy2\n",
        "!pip install einops"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: pymorphy2 in /usr/local/lib/python3.6/dist-packages (0.8)\n",
            "Requirement already satisfied: dawg-python>=0.7 in /usr/local/lib/python3.6/dist-packages (from pymorphy2) (0.7.2)\n",
            "Requirement already satisfied: pymorphy2-dicts<3.0,>=2.4 in /usr/local/lib/python3.6/dist-packages (from pymorphy2) (2.4.393442.3710985)\n",
            "Requirement already satisfied: docopt>=0.6 in /usr/local/lib/python3.6/dist-packages (from pymorphy2) (0.6.2)\n",
            "Requirement already satisfied: einops in /usr/local/lib/python3.6/dist-packages (0.2.0)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "okcgOCw7M-tL",
        "colab_type": "code",
        "outputId": "34f62c66-5bbe-40f6-bc59-8dc9179e77ab",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 164
        }
      },
      "source": [
        "%matplotlib inline\n",
        "import seaborn as sns\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "\n",
        "from __future__ import unicode_literals, print_function, division\n",
        "from io import open\n",
        "import unicodedata\n",
        "import string\n",
        "import re\n",
        "import random\n",
        "from collections import defaultdict\n",
        "import numpy as np\n",
        "from einops import rearrange\n",
        "\n",
        "import nltk\n",
        "from nltk.corpus import stopwords\n",
        "nltk.download('stopwords')\n",
        "nltk.download('punkt')\n",
        "\n",
        "from sklearn.datasets import fetch_20newsgroups\n",
        "\n",
        "from nltk import word_tokenize\n",
        "from nltk.tokenize import sent_tokenize, word_tokenize\n",
        "\n",
        "import pymorphy2\n",
        "import re\n",
        "\n",
        "\n",
        "import os\n",
        "import time\n",
        "import json\n",
        "import tqdm\n",
        "from bs4 import BeautifulSoup\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.optim import Adam\n",
        "import torch.nn.functional as F\n",
        "\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "print(f\"Using {device}\")\n",
        "\n",
        "\n",
        "def save_texts(texts, titles, filename):\n",
        "  dump_texts = os.path.join(root, filename)\n",
        "  with open(file=dump_texts, mode='w') as f:\n",
        "    json.dump({'texts': texts, 'titles': titles}, f)\n",
        "\n",
        "def load_texts(filename):\n",
        "  dump_texts = os.path.join(root, filename)\n",
        "  with open(file=dump_texts, mode='r') as f:\n",
        "    data = json.load(f)\n",
        "  return data['texts'], data['titles']"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.6/dist-packages/statsmodels/tools/_testing.py:19: FutureWarning: pandas.util.testing is deprecated. Use the functions in the public API at pandas.testing instead.\n",
            "  import pandas.util.testing as tm\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "[nltk_data] Downloading package stopwords to /root/nltk_data...\n",
            "[nltk_data]   Package stopwords is already up-to-date!\n",
            "[nltk_data] Downloading package punkt to /root/nltk_data...\n",
            "[nltk_data]   Package punkt is already up-to-date!\n",
            "Using cuda\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3qnoS4q_NCbT",
        "colab_type": "text"
      },
      "source": [
        "# Preprocess ria news"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_WgfjFFtNDmB",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# source_file = \"processed-ria.json\"\n",
        "# titles = []\n",
        "# texts = []\n",
        "\n",
        "# bad_text = []\n",
        "\n",
        "# with open(file=os.path.join(root, source_file), mode='r') as f:\n",
        "#   # lines = f.readlines()[:100]\n",
        "#   for _ in tqdm.tqdm(range(100000)):\n",
        "#     line = f.readline()\n",
        "#     try:\n",
        "#       line_j = json.loads(line)\n",
        "#       text = nltk.sent_tokenize(BeautifulSoup(line_j['text'], \"lxml\").text)[1]\n",
        "#       titles.append(line_j['title'])\n",
        "#       texts.append(text)\n",
        "#     except IndexError:\n",
        "#       bad_text.append(line_j)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CrVZU4kvNF9x",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# morph = pymorphy2.MorphAnalyzer()\n",
        "# pattern = re.compile(\"^[а-яА-Я]+$\")\n",
        "\n",
        "# texts_processed = []\n",
        "# titles_processed = []\n",
        "\n",
        "# def process_sentence(sentence):\n",
        "#     words = [w for w in nltk.word_tokenize(sentence) if w not in stopwords.words(\"russian\") and pattern.match(w)]\n",
        "#     words = [morph.parse(word)[0].normal_form for word in words]\n",
        "#     return words\n",
        "\n",
        "# for text_i in tqdm.tqdm(range(len(texts))):\n",
        "#   try:\n",
        "#     # sentences = nltk.sent_tokenize(texts[text_i])\n",
        "#     # texts_processed.append(process_sentence(sentences[0]))\n",
        "#     texts_processed.append(process_sentence(texts[text_i]))\n",
        "#     titles_processed.append(process_sentence(titles[text_i]))\n",
        "#   except IndexError as e:\n",
        "#     print(e)\n",
        "#     print(text_i)\n",
        "\n",
        "# texts = texts_processed\n",
        "# titles = titles_processed\n",
        "# save_texts(texts, titles)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pmpGQn5VNL_p",
        "colab_type": "text"
      },
      "source": [
        "# Load data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oMK1-OspDRL8",
        "colab_type": "code",
        "outputId": "27d2c4f7-2adf-4ac0-ef70-34005ce1a45f",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "texts, titles = load_texts(\"ria_text_titles.json\") # \n",
        "print(len(texts))"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "98148\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B-uC9447NOjo",
        "colab_type": "code",
        "outputId": "8754b787-aba6-4534-c9fb-0dcc1f7f19b7",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 430
        }
      },
      "source": [
        "texts, titles = load_texts(\"ria_text_titles.json\") # 98148\n",
        "# texts, titles = texts[:30000], titles[:30000]\n",
        "print(max([len(t) for t in texts]))\n",
        "\n",
        "\n",
        "\n",
        "SOS_token = 0 # start of string.\n",
        "EOS_token = 1 # end of string.\n",
        "PAD_token = 2\n",
        "UNK_token = 3\n",
        "\n",
        "class Lang: # like a vocabulary\n",
        "    def __init__(self, name):\n",
        "        self.name = name\n",
        "        self.word2index = {\"SOS\": 0, \"EOS\": 1, \"PAD\": PAD_token, 'UNK': UNK_token}\n",
        "        self.word2count = {}\n",
        "        self.index2word = {0: \"SOS\", 1: \"EOS\", PAD_token: \"PAD\", UNK_token: \"UNK\"}\n",
        "        self.n_words = len(self.word2index)\n",
        "        self.max_seq_length = None\n",
        "\n",
        "    def addSentence(self, sentence):\n",
        "        if self.max_seq_length is None:\n",
        "          self.max_seq_length = len(sentence)\n",
        "        elif self.max_seq_length < len(sentence):\n",
        "          self.max_seq_length = len(sentence)\n",
        "        for word in sentence:\n",
        "            self.addWord(word)\n",
        "\n",
        "    def addWord(self, word):\n",
        "        if word not in self.word2index:\n",
        "            self.word2index[word] = self.n_words\n",
        "            self.word2count[word] = 1\n",
        "            self.index2word[self.n_words] = word\n",
        "            self.n_words += 1\n",
        "        else:\n",
        "            self.word2count[word] += 1\n",
        "\n",
        "    def tensor2text(self, tensor):\n",
        "        return [self.index2word[token.item()] for token in tensor]\n",
        "        \n",
        "    def ids2text(self, numbers):\n",
        "      return [self.index2word[num] for num in numbers]\n",
        "\n",
        "    def text2ids(self, tokens):\n",
        "        return [self.word2index[token] if token in self.word2index else self.word2index['UNK'] for token in tokens]\n",
        "            \n",
        "def unicodeToAscii(s):\n",
        "    return ''.join(\n",
        "        c for c in unicodedata.normalize('NFD', s)\n",
        "        if unicodedata.category(c) != 'Mn'\n",
        "    )\n",
        "\n",
        "\n",
        "def normalizeString(s):\n",
        "    s = unicodeToAscii(s.lower().strip())\n",
        "    s = re.sub(r\"([.!?])\", r\" \\1\", s)\n",
        "    s = re.sub(r\"[^a-zA-Z.!?]+\", r\" \", s)\n",
        "    return s\n",
        "\n",
        "\n",
        "\n",
        "def prepareData(texts, titles):\n",
        "    pairs = [[txt, ttl] for txt, ttl in zip(texts, titles) if len(txt)>0 and len(ttl)>0 ]\n",
        "    lang = Lang('ru')\n",
        "\n",
        "    print(\"Read %s sentence pairs\" % len(pairs))\n",
        "    for pair in pairs:\n",
        "        lang.addSentence(pair[0])\n",
        "        lang.addSentence(pair[1])\n",
        "    print(\"Counted words:\")\n",
        "    print(lang.name, lang.n_words)\n",
        "    return lang, pairs\n",
        "\n",
        "def indexesFromSentence(lang, sentence):\n",
        "    return [lang.word2index[word] for word in sentence]\n",
        "\n",
        "\n",
        "def tensorFromSentence(lang, sentence):\n",
        "    indexes = indexesFromSentence(lang, sentence)\n",
        "    indexes.append(EOS_token)\n",
        "    return torch.tensor(indexes, dtype=torch.long, device=device)\n",
        "\n",
        "\n",
        "def tensorsFromPair(pair):\n",
        "    input_tensor = tensorFromSentence(lang, pair[0])\n",
        "    target_tensor_t = tensorFromSentence(lang, pair[1][:-1])\n",
        "    target_tensor_l = tensorFromSentence(lang, pair[1])\n",
        "    return (input_tensor, target_tensor_t, target_tensor_l)\n",
        "\n",
        "\n",
        "lang, pairs = prepareData(texts, titles)\n",
        "print(random.choice(pairs))\n",
        "\n",
        "lens = [len(t) for t in texts]\n",
        "sns.distplot(lens)\n",
        "\n",
        "i = 100\n",
        "print(texts[i])\n",
        "print(titles[i])\n",
        "print(lang.max_seq_length)"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "275\n",
            "Read 96153 sentence pairs\n",
            "Counted words:\n",
            "ru 55964\n",
            "[['главарь', 'боевик', 'али', 'тазиев', 'кличка', 'магас', 'задержать', 'северный', 'кавказ', 'сообщить', 'директор', 'фсб', 'россия', 'александр', 'бортник', 'среда', 'встреча', 'президент', 'дмитрий', 'медведев'], ['главарь', 'боевик', 'магас', 'задержать', 'северный', 'кавказ']]\n",
            "['роджер', 'федерер', 'четыре', 'свой', 'карьера', 'выиграть', 'открытый', 'чемпионат', 'австралия', 'теннис']\n",
            "['роджер', 'федерер', 'четыре', 'выиграть', 'открытый', 'чемпионат', 'австралия']\n",
            "275\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXoAAAD4CAYAAADiry33AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+j8jraAAAeP0lEQVR4nO3da4xc533f8e9/7nvhLrnk8iKKFCVRbkLbraPSUpraLhDBtmSgoYPIiKwXVgEBatAISGG4qIKggiKkLxSgNhBYaCtDQhWhiuTKNcIgdBXbcms0VWhSliyKkmmuaJrifUVyb9yd2bn8+2LOWY5Gs7tnuXM5O/P7AAvOnHN25nl2pN8++z/PeY65OyIi0r0SnW6AiIi0loJeRKTLKehFRLqcgl5EpMsp6EVEulyq0w2ot2nTJt+1a1enmyEisqa89tpr77v7aKN9sQv6Xbt2cfjw4U43Q0RkTTGzXy22T6UbEZEup6AXEelyCnoRkS6noBcR6XIKehGRLqegFxHpcgp6EZEup6AXEelyCvqIxi7O8E/+9O84enay000REVkRBX1E/+/d95mcK/Kd1850uikiIiuioI/oyOnqSP57b52jUtFduURk7VDQR3TkzCSZVIJzk3neOD3R6eaIiESmoI8gXyxz/OIM931yB+mk8b0j5zrdJBGRyBT0EbxzbopyxfmtWzfyqd2bOHDkPLqpuoisFQr6CN46U63Pf/zG9Xzh49s4MzHHm6c1+0ZE1gYFfQRHzkwyMpDhhuEcn9uzlVTCOPCWyjcisjYo6CM4cmaKj20fxswY7k/z69uG+Pm56U43S0QkEgX9MvLFMscvTPPx7UML27YM5bgwle9gq0REolPQL+Pn56cpVZyPbx9e2LZ1OMt5Bb2IrBEK+mUcCU7Efqwm6LesyzExWyRfLHeqWSIikSnol/H22Uk29KfZvr5vYduW4RwAF6cKnWqWiEhkCvpljE8X2Dbch5ktbNs6VA16lW9EZC1Q0C9jYrbIcF/6A9u2DivoRWTtUNAvY3KuyPr+Dwb9lnXVoL8wqaAXkfiLFPRmdreZHTOzMTN7pMH+rJm9GOw/aGa7gu27zGzOzN4Ivv5Lc5vfepNzHx7RD/WlyKUTmmIpImtCarkDzCwJPAl8FjgNHDKz/e7+ds1hDwJX3H23md0HPAH8frDvXXf/RJPb3TYTc0WG60b0ZsbWoZxKNyKyJkQZ0d8BjLn7CXefB14A9tUdsw94Nnj8EnCX1Z69XKPyxTLzpcqHRvSgi6ZEZO2IEvTbgfdqnp8OtjU8xt1LwCSwMdh3s5m9bmb/x8w+3egNzOwhMztsZofHx8dX1IFWmpgtAjQM+q3DGtGLyNrQ6pOx54Cd7v4bwFeB581sqP4gd3/K3fe6+97R0dEWNym6yblq0K/vy3xoX3VEX9ByxSISe1GC/gywo+b5jcG2hseYWQoYBi65e8HdLwG4+2vAu8BHVtvodpmYnQcaj+i3DOWYL1UWRv0iInG17MlY4BBwm5ndTDXQ7wPurztmP/AA8CpwL/CKu7uZjQKX3b1sZrcAtwEnmtb6FgtH9K+euMSpy7Mf2Fd70dSGgQ+P+EVE4mLZEX1Qc38YeBl4B/i2ux81s8fN7HeCw54GNprZGNUSTTgF8zPAm2b2BtWTtH/g7peb3YlWCYO+L5380L6tw1lAF02JSPxFGdHj7geAA3XbHq15nAe+1OD7vgN8Z5Vt7Jgw6PszHw76LUO6aEpE1gZdGbuEybkiBmRSH/4xbQ6vjtXCZiIScwr6JUzMFsmlkyQaXBKQSSXYOJBR6UZEYk9Bv4TJuWLDsk1IF02JyFqgoF/CxFyRviWCfutwjvOq0YtIzEU6GdurJueKDWfcADx/8BTT+RK/unSV5w+eWth+/50729U8EZFINKJfwtQyI/qhXIqr82VKlUobWyUisjIK+iVMzM4vOqIH6M9W/yDKFxX0IhJfCvpFVCpeLd0sMaLPBdMudZNwEYkzBf0iZuZLVLzxVbGhXLCvoBG9iMSYgn4Rk7OLXxUbyoYj+pJG9CISXwr6RSy1zk0oqxG9iKwBCvpFLAR9ZvEZqGGNvqARvYjEmIJ+ESsZ0etkrIjEmYJ+EeENRaLMuimUVLoRkfhS0C8iyog+lUyQTJjm0YtIrCnoFzExN08mlSCd/PDKlbWyqYRq9CISawr6RUzNFRnuS2MNliiulUsnVboRkVhT0C9iYrbY8Kbg9XKphE7GikisKegXMTlXZH2EoM+mk6rRi0isKegXMTkXbUSvGr2IxJ2CfhETs0WG+yOUblSjF5GYU9AvYmoFI3rV6EUkzhT0DZTKFaYLpWgnY9NJCsUK7t6GlomIrJyCvoHpfAmAoVy0EX3ZnVJFQS8i8aSgbyAM+nW55W+pu7CCper0IhJTCvoGpvLV5Q/WRRjRL6x3ozq9iMSUgr6Ba6Wb5Uf04V2m8hrRi0hMKegbmF7BiD6r+8aKSMwp6Bu4rhq9ro4VkZiKFPRmdreZHTOzMTN7pMH+rJm9GOw/aGa76vbvNLMZM/tac5rdWtdG9BFKN7rLlIjE3LJBb2ZJ4EngHmAP8GUz21N32IPAFXffDXwDeKJu/9eB762+ue1xbUQfba0bUI1eROIryoj+DmDM3U+4+zzwArCv7ph9wLPB45eAuyxY39fMvgj8EjjanCa33nShRDaVIJNa/sejWTciEndRgn478F7N89PBtobHuHsJmAQ2mtkg8O+BP13qDczsITM7bGaHx8fHo7a9ZabzxUijedBdpkQk/lp9MvYx4BvuPrPUQe7+lLvvdfe9o6OjLW7S8qbypUhTK0NawVJE4ixKmp0BdtQ8vzHY1uiY02aWAoaBS8CdwL1m9ufAeqBiZnl3/+aqW95C0/lSpBOxIa1gKSJxFiXNDgG3mdnNVAP9PuD+umP2Aw8ArwL3Aq94dZWvT4cHmNljwEzcQx5WVroB3WVKROJt2aB395KZPQy8DCSBZ9z9qJk9Dhx29/3A08BzZjYGXKb6y2DNms6X2DqUi3x8ViN6EYmxSPUJdz8AHKjb9mjN4zzwpWVe47HraF9HVEf0K6vRT84VW9giEZHrpytjG6jW6FdQutGIXkRiTEFfp1SuMDtfXvGIXjV6EYkrBX2dmUL0q2JDGtGLSJwp6OusZEGzUDaVoFxxSmWFvYjEj4K+TnjTkRVdMKX1bkQkxhT0dVayoFlI692ISJwp6OtcX+lGI3oRiS8FfZ2V3F0qlEtrRC8i8RV92Nrlnj94CoBX330fgO+/fYHBbLQfz8JdpjSiF5EY0oi+Tlh+yUVYiz6U031jRSTGFPR18sUyqYSRSkb/0SzMulHQi0gMKejr5IuVheCO6tp9Y1W6EZH4UdDXyRfLKyrbQPUuU+mkMTevEb2IxI+Cvk6hVCa3whE9VJdBmFPpRkRiSEFfJ1+sLEyXXIm+dFI1ehGJJQV9nXxRI3oR6S4K+jrVGv3Kg746otfJWBGJHwV9nXzpOks3GY3oRSSeFPQ1Ku7Ml1Y+vRKqyyBo1o2IxJGCvkYhKL1cb40+Xyzj7s1ulojIqijoa4SzZlY6jx6qNXrn2h2qRETiQkFfI18Kgv46RvR9wfdM5RX0IhIvCvoa+VWWbgCm5opNbZOIyGop6GsslG6uc9YNwKSCXkRiRkFf41qNXiN6EekeCvoa4Vr02etcAgE0oheR+FHQ15gv6mSsiHQfBX2NfKlCwiCVsBV/b/hXgEo3IhI3kYLezO42s2NmNmZmjzTYnzWzF4P9B81sV7D9DjN7I/j6mZn9bnOb31yFUplsKonZyoM+YUYunVDpRkRiZ9mgN7Mk8CRwD7AH+LKZ7ak77EHgirvvBr4BPBFsfwvY6+6fAO4G/quZxfaG5IVi5brq86FcOslUXkEvIvESJdXuAMbc/YS7zwMvAPvqjtkHPBs8fgm4y8zM3WfdPSxa54BYrw9QKFWua8ZNqC+dZGpONXoRiZcoQb8deK/m+elgW8NjgmCfBDYCmNmdZnYUOAL8QU3wx061dLPKEb1KNyISMy0/GevuB939o8AngT82s1z9MWb2kJkdNrPD4+PjrW7Sogql1ZVu+lS6EZEYipJqZ4AdNc9vDLY1PCaowQ8Dl2oPcPd3gBngY/Vv4O5Pufted987OjoavfVNVihWyK6ydKOTsSISN1GC/hBwm5ndbGYZ4D5gf90x+4EHgsf3Aq+4uwffkwIws5uAXwNONqXlLbD60k1CpRsRiZ1lZ8C4e8nMHgZeBpLAM+5+1MweBw67+37gaeA5MxsDLlP9ZQDwKeARMysCFeDfuPv7rehIMxRKldUFfSbJ1fkypXKFVFKXKIhIPESa6ujuB4ADddserXmcB77U4PueA55bZRvbouIe1OhXV7qB6tWxIwOZZjVNRGRVNOwMFMN1blYxou/TwmYiEkMK+sDCgmarOBm7sIKlZt6ISIwo6AOFYEGz1U6vBK1gKSLxoqAPFJpQusllwtJNbK8JE5EepKAPFJpQutGIXkTiSEEfKAQ3Bm/KyVjV6EUkRhT0gcIqbgweSieNVMI060ZEYkVBHwhH9JlVjOjNjOG+tEo3IhIrCvpAWKPPrSLoAYb60rqdoIjEioI+UChVSCZs1UsXDOVSKt2ISKwo6AP54uoWNAsNqXQjIjGjoA/Mr3JBs9BQX1ojehGJFQV9IF9a3Vr0oZH+DFdm55vQIhGR5lDQBwrF8qqWPwhtGMgwMVekXIn17XFFpIco6AOrXYs+NNKfxl1Xx4pIfCjoA9W7S62+dLMhWIf+8lWVb0QkHhT0gaaN6IOgV51eROJCQR8oFCurWv4gtKFfI3oRiRcFPVCuOPPlyqqWPwgtjOgV9CISEwp64Op8dcmC1S5/ADUjepVuRCQmFPTA1UI16JtxMrYvk6QvndSIXkRiQ0EPzASLkDVjHj1UyzeXFPQiEhMKemBmYUTfnB/HhoG0RvQiEhsKemqDfvWlG4CRgSyXZ3XBlIjEg4KeFpRu+jWiF5H4UNDT/BH9hoGMgl5EYkNBz7Wgb8b0SqiuYDldKDEf3LVKRKSTFPRcm16ZaVLpJlzvZkJz6UUkBhT0wHShRCphpBLNm14JumhKROIhUrKZ2d1mdszMxszskQb7s2b2YrD/oJntCrZ/1sxeM7Mjwb+/3dzmN8fVQqkpyx+EtN6NiMTJsulmZkngSeAeYA/wZTPbU3fYg8AVd98NfAN4Itj+PvAv3f3jwAPAc81qeDPN5EtNWdAsdG29G02xFJHOizKMvQMYc/cT7j4PvADsqztmH/Bs8Pgl4C4zM3d/3d3PBtuPAn1mlm1Gw5tpptCcG4OHNgykAZVuRCQeoqTbduC9mueng20Nj3H3EjAJbKw75veAn7p7of4NzOwhMztsZofHx8ejtr1pZgrF5gZ9v1awFJH4aMvJWDP7KNVyzr9utN/dn3L3ve6+d3R0tB1N+oCrhebcXSqUTiZYl0upRi8isRAl6M8AO2qe3xhsa3iMmaWAYeBS8PxG4LvAV9z93dU2uBXOT+UZzKaa+pojAxndZUpEYiFK0B8CbjOzm80sA9wH7K87Zj/Vk60A9wKvuLub2Xrgb4FH3P3vm9XoZpqYnWd8usDmoeaeOtjQn9GIXkRiYdmgD2ruDwMvA+8A33b3o2b2uJn9TnDY08BGMxsDvgqEUzAfBnYDj5rZG8HX5qb3YhXGLs4AsHldc4NeI3oRiYtI9Qp3PwAcqNv2aM3jPPClBt/3Z8CfrbKNLXV8IehzTX3dDf0Zjp2fbupriohcj56/Mvb4hRn60kmG+9NNfd2RgbRKNyISC809A7kGHb84ze7NgyTMmvJ6zx88BcCpS7PMFcv8t78/SSaV4P47dzbl9UVEVqrnR/RjF2e4bfNg01+3P5jFMxvceFxEpFN6Ouin80XOTebZvaX5QT+Qqc7Ln50vN/21RURWoqeDPpxxc9vmdU1/7XBe/nRe692ISGf1dNAfXwj65o/oh/qqJ3cn51S6EZHO6umgH7s4QyaVYMdIf9Nfe10ujQGTcxrRi0hn9XTQH78wza2jgyQTzZlxUyuZMNblUkwp6EWkw3o76Fs04yY03JfWiF5EOq5ng/5qocTpK3MKehHpej0b9CfGrwKwuw1B7+4tew8RkeX0bNCfvjIL0JITsaHhvjTz5Qr5YqVl7yEispyeDfqzk3kAbljf17L3WJhiqbn0ItJBPRv05ybmyKYSbGjyYma1hsOgn1XQi0jn9G7QT+a5YX0f1qTFzBoJg15TLEWkk3o26M9OzrFtuLlr0NcLL5qaUNCLSAf1bNCfm8izbbh19XnQRVMiEg89GfSlcoWL03luWN/aET0EUyx1MlZEOqgng/7CdIGK0/IRPVRn3uhkrIh0Uk8G/bmJOYC2jOjX66IpEemwngz6dsyhDw0FF01N5bVcsYh0Rm8GfTCib/WsG7g2xfJ88MtFRKTdeu7m4M8fPMX/PjZONpXgb352ruXvFwb9uck5/tHW5t/JSkRkOT05op+cKy4EcKtdC3qN6EWkM3o06OdZ38KlD2qFF00p6EWkU3oz6GfbN6IPL5oKzwuIiLRbzwV9sVzh6ny5bUEPsKE/s7AssohIu/Vc0IfLEQz3Zdr2nhsGMrx3WSN6EemMngv6iYWgb9+IfmQgw9nJOeZLugGJiLRfpKA3s7vN7JiZjZnZIw32Z83sxWD/QTPbFWzfaGY/MrMZM/tmc5t+fcJ7uK5vZ9D3Z3CHM6rTi0gHLBv0ZpYEngTuAfYAXzazPXWHPQhccffdwDeAJ4LteeA/AF9rWotXKQz64TbNuoFq6Qbg1GXV6UWk/aKM6O8Axtz9hLvPAy8A++qO2Qc8Gzx+CbjLzMzdr7r7/6Ua+LFw+eo8/Zkk6WT7qlYjCnoR6aAoabcdeK/m+elgW8Nj3L0ETAIbozbCzB4ys8Nmdnh8fDzqt61Yvljm7bNT3DI62LL3aGRdLkUmleC0gl5EOiAWJ2Pd/Sl33+vue0dHR1v2Pv/rrfPMFcvcsWukZe/RSMKMHRv6NKIXkY6IEvRngB01z28MtjU8xsxSwDBwqRkNbKa/+skpRgYy3DI60Pb33jnSr6AXkY6IsqjZIeA2M7uZaqDfB9xfd8x+4AHgVeBe4BXv8ALszx889YHn49MFDv7yMp/fs4VEC28IvpidI/0cPnkFd2/pDclFROotG/TuXjKzh4GXgSTwjLsfNbPHgcPuvh94GnjOzMaAy1R/GQBgZieBISBjZl8EPufubze/K0s7dPIyCYPbb9rQ7rcGYMdIP9OFEpNzRdb3t+9iLRGRSMsUu/sB4EDdtkdrHueBLy3yvbtW0b6mqLjz01NX+PVtQ6zLtW9aZa2dI/1AdeaNgl5E2ikWJ2Nb7WqhxOx8mVs2tb82H9pRE/QiIu3UE0E/U6jexq9To3lQ0ItI5/RE0E/nw6Dv3A21BrMpNg5keE9BLyJt1iNBX132oJMjeqiO6jWiF5F265Ggr47oB7OdvUWu5tKLSCf0xM3Bp/MlsqkEmVTnfq89f/AUU/kiZ67M8ZevniSVqLbl/jt3dqxNItIbemNEXyh1vGwDsHldlorD+9PznW6KiPSQ3gj6fLGjJ2JDW4f7ADg3qXXpRaR9eiLoZ/KlWAT96GCWZMI4PxmbVZtFpAf0RNBP50us6/CJWIBkwtgylOXclIJeRNqn64O+UCwzX67EokYPsG2oj3Ma0YtIG3V90E8XOn+xVK2twzmuFkoLc/tFRFqt+4M+nEMfk6DfNpwD0KheRNqmB4I+HlfFhrYFM290QlZE2qUHgj4o3cTgZCxAXybJcF9aUyxFpG16IuiTZvRnkp1uyoJtwzmVbkSkbbo+6GcKRQZzqVjdvm/bcI73ZwoUy5VON0VEekDXB/10TC6WqrV1uI+Kw8WpQqebIiI9oDeCPib1+VA48+b0hFayFJHW64GgLzIYkxk3oY0DGTYNZjl08jLu3unmiEiX6+qgL1ec2fly7Eo3Zsand2/i7ESeV09c6nRzRKTLdXXQzxRKOPG5KrbWJ3auZyCb4ls/PtHppohIl+vqoF+4WCobr9INQDqZ4J/dMsKPjo1z/MJ0p5sjIl2sq4N+JgY3BV/KnTdvJJdO8JRG9SLSQl0d9BNz4fIH8Qz6gWyK+++4iZd+eprXfnW5080RkS7V1UF/9OwkG/rTDPXFr3QT+urnPsINw338u//xJvliudPNEZEu1LVBf2V2nhPjV7l95wYSMboqtt5gNsUTv/ePOfH+Vb7+/V90ujki0oXiWdNogtdPXcGB23du6HRTlvT8wVMAfHLXBr714xOMTxf45K4R7r9zZ4dbJiLdoitH9O7OT09NcMumATYMZDrdnEi+8LFt3Lp5kO++fobvvn6GQkllHBFpjkhBb2Z3m9kxMxszs0ca7M+a2YvB/oNmtqtm3x8H24+Z2eeb1/TFHTp5hctX5/mnN8V7NF8rm07yr35rF//iI6McOnmZO/7jD/mjF15n/8/OMjmru1GJyPVbtnRjZkngSeCzwGngkJntd/e3aw57ELji7rvN7D7gCeD3zWwPcB/wUeAG4Adm9hF3b8lwtViu8PqpCf7ih8fJphJ89IbhVrxNyyTM+PxHt3Lr6CBvvDfBD96+wF+/cZaEwU0bBxhdl+Wf37qJzUNZNq/LsmkwSypppBIJkonq9ycTRsKMVNJImpFIXPs3lbi23wzcwXHCVRjCf81YOLZ21c9KxSm7kzAjYcRqRVCJtyhLfURZDSTKgiGR3ivC64T/P4WvWSw7ZtVrYKD6/0O+VCaVSJBJVbfli2Vm58v0Z5JkUwkqDpNzReaKZdb3penPJJkrljk7kadYrrBtOMdQLs2ZiTnGLs4wmEvxyV0jEVq3MlFq9HcAY+5+AsDMXgD2AbVBvw94LHj8EvBNq6bAPuAFdy8AvzSzseD1Xm1O86957VdX+MrTB7k6XyZh8Nu/tmXhh7/W7N48yO7Ng1TcOX15lp+fn+b4xRneOjPJT37Z3mmYYeBXgv/Qa5lR/SViRiL4RVOr/v83r/vfazXL/DTrd0yzwiXKQfX9b3hMzMKu15diSiaMdNKYL1Wo+LVtqYRRKF1bZjwVDKDmyx/cVnb/wM8wlTBKlQ/+UBPGwmt/bs+WjgX9duC9muengTsXO8bdS2Y2CWwMtv9D3fdur38DM3sIeCh4OmNmxyK1fglPw6an4f3Vvk5MbaI7+6Z+rT3d2reO9OtbwLceuO5vv2mxHbGYdePuTwFPNfM1zeywu+9t5mvGRbf2Tf1ae7q1b93Wryi1jTPAjprnNwbbGh5jZilgGLgU8XtFRKSFogT9IeA2M7vZzDJUT67urztmPxD+wXEv8IpXi4T7gfuCWTk3A7cBP2lO00VEJIplSzdBzf1h4GUgCTzj7kfN7HHgsLvvB54GngtOtl6m+suA4LhvUz1xWwL+sFUzbhpoaikoZrq1b+rX2tOtfeuqfpnucCQi0t3W5vxDERGJTEEvItLlujLol1uyYS0xs5NmdsTM3jCzw8G2ETP7vpkdD/5dE2s9mNkzZnbRzN6q2dawL1b1F8Fn+KaZ3d65li9tkX49ZmZngs/tDTP7Qs2+ti8Lcj3MbIeZ/cjM3jazo2b2R8H2Nf2ZLdGvNf+ZLcrdu+qL6gnjd4FbgAzwM2BPp9u1iv6cBDbVbftz4JHg8SPAE51uZ8S+fAa4HXhrub4AXwC+Bxjwm8DBTrd/hf16DPhag2P3BP9NZoGbg/9Wk53uwyL92gbcHjxeB/wiaP+a/syW6Nea/8wW++rGEf3Ckg3uPg+ESzZ0k33As8HjZ4EvdrAtkbn7j6nOyqq1WF/2AX/pVf8ArDezbe1p6cos0q/FLCwL4u6/BMJlQWLH3c+5+0+Dx9PAO1SvbF/Tn9kS/VrMmvnMFtONQd9oyYalPsS4c+DvzOy1YKkIgC3ufi54fB7Y0pmmNcVifemGz/HhoITxTE15bU32K1iR9jeAg3TRZ1bXL+iiz6xWNwZ9t/mUu98O3AP8oZl9pnanV/+27Io5st3UF+A/A7cCnwDOAf+ps825fmY2CHwH+LfuPlW7by1/Zg361TWfWb1uDPquWnbB3c8E/14Evkv1T8YL4Z/Ewb8XO9fCVVusL2v6c3T3C+5edvcK1bWqwj/111S/zCxNNQz/u7v/z2Dzmv/MGvWrWz6zRrox6KMs2bAmmNmAma0LHwOfA97ig0tOPAD8dWda2BSL9WU/8JVgJsdvApM15YLYq6tN/y7Vzw3W0LIgZmZUr3p/x92/XrNrTX9mi/WrGz6zRXX6bHArvqie/f8F1bPjf9Lp9qyiH7dQPdv/M+Bo2BeqS0D/EDgO/AAY6XRbI/bnr6j+SVykWud8cLG+UJ258WTwGR4B9na6/Svs13NBu9+kGhTbao7/k6Bfx4B7Ot3+Jfr1KaplmTeBN4KvL6z1z2yJfq35z2yxLy2BICLS5bqxdCMiIjUU9CIiXU5BLyLS5RT0IiJdTkEvItLlFPQiIl1OQS8i0uX+P4enJB+7ExcQAAAAAElFTkSuQmCC\n",
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ]
          },
          "metadata": {
            "tags": [],
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R0FOQBL0NU0Z",
        "colab_type": "text"
      },
      "source": [
        "# Datasets"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Br5sXOantT0z",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from torch.utils.data.dataset import Dataset\n",
        "from torch.utils.data import DataLoader\n",
        "\n",
        "\n",
        "class HeadlinesDataset(Dataset):\n",
        "    def __init__(self, pairs, lang, batch_size):\n",
        "        self.pairs = pairs\n",
        "        self.lang = lang\n",
        "        self.batch_size = batch_size\n",
        "        self.batch_num = len(self.pairs)//batch_size\n",
        "        \n",
        "    def __getitem__(self, index): # todo: make number of pads depend on current batch max text, title length\n",
        "        srcs = []\n",
        "        srcs_lns = []\n",
        "        tgts = []\n",
        "        tgts_lns = []\n",
        "        for i in range(self.batch_size):\n",
        "          src, tgt = self.pairs[index*self.batch_size+i]\n",
        "          srcs.append(src)\n",
        "          srcs_lns.append(len(src))\n",
        "          tgts.append(tgt)\n",
        "          tgts_lns.append(len(tgt))\n",
        "\n",
        "        max_src_ln = max(srcs_lns)          \n",
        "        max_tgt_ln = max(tgts_lns)    \n",
        "\n",
        "        converted_srcs = []    \n",
        "        converted_tgts = []    \n",
        "        srcs_padding_masks = []\n",
        "        tgts_padding_masks = []\n",
        "\n",
        "        for i in range(len(srcs)):\n",
        "          converted_srcs.append(self.lang.text2ids(srcs[i]) + [self.lang.word2index['PAD']]*(max_src_ln-len(srcs[i]))) \n",
        "          srcs_padding_masks.append([False]*len(srcs[i]) + [True]*(max_src_ln-len(srcs[i])))\n",
        "\n",
        "          converted_tgts.append(\n",
        "              [self.lang.word2index['SOS']] + \n",
        "                self.lang.text2ids(tgts[i]) + \n",
        "              [self.lang.word2index['EOS']] + \n",
        "              [self.lang.word2index['PAD']]*(max_tgt_ln-len(tgts[i]))\n",
        "              ) \n",
        "          tgts_padding_masks.append([False]*(len(tgts[i])+2) + [True]*(max_tgt_ln-len(tgts[i])))\n",
        "\n",
        "        return (np.array(converted_srcs), np.array(converted_tgts), np.array(srcs_padding_masks), np.array(tgts_padding_masks))\n",
        "\n",
        "    def __len__(self):\n",
        "        return self.batch_num\n",
        "\n",
        "\n",
        "# train_dataset = HeadlinesDataset(pairs[1024:], lang, batch_size=config['batch_size'])\n",
        "# test_dataset  = HeadlinesDataset(pairs[:1024], lang, batch_size=config['batch_size'])\n",
        "\n",
        "# train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=4, pin_memory=True)\n",
        "# test_loader = DataLoader(test_dataset, batch_size=1, shuffle=True, num_workers=4, pin_memory=True)\n",
        "\n",
        "# src, tgt, src_key_padding_mask, tgt_key_padding_mask =  next(iter(train_loader))\n",
        "# src, tgt, src_key_padding_mask, tgt_key_padding_mask =  next(iter(test_loader))\n",
        "# print(src)\n",
        "# print(tgt)\n",
        "# print(src_key_padding_mask)\n",
        "# print(tgt_key_padding_mask)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HTf3m8gdNcJr",
        "colab_type": "text"
      },
      "source": [
        "# Prepear word embeddings"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Wty7UgXlNIXZ",
        "colab_type": "code",
        "outputId": "165d8369-b288-4417-d7e8-a991c90b11c4",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 53
        }
      },
      "source": [
        "from random import shuffle\n",
        "\n",
        "union = []\n",
        "for pair in tqdm.tqdm(pairs):\n",
        "  union.append(['SOS']+pair[0]+['EOS']+['PAD']+['UNK'])\n",
        "\n",
        "for pair in tqdm.tqdm(pairs):\n",
        "  union.append(['SOS']+pair[1]+['EOS']+['PAD']+['UNK'])\n",
        "\n",
        "shuffle(union)\n",
        "\n",
        "\n",
        "from gensim.models.word2vec import Word2Vec\n",
        "from multiprocessing import cpu_count\n",
        "\n",
        "model = Word2Vec(union, \n",
        "                 min_count = 0, \n",
        "                 workers=cpu_count(),\n",
        "                 size=config['d_model'],\n",
        "                 window=10,\n",
        "                 negative=5,\n",
        "                 iter=10\n",
        "                )\n",
        "\n",
        "\n",
        "init_voc = set(lang.word2index.keys())\n",
        "wv_voc = set(model.wv.vocab.keys())\n",
        "print(len(wv_voc), len(init_voc))\n",
        "print(wv_voc - init_voc)\n",
        "print(init_voc - wv_voc)\n",
        "# assert(len(init_voc)==len(wv_voc))\n",
        "\n",
        "\n",
        "word_init_vectors = {}\n",
        "for w in model.wv.vocab:\n",
        "    word_init_vectors[w] = model[w].tolist()"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 96153/96153 [00:00<00:00, 232716.24it/s]\n",
            "100%|██████████| 96153/96153 [00:00<00:00, 270177.06it/s]\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5h7Rvg6vWyYI",
        "colab_type": "text"
      },
      "source": [
        "# Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nmcu_t772CVE",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import math \n",
        "# Source: https://pytorch.org/tutorials/beginner/transformer_tutorial.html\n",
        "class PositionalEncoding(nn.Module):\n",
        "    def __init__(self, d_model, dropout=0.1, max_len=200):\n",
        "        super(PositionalEncoding, self).__init__()\n",
        "        self.dropout = nn.Dropout(p=dropout)\n",
        "\n",
        "        pe = torch.zeros(max_len, d_model)\n",
        "        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n",
        "        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))\n",
        "        pe[:, 0::2] = torch.sin(position * div_term)\n",
        "        pe[:, 1::2] = torch.cos(position * div_term)\n",
        "        pe = pe.unsqueeze(0).transpose(0, 1)\n",
        "        self.register_buffer('pe', pe)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x + self.pe[:x.size(0), :]\n",
        "        return self.dropout(x)\n",
        "\n",
        "\n",
        "class ScheduledOptim():\n",
        "    '''A simple wrapper class for learning rate scheduling'''\n",
        "\n",
        "    def __init__(self, optimizer, d_model, n_warmup_steps):\n",
        "        self._optimizer = optimizer\n",
        "        self.n_warmup_steps = n_warmup_steps\n",
        "        self.n_current_steps = 0\n",
        "        self.init_lr = np.power(d_model, -0.5)\n",
        "\n",
        "    def step_and_update_lr(self):\n",
        "        \"Step with the inner optimizer\"\n",
        "        self._update_learning_rate()\n",
        "        self._optimizer.step()\n",
        "\n",
        "    def zero_grad(self):\n",
        "        \"Zero out the gradients by the inner optimizer\"\n",
        "        self._optimizer.zero_grad()\n",
        "\n",
        "    def _get_lr_scale(self):\n",
        "        return np.min([\n",
        "            np.power(self.n_current_steps, -0.5),\n",
        "            np.power(self.n_warmup_steps, -1.5) * self.n_current_steps])\n",
        "\n",
        "    def _update_learning_rate(self):\n",
        "        ''' Learning rate scheduling per step '''\n",
        "\n",
        "        self.n_current_steps += 1\n",
        "        lr = self.init_lr * self._get_lr_scale()\n",
        "\n",
        "        for param_group in self._optimizer.param_groups:\n",
        "            param_group['lr'] = lr\n",
        "\n",
        "\n",
        "def gen_nopeek_mask(length):\n",
        "    mask = rearrange(torch.triu(torch.ones(length, length)) == 1, 'h w -> w h')\n",
        "    mask = rearrange(torch.triu(torch.ones(length, length)) == 1, 'h w -> w h')\n",
        "    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))\n",
        "    return mask"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fkJPanlrMuZG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class LanguageTransformer(nn.Module):\n",
        "    def __init__(self, vocab_size, d_model, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, max_seq_length, pos_dropout, trans_dropout):\n",
        "        super().__init__()\n",
        "        self.d_model = d_model\n",
        "        self.embed_src = nn.Embedding(vocab_size, d_model)\n",
        "        self.embed_tgt = nn.Embedding(vocab_size, d_model)\n",
        "        self.pos_enc = PositionalEncoding(d_model, pos_dropout, max_seq_length)\n",
        "\n",
        "        self.transformer = nn.Transformer(d_model, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, trans_dropout)\n",
        "        self.fc = nn.Linear(d_model, vocab_size)\n",
        "\n",
        "    def init_embeddings(self, wv):\n",
        "        assert(self.embed_src.weight.data.shape == wv.shape)\n",
        "        assert(self.embed_tgt.weight.data.shape == wv.shape)\n",
        "        self.embed_src.weight.data = wv\n",
        "        self.embed_tgt.weight.data = wv\n",
        "\n",
        "    def forward(self, src, tgt, src_key_padding_mask, tgt_key_padding_mask, memory_key_padding_mask, tgt_mask):\n",
        "        src = src.permute(1, 0)\n",
        "        tgt = tgt.permute(1, 0)\n",
        "        src = self.pos_enc(self.embed_src(src) * math.sqrt(self.d_model)) # why every time multiply? just init it somewhere\n",
        "        tgt = self.pos_enc(self.embed_tgt(tgt) * math.sqrt(self.d_model))\n",
        "\n",
        "        output = self.transformer(src, tgt, tgt_mask=tgt_mask, src_key_padding_mask=src_key_padding_mask,\n",
        "                                  tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=memory_key_padding_mask)\n",
        "        output = output.permute(1, 0, 2)\n",
        "        return self.fc(output)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "R7JN7ZUbZAOC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# init\n",
        "model = LanguageTransformer(lang.n_words, config['d_model'], config['nhead'], \n",
        "                            config['num_encoder_layers'],config['num_decoder_layers'],\n",
        "                            config['dim_feedforward'], lang.max_seq_length,\n",
        "                            config['pos_dropout'], config['trans_dropout'])\n",
        "for p in model.parameters():\n",
        "    if p.dim() > 1:\n",
        "        nn.init.xavier_normal_(p)\n",
        "\n",
        "wv = torch.Tensor([word_init_vectors[lang.index2word[idx]] for idx in lang.index2word.keys()])\n",
        "model.init_embeddings(wv)\n",
        "\n",
        "model.to(device)\n",
        "# optimizer = Adam(model.parameters(), lr=0.0001)\n",
        "# optimizer = Adam(model.parameters(), betas=(0.9, 0.98), eps=1e-09)\n",
        "optimizer = ScheduledOptim(\n",
        "    Adam(model.parameters(), betas=(0.9, 0.98), eps=1e-09),\n",
        "    config['d_model'], config['n_warmup_steps'])\n",
        "\n",
        "criterion = nn.CrossEntropyLoss(ignore_index=lang.word2index['PAD'])\n",
        "train_losses = []\n",
        "test_losses = []"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sVSI7qYj0gds",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# train\n",
        "epochs = 10\n",
        "\n",
        "train_dataset = HeadlinesDataset(pairs[10000:], lang, batch_size=config['batch_size'])\n",
        "test_dataset  = HeadlinesDataset(pairs[:10000], lang, batch_size=config['batch_size'])\n",
        "\n",
        "for epoch_i in range(epochs):\n",
        "    train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=0, pin_memory=True)\n",
        "    test_loader = DataLoader(test_dataset, batch_size=1, shuffle=True, num_workers=0, pin_memory=True)\n",
        "    \n",
        "    train_loss = 0\n",
        "    model.train()\n",
        "    for src, tgt, src_key_padding_mask, tgt_key_padding_mask in tqdm.tqdm(iter(train_loader)):\n",
        "    # for src, tgt, src_key_padding_mask, tgt_key_padding_mask in iter(train_loader):\n",
        "        src, src_key_padding_mask = src[0].to(device), src_key_padding_mask[0].to(device) # to device\n",
        "        tgt, tgt_key_padding_mask = tgt[0].to(device), tgt_key_padding_mask[0].to(device) # to device\n",
        "        memory_key_padding_mask = src_key_padding_mask.clone()\n",
        "        tgt_inp, tgt_out = tgt[:, :-1], tgt[:, 1:]\n",
        "        tgt_key_padding_mask_inp = tgt_key_padding_mask[:, :-1]\n",
        "        tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to(device) # also to device\n",
        "\n",
        "        optimizer.zero_grad()\n",
        "        output = model(src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask_inp, memory_key_padding_mask, tgt_mask)\n",
        "        loss = criterion(rearrange(output, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))\n",
        "        loss.backward()\n",
        "        # optimizer.step()\n",
        "        optimizer.step_and_update_lr()\n",
        "        train_loss += loss.item()\n",
        "    train_losses.append((epoch_i, train_loss))\n",
        "\n",
        "\n",
        "    test_loss = 0\n",
        "    model.eval()\n",
        "    for src, tgt, src_key_padding_mask, tgt_key_padding_mask in tqdm.tqdm(iter(test_loader)):\n",
        "    # for src, tgt, src_key_padding_mask, tgt_key_padding_mask in iter(test_loader):\n",
        "        with torch.no_grad():\n",
        "            src, src_key_padding_mask = src[0].to(device), src_key_padding_mask[0].to(device) # to device\n",
        "            tgt, tgt_key_padding_mask = tgt[0].to(device), tgt_key_padding_mask[0].to(device) # to device\n",
        "            memory_key_padding_mask = src_key_padding_mask.clone()\n",
        "            tgt_inp, tgt_out = tgt[:, :-1], tgt[:, 1:]\n",
        "            tgt_key_padding_mask_inp = tgt_key_padding_mask[:, :-1]\n",
        "            tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to(device) # also to device\n",
        "\n",
        "            output = model(src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask_inp, memory_key_padding_mask, tgt_mask)\n",
        "            loss = criterion(rearrange(output, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))\n",
        "            test_loss += loss.item()\n",
        "    test_losses.append((epoch_i, test_loss))\n",
        "\n",
        "    print(f\"\\nEpoch: {epoch_i} Train loss: {train_loss} Test loss: {test_loss}\")\n",
        "    # break"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OFQdjoM9NB-s",
        "colab_type": "code",
        "outputId": "904385d7-0937-448b-b442-ef1a4e2077c3",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 53
        }
      },
      "source": [
        "print(train_losses)\n",
        "print(test_losses)"
      ],
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[(0, 10277.298031806946)]\n",
            "[(0, 1002.5683012008667)]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O8iHopzfzCPd",
        "colab_type": "text"
      },
      "source": [
        "# Inference"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VJivaP2o8357",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 92
        },
        "outputId": "0a7fc76a-8992-4b01-abdc-95840f0f7d11"
      },
      "source": [
        "src, tgt, src_key_padding_mask, tgt_key_padding_mask =  next(iter(train_loader))\n",
        "# print(src.shape)\n",
        "# print(tgt.shape)\n",
        "# print(src_key_padding_mask.shape)\n",
        "# print(tgt_key_padding_mask.shape)\n",
        "# print()\n",
        "\n",
        "text_i = 3\n",
        "src, src_key_padding_mask = src[0][text_i:text_i+1].to(device), src_key_padding_mask[0][text_i:text_i+1].to(device) # to device\n",
        "tgt, tgt_key_padding_mask = tgt[0][text_i:text_i+1].to(device), tgt_key_padding_mask[0][text_i:text_i+1].to(device) # to device\n",
        "# print(src.shape)\n",
        "# print(tgt.shape)\n",
        "# print(src_key_padding_mask.shape)\n",
        "# print(tgt_key_padding_mask.shape)\n",
        "\n",
        "\n",
        "def forced_inference(model, src, tgt, src_key_padding_mask, tgt_key_padding_mask):\n",
        "  # TODO: add accuracy / recall \n",
        "  memory_key_padding_mask = src_key_padding_mask.clone()  \n",
        "  tgt_inp, tgt_out = tgt[:, :-1], tgt[:, 1:]\n",
        "  tgt_key_padding_mask_inp = tgt_key_padding_mask[:, :-1]\n",
        "  tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to(device) # also to device\n",
        "  # print(src.shape)\n",
        "  # print(tgt_inp.shape)\n",
        "  # print(tgt_out.shape)\n",
        "  # print(tgt_key_padding_mask_inp.shape)\n",
        "  # print(tgt_mask.shape)\n",
        "  # print()\n",
        "  output = model(src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask_inp, memory_key_padding_mask, tgt_mask)\n",
        "  tokens = torch.argmax(output[0], axis=1)\n",
        "  return tokens\n",
        "\n",
        "def greedy_inference(model, src, tgt, src_key_padding_mask, tgt_key_padding_mask, lang):\n",
        "  memory_key_padding_mask = src_key_padding_mask.clone() \n",
        "  tgt_inp, tgt_out = tgt[:, :-1], tgt[:, 1:]\n",
        "  tgt_key_padding_mask_inp = tgt_key_padding_mask[:, :-1]\n",
        "  tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to(device) # also to device\n",
        "\n",
        "  tokens = [tgt_inp[0,0].item()]\n",
        "  tgt_inp_ = tgt_inp[:]\n",
        "  \n",
        "  for tgt_i in range(tgt.shape[1]-1):\n",
        "    tt = torch.Tensor(tokens)\n",
        "    tgt_inp_[:, :tt.shape[0]] = tt[:]\n",
        "    output = model(src, tgt_inp_, src_key_padding_mask, tgt_key_padding_mask_inp, memory_key_padding_mask, tgt_mask)\n",
        "    next_token = torch.argmax(output[0][tt.shape[0]+1]).item()\n",
        "    tokens.append(next_token)\n",
        "    if next_token == lang.word2index['EOS']:\n",
        "      break\n",
        "\n",
        "  return torch.Tensor(tokens)\n",
        "\n",
        "def inference(model, src, tgt, src_key_padding_mask, tgt_key_padding_mask, lang, mode):\n",
        "  \"\"\"\n",
        "  src: Tensor [1, S]\n",
        "  tgt: Tensor [1, T]\n",
        "  modes: forced / greedy [/ beam search]\n",
        "  \"\"\"\n",
        "  \n",
        "  if mode == \"forced\":\n",
        "    tokens = forced_inference(model, src, tgt, src_key_padding_mask, tgt_key_padding_mask)\n",
        "    \n",
        "  elif mode == \"greedy\":\n",
        "    tokens = greedy_inference(model, src, tgt, src_key_padding_mask, tgt_key_padding_mask, lang)\n",
        "\n",
        "  print(f\"Mode: {mode}\")\n",
        "  print(f'Source: {lang.tensor2text(src[0])}')\n",
        "  print(f'Target: {lang.tensor2text(tgt[0])}')\n",
        "  print(f'Generated: {lang.tensor2text(tokens)}')\n",
        "\n",
        "\n",
        "inference(model, src, tgt, src_key_padding_mask, tgt_key_padding_mask, lang, 'forced')\n",
        "inference(model, src, tgt, src_key_padding_mask, tgt_key_padding_mask, lang, 'greedy')"
      ],
      "execution_count": 80,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Source: ['наблюдение', 'квартира', 'который', 'высокий', 'вероятность', 'возникновение', 'пожар', 'планироваться', 'организовать', 'москва', 'рассказать', 'риа', 'новость', 'понедельник', 'источник', 'столичный', 'администрация', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD']\n",
            "Target: ['SOS', 'планировать', 'пожар', 'усилить', 'контроль', 'пожароопасный', 'квартира', 'EOS', 'PAD', 'PAD', 'PAD']\n",
            "Generated: ['SOS', 'планировать', 'пожар', 'EOS']\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tAwc2C0qVeAK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QMZ-3RmlVc6x",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oeMGFrfczXBd",
        "colab_type": "code",
        "outputId": "d2e8ab62-4ba6-4612-a284-2d662163356c",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 181
        }
      },
      "source": [
        "# failed sanity check\n",
        "\n",
        "src, tgt, src_key_padding_mask, tgt_key_padding_mask =  next(iter(train_loader))\n",
        "text_i = 3\n",
        "src, src_key_padding_mask = src[0][text_i:text_i+1].to(device), src_key_padding_mask[0][text_i:text_i+1].to(device) # to device\n",
        "tgt, tgt_key_padding_mask = tgt[0][text_i:text_i+1].to(device), tgt_key_padding_mask[0][text_i:text_i+1].to(device) # to device\n",
        "\n",
        "tgt_inp, tgt_out = tgt[:, :-1], tgt[:, 1:]\n",
        "tgt_key_padding_mask_inp = tgt_key_padding_mask[:, :-1]\n",
        "tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to(device) # also to device\n",
        "print(src.shape)\n",
        "print(tgt_inp.shape)\n",
        "print(tgt_out.shape)\n",
        "print(tgt_key_padding_mask_inp.shape)\n",
        "print(tgt_mask.shape)\n",
        "print()\n",
        "\n",
        "# tgt_key_padding_mask_inp = torch.Tensor([[False]+[True]*(tgt_inp.shape[1]-1)]).bool().to(device)\n",
        "model.eval()\n",
        "tgt_mask[:, 1:] = tgt_mask[0, 1] # sanity check\n",
        "output = model(src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask_inp, memory_key_padding_mask, tgt_mask)\n",
        "print(output.shape)\n",
        "print(output[0, :, 9]) # they should be same"
      ],
      "execution_count": 48,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "torch.Size([1, 39])\n",
            "torch.Size([1, 10])\n",
            "torch.Size([1, 10])\n",
            "torch.Size([1, 10])\n",
            "torch.Size([10, 10])\n",
            "\n",
            "torch.Size([1, 10, 55964])\n",
            "tensor([1.6116, 1.7882, 1.5525, 2.7367, 1.1095, 1.6610, 1.4240, 1.8079, 1.9798,\n",
            "        1.9811], device='cuda:0', grad_fn=<SelectBackward>)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TAaTFU_ENoVu",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 53
        },
        "outputId": "545811f1-979d-4fab-c8d1-f88758b514eb"
      },
      "source": [
        ""
      ],
      "execution_count": 42,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[2253,  150,  860,  689,  705,  835, 1093, 1728,    1,    2]],\n",
              "       device='cuda:0')"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 42
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ET5qGZpDN7Wt",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 162
        },
        "outputId": "3d9a1580-73b6-4101-e2cb-5c45e1cad7da"
      },
      "source": [
        "output = model(src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask_inp, memory_key_padding_mask, tgt_mask)\n",
        "print(output)"
      ],
      "execution_count": 30,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "tensor([[[-7.4522,  0.5211, -7.4363,  ..., -7.5318, -7.4330, -7.4693],\n",
            "         [-7.1339,  0.4152, -7.1078,  ..., -7.2298, -7.0972, -7.1576],\n",
            "         [-7.6243,  0.6096, -7.5907,  ..., -7.7260, -7.6079, -7.6481],\n",
            "         ...,\n",
            "         [-7.4810,  0.6389, -7.4677,  ..., -7.6108, -7.4503, -7.5029],\n",
            "         [-7.5305,  0.7622, -7.4795,  ..., -7.6175, -7.4777, -7.5463],\n",
            "         [-7.6285,  0.4827, -7.5915,  ..., -7.7473, -7.6052, -7.6280]]],\n",
            "       device='cuda:0', grad_fn=<AddBackward0>)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y8FiTkKpzLiQ",
        "colab_type": "code",
        "outputId": "985523e9-5d5c-4065-9911-80c73d8fb199",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 108
        }
      },
      "source": [
        "print(tgt_inp)\n",
        "print(tgt_key_padding_mask_inp)\n",
        "print()"
      ],
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "tensor([[   0, 2253,  150,  860,  689,  705,  835, 1093, 1728,    1]],\n",
            "       device='cuda:0')\n",
            "tensor([[False, False, False, False, False, False, False, False, False, False]],\n",
            "       device='cuda:0')\n",
            "tensor([[False,  True,  True,  True,  True,  True,  True,  True,  True,  True]])\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fld_JsGDzR84",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "loss = criterion(rearrange(output, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EqhVe5f084Gj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# inference\n",
        "\n",
        "\n",
        "src, tgt, src_key_padding_mask, tgt_key_padding_mask =  next(iter(test_loader))\n",
        "print(src)\n",
        "print(tgt)\n",
        "print(src_key_padding_mask)\n",
        "print(tgt_key_padding_mask)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WAOvRJG184Cp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eI-3Spl-0gls",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nRDpuv3ezUT_",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5MYZdOtCzUoi",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MCJHC_udqSww",
        "colab_type": "code",
        "outputId": "0590fa2f-b2c3-485a-c32d-ae674d898949",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 243
        }
      },
      "source": [
        "src_, src_key_padding_mask_ = src[0].to(device), src_key_padding_mask[0].to(device)\n",
        "tgt_, tgt_key_padding_mask_ = tgt[0].to(device), tgt_key_padding_mask[0].to(device)\n",
        "memory_key_padding_mask = src_key_padding_mask_.clone()\n",
        "tgt_inp, tgt_out = tgt_[:, :-1], tgt_[:, 1:]\n",
        "tgt_key_padding_mask_inp = tgt_key_padding_mask_[:, :-1]\n",
        "tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to(device)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "error",
          "ename": "IndexError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-13-77ac83874694>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0mtgt_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtgt_key_padding_mask_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtgt\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtgt_key_padding_mask\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mmemory_key_padding_mask\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msrc_key_padding_mask_\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclone\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m \u001b[0mtgt_inp\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtgt_out\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtgt_\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtgt_\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      5\u001b[0m \u001b[0mtgt_key_padding_mask_inp\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtgt_key_padding_mask_\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0mtgt_mask\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgen_nopeek_mask\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtgt_inp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mIndexError\u001b[0m: too many indices for tensor of dimension 1"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wvaPZRRvuOg6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print(tgt_key_padding_mask_.shape)\n",
        "print(src_.shape)\n",
        "print(tgt_inp.shape)\n",
        "print(src_key_padding_mask_.shape)\n",
        "print(memory_key_padding_mask.shape)\n",
        "print(tgt_mask.shape)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Mlsp_16Etdii",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "output = model(src_, tgt_inp, src_key_padding_mask_, tgt_key_padding_mask_inp, memory_key_padding_mask, tgt_mask)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mTX_erk1t2WD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        " torch.isnan(output).any()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b5_9Agftv8Vz",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "src_.shape"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qbyDvyKzwQyR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "src_[0][0]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xwtgvBOdw6Ft",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "src[0][9]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ddF9kC1Uw0B9",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "lang.tensor2text(tgt[0][9])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Wd-JydO0vuCn",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "for row_i in range(src_.shape[0]):\n",
        "  print(row_i, src_[row_i][1])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CEKnVJ-fp6Eh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "src, src_key_padding_mask = src[0].to(device), src_key_padding_mask[0].to(device) # to device\n",
        "tgt, tgt_key_padding_mask = tgt[0].to(device), tgt_key_padding_mask[0].to(device) # to device\n",
        "memory_key_padding_mask = src_key_padding_mask.clone()\n",
        "tgt_inp, tgt_out = tgt[:, :-1], tgt[:, 1:]\n",
        "tgt_key_padding_mask_inp = tgt_key_padding_mask[:, :-1]\n",
        "tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to(device) # also to device\n",
        "\n",
        "optimizer.zero_grad()\n",
        "output = model(src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask_inp, memory_key_padding_mask, tgt_mask)\n",
        "loss = criterion(rearrange(output, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))\n",
        "loss.backward()\n",
        "optimizer.step()\n",
        "train_loss += loss.item()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MiETl8cebCGq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "output = model(src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask_inp, memory_key_padding_mask, tgt_mask)\n",
        "criterion(rearrange(output, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QrOG2X88b8XL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "src_ = src.permute(1, 0)\n",
        "tgt_ = tgt_inp.permute(1, 0)\n",
        "# src_ = model.pos_enc(model.embed_src(src_) * math.sqrt(model.d_model)) # why every time multiply? just init it somewhere\n",
        "# tgt_ = model.pos_enc(model.embed_tgt(tgt_) * math.sqrt(model.d_model))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vMlRQSWTcT7p",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "src_[:, 0].unsqueeze(0)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KZd99ArrcC57",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print(model.embed_src(src_[:, 0].unsqueeze(0)))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8ei_HEKSelPz",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model.embed_src(torch.Tensor([[lang.word2index['PAD']]]).long().to(device) )"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8UklVda-bndB",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "\n",
        "output = self.transformer(src, tgt, tgt_mask=tgt_mask, src_key_padding_mask=src_key_padding_mask,\n",
        "                          tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=memory_key_padding_mask)\n",
        "output = output.permute(1, 0, 2)\n",
        "return self.fc(output)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "03qw5a46bieg",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "output"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bcwjAlMA1y7Y",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "global_embeddings = nn.Embedding(lang.n_words, config['d_model'])\n",
        "wv = torch.Tensor([word_init_vectors[lang.index2word[idx]] for idx in lang.index2word.keys()])\n",
        "global_embeddings.weight.data = wv\n",
        "\n",
        "\n",
        "criterion = nn.CrossEntropyLoss(ignore_index=lang.word2index['PAD'])\n",
        "\n",
        "d_model = config['d_model']\n",
        "pos_dropout = config['pos_dropout']\n",
        "nhead = config['nhead']\n",
        "num_encoder_layers = config['num_encoder_layers']\n",
        "num_decoder_layers = config['num_decoder_layers']\n",
        "dim_feedforward = config['dim_feedforward']\n",
        "trans_dropout = config['trans_dropout']\n",
        "\n",
        "\n",
        "\n",
        "embed_src = global_embeddings\n",
        "embed_tgt = global_embeddings\n",
        "pos_enc = PositionalEncoding(d_model, pos_dropout, lang.max_seq_length)\n",
        "transformer = nn.Transformer(d_model, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, trans_dropout)\n",
        "fc = nn.Linear(d_model, lang.n_words)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZHYTLN6Q55Bd",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "src, tgt, src_key_padding_mask, tgt_key_padding_mask =  next(iter(train_loader))\n",
        "src, src_key_padding_mask = src[0], src_key_padding_mask[0].to(device) # to device\n",
        "tgt, tgt_key_padding_mask = tgt[0], tgt_key_padding_mask[0].to(device) # to device\n",
        "memory_key_padding_mask = src_key_padding_mask.clone()\n",
        "tgt_inp, tgt_out = tgt[:, :-1], tgt[:, 1:]\n",
        "tgt_key_padding_mask_inp = tgt_key_padding_mask[:, :-1]\n",
        "tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to(device) # also to device\n",
        "\n",
        "print(src.shape)\n",
        "print(tgt.shape)\n",
        "print(src_key_padding_mask.shape)\n",
        "print(tgt_key_padding_mask.shape)\n",
        "print(tgt_inp.shape)\n",
        "print(tgt_out.shape)\n",
        "print(tgt_mask.shape)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BGmF1vKh7dCT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "src = src.permute(1, 0)\n",
        "tgt_inp = tgt_inp.permute(1, 0)\n",
        "src = pos_enc(embed_src(src) * math.sqrt(d_model))\n",
        "tgt_inp = pos_enc(embed_tgt(tgt_inp) * math.sqrt(d_model))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "f1ovGU5J90GA",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print(src.shape, tgt_inp.shape)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xtXVOdxY-lfP",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print(src.shape)\n",
        "print(tgt_inp.shape)\n",
        "print(tgt_mask.shape)\n",
        "print(src_key_padding_mask.shape)\n",
        "print(tgt_key_padding_mask_inp.shape)\n",
        "print(memory_key_padding_mask.shape)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "It8h9lga3Pii",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "output = transformer(src, tgt_inp, tgt_mask=tgt_mask, src_key_padding_mask=src_key_padding_mask,\n",
        "                          tgt_key_padding_mask=tgt_key_padding_mask_inp, memory_key_padding_mask=memory_key_padding_mask)\n",
        "output = output.permute(1, 0, 2)\n",
        "output = fc(output)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "L0-s6h7X9te1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "loss = criterion(rearrange(output, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "a_5b9bxKAJSN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C0LtWb5iW1A8",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import math\n",
        "\n",
        "import torch\n",
        "from torch import nn\n",
        "\n",
        "\n",
        "class LanguageTransformer(nn.Module):\n",
        "    def __init__(self, vocab_size, d_model, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, max_seq_length, pos_dropout, trans_dropout):\n",
        "        super().__init__()\n",
        "        self.d_model = d_model\n",
        "        self.embed_src = nn.Embedding(vocab_size, d_model)\n",
        "        self.embed_tgt = nn.Embedding(vocab_size, d_model)\n",
        "        self.pos_enc = PositionalEncoding(d_model, pos_dropout, max_seq_length)\n",
        "\n",
        "        self.transformer = nn.Transformer(d_model, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, trans_dropout)\n",
        "        self.fc = nn.Linear(d_model, vocab_size)\n",
        "\n",
        "    def forward(self, src, tgt, src_key_padding_mask, tgt_key_padding_mask, memory_key_padding_mask, tgt_mask):\n",
        "        # src = rearrange(src, 'n s -> s n')\n",
        "        src = src.permute(1, 0)\n",
        "        # tgt = rearrange(tgt, 'n t -> t n')\n",
        "        tgt = tgt.permute(1, 0)\n",
        "        src = self.pos_enc(self.embed_src(src) * math.sqrt(self.d_model)) # why every time multiply? just init it somewhere\n",
        "        tgt = self.pos_enc(self.embed_tgt(tgt) * math.sqrt(self.d_model))\n",
        "\n",
        "        output = self.transformer(src, tgt, tgt_mask=tgt_mask, src_key_padding_mask=src_key_padding_mask,\n",
        "                                  tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=memory_key_padding_mask)\n",
        "        # output = rearrange(output, 't n e -> n t e')\n",
        "        output = output.permute(1, 0, 2)\n",
        "        return self.fc(output)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Unw4BYAFZAVX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "        total_step += 1\n",
        "\n",
        "        src, src_key_padding_mask = src[0].to('cuda'), src_key_padding_mask[0].to('cuda')\n",
        "        tgt, tgt_key_padding_mask = tgt[0].to('cuda'), tgt_key_padding_mask[0].to('cuda')\n",
        "\n",
        "        memory_key_padding_mask = src_key_padding_mask.clone()\n",
        "        tgt_inp, tgt_out = tgt[:, :-1], tgt[:, 1:]\n",
        "        tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to('cuda')\n",
        "\n",
        "        optim.zero_grad()\n",
        "        outputs = model(src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask[:, :-1], memory_key_padding_mask, tgt_mask)\n",
        "        loss = criterion(rearrange(outputs, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))\n",
        "\n",
        "        loss.backward()\n",
        "        optim.step_and_update_lr()\n",
        "\n",
        "        total_loss += loss.item()\n",
        "        train_losses.append((step, loss.item()))\n",
        "        pbar.update(1)\n",
        "        if step % print_every == print_every - 1:\n",
        "            pbar.close()\n",
        "            print(f'Epoch [{epoch + 1} / {num_epochs}] \\t Step [{step + 1} / {len(train_loader)}] \\t '\n",
        "                  f'Train Loss: {total_loss / print_every}')\n",
        "            total_loss = 0\n",
        "\n",
        "            pbar = tqdm(total=print_every, leave=False)\n",
        "\n",
        "    pbar.close()\n",
        "    val_loss = validate(valid_loader, model, criterion)\n",
        "    val_losses.append((total_step, val_loss))\n",
        "    if val_loss < lowest_val:\n",
        "        lowest_val = val_loss\n",
        "        torch.save(model, 'output/transformer.pth')\n",
        "    print(f'Val Loss: {val_loss}')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "l5m8S0o8ZAiG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "f91eDhSkXFFD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CWMTJGMVeO0w",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import pickle\n",
        "import random\n",
        "import numpy as np\n",
        "from torch.utils.data import Dataset\n",
        "\n",
        "\n",
        "class ParallelLanguageDataset(Dataset):\n",
        "    def __init__(self, data_path_1, data_path_2, num_tokens, max_seq_length):\n",
        "        self.num_tokens = num_tokens\n",
        "        self.data_1, self.data_2, self.data_lengths = load_data(data_path_1, data_path_2, max_seq_length)\n",
        "\n",
        "        self.batches = gen_batches(num_tokens, self.data_lengths)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        src, src_mask = getitem(idx, self.data_1, self.batches, True)\n",
        "        tgt, tgt_mask = getitem(idx, self.data_2, self.batches, False)\n",
        "\n",
        "        return src, src_mask, tgt, tgt_mask\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.batches)\n",
        "\n",
        "    def shuffle_batches(self):\n",
        "        self.batches = gen_batches(self.num_tokens, self.data_lengths)\n",
        "\n",
        "\n",
        "def gen_batches(num_tokens, data_lengths):\n",
        "    # Shuffle all the indices\n",
        "    for k, v in data_lengths.items():\n",
        "        random.shuffle(v)\n",
        "\n",
        "    batches = []\n",
        "    prev_tokens_in_batch = 1e10\n",
        "    for k in sorted(data_lengths):\n",
        "        v = data_lengths[k]\n",
        "        total_tokens = (k[0] + k[1]) * len(v)\n",
        "\n",
        "        while total_tokens > 0:\n",
        "            tokens_in_batch = min(total_tokens, num_tokens) - min(total_tokens, num_tokens) % (k[0] + k[1])\n",
        "            sentences_in_batch = tokens_in_batch // (k[0] + k[1])\n",
        "\n",
        "            # Combine with previous batch?\n",
        "            if tokens_in_batch + prev_tokens_in_batch <= num_tokens:\n",
        "                batches[-1].extend(v[:sentences_in_batch])\n",
        "                prev_tokens_in_batch += tokens_in_batch\n",
        "            else:\n",
        "                batches.append(v[:sentences_in_batch])\n",
        "                prev_tokens_in_batch = tokens_in_batch\n",
        "            v = v[sentences_in_batch:]\n",
        "\n",
        "            total_tokens = (k[0] + k[1]) * len(v)\n",
        "    return batches\n",
        "\n",
        "\n",
        "def load_data(data_path_1, data_path_2, max_seq_length):\n",
        "    with open(data_path_1, 'rb') as f:\n",
        "        data_1 = pickle.load(f)\n",
        "    with open(data_path_2, 'rb') as f:\n",
        "        data_2 = pickle.load(f)\n",
        "\n",
        "    data_lengths = {}\n",
        "    for i, (str_1, str_2) in enumerate(zip(data_1, data_2)):\n",
        "        if 0 < len(str_1) <= max_seq_length and 0 < len(str_2) <= max_seq_length - 2:\n",
        "            if (len(str_1), len(str_2)) in data_lengths:\n",
        "                data_lengths[(len(str_1), len(str_2))].append(i)\n",
        "            else:\n",
        "                data_lengths[(len(str_1), len(str_2))] = [i]\n",
        "    return data_1, data_2, data_lengths\n",
        "\n",
        "\n",
        "def getitem(idx, data, batches, src):\n",
        "    sentence_indices = batches[idx]\n",
        "    if src:\n",
        "        batch = [data[i] for i in sentence_indices]\n",
        "    else:\n",
        "        batch = [[2] + data[i] + [3] for i in sentence_indices]\n",
        "\n",
        "    seq_length = 0\n",
        "    for sentence in batch:\n",
        "        if len(sentence) > seq_length:\n",
        "            seq_length = len(sentence)\n",
        "\n",
        "    masks = []\n",
        "    for i, sentence in enumerate(batch):\n",
        "        masks.append([False for _ in range(len(sentence))] + [True for _ in range(seq_length - len(sentence))])\n",
        "        batch[i] = sentence + [0 for _ in range(seq_length - len(sentence))]\n",
        "\n",
        "    return np.array(batch), np.array(masks)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e4IQmfdDeCPi",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class ScheduledOptim():\n",
        "    '''A simple wrapper class for learning rate scheduling'''\n",
        "\n",
        "    def __init__(self, optimizer, d_model, n_warmup_steps):\n",
        "        self._optimizer = optimizer\n",
        "        self.n_warmup_steps = n_warmup_steps\n",
        "        self.n_current_steps = 0\n",
        "        self.init_lr = np.power(d_model, -0.5)\n",
        "\n",
        "    def step_and_update_lr(self):\n",
        "        \"Step with the inner optimizer\"\n",
        "        self._update_learning_rate()\n",
        "        self._optimizer.step()\n",
        "\n",
        "    def zero_grad(self):\n",
        "        \"Zero out the gradients by the inner optimizer\"\n",
        "        self._optimizer.zero_grad()\n",
        "\n",
        "    def _get_lr_scale(self):\n",
        "        return np.min([\n",
        "            np.power(self.n_current_steps, -0.5),\n",
        "            np.power(self.n_warmup_steps, -1.5) * self.n_current_steps])\n",
        "\n",
        "    def _update_learning_rate(self):\n",
        "        ''' Learning rate scheduling per step '''\n",
        "\n",
        "        self.n_current_steps += 1\n",
        "        lr = self.init_lr * self._get_lr_scale()\n",
        "\n",
        "        for param_group in self._optimizer.param_groups:\n",
        "            param_group['lr'] = lr\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NV2jC2Isd8H1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def train(train_loader, valid_loader, model, optim, criterion, num_epochs):\n",
        "    print_every = 500\n",
        "    model.train()\n",
        "\n",
        "    lowest_val = 1e9\n",
        "    train_losses = []\n",
        "    val_losses = []\n",
        "    total_step = 0\n",
        "    for epoch in range(num_epochs):\n",
        "        pbar = tqdm(total=print_every, leave=False)\n",
        "        total_loss = 0\n",
        "\n",
        "        train_loader.dataset.shuffle_batches()\n",
        "        for step, (src, src_key_padding_mask, tgt, tgt_key_padding_mask) in enumerate(iter(train_loader)):\n",
        "            total_step += 1\n",
        "\n",
        "            src, src_key_padding_mask = src[0].to('cuda'), src_key_padding_mask[0].to('cuda')\n",
        "            tgt, tgt_key_padding_mask = tgt[0].to('cuda'), tgt_key_padding_mask[0].to('cuda')\n",
        "\n",
        "            memory_key_padding_mask = src_key_padding_mask.clone()\n",
        "            tgt_inp, tgt_out = tgt[:, :-1], tgt[:, 1:]\n",
        "            tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to('cuda')\n",
        "\n",
        "            optim.zero_grad()\n",
        "            outputs = model(src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask[:, :-1], memory_key_padding_mask, tgt_mask)\n",
        "            loss = criterion(rearrange(outputs, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))\n",
        "\n",
        "            loss.backward()\n",
        "            optim.step_and_update_lr()\n",
        "\n",
        "            total_loss += loss.item()\n",
        "            train_losses.append((step, loss.item()))\n",
        "            pbar.update(1)\n",
        "            if step % print_every == print_every - 1:\n",
        "                pbar.close()\n",
        "                print(f'Epoch [{epoch + 1} / {num_epochs}] \\t Step [{step + 1} / {len(train_loader)}] \\t '\n",
        "                      f'Train Loss: {total_loss / print_every}')\n",
        "                total_loss = 0\n",
        "\n",
        "                pbar = tqdm(total=print_every, leave=False)\n",
        "\n",
        "        pbar.close()\n",
        "        val_loss = validate(valid_loader, model, criterion)\n",
        "        val_losses.append((total_step, val_loss))\n",
        "        if val_loss < lowest_val:\n",
        "            lowest_val = val_loss\n",
        "            torch.save(model, 'output/transformer.pth')\n",
        "        print(f'Val Loss: {val_loss}')\n",
        "    return train_losses, val_losses\n",
        "\n",
        "\n",
        "def validate(valid_loader, model, criterion):\n",
        "    pbar = tqdm(total=len(iter(valid_loader)), leave=False)\n",
        "    model.eval()\n",
        "\n",
        "    total_loss = 0\n",
        "    for src, src_key_padding_mask, tgt, tgt_key_padding_mask in iter(valid_loader):\n",
        "        with torch.no_grad():\n",
        "            src, src_key_padding_mask = src[0].to('cuda'), src_key_padding_mask[0].to('cuda')\n",
        "            tgt, tgt_key_padding_mask = tgt[0].to('cuda'), tgt_key_padding_mask[0].to('cuda')\n",
        "            memory_key_padding_mask = src_key_padding_mask.clone()\n",
        "            tgt_inp = tgt[:, :-1]\n",
        "            tgt_out = tgt[:, 1:].contiguous()\n",
        "            tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to('cuda')\n",
        "\n",
        "            outputs = model(src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask[:, :-1], memory_key_padding_mask, tgt_mask)\n",
        "            loss = criterion(rearrange(outputs, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))\n",
        "\n",
        "            total_loss += loss.item()\n",
        "            pbar.update(1)\n",
        "\n",
        "    pbar.close()\n",
        "    model.train()\n",
        "    return total_loss / len(valid_loader)\n",
        "\n",
        "\n",
        "def gen_nopeek_mask(length):\n",
        "    mask = rearrange(torch.triu(torch.ones(length, length)) == 1, 'h w -> w h')\n",
        "    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))\n",
        "    return mask"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MwUsgjTygwm-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "kwargs = {\n",
        "    'num_epochs': 20,\n",
        "    'max_seq_length': 96,\n",
        "    'num_tokens': 2000,\n",
        "    'vocab_size': 10000+4,\n",
        "    'd_model': 512,\n",
        "    'num_encoder_layers': 6,\n",
        "    'num_decoder_layers': 6,\n",
        "    'dim_feedforward': 2048,\n",
        "    'nhead': 8,\n",
        "    'pos_dropout': 0.1,\n",
        "    'trans_dropout': 0.1,\n",
        "    'n_warmup_steps': 4000\n",
        "}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RspHKWOthXrs",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "train_dataset = ParallelLanguageDataset(root + 'data/processed/en/train.pkl',\n",
        "                                        root + 'data/processed/fr/train.pkl',\n",
        "                                        kwargs['num_tokens'], kwargs['max_seq_length'])\n",
        "train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=4, pin_memory=True)\n",
        "valid_dataset = ParallelLanguageDataset(root + 'data/processed/en/val.pkl',\n",
        "                                        root + 'data/processed/fr/val.pkl',\n",
        "                                        kwargs['num_tokens'], kwargs['max_seq_length'])\n",
        "valid_loader = DataLoader(valid_dataset, batch_size=1, shuffle=True, num_workers=4, pin_memory=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jnP2tyOGhgMO",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "src, src_key_padding_mask, tgt, tgt_key_padding_mask = train_dataset[100]\n",
        "print(src.shape)\n",
        "print(tgt.shape)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3S7ltY2jrR2m",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xnA9YQ9Mgvei",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model = LanguageTransformer(kwargs['vocab_size'], kwargs['d_model'], kwargs['nhead'], kwargs['num_encoder_layers'],\n",
        "                            kwargs['num_decoder_layers'], kwargs['dim_feedforward'], kwargs['max_seq_length'],\n",
        "                            kwargs['pos_dropout'], kwargs['trans_dropout']).to('cuda')\n",
        "for p in model.parameters():\n",
        "    if p.dim() > 1:\n",
        "        nn.init.xavier_normal_(p)\n",
        "\n",
        "optim = ScheduledOptim(\n",
        "    Adam(model.parameters(), betas=(0.9, 0.98), eps=1e-09),\n",
        "    kwargs['d_model'], kwargs['n_warmup_steps'])\n",
        "\n",
        "criterion = nn.CrossEntropyLoss(ignore_index=0)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "d3FvWIMWpRpl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print_every = 500\n",
        "model.train()\n",
        "\n",
        "lowest_val = 1e9\n",
        "train_losses = []\n",
        "val_losses = []\n",
        "total_step = 0"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ND4J-jJbxjCa",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# for epoch in range(kwargs['num_epochs']):\n",
        "epoch = 1"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EOkl_I_exku2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "pbar = tqdm(total=print_every, leave=False)\n",
        "total_loss = 0\n",
        "train_loader.dataset.shuffle_batches()\n",
        "\n",
        "# for step, (src, src_key_padding_mask, tgt, tgt_key_padding_mask) in enumerate(iter(train_loader)):\n",
        "step, (src, src_key_padding_mask, tgt, tgt_key_padding_mask) = next(enumerate(iter(train_loader)))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ze5Xan60zjN3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print(src.shape, tgt.shape)\n",
        "\n",
        "total_step += 1\n",
        "\n",
        "src, src_key_padding_mask = src[0].to('cuda'), src_key_padding_mask[0].to('cuda')\n",
        "tgt, tgt_key_padding_mask = tgt[0].to('cuda'), tgt_key_padding_mask[0].to('cuda')\n",
        "\n",
        "print(src.shape)\n",
        "print(tgt.shape)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wUjqMCvt3U_J",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "memory_key_padding_mask = src_key_padding_mask.clone()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G5kF69V45F0X",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "tgt_inp, tgt_out = tgt[:, :-1], tgt[:, 1:] # shift for 1 position"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mOf5mPfJ5HYK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "tgt_inp.shape"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5yJfGknU5Lh1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "tgt_out.shape"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E2aq9-DW57CB",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def gen_nopeek_mask(length):\n",
        "    mask = rearrange(torch.triu(torch.ones(length, length)) == 1, 'h w -> w h')\n",
        "    mask = rearrange(torch.triu(torch.ones(length, length)) == 1, 'h w -> w h')\n",
        "    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))\n",
        "    return mask\n",
        "\n",
        "gen_nopeek_mask(tgt_inp.shape[1])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E2MF0MHn8VKv",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to('cuda')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IFevNzG_yl_W",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "optim.zero_grad()\n",
        "outputs = model(src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask[:, :-1], memory_key_padding_mask, tgt_mask)\n",
        "loss = criterion(rearrange(outputs, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "J2mjXkUB8lxT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "outputs.shape"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K0Wo1KoS8d1G",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "rearrange(outputs, 'b t v -> (b t) v').shape"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XA9eENjWI5c9",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "rearrange(tgt_out, 'b o -> (b o)').shape"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aug6U-WOAzcR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "rearrange(np.array([]))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1CFZCC2f83Nt",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZNdLwMFc8cjO",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "loss.backward()\n",
        "optim.step_and_update_lr()\n",
        "\n",
        "total_loss += loss.item()\n",
        "train_losses.append((step, loss.item()))\n",
        "pbar.update(1)\n",
        "if step % print_every == print_every - 1:\n",
        "    pbar.close()\n",
        "    print(f'Epoch [{epoch + 1} / {num_epochs}] \\t Step [{step + 1} / {len(train_loader)}] \\t '\n",
        "          f'Train Loss: {total_loss / print_every}')\n",
        "    total_loss = 0\n",
        "\n",
        "    pbar = tqdm(total=print_every, leave=False)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QhUT2foyvnFy",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "pbar.close()\n",
        "val_loss = validate(valid_loader, model, criterion)\n",
        "val_losses.append((total_step, val_loss))\n",
        "if val_loss < lowest_val:\n",
        "    lowest_val = val_loss\n",
        "    torch.save(model, 'output/transformer.pth')\n",
        "print(f'Val Loss: {val_loss}')"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}