week08_transformers/homework.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lFqu64qVKd43"
   },
   "source": [
    "## Реализация архитектуры Transformer своими руками \n",
    "---\n",
    "**Разработчик: Денис Кузнеделев**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SeS6GPwzKd48"
   },
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/mryab/dl-hse-ami/blob/main/week08_transformers/homework.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OOqTyWa6Kd49"
   },
   "source": [
    "### Introduction\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rRIqFtZBKd49"
   },
   "source": [
    "В данной МДЗ мы будем реализовывать компоненты архитектуры Трансформер. \n",
    "\n",
    "Начиная со статьи [Attention Is All You Need](https://arxiv.org/abs/1706.03762), трансформеры применяются во всевозможных задачах и установили state-of-the-art на множестве бенчмарков. Первоначально они добились успеха в задачах NLP, но затем были успешно применены и в других областях - обработке сигналов, CV и даже RL.\n",
    "\n",
    "Ниже будет дано описание отдельных модулей и код, который нужно будет написать на основе описания."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "T4TpBpSwKd4-"
   },
   "source": [
    "### Required imports\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kFLKgdFGs-MJ"
   },
   "outputs": [],
   "source": [
    "import math\n",
    "import torch\n",
    "import numpy as np\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from check import check_task_1, check_task_2, check_task_3, check_task_4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "uhyPncgf2GCO"
   },
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "sns.set(font_scale=1.3)\n",
    "sns.set_style(\"darkgrid\", {\"axes.facecolor\": \".95\"})\n",
    "# set fonttype\n",
    "matplotlib.rcParams['pdf.fonttype'] = 42\n",
    "matplotlib.rcParams['ps.fonttype']  = 42\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "U4JtVV1E4qaS"
   },
   "outputs": [],
   "source": [
    "# do not forget to choose CUDA runtime!\n",
    "device = torch.device('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "LmwveSF-Kd5B"
   },
   "outputs": [],
   "source": [
    "# fix seeds for reproducbility\n",
    "np.random.seed(42);\n",
    "torch.random.manual_seed(42);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KK7UsmLdKd5C"
   },
   "source": [
    "### Scaled Dot Product Attention\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YcjXsPqgKd5C"
   },
   "source": [
    "В оригинальной работе [Attention Is All You Need](https://arxiv.org/abs/1706.03762) в качестве механизма внимания \n",
    "был использован scaled dot product attention - нормализованное скалярное произведение между key и query. \n",
    "На вход подается набор запросов $Q\\in\\mathbb{R}^{L\\times d_k}$, ключей $K\\in\\mathbb{R}^{L\\times d_k}$  и \n",
    "значений  $V\\in\\mathbb{R}^{L\\times d_v}$, где $L$ - длина последовательности, а $d_k, d_v$ - размерности query/key и value соотвественно.\n",
    "Значение attention от элемента $i$ на элемент $j$ зависит от похожести query $q_i$ и key $k_j$. Attention определяется по следующей формуле:\n",
    "\n",
    "$$ \\text{Attention}(Q,K,V)=\\text{softmax}\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right)V $$\n",
    "\n",
    "Произведение матриц $Q K^{T}$ составлено из всех попарных скалярных произведений ключей и значений.\n",
    "\n",
    "\n",
    "<p align=\"center\">\n",
    "  <img src=\"scaled_dot_product_attn.png\" width=\"300\"/>\n",
    "</p>\n",
    "\n",
    "Дополнительно в операцию может входить бинарная маска $M \\in \\{0, 1\\}$, зануляющая некоторые элементы в матрице attention, если по некоторой причине мы не хотим, чтобы токен $i$ из query взаимодействовал c $j$ из key. Это может быть полезно при генерации последовательностей, когда мы не хотим, чтобы данный токен смотрел вперед, на еще не сгенерированные элементы.\n",
    "\n",
    "\n",
    "Деление на $\\sqrt{d_k}$ необходимо, что выход операции сохранял дисперсию распределения, т.к:\n",
    "$$\n",
    "q_i \\sim \\mathcal{N}(0, 1), k_i \\sim \\mathcal{N}(0, 1) \\to \\text{Var}\\left(\\sum_{i=1}^{d_k} q_i\\cdot k_i\\right) = d_k\n",
    "$$\n",
    "\n",
    "**Задание 1 (0.1 балла):**. Реализуйте операцию scaled dot product."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rZut6BBmKd5D"
   },
   "outputs": [],
   "source": [
    "def scaled_softmax_attention(query, key, value):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        query: torch.Tensor (..., L, D)\n",
    "        key: torch.Tensor (..., L, D)\n",
    "        value: torch.Tensor (..., L, D)\n",
    "    Returns:\n",
    "        res: torch.Tensor (..., L, D), output of the attention layer (\\softmax(Q K^T / d) V\n",
    "        attention: torch.Tensor (..., L, L), attention weights (\\softmax(Q K^T / d))\n",
    "\n",
    "    L is the length of sequence, D is the embedding dimension\n",
    "    \"\"\"\n",
    "\n",
    "    return res, attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9MXJ925kKd5E"
   },
   "outputs": [],
   "source": [
    "query = torch.tensor([[ 0.3367,  0.1288,  0.2345,  0.2303],\n",
    "                      [-1.1229, -0.1863,  2.2082, -0.6380],\n",
    "                      [ 0.4617,  0.2674,  0.5349,  0.8094]])\n",
    "key   = torch.tensor([[ 1.1103, -1.6898, -0.9890,  0.9580],\n",
    "                      [ 1.3221,  0.8172, -0.7658, -0.7506],\n",
    "                      [ 1.3525,  0.6863, -0.3278,  0.7950]])\n",
    "value = torch.tensor([[ 0.2815,  0.0562,  0.5227, -0.2384],\n",
    "                      [-0.0499,  0.5263, -0.0085,  0.7291],\n",
    "                      [ 0.1331,  0.8640, -1.0157, -0.8887]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iuaz3cVwKd5F"
   },
   "outputs": [],
   "source": [
    "res, attn = scaled_softmax_attention(query, key, value)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hyOf98LTKd5G"
   },
   "source": [
    "**Sanity check**. \n",
    "\n",
    "Матрица attn должна быть размера $(L, L)$ и столбцы должны суммироваться в 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ym6UvVHkKd5G"
   },
   "outputs": [],
   "source": [
    "attn.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QkERgNoSKd5H"
   },
   "outputs": [],
   "source": [
    "attn.sum(dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "F4F9WLU5Kd5H"
   },
   "outputs": [],
   "source": [
    "check_task_1(res, attn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-Hcl0gv6Kd5I"
   },
   "source": [
    "### Multihead Attention Layer\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gX788mDyKd5I"
   },
   "source": [
    "Scaled dot product attention задает правило, по которому элементы последовательности взаимодейтсвуют друг с другом. Но может быть полезно задавать несколько различных правил взаимодействия. Поэтому последовательности $Q, K, V$ разбиваются на $h$\n",
    "частей вдоль размерности эмбединнга, и для каждой из них независимо считаем attention, а затем конкатенируем результат. К сконкатенированному результату применяется линейное преобразование $W_O$:\n",
    "\n",
    "$$\n",
    "\\begin{split}\\begin{split}\n",
    "    \\text{Multihead}(Q,K,V) & = \\text{Concat}(\\text{head}_1,...,\\text{head}_h)W^{O}\\\\\n",
    "    \\text{where } \\text{head}_i & = \\text{Attention}(QW_i^Q,KW_i^K, VW_i^V)\n",
    "\\end{split}\\end{split}\n",
    "$$\n",
    "\n",
    "Обучаемыми параметрами являются матрицы проекции $W_Q, W_K, W_V$ и $W_O$.\n",
    "Ниже приведен граф вычислений:\n",
    "<p align=\"center\">\n",
    "  <img src=\"multihead_attention.png\" width=\"300\"/>\n",
    "</p>\n",
    "\n",
    "**Задание 2 (0.1 балла):**. Реализуйте класс  `MultiheadAttention`, реализующий операции, описанные выше.\n",
    "\n",
    "<span style=\"color:red\">Замечание</span>.\n",
    "\n",
    "Обратите внимание, что слои `q_proj`, `k_proj`, `v_proj` создают последовательность для всех голов одноновременно, то есть их выходная размерность должна быть $\\mathbb{R}^{d_{embed}}$, а не $\\mathbb{R}^{d_{embed} / h}$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "p-YL41gHYBsz"
   },
   "outputs": [],
   "source": [
    "from util import hardcode_parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1xgxpmewKd5I"
   },
   "outputs": [],
   "source": [
    "class MultiheadAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, embed_dim, num_heads):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            embed_dim: dimensionality of embedding (total)\n",
    "            num_heads: number of heads (must divide embed_dim)\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        assert embed_dim % num_heads == 0, \"Embedding dimension must be 0 modulo number of heads.\"\n",
    "\n",
    "        self.embed_dim = embed_dim\n",
    "        self.num_heads = num_heads\n",
    "        self.head_dim = embed_dim // num_heads\n",
    "\n",
    "        self.q_proj = # TODO\n",
    "        self.k_proj = # TODO\n",
    "        self.v_proj = # TODO\n",
    "        self.o_proj = # TODO\n",
    "\n",
    "        self._reset_parameters()\n",
    "\n",
    "    # original implementation uses this initialization\n",
    "    def _reset_parameters(self):\n",
    "        for layer in self.modules():\n",
    "            if isinstance(layer, nn.Linear):\n",
    "                nn.init.xavier_uniform_(layer.weight)\n",
    "                if layer.bias is not None:\n",
    "                    layer.bias.data.fill_(0)\n",
    "\n",
    "    def forward(self, x, return_attention=False):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            x: torch.Tensor (B, L, D)\n",
    "            return_attention: If specified, returns attention along with outputs\n",
    "        Returns:\n",
    "            outputs: torch.Tensor (B, L, D)\n",
    "            attention: Optional[torch.Tensor] (B, num_heads, L, L)\n",
    "\n",
    "        B is batch size, L is the length of sequence, D is the embedding dimension\n",
    "        \"\"\"\n",
    "        # TODO\n",
    "        outputs, attention = None, None\n",
    "\n",
    "        if return_attention:\n",
    "            return outputs, attention\n",
    "        else:\n",
    "            return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "U9GIkvFeKd5J"
   },
   "outputs": [],
   "source": [
    "multihead_attention = MultiheadAttention(4, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UyYLmqIHXq1e"
   },
   "outputs": [],
   "source": [
    "hardcode_parameters(multihead_attention)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fGC2CGq1Kd5J"
   },
   "outputs": [],
   "source": [
    "inputs  = torch.stack([torch.cos(i * torch.ones(4)) for i in range(3)])[None, ...]\n",
    "\n",
    "with torch.no_grad():\n",
    "    outputs = multihead_attention(inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jomSUohNKd5J"
   },
   "outputs": [],
   "source": [
    "check_task_2(outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YC3Au9FqKd5K"
   },
   "source": [
    "### Encoder Block\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OBXYUCM0Kd5K"
   },
   "source": [
    "Архитектура трансформера в оригинальной статье состоит из последовательности блоков энкодера и декодера. \n",
    "\n",
    "<p align=\"center\">\n",
    "  <img src=\"transformer_architecture.png\" width=\"300\" />\n",
    "</p>\n",
    "\n",
    "В данной МДЗ необходимо будет реализовать только энкодер. \n",
    "\n",
    "Блок энкодера состоит из операции `MultiheadAttention` и применения `FeedForward` сети к каждому токену по отдельности. \n",
    "\n",
    "К выходу `MultiheadAttention` и `FeedForward`\n",
    "прибавляется skip connection, и к полученной сумме применяется `LayerNormalization`.\n",
    "\n",
    "В качестве `FeedForward` сети берется простая двуслойная сеть с некоторой активацией (обычно  `ReLU` или `GELU`).\n",
    "\n",
    "Таким образом, энкодер выполняет следующее:\n",
    "$$\n",
    "x = \\text{LayerNorm}(x+\\text{MultiheadAttention}(x,x,x)) \n",
    "$$\n",
    "$$\n",
    "\\begin{split}\\begin{split}\n",
    "    \\text{FFN}(x) & = \\mathrm{Act}(x W_1 + b_1) W_2 + b_2\\\\\n",
    "    x & = \\text{LayerNorm}(x + \\text{FFN}(x))\n",
    "\\end{split}\\end{split}\n",
    "$$\n",
    "\n",
    "В целях регуляризации на выход `MultiheadAttention` и `FeedForward`, но перед `LayerNorm` можно накинуть `Dropout`.\n",
    "\n",
    "**Задание 3 (0.1 балла):**. Реализуйте класс  `EncoderBlock`, реализующий операции, описанные выше."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cC6UzD5FKd5L"
   },
   "outputs": [],
   "source": [
    "class EncoderBlock(nn.Module):\n",
    "\n",
    "    def __init__(self, embed_dim, num_heads, feedforward_dim, activation=nn.ReLU, dropout=0.0):\n",
    "        \"\"\"\n",
    "        Inputs:\n",
    "            embed_dim - Dimensionality of the input\n",
    "            num_heads - Number of heads to use in the attention block\n",
    "            feedforward_dim - Dimensionality of the hidden layer in the MLP\n",
    "            activation - activation function in FFN\n",
    "            dropout - Dropout probability to use in the dropout layers\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "\n",
    "        # TODO\n",
    "\n",
    "    def forward(self, x, return_attention=False):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            x: torch.Tensor (B, L, D)\n",
    "        Returns:\n",
    "            outputs: torch.Tensor (B, L, D)\n",
    "            attention: Optional[torch.Tensor] (B, num_heads, L, L)\n",
    "        \"\"\"\n",
    "        outputs, attention = None\n",
    "\n",
    "        if return_attention:\n",
    "            return outputs, attention\n",
    "        else:\n",
    "            return outputs\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lPurRfeNKd5L"
   },
   "source": [
    "### Positional Encoding\n",
    "---\n",
    "\n",
    "Описанная выше конструкция очень гибкая и универсальная в плане возможности преобразования последовательностей. \n",
    "Но есть один нюанс: абсолютное положение токенов никак не определено в текущей форме, и операция `EncoderBlock`\n",
    "обладает перестановочной симметрией. То есть, если переставить токены в последовательности, то выход от переставленной последовательности будет таким же, как если прогнать исходную последовательность, а затем ее переставить.\n",
    "\n",
    "Сгенерируем рандомную последовательность."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3qN7EmP8Kd5L"
   },
   "outputs": [],
   "source": [
    "encoder_block = EncoderBlock(embed_dim=24, num_heads=3, feedforward_dim=24 * 4, dropout=0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "N0avWv-zKd5M"
   },
   "outputs": [],
   "source": [
    "inputs = torch.randn(1, 16, 24)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hCck39HfKd5M"
   },
   "source": [
    "Выход от исходной последовательности"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "E3U5dqsNKd5M"
   },
   "outputs": [],
   "source": [
    "outputs = encoder_block(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rNwGn54zKd5N"
   },
   "source": [
    "Сгенерируем случайную перестановку"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CQ4VtFJGKd5O"
   },
   "outputs": [],
   "source": [
    "ids_perm = torch.randperm(inputs.size(1))\n",
    "shuffled_inputs = inputs[:, ids_perm, :]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UOC4e3OgKd5P"
   },
   "source": [
    "Посчитаем выход для переставленной последовательности"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IP_HmExhKd5P"
   },
   "outputs": [],
   "source": [
    "shuffled_outputs = encoder_block(shuffled_inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ncG58JOXKd5Q"
   },
   "source": [
    "И сравним переставленный выход исходной последовательности с выходом переставленной"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oT-b6FHPKd5R"
   },
   "outputs": [],
   "source": [
    "torch.allclose(outputs[:, ids_perm, :], shuffled_outputs, atol=1e-4, rtol=1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wT_qun9HKd5R"
   },
   "source": [
    "Вуаля! Совпадают.\n",
    "\n",
    "Но во многих задачах важен порядок, и чтобы его каким-то образом учесть добавляют так называемый PositionalEncoding, который явно задает информацию о положении токена в последовательности. Он может быть как обучаемым, так и зафиксированным.\n",
    "В оригинальной работе Attention is all you need был выбран следующий энкодинг и синусов и косинусов разных частот:\n",
    "\n",
    "$$\n",
    "\\begin{split}PE_{(pos,i)} = \\begin{cases}\n",
    "    \\sin\\left(\\frac{pos}{10000^{i/d_{\\text{model}}}}\\right) & \\text{if}\\hspace{3mm} i \\text{ mod } 2=0\\\\\n",
    "    \\cos\\left(\\frac{pos}{10000^{(i-1)/d_{\\text{model}}}}\\right) & \\text{otherwise}\\\\\n",
    "\\end{cases}\\end{split}\n",
    "$$\n",
    "\n",
    "$PE_{(pos,i)}$ обозначает позиционный энкодинг токена в позиции $pos$, а $i$ нумерует размерность эмбеддинга.\n",
    "\n",
    "**Задание 4 (0.1 балла):**. Реализуйте класс  `PositionalEmbedding`, добавляющий позиционный энкодинг определенный выше к входной последовательности."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fJJhICZX5x7R"
   },
   "outputs": [],
   "source": [
    "class PositionalEncoding(nn.Module):\n",
    "\n",
    "    def __init__(self, embed_dim, max_len: int = 5000):\n",
    "        \"\"\"\n",
    "        Inputs\n",
    "            embed_dim - Hidden dimensionality of the input.\n",
    "            max_len - Maximum length of a sequence to expect.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        # TODO\n",
    "        pe =  # here should be a tensor of size (1, max_len, embed_dim), dummy dimension is needed for proper addition\n",
    "\n",
    "        # register_buffer => Tensor which is not a parameter, but should be part of the modules state.\n",
    "        # Used for tensors that need to be on the same device as the module.\n",
    "        # persistent=False tells PyTorch to not add the buffer to the state dict (e.g. when we save the model)\n",
    "        self.register_buffer('pe', pe, persistent=False)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = None\n",
    "        # TODO\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-ZeXXKasKd5S"
   },
   "source": [
    "Визуализируем энкодинг"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XIbOv_QxKd5S"
   },
   "outputs": [],
   "source": [
    "positional_encoding = PositionalEncoding(embed_dim=64, max_len=128)\n",
    "pe = positional_encoding.pe[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "N2-_ifqPKd5T"
   },
   "outputs": [],
   "source": [
    "check_task_3(pe)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sPr_XImKKd5T"
   },
   "outputs": [],
   "source": [
    "pe = positional_encoding.pe.numpy()[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "o2BVbmUiKd5T"
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(12, 5))\n",
    "\n",
    "im = ax.imshow(pe.T)\n",
    "ax.grid(False)\n",
    "\n",
    "ax.set_ylabel(r'Embed dim')\n",
    "ax.set_xlabel(r'Position in sequence')\n",
    "\n",
    "fig.colorbar(im, ax=ax);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Implementation check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "check_task_4(positional_encoding)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Zc_OXXnoKd5U"
   },
   "source": [
    "### Transformer \n",
    "---\n",
    "\n",
    "Теперь у нас есть все необходимые компоненты, чтобы собрать трансформер своими руками.\n",
    "Сам по себе трансформер может выполнять великое множество задач, но в рамках данной МДЗ мы ограничимся задачей классификации последовательностей.\n",
    "\n",
    "Архитектура состоит из следующих компонент:\n",
    "- Линейное преобразование входной последовательности : $\\mathbb{R}^{d_{in}} \\rightarrow \\mathbb{R}^{d_{embed}}$, примененное поэлементно к каждому токену\n",
    "- Один или несколько блоков `EncoderBlock`\n",
    "- `PositionalEmbedding` \n",
    "- Для задачи классификации создается специальный токен `[CLS]`, который прибавляется в начало (или конец последовательности)\n",
    "- `[CLS]` токен, пропущенный через последовательность, подается на вход классификатора (скажем, линейного слоя $\\mathbb{R}^{d_{embed}} \\rightarrow \\mathbb{R}^{|C|}$, $|C|$ - число классов).\n",
    "\n",
    "**Задание 5 (0.2 балла):** Реализуйте класс  `TransformerForSequenceClassification`, принимающий на вход последовательность и предсказывающий ее класс."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Tk0AH3cW5p7h"
   },
   "outputs": [],
   "source": [
    "class TransformerForSequenceClassification(nn.Module):\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_dim: int,\n",
    "        embed_dim: int,\n",
    "        num_classes: int,\n",
    "        num_heads: int,\n",
    "        feedforward_dim: int,\n",
    "        num_layers: int,\n",
    "        activation = nn.GELU,\n",
    "        max_len: int = 5000,\n",
    "        dropout: float = 0.0\n",
    "    ):\n",
    "        super().__init__()\n",
    "        # define layers\n",
    "        self.cls_token = None # TODO create vector of size (embed_dim,) from N(0, 1)\n",
    "        self.input_embedding = None # TODO\n",
    "        self.positional_encoding = None # TODO\n",
    "\n",
    "        encoder_blocks = None # TODO\n",
    "        self.encoder = None # TODO\n",
    "\n",
    "        self.classifier = None # TODO\n",
    "\n",
    "    def forward_attention(self, x):\n",
    "        \"\"\"\n",
    "        This method returns attention maps from all encoder blocks\n",
    "\n",
    "        Args:\n",
    "            x: torch.Tensor (B, L, |V|)\n",
    "        Returns:\n",
    "            attn: List[torch.Tensor] (B, num_heads, L+1, L+1) x num_layers\n",
    "        \"\"\"\n",
    "        pass # TODO\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            x: torch.Tensor (B, L, |V|)\n",
    "        Returns:\n",
    "            x: torch.Tensor (B, |C|)\n",
    "        \"\"\"\n",
    "        pass # TODO\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_0NEZ5okKd5U"
   },
   "source": [
    "### Data\n",
    "---\n",
    "В данной МДЗ в качестве задачи, которой мы будем обучать трансформер, является определение того, является ли строка палиндромом. То есть тождественна ли строка, написанная задом наперед, исходной строке.\n",
    "\n",
    "$$\n",
    "a b f g a a g f b a \\quad \\mathrm{\\ is \\ a \\ palindrome}\n",
    "$$\n",
    "\n",
    "Последовательности генерируются случайным образом из некоторого слова размера `vocab_length`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Xd05RqIu8C09"
   },
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "from torch.utils.data import DataLoader\n",
    "from dataset import PalindromeDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bImOPPIB8jE3"
   },
   "outputs": [],
   "source": [
    "vocab_length = 33\n",
    "sequence_length = 256\n",
    "\n",
    "make_dataset = partial(\n",
    "    PalindromeDataset,\n",
    "    vocab_length=vocab_length,\n",
    "    sequence_length=sequence_length\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zYeYo7fxKd5W"
   },
   "source": [
    "Созданим обучающую и тестовую выборку"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gb6VIP0y8e4C"
   },
   "outputs": [],
   "source": [
    "train_dataset = make_dataset(size=50000)\n",
    "val_dataset  = make_dataset(size=10000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "s411xpdoKd5W"
   },
   "outputs": [],
   "source": [
    "# training hyperparameters\n",
    "batch_size = 128\n",
    "num_workers = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TDEUki4T8Ysv"
   },
   "outputs": [],
   "source": [
    "train_loader = DataLoader(\n",
    "    train_dataset,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True,\n",
    "    drop_last=True,\n",
    "    num_workers=num_workers,\n",
    "    pin_memory=True\n",
    ")\n",
    "val_loader  = DataLoader(\n",
    "    val_dataset,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=False,\n",
    "    num_workers=num_workers\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dBcLdqwgKd5X"
   },
   "outputs": [],
   "source": [
    "# model hyperparams\n",
    "embed_dim = # TODO\n",
    "num_heads = # TODO\n",
    "feedforward_dim = # TODO\n",
    "num_layers =  # TODO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BeWqciEx_MCh"
   },
   "outputs": [],
   "source": [
    "model = TransformerForSequenceClassification(\n",
    "    num_classes=1,\n",
    "    input_dim=vocab_length,\n",
    "    embed_dim=embed_dim,\n",
    "    num_heads=num_heads,\n",
    "    feedforward_dim=feedforward_dim,\n",
    "    activation=nn.ReLU,\n",
    "    num_layers=num_layers\n",
    ")\n",
    "# put model on device\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NmiEkjNIKd5Y"
   },
   "source": [
    "Есть мнение, что если обучать сразу Transformer без предварительного 'разогрева', то обучение может разойтись, или долго и тяжело сходиться к оптимуму. На рисунке ниже приведено поведение кривых обучения без и с разогревом (синяя, без). \n",
    "\n",
    "<p align=\"center\">\n",
    "  <img src=\"warmup_loss_plot.png\" width=\"500\" />\n",
    "</p>\n",
    "\n",
    " Поэтому на практике при обучение трансформеров принято использовать расписание с 'разогревом', когда первое время, заданное количество шагов оптимизатора или эпох, `learning rate` сначала линейно растет, а затем затухает. Ниже мы будем использовать косинусное затухание."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nZYua9IACNFR"
   },
   "outputs": [],
   "source": [
    "from util import CosineAnnealingWithWarmupLR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yCGWcUF3Kd5Y"
   },
   "source": [
    "Визуализируем расписание"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Sa_MESRRKd5Z"
   },
   "outputs": [],
   "source": [
    "# Needed for initializing the lr scheduler\n",
    "p = nn.Parameter(torch.empty(4,4))\n",
    "optimizer = torch.optim.Adam([p], lr=1e-3)\n",
    "lr_scheduler = CosineAnnealingWithWarmupLR(optimizer=optimizer, warmup_steps=100, max_steps=2000)\n",
    "\n",
    "# Plotting\n",
    "fig, ax = plt.subplots(figsize=(10, 4))\n",
    "steps = range(2000)\n",
    "\n",
    "ax.plot(steps, [lr_scheduler.get_lr_factor(e) for e in steps])\n",
    "ax.set_ylabel(\"Learning rate factor\")\n",
    "ax.set_xlabel(\"Optimizer steps\")\n",
    "ax.set_title(\"Cosine Warm-up Learning Rate Scheduler\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sngnTHDVKd5Z"
   },
   "source": [
    "**Задание 6 (0.1 балла):** обучите модель-классификатор палиндромов."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fXzhBPgeKd5Z"
   },
   "outputs": [],
   "source": [
    "from util import train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1qoikuHcKd5a"
   },
   "source": [
    "Зададим параметры обучения и оптимизатора"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iKu7i5m6Kd5a"
   },
   "outputs": [],
   "source": [
    "num_epochs = # TODO\n",
    "warmup_steps = # TODO\n",
    "lr = 5e-5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "w6bWXvqoEZRE"
   },
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=0.0)\n",
    "\n",
    "scheduler = CosineAnnealingWithWarmupLR(\n",
    "    optimizer,\n",
    "    warmup_steps=warmup_steps,\n",
    "    max_steps=int(num_epochs * len(train_loader))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ra0fG49gFfVP"
   },
   "outputs": [],
   "source": [
    "train(\n",
    "    model,\n",
    "    num_epochs=num_epochs,\n",
    "    optimizer=optimizer,\n",
    "    scheduler=scheduler,\n",
    "    train_loader=train_loader,\n",
    "    val_loader=val_loader,\n",
    "    device=device\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "X50mAkjCKd5b"
   },
   "source": [
    "Если все было реализовано верно, то модель должна выдавать качество > 95%"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "K8NDS7jVKd5b"
   },
   "source": [
    "### Linformer\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IzP9aS_FKd5b"
   },
   "source": [
    "Механизм Attention очень мощный, но имеет существенный недостаток при работе с длинными последовательностями.\n",
    "\n",
    "Как нетрудно заметить, операция $Q K^{T}$ квадратична по длине последовательности по сложности вычислений и занимаемой памяти.\n",
    "\n",
    "Что может быть критично, если контекст важный для данного токена имеет длину более 1000 токенов. Было предложено множество подходов по замене или приближению операции  Attention на нечто имеющее субквадратичную сложность по длине последовательности. Для желающих узнать больше есть [хороший обзор по теме](https://lilianweng.github.io/posts/2020-04-07-the-transformer-family/), а в данной МДЗ будет предложено реализовать [Linformer](https://arxiv.org/abs/2006.04768).\n",
    "\n",
    "Возьмем ранее обученную модель и ее матрицу `Attention`. Одной из основных характеристик матрицы является ее ранг и разложение по сингулярным числам. Если большая часть массы собственных/сингулярных значений концентрируется на первых нескольких собственных векторах, то эффективно матрица является отображением в пространство низкой размерности. \n",
    "\n",
    "Здесь можно вспомнить принцип работы `PCA`, где уменьшение размерности достигается за счет проекция на несколько первых сингулярных векторов. Построим нормализованную кумулятивную сумму первых $k$-собственных значений (известную в литературе как Explained Variance)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "q4wO95xkKd5c"
   },
   "outputs": [],
   "source": [
    "sequence, label = val_dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "13gi_b2JKd5c"
   },
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    # we take attention only from the last layer\n",
    "    attn = model.forward_attention(sequence[None, ...].to(device))[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QazU0RhiKd5c"
   },
   "outputs": [],
   "source": [
    "# select number such that num_heads divides it\n",
    "heads_per_row = # TODO\n",
    "\n",
    "fig, ax = plt.subplots(\n",
    "    nrows=num_heads // heads_per_row,\n",
    "    ncols=heads_per_row,\n",
    "    figsize=(heads_per_row * 5, (num_heads // heads_per_row) * 5),\n",
    "    squeeze=False\n",
    ")\n",
    "\n",
    "for head_idx in range(num_heads):\n",
    "    row, col = head_idx // heads_per_row, head_idx % heads_per_row\n",
    "    sing_vals = torch.linalg.svdvals(attn[0, head_idx, 1:, 1:]).cpu().numpy()\n",
    "    expl_var  = np.cumsum(sing_vals ** 2, axis=0)\n",
    "    expl_var /= expl_var[-1]\n",
    "    ax[row, col].plot(expl_var)\n",
    "    ax[row, col].axhline(0.99, linestyle='--', color='red')\n",
    "    ax[row, col].set_xlabel(r'$k$')\n",
    "    ax[row, col].set_ylabel(r'Explained variance')\n",
    "    ax[row, col].set_xscale('log')\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "w4Sg09YHKd5c"
   },
   "source": [
    "Что вы наблюдаете?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WFtesqUwKd5d"
   },
   "source": [
    "Отсюда и берет начало идея `Linformer`. \n",
    "\n",
    "Так как матрица Attention низкоранговая, вместо того чтобы считать большую квадратную матрицу $L \\times L$, предлагается вычислять прямоугольную\n",
    "$L \\times k$, где $k \\ll L$. \n",
    "\n",
    "Последовательности $K$ и $V$ (после прогонки через $W_K$, $W_V$) отображаются вдоль оси, отвечающей длине последовательности из $L$ - мерного пространства в $k$ - мерное. И только затем вычисляется скалярное произведение уже с $k$ ключами для всех query.\n",
    "\n",
    "<p align=\"center\">\n",
    "  <img src=\"LinformerAttention.png\" width=\"300\" />\n",
    "</p>\n",
    "\n",
    "Вычислительная сложность уменьшается с $\\mathcal{O}(L^2)$ до $\\mathcal{O}(L k)$.\n",
    "\n",
    "В конечном итоге операция `LinformerAttention` имеет следующий вид:\n",
    "\n",
    "$$ \\text{LinformerAttention}(Q,K,V)=\\text{softmax}\\left(\\frac{Q (E K)^T}{\\sqrt{d_k}}\\right) F V $$ \n",
    "\n",
    "Выше $E$ и $F$ - обучаемые матрицы проекции.\n",
    "\n",
    "<span style=\"color:red\">Замечание</span>.\n",
    "\n",
    "Описанный подход ограничивает применимость `Linformer` на последовательности фиксированной длины.\n",
    "\n",
    "К счастью, в текущей задаче с такими мы и работаем.\n",
    "\n",
    "**Задание 7 (0.2 балла):**. Реализуйте класс  `LinformerAttention` - реализующий операции описанные выше, и\n",
    "`LinformerBlock`, `LinformerForSequenceClassification`, повторяющих функционал `EncoderBlock`, `TransformerForSequenceClassification`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iSYmdBltKd5d"
   },
   "outputs": [],
   "source": [
    "class LinformerAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, embed_dim, num_heads, sequence_length, proj_dim):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            embed_dim: dimensionality of embedding (total)\n",
    "            num_heads: number of heads (must divide embed_dim)\n",
    "            sequence_length: length of the input sequence\n",
    "            proj_dim: length of  the projected sequence\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        assert embed_dim % num_heads == 0, \"Embedding dimension must be 0 modulo number of heads.\"\n",
    "\n",
    "        # TODO\n",
    "\n",
    "        self.embed_dim = embed_dim\n",
    "        self.num_heads = num_heads\n",
    "        self.head_dim = embed_dim // num_heads\n",
    "\n",
    "        self.q_proj = # TODO\n",
    "        self.k_proj = # TODO\n",
    "        self.v_proj = # TODO\n",
    "        self.o_proj = # TODO\n",
    "        # projections on sequence length\n",
    "        self.e_proj = # TODO\n",
    "        self.f_proj = # TODO\n",
    "\n",
    "        self._reset_parameters()\n",
    "\n",
    "    def _reset_parameters(self):\n",
    "        for layer in self.modules():\n",
    "            if isinstance(layer, nn.Linear):\n",
    "                nn.init.xavier_uniform_(layer.weight)\n",
    "                if layer.bias is not None:\n",
    "                    layer.bias.data.fill_(0)\n",
    "\n",
    "    def forward(self, x, return_attention=False):\n",
    "        outputs, attention = None, None\n",
    "\n",
    "        # TODO\n",
    "\n",
    "        if return_attention:\n",
    "            return outputs, attention\n",
    "        else:\n",
    "            return outputs\n",
    "\n",
    "\n",
    "class LinformerBlock(EncoderBlock):\n",
    "\n",
    "    def __init__(self, embed_dim, num_heads, feedforward_dim, sequence_length, proj_dim, activation=nn.ReLU, dropout=0.0):\n",
    "        \"\"\"\n",
    "        Inputs:\n",
    "            embed_dim - Dimensionality of the input\n",
    "            num_heads - Number of heads to use in the attention block\n",
    "            feedforward_dim - Dimensionality of the hidden layer in the MLP\n",
    "            activation - activation function in FFN\n",
    "            sequence_length: length of the input sequence\n",
    "            proj_dim: length of  the projected sequence\n",
    "            dropout - Dropout probability to use in the dropout layers\n",
    "        \"\"\"\n",
    "        super(EncoderBlock, self).__init__()\n",
    "\n",
    "        # TODO\n",
    "\n",
    "\n",
    "class LinformerForSequenceClassification(TransformerForSequenceClassification):\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_dim: int,\n",
    "        embed_dim: int,\n",
    "        num_classes: int,\n",
    "        num_heads: int,\n",
    "        feedforward_dim: int,\n",
    "        num_layers: int,\n",
    "        sequence_length: int,\n",
    "        proj_dim: int,\n",
    "        activation = nn.GELU,\n",
    "        max_len: int = 5000,\n",
    "        dropout: float = 0.0\n",
    "    ):\n",
    "        super(TransformerForSequenceClassification, self).__init__()\n",
    "        # define layers\n",
    "        self.cls_token = None # TODO\n",
    "        self.input_embedding =  None # TODO\n",
    "        self.positional_encoding = None # TODO\n",
    "\n",
    "        encoder_blocks = None # TODO\n",
    "        self.encoder = None # TODO\n",
    "\n",
    "        self.classifier =  None # TODO"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "D-zyJel2Kd5d"
   },
   "source": [
    "Сравним зависимость расхода памяти от длины последовательности\n",
    "для `Transformer` и `Linformer`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "V4J7XHUPKd5e"
   },
   "outputs": [],
   "source": [
    "embed_dim = 8\n",
    "sequence_lengths = [64, 128, 256, 512, 1024, 2048, 4096]\n",
    "\n",
    "mem_usages = {'transformer': [], 'linformer': []}\n",
    "\n",
    "for sequence_length in sequence_lengths:\n",
    "    # generate input\n",
    "    input = torch.randn(32, sequence_length, embed_dim, device=device)\n",
    "    # define transformer block\n",
    "    transformer_attn = MultiheadAttention(\n",
    "        embed_dim=embed_dim,\n",
    "        num_heads=1\n",
    "    ).to(device)\n",
    "    # define linformer block\n",
    "    linformer_attn = LinformerAttention(\n",
    "        embed_dim=embed_dim,\n",
    "        num_heads=1,\n",
    "        sequence_length=sequence_length,\n",
    "        proj_dim=8\n",
    "    ).to(device)\n",
    "\n",
    "    # get memory usage for model\n",
    "    for model_name, attn_layer in zip(\n",
    "        ['transformer', 'linformer'],\n",
    "        [transformer_attn, linformer_attn]\n",
    "    ):\n",
    "        mem_alloc_pre = torch.cuda.memory_allocated()\n",
    "        output = attn_layer(input)\n",
    "        mem_alloc_aft = torch.cuda.memory_allocated()\n",
    "        # convert to Mb\n",
    "        mem_usages[model_name].append((mem_alloc_aft - mem_alloc_pre) / 2 ** 20)\n",
    "        # free output\n",
    "        del output\n",
    "        torch.cuda.empty_cache()\n",
    "\n",
    "    # free memory\n",
    "    del input, transformer_attn, linformer_attn\n",
    "    torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jjSAz6r3Kd5e"
   },
   "source": [
    "Построим графики"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zu89W9tAKd5e"
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(9, 6))\n",
    "\n",
    "ax.plot(sequence_lengths, mem_usages['transformer'], '-v',\n",
    "        color='tomato', label='transformer')\n",
    "ax.plot(sequence_lengths, mem_usages['linformer'], '-v',\n",
    "        color='navy', label='linformer')\n",
    "# define x,y label\n",
    "ax.set_xlabel('Sequence length')\n",
    "ax.set_ylabel('Memory usage (Mb)')\n",
    "# plot in logspace\n",
    "ax.set_xscale('log');\n",
    "ax.set_yscale('log');\n",
    "# plot legend\n",
    "ax.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Vw7wMv0wKd5f"
   },
   "source": [
    "Как видно, используемая Трансформером память растет значительно быстрее с ростом длины последовательности."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TtCdFptUKd5f"
   },
   "source": [
    "**Задание 8 (0.1 балла):** обучите `Linformer` на том же самом датасете."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mniVdR1GKd5f"
   },
   "outputs": [],
   "source": [
    "# model hyperparams\n",
    "embed_dim = # TODO\n",
    "num_heads = # TODO\n",
    "feedforward_dim = # TODO\n",
    "num_layers = # TODO\n",
    "proj_dim = # TODO\n",
    "# the same as for transformer task\n",
    "sequence_length = 256"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Q-iboxLtKd5f"
   },
   "outputs": [],
   "source": [
    "model = LinformerForSequenceClassification(\n",
    "    num_classes=1,\n",
    "    input_dim=vocab_length,\n",
    "    embed_dim=embed_dim,\n",
    "    num_heads=num_heads,\n",
    "    feedforward_dim=feedforward_dim,\n",
    "    sequence_length=(sequence_length + 1),\n",
    "    proj_dim=proj_dim,\n",
    "    activation=nn.ReLU,\n",
    "    num_layers=num_layers,\n",
    ")\n",
    "# put model on device\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "h_2uO-gMKd5f"
   },
   "outputs": [],
   "source": [
    "num_epochs = # TODO\n",
    "warmup_steps = # TODO\n",
    "lr = 5e-5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "H5i8DCknKd5g"
   },
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ax2ejW9oKd5g"
   },
   "outputs": [],
   "source": [
    "scheduler = CosineAnnealingWithWarmupLR(\n",
    "    optimizer,\n",
    "    warmup_steps=warmup_steps,\n",
    "    max_steps=int(num_epochs * len(train_loader))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Vdr3US4QKd5g"
   },
   "outputs": [],
   "source": [
    "train(\n",
    "    model,\n",
    "    num_epochs=num_epochs,\n",
    "    optimizer=optimizer,\n",
    "    scheduler=scheduler,\n",
    "    train_loader=train_loader,\n",
    "    val_loader=val_loader,\n",
    "    device=device\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "J423rtjuKd5g"
   },
   "source": [
    "Если все было реализовано правильно, `Linformer` должен так же успешно справляться с задачей определения палиндрома."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}