notebooks/Talk.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[nltk_data] Downloading package wordnet to /home/carson/nltk_data...\n",
      "[nltk_data]   Package wordnet is already up-to-date!\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "importing Jupyter notebook from MoveData.ipynb\n",
      "importing Jupyter notebook from EncoderDecoder.ipynb\n",
      "importing Jupyter notebook from Elements.ipynb\n"
     ]
    }
   ],
   "source": [
    "import re, math\n",
    "import numpy as np\n",
    "\n",
    "import nltk\n",
    "nltk.download('wordnet') \n",
    "from nltk.corpus import wordnet\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.distributions.categorical import Categorical\n",
    "from torch.autograd import Variable\n",
    "import torch.nn.functional as F \n",
    "import import_ipynb\n",
    "from MoveData import Options, json2datatools, num_batches, nopeak_mask, create_masks\n",
    "from EncoderDecoder import Encoder, Decoder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Neural Network as an Agent\n",
    "\n",
    "In this section we will be putting everything together. We learned AI to make AI robots so let finally do that. \n",
    "\n",
    "If you are not using this notebook to learn, change the below variable `teaching` to `False` so that other notebooks can import the functions defined in this notebook without running all the examples, if you are here to learn and interact with the notebook, change it to `True`\n",
    "\n",
    "The cell below you have seen before, we will need the input and output vocabulary fields `infield, outfield` for our demonstration of how a sequence of words is represented by the transformer and the role that probability plays in the model's outputs. Note, for the time being, we are assuming that the vocabulary consists of all the words in the training set, nothing more nothing less, so if since running START_HERE you have added some lines of data, simply retrain the model from START_HERE before moving on"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "teaching = False\n",
    "\n",
    "if teaching:\n",
    "    opt = Options(batchsize=2, device = torch.device(\"cpu\"), epochs=25, lr=0.01, \n",
    "                  max_len = 25, save_path = '../saved/weights/transformer_example_weights')\n",
    "\n",
    "    data_iter, infield, outfield, opt = json2datatools(path='../saved/pairs.json', opt=opt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Until we built more capabilities into chloe, we will have to reply on a few software tools. The first tool is a tool for expanding chloe's vocabulary without having to learn them from scratch.\n",
    "\n",
    "nltk is the [Natural Language Toolkit](https://www.nltk.org/) that we will be using for things such as synonym matching, that way when you say \"adore\", Chloe knows it means the same thing as \"like\", even if \"adore\" is not in Chloe's vocabulary. \n",
    "\n",
    "Run the cell below to see how we go from string words to integers and how when we cannot find a word in our vocabulary, we try to find a synonym for that word that is in our vocabulary. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_synonym(word, field, explain=False):\n",
    "    syns = wordnet.synsets(word)\n",
    "    for s in syns:\n",
    "        if explain: print('synonym:', s.name())\n",
    "        for l in s.lemmas():\n",
    "            if explain: print('-lemma:', l.name())\n",
    "            if field.vocab.stoi[l.name()] != 0:\n",
    "                if explain: print('found in vocab', l.name())\n",
    "                return field.vocab.stoi[l.name()]\n",
    "    return 0 # if we cannot find a synonym, return 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "synonym: fine.n.01\n",
      "-lemma: fine\n",
      "-lemma: mulct\n",
      "-lemma: amercement\n",
      "synonym: ticket.v.01\n",
      "-lemma: ticket\n",
      "-lemma: fine\n",
      "synonym: all_right.s.01\n",
      "-lemma: all_right\n",
      "-lemma: fine\n",
      "-lemma: o.k.\n",
      "-lemma: ok\n",
      "found in vocab ok\n",
      "token =  13\n"
     ]
    }
   ],
   "source": [
    "if teaching:\n",
    "    print('token = ', get_synonym(\"fine\", infield, explain=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next function takes your sentence in the form of text and converts it to a sequence of tokens within a torch tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def string2tensor(string, inputfield, explain=False):\n",
    "    '''\n",
    "    input:\n",
    "        string (str) input sentence\n",
    "        inputfield a PyTorch torchtext.data.Field object\n",
    "        explain, set this to True if you want to see how the sentence was split \n",
    "    output:\n",
    "        sequence of tokens (torch tensor of integers) shape  \n",
    "    '''\n",
    "    sentence = inputfield.preprocess(string)\n",
    "    if explain: print(sentence)\n",
    "    integer_sequence = []\n",
    "    for tok in sentence:\n",
    "        if inputfield.vocab.stoi[tok] != 0:\n",
    "            integer_sequence.append(inputfield.vocab.stoi[tok])\n",
    "        else:\n",
    "            integer_sequence.append(get_synonym(tok, inputfield))\n",
    "    return torch.LongTensor([integer_sequence])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['ok', ',', \"aren't\", 'you', 'a', 'robot', '?']\n",
      "tensor([[13,  0,  0,  3, 15, 37,  2]])\n"
     ]
    }
   ],
   "source": [
    "if teaching:\n",
    "    input_sequence = string2tensor(\"ok, aren't -you [a robot?\", infield, explain=True)\n",
    "    print(input_sequence)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets assemble the Encoder and Decoder into the Transformer, like older sequence to sequence models, Transformers also encode the sentence into vector representations and pass those representations along to the decoder to generate the response/reply/output/translation/etc. The Encoder and Decoder we can use separately such as `model.encoder(arguments)`. The very last part of the transformer is the mapping of each vector in the decoder output to logits for each token in the output vocabulary `output = self.out(d_output)`. There are as many logits as there are tokens in the output vocabulary. In the conceptual diagram we pretend that each decoder output is represented with 4 dimensional vectors and the last linear layer maps this vector to the output vocabulary which is only 5 tokens including the end of sentence `<eos>` token. In a more realistic model, the decoder might output 512 dimensional vectors and the last linear layer maps this vector to the output vocabulary which is a few thousand tokens wide, including words, punctuation marks and the end of sentence, unknown etc tokens. The softmax later on will balance all these logits to sum to 1.0 so that you can treat this as a probability distribution over the vocabulary from which the agent will draw/sample its next word. \n",
    "\n",
    "<img src=\"../saved/images/vec2vocab.png\" height=500 width=600>\n",
    "\n",
    "Define the Transformer class, instantiate a model and load the weights you trained in START_HERE into that model by running the next 2 cells. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "    def __init__(self, in_vocab_size, out_vocab_size, emb_dim, n_layers, heads, dropout):\n",
    "        super().__init__()\n",
    "        self.encoder = Encoder(in_vocab_size, emb_dim, n_layers, heads, dropout)\n",
    "        self.decoder = Decoder(out_vocab_size, emb_dim, n_layers, heads, dropout)\n",
    "        self.out = nn.Linear(emb_dim, out_vocab_size)\n",
    "    def forward(self, src_seq, trg_seq, src_mask, trg_mask):\n",
    "        e_output = self.encoder(src_seq, src_mask)\n",
    "        d_output = self.decoder(trg_seq, e_output, src_mask, trg_mask)\n",
    "        output = self.out(d_output)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "if teaching:\n",
    "    emb_dim, n_layers, heads, dropout = 32, 2, 8, 0.1 \n",
    "    chloe = Transformer(len(infield.vocab), len(outfield.vocab), \n",
    "                        emb_dim, n_layers, heads, dropout)\n",
    "    chloe.load_state_dict(torch.load(opt.save_path))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The way we use the transformer is that we first use the encoder to modify the input sequence into a new sequence of vectors where each vector uses the other vectors in the sentence as context to re-represented itself. \n",
    "\n",
    "` encoding = model.encoder(input_sequence, input_mask)`\n",
    "\n",
    "Then we run a loop. The loop will keep running until it gets to `max_len` or until it believes it is finished with it's sentence. \n",
    "\n",
    "To get started talking, the decoder will need two sequences. The sentence it is responding to `encoding` and a start token.\n",
    "\n",
    "`init_tok = outfield.vocab.stoi['<sos>']` \n",
    "\n",
    "`decoder_input = torch.LongTensor([[init_tok]])` \n",
    "\n",
    "Then the loops starts. Each cycle the decoder takes the decoder input, which are the words spoken thus far, and outputs the next word for each of those decoder input tokens. For example, if you input `<sos>`, you get `hi`.\n",
    "\n",
    "`<sos>` -> `hi`\n",
    "\n",
    "`<sos> hi` -> `hi I`\n",
    "\n",
    "`<sos> hi I` -> `hi I am`\n",
    "\n",
    "`<sos> hi I am` -> `hi I am chloe`\n",
    "\n",
    "`<sos> hi I am chloe` -> `hi I am chloe <eos>`\n",
    "\n",
    "Lets take a look at the second step to see how we got `<sos> hi` -> `hi I`. The decoder input of 2 vectors goes in and 2 vectors come out. Those two vectors are a shift to the right, or a shift into the future by one step, thus `<sos>` again produces `hi` and `hi` produces `I`. \n",
    "\n",
    "Actually, when `<sos>`, the vector, is inputted to the decoder, the decoder takes this vector and using `encoding` as the context, produces a vector that is the size of the vocabulary. The code `self.out = nn.Linear(emb_dim, out_vocab_size)` shows you that `output = self.out(d_output)` is tranforming a vector of size `emb_dim` and expanding that out to size `out_vocab_size`. \n",
    "\n",
    "`softout = F.softmax(out, dim=-1)` takes the model output (the logits) and balances/smoothens the these numbers such that they add up to 1.0 like a real distribution. The distribution `distr = Categorical(probs=softout)` is over every token in the output vocabulary.  In this line:\n",
    "\n",
    "`action = distr.sample()[:,-1].unsqueeze(0)`\n",
    "\n",
    "We pick a word from the output vocabulary while giving words with more larger more positive logits a higher chance of being picked, aka sampling from the distribution. If you have not taken probability theory, this subject is considered a pre-requisite for machine learning, but I dont think it should stop you from understanding this lesson. Imagine we are having a raffle. The logits are like how many raffle tickets you have. If you are a word token with twice as many raffle tickets as me, you have about double the chance of being picked than I do, but I might still be picked because its random. \n",
    "\n",
    "`[:,-1]` is just saying we are only going to take the new output, the output of the leading token, For `<sos> hi` that would be the word `hi` whose output is `I`. `.unsqueeze(0)` adds the batch dimension. `decoder_input = torch.cat((decoder_input, action), dim=1)` appends the next token to the sequence we are building so that the next input to the decoder is `<sos> hi I`\n",
    "\n",
    "`de_str = ' '.join([outfield.vocab.itos[tok] for tok in decoder_input[0][1:-1]])` joins the tokens in their string form together, separated by spaces `' '`, to give us our string readable output response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def talk_to_chloe(input_str, model, opt, infield, outfield):\n",
    "    '''\n",
    "    input:\n",
    "        input_str is a string, it is what you want to say to the dialogue model\n",
    "        model is a Transformer model with encoder, decoder and a last layer linear transformation\n",
    "        opt is an options object with the maximum length of the output sequence opt.max_len\n",
    "        infield and outfield are the data.fields that store the vocabulary\n",
    "    output:\n",
    "        an output string response from the dialogue model\n",
    "    \n",
    "    Note: this version assumes we are evaluating the model on CPU \n",
    "    '''\n",
    "    model.eval()\n",
    "    model.cpu()\n",
    "    input_sequence = string2tensor(input_str, infield) # string to tensor \n",
    "    input_mask = (input_sequence != infield.vocab.stoi['<pad>']).unsqueeze(-2) #make input mask\n",
    "    encoding = model.encoder(input_sequence, input_mask) # use the encoder rerepresent the input\n",
    "    init_tok = outfield.vocab.stoi['<sos>'] # this is the integer for the start token\n",
    "    decoder_input = torch.LongTensor([[init_tok]]) # use start token to initiate the decoder\n",
    "    \n",
    "    # continue obtaining the next decoder token until decoder outputs and end token or til max_len \n",
    "    for pos in range(opt.max_len):\n",
    "        decoder_input_mask = nopeak_mask(size=pos+1, opt=opt) # make target mask, pos+1 casue pos starts at 0\n",
    "        # the out vector contains the logits that are rebalanced by the softmax\n",
    "        out = model.out(model.decoder(decoder_input, decoder_input_mask, encoding, input_mask))\n",
    "        softout = F.softmax(out, dim=-1) \n",
    "        #softout is a categorical probability distribution over the output vocab\n",
    "        distr = Categorical(probs=softout)\n",
    "        action = distr.sample()[:,-1].unsqueeze(0) # sample from that distribution to get next token\n",
    "        # concatenate that token to our running list of output tokens \n",
    "        decoder_input = torch.cat((decoder_input, action), dim=1) \n",
    "        # if the model outputs an end of sentence token, it is done with this sentence\n",
    "        if outfield.vocab.itos[action] == '<eos>':\n",
    "            # [0] because we are assuming batch size of 1 \n",
    "            # [1:-1] excludes the start and end token from the output string \n",
    "            de_str = ' '.join([outfield.vocab.itos[tok] for tok in decoder_input[0][1:-1]])\n",
    "            return de_str\n",
    "        \n",
    "    de_str = ' '.join([outfield.vocab.itos[tok] for tok in decoder_input[0]])\n",
    "    return de_str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "meowci beaucoup !\n"
     ]
    }
   ],
   "source": [
    "if teaching:\n",
    "    print(talk_to_chloe(\"how?\", chloe, opt, infield, outfield))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Yay! you have completed building a thing that can talk to you. Go to MoveData.ipynb for the next lesson\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def talk_to_model(input_str, model, opt, infield, outfield):\n",
    "    '''\n",
    "    input:\n",
    "        input_str is a string, it is what you want to say to the dialogue model\n",
    "        model is a Transformer model with encoder, decoder and a last layer linear transformation\n",
    "        opt is an options object with the maximum length of the output sequence opt.max_len\n",
    "        infield and outfield are the data.fields that store the vocabulary\n",
    "    output:\n",
    "        an output string response from the dialogue model\n",
    "    \n",
    "    Note: this version assumes we are evaluating the model on CPU \n",
    "    '''\n",
    "    model.eval()\n",
    "    model.cpu()\n",
    "    input_sequence = string2tensor(input_str, infield) # string to tensor \n",
    "    input_mask = (input_sequence != infield.vocab.stoi['<pad>']).unsqueeze(-2) #make input mask\n",
    "    #encoding = model.encoder(input_sequence, input_mask, model.memory, model.mem_mask) # use the encoder rerepresent the input\n",
    "    encoding = model.encoder(input_sequence, input_mask)\n",
    "    init_tok = outfield.vocab.stoi['<sos>'] # this is the integer for the start token\n",
    "    decoder_input = torch.LongTensor([[init_tok]]) # use start token to initiate the decoder\n",
    "    \n",
    "    # continue obtaining the next decoder token until decoder outputs and end token or til max_len \n",
    "    for pos in range(opt.max_len):\n",
    "        decoder_input_mask = nopeak_mask(size=pos+1, opt=opt) # make target mask, pos+1 casue pos starts at 0\n",
    "        # the out vector contains the logits that are rebalanced by the softmax\n",
    "        out = model.out(model.decoder(decoder_input, decoder_input_mask, encoding, input_mask))\n",
    "        softout = F.softmax(out, dim=-1) \n",
    "        #softout is a categorical probability distribution over the output vocab\n",
    "        distr = Categorical(probs=softout)\n",
    "        action = distr.sample()[:,-1].unsqueeze(0) # sample from that distribution to get next token\n",
    "        # concatenate that token to our running list of output tokens \n",
    "        decoder_input = torch.cat((decoder_input, action), dim=1) \n",
    "        # if the model outputs an end of sentence token, it is done with this sentence\n",
    "        if outfield.vocab.itos[action] == '<eos>':\n",
    "            # [0] because we are assuming batch size of 1 \n",
    "            # [1:-1] excludes the start and end token from the output string \n",
    "            de_str = ' '.join([outfield.vocab.itos[tok] for tok in decoder_input[0][1:-1]])\n",
    "            return de_str\n",
    "        \n",
    "    de_str = ' '.join([outfield.vocab.itos[tok] for tok in decoder_input[0]])\n",
    "    return de_str"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}