week13/lstm_language_model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Генерация текста при помощи LSTM\n",
    "\n",
    "На этом семинаре поговорим про RNN. Здесь мы обучим модель на тексте Анны Карениной, а затем будем генерировать новый текст. **Эта модель сможет генерировать новый текст на основе нашего текста!**\n",
    "\n",
    "Есть интересная [статья об RNNs](http://karpathy.github.io/2015/05/21/rnn-effectiveness/), показывающая впечатляющие примеры того, что рекуррентные сети умеют делать. Ниже представлена ​​общая архитектура RNN.\n",
    "\n",
    "<img src=\"https://github.com/udacity/deep-learning-v2-pytorch/blob/master/recurrent-neural-networks/char-rnn/assets/charseq.jpeg?raw=1\" width=\"500\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Сначала загрузим ресурсы, необходимые для загрузки данных и создания модели."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "from pathlib import Path\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "import requests\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.tensorboard import SummaryWriter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_datetime():\n",
    "    return datetime.datetime.now().isoformat(sep='_', timespec='milliseconds').replace(':', '-')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Загружаем данные\n",
    "\n",
    "Затем мы загрузим текстовый файл Анны Карениной."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text_path = Path('anna.txt')\n",
    "\n",
    "if not text_path.exists():\n",
    "    !wget https://github.com/hse-ds/iad-deep-learning/blob/66fb012/sem10/anna.txt -O $text_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# open text file and read in data as `text`\n",
    "with text_path.open() as fp:\n",
    "    text = fp.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Давайте проверим первые 100 символов, убедимся, что все красиво. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text[:100]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Преобразуем символы в числа\n",
    "\n",
    "В ячейках ниже постараемся создать пару **словарей** для преобразования символов в целые числа и обратно. Кодирование символов как целых чисел упрощает их использование в качестве входных данных в сети."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# encode the text and map each character to an integer and vice versa\n",
    "\n",
    "# we create two dictionaries:\n",
    "# 1. int2char, which maps integers to characters\n",
    "# 2. char2int, which maps characters to unique integers\n",
    "chars = sorted(set(text))\n",
    "print(f'Number of distinct chars in text: {len(chars)}')\n",
    "print(repr(''.join(chars)))\n",
    "\n",
    "int2char = {i: ch for i, ch in enumerate(chars)}\n",
    "char2int = {ch: i for i, ch in int2char.items()}\n",
    "\n",
    "# encode the text\n",
    "text_encoded = np.array([char2int[ch] for ch in text])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "И мы можем видеть те же самые символы сверху, закодированные как целые числа."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text_encoded[:100]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Вопрос:** почему нельзя было просто воспользоваться функцией `ord()`?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print([ord(c) for c in chars])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## One-hot encoding\n",
    "\n",
    "Как вы можете видеть на изображении char-RNN выше, наш LSTM ожидает ввода, который  **one-hot encoded** , что означает, что каждый символ преобразуется в целое число (через наш созданный словарь), а затем преобразуется в столбец вектор, где только соответствующий ему целочисленный индекс будет иметь значение 1, а остальная часть вектора будет заполнена нулями. Давайте создадим для этого функцию!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def one_hot_encode(arr, n_labels):\n",
    "    # Initialize the the encoded array\n",
    "    one_hot = np.zeros((arr.size, n_labels), dtype=np.float32)\n",
    "    \n",
    "    # Fill the appropriate elements with ones\n",
    "    one_hot[np.arange(one_hot.shape[0]), arr.flatten()] = 1.\n",
    "    \n",
    "    # Finally reshape it to get back to the original array\n",
    "    one_hot = one_hot.reshape((*arr.shape, n_labels))\n",
    "    \n",
    "    return one_hot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check that the function works as expected\n",
    "test_seq = np.array([\n",
    "    [3, 5, 1],\n",
    "    [7, 2, 5],\n",
    "])\n",
    "one_hot = one_hot_encode(test_seq, 8)\n",
    "\n",
    "print(one_hot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Делаем батч-генератор\n",
    "\n",
    "\n",
    "Для обучения на этих данных нужно создать мини-батчи для обучения. На простом примере наши батчи будут выглядеть так:\n",
    "\n",
    "<img src=\"https://github.com/udacity/deep-learning-v2-pytorch/blob/master/recurrent-neural-networks/char-rnn/assets/sequence_batching@1x.png?raw=1\" width=500px>\n",
    "\n",
    "\n",
    "<br>\n",
    "\n",
    "В этом примере возьмем закодированные символы (переданные как параметр `arr`) и разделим их на несколько последовательностей, заданных параметром `batch_size`. Каждая из наших последовательностей будет иметь длину `seq_length`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Собираем батчи\n",
    "\n",
    "Наш батч-генератор будет работать следующим образом. Сначала мы выделим в тексте позиции, с которых будут начинаться подпоследовательности, которые станут элементами батчей. Такие подпоследовательности будут иметь одинаковую длину:\n",
    "\n",
    "```\n",
    "Everything was in confusion in the Oblonskys' house. The wife had discovered that the husband was carrying on an intrigue\n",
    "↑               ↑               ↑               ↑               ↑               ↑               ↑               ↑               \n",
    "```\n",
    "\n",
    "Мы будем воспринимать текст как набор из таких подпоследовательностей:\n",
    "\n",
    "```\n",
    "Everything was i\n",
    "n confusion in t\n",
    "he Oblonskys' ho\n",
    "use. The wife ha\n",
    "d discovered tha\n",
    "t the husband wa\n",
    "s carrying on an\n",
    " intrigue\n",
    "```\n",
    "\n",
    "Последнюю подпоследовательность, которая будет слишком короткой, мы просто выбросим:\n",
    "\n",
    "```\n",
    "Everything was i\n",
    "n confusion in t\n",
    "he Oblonskys' ho\n",
    "use. The wife ha\n",
    "d discovered tha\n",
    "t the husband wa\n",
    "s carrying on an\n",
    "```\n",
    "\n",
    "Оставшиеся подпоследовательности мы перемешаем в случайном порядке:\n",
    "\n",
    "```\n",
    "use. The wife ha\n",
    "t the husband wa\n",
    "he Oblonskys' ho\n",
    "d discovered tha\n",
    "n confusion in t\n",
    "Everything was i\n",
    "s carrying on an\n",
    "```\n",
    "\n",
    "Наконец, мы разобьём их на батчи:\n",
    "\n",
    "```\n",
    "use. The wife ha\n",
    "t the husband wa\n",
    "he Oblonskys' ho\n",
    "\n",
    "d discovered tha\n",
    "n confusion in t\n",
    "Everything was i\n",
    "\n",
    "s carrying on an\n",
    "```\n",
    "\n",
    "и для простоты выбросим слишком короткие батчи:\n",
    "\n",
    "```\n",
    "use. The wife ha\n",
    "t the husband wa\n",
    "he Oblonskys' ho\n",
    "\n",
    "d discovered tha\n",
    "n confusion in t\n",
    "Everything was i\n",
    "```\n",
    "\n",
    "Это то, что будет подаваться в модель на вход, то есть `x`. В нашем случае `y` будет теми же последовательностями, но сдвинутыми на один символ:\n",
    "\n",
    "```\n",
    "use. The wife ha         se. The wife had\n",
    "t the husband wa    →     the husband was\n",
    "he Oblonskys' ho         e Oblonskys' hou\n",
    "\n",
    "d discovered tha          discovered that\n",
    "n confusion in t    →     confusion in th\n",
    "Everything was i         verything was in\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_batches(arr, batch_size, seq_length):\n",
    "    '''Create a generator that returns batches of size\n",
    "       batch_size x seq_length from arr.\n",
    "       \n",
    "       Arguments\n",
    "       ---------\n",
    "       arr: Array you want to make batches from\n",
    "       batch_size: Batch size, the number of sequences per batch\n",
    "       seq_length: Number of encoded chars in a sequence\n",
    "    '''\n",
    "    assert len(arr.shape) == 1\n",
    "\n",
    "    # Сделайте массив start_indices из позиций начал подпоследовательностей в тексте\n",
    "    <YOUR CODE>\n",
    "    assert start_indices[-1] + seq_length + 1 <= len(arr)\n",
    "\n",
    "    # Перемешайте его\n",
    "    <YOUR CODE>\n",
    "\n",
    "    # Сделайте так, чтобы количество элементов в нём делилось на batch_size\n",
    "    <YOUR CODE>\n",
    "\n",
    "    # Пройдите по start_indices скользящим окном с шириной и шагом batch_size\n",
    "    for i in <YOUR CODE>:\n",
    "        batch_start_indices = <YOUR CODE>\n",
    "\n",
    "        # Сформируйте x и y\n",
    "        <YOUR CODE>\n",
    "\n",
    "        assert x.shape == (batch_size, seq_length), x.shape\n",
    "        assert y.shape == (batch_size, seq_length), y.shape\n",
    "        \n",
    "        yield x, y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Тестируем батч-генератор\n",
    "\n",
    "Теперь создадим несколько наборов данных, и проверим, что происходит, когда мы батчуем данные."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batches = get_batches(text_encoded, 8, 50)\n",
    "x, y = next(batches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# printing out the first 10 items in a sequence\n",
    "print('x\\n', x[:10, :10])\n",
    "print('\\ny\\n', y[:10, :10])\n",
    "assert (x[:, 1:] == y[:, :-1]).all(), 'y does not seem to be a shifted copy of x!'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Если вы правильно реализовали get_batches, результат должен выглядеть примерно так:\n",
    "```\n",
    "x\n",
    " [[25  8 60 11 45 27 28 73  1  2]\n",
    " [17  7 20 73 45  8 60 45 73 60]\n",
    " [27 20 80 73  7 28 73 60 73 65]\n",
    " [17 73 45  8 27 73 66  8 46 27]\n",
    " [73 17 60 12 73  8 27 28 73 45]\n",
    " [66 64 17 17 46  7 20 73 60 20]\n",
    " [73 76 20 20 60 73  8 60 80 73]\n",
    " [47 35 43  7 20 17 24 50 37 73]]\n",
    "\n",
    "y\n",
    " [[ 8 60 11 45 27 28 73  1  2  2]\n",
    " [ 7 20 73 45  8 60 45 73 60 45]\n",
    " [20 80 73  7 28 73 60 73 65  7]\n",
    " [73 45  8 27 73 66  8 46 27 65]\n",
    " [17 60 12 73  8 27 28 73 45 27]\n",
    " [64 17 17 46  7 20 73 60 20 80]\n",
    " [76 20 20 60 73  8 60 80 73 17]\n",
    " [35 43  7 20 17 24 50 37 73 36]]\n",
    " ```\n",
    " хотя точные цифры могут отличаться. Убедитесь, что данные сдвинуты на один шаг для `y`!!!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Создаём модель\n",
    "\n",
    "\n",
    "<img src=\"https://github.com/udacity/deep-learning-v2-pytorch/blob/master/recurrent-neural-networks/char-rnn/assets/charRNN.png?raw=1\" width=500px>\n",
    "\n",
    "Затем используем PyTorch для определения архитектуры сети. Начнем с определения слоев и операций, методов прямого прохода. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Устройство модели\n",
    "\n",
    "В шаблоне `__init__()` уже прописано создание и сохранение необходимых словарей. Далее там предлагается описать следующую архитектуру:\n",
    "\n",
    "* Слой LSTM, который принимает в качестве параметров:\n",
    "    * `input_size`: размерность входа. В нашем случае это количество символов\n",
    "    * `hidden_size`: размер скрытого слоя `hidden_size`\n",
    "    * `num_layers`: количество слоёв ` num_layers`\n",
    "    * `dropout`: вероятность дропаута `dropout_prob`\n",
    "    * флаг `batch_first=True` (потому что модель на вход будет получать батчи с шейпом `batch x time x units`, а не `time x batch x units`)\n",
    "* Слой Dropout с параметром `dropout_prob`\n",
    "* Полносвязный слой с параметрами: размер ввода `hidden_size`, размер выхода — количество символов\n",
    "\n",
    "Обратите внимание, что некоторые параметры были названы и указаны в функции `__init__`, их нужно сохранить и использовать, выполняя что-то вроде` self.dropout_prob = dropout_prob`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check if GPU is available\n",
    "device = 'cuda:0' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "if device == 'cpu':\n",
    "    print('No GPU available, training on CPU; consider making n_epochs very small.')\n",
    "else:\n",
    "    print('Training on GPU!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CharRNN(nn.Module):\n",
    "    def __init__(self, tokens, hidden_size=256, num_layers=2, dropout_prob=0.5):\n",
    "        super().__init__()\n",
    "        self.dropout_prob = dropout_prob\n",
    "        self.num_layers = num_layers\n",
    "        self.hidden_size = hidden_size\n",
    "        \n",
    "        # creating character dictionaries\n",
    "        self.chars = tokens\n",
    "        self.int2char = {i: ch for i, ch in enumerate(self.chars)}\n",
    "        self.char2int = {ch: i for i, ch in self.int2char.items()}\n",
    "        \n",
    "        ## TODO: define the layers of the model\n",
    "        \n",
    "        <YOUR CODE>\n",
    "      \n",
    "    \n",
    "    def forward(self, x, hidden):\n",
    "        ''' Forward pass through the network. \n",
    "            These inputs are x, and the hidden/cell state `hidden`. '''\n",
    "                \n",
    "        ## TODO: Get the outputs and the new hidden state from the lstm\n",
    "\n",
    "        <YOUR CODE>\n",
    "        \n",
    "        # return the final output and the hidden state\n",
    "        return out, hidden\n",
    "    \n",
    "    \n",
    "    def init_hidden(self, batch_size):\n",
    "        ''' Initializes hidden state '''\n",
    "        # Create two new tensors with sizes num_layers x batch_size x hidden_size,\n",
    "        # initialized to zero, for hidden state and cell state of LSTM\n",
    "\n",
    "        hidden_shape = (self.num_layers, batch_size, self.hidden_size)\n",
    "        h0 = torch.zeros(hidden_shape, device=device)\n",
    "        c0 = torch.zeros(hidden_shape, device=device)\n",
    "        \n",
    "        return h0, c0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## TODO: set your model hyperparameters\n",
    "# define and print the net\n",
    "hidden_size = <YOUR CODE>\n",
    "num_layers = <YOUR CODE>\n",
    "\n",
    "net = CharRNN(chars, hidden_size=hidden_size, num_layers=num_layers)\n",
    "print(net)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Обучающий цикл\n",
    "\n",
    "Во время обучения нужно установить количество эпох, скорость обучения и другие параметры.\n",
    "\n",
    "Используем оптимизатор Адама и кросс-энтропию, считаем loss и, как обычно, выполняем back propagation!\n",
    "\n",
    "Пара подробностей об обучении:\n",
    "\n",
    "* На каждой итерации мы делаем `.detach()` на скрытом состоянии LSTM, чтобы не делать бэкпроп через всю историю обучения.\n",
    "* Мы используем [`clip_grad_norm_`](https://pytorch.org/docs/stable/_modules/torch/nn/utils/clip_grad.html), чтобы избавиться от взрывающегося градиента."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cross_entropy_batch_time_class(input, target):\n",
    "    return F.cross_entropy(input.transpose(1, 2), target)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(net, data, tb_dir, tb_tag=None, epochs=10, batch_size=10, seq_length=50, lr=0.001, clip=5, val_frac=0.1, val_frequency=10):\n",
    "    ''' Training a network \n",
    "    \n",
    "        Arguments\n",
    "        ---------\n",
    "        \n",
    "        net: CharRNN network\n",
    "        data: text data to train the network\n",
    "        epochs: Number of epochs to train\n",
    "        batch_size: Number of mini-sequences per mini-batch, aka batch size\n",
    "        seq_length: Number of character steps per mini-batch\n",
    "        lr: learning rate\n",
    "        clip: gradient clipping\n",
    "        val_frac: Fraction of data to hold out for validation\n",
    "        val_frequency: Number of steps for printing training and validation loss\n",
    "    \n",
    "    '''\n",
    "    net.train()\n",
    "    \n",
    "    opt = torch.optim.Adam(net.parameters(), lr=lr)\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    \n",
    "    # create training and validation data\n",
    "    val_idx = int(len(data)*(1-val_frac))\n",
    "    data, val_data = data[:val_idx], data[val_idx:]\n",
    "\n",
    "    net = net.to(device)\n",
    "\n",
    "    if tb_tag is None:\n",
    "        tb_run_name = get_datetime()\n",
    "    else:\n",
    "        tb_run_name = f'{get_datetime()}_{tb_tag}'\n",
    "    \n",
    "    num_samples = 0\n",
    "    with SummaryWriter(log_dir=str(tb_dir / tb_run_name)) as writer:\n",
    "        for e in range(epochs):\n",
    "            # initialize hidden state\n",
    "            hidden = net.init_hidden(batch_size)\n",
    "            \n",
    "            with tqdm(get_batches(data, batch_size, seq_length), total=len(data) // (batch_size * seq_length)) as progress_bar:\n",
    "                for x, y in progress_bar:\n",
    "                    num_samples += batch_size\n",
    "\n",
    "                    net.train()\n",
    "                    \n",
    "                    # One-hot encode our data and make them Torch tensors\n",
    "                    x = one_hot_encode(x, len(net.chars))\n",
    "                    inputs = torch.tensor(x, device=device)\n",
    "                    targets = torch.tensor(y, device=device, dtype=torch.int64)\n",
    "\n",
    "                    # get the output from the model\n",
    "                    output, hidden = net(inputs, hidden)\n",
    "\n",
    "                    # Creating new variables for the hidden state, otherwise\n",
    "                    # we'd backprop through the entire training history\n",
    "                    hidden = tuple(h.detach() for h in hidden)\n",
    "                    \n",
    "                    # calculate the loss and perform backprop\n",
    "                    loss = cross_entropy_batch_time_class(output, targets)\n",
    "\n",
    "                    opt.zero_grad()  # zero accumulated gradients\n",
    "                    loss.backward()\n",
    "                    # `clip_grad_norm` helps prevent the exploding gradient problem in RNNs / LSTMs.\n",
    "                    nn.utils.clip_grad_norm_(net.parameters(), clip)\n",
    "                    opt.step()\n",
    "\n",
    "                    writer.add_scalar('loss/train', loss.item(), num_samples)\n",
    "                    \n",
    "                    # loss stats\n",
    "                    if (num_samples // batch_size) % val_frequency == 0:\n",
    "                        with torch.no_grad():\n",
    "                            net.eval()\n",
    "\n",
    "                            # Get validation loss\n",
    "                            val_hidden = net.init_hidden(batch_size)\n",
    "                            val_losses = []\n",
    "                            for x, y in get_batches(val_data, batch_size, seq_length):\n",
    "                                # One-hot encode our data and make them Torch tensors\n",
    "                                x = one_hot_encode(x, len(net.chars))\n",
    "                                inputs = torch.tensor(x, device=device)\n",
    "                                targets = torch.tensor(y, device=device, dtype=torch.int64)\n",
    "\n",
    "                                output, val_hidden = net(inputs, val_hidden)\n",
    "                                val_loss = cross_entropy_batch_time_class(output, targets)\n",
    "                            \n",
    "                                val_losses.append(val_loss.item())\n",
    "                            \n",
    "                            writer.add_scalar('loss/valid', np.mean(val_losses), num_samples)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Создаём модель"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Задаём гиперпараметры\n",
    "\n",
    "Теперь мы можем создать модель с заданными гиперпараметрами. Определим размеры мини-батчей!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = <YOUR CODE>\n",
    "seq_length = <YOUR CODE>\n",
    "n_epochs = <YOUR CODE>  # start small if you are just testing initial behavior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tb_dir = Path('tb_logs')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%tensorboard --port 6006 --logdir $tb_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train the model\n",
    "train(\n",
    "    net, text_encoded, tb_dir,\n",
    "    tb_tag=f'bs{batch_size}_sl{seq_length}_hs{hidden_size}_nl{num_layers}',\n",
    "    epochs=n_epochs, batch_size=batch_size, seq_length=seq_length, val_frequency=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Улучшаем модель\n",
    "\n",
    "Чтобы выбрать оптимальные гиперпараметры, нужно посмотреть на значения функции потерь на обучении и валидации. Если ваш лосс на обучении намного ниже, чем на тесте, то модель переобучена. Увеличьте регуляризацию или уменьшите число слоев в сети."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Гиперпараметры\n",
    "\n",
    "Гиперпараметры сети:\n",
    "\n",
    "* `hidden_size` - Количество юнитов в скрытых слоях.\n",
    "* `num_layers` - Количество используемых скрытых слоев LSTM.\n",
    "\n",
    "В нашем примере вероятность отсева и скорость обучения сохраняется.\n",
    "\n",
    "Для обучения:\n",
    "* `batch_size` - количество объектов в батче, проходящих по сети за один проход.\n",
    "* `seq_length` - Количество символов в последовательности, на которой обучается сеть. Обычно чем больше, тем лучше, сеть будет изучать более дальние зависимости.\n",
    "* `lr` - learning rate.\n",
    "\n",
    "\n",
    " ## Советы и хитрости\n",
    "\n",
    "> -  В глубоком обучении очень распространено запускать множество различных моделей с множеством различных настроек гиперпараметров и, в конце концов, использовать любую контрольную точку, дающую наилучшее качество на валидации.\n",
    "\n",
    "> - Кстати, размер ваших тренировочных и проверочных разделов также является параметрами. Убедитесь, что у вас есть приличный объем данных в вашей валидационной выборки, иначе производительность проверки будет шумной и не очень информативной."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checkpoint\n",
    "\n",
    "После обучения сохраним модель, чтобы можно было загрузить ее позже. Здесь сохраняются параметры, необходимые для создания той же архитектуры, гиперпараметры скрытого слоя и токены."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# change the name, for saving multiple files\n",
    "model_path = Path('lstm.pth')\n",
    "\n",
    "checkpoint = {\n",
    "    'hidden_size': net.hidden_size,\n",
    "    'num_layers': net.num_layers,\n",
    "    'state_dict': net.state_dict(),\n",
    "    'tokens': net.chars,\n",
    "}\n",
    "\n",
    "with model_path.open('wb') as fp:\n",
    "    torch.save(checkpoint, fp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Делаем предсказания\n",
    "\n",
    "Теперь, когда модель обучена, сделаем предсказание следующих символов! Для предсказания мы передаем последний символ обучения, и сеть предсказывает следующий символ, который мы потом передаем обратно и получаем еще один предсказанный символ и так далее...\n",
    "\n",
    "Наши прогнозы основаны на категориальном распределении вероятностей по всем возможным символам. Мы можем ограничить число символов, чтобы сделать получаемый предсказанный текст более разумным, рассматривая только некоторые наиболее вероятные символы $K$. Это не позволит сети выдавать нам совершенно абсурдные символы, а также позволит внести некоторый шум и случайность в выбранный текст. Узнать больше [можно здесь](https://pytorch.org/docs/stable/torch.html#torch.topk).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(net, char, hidden, top_k=None):\n",
    "        ''' Given a character, predict the next character.\n",
    "            Returns the predicted character and the hidden state.\n",
    "        '''\n",
    "        \n",
    "        # tensor inputs\n",
    "        x = np.array([[net.char2int[char]]])\n",
    "        x = one_hot_encode(x, len(net.chars))\n",
    "        inputs = torch.from_numpy(x).to(device)\n",
    "        \n",
    "        # detach hidden state from history\n",
    "        hidden = tuple(h.detach() for h in hidden)\n",
    "        # get the output of the model\n",
    "        out, hidden = net(inputs, hidden)\n",
    "\n",
    "        # get the character probabilities\n",
    "        char_probs = F.softmax(out, dim=-1).cpu().detach()\n",
    "\n",
    "        # get top characters\n",
    "        if top_k is None:\n",
    "            top_ch = np.arange(len(net.chars))\n",
    "        else:\n",
    "            # https://pytorch.org/docs/stable/generated/torch.topk.html\n",
    "            char_probs, top_ch = char_probs.topk(top_k)\n",
    "            top_ch = top_ch.numpy().squeeze()\n",
    "        \n",
    "        # select the likely next character with some element of randomness\n",
    "        char_probs = char_probs.numpy().squeeze(axis=(0, 1))\n",
    "        char = np.random.choice(top_ch, p=char_probs/char_probs.sum())\n",
    "        \n",
    "        # return the encoded value of the predicted char and the hidden state\n",
    "        return net.int2char[char], hidden"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predict(net, 'A', net.init_hidden(batch_size=1), top_k=5)[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Инициализируем скрытое состояние и генерируем текст\n",
    "\n",
    "Нужно создать скрытое состояние, чтобы сеть не генерировала произвольные символы. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample(net, size, prime='The', top_k=None):\n",
    "    net = net.to(device)\n",
    "    net.eval()\n",
    "    \n",
    "    # First off, run through the prime characters\n",
    "    chars = list(prime)\n",
    "    hidden = net.init_hidden(batch_size=1)\n",
    "    for ch in prime:\n",
    "        char, h = predict(net, ch, hidden, top_k=top_k)\n",
    "\n",
    "    chars.append(char)\n",
    "    \n",
    "    # Now pass in the previous character and get a new one\n",
    "    for _ in range(size):\n",
    "        char, hidden = predict(net, chars[-1], hidden, top_k=top_k)\n",
    "        chars.append(char)\n",
    "\n",
    "    return ''.join(chars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(sample(net, 1000, prime='Anna', top_k=5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Загружаем чекпойнт"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we have loaded in a model that trained over 20 epochs `rnn_20_epoch.net`\n",
    "with Path('lstm.pth').open('rb') as fp:\n",
    "    checkpoint = torch.load(fp)\n",
    "    \n",
    "loaded = CharRNN(checkpoint['tokens'], hidden_size=checkpoint['hidden_size'], num_layers=checkpoint['num_layers'])\n",
    "loaded.load_state_dict(checkpoint['state_dict'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sample using a loaded model\n",
    "print(sample(loaded, 2000, top_k=5, prime=\"And Levin said\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Acknowledgements\n",
    "\n",
    "Ноутбук основан на [ноутбуке для десятого семинара](https://github.com/hse-ds/iad-deep-learning/blob/66fb0128da4e65cb3260c088e2d462eb9d0c5eb1/sem10/sem10.ipynb) курса DL на майноре ИАД ВШЭ, который, в свою очередь, основан на [ноутбуке](https://github.com/udacity/deep-learning-v2-pytorch/blob/704a3e54b78afb8cd64d001d5160db93c986f63a/recurrent-neural-networks/char-rnn/Character_Level_RNN_Solution.ipynb) из курса Udacity Deep Learning."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}