{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generative networks\n",
    "\n",
    "Recurrent Neural Networks (RNNs) and their gated cell variants such as Long Short Term Memory Cells (LSTMs) and Gated Recurrent Units (GRUs) provided a mechanism for language modeling, i.e. they can learn word ordering and provide predictions for next word in a sequence. This allows us to use RNNs for **generative tasks**, such as ordinary text generation, machine translation, and even image captioning.\n",
    "\n",
    "In RNN architecture we discussed in the previous unit, each RNN unit produced next next hidden state as an output. However, we can also add another output to each recurrent unit, which would allow us to output a **sequence** (which is equal in length to the original sequence). Moreover, we can use RNN units that do not accept an input at each step, and just take some initial state vector, and then produce a sequence of outputs.\n",
    "\n",
    "In this notebook, we will focus on simple generative models that help us generate text. For simplicity, let's build **character-level network**, which generates text letter by letter. During training, we need to take some text corpus, and split it into letter sequences. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading dataset...\n",
      "Building vocab...\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torchtext\n",
    "import numpy as np\n",
    "from torchnlp import *\n",
    "train_dataset,test_dataset,classes,vocab = load_dataset()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building character vocabulary\n",
    "\n",
    "To build character-level generative network, we need to split text into individual characters instead of words. This can be done by defining a different tokenizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vocabulary size = 84\n",
      "Encoding of 'a' is 4\n",
      "Character with code 13 is h\n"
     ]
    }
   ],
   "source": [
    "def char_tokenizer(words):\n",
    "    return list(words) #[word for word in words]\n",
    "\n",
    "counter = collections.Counter()\n",
    "for (label, line) in train_dataset:\n",
    "    counter.update(char_tokenizer(line))\n",
    "vocab = torchtext.vocab.Vocab(counter)\n",
    "\n",
    "vocab_size = len(vocab)\n",
    "print(f\"Vocabulary size = {vocab_size}\")\n",
    "print(f\"Encoding of 'a' is {vocab.stoi['a']}\")\n",
    "print(f\"Character with code 13 is {vocab.itos[13]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see the example of how we can encode the text from our dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([43,  4, 11, 11,  2, 26,  5, 23,  2, 38,  3,  4, 10,  9,  2, 31, 11,  4,\n",
       "        21,  2, 38,  4, 14, 25,  2, 34,  8,  5,  6,  2,  5, 13,  3,  2, 38, 11,\n",
       "         4, 14, 25,  2, 55, 37,  3, 15,  5,  3, 10,  9, 56,  2, 37,  3, 15,  5,\n",
       "         3, 10,  9,  2, 29,  2, 26, 13,  6, 10,  5, 29,  9,  3, 11, 11,  3, 10,\n",
       "         9, 27,  2, 43,  4, 11, 11,  2, 26,  5, 10,  3,  3,  5, 58,  9,  2, 12,\n",
       "        21,  7,  8, 12, 11,  7,  8, 18, 61, 22,  4,  8, 12,  2,  6, 19,  2, 15,\n",
       "        11,  5, 10,  4, 29, 14, 20,  8,  7, 14,  9, 27,  2,  4, 10,  3,  2,  9,\n",
       "         3,  3,  7,  8, 18,  2, 18, 10,  3,  3,  8,  2,  4, 18,  4,  7,  8, 23])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def enc(x):\n",
    "    return torch.LongTensor(encode(x,voc=vocab,tokenizer=char_tokenizer))\n",
    "\n",
    "enc(train_dataset[0][1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training a generative RNN\n",
    "\n",
    "The way we will train RNN to generate text is the following. On each step, we will take a sequence of characters of length `nchars`, and ask the network to generate next output character for each input character:\n",
    "\n",
    "![Image showing an example RNN generation of the word 'HELLO'.](images/rnn-generate.png)\n",
    "\n",
    "Depending on the actual scenario, we may also want to include some special characters, such as *end-of-sequence* `<eos>`. In our case, we just want to train the network for endless text generation, thus we will fix the size of each sequence to be equal to `nchars` tokens. Consequently, each training example will consist of `nchars` inputs and `nchars` outputs (which are input sequence shifted one symbol to the left). Minibatch will consist of several such sequences.\n",
    "\n",
    "The way we will generate minibatches is to take each news text of length `l`, and generate all possible input-output combinations from it (there will be `l-nchars` such combinations). They will form one minibatch, and size of minibatches would be different at each training step. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[43,  4, 11,  ..., 18, 61, 22],\n",
       "         [ 4, 11, 11,  ..., 61, 22,  4],\n",
       "         [11, 11,  2,  ..., 22,  4,  8],\n",
       "         ...,\n",
       "         [37,  3, 15,  ...,  4, 18,  4],\n",
       "         [ 3, 15,  5,  ..., 18,  4,  7],\n",
       "         [15,  5,  3,  ...,  4,  7,  8]], device='cuda:0'),\n",
       " tensor([[ 4, 11, 11,  ..., 61, 22,  4],\n",
       "         [11, 11,  2,  ..., 22,  4,  8],\n",
       "         [11,  2, 26,  ...,  4,  8, 12],\n",
       "         ...,\n",
       "         [ 3, 15,  5,  ..., 18,  4,  7],\n",
       "         [15,  5,  3,  ...,  4,  7,  8],\n",
       "         [ 5,  3, 10,  ...,  7,  8, 23]], device='cuda:0'))"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nchars = 100\n",
    "\n",
    "def get_batch(s,nchars=nchars):\n",
    "    ins = torch.zeros(len(s)-nchars,nchars,dtype=torch.long,device=device)\n",
    "    outs = torch.zeros(len(s)-nchars,nchars,dtype=torch.long,device=device)\n",
    "    for i in range(len(s)-nchars):\n",
    "        ins[i] = enc(s[i:i+nchars])\n",
    "        outs[i] = enc(s[i+1:i+nchars+1])\n",
    "    return ins,outs\n",
    "\n",
    "get_batch(train_dataset[0][1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's define generator network. It can be based on any recurrent cell which we discussed in the previous unit (simple, LSTM or GRU). In our example we will use LSTM.\n",
    "\n",
    "Because the network takes characters as input, and vocabulary size is pretty small, we do not need embedding layer, one-hot-encoded input can directly go to LSTM cell. However, because we pass character numbers as input, we need to one-hot-encode them before passing to LSTM. This is done by calling `one_hot` function during `forward` pass. Output encoder would be a linear layer that will convert hidden state into one-hot-encoded output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LSTMGenerator(torch.nn.Module):\n",
    "    def __init__(self, vocab_size, hidden_dim):\n",
    "        super().__init__()\n",
    "        self.rnn = torch.nn.LSTM(vocab_size,hidden_dim,batch_first=True)\n",
    "        self.fc = torch.nn.Linear(hidden_dim, vocab_size)\n",
    "\n",
    "    def forward(self, x, s=None):\n",
    "        x = torch.nn.functional.one_hot(x,vocab_size).to(torch.float32)\n",
    "        x,s = self.rnn(x,s)\n",
    "        return self.fc(x),s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "During training, we want to be able to sample generated text. To do that, we will define `generate` function that will produce output string of length `size`, starting from the initial string `start`.\n",
    "\n",
    "The way it works is the following. First, we will pass the whole start string through the network, and take output state `s` and next predicted character `out`. Since `out` is one-hot encoded, we take `argmax` to get the index of the character `nc` in the vocabulary, and use `itos` to figure out the actual character and append it to the resulting list of characters `chars`. This process of generating one character is repeated `size` times to generate required number of characters.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate(net,size=100,start='today '):\n",
    "        chars = list(start)\n",
    "        out, s = net(enc(chars).view(1,-1).to(device))\n",
    "        for i in range(size):\n",
    "            nc = torch.argmax(out[0][-1])\n",
    "            chars.append(vocab.itos[nc])\n",
    "            out, s = net(nc.view(1,-1),s)\n",
    "        return ''.join(chars)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's do the training! Training loop is almost the same as in all our previous examples, but instead of accuracy we print sampled generated text every 1000 epochs.\n",
    "\n",
    "Special attention needs to be paid to the way we compute loss. We need to compute loss given one-hot-encoded output `out`, and expected text `text_out`, which is the list of character indices. Luckily, the `cross_entropy` function expects unnormalized network output as first argument, and class number as the second, which is exactly what we have. It also performs automatic averaging over minibatch size.\n",
    "\n",
    "We also limit the training by `samples_to_train` samples, in order not to wait for too long. We encourage you to experiment and try longer training, possibly for several epochs (in which case you would need to create another loop around this code)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Current loss = 4.442246913909912\n",
      "today ggrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrgrg\n",
      "Current loss = 2.1178359985351562\n",
      "today and a could a the to the to the to the to the to the to the to the to the to the to the to the to th\n",
      "Current loss = 1.6465336084365845\n",
      "today on Tuesday the company to the United States and a policing to the United States and a policing to th\n",
      "Current loss = 2.3716814517974854\n",
      "today to the United States and a new men to the United States and a new men to the United States and a new\n",
      "Current loss = 1.6844098567962646\n",
      "today of the first the first the first the first the first the first the first the first the first the fir\n",
      "Current loss = 1.702707052230835\n",
      "today of the United States a said the United States a said the United States a said the United States a sa\n",
      "Current loss = 1.9633255004882812\n",
      "today of the first the first the first the first the first the first the first the first the first the fir\n",
      "Current loss = 1.8642014265060425\n",
      "today of the second a second a second a second a second a second a second a second a second a second a sec\n",
      "Current loss = 1.7720613479614258\n",
      "today and and and the company of the company of the company of the company of the company of the company o\n",
      "Current loss = 1.52818763256073\n",
      "today and the company of the company of the company of the company of the company of the company of the co\n",
      "Current loss = 1.5444810390472412\n",
      "today and the counters to the first the counters to the first the counters to the first the counters to th\n"
     ]
    }
   ],
   "source": [
    "net = LSTMGenerator(vocab_size,64).to(device)\n",
    "\n",
    "samples_to_train = 10000\n",
    "optimizer = torch.optim.Adam(net.parameters(),0.01)\n",
    "loss_fn = torch.nn.CrossEntropyLoss()\n",
    "net.train()\n",
    "for i,x in enumerate(train_dataset):\n",
    "    # x[0] is class label, x[1] is text\n",
    "    if len(x[1])-nchars<10:\n",
    "        continue\n",
    "    samples_to_train-=1\n",
    "    if not samples_to_train: break\n",
    "    text_in, text_out = get_batch(x[1])\n",
    "    optimizer.zero_grad()\n",
    "    out,s = net(text_in)\n",
    "    loss = torch.nn.functional.cross_entropy(out.view(-1,vocab_size),text_out.flatten()) #cross_entropy(out,labels)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    if i%1000==0:\n",
    "        print(f\"Current loss = {loss.item()}\")\n",
    "        print(generate(net))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example already generates some pretty good text, but it can be further improved in several ways:\n",
    "* **Better minibatch generation**. The way we prepared data for training was to generate one minibatch from one sample. This is not ideal, because minibatches are all of different sizes, and some of them even cannot be generated, because the text is smaller than `nchars`. Also, small minibatches do not load GPU sufficiently enough. It would be wiser to get one large chunk of text from all samples, then generate all input-output pairs, shuffle them, and generate minibatches of equal size.\n",
    "* **Multilayer LSTM**. It makes sense to try 2 or 3 layers of LSTM cells. As we mentioned in the previous unit, each layer of LSTM extracts certain patterns from text, and in case of character-level generator we can expect lower LSTM level to be responsible for extracting syllables, and higher levels - for words and word combinations. This can be simply implemented by passing number-of-layers parameter to LSTM constructor.\n",
    "* You may also want to experiment with **GRU units** and see which ones perform better, and with **different hidden layer sizes**. Too large hidden layer may result in overfitting (e.g. network will learn exact text), and smaller size might not produce good result."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Soft text generation and temperature\n",
    "\n",
    "In the previous definition of `generate`, we were always taking the character with highest probability as the next character in generated text. This resulted in the fact that the text often \"cycled\" between the same character sequences again and again, like in this example:\n",
    "```\n",
    "today of the second the company and a second the company ...\n",
    "```\n",
    "\n",
    "However, if we look at the probability distribution for the next character, it could be that the difference between a few highest probabilities is not huge, e.g. one character can have probability 0.2, another - 0.19, etc. For example, when looking for the next character in the sequence '*play*', next character can equally well be either space, or **e** (as in the word *player*).\n",
    "\n",
    "This leads us to the conclusion that it is not always \"fair\" to select the character with higher probability, because choosing the second highest might still lead us to meaningful text. It is more wise to **sample** characters from the probability distribution given by the network output.\n",
    "\n",
    "This sampling can be done using `multinomial` function that implements so-called **multinomial distribution**. A function that implements this **soft** text generation is defined below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Temperature = 0.3\n",
      "Today and to has a software to in the first the power the gold medal was of the first and succer to the company will a report the first the and the gain the company in the and a new a report a pack of the four the first the company of the such with the half to a security to the and a success the first she\n",
      "\n",
      "--- Temperature = 0.8\n",
      "Today drud out of the three-rent possiem that sales purssion has finminiaty women's from NAC Inc. (AP) -- Shimbon has weel with a may stelight first three flaw gold from their a scent, big study with a nighting sovicturner has slarh football at a hour of Angelage discression, into cubs, US year player sor\n",
      "\n",
      "--- Temperature = 1.0\n",
      "Today by compoy, said to hup the couns ay rrope iist\\fill sinie-5-1- than he of a fightier Corp. the Vew, Mkli Unite Hold Austria on Tuesday resfare rextarted in the new has buy thisnillials thrust first capuration of the it larget expected the ir edulagy Airin Penny after Emonet Cuc Washieve an are Gurry\n",
      "\n",
      "--- Temperature = 1.3\n",
      "Today cluscy,, wangled and-ox they, stee of as;\\seculity dillancrile inmution svanse gall ATHEYS today a first oresift 6-Jalf mangback explymate that wrook\" haffic illowbre overwage in Tecrian Hunrieleers to attowny service Adching, blanks governine? Aug. : : NE: Sir NFP (P2AAU) Bow SWDE: The ex2\"cut Pmoc\n",
      "\n",
      "--- Temperature = 1.8\n",
      "Today sas gom, twing hWe a Dajfcou hamb--5 to bemolecresem ig irkembets plentll repws, scatchey: Actuss.io Theffouge, cirw biggemed Goiga propperinut you racive #5-Aeia:riato..Lf. N7TNap:,ser,wploy a Fraull tbashonerdlantuanseve /bBT -$06 Wemob-e.P EvVlaicy), ZOf0 cUSeballd sturk out houselty, TARZM) AbAe\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def generate_soft(net,size=100,start='today ',temperature=1.0):\n",
    "        chars = list(start)\n",
    "        out, s = net(enc(chars).view(1,-1).to(device))\n",
    "        for i in range(size):\n",
    "            #nc = torch.argmax(out[0][-1])\n",
    "            out_dist = out[0][-1].div(temperature).exp()\n",
    "            nc = torch.multinomial(out_dist,1)[0]\n",
    "            chars.append(vocab.itos[nc])\n",
    "            out, s = net(nc.view(1,-1),s)\n",
    "        return ''.join(chars)\n",
    "    \n",
    "for i in [0.3,0.8,1.0,1.3,1.8]:\n",
    "    print(f\"--- Temperature = {i}\\n{generate_soft(net,size=300,start='Today ',temperature=i)}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have introduced one more parameter called **temperature**, which is used to indicate how hard we should stick to the highest probability. If temperature is 1.0, we do fair multinomial sampling, and when temperature goes to infinity - all probabilities become equal, and we randomly select next character. In the example below we can observe that the text becomes meaningless when we increase the temperature too much, and it resembles \"cycled\" hard-generated text when it becomes closer to 0. "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py37_pytorch",
   "language": "python",
   "name": "conda-env-py37_pytorch-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
