sem18/main.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.3"
    },
    "colab": {
      "name": "main.ipynb",
      "provenance": []
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "ab1befa3fa404373afc824f2824cb1a0": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_6639bf4ef82a45a2a116cf1e7910dee7",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_6ad04290250b4439a3ba36dcce626608",
              "IPY_MODEL_03906703a09c48aba9a06c15f552927c",
              "IPY_MODEL_19bc31c66b8f46928dd91084d8691f2f"
            ]
          },
          "model_module_version": "1.5.0"
        },
        "6639bf4ef82a45a2a116cf1e7910dee7": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          },
          "model_module_version": "1.2.0"
        },
        "6ad04290250b4439a3ba36dcce626608": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_9682fcbcb72e44d38c20b4cfac5407f3",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": "Downloading: 100%",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_a7a9e6f1dcfa4e6a911f6f440c46d890"
          },
          "model_module_version": "1.5.0"
        },
        "03906703a09c48aba9a06c15f552927c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_b0772ff6976d402fb3fa61f0778cfd48",
            "_dom_classes": [],
            "description": "",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 1888175414,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 1888175414,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_c207a78178824e3e808cb65f54f45994"
          },
          "model_module_version": "1.5.0"
        },
        "19bc31c66b8f46928dd91084d8691f2f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_0b2d4f3f73074592aa6a94a33e7973a2",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 1.89G/1.89G [00:32&lt;00:00, 63.4MB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_401cca8b41f04dfdb83e0decc1b50881"
          },
          "model_module_version": "1.5.0"
        },
        "9682fcbcb72e44d38c20b4cfac5407f3": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          },
          "model_module_version": "1.5.0"
        },
        "a7a9e6f1dcfa4e6a911f6f440c46d890": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          },
          "model_module_version": "1.2.0"
        },
        "b0772ff6976d402fb3fa61f0778cfd48": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          },
          "model_module_version": "1.5.0"
        },
        "c207a78178824e3e808cb65f54f45994": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          },
          "model_module_version": "1.2.0"
        },
        "0b2d4f3f73074592aa6a94a33e7973a2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          },
          "model_module_version": "1.5.0"
        },
        "401cca8b41f04dfdb83e0decc1b50881": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          },
          "model_module_version": "1.2.0"
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gFDEdM3AeCzx"
      },
      "source": [
        "# Attention is all you need. Трансформеры."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ncRUyPBReCzz"
      },
      "source": [
        "## Библиотеки"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "D1YShLCYtnmK"
      },
      "source": [
        "!pip install --quiet -U dvc[gdrive]==1.11.1 transformers torch sentencepiece torchtext"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QQ7M5xkjeCzz"
      },
      "source": [
        "import matplotlib.pyplot as plt\n",
        "from mpl_toolkits import mplot3d\n",
        "from matplotlib import gridspec\n",
        "\n",
        "from tqdm.notebook import tqdm\n",
        "import numpy as np\n",
        "import  pandas as pd\n",
        "import seaborn as sns\n",
        "import torch\n",
        "import scipy\n",
        "from transformers import AutoModel, AutoTokenizer\n",
        "from transformers import MarianMTModel, MarianTokenizer\n",
        "\n",
        "from sklearn.metrics import classification_report\n",
        "from torch.utils.tensorboard import SummaryWriter\n",
        "\n",
        "import io\n",
        "import math\n",
        "import time\n",
        "from tqdm import tqdm\n",
        "from collections import Counter\n",
        "\n",
        "import torch\n",
        "import numpy as np\n",
        "import scipy.spatial\n",
        "import torch.nn as nn\n",
        "from torch import Tensor\n",
        "from torch.nn.utils.rnn import pad_sequence\n",
        "from torch.utils.data import DataLoader\n",
        "from torch.nn import (TransformerEncoder, TransformerDecoder,\n",
        "                      TransformerEncoderLayer, TransformerDecoderLayer)\n",
        "\n",
        "import torchtext\n",
        "from torchtext.vocab import vocab\n",
        "from torchtext.data.utils import get_tokenizer\n",
        "from torchtext.utils import download_from_url, extract_archive\n",
        "\n",
        "from transformers import AutoModel, AutoTokenizer"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YOMr4V9weCz0"
      },
      "source": [
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MJipdlV3eCz0",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 37
        },
        "outputId": "54efb475-7693-4c49-f760-5f19732dd5f7"
      },
      "source": [
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "device"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'cuda'"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FMLHstClxJjd"
      },
      "source": [
        "## Выборка паралельных текстов"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xJuyKE1E3HWr"
      },
      "source": [
        "!dvc get https://github.com/andriygav/MachineLearningSeminars sem18/data/Wikipedia.en-ru.ru\n",
        "!dvc get https://github.com/andriygav/MachineLearningSeminars sem18/data/Wikipedia.en-ru.en"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "78ZguASaxN2w"
      },
      "source": [
        "with open('./Wikipedia.en-ru.ru') as f:\n",
        "    ru_all_texts = f.read().splitlines()\n",
        "\n",
        "with open('./Wikipedia.en-ru.en') as f:\n",
        "    en_all_texts = f.read().splitlines()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0tIZBvxQJsbQ"
      },
      "source": [
        "ru_texts = []\n",
        "en_texts = []\n",
        "\n",
        "for ru_text, en_text in zip(ru_all_texts, en_all_texts):\n",
        "    if len(ru_text) < 100 and len(en_text) < 100:\n",
        "        ru_texts.append(ru_text.lower())\n",
        "        en_texts.append(en_text.lower())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Dmp0a-hG3nR4",
        "outputId": "2a806922-0d7b-4072-c8b9-1d52fcffe898"
      },
      "source": [
        "len(ru_texts), len(en_texts)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(306887, 306887)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kTOMF7UqeCz1"
      },
      "source": [
        "## Код для обучения"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w7io7qIUIsp_"
      },
      "source": [
        "def train_on_batch(model, x_batch, y_batch, optimizer, loss_function):\n",
        "    model.train()\n",
        "    optimizer.zero_grad()\n",
        "    \n",
        "    output = model(x_batch.to(model.device), y_batch.to(model.device))\n",
        "    \n",
        "    loss = loss_function(output.transpose(1,2), \n",
        "                         y_batch.to(model.device))\n",
        "    loss.backward()\n",
        "\n",
        "    optimizer.step()\n",
        "    return loss.cpu().item()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rz7sgZHWIvaY"
      },
      "source": [
        "def train_epoch(train_generator, model, loss_function, optimizer, callback = None):\n",
        "    epoch_loss = 0\n",
        "    total = 0\n",
        "    for it, (batch_of_x, batch_of_y) in enumerate(train_generator):\n",
        "        batch_loss = train_on_batch(model, batch_of_x, batch_of_y, optimizer, loss_function)\n",
        "        \n",
        "        if callback is not None:\n",
        "            with torch.no_grad():\n",
        "                callback(model, batch_loss)\n",
        "            \n",
        "        epoch_loss += batch_loss*len(batch_of_x)\n",
        "        total += len(batch_of_x)\n",
        "    \n",
        "    return epoch_loss/total"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mLe12zzAIx_d"
      },
      "source": [
        "def trainer(count_of_epoch, \n",
        "            batch_size, \n",
        "            dataset,\n",
        "            model, \n",
        "            loss_function,\n",
        "            optimizer,\n",
        "            lr = 0.001,\n",
        "            callback = None):\n",
        "\n",
        "    optima = optimizer(model.parameters(), lr=lr)\n",
        "    \n",
        "    iterations = tqdm(range(count_of_epoch), desc='epoch')\n",
        "    iterations.set_postfix({'train epoch loss': np.nan})\n",
        "    for it in iterations:\n",
        "        batch_generator = tqdm(\n",
        "            torch.utils.data.DataLoader(dataset=dataset, \n",
        "                                        batch_size=batch_size, \n",
        "                                        shuffle=True, pin_memory=True), \n",
        "            leave=False, total=len(dataset)//batch_size+(len(dataset)%batch_size>0))\n",
        "        \n",
        "        epoch_loss = train_epoch(train_generator=batch_generator, \n",
        "                    model=model, \n",
        "                    loss_function=loss_function, \n",
        "                    optimizer=optima, \n",
        "                    callback=callback)\n",
        "        \n",
        "        iterations.set_postfix({'train epoch loss': epoch_loss})"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wiamD1GkeCz1"
      },
      "source": [
        "## Модель внимания в рекурентных моделях"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X_AbmZii82Qr"
      },
      "source": [
        "### Определение модели"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5pypFq8z3xWa"
      },
      "source": [
        "class Encoder(torch.nn.Module):\n",
        "    @property\n",
        "    def device(self):\n",
        "        return next(self.parameters()).device\n",
        "\n",
        "    def __init__(self, vocab_size, emb_dim=30, hidden_dim=30):\n",
        "        super(Encoder, self).__init__()\n",
        "        self.embedding = torch.nn.Embedding(vocab_size, emb_dim)\n",
        "        self.lstm = torch.nn.LSTM(emb_dim, hidden_dim, batch_first=True)\n",
        "\n",
        "    def forward(self, input):\n",
        "        r'''\n",
        "        :param input: тензор размера batch_size x seq_len --- список токенов\n",
        "\n",
        "        '''\n",
        "        act = self.embedding(input)\n",
        "        act, hidden = self.lstm(act)\n",
        "        return act, hidden"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-AEZeHhq5U4k"
      },
      "source": [
        "class Decoder(torch.nn.Module):\n",
        "    @property\n",
        "    def device(self):\n",
        "        return next(self.parameters()).device\n",
        "\n",
        "    def __init__(self, vocab_size, emb_dim=30, hidden_dim=30):\n",
        "        super(Decoder, self).__init__()\n",
        "        self.embedding = torch.nn.Embedding(vocab_size, emb_dim)\n",
        "        self.attention = torch.nn.MultiheadAttention(emb_dim, 1)\n",
        "\n",
        "        self.lstm = torch.nn.LSTM(emb_dim, hidden_dim, batch_first=True)\n",
        "\n",
        "        self.linear = torch.nn.Linear(hidden_dim, vocab_size)\n",
        "\n",
        "    def forward(self, input, encoder_outputs, hidden):\n",
        "        r'''\n",
        "        :param input: тезор размера batch_size x seq_len\n",
        "        '''\n",
        "        act = self.embedding(input)\n",
        "\n",
        "        act, _ = self.attention(act.transpose(0, 1), \n",
        "                                encoder_outputs.transpose(0, 1), \n",
        "                                encoder_outputs.transpose(0, 1))\n",
        "\n",
        "        act = act.transpose(0, 1)\n",
        "        act, hidden = self.lstm(act, hidden)\n",
        "        return self.linear(act), hidden"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mN_1fGTh_zXl"
      },
      "source": [
        "class seq2seq(torch.nn.Module):\n",
        "    @property\n",
        "    def device(self):\n",
        "        return next(self.parameters()).device\n",
        "\n",
        "    def __init__(self, vocab_size, emb_dim=30, hidden_dim=30):\n",
        "        super(seq2seq, self).__init__()\n",
        "        self.vocab_size = vocab_size\n",
        "        self.encoder = Encoder(vocab_size, emb_dim, hidden_dim)\n",
        "        self.decoder = Decoder(vocab_size, emb_dim, hidden_dim)\n",
        "\n",
        "    def forward(self, input, decoder_input=None, max_seq_len=64):\n",
        "        r'''\n",
        "        '''\n",
        "        encoder_output, hidden = self.encoder(input)\n",
        "\n",
        "        if decoder_input is None:\n",
        "            translated_scores = torch.zeros(len(input), \n",
        "                                            max_seq_len, \n",
        "                                            self.vocab_size).to(self.device)\n",
        "            translated_scores[:, 0, input[:, 0]] = 1.\n",
        "            for i in range(1, max_seq_len):\n",
        "                translated_scores[:, i:i+1], hidden = self.decoder(\n",
        "                    torch.argmax(translated_scores[:, i-1:i], axis=-1), \n",
        "                    encoder_output, \n",
        "                    hidden)\n",
        "        else:\n",
        "            translated_scores, _ = self.decoder(\n",
        "                decoder_input, encoder_output, hidden)\n",
        "\n",
        "        return translated_scores\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8K0atJo18zaG"
      },
      "source": [
        "### Инициализация модели"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "W24s1W2VtEaT",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "3f53f2c5-838a-4896-d2ad-f34ffbd54494"
      },
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained('bert-base-multilingual-uncased', \n",
        "                                          verbose=False)\n",
        "tokenizer.vocab_size"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "105879"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 13
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-mS5tzMbPS0W"
      },
      "source": [
        "en_texts_tensors = tokenizer(en_texts[:20000], \n",
        "                             return_tensors='pt', max_length=64, padding=True)\n",
        "ru_texts_tensors = tokenizer(ru_texts[:20000], \n",
        "                             return_tensors='pt', max_length=64, padding=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T8bX1OlvLp9I"
      },
      "source": [
        "dataset_train_pt = torch.utils.data.TensorDataset(\n",
        "    en_texts_tensors['input_ids'][:19000], ru_texts_tensors['input_ids'][:19000])\n",
        "dataset_test_pt = torch.utils.data.TensorDataset(\n",
        "    en_texts_tensors['input_ids'][19000:], ru_texts_tensors['input_ids'][19000:])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8Dg8VO66MnsR",
        "outputId": "597e4542-c099-4c12-9e44-5ccb3e909e62"
      },
      "source": [
        "model = seq2seq(tokenizer.vocab_size)\n",
        "model.to(device)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "seq2seq(\n",
              "  (encoder): Encoder(\n",
              "    (embedding): Embedding(105879, 30)\n",
              "    (lstm): LSTM(30, 30, batch_first=True)\n",
              "  )\n",
              "  (decoder): Decoder(\n",
              "    (embedding): Embedding(105879, 30)\n",
              "    (attention): MultiheadAttention(\n",
              "      (out_proj): _LinearWithBias(in_features=30, out_features=30, bias=True)\n",
              "    )\n",
              "    (lstm): LSTM(30, 30, batch_first=True)\n",
              "    (linear): Linear(in_features=30, out_features=105879, bias=True)\n",
              "  )\n",
              ")"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 16
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E_ruz4ZqMt8-"
      },
      "source": [
        "loss_function = torch.nn.CrossEntropyLoss(ignore_index=0)\n",
        "optimizer = torch.optim.Adam"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Zob6-9yKM0a6"
      },
      "source": [
        "trainer(count_of_epoch=15, \n",
        "        batch_size=64, \n",
        "        dataset=dataset_train_pt,\n",
        "        model=model, \n",
        "        loss_function=loss_function,\n",
        "        optimizer = optimizer,\n",
        "        lr=0.001,\n",
        "        callback=None)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 37
        },
        "id": "2vktEqR3VKsh",
        "outputId": "0d495d4d-819e-449b-cfd5-9a694de9daaf"
      },
      "source": [
        "tokens_en = tokenizer(['in 1924, most of the site was converted to a museum.'], return_tensors='pt')\n",
        "tokens_ru = tokenizer(['в 1924 году крепость стала музеем.'], return_tensors='pt')\n",
        "with torch.no_grad():\n",
        "    answer = torch.argmax(\n",
        "        model(tokens_en['input_ids'].to(model.device), \n",
        "              tokens_ru['input_ids'].to(model.device), max_seq_len=10), \n",
        "        axis=-1)\n",
        "tokenizer.decode(answer[0])"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'. на [SEP]с [SEP]с [SEP] [SEP] [SEP] в'"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 70
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "StDlin7kV1Gq"
      },
      "source": [
        "with torch.no_grad():\n",
        "    encoder_outputs, _ = model.encoder(\n",
        "        tokens_en['input_ids'].to(model.device))\n",
        "\n",
        "    decoder_embedding = model.decoder.embedding(\n",
        "        tokens_ru['input_ids'].to(model.device))\n",
        "\n",
        "    act, act_weight = model.decoder.attention(\n",
        "        decoder_embedding.transpose(0, 1), \n",
        "        encoder_outputs.transpose(0, 1), \n",
        "        encoder_outputs.transpose(0, 1))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "i_wQB-zvYo3i"
      },
      "source": [
        "attention_frame = pd.DataFrame(\n",
        "    act_weight[0].cpu().numpy(),\n",
        "    index=tokenizer.convert_ids_to_tokens(\n",
        "        tokens_ru['input_ids'][0].cpu().numpy()),\n",
        "    columns=tokenizer.convert_ids_to_tokens(\n",
        "        tokens_en['input_ids'][0].cpu().numpy()))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 306
        },
        "id": "phu_44JZW_R0",
        "outputId": "6bb60a51-7ca5-4877-e2d5-c69cafa63d07"
      },
      "source": [
        "sns.heatmap(data=attention_frame)\n",
        "\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYMAAAEhCAYAAACdsMz3AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAgAElEQVR4nO3deZxddX3/8dc7A2FfRBaVYIMYNlkCBhCLCwIKUqFUZVXKomlVqIXWCi6oqLUqlVpKUUSWUpVV+ouaCigCLogJEJaExbBUwAUQUJZAkpn3749zRm4uM3PvnTlz72Hm/eRxHtzzved+7/fOZM7nfnfZJiIiJrcpvS5ARET0XoJBREQkGERERIJBRESQYBARESQYREQEsFKvC9ArK03duJIxtStN6asim0ryWXWllSsoCey33nbV5PPM1EryOX757ZXk8+CTj1SSj1Al+fQP9I85jzWnrlZBSUCq5jOtXNHfwwDVDHlfsmxpJfk8+fS9Y/4BLXvknrY/1Mrrv6KaX0gHJm0wiIjoqgqC/3hKMIiI6AYP9LoEI0owiIjohoEEg4iISc+pGURERGoGEREB/ct6XYIRJRhERHRDzZuJejbpTNJ0SUskLSjPXyLpAkl3S7pB0lxJm5fX3TbE618j6XpJCyTdLumTZfpBkhZL+m6XP1JExPAGBto/eqDXNYO7bc9UMePlMuA82wcDSNoe2Ai4f5jXngccaPtmSX3AFgC2L5T0O+Afx7/4ERHtSQdye3YHltn+ymCC7ZuhqEEM85oNgd+U1/YDi8a3iBERY1DzDuS6rE20DXBDh685FbhT0mWS/kbSqq1eIGm2pPmS5g8MPDWqgkZEjIoH2j96oC7BoGO2TwZmAVcAhwLfb+M1Z9qeZXvWlClrjHcRIyKe07+s/aMH6hIMFgKv7vRFtu+2fQawB7C9pBdXXrKIiCrUvAO5LsHgKmAVSbMHEyRtJ+l1w71A0r56bqnFGUA/8Pj4FjMiYpRq3kxUiw5k25Z0APBvkj4MPAPcB/x9eckWkh5oeMlxwNuBUyU9DSwHDis7kiMi6qfmHci1CAYAtn8NHDjM00Mt1H/xOBYnIqJSVX5XlbQ38GWgDzjL9r80Pf9yiuH365bXnGB77kh59rKZqB9YZ3DSWVUkHQT8J/BYlflGRIxJ//L2jxGU86pOB/YBtgYOkbR102UfAy6yvQNwMMU9cUQ9qxnYvh/YZBzyvRC4sOp8IyLGpLq+gJ2BxbbvAZB0AbA/K861MrB2+Xgd4NetMq1NM1FExIRW3U5nG7PiygwPALs0XfNJ4ApJxwJrAHu2yrQuo4kiIia2DkYTNU6QLY/Zrd9gBYcA59qeBrwVOF/SiPf7SVszqGq36bVWqWZD8r6Rf09tqWqj9l/1P1FJPhdW9LNZ8syzleQzpYKfcZFPNT/nKjahn9pXzZ9wf0VNGFVtZF/F3wPA7uu/qpJ8KtHBaCLbZwJnDvP0g6zYxD6tTGt0NLB3mdd15QoN6wMPDfeeqRlERHRDdfMM5gEzJG0qaSpFB/Gcpmt+RTEZF0lbAasCD4+U6aStGUREdNXykUcJtcv2cknHAJdTDBs92/ZCSScD823PAf4B+Jqk4yg6k4+wPWK1LcEgIqILqpxnUM4ZmNuUdlLD40XAn3eSZ4JBREQ3ZAZyRETUfdvLBIOIiG5IzSAiIlIziIiIlmsO9VqCQUREN9S8mWjCTDqTNF3SEkkLJN0j6ZRelyki4k+y01lX3W17JrArcESPyxIR8ZzsdNZVm5X7I2wKPK9mUC72NBtgSt86TJmyRpeLFxGTVs2biSZaMLjb9kxJqwPzJZ1b7psArLj408pTN65mRa2IiHakA7knnqXYSe1FrLjud0REb2RoaVcNNhOtAlxp+5ZeFygiAkgzUbfYvg+oZgH9iIiqJRhERAQjryDdcwkGERHdkJpBRETUfTTRRJt0FhFRTxXOQJa0t6Q7JS2WdMIQz59arsawQNJdkh5vlWdqBhER3VBRn4GkPuB0YC/gAWCepDnl7mblW/m4huuPBXZole+kDQZ/vOxDleSz9gFfrCSfOnUtXf30H3pdhElhijTmPH6/5IkKSjJxbbDuqr0uwnOq6zPYGVhs+x4ASRcA+wOLhrn+EOATrTKdtMEgIqKrqgsGG7PiZNoHgF2GulDSn1Esz3NVq0wTDCIiusD9/W1f27iOWunMcjmdTh0MXGK75ZsnGEREdEMHNYPGddSG8CCwScP5tDJtKAcDH2jnPTOaKCKiG6pbwnoeMEPSppKmUtzw5zRfJGlLivXZrmuneKkZRER0w0A1w0RsL5d0DHA50AecbXuhpJOB+bYHA8PBwAV2e8OYEgwiIrqhwhnItucCc5vSTmo6/2QneSYYRER0Q82Xo+hpn4GksyU9JOm2hrTtJV0n6VZJ35G0dpm+l6QbyvQbJL1piPzmNOYVEVEb/f3tHz3Q6w7kc4G9m9LOAk6wvS1wGTA4O+wR4G1l+l8D5ze+SNJfAU+Oa2kjIkZrwO0fPdDTYGD7WuDRpuTNgWvLx1cCby+vvcn2r8v0hcBqklYBkLQmcDzwmXEvdETEaFQ3mmhc9LpmMJSFFFOrAd7JiuNpB70duNH2s+X5p4F/BZ4eKWNJsyXNlzT/69//WVXljYhoLTWDjh0FvF/SDcBawNLGJyW9Cvg88Dfl+UxgM9uXtcrY9pm2Z9medfTer62+5BERw/DAQNtHL9RuNJHtO4A3A0jaHNh38DlJ0yj6EQ63fXeZvCswS9J9FJ9nQ0lX235jN8sdETGiHn3jb1ftagaSNiz/PwX4GPCV8nxd4HsUncs/Hbze9hm2X2Z7OrAbcFcCQUTUTkYTDU/StyimSm8h6QFJRwOHSLoLuAP4NXBOefkxwCuBkxo2bdiwJwWPiOhUhZvbjIeeNhPZPmSYp748xLWfocVoIdv3AduMvWQRERWreTNR7foMIiImpB4NGW1XgkFERDekZhAREV7em47hdiUYRER0Q2oGERGRPoOamvHusyrJZ6M1X1RJPlV4ZvnS1he1oU/1mn6ypKLPtWygmmr6FKmSfFZbaWol+VRhaf/yXhdhBVP7qrk1Hbakr5J8KlHzmkG9/uojIiYoD7jtoxVJe0u6U9JiSScMc82BkhZJWijpm63ynLQ1g4iIrqqoZiCpDzgd2At4AJgnaY7tRQ3XzABOBP7c9mPtTNBNMIiI6IbqRhPtDCy2fQ+ApAsoVnpe1HDNe4HTbT8GYPuhVpmmmSgiohuqW8J6Y+D+hvMHyrRGmwObS/qppJ9Lat5E7HlSM4iI6AK7/WYiSbOB2Q1JZ9o+s4O3WwmYAbwRmAZcK2lb24+P9IKIiBhvHfQZlDf+4W7+D7Lipl/TyrRGDwDX214G3Fsu/jkDmDfce6aZKCKiG6prJpoHzJC0qaSpwMHAnKZr/oeiVoCk9Smaje4ZKdPUDCIiuqCdIaNt5WMvl3QMcDnQB5xte6Gkk4H5tueUz71Z0iKgH/iQ7d+PlG/Pg4Gk6cDtwJ1l0i3AecApFOWbB7xvcL9jSbdR1GiWAq+0vaako4DtbP99ec17ga1tH9fFjxIRMbzl1U06sz0XmNuUdlLDYwPHl0db6tJMdLftmbZnUnSanAscZHtbioDwvoZr+4C9y2sHXQS8TdLK5fmRwNnjX+yIiPZUOelsPNQlGDTaArjX9l3l+XnA6xueXxN4tPEFtp8ErgL+QtKWwMq2b23OWNJsSfMlzX/q2Uebn46IGD/V9RmMi543E3VC0qrAquXNv9lZwEcotss8Z4jnV+ihn7beNvVeKCQiJpZ6r1NXy2BwJzBd0ittLwbeDVxTPncARcfI89i+XtImwI7Adl0paUREm3rV/NOu2gUD289IOhK4WNJgB/JXJM0Cvg48KmlBeflqkk5u6Di5CJg5OAU7IqIuXGEH8njoeTAYahN72z8EdmhMk7Qm8AXbn2xK+4+Gy3YDTh2vskZEjFqaiSqzCHikKe0Z4AxJ6wK/AG4uA0lERK3UfG+bF04wKFfde6gpbTlwfXm6edcLFRHRrgSDiIhIzSAiIlIziIgIGKjXNtPPM2mDwb13NS/yNzqrvex1leQTk8+zy5f1ugi19fSyZyvJ50NrNK/sPDo3VpBHmokiIgKsXpdgRAkGERFdkJpBRETggXrXDOq4amlExITjgfaPViTtLelOSYslnTDE80dIeljSgvJ4T6s8UzOIiOiCgf5qagaS+oDTgb0o9jqeJ2mO7UVNl15o+5h2800wiIjoggqbiXYGFtu+B0DSBcD+FEv2jFqaiSIiusBu/2hhY+D+hvMHyrRmb5d0i6RLyuX9R5RgEBHRBR5Q20fjrozlMbvDt/sOMN32dsCVFDtGjijNRBERXdBJM1HjroxDeBBo/KY/rUxrfP3vG07PAr7Q6j1b1gwkTZd0W/l4K0k3S3qdpDskfUPS7WU1ZPXymldLukbSDZIul/TSMv3qsvd7sHe7v0yXpC9Kuk3SrZIOanjvD5dpN0v6l/J9F0haJGnJYF7ltfdJWr/V54mI6IWBfrV9tDAPmCFpU0lTgYOBFZZUGLzvlvYDbm+Vads1A0kbA98CDgWeoti4/mjbP5V0NvB+SV8GTgP2t/1weWP/LHBUmc1htueX+Q3uY/xXwExge2B9ip7xa8u0/YFdbD8taT3bjwIzJU0Hvmt7Zrvlj4joJVc0A9n2cknHUGwB3AecbXuhpJOB+bbnAH8naT9gOfAocESrfNsNBmsC3weuKt90OnC/7Z+Wz/838HflNdsAV0qiLOhvWuS9G/At2/3A7yRdA+wEvAE4x/bTAGUgaOVHkgaAW4H32l7S+GTZ7jYb4D//9TO85/BD2sgyImLsqpyBbHsuMLcp7aSGxycCJ3aSZ7vBYBPgXcCJkrYClgDNfd4GBCy0vWsnhajQ7sDvgf8C3k1Tm1tjO9yyR+6p94akETGhDNR8baJ2RxPdbvtbwLHAVylu+i+XNHjTPxT4CXAnsMFguqSVJb2qRd4/Bg6S1CdpA+D1FFtYXgkc2dAXsV47BbVtimrR1DY/W0TEuLPV9tELHY0msn2NpDuAfShu/B8o+wsWAWfYXirpHcC/S1qnzP/fgIUjZHsZsCtwM0Xt4p9s/xb4vqSZwHxJSymqRB9pUcTvls1ETwIntbg2IqJr6r42kdzGDIfnvei5Dtxtqi5Qt1TVTJT9DCLqa7sXb1pJPjf+5idjvpMv2mzftu85W9/9va5HjswziIjogrr3GYwqGNi+j2LUUEREtKFXfQHtSs0gIqILRtEi31UJBhERXTAhm4kiIqIzaSaqqX13eH8l+ay9yuqV5DOaUV3NlixfWkFJoJw9PmZVfCYAP29+48Sw1tTVxpxHVb/z/oFqpsdW9btafaVVKsnnipl9leRThf6aDy2dtMEgIqKbUjOIiIj0GURExPMXc6ubBIOIiC5IzSAiIuhPMIiICFPvYNDuEtYRETEGA27/aEXS3uU2woslnTDCdW+XZEmzWuWZmkFERBcMVFQzkNQHnA7sBTxAsVXwHNuLmq5bC/ggcH07+aZmEBHRBUZtHy3sDCy2fY/tpcAFFPvFN/s08HngmXbK19VgIOlwSbdIulnS+ZIWlEd/w+OXSXqvpHnldZcO7nZW5nGupHvLa5dKWl/SmpJ+KOlGSbdKGuoHExHRMwMdHC1sDNzfcP5AmfYnknYENrH9vXbL17VmonL7y48Br7X9iKT1Bje5l/Sk7ZkN137b9tfKx58BjgZOK5/uA/7B9rcl3VemPQMcYPuPktYHfl5Wm+o+tDciJon+DpqJJM0GZjcknVnu4d7Oa6cAXwKO6KR83ewzeBNwse1HAAYDwTC2KYPAusCawOUNz63G86s9Av5Z0uspAuvGwEbAb1e4qOEHvNW6WzNtzU1G/2kiIjrQyepP5Y1/uJv/g0DjzWtamTZoLYr9Zq4u1xl7CTBH0n625w/3nnXtMzgXOMb2tsCngFUbnnsZ8Oum6w8DNgBeXdYwftf0GqD4AdueZXtWAkFEdFOFfQbzgBmSNpU0FTgYmPOn97H/YHt929NtTwd+DowYCKC7weAq4J2SXgwgab0Rrl0L+I2klSlu9JSveSUwHVjUdP06wEO2l0naHfizKgseETFWA2r/GInt5cAxFC0mtwMX2V4o6WRJ+422fF1rJioL+1ngGkn9wE0M36b1cYrhUA+X/19L0suA/wfMLnvQG30D+I6kW4H5wB3j8BEiIkatqqGlALbnAnOb0k4a5to3tpNnV+cZ2D4POG+I9DWbzs8Azhgii1c1XTe94XTXCooYETEu+ntdgBYy6SwiogsGKto0arwkGEREdEHdx7knGEREdEE1G4uOnwSDiIguqPkWyAkGERHdUOVoovEwaYPBV9atptK25UNLKsmnipUz6t4mGSt6atmzY85jWf/yCkpSP08urebv6jd3rF1JPutXkEd/vWPB5A0GERHdlD6DiIiofc09wSAiogvSgRwREWkmioiIBIOIiCCjiSIigtQMIiKCjCaKiAjqP5qoKzudSZouyZL+tjzvk/SgpEsl3VvuaIaktQfPJZ0m6UZJd5T7IQ++7ouS5km6RdLfNLzHhxrSP9WNzxUR0a6BDo5WJO0t6U5JiyWdMMTzfyvpVkkLJP1E0tat8uzmtpeLgb8sH+8N3A88AVwN7FumHwx82/Yy28fa3pFi05oPSloVOBr4g+2dgJ2A95b7gL4ZmAHsDMwEXi3p9c0FkDRb0nxJ87/16APj9kEjIpr1d3CMRFIfcDqwD7A1cMgQN/tv2t623BP+C8CXWpWvm8HgWWCxpFcB7wbOL9PPAo4sHx8JnDP4AknfAR4EzrD9DPBm4HBJCyi2w3wxRRB4c3ncBNwIbFmmr8D2mbZn2Z51yHrTqv+EERHDqGoPZIovvYtt31NuAXwBsH/jBbb/2HC6Bm10WXS7z+Ac4J/K9/0dgO2fls1IbwT6bN82eLHtt0l6MXC5pLUBAcfavrwxU0lvAT5n+6td+hwRER3pZDSRpNnA7IakM22fWT7emKJlZdADwC5D5PEB4HhgKvCmVu/ZzZoBtm8ANqTh23/pv4BvsmKtYN3y4TJgI4pawOXA+xr6GDaXtEaZfpSkNcv0jSVtOJ6fJSKiE+7kaGjFKI8zh8l2+PezT7e9GfBh4GOtru/6aCLb+wBIekdD8jeAzwDfaki7uLyhrw583fa9ks4CpgM3ShLwMPCXtq+QtBVwXZHMk8C7gIfG+/NERLRjoLrBpQ8CmzScTyvThnMBcEarTLsSDGzfB2zTlHYJcEl5uhtwie3HG57fa4h8BoCPlEfzc18GvlxdqSMiqlPhpLN5wAxJm1IEgYOBQxsvkDTD9i/L032BX9JCz+cZSDqNolf8rb0uS0TEeGk1SqhdtpdLOoaiebwPONv2QkknA/NtzwGOkbQnRTP7Y8Bft8q358HA9rG9LkNExHirctKZ7bnA3Ka0kxoef7DTPHseDCIiJoMK+wzGRYJBREQX1DsUJBhERHRFVi2tqc3vXNjrIlRupSl9leTzlg23qySfb7yjmmksO/33byvJZ/HjI42+675l/cvHnMcUVdMQ3VfRv53VVppaST5PLl1SST47PHhjJfmM/TeVZqKIiKC60UTjJcEgIqILUjOIiIiah4IEg4iIrkgHckRE4JrXDRIMIiK6YHnNg0FXl7BuRdLnJO0u6S8lndj03BaSzpM0RdJ1vSpjRMRodLKEdS/UKhhQbNDwc+ANwLVNz72uTNsWuI2IiBeQAdz20Qu1aCaS9EXgLcCmwHXAZsAeki4BfgScBrycYne0tYABSfNtz5J0BMWmOFvZvqPc12ARcKTtc7v+YSIihlD3DuRa1Axsf4his/tzKTa6v8X2drZPtv3jclPnOyk2f74S2Mf2rIYsfgEcVT4+imJ/5IiI2nAH//VCLWoGpR2Bmyk2s7+98QlJqwPP2rakGRSBodE8YAdJqwIzgflDvUHjvqLqW4cpU9ao9hNERAwjNYMWJM2UtAD4LPCPwPeAt0haIGk1SXOABcDWkm4BtgPmSzqoKavvUzQn/e9w79W4r2gCQUR0Uz9u+2hF0t6S7pS0WNIJQzx/vKRFkm6R9ENJf9Yqz54HA9sLymaguyiaga4C3mJ7pu0ltvcDvga8D/g74Cvlcxc2ZXU+8Frgv7tY/IiItgzYbR8jkdQHnE6xQ+TWwCGStm667CZglu3tKLYX/kKr8vU8GABI2gB4rNzjeEvbi5oueT3wE4oRRdcMlYfth2y/yvZD41vaiIjOVTi0dGdgse17bC+l2PB+/xXey/6R7afL058D01plWos+A9sPU2zajO3XDPH828qHnx7iuXMpOp4b046pvJAREWNQ4ZDRjYH7G84foBiWP5yjGaH5fFAtgkFExETXySihxsEupTNtn9npe0p6FzCLYu7WiBIMIiK6oJPRROWNf7ib/4PAJg3n08q0FUjaE/go8Abbz7Z6zwSDiIgu6K9ucOk8YIakTSmCwMHAoY0XSNoB+Cqwd7v9qAkGERFdUFUosL1c0jHA5UAfcLbthZJOBubbngN8EVgTuFjF1qi/KkdmDivBICKiC9xiyGiHec0F5jalndTweM9O80wwmECWD1Szy+r9yx6vJJ+VDn/eXJhReeTr/1BJPvVeQHh0Wo1Jb5f7q9jyHZ5yNd9/q/pcdZJtLyMiovbLUSQYRER0QYUdyOMiwSAioguq7DMYDwkGERFdUO96QYJBRERX9GqfgnYlGEREdEFGE0VERPoMIiIio4kiIoL6T6RLMIiI6IJ6h4IEg4iIrkgHco00bhihvnWYMmWNHpcoIiaLBIMaadwwYqWpG9f7NxMRE0p/RYv4jZdJFQwiInql7pPOpvS6AONB0lxJL+t1OSIiBtlu++iFCRkMbL/V9q97XY6IiEEDuO2jFUl7S7pT0mJJz9s4RNLrJd0oabmkd7RTvgkZDCIi6qaqmoGkPuB0YB9ga+AQSVs3XfYr4Ajgm+2WL30GERFdUOFoop2BxbbvAZB0AbA/sGjwAtv3lc+13WudYBAR0QUVjibaGLi/4fwBYJexZppmooiILnAH/0maLWl+wzF7vMuXmkFERBd0sjZR45yoITwIbNJwPq1MG5MEgzFaZaWVK8mnT2OvpFVVlteusnEl+Vz61gsqyWetlVevJJ8/Pvt0JfkIVZJP/0D/mPNYfeqqFZSkOitP6askn7ov6jYaFc4zmAfMkLQpRRA4GDh0rJmmmSgiogsG7LaPkdheDhwDXA7cDlxke6GkkyXtByBpJ0kPAO8EvippYavypWYQEdEFVc5Atj0XmNuUdlLD43kUzUdtSzCIiOiCrE0UERE4wSAiIrKEdURE9GwBunYlGEREdEFqBhERQf9AvfsMejbPQNJ0SUskLSjPPyppoaRbJC2QtEuZfnW5VOuC8rikTP+kpAfLtNsaxtceJ+lXkv6jV58tIqJZJ8tR9EKvawZ3254paVfgL4AdbT8raX1gasN1h9meP8TrT7V9iqStgB9L2tD2qZIeA2Z1ofwREW1Jn0F7Xgo8YvtZANuPdPJi27dLWg6sDzw0DuWLiBiTuvcZ1GU5iiuATSTdJek/Jb2h6flvNDQTfbH5xWWT0gDw8Ehv0rgS4MDAU9WVPiKihbpve1mLmoHtJyW9GngdsDtwoaQTbJ9bXjJcM9Fxkt4FPAEc5BY/xcaVAFeaunG9w3RETCh1X3yvFsEAwHY/cDVwtaRbgb8Gzm3xslNtnzLORYuIGLMsR9EGSVsAA7Z/WSbNBP6vh0WKiKhUOpDbsyZwmqR1geXAYqBxZ59vSFpSPn7E9p7dLmBExFikmagNtm8AXjvMc28cJv2T41ikiIhK9Wr+QLt6OZqoH1hncNJZVSQdB5wI/LHKfCMixqKqzW3GS8+Cge37bW9ie2bF+Z5qewvbH6ky34iIsahyaKmkvcuVGRZLOmGI51eRdGH5/PWSprfKsy7zDCIiJrQBD7R9jERSH3A6sA+wNXCIpK2bLjsaeMz2K4FTgc+3Kl+CQUREF1RYM9gZWGz7HttLgQuA/Zuu2R84r3x8CbCHJFVWwMl2ALOTzwsjnzqVJflMzt95lQfFaMr5DcfshufeAZzVcP5u4D+aXn8bMK3h/G5g/ZHeMzWDkc1ufUnyqUk+dSpL8ulOPnUqS6Vsn2l7VsNx5ni/Z4JBRMQLy4PAJg3n08q0Ia+RtBKwDvD7kTJNMIiIeGGZB8yQtKmkqcDBwJyma+ZQLOkDRbPSVS7bi4ZTi0lnNVZV1Sz5jH8+dSpL8ulOPnUqS9fYXi7pGOByoA842/ZCSScD823PAb4OnC9pMfAoRcAYkVoEi4iImATSTBQREQkGERGRYBARESQYTFqSXipplVG87oPtpA3z2vM7uT7qQdKLJO0s6fWDR6/LFNVLB3IDSc3Ds4byqO0j2szvtcB0GkZt2f6vURWuYpJ+AGwGXGr7Hzt43Y22d2xKu8n2Dm28dhGwJ/C/wBuBFabH23603XKU+W0E/DPwMtv7lOuz7Gr76x3mI+Aw4BW2T5b0cuAltn/RYT5fAD4DLAG+D2wHHGf7v9t47fEjPW/7S52UpcxzI2Cn8vQXth8aRR7vAT5IMZZ9AfAa4Drbb+own3WBw3n+38PfdVqmYfJ/ie3ftnntv7dx2R9tf2yMxXpBydDSFW0FvGeE50WxQFRL5bfgzSj+gPrLZAOjDgaSrur0j3A4tvcsb4LNC1wN996HAIcCmzYFzbUphq614yvAD4FXADc0Zk/xs3lFm/kMOhc4B/hoeX4XcCHFsLpO/CcwALwJOJliT+1Lee5G2q432/4nSQcA9wF/BVwLtAwGwFrl/7co33fwZ/w2oKOgBCDpQOCLFFvJimLzqA/ZvqTDrD5YlufntneXtCVFAO7UXODnwK0UP+uqfR3Yt81r9wdOanHNCUCCwST2UdvXjHSBpE+1mdcsYOtWEz1GeJ9bmpOAzQfTbW83mnwblWVb2OblPwN+A6wP/GtD+hNAc1mHe79/B/5d0hkUgWGwueFa2ze3WY5G69u+SNKJZf7LJfW3etEQdrG9o6SbynweKyfzdGrw72lf4GLbf2i1Ntgg258CkHQtsHLKt9YAAAzJSURBVKPtJ8rzTwLfG0VZPgrsNFgbkLQB8AOKRcs68YztZyQhaRXbd5Tb1HZqVdsj1n7Gwna7gQCKvdPPG+kCSS8aY5FecBIMGti+qDmt/Efx+OBNfahrhnEb8BKKG+ho3EexQc9gs4OAH1N8U+w62/8H/J+kPYEltgckbQ5sSfFtrxN3UHxb/jbF5zpf0tdsn9ZhPk9JejFFrQJJrwH+0GEeAMvKZYEH89mA0X17/a6kOyh+X+8r83mmwzw2ApY2nC8t0zo1palZ6PeMro/wgbKJ53+AKyU9xuj2Jz9f0nuB7wLPDiZ22jRYBdv/BiBpfduPjHTNZJI+gwaSTgIuKr/9rELR7rs9xb7Mh9r+QQd5/QiYSVHFb/zHv18HeRwAHAecYnuOpHtsd9qUUilJNwCvA14E/JRiavxS24d1kMctFG37T5Xna1C0Q3dU25G0I3AasA1F8N0AeIfttmoqDfkcBhwE7Eix7O87gI93EPgb81oP+IPtfkmrA2u325Zdvv6jwIHAZWXSX1L8m+yoaabsv9ge+FaZdBBwi+0Pd5JPU55voFjj5vsulk7u5LUfAD4LPA5/2v/Rvfj3LOltwNkUf9f9wIG2f9btctRNgkEDSQuBbWxb0mzgEIoOz82B82zv3EFebxgqvVUz1BD5rAF8mqL/4dW2p3Xy+qoNdiBLOhZYzfYXJC1wBzvWSbqVognjmfJ8VWCe7W1HUZ6VKNrZBdxpe1mneZT5bAnsUebzQ9u3jzKfbSj6YVYdTOt00EAZ5F5Xnl5r+6ZRlOPzwPXAbmXSj4HXjCUYjIWke4Cdh/sm3uWy3EIRAO6QtAvwBdtD/r1OJmkmWtHShjb+twAX2O4Hbi9vOm3r9KY/Qj5PAcdL2h7YtYo8x0iSdqUYfXN0mdbXYR7nANdLavz222mn76CdeW6Eyo6SRnPzPd/2uymar5rTOsnnExSjpLam6DDdB/gJnQ8aWJ1iNMs5kjaQtKntezvMY6/yxv/thvJ9CuhJMAAWA0/36L2bLbd9B4Dt6yWt1eoFk0GCwYqeLb/Z/Q7YHWgccrlGOxlI+ont3SQ9wXPVYShHzNheu5MCSVrZ9rKyg/XmMm3Yts4u+HvgROCycnGsVwA/6iQD21+SdDXPfWs9cpTffqsasfWqpnz7gFd3Wh6K5qXtgZtsH1kO7WxnJFHje3+CYvDBFhRBc+Uyjz9v8/XvA94PvKJpEMJaFM16vfIUsKBsPm1sNq1kaGmHNmwayrvC+WiG8U4ECQYr+iDFaIsNKEYc3Asg6a3Aje1kYHu38v9j+rYhaXfgfGBVSTdS7HR0X/n0FRTt211X1niukbSmpDVt3wN0/Adt+0ba/JmOYKwjtk4EPgKsJumPPDfvYSmjW8nymbJjfbmktYGHWHHd+XYcAOxA+bOx/esOv7l+k2Iex+cohkcOeqIXnbUN/qc86uBrPDeUd6jzSSnBoIHt6ylGxzSnz5W0WpeL8wXgLeW373dQjOR4t+2f0zRZq5skbUvxzXu94lQPA4fbbneIapXGNGLL9ueAz0n6nO0TKyjPvHLkzdco5lE8CVzXYR5Lyz6rwZFNbdVIB9n+A8WIqkM6fN9x1WooZzcNDuONFSUYtO9UiolI3TJ18AZr+xJJtwPflvRhVmx+6ravAsfb/hGApDdS3Pxe260CSPoOxc9gLWCRpFGN2JK0Zdl2fHHZabuCsvbSibWBd1JM9Po+xUiijkY2ARdJ+iqwbjkU8yjgrA7zqB1J9zLEv9sejSa6yPaB5ePPN3aqS7rC9pu7XaY6SDBoX7e/jS9TwxT7soawB8U47c26XJZGawwGAgDbV3f67bUCp1D8Pj5P0fk8aDCtXcdT7H/bOImu8YbV6Wzvr1OMAjqN4nd0k6RrbX+53QxsnyJpL4o5JlsAJ9m+ssNy1NGshserUgTN9XpUlhkNj/dixU71DbpcltrI0NI2SfqV7Zd38f32BB5208zcshniA7Y/262yNL3/ZRTt2eeXSe+iGPJ6QA/KMtQ6SbeMYr7CgRRj5/8o6eMU/TGfHkXNYLDzeSeKAQh/SzFB73lNjyO8foVvqsOlTQSSbrA9mo76sb7vn/7dNP8bGurf1GSRmkGDcvz7UNFRjG4W6KgNN8HN9uMUk3d65SjgUzzXZPZj4MhuFmAcRsx8zMWyFrtR1AZOAc4AdumwXD+kGHV2HcXP5U/LQXSg+ZsqFENUX9DBoKkZbgpFTaFX95/VJe1QlmO18rHKo9t9g7WRYLCiv+h1AQZJWhP4J+DtFCtGLgXuBr5i+9weFm0zihEyUyj+/exBcQMd81pJHah6xMzgsNR9ga/Z/p6kz4win1sohqRuQ9GJ+7ik62wvafXChgC3Wc2GhFalsSluOcVyKwf2pij8FvjSEI8HzyelNBM1kPRKYCPbP21K/3Pgt7bv7mJZ/h/FkgQ/oPijWQO4gGIlxQdtf6RbZWkq150U8y9uo2H9HhdrF70gSfou8CDFt/IdKdYW+oXt7UeZ31rAERQ/p5fYbrlvhKR1KJb4qNuQ0JgkEgwalDeFE23f2pS+LfDPtru2SJykmxtvRpLm2d5J0hRgUSft0BWX6yeDcykmChVrCO0N3Gr7l5JeCmxr+4oO8zmGogP51RTffH8M/Nj2VW2+vg9Y2Kvf7XhSRXtPVFSWnYD7BwdnSDqcogb+f8AnJ2vwTTPRijZqDgQAtm+VNL3LZXlK0m62fyJpP8o9A8pJTT2bZwB8QtJZFPsSNA7n/PbwL6k320/TsGyD7d8wurkLq1I0Odxge/koytEv6U5JL7f9q1G8f52dSzV7T1ThqxRrjqFi17Z/AY6lWFjyTIqZ5JNOgsGK1h3huW53LP0tcJakGRR7DhwFf1peua0NdsbJkRQT81bmuWYi03Aznaxsn1JBNi8CFpZzJ55qyLvt1W5rqqq9J6rQ1/Dt/yDgTNuXApdKWtCjMvVcgsGK5kt6r+2vNSaq2PrvhmFeMy7KyUrPWyXV9sPluke9spPt0WxuEu35eK8LME6q2nuiCn2SViprb3tQzDUZNGnviekzaFC2a15GMXJn8OY/C5gKHOAO1qUfT92e89D03ucAX7S9qBfvPxlI+jNghu0flP0ZfS53PnuhUkV7T1RUlo8CbwUeAV5OsbOcywEk59lua1HAiSbBYAjlInHblKcL2+0ArLgMw/2RCNi8nREq46FcFmMz4F6KPoPB1Vi7ObR0wiqXoJgNrGd7s7KZ8Cu29+hx0cZMFe09UVFZXgO8FLjCz22ytDmw5mgmG04ECQYN2pl92K0ZipJ+R7GnwmPNTwE/s/2y8S7DUMpvrc/zQh5aWidlm/XOwPW2dyjTbvUoNv6pE0nvpJjl/YSkj1EM4f1ML268dfo7r5NJ2z42jK1G+EYOxY14nS6V5bsU31Ke16GlYi+AnshNf9w9a3vp4ICx8tv0RPjG9nHbF5ezvPdglLO8K1Knv/PaSDBYUTvju7syAsL20SM8d2g3yhA9cY2kwT0W9qKYlfydHpepClXN8q5Cbf7O6yTNRBE1Uk4qPBp4M8U31MuBs/wC/0OtepZ3VC/BIKJGJP0V8D3bz7a8+AWkqlneMX6m9LoAEbGCtwF3STpf0l+UfQYTwfrAfIp9xl9OMWnxjt4WKRqlZhBRM5JWpli2+iBgN+BK2+/pbanGpmF5eFEs27EpxfDSV/W0YPEnE+VbR8SEYXuZpP+luHmuRrGb2ws6GDQPjS0nob2/R8WJIaSZKKJGJO0j6VzglxQraZ4FvKSnhRoH5fyCXgwrjWGkZhBRL4dTrOb5NxOpE1nS8Q2nUyiW+f51j4oTQ0ifQUSMO0mf4LnJc4M7nV06kQLeC12CQUSNlENLPw9syHP78tr22j0t2BiVG8p8BJjOcy0SWdOqRhIMImpE0mLgbbZv73VZqjQRt0udaNJnEFEvv5togaD0sO2JsKzGhJWaQUSNSPoyxeih/2GCbCsKIGkP4BAm0HapE01qBhH1sjbwNMXaRIMmwrai2S615lIziIhxJ+nObJdab5l0FlEjkqZJukzSQ+VxqaRpvS5XBX4maeteFyKGl5pBRI1IuhL4JnB+mfQu4DDbe/WuVGOX7VLrL8EgokYkLbA9s1XaC022S62/dCBH1MvvJb0L+FZ5fgjw+x6WpxK56ddfagYRNVJ+gz4N2JVitM3PgGNt39/TgsWEl2AQUSOSzgP+3vZj5fl6wCm2j+ptyWKiy2iiiHrZbjAQANh+FNihh+WJSSLBIKJepkh60eBJWTNI316Mu/wji6iXfwWuk3Rxef5O4LM9LE9MEukziKiZcnLWm8rTq2wv6mV5YnJIMIiIiPQZREREgkFERJBgEBERJBhERAQJBhERAfx/savf3h3cvyoAAAAASUVORK5CYII=\n",
            "text/plain": [
              "<Figure size 432x288 with 2 Axes>"
            ]
          },
          "metadata": {
            "tags": [],
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r6xRmj88eCz1"
      },
      "source": [
        "## Трансформер "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v7pNT8SUbH4I"
      },
      "source": [
        "### Определение модели"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Bxa45fl0alnH"
      },
      "source": [
        "class Encoder(torch.nn.Module):\n",
        "    @property\n",
        "    def device(self):\n",
        "        return next(self.parameters()).device\n",
        "\n",
        "    def __init__(self, vocab_size, emb_dim=30, hidden_dim=30):\n",
        "        super(Encoder, self).__init__()\n",
        "\n",
        "        self.layers = torch.nn.Sequential()\n",
        "        \n",
        "        self.embedding = torch.nn.Embedding(vocab_size, emb_dim)\n",
        "        self.attention1 = torch.nn.MultiheadAttention(emb_dim, 1)\n",
        "        self.linear1 = torch.nn.Linear(emb_dim, hidden_dim)\n",
        "        self.relu = torch.nn.ReLU()\n",
        "        self.attention2 = torch.nn.MultiheadAttention(hidden_dim, 1)\n",
        "        self.linear2 = torch.nn.Linear(hidden_dim, hidden_dim)\n",
        "\n",
        "\n",
        "    def forward(self, input):\n",
        "        r'''\n",
        "        :param input: тензор размера batch_size x seq_len --- список токенов\n",
        "\n",
        "        '''\n",
        "        input = input.transpose(0, 1)\n",
        "        act = self.embedding(input)\n",
        "        act, _ = self.attention1(act, act, act)\n",
        "        act = self.relu(act)\n",
        "        act = self.linear1(act)\n",
        "        act = self.relu(act)\n",
        "        act, _ = self.attention2(act, act, act)\n",
        "        act = self.relu(act)\n",
        "        act = self.linear2(act)\n",
        "        return act.transpose(0, 1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8esH9fEzasPV"
      },
      "source": [
        "class Decoder(torch.nn.Module):\n",
        "    @property\n",
        "    def device(self):\n",
        "        return next(self.parameters()).device\n",
        "\n",
        "    def __init__(self, vocab_size, emb_dim=30, hidden_dim=30):\n",
        "        super(Decoder, self).__init__()\n",
        "        self.embedding = torch.nn.Embedding(vocab_size, emb_dim)\n",
        "        self.attention = torch.nn.MultiheadAttention(emb_dim, 1)\n",
        "        self.linear1 = torch.nn.Linear(emb_dim, hidden_dim)\n",
        "        self.relu = torch.nn.ReLU()\n",
        "\n",
        "        self.attention2 = torch.nn.MultiheadAttention(hidden_dim, 1)\n",
        "        self.linear2 = torch.nn.Linear(hidden_dim, hidden_dim)\n",
        "\n",
        "        self.linear = torch.nn.Linear(hidden_dim, vocab_size)\n",
        "\n",
        "    def forward(self, input, encoder_outputs):\n",
        "        r'''\n",
        "        :param input: тезор размера batch_size x seq_len\n",
        "        '''\n",
        "        input = input.transpose(0, 1)\n",
        "        act = self.embedding(input)\n",
        "\n",
        "        act, _ = self.attention(act, \n",
        "                                encoder_outputs.transpose(0, 1), \n",
        "                                encoder_outputs.transpose(0, 1))\n",
        "        \n",
        "        act = self.relu(act)\n",
        "        act = self.linear1(act)\n",
        "        act = self.relu(act)\n",
        "        act, _ = self.attention2(act, \n",
        "                                 encoder_outputs.transpose(0, 1), \n",
        "                                 encoder_outputs.transpose(0, 1))\n",
        "        act = self.relu(act)\n",
        "        act = self.linear2(act)\n",
        "\n",
        "        return self.linear(act).transpose(0, 1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "spW1dQbmaszI"
      },
      "source": [
        "class seq2seq(torch.nn.Module):\n",
        "    @property\n",
        "    def device(self):\n",
        "        return next(self.parameters()).device\n",
        "\n",
        "    def __init__(self, vocab_size, emb_dim=30, hidden_dim=30):\n",
        "        super(seq2seq, self).__init__()\n",
        "        self.vocab_size = vocab_size\n",
        "        self.encoder = Encoder(vocab_size, emb_dim, hidden_dim)\n",
        "        self.decoder = Decoder(vocab_size, emb_dim, hidden_dim)\n",
        "\n",
        "    def forward(self, input, decoder_input=None, max_seq_len=64):\n",
        "        r'''\n",
        "        '''\n",
        "        encoder_output = self.encoder(input)\n",
        "\n",
        "        if decoder_input is None:\n",
        "            translated_scores = torch.zeros(len(input), \n",
        "                                            max_seq_len, \n",
        "                                            self.vocab_size).to(self.device)\n",
        "            translated_scores[:, 0, input[:, 0]] = 1.\n",
        "            for i in range(1, max_seq_len):\n",
        "                translated_scores[:, i:i+1] = self.decoder(\n",
        "                    torch.argmax(translated_scores, axis=-1), \n",
        "                    encoder_output)[:, i:i+1]\n",
        "        else:\n",
        "            translated_scores = self.decoder(\n",
        "                decoder_input, encoder_output)\n",
        "\n",
        "        return translated_scores\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "79eZIWWZfhDm"
      },
      "source": [
        "### Инициализация модели"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "czRgHj18fjvn",
        "outputId": "cb4e7300-2b7f-4bc1-9c91-3b16f6683e5d"
      },
      "source": [
        "model = seq2seq(tokenizer.vocab_size)\n",
        "model.to(device)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "seq2seq(\n",
              "  (encoder): Encoder(\n",
              "    (layers): Sequential()\n",
              "    (embedding): Embedding(105879, 30)\n",
              "    (attention1): MultiheadAttention(\n",
              "      (out_proj): _LinearWithBias(in_features=30, out_features=30, bias=True)\n",
              "    )\n",
              "    (linear1): Linear(in_features=30, out_features=30, bias=True)\n",
              "    (relu): ReLU()\n",
              "    (attention2): MultiheadAttention(\n",
              "      (out_proj): _LinearWithBias(in_features=30, out_features=30, bias=True)\n",
              "    )\n",
              "    (linear2): Linear(in_features=30, out_features=30, bias=True)\n",
              "  )\n",
              "  (decoder): Decoder(\n",
              "    (embedding): Embedding(105879, 30)\n",
              "    (attention): MultiheadAttention(\n",
              "      (out_proj): _LinearWithBias(in_features=30, out_features=30, bias=True)\n",
              "    )\n",
              "    (linear1): Linear(in_features=30, out_features=30, bias=True)\n",
              "    (relu): ReLU()\n",
              "    (attention2): MultiheadAttention(\n",
              "      (out_proj): _LinearWithBias(in_features=30, out_features=30, bias=True)\n",
              "    )\n",
              "    (linear2): Linear(in_features=30, out_features=30, bias=True)\n",
              "    (linear): Linear(in_features=30, out_features=105879, bias=True)\n",
              "  )\n",
              ")"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 84
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 37
        },
        "id": "uw7nACxhfnP7",
        "outputId": "5c920cc9-1672-4571-9134-4736154b86ff"
      },
      "source": [
        "tokens_en = tokenizer(['in 1924, most of the site was converted to a museum.'], return_tensors='pt')\n",
        "tokens_ru = tokenizer(['в 1924 году крепость стала музеем.'], return_tensors='pt')\n",
        "with torch.no_grad():\n",
        "    answer = torch.argmax(\n",
        "        model(tokens_en['input_ids'].to(model.device), \n",
        "              tokens_ru['input_ids'].to(model.device)), \n",
        "        axis=-1)\n",
        "tokenizer.decode(answer[0])"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'##綸綸綸綸綸綸綸綸綸綸'"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 85
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ukw7dD4agDii"
      },
      "source": [
        "loss_function = torch.nn.CrossEntropyLoss(ignore_index=0)\n",
        "optimizer = torch.optim.Adam"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZyXmosIHgHAC"
      },
      "source": [
        "trainer(count_of_epoch=5, \n",
        "        batch_size=64, \n",
        "        dataset=dataset_train_pt,\n",
        "        model=model, \n",
        "        loss_function=loss_function,\n",
        "        optimizer = optimizer,\n",
        "        lr=0.001,\n",
        "        callback=None)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 37
        },
        "id": "2GeTQJgogOPN",
        "outputId": "440d55db-84e8-4d30-dbb9-7eadfc94b2ab"
      },
      "source": [
        "tokens_en = tokenizer(['in 1924, most of the site was converted to a museum.'], return_tensors='pt')\n",
        "tokens_ru = tokenizer(['в 1924 году крепость стала музеем.'], return_tensors='pt')\n",
        "with torch.no_grad():\n",
        "    answer = torch.argmax(\n",
        "        model(tokens_en['input_ids'].to(model.device), \n",
        "              tokens_ru['input_ids'].to(model.device)), \n",
        "        axis=-1)\n",
        "tokenizer.decode(answer[0])"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'##аааааааааа'"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 88
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Переводчик из Encoder блоков\n",
        "\n",
        "[На основе кода по генерации перевода.](https://pytorch.org/tutorials/beginner/translation_transformer.html)\n",
        "\n",
        "[Статья 2021 года. ParaSCI: A Large Scientific Paraphrase Dataset for Longer Paraphrase Generation.](https://github.com/dqxiu/ParaSCI)"
      ],
      "metadata": {
        "id": "rCCfd8PkEyOq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Убирем рандом\n",
        "_ = torch.manual_seed(0)"
      ],
      "metadata": {
        "id": "bDPwZ0hhHwbr"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "DEVICE"
      ],
      "metadata": {
        "id": "t2QIgcuiHyz7",
        "outputId": "4224e4c2-e79d-42a9-ef67-133c1276ce05",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 36
        }
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'cuda'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Скачиваем данные"
      ],
      "metadata": {
        "id": "435FOrWzFFC-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "url_base = 'https://raw.githubusercontent.com/dqxiu/ParaSCI/master/Data/ParaSCI-ACL/'\n",
        "train_urls = ('train/train.src', 'train/train.tgt')\n",
        "val_urls = ('val/val.src', 'val/val.tgt')\n",
        "test_urls = ('test/test.src', 'test/test.tgt')\n",
        "\n",
        "train_filepaths = [download_from_url(url_base + url) for url in train_urls]\n",
        "val_filepaths = [download_from_url(url_base + url) for url in val_urls]\n",
        "test_filepaths = [download_from_url(url_base + url) for url in test_urls]"
      ],
      "metadata": {
        "id": "fDHCtzgqE21l"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Строим словарь"
      ],
      "metadata": {
        "id": "7CDXR6G6FG_R"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer = get_tokenizer('spacy', language='en_core_web_sm')\n",
        "\n",
        "def build_vocab(filepaths, tokenizer):\n",
        "    counter = Counter()\n",
        "    for filepath in filepaths:\n",
        "        with io.open(filepath, encoding=\"utf8\") as f:\n",
        "            for string_ in f:\n",
        "                counter.update(tokenizer(string_))\n",
        "    return vocab(counter, specials=['<unk>', '<pad>', '<bos>', '<eos>'])\n",
        "\n",
        "vocab = build_vocab(train_filepaths, tokenizer)\n",
        "vocab.set_default_index(vocab['<unk>'])"
      ],
      "metadata": {
        "id": "q22j8lNKE6fV"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "len(vocab)"
      ],
      "metadata": {
        "id": "qx8kCVxkFWnY",
        "outputId": "c8cca904-0711-47ce-8078-acb6533efbf7",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "14822"
            ]
          },
          "metadata": {},
          "execution_count": 16
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Предобработка данных"
      ],
      "metadata": {
        "id": "TCOlOkeJFXHM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def data_process(filepaths):\n",
        "    raw_src_iter = iter(io.open(filepaths[0], encoding=\"utf8\"))\n",
        "    raw_tgt_iter = iter(io.open(filepaths[1], encoding=\"utf8\"))\n",
        "    data = []\n",
        "    for (raw_src, raw_tgt) in zip(raw_src_iter, raw_tgt_iter):\n",
        "        src_tensor_ = torch.tensor(\n",
        "            [vocab[token] for token in tokenizer(raw_src.rstrip(\"\\n\"))],\n",
        "            dtype=torch.long)\n",
        "        tgt_tensor_ = torch.tensor(\n",
        "            [vocab[token] for token in tokenizer(raw_tgt.rstrip(\"\\n\"))],\n",
        "            dtype=torch.long)\n",
        "        data.append((src_tensor_, tgt_tensor_))\n",
        "    return data\n",
        "\n",
        "train_data = data_process(train_filepaths)\n",
        "val_data = data_process(val_filepaths, )\n",
        "test_data = data_process(test_filepaths)"
      ],
      "metadata": {
        "id": "vnQ0w42uFZUp"
      },
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "BATCH_SIZE = 16\n",
        "PAD_IDX = vocab['<pad>']\n",
        "BOS_IDX = vocab['<bos>']\n",
        "EOS_IDX = vocab['<eos>']"
      ],
      "metadata": {
        "id": "D5GWsa83Fc7i"
      },
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def generate_batch(data_batch):\n",
        "    src_batch, tgt_batch = [], []\n",
        "    for (src_item, tgt_item) in data_batch:\n",
        "        src_batch.append(torch.cat([torch.tensor([BOS_IDX]), src_item, torch.tensor([EOS_IDX])], dim=0))\n",
        "        tgt_batch.append(torch.cat([torch.tensor([BOS_IDX]), tgt_item, torch.tensor([EOS_IDX])], dim=0))\n",
        "    src_batch = pad_sequence(src_batch, padding_value=PAD_IDX)\n",
        "    tgt_batch = pad_sequence(tgt_batch, padding_value=PAD_IDX)\n",
        "    return src_batch, tgt_batch\n",
        "\n",
        "train_iter = DataLoader(train_data, batch_size=BATCH_SIZE,\n",
        "                        shuffle=True, collate_fn=generate_batch)\n",
        "valid_iter = DataLoader(val_data, batch_size=BATCH_SIZE,\n",
        "                        shuffle=True, collate_fn=generate_batch)\n",
        "test_iter = DataLoader(test_data, batch_size=BATCH_SIZE,\n",
        "                       shuffle=True, collate_fn=generate_batch)"
      ],
      "metadata": {
        "id": "0MF7CmUDFdUU"
      },
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Определение самой модели"
      ],
      "metadata": {
        "id": "WnHQm23FFf0A"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Seq2SeqTransformer(nn.Module):\n",
        "    def __init__(self, num_encoder_layers: int, num_decoder_layers: int,\n",
        "                 emb_size: int, src_vocab_size: int, tgt_vocab_size: int,\n",
        "                 dim_feedforward:int = 512, dropout:float = 0.1):\n",
        "        super(Seq2SeqTransformer, self).__init__()\n",
        "        encoder_layer = TransformerEncoderLayer(d_model=emb_size, nhead=NHEAD,\n",
        "                                                dim_feedforward=dim_feedforward)\n",
        "        self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)\n",
        "        decoder_layer = TransformerDecoderLayer(d_model=emb_size, nhead=NHEAD,\n",
        "                                                dim_feedforward=dim_feedforward)\n",
        "        self.transformer_decoder = TransformerDecoder(decoder_layer, num_layers=num_decoder_layers)\n",
        "\n",
        "        self.generator = nn.Linear(emb_size, tgt_vocab_size)\n",
        "        self.src_tok_emb = TokenEmbedding(src_vocab_size, emb_size)\n",
        "        self.tgt_tok_emb = TokenEmbedding(tgt_vocab_size, emb_size)\n",
        "        self.positional_encoding = PositionalEncoding(emb_size, dropout=dropout)\n",
        "\n",
        "    def forward(self, src: Tensor, trg: Tensor, src_mask: Tensor,\n",
        "                tgt_mask: Tensor, src_padding_mask: Tensor,\n",
        "                tgt_padding_mask: Tensor, memory_key_padding_mask: Tensor):\n",
        "        src_emb = self.positional_encoding(self.src_tok_emb(src))\n",
        "        tgt_emb = self.positional_encoding(self.tgt_tok_emb(trg))\n",
        "        memory = self.transformer_encoder(src_emb, src_mask, src_padding_mask)\n",
        "        outs = self.transformer_decoder(tgt_emb, memory, tgt_mask, None,\n",
        "                                        tgt_padding_mask, memory_key_padding_mask)\n",
        "        return self.generator(outs)\n",
        "\n",
        "    def encode(self, src: Tensor, src_mask: Tensor):\n",
        "        return self.transformer_encoder(self.positional_encoding(\n",
        "                            self.src_tok_emb(src)), src_mask)\n",
        "\n",
        "    def decode(self, tgt: Tensor, memory: Tensor, tgt_mask: Tensor):\n",
        "        return self.transformer_decoder(self.positional_encoding(\n",
        "                          self.tgt_tok_emb(tgt)), memory,\n",
        "                          tgt_mask)\n",
        "\n",
        "    \n",
        "class PositionalEncoding(nn.Module):\n",
        "    def __init__(self, emb_size: int, dropout, maxlen: int = 5000):\n",
        "        super(PositionalEncoding, self).__init__()\n",
        "        den = torch.exp(- torch.arange(0, emb_size, 2) * math.log(10000) / emb_size)\n",
        "        pos = torch.arange(0, maxlen).reshape(maxlen, 1)\n",
        "        pos_embedding = torch.zeros((maxlen, emb_size))\n",
        "        pos_embedding[:, 0::2] = torch.sin(pos * den)\n",
        "        pos_embedding[:, 1::2] = torch.cos(pos * den)\n",
        "        pos_embedding = pos_embedding.unsqueeze(-2)\n",
        "\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "        self.register_buffer('pos_embedding', pos_embedding)\n",
        "\n",
        "    def forward(self, token_embedding: Tensor):\n",
        "        return self.dropout(token_embedding +\n",
        "                            self.pos_embedding[:token_embedding.size(0),:])\n",
        "\n",
        "class TokenEmbedding(nn.Module):\n",
        "    def __init__(self, vocab_size: int, emb_size):\n",
        "        super(TokenEmbedding, self).__init__()\n",
        "        self.embedding = nn.Embedding(vocab_size, emb_size)\n",
        "        self.emb_size = emb_size\n",
        "    def forward(self, tokens: Tensor):\n",
        "        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)\n",
        "    \n",
        "# Делаем, так чтобы в обучении не было заглядывания на дальнешие слова\n",
        "def generate_square_subsequent_mask(sz):\n",
        "    mask = (torch.triu(torch.ones((sz, sz), device=DEVICE)) == 1).transpose(0, 1)\n",
        "    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))\n",
        "    return mask\n",
        "\n",
        "def create_mask(src, tgt):\n",
        "    src_seq_len = src.shape[0]\n",
        "    tgt_seq_len = tgt.shape[0]\n",
        "\n",
        "    tgt_mask = generate_square_subsequent_mask(tgt_seq_len)\n",
        "    src_mask = torch.zeros((src_seq_len, src_seq_len), device=DEVICE).type(torch.bool)\n",
        "\n",
        "    src_padding_mask = (src == PAD_IDX).transpose(0, 1)\n",
        "    tgt_padding_mask = (tgt == PAD_IDX).transpose(0, 1)\n",
        "    return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask"
      ],
      "metadata": {
        "id": "HRMVt1D-Fh2p"
      },
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Определяем разные варианты декодирования токенов"
      ],
      "metadata": {
        "id": "fcLEg_LmFqBc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def greedy_decode(model, src, src_mask, max_len, start_symbol, num_samples=1):\n",
        "    src = src.to(DEVICE)\n",
        "    src = torch.cat([src]*num_samples, dim=1)\n",
        "    src_mask = src_mask.to(DEVICE)\n",
        "\n",
        "    memory = model.encode(src, src_mask)\n",
        "    \n",
        "    ys = torch.ones(1, num_samples).fill_(start_symbol).type(torch.long).to(DEVICE)\n",
        "    for i in range(max_len-1):\n",
        "        memory = memory.to(DEVICE)\n",
        "        memory_mask = torch.zeros(ys.shape[0], memory.shape[0]).to(DEVICE).type(torch.bool)\n",
        "        tgt_mask = (generate_square_subsequent_mask(ys.size(0))\n",
        "                                    .type(torch.bool)).to(DEVICE)\n",
        "        out = model.decode(ys, memory, tgt_mask)\n",
        "        out = out.transpose(0, 1)\n",
        "        prob = model.generator(out[:, -1])\n",
        "        _, next_word = torch.max(prob, dim = 1)\n",
        "        next_word = next_word.detach()\n",
        "\n",
        "        ys = torch.cat([ys,\n",
        "                        next_word.view(1, -1)], dim=0)\n",
        "    return ys.transpose(0,1)\n",
        "\n",
        "def sampling_decode(model, src, src_mask, max_len, start_symbol, num_samples=1):\n",
        "    src = src.to(DEVICE)\n",
        "    src = torch.cat([src]*num_samples, dim=1)\n",
        "    src_mask = src_mask.to(DEVICE)\n",
        "\n",
        "    memory = model.encode(src, src_mask)\n",
        "    \n",
        "    ys = torch.ones(1, num_samples).fill_(start_symbol).type(torch.long).to(DEVICE)\n",
        "    for i in range(max_len-1):\n",
        "        memory = memory.to(DEVICE)\n",
        "        memory_mask = torch.zeros(ys.shape[0], memory.shape[0]).to(DEVICE).type(torch.bool)\n",
        "        tgt_mask = (generate_square_subsequent_mask(ys.size(0))\n",
        "                                    .type(torch.bool)).to(DEVICE)\n",
        "        out = model.decode(ys, memory, tgt_mask)\n",
        "        out = out.transpose(0, 1)\n",
        "        prob = model.generator(out[:, -1])\n",
        "        next_word = torch.multinomial(torch.nn.functional.softmax(prob, dim=-1), 1)\n",
        "        next_word = next_word.detach()\n",
        "\n",
        "        ys = torch.cat([ys,\n",
        "                        next_word.view(1, -1)], dim=0)\n",
        "    return ys.transpose(0,1)\n",
        "\n",
        "def paraphrase(model, \n",
        "              srcs, \n",
        "              src_vocab, \n",
        "              tgt_vocab, \n",
        "              src_tokenizer, \n",
        "              decoder=greedy_decode, \n",
        "              ret_tokens=False, \n",
        "              ret_idx=False, \n",
        "              max_len_add=10,\n",
        "              input_idx=False,\n",
        "              **argv):\n",
        "    model.eval()\n",
        "    global_answers = []\n",
        "    for src in srcs:\n",
        "        if not input_idx:\n",
        "            tokens = [BOS_IDX] + [src_vocab.get_stoi()[tok] for tok in src_tokenizer(src)]+ [EOS_IDX]\n",
        "            src = torch.LongTensor(tokens)\n",
        "        num_tokens = len(src)\n",
        "        src = src.reshape(num_tokens, 1)\n",
        "        \n",
        "        src_mask = (torch.zeros(num_tokens, num_tokens)).type(torch.bool)\n",
        "        tgt_tokens = decoder(model, src, src_mask, max_len=num_tokens + max_len_add, start_symbol=BOS_IDX, **argv)\n",
        "\n",
        "        answers = []\n",
        "        for tgt_token in tgt_tokens:\n",
        "            if not ret_idx:\n",
        "                reference = []\n",
        "                for tok in tgt_token:\n",
        "                    if tok.item() == tgt_vocab['<eos>']:\n",
        "                        break\n",
        "                    if tok.item() not in {tgt_vocab['<eos>'], tgt_vocab['<bos>'], tgt_vocab['<pad>']}:\n",
        "                        reference.append(tgt_vocab.get_itos()[tok])\n",
        "                answers.append(\" \".join(reference).strip())\n",
        "                if ret_tokens:\n",
        "                    answers[-1] = answers[-1].split(\" \")\n",
        "            else:\n",
        "                reference = []\n",
        "                for tok in tgt_token:\n",
        "                    if tok.item() == tgt_vocab['<eos>']:\n",
        "                        break\n",
        "                    if tok.item() not in {tgt_vocab['<eos>'], tgt_vocab['<bos>'], tgt_vocab['<pad>']}:\n",
        "                        reference.append(tok.item())\n",
        "                        \n",
        "                answers.append(reference)\n",
        "        global_answers.append(answers)\n",
        "    return global_answers"
      ],
      "metadata": {
        "id": "u5LDhoxwFt25"
      },
      "execution_count": 31,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def evaluate(model, val_iter):\n",
        "    model.eval()\n",
        "    losses = 0\n",
        "    for idx, (src, tgt) in (enumerate(valid_iter)):\n",
        "        src = src.to(DEVICE)\n",
        "        tgt = tgt.to(DEVICE)\n",
        "\n",
        "        tgt_input = tgt[:-1, :]\n",
        "\n",
        "        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)\n",
        "\n",
        "        logits = model(src, tgt_input, src_mask, tgt_mask,\n",
        "                              src_padding_mask, tgt_padding_mask, src_padding_mask)\n",
        "        \n",
        "        tgt_out = tgt[1:,:]\n",
        "        loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))\n",
        "        losses += loss.item()\n",
        "    return losses / len(val_iter)"
      ],
      "metadata": {
        "id": "59Y9vAErFw8m"
      },
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Функции обучения"
      ],
      "metadata": {
        "id": "eJFULvEJFyX9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def train_epoch(model, train_iter, optimizer, loss_fn):\n",
        "    model.train()\n",
        "    losses = 0\n",
        "    for idx, (src, tgt) in enumerate(train_iter):\n",
        "        src = src.to(DEVICE)\n",
        "        tgt = tgt.to(DEVICE)\n",
        "\n",
        "        tgt_input = tgt[:-1, :]\n",
        "\n",
        "        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)\n",
        "\n",
        "        logits = model(src, \n",
        "                       tgt_input, \n",
        "                       src_mask, \n",
        "                       tgt_mask,\n",
        "                       src_padding_mask, \n",
        "                       tgt_padding_mask, \n",
        "                       src_padding_mask)\n",
        "\n",
        "        optimizer.zero_grad()\n",
        "\n",
        "        tgt_out = tgt[1:,:]\n",
        "        loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))\n",
        "        loss.backward()\n",
        "\n",
        "        optimizer.step()\n",
        "        losses += loss.item()\n",
        "    return losses / len(train_iter)"
      ],
      "metadata": {
        "id": "iK0uKzFQF2zR"
      },
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Обучение модели"
      ],
      "metadata": {
        "id": "VTP57JcRF7GT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "SRC_VOCAB_SIZE = len(vocab)\n",
        "TGT_VOCAB_SIZE = len(vocab)\n",
        "EMB_SIZE = 512\n",
        "NHEAD = 8\n",
        "FFN_HID_DIM = 512\n",
        "NUM_ENCODER_LAYERS = 3\n",
        "NUM_DECODER_LAYERS = 3\n",
        "NUM_EPOCHS = 5\n",
        "\n",
        "\n",
        "transformer = Seq2SeqTransformer(NUM_ENCODER_LAYERS, \n",
        "                                 NUM_DECODER_LAYERS,\n",
        "                                 EMB_SIZE, SRC_VOCAB_SIZE, \n",
        "                                 TGT_VOCAB_SIZE,\n",
        "                                 FFN_HID_DIM)\n",
        "\n",
        "for p in transformer.parameters():\n",
        "    if p.dim() > 1:\n",
        "        nn.init.xavier_uniform_(p)\n",
        "\n",
        "transformer = transformer.to(DEVICE)\n",
        "\n",
        "loss_fn = torch.nn.CrossEntropyLoss(ignore_index=PAD_IDX)\n",
        "\n",
        "optimizer = torch.optim.Adam(\n",
        "    transformer.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9\n",
        ")\n",
        "\n",
        "\n",
        "for epoch in range(1, NUM_EPOCHS+1):\n",
        "    start_time = time.time()\n",
        "    train_loss = train_epoch(transformer, train_iter, optimizer, loss_fn)\n",
        "    end_time = time.time()\n",
        "    val_loss = evaluate(transformer, valid_iter)\n",
        "    all_time = time.time()\n",
        "    print(f\"Epoch: {epoch}, \"\n",
        "          f\"Train loss: {train_loss:.3f}, \"\n",
        "          f\"Val loss: {val_loss:.3f}, \"\n",
        "          f\"Epoch time = {(end_time - start_time):.3f}s, \"\n",
        "          f\"All time = {(all_time - start_time):.3f}s\")"
      ],
      "metadata": {
        "id": "6amMr5afF825",
        "outputId": "e32ddded-5ed3-4ae1-c8b6-339ceb910abe",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch: 1, Train loss: 4.584, Val loss: 2.931, Epoch time = 90.025s, All time = 93.271s\n",
            "Epoch: 2, Train loss: 3.442, Val loss: 2.355, Epoch time = 85.002s, All time = 88.477s\n",
            "Epoch: 3, Train loss: 3.002, Val loss: 2.026, Epoch time = 86.008s, All time = 90.066s\n",
            "Epoch: 4, Train loss: 2.691, Val loss: 1.768, Epoch time = 85.935s, All time = 88.315s\n",
            "Epoch: 5, Train loss: 2.446, Val loss: 1.597, Epoch time = 84.769s, All time = 87.147s\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Пример работы (greedy search)"
      ],
      "metadata": {
        "id": "37I3jdEIGFZ9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "paraphrase(transformer, [\"in our work , we focus on supervised domain adaptation .\"], \n",
        "          vocab, \n",
        "          vocab, \n",
        "          tokenizer, \n",
        "          decoder=greedy_decode, num_samples=5)[0]"
      ],
      "metadata": {
        "id": "lIeSk8aEGBxE",
        "outputId": "c9665d52-700e-44c3-c42a-53cdf9b8092e",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 32,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 548 ms, sys: 4.02 ms, total: 552 ms\n",
            "Wall time: 563 ms\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['in this work , we propose a new method for domain adaptation .',\n",
              " 'in this work , we propose a new method for domain adaptation .',\n",
              " 'in this work , we propose a new method for domain adaptation .',\n",
              " 'in this work , we propose a new method for domain adaptation .',\n",
              " 'in this work , we propose a new method for domain adaptation .']"
            ]
          },
          "metadata": {},
          "execution_count": 32
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Пример работы (multinominal sampling)"
      ],
      "metadata": {
        "id": "47P0RPlXGKb1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "paraphrase(transformer, [\"in our work , we focus on supervised domain adaptation .\"], \n",
        "          vocab, \n",
        "          vocab, \n",
        "          tokenizer, \n",
        "          decoder=sampling_decode, num_samples=5)[0]"
      ],
      "metadata": {
        "id": "Kz0rg6KpGIy4",
        "outputId": "594d3dc5-49d2-40c9-cc99-669f78909200",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 33,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 365 ms, sys: 25 µs, total: 365 ms\n",
            "Wall time: 367 ms\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['we will attempt to perform better domain data .',\n",
              " 'we adopted a very state of the art results .',\n",
              " 'in the first , it is given a large way of domain adaptation .',\n",
              " 'by the predicted x appropriate expressed by the domain behavior .',\n",
              " 'in this work , by skipgram training on domain adaptation .']"
            ]
          },
          "metadata": {},
          "execution_count": 33
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xdaa2cKZeCz2"
      },
      "source": [
        "## Переводчик Tensor2Tensor"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "olM9pHpngVnu",
        "outputId": "d5961eab-89f6-45fa-e558-e24d7c8ef5c6"
      },
      "source": [
        "tokenizer = MarianTokenizer.from_pretrained('Helsinki-NLP/opus-mt-ru-en')\n",
        "model = MarianMTModel.from_pretrained('Helsinki-NLP/opus-mt-ru-en')\n",
        "model.to(device)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "MarianMTModel(\n",
              "  (model): BartModel(\n",
              "    (shared): Embedding(62518, 512, padding_idx=62517)\n",
              "    (encoder): BartEncoder(\n",
              "      (embed_tokens): Embedding(62518, 512, padding_idx=62517)\n",
              "      (embed_positions): SinusoidalPositionalEmbedding(512, 512)\n",
              "      (layers): ModuleList(\n",
              "        (0): EncoderLayer(\n",
              "          (self_attn): Attention(\n",
              "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "          )\n",
              "          (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "          (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
              "          (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
              "          (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "        )\n",
              "        (1): EncoderLayer(\n",
              "          (self_attn): Attention(\n",
              "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "          )\n",
              "          (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "          (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
              "          (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
              "          (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "        )\n",
              "        (2): EncoderLayer(\n",
              "          (self_attn): Attention(\n",
              "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "          )\n",
              "          (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "          (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
              "          (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
              "          (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "        )\n",
              "        (3): EncoderLayer(\n",
              "          (self_attn): Attention(\n",
              "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "          )\n",
              "          (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "          (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
              "          (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
              "          (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "        )\n",
              "        (4): EncoderLayer(\n",
              "          (self_attn): Attention(\n",
              "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "          )\n",
              "          (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "          (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
              "          (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
              "          (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "        )\n",
              "        (5): EncoderLayer(\n",
              "          (self_attn): Attention(\n",
              "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "          )\n",
              "          (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "          (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
              "          (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
              "          (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "        )\n",
              "      )\n",
              "      (layernorm_embedding): Identity()\n",
              "    )\n",
              "    (decoder): BartDecoder(\n",
              "      (embed_tokens): Embedding(62518, 512, padding_idx=62517)\n",
              "      (embed_positions): SinusoidalPositionalEmbedding(512, 512)\n",
              "      (layers): ModuleList(\n",
              "        (0): DecoderLayer(\n",
              "          (self_attn): Attention(\n",
              "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "          )\n",
              "          (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "          (encoder_attn): Attention(\n",
              "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "          )\n",
              "          (encoder_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "          (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
              "          (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
              "          (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "        )\n",
              "        (1): DecoderLayer(\n",
              "          (self_attn): Attention(\n",
              "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "          )\n",
              "          (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "          (encoder_attn): Attention(\n",
              "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "          )\n",
              "          (encoder_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "          (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
              "          (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
              "          (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "        )\n",
              "        (2): DecoderLayer(\n",
              "          (self_attn): Attention(\n",
              "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "          )\n",
              "          (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "          (encoder_attn): Attention(\n",
              "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "          )\n",
              "          (encoder_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "          (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
              "          (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
              "          (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "        )\n",
              "        (3): DecoderLayer(\n",
              "          (self_attn): Attention(\n",
              "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "          )\n",
              "          (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "          (encoder_attn): Attention(\n",
              "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "          )\n",
              "          (encoder_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "          (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
              "          (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
              "          (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "        )\n",
              "        (4): DecoderLayer(\n",
              "          (self_attn): Attention(\n",
              "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "          )\n",
              "          (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "          (encoder_attn): Attention(\n",
              "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "          )\n",
              "          (encoder_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "          (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
              "          (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
              "          (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "        )\n",
              "        (5): DecoderLayer(\n",
              "          (self_attn): Attention(\n",
              "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "          )\n",
              "          (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "          (encoder_attn): Attention(\n",
              "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
              "          )\n",
              "          (encoder_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "          (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
              "          (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
              "          (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
              "        )\n",
              "      )\n",
              "      (layernorm_embedding): Identity()\n",
              "    )\n",
              "  )\n",
              ")"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VFsqY73Wj9Bx"
      },
      "source": [
        "batch_x = tokenizer.prepare_seq2seq_batch(src_texts=['Привет мир, меня зовут Трансформер'], return_tensors=\"pt\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UZsTxgZkkDvV",
        "outputId": "5c0abec8-43e9-42fe-ddf9-6df3a28e55d7"
      },
      "source": [
        "tokenizer.batch_decode(model.generate(**batch_x.to(device)))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['<pad> Hey, world, my name is Transformer.']"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z-UYkvAaeCz2"
      },
      "source": [
        "## Токенизация BPE (Byte Pair Encoding)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BEPmEk_pkgkm",
        "outputId": "0dfafb41-0378-4eeb-c149-6cf507c6d1f4"
      },
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/LaBSE', \n",
        "                                          verbose=False)\n",
        "tokenizer.vocab_size"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "501153"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 20
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OTKU4KwBk1Ms",
        "outputId": "b06cd282-3a7c-4402-cd1e-f4af4c26877a"
      },
      "source": [
        "tokenizer.tokenize('Hello Mathematic, my Name is Andrey, how are you?')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['Hello',\n",
              " 'Math',\n",
              " '##emat',\n",
              " '##ic',\n",
              " ',',\n",
              " 'my',\n",
              " 'Name',\n",
              " 'is',\n",
              " 'Andrey',\n",
              " ',',\n",
              " 'how',\n",
              " 'are',\n",
              " 'you',\n",
              " '?']"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 29
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BUkud37HeCz2"
      },
      "source": [
        "## BERT (Bidirectional Encoder Representations from Transformers)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "V91ZW_IclYB8",
        "outputId": "61f9162f-6fea-4c5a-b4ff-e9b5618c0b88"
      },
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained('bert-base-multilingual-cased', \n",
        "                                          verbose=False)\n",
        "model = AutoModel.from_pretrained('bert-base-multilingual-cased')\n",
        "model.to(device)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "BertModel(\n",
              "  (embeddings): BertEmbeddings(\n",
              "    (word_embeddings): Embedding(119547, 768, padding_idx=0)\n",
              "    (position_embeddings): Embedding(512, 768)\n",
              "    (token_type_embeddings): Embedding(2, 768)\n",
              "    (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "    (dropout): Dropout(p=0.1, inplace=False)\n",
              "  )\n",
              "  (encoder): BertEncoder(\n",
              "    (layer): ModuleList(\n",
              "      (0): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (1): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (2): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (3): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (4): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (5): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (6): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (7): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (8): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (9): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (10): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (11): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "    )\n",
              "  )\n",
              "  (pooler): BertPooler(\n",
              "    (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "    (activation): Tanh()\n",
              "  )\n",
              ")"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 31
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dyZ_ZNUFmJLR"
      },
      "source": [
        "### Архитектура BERT\n",
        "Разделяется условно на три части:\n",
        "\n",
        "- Tokens Embedding\n",
        "- Self-Attention\n",
        "- Pooler"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JpTzQVm8mVaV"
      },
      "source": [
        "### Математическая модель BERT\n",
        "Пусть задано множество токенов:\n",
        "$$\n",
        "\\mathcal{I} = \\{\\mathbf{i}| \\mathbf{i} = [0, \\cdots, 0, 1, 0,  \\cdots, 0]^{\\mathsf{T}}\\}\n",
        "$$\n",
        "\n",
        "Задано множество предложений и множество типов токенов в предложении:\n",
        "$$\n",
        "\\mathcal{S} = \\mathcal{I}^n, \\quad \\mathcal{T} = \\{[0,1]^{\\mathsf{T}}, [1,0]^{\\mathsf{T}}\\}^n\n",
        "$$\n",
        "\n",
        "Отображения:\n",
        "$$\n",
        "BM_1: \\mathbb{R}^{n\\times L}\\times \\mathbb{R}^{2\\times L} \\to \\mathbb{R}^{n \\times l}\n",
        "$$\n",
        "$$\n",
        "BM_2: \\mathbb{R}^{n\\times L}\\times \\mathbb{R}^{2\\times L} \\to \\mathbb{R}^{1 \\times l}\n",
        "$$\n",
        "\n",
        "Суперпозиция отображений:\n",
        "$$\n",
        "BM_1 = BL_m \\circ \\cdots \\circ BL_1 \\circ BSE\n",
        "$$\n",
        "$$\n",
        "BM_2 = BP \\circ BL_m \\circ \\cdots \\circ BL_1 \\circ BSE\n",
        "$$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "STvhmoi6mleu"
      },
      "source": [
        "### Математическая модель BERT: BSE\n",
        "\n",
        "Функция $BSE$:\n",
        "$$\n",
        "BSE:\\mathbb{R}^{n\\times L} \\times \\mathbb{R}^{n\\times 2} \\to \\mathbb{R}^{n\\times l}.\n",
        "$$\n",
        "\n",
        "Для произвольной матрицы $\\mathbf{s} \\in \\mathcal{S} \\subset \\mathbb{R}^{n\\times L}$ и матрицы $\\mathbf{t} \\in \\mathcal{T}\\subset \\mathbb{R}^{n\\times 2}$ отображение $BSE$ принимает следующий вид:\n",
        "$$\n",
        "BSE\\bigr(\\mathbf{s}, \\mathbf{t}\\bigr) = \\frac{\\mathbf{h}_{bse} - \\mathsf{E}\\mathbf{h}_{bse}}{\\sqrt{\\mathsf{D}\\mathbf{h}_{bse}+\\varepsilon}}\\cdot\\textbf{w}_1 + \\textbf{w}_2, \\quad \\mathbf{h}_{bse} = \\mathbf{s}\\mathbf{W}_1 + \\mathbf{1}_{n\\times n}\\mathbf{W}_2 + \\mathbf{t}\\mathbf{W}_3,\n",
        "$$\n",
        "где $\\mathbf{W}_1 \\in \\mathbb{R}^{L\\times l},~\\mathbf{W}_2 \\in \\mathbb{R}^{n\\times l},~\\mathbf{W}_3 \\in \\mathbb{R}^{2\\times l}.$\n",
        "\n",
        "Функция $BSE$ имеет настриваемые параметры: $\\mathbf{W}_1, \\mathbf{W}_2, \\mathbf{W}_3, \\mathbf{w}_1, \\mathbf{w}_2.$\n",
        "\n",
        "Результат работы функции $BSE$ обозначим:\n",
        "$$\n",
        "\\mathbf{h}_0 = BSE\\bigr(\\mathbf{s}, \\mathbf{t}\\bigr),\n",
        "$$\n",
        "где $\\mathbf{h} \\in \\mathbb{R}^{n\\times l}.$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fLmx--Alm0ps"
      },
      "source": [
        "### Математическая модель BERT: BL\n",
        "\n",
        "Функция $BL$:\n",
        "$$\n",
        "    BL: \\mathbb{R}^{n\\times l} \\to \\mathbb{R}^{n\\times l}.\n",
        "$$\n",
        "Для матрицы $\\mathbf{h} \\in \\mathbb{R}^{n\\times l}$ BL принимает следующий вид:\n",
        "$$\n",
        "    BL_q\\bigr(\\mathbf{h}\\bigr) = \\frac{\\mathbf{u}\\mathbf{W}_{3+6q} + \\mathbf{a} - \\mathsf{E}\\left(\\mathbf{u}\\mathbf{W}_{3+6q} + \\mathbf{a}\\right)}{\\sqrt{\\mathsf{D}\\left(\\mathbf{u}\\mathbf{W}_{3+6q} + \\mathbf{a}\\right)+\\varepsilon}}\\cdot\\textbf{w}_{3+4q} + \\textbf{w}_{4+4q},\n",
        "$$\n",
        "$$\n",
        "\\mathbf{u} = \\sigma\\bigr(\\mathbf{a}\\mathbf{W}_{4+6q}\\bigr), \\quad \\mathbf{a} = \\frac{\\mathbf{c}\\mathbf{W}_{5+6q} - \\mathsf{E}\\mathbf{c}\\mathbf{W}_{5+6q}}{\\sqrt{\\mathsf{D}\\mathbf{c}\\mathbf{W}_{5+6q}+\\varepsilon}}\\cdot\\textbf{w}_{5+4q} + \\textbf{w}_{6+4q}\n",
        "$$\n",
        "$$\n",
        "\\mathbf{c} = [\\mathbf{c}_1, \\cdots \\mathbf{c}_{r_2}]\n",
        "$$\n",
        "$$\n",
        "\\mathbf{c}_j = \\text{softmax}\\bigr(\\mathbf{h}\\mathbf{W}^{j}_{6+6q}\\odot\\mathbf{h}\\mathbf{W}^{j}_{7+6q}\\bigr)\\odot\\mathbf{h}\\mathbf{W}^{j}_{8+6q}\n",
        "$$\n",
        "где для всех $q$ матрицы $\\mathbf{W}^{j}_{6+6q}, \\mathbf{W}^{j}_{7+6q}, \\mathbf{W}^{j}_{6+6q} \\in \\mathbb{R}^{l \\times r}$, для всех $j$ матрицы $\\mathbf{c}_j \\in \\mathbb{R}^{n\\times r}$, для всех $q$ матрицы $\\mathbf{W}^{j}_{5+6q} \\in \\mathbb{R}^{l \\times l}, \\mathbf{W}^{j}_{4+6q} \\in \\mathbb{R}^{l \\times p}, \\mathbf{W}^{j}_{3+6q} \\in \\mathbb{R}^{p \\times l}$, матрица $\\mathbf{c}, \\mathbf{a} \\in \\mathbb{R}^{n \\times l}$, матрица $\\mathbf{u} \\in \\mathbb{R}^{n \\times p}$.\n",
        "\n",
        "Настраиваемые параметры: $\\mathbf{W}^{j}_{3+6q}, \\mathbf{W}^{j}_{4+6q}, \\mathbf{W}^{j}_{5+6q}, \\mathbf{W}^{j}_{6+6q}, \\mathbf{W}^{j}_{7+6q}, \\mathbf{W}^{j}_{8+6q}, \\textbf{w}_{3+4q}, \\textbf{w}_{4+4q}, \\textbf{w}_{5+4q}, \\textbf{w}_{6+4q}$\n",
        "\n",
        "Результат работы функции $BL_q$:\n",
        "$$\n",
        "\\forall q \\in \\{1, \\cdots m\\} \\quad \\mathbf{h}_{q} = BL_q\\bigr(\\mathbf{h}_{q-1}\\bigr).\n",
        "$$\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZRI0Gz_enAwL"
      },
      "source": [
        "### Математическая модель BERT: BP\n",
        "\n",
        "Функция $BP$:\n",
        "$$\n",
        "    BP: \\mathbb{R}^{n\\times l} \\to \\mathbb{R}^{n\\times l}.\n",
        "$$\n",
        "Для матрицы $\\mathbf{h}_{m} \\in \\mathbb{R}^{n \\times l}$ BP принимает следующий вид:\n",
        "$$\n",
        "BP\\bigr(\\mathbf{h}_{m}\\bigr) = \\sigma\\bigr(\\mathbf{h}_m^{1}\\mathbf{W}_{9+6m}\\bigr),\n",
        "$$\n",
        "где $\\mathbf{h}_m^{1}$ первая строка матрицы $\\mathbf{h}_{m}$, а матрица $\\mathbf{W}_{9+6m} \\in \\mathbb{R}^{l\\times l}$\n",
        "\n",
        "Функция $BP$ имеет настраиваемые параметры $\\mathbf{W}_{9+6m}$\n",
        "\n",
        "Результат работы функции $BP$:\n",
        "$$\n",
        "\\mathbf{h} = BP\\bigr(\\mathbf{h}_m\\bigr).\n",
        "$$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "My27CiQOnJuD"
      },
      "source": [
        "### Снова суперпозиция\n",
        "\n",
        "Вернемся к суперпозициям $BM_1, BM_2$:\n",
        "$$\n",
        "\\mathbf{h} = BP\\bigr(\\mathbf{h}_m\\bigr), \\quad\\mathbf{h}_q = BM\\bigr(\\mathbf{h}_{q-1} \\bigr), \\quad \\mathbf{h}_0 = BSE\\bigr(\\mathbf{s}, \\mathbf{t}\\bigr)\n",
        "$$\n",
        "\n",
        "Получаем вектор эмбедингов слов:\n",
        "$$\n",
        "BM_1\\bigr(\\mathbf{s}, \\mathbf{t}\\bigr) = \\mathbf{h}_m,\n",
        "$$\n",
        "Получаем вектор эмбединга предложения:\n",
        "$$\n",
        "BM_2\\bigr(\\mathbf{s}, \\mathbf{t}\\bigr) = \\mathbf{h}\n",
        "$$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "taYtx4v9nP0X"
      },
      "source": [
        "### Multitask Learning\n",
        "\n",
        "LM модель:\n",
        "$$\n",
        "\\mathbf{v} = \\text{softmax}\\bigr(\\mathbf{h}_m\\mathbf{W}_{LM}\\bigr),\n",
        "$$\n",
        "где $\\mathbf{W}_{LM} \\in \\mathbb{R}^{l \\times L}$, а $\\mathbf{v}$ это вероятность каждого токена.\n",
        "\n",
        "NSP модель:\n",
        "$$\n",
        "z = \\sigma\\bigr(\\mathbf{h}\\mathbf{W}_{NSP}\\bigr),\n",
        "$$\n",
        "где $\\mathbf{W}_{NSP} \\in \\mathbb{R}^{l \\times 1}$, а $z$ это вероятность класса $1$.\n",
        "\n",
        "Функция ошибки:\n",
        "$$\n",
        "L\\bigr(\\mathbf{S}, \\mathbf{y}\\bigr) = \\sum_{\\mathbf{s}_i, \\mathbf{t}_i \\in \\mathbf{S}}CrossEntropy\\bigr(\\mathbf{v}_i, \\mathbf{s}_i\\bigr) + \\sum_{\\mathbf{s}_i, \\mathbf{t}_i \\in \\mathbf{S}, y_i \\mathbf{y}}CrossEntropyLoss\\bigr(z_i, y_i\\bigr)\n",
        "$$\n",
        "Задача оптимизации:\n",
        "$$\n",
        "L\\bigr(\\mathbf{S}, \\mathbf{y}\\bigr) \\to \\min_{\\mathbf{W}_{all}}\n",
        "$$\n",
        "\n",
        "Все параметры:\n",
        "$$\n",
        "\\mathbf{W}_{all} = [\\mathbf{W}_{LM}, \\mathbf{W}_{NSP}, \\mathbf{W}_{9+6m}\n",
        "\\mathbf{W}^{j}_{3+6q}, \\mathbf{W}^{j}_{4+6q}, \\mathbf{W}^{j}_{5+6q}, \\mathbf{W}^{j}_{6+6q}, \\mathbf{W}^{j}_{7+6q}, \\mathbf{W}^{j}_{8+6q}, \\textbf{w}_{3+4q}, \\textbf{w}_{4+4q}, \\textbf{w}_{5+4q}, \\textbf{w}_{6+4q},\n",
        "\\mathbf{W}_1, \\mathbf{W}_2, \\mathbf{W}_3, \\mathbf{w}_1, \\mathbf{w}_2]\n",
        "$$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z_zdJ-cGmxYx"
      },
      "source": [
        "tokens = tokenizer(['Hello World', 'Andrey Grabovoy'], return_tensors='pt', padding=True)\n",
        "\n",
        "with torch.no_grad():\n",
        "    output = model(**tokens.to(device))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Jqr9nnkAoVuJ",
        "outputId": "f6cafa0c-e7af-4653-a50d-fe190765cd90"
      },
      "source": [
        "tokens['input_ids']"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[  101, 31178, 10315,   102,     0,     0,     0],\n",
              "        [  101, 25188, 10157, 61020, 28194, 10157,   102]], device='cuda:0')"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 49
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0VvmbaWfoN1W",
        "outputId": "9ad82b3a-0f73-4cda-e353-7402f12a22f1"
      },
      "source": [
        "output[0].shape, output[1].shape"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(torch.Size([2, 7, 768]), torch.Size([2, 768]))"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 51
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7KUZ6byVov7M",
        "outputId": "c318d73c-db55-4851-ed75-580e566d4ddd"
      },
      "source": [
        "tokens"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'input_ids': tensor([[  101, 31178, 10315,   102,     0,     0,     0],\n",
              "        [  101, 25188, 10157, 61020, 28194, 10157,   102]], device='cuda:0'), 'token_type_ids': tensor([[0, 0, 0, 0, 0, 0, 0],\n",
              "        [0, 0, 0, 0, 0, 0, 0]], device='cuda:0'), 'attention_mask': tensor([[1, 1, 1, 1, 0, 0, 0],\n",
              "        [1, 1, 1, 1, 1, 1, 1]], device='cuda:0')}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 52
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "keUHDcaXeCz2"
      },
      "source": [
        "## Модель LaBSE (Language-agnostic BERT Sentence Embedding)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000,
          "referenced_widgets": [
            "ab1befa3fa404373afc824f2824cb1a0",
            "6639bf4ef82a45a2a116cf1e7910dee7",
            "6ad04290250b4439a3ba36dcce626608",
            "03906703a09c48aba9a06c15f552927c",
            "19bc31c66b8f46928dd91084d8691f2f",
            "9682fcbcb72e44d38c20b4cfac5407f3",
            "a7a9e6f1dcfa4e6a911f6f440c46d890",
            "b0772ff6976d402fb3fa61f0778cfd48",
            "c207a78178824e3e808cb65f54f45994",
            "0b2d4f3f73074592aa6a94a33e7973a2",
            "401cca8b41f04dfdb83e0decc1b50881"
          ]
        },
        "id": "QEsAYf4DlYae",
        "outputId": "1bc7ca3e-5362-4b29-bcb2-c4ad2f690942"
      },
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/LaBSE', \n",
        "                                          verbose=False)\n",
        "model = AutoModel.from_pretrained('sentence-transformers/LaBSE')\n",
        "model.to(device)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "ab1befa3fa404373afc824f2824cb1a0",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "Downloading:   0%|          | 0.00/1.89G [00:00<?, ?B/s]"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "BertModel(\n",
              "  (embeddings): BertEmbeddings(\n",
              "    (word_embeddings): Embedding(501153, 768, padding_idx=0)\n",
              "    (position_embeddings): Embedding(512, 768)\n",
              "    (token_type_embeddings): Embedding(2, 768)\n",
              "    (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "    (dropout): Dropout(p=0.1, inplace=False)\n",
              "  )\n",
              "  (encoder): BertEncoder(\n",
              "    (layer): ModuleList(\n",
              "      (0): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (1): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (2): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (3): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (4): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (5): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (6): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (7): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (8): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (9): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (10): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "      (11): BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "    )\n",
              "  )\n",
              "  (pooler): BertPooler(\n",
              "    (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "    (activation): Tanh()\n",
              "  )\n",
              ")"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 53
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JFHMMRvJo0i6"
      },
      "source": [
        "tokens = tokenizer(['Moscow Institute of Physics and Technology', \n",
        "                    'Московский Физико-Технический Институт', \n",
        "                    'Московский Государственный Университет', \n",
        "                    'Moscow State University'], \n",
        "                   return_tensors='pt', padding=True)\n",
        "\n",
        "with torch.no_grad():\n",
        "    output = model(**tokens.to(device))\n",
        "\n",
        "sentence_embeding = output[1].cpu().numpy()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9PVQeA1NpUR3",
        "outputId": "45a1950f-d875-449c-c6c5-b592038a770f"
      },
      "source": [
        "sentence_embeding.shape"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(4, 768)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 66
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Aqhb5ZKMpWIy"
      },
      "source": [
        "frame = pd.DataFrame(np.round(\n",
        "    scipy.spatial.distance.cdist(sentence_embeding, sentence_embeding, \n",
        "                                 metric='cosine'), 2))\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 265
        },
        "id": "KjgPCRq2pd32",
        "outputId": "529c9efc-9a76-4710-d583-e390b4096f14"
      },
      "source": [
        "sns.heatmap(data=frame)\n",
        "\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAWQAAAD4CAYAAADbyJysAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAATJ0lEQVR4nO3df6zdd33f8ecrBicb2Whaqjaz02CKWRvWKRGJ+QMtoJIQU6o4VYNquq7pFM3rhDcqVK1G0LCasdJ2Yvtj6RprWCpVi5uSqlxRd1EKYVPHAtchaaidedyYKrkWLQrJyqLQJPfed/+4X5Lj2+t7zrXP9fmc730+rj7y9/v5/vroyH7ft9/fz/d7UlVIkibvokkPQJK0zIAsSY0wIEtSIwzIktQIA7IkNeJlG36BrducxtH5wOVvmfQQmvH+Bz806SE0Y/GrD016CM245Lofz/me44UnT40cc17+qtec9/XGyQxZkhqx4RmyJF1QS4uTHsE5MyBL6pfFhUmP4JwZkCX1StXSpIdwzgzIkvplyYAsSW0wQ5akRnhTT5IaYYYsSW0oZ1lIUiO8qSdJjbBkIUmN8KaeJDXCDFmSGuFNPUlqhDf1JKkNVdaQJakN1pAlqRGWLCSpEWbIktSIxRcmPYJzZkCW1C+WLCSpEVNcsvBbpyX1y9LS6G2IJLuTnEwyl+TAKtt/NsmXkzyc5E+SXDWw7X3dcSeT3DTK0M2QJfXLmEoWSbYAdwI3AvPAbJKZqjoxsNvvVNVvdPvfDHwU2N0F5r3A64F/APxxktfVkEnSZsiSeqUWXxi5DbELmKuqU1X1PHAE2HPGtaq+ObD6CqC65T3Akap6rqq+Csx151vT0Aw5yQ90J9/WdZ0GZqrq0WHHStIFt44acpJ9wL6BrkNVdahb3gY8MbBtHnjjKud4N/BeYCvwwwPHPrDi2G0MsWaGnOQXWP6tEOCLXQvwidXqKZI0ceuoIVfVoaq6dqAdGn6BM1XVnVX1/cAvAB84n6EPy5BvB15fVWfk9kk+ChwHPrLaQYO/dbLllVx00SvOZ4ySNLrxzbI4DVwxsL696zubI8B/PcdjgeE15CWWC9IrXd5tW9Xgbx2DsaQLanyzLGaBnUl2JNnK8k26mcEdkuwcWH0H8JVueQbYm+TiJDuAnSxXGNY0LEP+OeAzSb7CS7WU7wNeC+wfdnJJuuDGlCFX1UKS/cC9wBbgcFUdT3IQOFZVM8D+JDcALwBPA7d1xx5PcjdwAlgA3j1shgUMCchV9d+TvI7lu4ODN/VmRzm5JF1wC+N7QX1VHQWOrui7Y2D5PWsc+2Hgw+u53tBZFlW1xJl3CyWpXVP8pJ4PhkjqF99lIUmNMEOWpEaYIUtSI8yQJakRY5xlcaEZkCX1S9XwfRplQJbUL9aQJakRBmRJaoQ39SSpEYvT+1YHA7KkfrFkIUmNMCBLUiOsIUtSG2rJeciS1AZLFpLUCGdZSFIjzJAlqREGZElqhC8XkqRGmCFLUiOc9nZ2H7j8LRt9ianx77/2uUkPoRn/9uO/POkhNGPLTXsnPYR+meJZFhdNegCSNE61tDRyGybJ7iQnk8wlObDK9vcmOZHkkSSfSXLlwLbFJA93bWaUsVuykNQvYypZJNkC3AncCMwDs0lmqurEwG4PAddW1bNJ/hXwq8BPdNu+VVVXr+eaZsiS+qWWRm9r2wXMVdWpqnoeOALsOeNSVfdX1bPd6gPA9vMZugFZUr8s1cgtyb4kxwbavoEzbQOeGFif7/rO5nbgjwbWL+nO+UCSW0YZuiULSf2yMPpNvao6BBw630sm+SngWuDNA91XVtXpJK8BPpvky1X12FrnMSBL6pfxvX7zNHDFwPr2ru8MSW4A3g+8uaqee3EYVae7P08l+RxwDbBmQLZkIalf1lGyGGIW2JlkR5KtwF7gjNkSSa4B7gJurqqvD/RfluTibvlVwJuAwZuBqzJDltQro0xnG+k8VQtJ9gP3AluAw1V1PMlB4FhVzQC/BlwK/F4SgMer6mbgB4G7kiyxnPh+ZMXsjFUZkCX1yxif1Kuqo8DRFX13DCzfcJbjPg/80HqvZ0CW1C8+Oi1JjZjiR6cNyJJ6xe/Uk6RWGJAlqRG+D1mSGmGGLEmNMCBLUhtq0ZKFJLXBDFmS2uC0N0lqhQFZkhoxvSVkA7KkfqmF6Y3IBmRJ/TK98diALKlfvKknSa0wQ5akNkxzhnzO36mX5J+PcyCSNBZL62iNOZ8vOf2ls21Isi/JsSTHHnxm7jwuIUnrUwujt9asWbJI8sjZNgHfc7bjquoQcAjg3135T6f3/w+Spk41mPmOalgN+XuAm4CnV/QH+PyGjEiSzkePA/KngUur6uGVG5J8bkNGJEnnobcZclXdvsa2nxz/cCTp/ExzQD6fm3qS1JxazMhtmCS7k5xMMpfkwCrb35vkRJJHknwmyZUD225L8pWu3TbK2A3IknqllkZva0myBbgTeDtwFfCuJFet2O0h4Nqq+sfAJ4Ff7Y79TuCDwBuBXcAHk1w2bOwGZEm9UksZuQ2xC5irqlNV9TxwBNhzxrWq7q+qZ7vVB4Dt3fJNwH1V9VRVPQ3cB+wedkEDsqReWU+GPPjMRNf2DZxqG/DEwPp813c2twN/dI7HAj46LalnqobXhl/a96VnJs5Hkp8CrgXefD7nMUOW1CvjqiEDp4ErBta3d31nSHID8H7g5qp6bj3HrmRAltQrS4sZuQ0xC+xMsiPJVmAvMDO4Q5JrgLtYDsZfH9h0L/C2JJd1N/Pe1vWtyZKFpF4Z4WbdaOepWkiyn+VAugU4XFXHkxwEjlXVDPBrwKXA7yUBeLyqbq6qp5J8iOWgDnCwqp4adk0DsqReGVdABqiqo8DRFX13DCzfsMaxh4HD67meAVlSr9QUv87MgCypV8aZIV9oBmRJvbKeaW+tMSBL6pXFEd5R0SoDsqReMUOWpEZYQ5akRjjLQpIaYYYsSY1YXJreN0IYkCX1iiULSWrEkrMsJKkNTnuTpEZMc8kitcGjf+HJU1P88YzXCx//5UkPoRl//8DR4TttEu/43msmPYRmfOrxT593ents+y0jx5xr5/+gqXTaDFlSrzjLQpIaMc3/JTcgS+oVZ1lIUiOcZSFJjRj+ZdLtMiBL6pXCDFmSmrBgyUKS2mCGLEmNmOYa8vTOoJakVRQZuQ2TZHeSk0nmkhxYZfv1Sb6UZCHJrSu2LSZ5uGszo4zdDFlSr4wrQ06yBbgTuBGYB2aTzFTViYHdHgd+Bvj5VU7xraq6ej3XNCBL6pXF8dWQdwFzVXUKIMkRYA/wYkCuqj/vto3l94AlC0m9spTRW5J9SY4NtH0Dp9oGPDGwPt/1jeqS7pwPJLlllAPMkCX1ytI6MuSqOgQc2qChXFlVp5O8Bvhski9X1WNrHWCGLKlXah1tiNPAFQPr27u+0cZRdbr78xTwOWDoe1YNyJJ6ZWkdbYhZYGeSHUm2AnuBkWZLJLksycXd8quANzFQez4bA7KkXllKRm5rqaoFYD9wL/AocHdVHU9yMMnNAEmuSzIPvBO4K8nx7vAfBI4l+VPgfuAjK2ZnrMoasqReWRzjuarqKHB0Rd8dA8uzLJcyVh73eeCH1ns9A7KkXlma3ienDciS+mU9syxaY0CW1Ct+hZMkNcKShSQ1Yprf9mZAltQri1OcIQ+dh5zkB5K8NcmlK/p3b9ywJOncjPHBkAtuzYCc5N8AnwL+NfBnSfYMbP4PGzkwSToXvQ3IwL8A3lBVtwBvAX4xyXu6bWf9j8HgG5T+28c/MZ6RStIIKqO31gyrIV9UVc/A8ns/k7wF+GSSK1kjIA++QemFJ09N8ywUSVOmxcx3VMMy5L9M8uIb77vg/KPAqziHxwIlaaMtrqO1ZliG/NPAwmBH98KNn05y14aNSpLOUW/nIVfV/Brb/tf4hyNJ52eaSxbOQ5bUKwZkSWrENM8iMCBL6pXe1pAladq0OHtiVAZkSb2yNMVFCwOypF7xpp4kNWJ682MDsqSeMUOWpEYsZHpzZAOypF6Z3nA8wgvqJWmajPN9yEl2JzmZZC7JgVW2X5/kS0kWkty6YtttSb7StdtGGbsZsqReGde0tyRbgDuBG4F5YDbJTFWdGNjtceBngJ9fcex3Ah8ErmU5aX+wO/bpta5phiypV2odbYhdwFxVnaqq54EjwOC3JlFVf15Vj/C3E+6bgPuq6qkuCN8HDP3aOwOypF5ZT8li8NuNurZv4FTbgCcG1ue7vlGc07GWLCT1yuI6ShaD327UAjNkSb0yxpt6p4ErBta3d32jOKdjDciSeqXW8TPELLAzyY4kW4G9wMyIw7gXeFuSy5JcBryt61uTAVlSr4wrQ+6+rm4/y4H0UeDuqjqe5GCSmwGSXJdkHngncFeS492xTwEfYjmozwIHu741WUOW1CvjfNtbVR0Fjq7ou2NgeZblcsRqxx4GDq/negZkSb0yzU/qGZAl9crCFIdkA7KkXhnhZl2zNjwgL371oY2+xNTYctPeSQ+hGe/4z1+b9BCa8Yd/4b+RcfL1m5LUCDNkSWqEGbIkNWKxzJAlqQl+67QkNcIasiQ1whqyJDXCkoUkNcKShSQ1wlkWktQISxaS1Ahv6klSI6whS1IjLFlIUiPKm3qS1IZFM2RJaoMlC0lqhCULSWqEGbIkNWKap71dNOkBSNI4LVaN3IZJsjvJySRzSQ6ssv3iJL/bbf9Ckld3/a9O8q0kD3ftN0YZuxmypF4ZV8kiyRbgTuBGYB6YTTJTVScGdrsdeLqqXptkL/ArwE902x6rqqvXc00zZEm9skSN3IbYBcxV1amqeh44AuxZsc8e4De75U8Cb02Scx27AVlSr1TVyC3JviTHBtq+gVNtA54YWJ/v+lhtn6paAP4K+K5u244kDyX5H0n+yShjt2QhqVfWU7KoqkPAoQ0YxteA76uqbyR5A/AHSV5fVd9c6yAzZEm9Uuv4GeI0cMXA+vaub9V9krwMeCXwjap6rqq+AVBVDwKPAa8bdsGhATnJriTXdctXJXlvkh8ZdpwkTcJiLY3chpgFdibZkWQrsBeYWbHPDHBbt3wr8NmqqiTf3d0UJMlrgJ3AqWEXXLNkkeSDwNuBlyW5D3gjcD9wIMk1VfXhYReQpAtpXE/qVdVCkv3AvcAW4HBVHU9yEDhWVTPAx4DfSjIHPMVy0Aa4HjiY5AWWX9H8s1X11LBrDqsh3wpcDVwM/AWwvaq+meQ/Al8AVg3IXWF8H8B/ed+/5PYfu3HYOCRpLMb5pF5VHQWOrui7Y2D5r4F3rnLcPcA9673esIC8UFWLwLNJHvt2QbqqvpXkrPn+YKH8r2fvmd7HZiRNnWl+Um9YQH4+yd+tqmeBN3y7M8krme5vSpHUU0s9frnQ9VX1HEDVGRXwl/NSIVuSmtHbDPnbwXiV/ieBJzdkRJJ0HkaYPdEsHwyR1Ct9LllI0lTpbclCkqaNGbIkNcIMWZIasViLkx7COTMgS+oVv+RUkhrhl5xKUiPMkCWpEc6ykKRGOMtCkhrho9OS1AhryJLUCGvIktQIM2RJaoTzkCWpEWbIktQIZ1lIUiO8qSdJjZjmksVFkx6AJI1TreNnmCS7k5xMMpfkwCrbL07yu932LyR59cC293X9J5PcNMrYDciSeqWqRm5rSbIFuBN4O3AV8K4kV63Y7Xbg6ap6LfCfgF/pjr0K2Au8HtgN/Hp3vjUZkCX1ylLVyG2IXcBcVZ2qqueBI8CeFfvsAX6zW/4k8NYk6fqPVNVzVfVVYK4735o2vIZ8yXU/no2+xiiS7KuqQ5MeRwta+Cw+9fhbJ3n5F7XwWbSiL5/FwvOnR445SfYB+wa6Dg18BtuAJwa2zQNvXHGKF/epqoUkfwV8V9f/wIpjtw0bz2bKkPcN32XT8LN4iZ/FSzbdZ1FVh6rq2oE20V9ImykgS9J6nAauGFjf3vWtuk+SlwGvBL4x4rF/iwFZklY3C+xMsiPJVpZv0s2s2GcGuK1bvhX4bC3fLZwB9nazMHYAO4EvDrvgZpqHPPW1sTHys3iJn8VL/CwGdDXh/cC9wBbgcFUdT3IQOFZVM8DHgN9KMgc8xXLQptvvbuAEsAC8u2r412FnmidRS1KfWLKQpEYYkCWpEb0PyMMefdxMkhxO8vUkfzbpsUxSkiuS3J/kRJLjSd4z6TFNSpJLknwxyZ92n8UvTXpMm1mva8jdo4r/F7iR5YnZs8C7qurERAc2IUmuB54BPl5V/2jS45mUJJcDl1fVl5L8PeBB4JbN+Peie6rsFVX1TJKXA38CvKeqHhhyqDZA3zPkUR593DSq6n+yfCd4U6uqr1XVl7rl/w88yghPUfVRLXumW3151/qbpTWu7wF5tUcfN+U/PK2uezvXNcAXJjuSyUmyJcnDwNeB+6pq034Wk9b3gCydVZJLgXuAn6uqb056PJNSVYtVdTXLT5PtSrJpy1mT1veAfE6PL6r/unrpPcBvV9XvT3o8Laiq/wfcz/LrIjUBfQ/Iozz6qE2mu5H1MeDRqvropMczSUm+O8l3dMt/h+Ub4P9nsqPavHodkKtqAfj2o4+PAndX1fHJjmpyknwC+N/AP0wyn+T2SY9pQt4E/DPgh5M83LUfmfSgJuRy4P4kj7CcwNxXVZ+e8Jg2rV5Pe5OkadLrDFmSpokBWZIaYUCWpEYYkCWpEQZkSWqEAVmSGmFAlqRG/A3AmFDYCU1hUAAAAABJRU5ErkJggg==\n",
            "text/plain": [
              "<Figure size 432x288 with 2 Axes>"
            ]
          },
          "metadata": {
            "tags": [],
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "D3isf_RGqKQS"
      },
      "source": [],
      "execution_count": null,
      "outputs": []
    }
  ]
}