notebooks/machine_translation/MT_transformer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a0086f61",
   "metadata": {
    "id": "a0086f61"
   },
   "source": [
    "# Дисклеймер\n",
    "Эту тетрадку нужно запускать в колабе или на машине с gpu."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3b0ea3c",
   "metadata": {
    "id": "e3b0ea3c"
   },
   "outputs": [],
   "source": [
    "# !apt-get install unzip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "a9220b9d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.1.1\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ[\"KERAS_BACKEND\"] = \"torch\"\n",
    "\n",
    "import keras\n",
    "import torch\n",
    "print(keras.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d650e9eb",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "d650e9eb",
    "outputId": "16aecd0e-9060-4289-f372-c1a2372d1953",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# %pip install tokenizers matplotlib scikit-learn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b96ff225",
   "metadata": {},
   "source": [
    "Большая часть этого семинара пересекается с видео от Andrey Karpathy - https://www.youtube.com/watch?v=kCc8FmEb1nY Возможно его объяснения покажутся вам лучше. \n",
    "Его видео более детальное и объемное, но фокус в нем скорее на GPT (декодере) для простой генерации текста, и еще оно на torch. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8559a385",
   "metadata": {
    "id": "8559a385"
   },
   "source": [
    "# Транформеры для машинного перевода"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "WjM0-Q75xtgk",
   "metadata": {
    "id": "WjM0-Q75xtgk"
   },
   "source": [
    "На huggingface очень много предобученных трансформеров, которые позволят вам решить очень большой процент рабочих/исследовательских задач. Однако бывают ситуации, когда нужной предобученной модели нет или она работает не очень хорошо. Обучить очень большую модель вряд ли получится (нужны видеокарты), но вот для средних моделей и специфичных задач может хватить даже колаба/каггла."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f19d2aa7",
   "metadata": {
    "id": "f19d2aa7"
   },
   "source": [
    "Давайте обучим свой небольшой трансформер на задаче машинного перевода. Это самая классическая seq2seq задача и подход, который мы размерем, применим и для других аналогичных задач. Для того, чтобы обучить модель решать какую-то другую seq2seq задачу вам нужно будет только подставить другие данные. Под seq2seq вообще при желании можно подвести любую задачу (классификация - это предсказание последовательности длинной 1 или можно предсказыть названия классов текстом). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "947b3313",
   "metadata": {
    "id": "947b3313"
   },
   "outputs": [],
   "source": [
    "# import tensorflow as tf\n",
    "import keras\n",
    "from tokenizers import BertWordPieceTokenizer\n",
    "\n",
    "from tokenizers import Tokenizer\n",
    "from tokenizers.models import WordPiece\n",
    "from tokenizers.pre_tokenizers import Whitespace\n",
    "from tokenizers import normalizers\n",
    "from tokenizers.normalizers import Lowercase\n",
    "from tokenizers.trainers import WordPieceTrainer\n",
    "from tokenizers import decoders\n",
    "\n",
    "import os\n",
    "import re\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.model_selection import StratifiedShuffleSplit, train_test_split\n",
    "from string import punctuation\n",
    "from collections import Counter\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9c8a8cc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "508e8ae0",
   "metadata": {
    "id": "508e8ae0"
   },
   "source": [
    "Данные взяты вот отсюда - https://opus.nlpl.eu/opus-100.php (раздел с отдельными языковыми парами)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 337,
   "id": "kvUX7tdtc5pZ",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kvUX7tdtc5pZ",
    "outputId": "3a1609b9-4478-442e-9437-3d7b65544c84"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-03-20 21:26:41--  https://data.statmt.org/opus-100-corpus/v1.0/supervised/en-ru/opus.en-ru-test.ru\n",
      "Resolving data.statmt.org (data.statmt.org)... 129.215.32.28\n",
      "Connecting to data.statmt.org (data.statmt.org)|129.215.32.28|:443... "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 305669 (299K)\n",
      "Saving to: ‘opus.en-ru-test.ru’\n",
      "\n",
      "opus.en-ru-test.ru  100%[===================>] 298.50K   345KB/s    in 0.9s    \n",
      "\n",
      "2024-03-20 21:26:43 (345 KB/s) - ‘opus.en-ru-test.ru’ saved [305669/305669]\n",
      "\n",
      "--2024-03-20 21:26:43--  https://data.statmt.org/opus-100-corpus/v1.0/supervised/en-ru/opus.en-ru-test.en\n",
      "Resolving data.statmt.org (data.statmt.org)... 129.215.32.28\n",
      "Connecting to data.statmt.org (data.statmt.org)|129.215.32.28|:443... "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 173307 (169K)\n",
      "Saving to: ‘opus.en-ru-test.en’\n",
      "\n",
      "opus.en-ru-test.en  100%[===================>] 169.25K   245KB/s    in 0.7s    \n",
      "\n",
      "2024-03-20 21:26:44 (245 KB/s) - ‘opus.en-ru-test.en’ saved [173307/173307]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# !wget https://data.statmt.org/opus-100-corpus/v1.0/supervised/en-ru/opus.en-ru-train.ru\n",
    "# !wget https://data.statmt.org/opus-100-corpus/v1.0/supervised/en-ru/opus.en-ru-train.en\n",
    "!wget https://data.statmt.org/opus-100-corpus/v1.0/supervised/en-ru/opus.en-ru-test.ru\n",
    "!wget https://data.statmt.org/opus-100-corpus/v1.0/supervised/en-ru/opus.en-ru-test.en"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "38911d06",
   "metadata": {
    "id": "38911d06"
   },
   "outputs": [],
   "source": [
    "# в русскоязычных данных есть \\xa0 вместо пробелов, он может некорректно обрабатываться токенизатором\n",
    "text = open('opus.en-ru-train.ru').read().replace('\\xa0', ' ')\n",
    "f = open('opus.en-ru-train.ru', 'w')\n",
    "f.write(text)\n",
    "f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 412,
   "id": "e110ff04",
   "metadata": {
    "id": "e110ff04"
   },
   "outputs": [],
   "source": [
    "en_sents = open('opus.en-ru-train.en').read().lower().splitlines()\n",
    "ru_sents = open('opus.en-ru-train.ru').read().lower().splitlines()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35b8fae3",
   "metadata": {
    "id": "35b8fae3"
   },
   "source": [
    "Пример перевода с английского на русский"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0eb9b498",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0eb9b498",
    "outputId": "277adddf-1f07-4960-9b0f-b8377f6add0b"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('so what are you thinking?', 'ну и что ты думаешь?')"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "en_sents[-1], ru_sents[-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31bbf529",
   "metadata": {
    "id": "31bbf529"
   },
   "source": [
    "Как обычно нам нужен токенизатор, а точнее даже 2, т.к. у нас два корпуса. Будем использовать wordpiece токенайзер из библиотеки tokenizers от huggingface. Wordpiece разбивает текст на токены, которые могут быть как целыми словами, так и кусками слов и даже отдельными символами. \n",
    "В токенизатор мы передаем еще несколько специальных символов - UNK (это символ замена для неизвестных слов), START - этот токен будет добавлен в начало текста, END - этот символ будет добавлен в конце, PAD - это индекс паддинга\n",
    "\n",
    "В токенезаторе для целевого языка по идее нам не нужен UNK, потому что в переводном тексте технических токенов быть не должно, но мы не можем без него обойтись, потому что мы ограничиваем размер словаря. Можно было бы построить словарь так, чтобы в нем были все возможные токены и тогда UNK был бы не нужен, но тогда размер словаря будет слишком большим. Но мы всегда можем отфильтровать UNK из перевода, просто не разрешая модели предсказывать UNK."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0be060a4",
   "metadata": {
    "id": "0be060a4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "tokenizer_en = Tokenizer(WordPiece(), )\n",
    "tokenizer_en.normalizer = normalizers.Sequence([Lowercase()])\n",
    "tokenizer_en.pre_tokenizer = Whitespace()\n",
    "\n",
    "trainer_en = WordPieceTrainer(\n",
    "          vocab_size=30000, special_tokens=[\"[UNK]\", \"[PAD]\"])\n",
    "tokenizer_en.train(files=[\"opus.en-ru-train.en\"], trainer=trainer_en )\n",
    "\n",
    "tokenizer_ru = Tokenizer(WordPiece(), )\n",
    "tokenizer_ru.normalizer = normalizers.Sequence([Lowercase()])\n",
    "tokenizer_ru.pre_tokenizer = Whitespace()\n",
    "\n",
    "trainer_ru = WordPieceTrainer(\n",
    "          vocab_size=30000, special_tokens=[\"[UNK]\", \"[PAD]\", \"[START]\", \"[END]\", ])\n",
    "tokenizer_ru.train(files=[\"opus.en-ru-train.ru\"], trainer=trainer_ru )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "xWX3xnMUzbdt",
   "metadata": {
    "id": "xWX3xnMUzbdt"
   },
   "outputs": [],
   "source": [
    "tokenizer_en.decoder = decoders.WordPiece()\n",
    "tokenizer_ru.decoder = decoders.WordPiece()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "FZLOU_Qe4Tw-",
   "metadata": {
    "id": "FZLOU_Qe4Tw-"
   },
   "source": [
    "Когда мы будет раскодировать предсказанные индексы в слова нам пригодится декодер. Лучше использовать готовый из той же библиотеки. Он сам склеит символьные нграммы в слова. Без него нам нужно было бы писать правила, чтобы избавиться от вот такого"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "EVuGuhHg3-Tq",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "EVuGuhHg3-Tq",
    "outputId": "4e66aeac-a7b6-46d0-a849-23e1667e13ce"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'пример текста с ред ##ким словом'"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# без декодера\n",
    "' '.join(tokenizer_ru.encode('Пример текста с редким словом').tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f01e7bba",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'пример текста с редким словом'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# c готовым декодером\n",
    "tokenizer_ru.decoder.decode(tokenizer_ru.encode('Пример текста с редким словом').tokens)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "605ea30c",
   "metadata": {
    "id": "605ea30c"
   },
   "source": [
    "### ВАЖНО!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67f07127",
   "metadata": {
    "id": "67f07127"
   },
   "source": [
    "Токенизатор - это неотъемлимая часть модели, поэтому не забывайте сохранять токенизатор вместе с моделью. Если вы забудете про это и переобучите токенизатор, то индексы токенов разойдутся и веса модели будут бесполезны. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "496d0ea7",
   "metadata": {
    "id": "496d0ea7"
   },
   "outputs": [],
   "source": [
    "# раскоментируйте эту ячейку при обучении токенизатора\n",
    "# а потом снова закоментируйте чтобы при перезапуске не перезаписать токенизаторы\n",
    "# tokenizer_en.save('tokenizer_en')\n",
    "# tokenizer_ru.save('tokenizer_ru')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f4661964",
   "metadata": {
    "id": "f4661964"
   },
   "outputs": [],
   "source": [
    "tokenizer_en = Tokenizer.from_file(\"tokenizer_en\")\n",
    "tokenizer_ru = Tokenizer.from_file(\"tokenizer_ru\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b0685f9",
   "metadata": {
    "id": "2b0685f9"
   },
   "source": [
    "##### Переводим текст в индексы. \n",
    "\n",
    "Для исходного языка нам не нужно никаких дополнительных токенов, поэтому просто вызываем токенизатор.   \n",
    "А вот для целевого в начало нужно добавить токен '[START]', а в конец '[END]'.  \n",
    "Эти теги понадобятся, когда мы будем генерировать тексты с уже обученной моделью - тэг [START] мы будем подавать в самом начале перевода, чтобы состояние декодера не было пустым (ниже вы увидите, почему оно не может быть пустым), а [END] будем использовать, чтобы понимать когда остановить перевод (если модель предсказала [END] - значит перевод окончен)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 413,
   "id": "dc003758",
   "metadata": {
    "id": "dc003758"
   },
   "outputs": [],
   "source": [
    "def encode(text, tokenizer, target=False):\n",
    "    if target:\n",
    "        return [tokenizer.token_to_id('[START]')] + tokenizer.encode(text).ids + \\\n",
    "                [tokenizer.token_to_id('[END]')]\n",
    "    else:\n",
    "        return tokenizer.encode(text).ids \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff04bf47",
   "metadata": {
    "id": "ff04bf47"
   },
   "source": [
    "Кодируем и паддим"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 584,
   "id": "7fc2dae1",
   "metadata": {
    "id": "7fc2dae1"
   },
   "outputs": [],
   "source": [
    "X_en = [encode(t, tokenizer_en) for t in en_sents]\n",
    "X_ru = [encode(t, tokenizer_ru, True) for t in ru_sents]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa68238e",
   "metadata": {},
   "source": [
    "Как обычно мы должны выбрать максимальную допустимую длину. Максимальные значения слишком большие"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 415,
   "id": "281b5b90",
   "metadata": {
    "id": "281b5b90"
   },
   "outputs": [],
   "source": [
    "max_len_en = np.max([len(x) for x in X_en])\n",
    "max_len_ru = np.max([len(x) for x in X_ru])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 416,
   "id": "5984886a",
   "metadata": {
    "id": "5984886a"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(17863, 19389)"
      ]
     },
     "execution_count": 416,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "max_len_en, max_len_ru"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ad8e2ae",
   "metadata": {},
   "source": [
    "Если посмотреть на перцентили, то будет видно, что это скорее всего выбросы. Большая часть текстов короткие."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 417,
   "id": "49f2f2ec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Исходный язык\n",
      "50 % текстов <=  10.0\n",
      "75 % текстов <=  18.0\n",
      "90 % текстов <=  33.0\n",
      "95 % текстов <=  45.0\n",
      "99 % текстов <=  75.0\n"
     ]
    }
   ],
   "source": [
    "print('Исходный язык')\n",
    "print('50 % текстов <= ', np.percentile([len(x) for x in X_en], 50))\n",
    "print('75 % текстов <= ', np.percentile([len(x) for x in X_en], 75))\n",
    "print('90 % текстов <= ', np.percentile([len(x) for x in X_en], 90))\n",
    "print('95 % текстов <= ', np.percentile([len(x) for x in X_en], 95))\n",
    "print('99 % текстов <= ', np.percentile([len(x) for x in X_en], 99))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 418,
   "id": "6ab5bcd6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Целевой язык\n",
      "50 % текстов <=  11.0\n",
      "75 % текстов <=  20.0\n",
      "90 % текстов <=  36.0\n",
      "95 % текстов <=  48.0\n",
      "99 % текстов <=  79.0\n"
     ]
    }
   ],
   "source": [
    "print('Целевой язык')\n",
    "print('50 % текстов <= ', np.percentile([len(x) for x in X_ru], 50))\n",
    "print('75 % текстов <= ', np.percentile([len(x) for x in X_ru], 75))\n",
    "print('90 % текстов <= ', np.percentile([len(x) for x in X_ru], 90))\n",
    "print('95 % текстов <= ', np.percentile([len(x) for x in X_ru], 95))\n",
    "print('99 % текстов <= ', np.percentile([len(x) for x in X_ru], 99))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ab64c02",
   "metadata": {},
   "source": [
    "Давайте для каждого из языков возьмем значения, которые покрывают 95 % всех данных"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 419,
   "id": "5cc0a376",
   "metadata": {
    "id": "5cc0a376"
   },
   "outputs": [],
   "source": [
    "# обратите внимание, что в seq2seq длины могут быть разными\n",
    "max_len_en, max_len_ru = 45, 48"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bdc8070",
   "metadata": {},
   "source": [
    "По умолчанию паддинг делается нулями, но нам это не подходит, так как в постороенных словарях 0 - это тэг UNK. Вытащим тэг паддинга из словаря и будет использовать его"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 420,
   "id": "3f4f31fa",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "3f4f31fa",
    "outputId": "85345521-b21a-4bc1-96e2-028ecd737034"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 420,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "PAD_IDX = tokenizer_ru.token_to_id('[PAD]')\n",
    "PAD_IDX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 585,
   "id": "49ae735e",
   "metadata": {
    "id": "49ae735e"
   },
   "outputs": [],
   "source": [
    "X_en = keras.preprocessing.sequence.pad_sequences(\n",
    "              X_en, maxlen=max_len_en, padding='post', value=PAD_IDX)\n",
    "\n",
    "X_ru_out = keras.preprocessing.sequence.pad_sequences(\n",
    "              [x[1:] for x in X_ru], maxlen=max_len_ru-1, padding='post', \n",
    "              value=PAD_IDX)\n",
    "\n",
    "X_ru_dec = keras.preprocessing.sequence.pad_sequences(\n",
    "              [x[:-1] for x in X_ru], maxlen=max_len_ru-1, \n",
    "              padding='post', value=PAD_IDX)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cae14749",
   "metadata": {},
   "source": [
    "На вход модели мы будем подавать три матрицы.\n",
    "\n",
    "Первая - X_en - это индексы токенов в английских текстах (тут все как и в предыдущих семинарах)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 586,
   "id": "62d4e3da",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"yeah, that's not exactly...\""
      ]
     },
     "execution_count": 586,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "en_sents[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 587,
   "id": "597680b6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([3280,   13, 2763,    8,   58, 2808, 5148, 2856,    1,    1,    1,\n",
       "          1,    1,    1,    1,    1,    1,    1,    1,    1,    1,    1,\n",
       "          1,    1,    1,    1,    1,    1,    1,    1,    1,    1,    1,\n",
       "          1,    1,    1,    1,    1,    1,    1,    1,    1,    1,    1,\n",
       "          1], dtype=int32)"
      ]
     },
     "execution_count": 587,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_en[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d3f4c4e",
   "metadata": {},
   "source": [
    "А вот русскоязычные тексты мы будем подавать через две матрицы - одна будет использоваться как вход для декодера (X_ru_dec), а вторая как правильные ответы (X_ru_out). \n",
    "\n",
    "Если присмотреться, то можно заметить что X_ru_dec и X_ru_out отличаются на 1 шаг (X_ru_out это X_ru_dec сдвинутый вправо)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 588,
   "id": "d751b1d5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'да, но не совсем...'"
      ]
     },
     "execution_count": 588,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ru_sents[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 589,
   "id": "ed95081e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([   2, 2588,   15, 2589, 2513, 5362, 2643,    1,    1,    1,    1,\n",
       "          1,    1,    1,    1,    1,    1,    1,    1,    1,    1,    1,\n",
       "          1,    1,    1,    1,    1,    1,    1,    1,    1,    1,    1,\n",
       "          1,    1,    1,    1,    1,    1,    1,    1,    1,    1,    1,\n",
       "          1,    1,    1], dtype=int32)"
      ]
     },
     "execution_count": 589,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_ru_dec[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 590,
   "id": "14d7c918",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([2588,   15, 2589, 2513, 5362, 2643,    3,    1,    1,    1,    1,\n",
       "          1,    1,    1,    1,    1,    1,    1,    1,    1,    1,    1,\n",
       "          1,    1,    1,    1,    1,    1,    1,    1,    1,    1,    1,\n",
       "          1,    1,    1,    1,    1,    1,    1,    1,    1,    1,    1,\n",
       "          1,    1,    1], dtype=int32)"
      ]
     },
     "execution_count": 590,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_ru_out[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79452bf4",
   "metadata": {},
   "source": [
    "Такая структура внутри модели будет использована, чтобы предсказывать следующее слово по текущему контексту. Грубо говоря, модель будет интерпретировать эти две последовательности индексов, как пары контекст-продолжение. \n",
    "\n",
    "Вот это:  \n",
    "2, 2588, 15, 2589, 2513, 5362, 2643   \n",
    "2588, 15, 2589, 2513, 5362, 2643,    3\n",
    "\n",
    "Превратится для модели в:  \n",
    "2 - 2588  \n",
    "2, 2588 - 15  \n",
    "2, 2588, 15, 2589 - 2513  \n",
    "2, 2588, 15, 2589, 2513 - 5362  \n",
    "2, 2588, 15, 2589, 2513, 5362 - 2643  \n",
    "2, 2588, 15, 2589, 2513, 5362, 2643 - 3  \n",
    "\n",
    "(как это будет сделано показано дальше в разделе про маски)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 496,
   "id": "824dcb29",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "824dcb29",
    "outputId": "34899386-a168-4cd3-b063-71e305cc1c22"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((1000000, 45), (1000000, 47))"
      ]
     },
     "execution_count": 496,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# миллион примеров \n",
    "X_en.shape, X_ru_out.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acfa4331",
   "metadata": {
    "id": "acfa4331"
   },
   "source": [
    "Разделяем на трейн и тест"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 591,
   "id": "25fa5f05",
   "metadata": {
    "id": "25fa5f05"
   },
   "outputs": [],
   "source": [
    "(X_en_train, X_en_valid, \n",
    "X_ru_dec_train, X_ru_dec_valid, \n",
    "X_ru_out_train, X_ru_out_valid) = train_test_split(X_en, \n",
    "                                                  X_ru_dec, \n",
    "                                                  X_ru_out, \n",
    "                                                  test_size=0.05)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9faa719",
   "metadata": {
    "id": "d9faa719"
   },
   "source": [
    "Дальше код модели, он взят вот отсюда (с небольшими изменениями) - https://www.tensorflow.org/text/tutorials/transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f045ac37",
   "metadata": {
    "id": "f045ac37"
   },
   "source": [
    "## Имплементация трансформера"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4270b2f",
   "metadata": {},
   "source": [
    "Главное в трансформере это механизм внимания. Функция scaled_dot_product_attention преобразует query, key, value вектора в обновленные вектора с помощью механизма внимания.\n",
    "\n",
    "Визуализация этих шагов (источник: https://jalammar.github.io/illustrated-transformer/):\n",
    "![](https://jalammar.github.io/images/t/self-attention-output.png \"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56e70b1a",
   "metadata": {},
   "source": [
    "Пройдитесь по каждой строчке в этой функции и сопоставьте с визуализацией."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 429,
   "id": "656be820",
   "metadata": {
    "id": "656be820"
   },
   "outputs": [],
   "source": [
    "def scaled_dot_product_attention(query, key, value, mask):\n",
    "    # Считаем скалярное произведение между запросом (query) и ключом (key), транспонируя ключ\n",
    "    matmul_qk = keras.ops.matmul(query, keras.ops.transpose(key, axes=[0, 1, 3, 2]))\n",
    "\n",
    "    # Получаем глубину (размерность) ключа и преобразуем ее во float\n",
    "    depth = keras.ops.cast(key.shape[-1], torch.float32)\n",
    "\n",
    "    # Делим результат скалярного произведения на квадратный корень из глубины\n",
    "    # Это делается для уменьшения влияния больших значений и стабилизации градиентов во время обучения\n",
    "    logits = matmul_qk / keras.ops.sqrt(depth)\n",
    "\n",
    "    # Если есть маска, применяем ее к логитам, чтобы обнулить нежелательные значения\n",
    "    if mask is not None:\n",
    "        logits += (mask * -1e9)\n",
    "\n",
    "    # Применяем функцию softmax для получения весов внимания\n",
    "    attention_weights = keras.ops.softmax(logits, axis=-1)\n",
    "\n",
    "    # Умножаем веса внимания на значения (value) для получения итогового результата\n",
    "    output = keras.ops.matmul(attention_weights, value)\n",
    "\n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3efb7349",
   "metadata": {},
   "source": [
    "В этой функции query, key и value вектора подаются на вход, но как они получаются из эмбедингов?\n",
    "\n",
    "Это реализовано в классе MultiHeadAttention, который уже является самостоятельным слоем, на котором потом будет построен блок трансформера как в энкодере так и в декодере.\n",
    "\n",
    "В этом слое можно увидеть три отдельных полносвязных слоя одинаковой размерности (d_model), которые применяются к inputs, чтобы получить q, k, v вектора.  \n",
    "Можно заметить, что в функции call, разделение на q,k,v происходит еще раньше - они просто достаются из inputs по ключу, как из словаря. Но если немного промотать вперед до encoder_layer , то можно убедиться, что при вызове MultiheadAttention в эти переменные подставляются одно и то же (нажмите CTRL+F и введите #call_mha, чтобы перейти к этому моменту и потом вернутся сюда). \n",
    "\n",
    "А если сделать еще один шаг до encoder, то можно увидеть, что такое inputs - это эмбединги слов (CTRL+F и введите #inputs). \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 430,
   "id": "f4b51870",
   "metadata": {
    "id": "f4b51870"
   },
   "outputs": [],
   "source": [
    "class MultiHeadAttention(keras.layers.Layer):\n",
    "\n",
    "    def __init__(self, d_model, num_heads, name=\"multi_head_attention\"):\n",
    "        super(MultiHeadAttention, self).__init__(name=name)\n",
    "        self.num_heads = num_heads  # количество голов для внимания\n",
    "        self.d_model = d_model  # размерность вектора модели\n",
    "\n",
    "        # Убеждаемся, что размерность модели делится нацело на количество голов\n",
    "        assert d_model % self.num_heads == 0\n",
    "\n",
    "        self.depth = d_model // self.num_heads  # размерность каждой головы\n",
    "\n",
    "        # Создаем полносвязные слои для запроса, ключа и значения\n",
    "        self.query_dense = keras.layers.Dense(units=d_model)\n",
    "        self.key_dense = keras.layers.Dense(units=d_model)\n",
    "        self.value_dense = keras.layers.Dense(units=d_model)\n",
    "\n",
    "        # Создаем последний полносвязный слой\n",
    "        self.dense = keras.layers.Dense(units=d_model)\n",
    "\n",
    "    def split_heads(self, inputs, batch_size):\n",
    "        # Разделяем входные данные на головы\n",
    "        inputs = keras.ops.reshape(\n",
    "            inputs, newshape=(batch_size, -1, self.num_heads, self.depth))\n",
    "        return keras.ops.transpose(inputs, axes=[0, 2, 1, 3])\n",
    "\n",
    "    def call(self, inputs):\n",
    "        query, key, value, mask = inputs['query'], inputs['key'], inputs[\n",
    "            'value'], inputs.get('mask', None)\n",
    "        batch_size = query.shape[0]\n",
    "\n",
    "        # Пропускаем запрос, ключ и значение через соответствующие полносвязные слои\n",
    "        query = self.query_dense(query)\n",
    "        key = self.key_dense(key)\n",
    "        value = self.value_dense(value)\n",
    "\n",
    "        # Разделяем запрос, ключ и значение на головы \n",
    "        # то есть просто разрезаем вектора на num_heads частей \n",
    "        # и сравниваем все части между собой\n",
    "        query = self.split_heads(query, batch_size)\n",
    "        key = self.split_heads(key, batch_size)\n",
    "        value = self.split_heads(value, batch_size)\n",
    "\n",
    "        # Выполняем механизм внимания с масштабированным скалярным произведением\n",
    "        scaled_attention = scaled_dot_product_attention(query, key, value, mask)\n",
    "\n",
    "        scaled_attention = keras.ops.transpose(scaled_attention, axes=[0, 2, 1, 3])\n",
    "\n",
    "        # Объединяем головы вместе (склеиваем векторы в один)\n",
    "        concat_attention = keras.ops.reshape(scaled_attention,\n",
    "                                      (batch_size, -1, self.d_model))\n",
    "\n",
    "        # Пропускаем объединенное внимание через дополнительный полносвязный слой\n",
    "        # Он просто добавляет сложности нашей модели\n",
    "        outputs = self.dense(concat_attention)\n",
    "\n",
    "        return outputs\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2f10e40",
   "metadata": {},
   "source": [
    "#### Теперь наверное самый сложный момент - почему это MultiHead, а не просто Attention layer? \n",
    "\n",
    "В MultiheadAttention используются не по одному q,k,v вектору, а сразу по N векторов, но меньшего размера. В классе выше N задается через парметр num_heads и размерность отдельного q,k или v вектора будет равна d_model/num_heads. \n",
    "\n",
    "Внимание рассчитывается между соответствующими парами q и v внутри каждой головы, а итоговый вектор получается объединением результата внимания от всех голов (его размерность будет равна d_model).  \n",
    "\n",
    "В статье illustrated transformer этот момент визуализируется вот так:\n",
    "![](https://jalammar.github.io/images/t/transformer_multi-headed_self-attention-recap.png)\n",
    "\n",
    "Но реализация в TF немного отличается. На этой картинке на шаге 3 показано, что полносвязные слои, через которые получаются q,k,v векторы у каждой головы свои - и их количество не 3, а 3 умножить на количество голов. А вот в TF полносвязных слоя всего 3, но дальше вектора просто разрезаются на куски и начаная с шага 4, все уже совпадает. \n",
    "\n",
    "Если задуматься, то разницы между этими подходами нет. Полносвязный слой это матрица, на которую умножается входной вектор, в результате чего получается другой вектор. И каждое значение в новом векторе - это результат попарного умножения и суммы значений в исходном векторе и отдельной колонки в матрице. Таким образом, уможить на матрицу и разрезать результат на три куска - это то же самое, что разрезать матрицу на три куска и умножить вектор на каждый из кусков отдельно:\n",
    "\n",
    "![](https://i.ibb.co/9Yg580J/image.png)\n",
    "\n",
    "\n",
    "Давайте еще отдельно посмотрим как происходит разбиение на головы и как потом выглядит скалярное произведение для отдельного куска."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5760c73c",
   "metadata": {},
   "source": [
    "Инициализируем класс"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 431,
   "id": "19dd941b",
   "metadata": {},
   "outputs": [],
   "source": [
    "mha = MultiHeadAttention(d_model=10, num_heads=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecd8fa2d",
   "metadata": {},
   "source": [
    "Сгенерируем батч из одного текста с двумя словами и размером эбмединга 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 432,
   "id": "7c25cb11",
   "metadata": {},
   "outputs": [],
   "source": [
    "inp = np.random.normal(size=(1, 2, 3)) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 433,
   "id": "2fa2e4fc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[-1.15717704,  1.47875626, -0.21808616],\n",
       "        [ 0.79722457, -0.69800946,  1.72175378]]])"
      ]
     },
     "execution_count": 433,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f14f09c0",
   "metadata": {},
   "source": [
    "Сформируем входной словарь для класса"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 434,
   "id": "7e866e38",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = {'query': inp,'key': inp, 'value': inp,}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "819a35c3",
   "metadata": {},
   "source": [
    "Применим слой к входному словарю"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "id": "d0b7823f",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = mha(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61747204",
   "metadata": {},
   "source": [
    "На выходе получилось, что для каждого слова у нас теперь эмдединг размером 10 (мы указали такой d_model выше)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "id": "8e5bb366",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 2, 10])"
      ]
     },
     "execution_count": 122,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62b016f9",
   "metadata": {},
   "source": [
    "Давайте посмотрим как это вектор вычислялся внутри класса по шагам"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 123,
   "id": "0cdee43b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# преобразуем эмбединг в q,k,v вектора\n",
    "q = mha.query_dense(inputs['query'])\n",
    "k = mha.query_dense(inputs['key'])\n",
    "v = mha.query_dense(inputs['value'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "id": "33465f95",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 2, 10])"
      ]
     },
     "execution_count": 124,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# размерность не изменилась так как мы еще не порезали на головы\n",
    "q.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "id": "1a818c70",
   "metadata": {},
   "outputs": [],
   "source": [
    "# разрезаем каждый вектор на куски (головы)\n",
    "qh = mha.split_heads(q, q.shape[0])\n",
    "kh = mha.split_heads(k, k.shape[0])\n",
    "vh = mha.split_heads(v, v.shape[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f0a7f8b",
   "metadata": {},
   "source": [
    "В результате получается тензор с размерностями (количество текстов x количество слов x количество голов x размерность вектора головы)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "id": "1c182c26",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 2, 2, 5])"
      ]
     },
     "execution_count": 127,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# мы указали num_heads 2 - поэтому размерность каждого куска 5 (10 разделить на два)\n",
    "qh.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "id": "6cbb9578",
   "metadata": {},
   "outputs": [],
   "source": [
    "# keras.ops.transpose(kh).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e07c7527",
   "metadata": {},
   "source": [
    "Далее вычисляется скалярное произведение между q и k кусками. На вход подаются тензоры, но тензорфлоу умеет работать с таким форматом тоже"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 133,
   "id": "a7f9d59a",
   "metadata": {},
   "outputs": [],
   "source": [
    "matmul_qk = keras.ops.matmul(qh, keras.ops.transpose(kh, axes=[0, 1, 3, 2]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 134,
   "id": "5a2ef8d5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 2, 2, 2])"
      ]
     },
     "execution_count": 134,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# На выходе получится тензор вот такой размерности\n",
    "# 5 изменилось на 2\n",
    "# так как для каждого куска у нас по две близости (до каждого слова включая себя)\n",
    "matmul_qk.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e837c492",
   "metadata": {},
   "source": [
    "Применим софтмакс, чтобы нагляднее увидеть близости"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 137,
   "id": "85a6d549",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[[0.78, 0.22],\n",
       "         [0.36, 0.64]],\n",
       "\n",
       "        [[0.53, 0.47],\n",
       "         [0.21, 0.79]]]], dtype=float32)"
      ]
     },
     "execution_count": 137,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "keras.ops.softmax(matmul_qk, axis=-1).detach().cpu().numpy().round(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c01fb05",
   "metadata": {},
   "source": [
    "Далее эти близости используются как веса для суммы value кусков и потом результат по каждому кусочку собирается обратно в один вектор размерности (1, 2, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "858e638a",
   "metadata": {},
   "source": [
    "### Маски\n",
    "\n",
    "Еще один важный момент это то как используются маски, чтобы занулить некоторые близости.\n",
    "Есть два вида масок - маска для паддинга и маска для слов в будущем. \n",
    "Маска для паддинга используется и в энкодере и в декодере, а маска для слов в будущем только для декодера."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "id": "5c48cea2",
   "metadata": {
    "id": "5c48cea2"
   },
   "outputs": [],
   "source": [
    "def create_padding_mask(x):\n",
    "    mask = keras.ops.cast(keras.ops.equal(x, PAD_IDX), torch.float32)\n",
    "    # (batch_size, 1, 1, sequence length)\n",
    "    return mask[:, None, None, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "id": "64790293",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[1., 0., 0.]]]], device='cuda:0')"
      ]
     },
     "execution_count": 139,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "create_padding_mask([[1,2,3]]) # 1 подставилась для индекса который равен PAD_IDX"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b0ab9ad",
   "metadata": {},
   "source": [
    "Маска для будущего создает треугольную матрицу, которая накладывается на близости в self-attention части декодера. Как раз с помощью этой маски и реализуется разбиение входных последовательностей на контекст-продолжение - сначала вычисляется внимание до всех слов, даже в будущем, но потом полученные близости слов в будущем просто зануляются и не используются."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 153,
   "id": "21c0c899",
   "metadata": {
    "id": "21c0c899"
   },
   "outputs": [],
   "source": [
    "def create_look_ahead_mask(x):\n",
    "    # эта функция немножко сложная, но суть у нее достаточно простая\n",
    "    # нужно создать треугольную маску, с помощью которой мы закроем \n",
    "    # для каждого токена все последующие токены\n",
    "    seq_len = x.shape[1]\n",
    "    ones_mask = keras.ops.ones((1, seq_len, seq_len), dtype=\"int32\")\n",
    "    row_index = keras.ops.cumsum(ones_mask, axis=-2)\n",
    "    col_index = keras.ops.cumsum(ones_mask, axis=-1)\n",
    "    look_ahead_mask = ~ keras.ops.greater_equal(row_index, col_index)\n",
    "    padding_mask = create_padding_mask(x)\n",
    "    return keras.ops.maximum(look_ahead_mask, padding_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 160,
   "id": "b6f988e6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[False,  True,  True],\n",
       "         [False, False,  True],\n",
       "         [False, False, False]]], device='cuda:0')"
      ]
     },
     "execution_count": 160,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# для декодера нам нужно замаскировать следующие токены\n",
    "# так как мы пытаемся их сгенерировать\n",
    "# для этого создается вот такая маска\n",
    "ones_mask = keras.ops.ones((1, 3, 3), dtype=\"int32\")\n",
    "row_index = keras.ops.cumsum(ones_mask, axis=-2)\n",
    "col_index = keras.ops.cumsum(ones_mask, axis=-1)\n",
    "mask = ~ keras.ops.greater_equal(row_index, col_index)\n",
    "mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 161,
   "id": "499d2f3d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[   0, -100, -100],\n",
       "         [   0,    0, -100],\n",
       "         [   0,    0,    0]]], device='cuda:0')"
      ]
     },
     "execution_count": 161,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# в слое с вниманием эта маска будет использоваться чтобы занулить близости с токенами из будущего\n",
    "(mask * -100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 162,
   "id": "9c0f62e3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[1., 1., 1.],\n",
       "          [1., 0., 1.],\n",
       "          [1., 0., 0.]]]], device='cuda:0')"
      ]
     },
     "execution_count": 162,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# эта функция также проверяет и паддинг и если он есть то и его тоже замаскирует\n",
    "# 1 - это паддинг айди\n",
    "create_look_ahead_mask(torch.LongTensor([[1,3,2]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "id": "7a1f209e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[0., 1., 1.],\n",
       "          [0., 0., 1.],\n",
       "          [0., 0., 0.]]]], device='cuda:0')"
      ]
     },
     "execution_count": 163,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "create_look_ahead_mask(torch.Tensor([[2,4,3]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 164,
   "id": "41493f7e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[0., 1., 1.],\n",
       "          [0., 0., 1.],\n",
       "          [0., 0., 0.]]]], device='cuda:0')"
      ]
     },
     "execution_count": 164,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "create_look_ahead_mask(torch.Tensor([[2,4,3]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c019796a",
   "metadata": {},
   "source": [
    "### Позиционное кодирование\n",
    "\n",
    "К каждому эбмедингу слова добавляется вектор, который зависит от позиции слова. Таким образом в трансформер передается информация о порядке слов. Обратите внимание, что выше вектора сравниваются через скалярное произведение просто попарно и все вектора, которые не замаскированы, можно переставить местами и результат не изменится. Добавление информации о позиции в сами вектора изменяет результат.\n",
    "\n",
    "Класс PositionalEncoding генерирует вектор для каждой позиции заданной размерности. Размерность задается также через d_model, чтобы можно было просто прибавить к вектору слова"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 258,
   "id": "e120bbe5",
   "metadata": {
    "id": "e120bbe5"
   },
   "outputs": [],
   "source": [
    "class PositionalEncoding(keras.layers.Layer):\n",
    "\n",
    "    def __init__(self, position, d_model):\n",
    "        super(PositionalEncoding, self).__init__()\n",
    "        self.pos_encoding = self.positional_encoding(position, d_model)\n",
    "\n",
    "    def get_angles(self, position, i, d_model):\n",
    "        angles = 1 / keras.ops.power(10000, (2 * (i // 2)) / d_model)\n",
    "        return keras.ops.multiply(position, angles)\n",
    "\n",
    "    def positional_encoding(self, position, d_model):\n",
    "        angle_rads = self.get_angles(\n",
    "            position=torch.arange(0, position, dtype=torch.float32)[:, None],\n",
    "            i=torch.arange(0, d_model, dtype=torch.float32)[None, :],\n",
    "            d_model=d_model)\n",
    "        sines = keras.ops.sin(angle_rads[:, 0::2])\n",
    "        cosines = keras.ops.cos(angle_rads[:, 1::2])\n",
    "\n",
    "        pos_encoding = keras.ops.concatenate([sines, cosines], axis=-1)\n",
    "        pos_encoding = pos_encoding[None, ...]\n",
    "        return keras.ops.cast(pos_encoding, torch.float32)\n",
    "\n",
    "    def call(self, inputs):\n",
    "        return inputs + self.pos_encoding[:, :inputs.shape[1], :]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c7fb39e",
   "metadata": {},
   "source": [
    "### Encoder"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7763b910",
   "metadata": {},
   "source": [
    "Дальше кастомные функции и классы, которые были определены выше и стандартные инструменты tf.keras объединяются, чтобы собрать эндкодер и декодер. Энкодер устроен вот так (источник - https://jalammar.github.io/illustrated-transformer/)\n",
    "![](https://jalammar.github.io/images/t/transformer_resideual_layer_norm.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3039bada",
   "metadata": {},
   "source": [
    "Выше мы разобрались с self-attention и positional encoding. Теперь к этому нужно добавить еще один полносвязный слой (а точнее даже два), а также нормализацию и skip связи (пунктирные стрелочки). Ну и Dropout слои, но это просто случайное зануление некоторых векторов, чтобы модель лучше обучалась, а не что-то существенное, поэтому их даже нет на картинке.\n",
    "\n",
    "Нормализация делается с помощью LayerNormalization слоя. Она нужна, чтобы значения получаемых векторов оставались в каких-то разумных пределах. А skip-связи реализованы просто через суммирование - обратите внимание, как переменная inputs сначала передается в MultiHeadAttention, а потом полученный результат прибавляется к исходной inputs; и далее attention (который является суммой inputs и MHA(inputs)) передается в два полносвязных слоя и полученный результат прибавляется к изначальной attention переменной. В итоге получается, что информация как бы постепенно добавляется к изначальным векторам, а не перезаписывается."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 498,
   "id": "68472627",
   "metadata": {
    "id": "68472627"
   },
   "outputs": [],
   "source": [
    "def encoder_layer(units, d_model, num_heads, dropout, name=\"encoder_layer\"):\n",
    "    inputs = keras.Input(shape=(None, d_model), name=\"inputs\")\n",
    "    padding_mask = keras.Input(shape=(1, 1, None), name=\"padding_mask\")\n",
    "\n",
    "    #call_mha\n",
    "    attention = MultiHeadAttention(\n",
    "      d_model, num_heads, name=\"attention\")({\n",
    "          'query': inputs,\n",
    "          'key': inputs,\n",
    "          'value': inputs,\n",
    "          'mask': padding_mask\n",
    "      })\n",
    "    attention = keras.layers.Dropout(rate=dropout)(attention)\n",
    "    attention = keras.layers.LayerNormalization(\n",
    "      epsilon=1e-6)(inputs + attention)\n",
    "\n",
    "    outputs = keras.layers.Dense(units=units, activation='relu')(attention)\n",
    "    outputs = keras.layers.Dense(units=d_model)(outputs)\n",
    "    outputs = keras.layers.Dropout(rate=dropout)(outputs)\n",
    "    outputs = keras.layers.LayerNormalization(\n",
    "      epsilon=1e-6)(attention + outputs)\n",
    "\n",
    "    return keras.Model(\n",
    "      inputs=[inputs, padding_mask], outputs=outputs, name=name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22c5f72e",
   "metadata": {},
   "source": [
    "Блоков энкодера может быть сколько угодно, так как один такой блок принимает на вход векторы размерности d_model и выдает как результат вектора размерности d_model, т.е. можно скармливать результаты одного блока другому и так далее. За количество последовательных транформерных блоков в энкодере отвечает параметр num_layers. Обратите внимание, что ниже все это реализовано просто через цикл с перезаписыванием переменной."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 499,
   "id": "89dcc42e",
   "metadata": {
    "id": "89dcc42e"
   },
   "outputs": [],
   "source": [
    "def encoder(vocab_size,\n",
    "            num_layers,\n",
    "            units,\n",
    "            d_model,\n",
    "            num_heads,\n",
    "            dropout,\n",
    "            max_len,\n",
    "            name=\"encoder\"):\n",
    "    inputs = keras.Input(shape=(max_len,), name=\"inputs\")\n",
    "    padding_mask = keras.Input(shape=(1, 1, None), name=\"padding_mask\")\n",
    "\n",
    "    embeddings = keras.layers.Embedding(vocab_size, d_model)(inputs)\n",
    "    embeddings *= keras.ops.sqrt(keras.ops.cast(d_model, torch.float32))\n",
    "    embeddings = PositionalEncoding(max_len, d_model)(embeddings)\n",
    "\n",
    "    #inputs (они тут называются outputs но это просто такой нейминг, \n",
    "    # этот параметр передается в encoder_layer первым и encoder_layer будет считать его inputs)\n",
    "    # outputs он тут называется для удобства, так как он будет перезаписываться\n",
    "    # чтобы на вход следующему блоку подавать уже не эмбединги,\n",
    "    # а то что получится как результат предыдущего блока\n",
    "    outputs = keras.layers.Dropout(rate=dropout)(embeddings)\n",
    "\n",
    "    for i in range(num_layers):\n",
    "        outputs = encoder_layer(\n",
    "            units=units,\n",
    "            d_model=d_model,\n",
    "            num_heads=num_heads,\n",
    "            dropout=dropout,\n",
    "            name=\"encoder_layer_{}\".format(i),\n",
    "        )([outputs, padding_mask])\n",
    "\n",
    "    return keras.Model(inputs=[inputs, padding_mask], outputs=outputs, name=name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d048be7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "c2380074",
   "metadata": {},
   "source": [
    "## Decoder\n",
    "\n",
    "Декодер построен из точно таких же частей с одним важным отличием - в одном блоке декодера два MultiHeadAttention слоя. Первый из них такой же как в энкодере и реализует self-attention, а во втором реализован cross-attention. \n",
    "\n",
    "![](https://jalammar.github.io/images/t/transformer_resideual_layer_norm_3.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "809b26df",
   "metadata": {},
   "source": [
    "В коде разница выражает в параметрах, которые передаются в MultiHeadAttention. Обратите внимание, что ниже первый MultiHeadAttention вызывается также как и в энкодере (inputs={'query': inputs, 'key': inputs, 'value': inputs, 'mask': look_ahead_mask}), а вот на вход второму передаются inputs={ 'query': attention1, 'key': enc_outputs, 'value': enc_outputs, 'mask': padding_mask}, то есть query часть тут это inputs декодера уже обогащенный self-attention, а key и value это последние значения этих векторов в энкодере. Еще важно заметить, что в первом MultiHeadAttention используется look_ahead_mask, а во втором только padding_mask.\n",
    "\n",
    "Таким образом, в декодере обращается внимание на уже сгенерированный перевод и на весь исходный текст. \n",
    "\n",
    "Обратите внимание, что вход декодера не может быть пустым, потому что тогда inputs будет пустым и не будет query, key, value векторов между которыми можно расчитать внимание и соответственно не из чего делать предсказания."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 500,
   "id": "a2f90e7c",
   "metadata": {
    "id": "a2f90e7c"
   },
   "outputs": [],
   "source": [
    "def decoder_layer(units, d_model, num_heads, dropout, name=\"decoder_layer\"):\n",
    "    inputs = keras.Input(shape=(None, d_model), name=\"inputs\")\n",
    "    enc_outputs = keras.Input(shape=(None, d_model), name=\"encoder_outputs\")\n",
    "    look_ahead_mask = keras.Input(\n",
    "      shape=(1, None, None), name=\"look_ahead_mask\")\n",
    "    padding_mask = keras.Input(shape=(1, 1, None), name='padding_mask')\n",
    "\n",
    "    attention1 = MultiHeadAttention(\n",
    "      d_model, num_heads, name=\"attention_1\")(inputs={\n",
    "          'query': inputs,\n",
    "          'key': inputs,\n",
    "          'value': inputs,\n",
    "          'mask': look_ahead_mask\n",
    "      })\n",
    "    attention1 = keras.layers.LayerNormalization(\n",
    "      epsilon=1e-6)(attention1 + inputs)\n",
    "\n",
    "    attention2 = MultiHeadAttention(\n",
    "      d_model, num_heads, name=\"attention_2\")(inputs={\n",
    "          'query': attention1,\n",
    "          'key': enc_outputs,\n",
    "          'value': enc_outputs,\n",
    "          'mask': padding_mask\n",
    "      })\n",
    "    attention2 = keras.layers.Dropout(rate=dropout)(attention2)\n",
    "    attention2 = keras.layers.LayerNormalization(\n",
    "      epsilon=1e-6)(attention2 + attention1)\n",
    "\n",
    "    outputs = keras.layers.Dense(units=units, activation='relu')(attention2)\n",
    "    outputs = keras.layers.Dense(units=d_model)(outputs)\n",
    "    outputs = keras.layers.Dropout(rate=dropout)(outputs)\n",
    "    outputs = keras.layers.LayerNormalization(\n",
    "      epsilon=1e-6)(outputs + attention2)\n",
    "\n",
    "    return keras.Model(\n",
    "      inputs=[inputs, enc_outputs, look_ahead_mask, padding_mask],\n",
    "      outputs=outputs,\n",
    "      name=name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 501,
   "id": "1c0dd6a2",
   "metadata": {
    "id": "8e9f89b8"
   },
   "outputs": [],
   "source": [
    "def decoder(vocab_size,\n",
    "            num_layers,\n",
    "            units,\n",
    "            d_model,\n",
    "            num_heads,\n",
    "            dropout,\n",
    "            max_len,\n",
    "            name='decoder'):\n",
    "    inputs = keras.Input(shape=(max_len,), name='inputs')\n",
    "    enc_outputs = keras.Input(shape=(None, d_model), name='encoder_outputs')\n",
    "    look_ahead_mask = keras.Input(\n",
    "      shape=(1, None, None), name='look_ahead_mask')\n",
    "    padding_mask = keras.Input(shape=(1, 1, None), name='padding_mask')\n",
    "\n",
    "    embeddings = keras.layers.Embedding(vocab_size, d_model)(inputs)\n",
    "    embeddings *= keras.ops.sqrt(keras.ops.cast(d_model, torch.float32))\n",
    "    embeddings = PositionalEncoding(max_len, d_model)(embeddings)\n",
    "\n",
    "    outputs = keras.layers.Dropout(rate=dropout)(embeddings)\n",
    "\n",
    "    for i in range(num_layers):\n",
    "        outputs = decoder_layer(\n",
    "            units=units,\n",
    "            d_model=d_model,\n",
    "            num_heads=num_heads,\n",
    "            dropout=dropout,\n",
    "            name='decoder_layer_{}'.format(i),\n",
    "        )(inputs=[outputs, enc_outputs, look_ahead_mask, padding_mask])\n",
    "\n",
    "    return keras.Model(\n",
    "      inputs=[inputs, enc_outputs, look_ahead_mask, padding_mask],\n",
    "      outputs=outputs,\n",
    "      name=name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbd2160e",
   "metadata": {},
   "source": [
    "### Модель целиком\n",
    "\n",
    "Общая модель состоит из энкодера, декодера и еще 1 полносвязного слоя. Дополнительный последний слой нужен, чтобы из векторов получить вероятности классов (=слов)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 502,
   "id": "e356741b",
   "metadata": {
    "id": "e356741b"
   },
   "outputs": [],
   "source": [
    "def transformer(vocab_size,\n",
    "                num_layers,\n",
    "                units,\n",
    "                d_model,\n",
    "                num_heads,\n",
    "                dropout,\n",
    "                max_len,\n",
    "                name=\"transformer\"):\n",
    "    inputs = keras.Input(shape=(max_len[0],), name=\"inputs\")\n",
    "    dec_inputs = keras.Input(shape=(max_len[1]-1,), name=\"dec_inputs\")\n",
    "\n",
    "    enc_padding_mask = keras.layers.Lambda(\n",
    "      create_padding_mask, output_shape=(1, 1, None),\n",
    "      name='enc_padding_mask')(inputs)\n",
    "    \n",
    "    look_ahead_mask = keras.layers.Lambda(\n",
    "      create_look_ahead_mask,\n",
    "      output_shape=(1, None, None),\n",
    "      name='look_ahead_mask')(dec_inputs)\n",
    "    \n",
    "    dec_padding_mask = keras.layers.Lambda(\n",
    "      create_padding_mask, output_shape=(1, 1, None),\n",
    "      name='dec_padding_mask')(inputs)\n",
    "\n",
    "    enc_outputs = encoder(\n",
    "                          vocab_size=vocab_size[0],\n",
    "                          num_layers=num_layers,\n",
    "                          units=units,\n",
    "                          d_model=d_model,\n",
    "                          num_heads=num_heads,\n",
    "                          dropout=dropout,\n",
    "                          max_len=max_len[0],\n",
    "                        )(inputs=[inputs, enc_padding_mask])\n",
    "\n",
    "    dec_outputs = decoder(\n",
    "                          vocab_size=vocab_size[1],\n",
    "                          num_layers=num_layers,\n",
    "                          units=units,\n",
    "                          d_model=d_model,\n",
    "                          num_heads=num_heads,\n",
    "                          dropout=dropout,\n",
    "                          max_len=max_len[1]-1,\n",
    "                        )(inputs=[dec_inputs, enc_outputs, \n",
    "                                  look_ahead_mask, dec_padding_mask])\n",
    "\n",
    "    outputs = keras.layers.Dense(units=vocab_size[1], name=\"outputs\")(dec_outputs)\n",
    "\n",
    "    return keras.Model(inputs=[inputs, dec_inputs], outputs=outputs, name=name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6adc39ee",
   "metadata": {},
   "source": [
    "Функция потерь тут реализована отдельно, чтобы не учитывать паддинги в расчетах. Поэтому выше последний полносвязный слой был без софтмакса - если применить софтмакс, а потом какие-то вероятности занулить, то распределение сломается (не будет суммироваться в 1). Поэтому указан параметр from_logits=True, который как раз означает, что на вход подаются не вероятности, а логиты, то есть просто результат полносвязного слоя без активации.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 592,
   "id": "d23ece2d",
   "metadata": {
    "id": "6c35ce0f"
   },
   "outputs": [],
   "source": [
    "\n",
    "L  = keras.losses.SparseCategoricalCrossentropy(\n",
    "      from_logits=True, reduction='none',)\n",
    "\n",
    "def loss_function(y_true, y_pred):\n",
    "    loss = L(y_true, y_pred)\n",
    "\n",
    "    mask = keras.ops.cast(keras.ops.not_equal(y_true, PAD_IDX), torch.float32)\n",
    "    loss = keras.ops.multiply(loss, mask)\n",
    "\n",
    "    return keras.ops.mean(loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4743513",
   "metadata": {
    "id": "a4743513"
   },
   "source": [
    "Определяем параметры модели. Есть стандартные наборы параметров для обучения маленькой, средней или большой модели. Их можно посмотреть в оригинальной статье про трансформер - https://arxiv.org/pdf/1706.03762.pdf\n",
    "\n",
    "Мы возьмем маленький, т.к. он требует меньше ресурсов. НО для практической задачи есть смысл использовать модель побольше, есть статьи в которые показано, что большие трансформеры требуют даже меньше данных для обучения -  https://arxiv.org/pdf/2002.11794.pdf."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 593,
   "id": "542fcc43",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "542fcc43",
    "outputId": "c9359c72-3dfb-4985-b26b-a832fc6fb080"
   },
   "outputs": [],
   "source": [
    "keras.backend.clear_session()\n",
    "\n",
    "# small model\n",
    "NUM_LAYERS = 2\n",
    "D_MODEL = 256\n",
    "NUM_HEADS = 8\n",
    "\n",
    "UNITS = 512\n",
    "DROPOUT = 0.1\n",
    "\n",
    "\n",
    "# average model\n",
    "# NUM_LAYERS = 6\n",
    "# D_MODEL = 512\n",
    "# NUM_HEADS = 8\n",
    "# UNITS = 2048\n",
    "# DROPOUT = 0.1\n",
    "\n",
    "\n",
    "model = transformer(\n",
    "    vocab_size=(tokenizer_en.get_vocab_size(),tokenizer_ru.get_vocab_size()),\n",
    "    num_layers=NUM_LAYERS,\n",
    "    units=UNITS,\n",
    "    d_model=D_MODEL,\n",
    "    num_heads=NUM_HEADS,\n",
    "    dropout=DROPOUT,\n",
    "    max_len=[max_len_en, max_len_ru])\n",
    "\n",
    "\n",
    "optimizer = keras.optimizers.Adam(\n",
    "    0.001, beta_1=0.9, beta_2=0.98, epsilon=1e-9)\n",
    "\n",
    "def accuracy(y_true, y_pred):\n",
    "    return keras.metrics.sparse_categorical_accuracy(y_true, y_pred)\n",
    "\n",
    "\n",
    "model.compile(optimizer=optimizer, loss=loss_function, metrics=[accuracy])\n",
    "checkpoint = keras.callbacks.ModelCheckpoint('model_ruen.weights.h5',\n",
    "                                            monitor='val_loss',\n",
    "                                            verbose=1,\n",
    "                                            save_weights_only=True,\n",
    "                                            save_best_only=True,\n",
    "                                            mode='min',\n",
    "                                            save_freq='epoch')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6590325",
   "metadata": {
    "id": "c6590325"
   },
   "source": [
    "Трансформеры - большие модели, обучаться они должны долго. Размер батча здесь уже имеет большое значение: сликом маленький батч приведет к сильному увеличению времени на обучение, но сильно большим его поставить не получится, т.к. модель большая и быстро займет всю видеокарту."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 594,
   "id": "439f7e12",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 594,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import gc\n",
    "torch.cuda.empty_cache()\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 595,
   "id": "1b2673cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.load_weights('model_ruen.weights.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 596,
   "id": "308a8e81",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 496
    },
    "id": "308a8e81",
    "outputId": "981be9b3-524f-4e0a-fc93-c44bddd4bf3c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/3\n",
      "\u001b[1m4750/4750\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m0s\u001b[0m 144ms/step - accuracy: 0.0894 - loss: 1.6173\n",
      "Epoch 1: val_loss improved from inf to 0.99290, saving model to model_ruen.weights.h5\n",
      "\u001b[1m4750/4750\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m708s\u001b[0m 149ms/step - accuracy: 0.0894 - loss: 1.6173 - val_accuracy: 0.1475 - val_loss: 0.9929\n",
      "Epoch 2/3\n",
      "\u001b[1m4750/4750\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m0s\u001b[0m 144ms/step - accuracy: 0.1476 - loss: 0.9682\n",
      "Epoch 2: val_loss improved from 0.99290 to 0.89058, saving model to model_ruen.weights.h5\n",
      "\u001b[1m4750/4750\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m708s\u001b[0m 149ms/step - accuracy: 0.1476 - loss: 0.9682 - val_accuracy: 0.1599 - val_loss: 0.8906\n",
      "Epoch 3/3\n",
      "\u001b[1m4750/4750\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m0s\u001b[0m 145ms/step - accuracy: 0.1591 - loss: 0.8709\n",
      "Epoch 3: val_loss improved from 0.89058 to 0.85019, saving model to model_ruen.weights.h5\n",
      "\u001b[1m4750/4750\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m708s\u001b[0m 149ms/step - accuracy: 0.1591 - loss: 0.8709 - val_accuracy: 0.1654 - val_loss: 0.8502\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.src.callbacks.history.History at 0x7f6bddb75180>"
      ]
     },
     "execution_count": 596,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.fit((X_en_train, X_ru_dec_train), X_ru_out_train, \n",
    "             validation_data=((X_en_valid, X_ru_dec_valid), X_ru_out_valid),\n",
    "             batch_size=200,\n",
    "             epochs=3,\n",
    "             callbacks=[checkpoint]\n",
    "             )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d61cb6ed",
   "metadata": {
    "id": "d61cb6ed"
   },
   "source": [
    "Можно попробовать перевести рандомное предложение и оценить качество модели самому"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 597,
   "id": "63762869",
   "metadata": {
    "id": "63762869"
   },
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def translate(text):\n",
    "    input_ids = encode(text, tokenizer_en, target=False)\n",
    "\n",
    "    input_ids = keras.ops.cast(keras.preprocessing.sequence.pad_sequences(\n",
    "                                      [input_ids], maxlen=max_len_en, padding='post', \n",
    "                                       # важно не забыть паддинг с нужным id\n",
    "                                       value=PAD_IDX), torch.int32)\n",
    "\n",
    "    \n",
    "    \n",
    "    output_ids = [tokenizer_ru.token_to_id('[START]') ]\n",
    "    \n",
    "    pred = model((input_ids, keras.ops.cast(keras.preprocessing.sequence.pad_sequences(\n",
    "                                      [output_ids], maxlen=max_len_ru-1, padding='post', \n",
    "                                       # важно не забыть паддинг с нужным id\n",
    "                                       value=PAD_IDX), torch.int32))).cpu().numpy()\n",
    " \n",
    "    \n",
    "    while pred.argmax(2)[0][-1] not in [tokenizer_ru.token_to_id('[END]')]:\n",
    "        if len(output_ids) >= max_len_ru:\n",
    "            break\n",
    "        # можно занизить скор тэга UNK чтобы он никогда не генерировался\n",
    "        pred[:, :, tokenizer_ru.token_to_id('[UNK]')] = -100\n",
    "\n",
    "        output_ids.append(pred.argmax(2)[0][-1])\n",
    "        pred = model((input_ids, keras.ops.cast(keras.preprocessing.sequence.pad_sequences(\n",
    "                                      [output_ids], maxlen=max_len_ru-1, padding='post', \n",
    "                                       # важно не забыть паддинг с нужным id\n",
    "                                       value=PAD_IDX), torch.int32))).cpu().numpy()\n",
    "\n",
    "    return tokenizer_ru.decode(output_ids[1:], )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 600,
   "id": "5f16e8b5",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "5f16e8b5",
    "outputId": "b974fb32-2c36-44bc-8de9-df7007515f58"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'трансформатортортортортор'"
      ]
     },
     "execution_count": 600,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "translate(\"Transformer\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 601,
   "id": "7d0ee1b8",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "7d0ee1b8",
    "outputId": "9fd418f9-1636-4b89-f931-599086753df9"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'можешь перевести это?'"
      ]
     },
     "execution_count": 601,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "translate(\"can you translate this sentence?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 602,
   "id": "aa0b57d2",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "aa0b57d2",
    "outputId": "cdeea84c-8dae-42f4-c63b-b09b2938c697"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'пожалуйста, перевести это на это в русский'"
      ]
     },
     "execution_count": 602,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "translate(\"please translate this sentence into russian\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 395,
   "id": "a5e2071b",
   "metadata": {
    "id": "a5e2071b"
   },
   "outputs": [],
   "source": [
    "# почему без старт и енд в енкодере хуже???\n",
    "# ссылки на доп материалы\n",
    "# как примменяется паддинг маск к атеншену"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "-ZPgGsEl8Cul",
   "metadata": {
    "id": "-ZPgGsEl8Cul"
   },
   "source": [
    "# BLEU\n",
    "\n",
    "Оценивание машинного перевода очень сложная задача. Главная сложность в том, что хороших переводов одного текста может быть несколько, а есть обычно только 1. Технического решения этой проблемы пока не придумали и если нужна достоверная оценка, то используют людей или даже профессиональных переводчиков. Для автоматической оценки стандартно используется простая метрика BLEU. Она считает пересечения между двумя переводами на уровне отдельных токенов, биграммов, триграммов и четырехграммов. Чем больше полных пересечений, тем лучше. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "MX9WvnhW_UNV",
   "metadata": {
    "id": "MX9WvnhW_UNV"
   },
   "source": [
    "Реализация bleu есть в nltk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 396,
   "id": "2735e9d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %pip install nltk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 603,
   "id": "YbuJQX9S8E0o",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "YbuJQX9S8E0o",
    "outputId": "89d78b1a-8c15-44c3-c54b-7a2fc45dcbea"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.4548019047027907\n"
     ]
    }
   ],
   "source": [
    "import nltk\n",
    "\n",
    "hypothesis = ['It', 'is', 'a', 'cat', 'at', 'room']\n",
    "reference = ['It', 'is', 'a', 'cat', 'inside', 'the', 'room']\n",
    "\n",
    "\n",
    "BLEUscore = nltk.translate.bleu_score.sentence_bleu([reference], hypothesis, auto_reweigh=True)\n",
    "print(BLEUscore)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "Rdv_0kXKH94I",
   "metadata": {
    "id": "Rdv_0kXKH94I"
   },
   "source": [
    "Давайте расчитаем BLEU на 10 примерах из тестового корпуса"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 609,
   "id": "1e75c215",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = open('opus.en-ru-test.ru').read().replace('\\xa0', ' ')\n",
    "f = open('opus.en-ru-test.ru', 'w')\n",
    "f.write(text)\n",
    "f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 610,
   "id": "TbLvCxxlFCn6",
   "metadata": {
    "id": "TbLvCxxlFCn6"
   },
   "outputs": [],
   "source": [
    "en_sents_test = open('opus.en-ru-test.en').read().lower().splitlines()\n",
    "ru_sents_test = open('opus.en-ru-test.ru').read().lower().splitlines()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 611,
   "id": "Dy0cRSoLFC0k",
   "metadata": {
    "id": "Dy0cRSoLFC0k"
   },
   "outputs": [],
   "source": [
    "translations = []\n",
    "\n",
    "for i in range(10):\n",
    "    translations.append(translate(en_sents_test[i]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 612,
   "id": "Cuj_lcoUFg9W",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Cuj_lcoUFg9W",
    "outputId": "c5949743-1846-46b1-b758-e4c9471bc956"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['if you only stay there.',\n",
       " \"i don't know how you do it, pop, carrying these boxes around every day.\",\n",
       " 'we might have a slight edge in mediation.',\n",
       " 'how long is it going to take you to get him what he needs?',\n",
       " \"on 1 april president of the nagorno karabagh republic bako sahakyan met head of the general staff of the republic of armenia's armed forces colonel-general yuri khachaturov.\",\n",
       " 'mr priesner also noted that the e-justice management system has not only improved case management, but has also led to a significant streamlining in procedures.',\n",
       " \"you don't like chicken noodle soup?\",\n",
       " 'posted: 14 may 2005, 20:31',\n",
       " 'now, for a minute, i thought maybe he was being tailed.',\n",
       " '« : 26 октябрь 2017, 06:50:24 »']"
      ]
     },
     "execution_count": 612,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "en_sents_test[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 613,
   "id": "pRY7yBESGYxa",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "pRY7yBESGYxa",
    "outputId": "7fbaf103-1e48-41b3-f616-472821d34765"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['только бы не вылететь.',\n",
       " 'и как ты только справляешься, папа, таская эти коробки взад-вперед целый день.',\n",
       " 'возможно, у нас есть небольшое преимущество в переговорах.',\n",
       " 'сколько времени вы будете делать то, что ему нужно?',\n",
       " '1 апреля президент нкр бако саакян принял начальника генштаба вооруженных сил республики армения генерал-полковника юрия хачатурова.',\n",
       " 'г-н приснер также упомянул, что система электронного правосудия не только позволила улучшить процесс ведения дел, но также способствует значительному упорядочению процедур.',\n",
       " '- неплохо, да.',\n",
       " 'posted: 15 dec 2006, 00:07',\n",
       " 'и на минутку я подумал, что за ним могут следить.',\n",
       " '«: 11 октябрь 2011, 17:15:34»']"
      ]
     },
     "execution_count": 613,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ru_sents_test[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 614,
   "id": "6466e6cf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['если ты там не.',\n",
       " 'я не знаю, как ты это делаешь,, все все...... я не знаю................ я не знаю.....',\n",
       " 'у возможно небольшое в посредни дело. мы.,,, возможно,,,,,,,,, возможно,,,,,,,,,,,,,,,,,,.....',\n",
       " 'сколько времени тебя нужно?',\n",
       " '1 президент нагорно карабахской республики бак бако саакянянянянянянянянянянянянянянянянааааааааяяяя..',\n",
       " 'г - - - петербург г - - - - н шустернернерррр такжерр подчеркнулрррр,р подчеркнул,бкркррррр,.....',\n",
       " 'тебе не нравится куриныйный суп??',\n",
       " 'posted : 14 14 20 feb 2005,, 10 :',\n",
       " 'я подумал,,, может он был,, он был на. и я думаю, он не не... я не не не не не не... не..........',\n",
       " '« ответ : # 30 10 10 2011 10, 11 10 11 11 13 14 25 25 25 25 25 25 25 25 25 25 25 25 25 25 25 25 25 25 25 25 25 25 25 25 25 25']"
      ]
     },
     "execution_count": 614,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "translations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "YzWaYs0lIB80",
   "metadata": {
    "id": "YzWaYs0lIB80"
   },
   "source": [
    "BLEU очень сильно зависит от токенизации. В нашем случае это не проблема так как мы оцениваем одну модель. Но если бы мы хотели сравнить несколько, то нужно было бы создавать 1 общий токенизатор и рассчитывать через него"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 615,
   "id": "hx3kUhHQFhMQ",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "hx3kUhHQFhMQ",
    "outputId": "82262875-3ef2-4232-caf3-36ad77f76320"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/.pyenv/versions/3.10.4/lib/python3.10/site-packages/nltk/translate/bleu_score.py:552: UserWarning: \n",
      "The hypothesis contains 0 counts of 2-gram overlaps.\n",
      "Therefore the BLEU score evaluates to 0, independently of\n",
      "how many N-gram overlaps of lower order it contains.\n",
      "Consider using lower n-gram order or use SmoothingFunction()\n",
      "  warnings.warn(_msg)\n",
      "/home/ubuntu/.pyenv/versions/3.10.4/lib/python3.10/site-packages/nltk/translate/bleu_score.py:552: UserWarning: \n",
      "The hypothesis contains 0 counts of 3-gram overlaps.\n",
      "Therefore the BLEU score evaluates to 0, independently of\n",
      "how many N-gram overlaps of lower order it contains.\n",
      "Consider using lower n-gram order or use SmoothingFunction()\n",
      "  warnings.warn(_msg)\n",
      "/home/ubuntu/.pyenv/versions/3.10.4/lib/python3.10/site-packages/nltk/translate/bleu_score.py:552: UserWarning: \n",
      "The hypothesis contains 0 counts of 4-gram overlaps.\n",
      "Therefore the BLEU score evaluates to 0, independently of\n",
      "how many N-gram overlaps of lower order it contains.\n",
      "Consider using lower n-gram order or use SmoothingFunction()\n",
      "  warnings.warn(_msg)\n"
     ]
    }
   ],
   "source": [
    "bleus = []\n",
    "\n",
    "for i, t in enumerate(translations):\n",
    "    reference = tokenizer_ru.encode(t).tokens\n",
    "    hypothesis = tokenizer_ru.encode(ru_sents_test[i]).tokens\n",
    "\n",
    "    bleus.append(round(nltk.translate.bleu_score.sentence_bleu([reference], hypothesis,  auto_reweigh=True), 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 616,
   "id": "rJDX-Z-3GKDE",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "rJDX-Z-3GKDE",
    "outputId": "ee22d497-8f53-4522-89c2-aeb5e329cbed"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.77"
      ]
     },
     "execution_count": 616,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(sum(bleus)/len(bleus))*100"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "Lwuh1mO8Icw0",
   "metadata": {
    "id": "Lwuh1mO8Icw0"
   },
   "source": [
    "BLEU обычно показывают не от 0 до 1, а от 0 до 100. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 617,
   "id": "14270f7d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.0, 0.0, 0.0, 0.0, 0.077, 0.0, 0, 0.0, 0.0, 0.0]"
      ]
     },
     "execution_count": 617,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bleus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66c9170b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}