2/task_2.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "task-2-artificial-self-AMLD2020.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yX2V-RTlOCMp",
        "colab_type": "text"
      },
      "source": [
        "# Meet your Artificial Self\n",
        "\n",
        "> This notebook will be used for the [AMLD 2020](https://appliedmldays.org/) workshop [**\"Meet your Artificial Self\"**](https://appliedmldays.org/workshops/meet-your-artificial-self-generate-text-that-sounds-like-you), taking place January 25 in Lausanne, Switzerland.\n",
        "\n",
        "## Task 2\n",
        "In task 1 we learnt how to fine-tune a language model and we saw how style transfer works. Conversations are a different beast however! \n",
        "\n",
        "In this task we will try our first approach at training a conversational model.\n",
        "\n",
        "## Important resources\n",
        "* [Workshop Github repo](https://github.com/mar-muel/artificial-self-AMLD-2020/tree/master/2)\n",
        "* [PyTorch documentation](https://pytorch.org/docs/stable/index.html)\n",
        "* Huggingface transformers library [ [Github](https://github.com/huggingface/transformers) | [Docs](https://huggingface.co/transformers/) ]\n",
        "\n",
        "## Approach\n",
        "In this task we will try a naive approach to getting conversational style by simply feeding the model \"raw\" conversation data of the form:\n",
        "```\n",
        "<speaker1> Hi\n",
        "<speaker2> Hey - how are you?\n",
        "<speaker1> Great, thanks!\n",
        "...\n",
        "```\n",
        "Our hope is that the model will simply learn this structure and we will be able to query the model with an input of the form:\n",
        "\n",
        "```\n",
        "<speaker2> Am I speaking to a bot?\n",
        "<speaker1>\n",
        "```\n",
        "We then expect the model to extend the text from this prefix.\n",
        "\n",
        "This notebook will run you through all the steps from collecting the training data until interacting with the final model.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PUyajaGkR1JM",
        "colab_type": "text"
      },
      "source": [
        "# Setting things up\n",
        "The following cells will clone the repository and install all the necessary dependencies\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cC6KTnYXD3x4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!nvidia-smi | grep -q 'failed' && echo \"STOP! You are using a runtime without a GPU. Change the runtime type before going further!\""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H5PGOly6SB12",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!git clone https://github.com/mar-muel/artificial-self-AMLD-2020.git"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B8LoBjs4SOhu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Set working directory\n",
        "%cd /content/artificial-self-AMLD-2020/2"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E2e-ZKZISylw",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Install all dependencies for this task\n",
        "!pip install -r requirements-colab.txt"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0pM82B4aUXg8",
        "colab_type": "text"
      },
      "source": [
        "# Import training data\n",
        "In this step we will add the data to Colab to train the model. You are free to choose from two options:\n",
        "1. Use your own chat logs (\"Chat logs\")\n",
        "2. Use conversational data from well known people (\"Alternative data\")\n",
        "\n",
        "# Your own chat logs\n",
        "In order to obtain your chat logs you need to first request them and then parse them into a format we can use. For this we will use the [Chatistics](https://github.com/MasterScrat/Chatistics) repository.\n",
        "\n",
        "Do the following steps on your local machine:\n",
        "1. **Download** the raw data as explained in the [Chatistics](https://github.com/MasterScrat/Chatistics) repository\n",
        "2. **Parse** the chat logs (as described in the repo)\n",
        "3. **Export** the data you wish to use for this workshop. For an overview of the filter options run `python export.py --help`. Once decided run the following command to export the data to JSON\n",
        "```\n",
        "python export.py -f json --compress <your filter options>\n",
        "```\n",
        "4. **Mount** your Google Drive in this notebook so you can store the exported data there"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nQQdr1g_VvfK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Mount Google Drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VmK3IG0wPKCF",
        "colab_type": "text"
      },
      "source": [
        "5. **Upload** this JSON file now into the folder `/content/drive/My Drive/AMLD/chatistics_data/chatistics_export_YYYY-MM-DD_HH-MM-SS.json`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IAE4MrVSVmV1",
        "colab_type": "text"
      },
      "source": [
        "# Alternative dataset 1: world leader interviews\n",
        "\n",
        "Instead of using your own data, you can use some datasets we prepared for you. The first option is a dataset of interviews of world leaders: Barack Obama and Vladimir Putin. These interviews will be treated as chat conversations, where the interlocutors are the reporters.\n",
        "\n",
        "To use those, copy the conversation you want from the `datasets` folder to the task2 `data` folder:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_9PmICPoWt98",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Barack Obama\n",
        "!cp ../datasets/barack_obama_interviews.json data\n",
        "\n",
        "# Vladimir Putin\n",
        "#!cp ../datasets/vladimir_putin_interviews.json data"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lQCXoLLhO2FG",
        "colab_type": "text"
      },
      "source": [
        "# Alternative dataset 2: movie quotes\n",
        "\n",
        "Another option is to use quotes from movies. You can use the [Cornell Movie-Dialogs Corpus](https://), which contains **220,579 conversational exchanges** between 10,292 pairs of movie characters."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pRYuNd85O5cl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!cp ../datasets/cornell_movie_dialogs_corpus.json.zip data"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nxNG4fmqcHGt",
        "colab_type": "text"
      },
      "source": [
        "# Prepare the data\n",
        "For this task we will use the transfomers library by Huggingface. The transformers library implements many recent NLP models (such as BERT and GPT-2).\n",
        "\n",
        "Our data is currently in JSON format, but can easily be read into a Pandas Dataframe.\n",
        "\n",
        "As a first step we want to get our conversation data from this format:\n",
        "\n",
        "| timestamp | conversationId | conversationWithName | senderName | outgoing | text | language | platform |\n",
        "| :----: | :----: | :----: | :----: | :----: | :----: | :----: | :----: |\n",
        "| 1575463019 | 693342290 | Alice | Bob | True | Hi Alice! | en | whatsapp |\n",
        "| 1575463025 | 693342290 | Alice | Alice | False | Hi Bob! How are you these days? | en | whatsapp |\n",
        "| 1575463030 | 693342290 | Alice | Bob | True | Great! Thanks | en | whatsapp |\n",
        "\n",
        "and get into a text file of this format:\n",
        "```\n",
        "<person1> Hi Alice!\n",
        "<person2> Hi Bob! How are you these days?\n",
        "<person1> Great! Thanks\n",
        "...\n",
        "```\n",
        "\n",
        "Our model will try to generate this structure from the data. The tags `<person1>` and `<person2>` have nothing special to them and could in theory be replaced by something else as well (more about this below)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MLHMIoUbTFp-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from utils import generate_input_task2\n",
        "import os"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Bvdcn7HBTGgG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Set the correct path here\n",
        "input_data = \"/content/drive/My Drive/AMLD/chatistics_data/chatistics_export_2020-01-19_02-18-17.json\" #@param {type:\"string\"}\n",
        "\n",
        "assert os.path.isfile(input_data)\n",
        "generate_input_task2(input_data)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "00KS5e8KYpWz",
        "colab_type": "text"
      },
      "source": [
        "The script has now generated an input file `cached_input_task2.txt`. You can inspect it with the `head` command:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kUh1_4LBpA05",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!head cached_input_task2.txt"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f9gLQPF6Y-K4",
        "colab_type": "text"
      },
      "source": [
        "# Train the model\n",
        "We can now start training the model. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2GsvtlOlZSAF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from collections import defaultdict\n",
        "from itertools import chain\n",
        "import random\n",
        "import numpy as np\n",
        "import torch\n",
        "from torch.utils.data import Dataset, DataLoader, TensorDataset\n",
        "from tqdm import tqdm, trange\n",
        "from transformers import (\n",
        "    WEIGHTS_NAME,\n",
        "    AdamW,\n",
        "    GPT2Config,\n",
        "    GPT2LMHeadModel,\n",
        "    GPT2Tokenizer,\n",
        "    OpenAIGPTConfig,\n",
        "    OpenAIGPTLMHeadModel,\n",
        "    OpenAIGPTTokenizer,\n",
        "    PreTrainedTokenizer,\n",
        "    get_linear_schedule_with_warmup,\n",
        ")\n",
        "from utils import get_input_task2, set_seed, add_special_tokens_\n",
        "import logging\n",
        "import torch.nn.functional as F\n",
        "\n",
        "# set up logging\n",
        "logger = logging.getLogger(__name__)\n",
        "logging.basicConfig(level=logging.DEBUG, format='%(asctime)s [%(levelname)-5.5s] [%(name)-12.12s]: %(message)s')\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YYeG8QI0ce_W",
        "colab_type": "text"
      },
      "source": [
        "### Constants"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ckAH0GK5aH0t",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "run_name = 'run1'               # The name of the run (subdirectory in ./runs)\n",
        "model_type = 'openai-gpt'       # Initialize model from path to checkpoint or with model name (\"openai-gpt\" or \"gpt2\")\n",
        "save_every = 50                 # Save checkpoint every n updates steps.\n",
        "max_input_length = 400          # Number of tokens which will be fed into the model (reduce this number if you have memory constraints)\n",
        "weight_decay = 0                # Weight decay if we apply some.\n",
        "train_batch_size = 4            # Batch size for training\n",
        "gradient_accumulation_steps = 8 # Accumulate gradients on several steps\n",
        "lr = 5e-5                       # Learning rate\n",
        "adam_epsilon = 1e-8             # Epsilon for Adam optimizer.\n",
        "max_norm = 1                    # Clipping gradient norm\n",
        "n_epochs = 2                    # Number of training epochs\n",
        "device = 'cuda'                 # Device (cuda or cpu)\n",
        "warmup_steps = 0                # Linear warmup over warmup_steps.\n",
        "seed = 42                       # random seed for initialization"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HFnOq8n4cib3",
        "colab_type": "text"
      },
      "source": [
        "### Data loading\n",
        "In PyTorch we can define a class which inherits from the Dataset class containing an initializer method `__init__()` in which we \n",
        "\n",
        "1. Read the text file which we generated before\n",
        "2. Tokenize the text (split it the text into smaller words/character pairs) and convert the tokens into vocabulary IDs (the positions of the tokens in the vocabulary. GPT-2 uses so-called BPE (byte-pair encoding). If you're interested how it works you can read more about it in [this blog post](https://leimao.github.io/blog/Byte-Pair-Encoding/).\n",
        "3. Cut up the array of token IDs into chunks of size `max_input_length` (usually chosen to the maximum of what the memory/model allows for)\n",
        "4. Append the generated training example as a list\n",
        "\n",
        "The `__get_item__()` simply implements the retrieval of a new training example."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xeBZosNKZYUa",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class TextDataset(Dataset):\n",
        "    def __init__(self, tokenizer):\n",
        "        # load the text data generated from before into memory\n",
        "        text = get_input_task2(input_data)\n",
        "        logger.info(\"Tokenizing and building input...\")\n",
        "        # tokenize the whole file\n",
        "        tokenized_text = tokenizer.convert_tokens_to_ids(tokenizer.tokenize(text))\n",
        "        # generate training examples by cutting the text into blocks of size max_input_length\n",
        "        self.examples = []\n",
        "        block_size = max_input_length\n",
        "        if block_size < 0:\n",
        "            # use maximum possible input block size\n",
        "            block_size = tokenizer.max_len_single_sentence\n",
        "        for i in range(0, len(tokenized_text) - block_size + 1, block_size):  # Truncate in block of block_size\n",
        "            # Note that we are loosing the last truncated example here for the sake of simplicity (no padding)\n",
        "            self.examples.append(tokenizer.build_inputs_with_special_tokens(tokenized_text[i : i + block_size]))\n",
        "    def __len__(self):\n",
        "        return len(self.examples)\n",
        "    def __getitem__(self, item):\n",
        "        return torch.tensor(self.examples[item])\n",
        "\n",
        "def get_data_loader(tokenizer):\n",
        "    \"\"\" Prepare the dataset for training and evaluation \"\"\"\n",
        "    dataset = TextDataset(tokenizer)\n",
        "    logger.info(\"Train dataset: {:,} samples\".format(len(dataset)))\n",
        "    logger.info(\"Build dataloaders\")\n",
        "    data_loader = DataLoader(dataset, batch_size=train_batch_size, shuffle=True)\n",
        "    return data_loader\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z7voPNvpgjL2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Setting the same seed allows for some reproducibility of the experiments\n",
        "set_seed(seed)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fXbo_iLqiEjf",
        "colab_type": "text"
      },
      "source": [
        "### Load models and tokenizers\n",
        "The transformers library comes with a built in method `from_pretrained(model_type)`. `model_type` can specify either\n",
        "1. One of the [pretrained model architectures](https://huggingface.co/transformers/pretrained_models.html). In this case the model will be downloaded and cached on disk before loaded into memory.\n",
        "2. The path to a folder with an existing model checkpoint.\n",
        "\n",
        "This allows us to use the same syntax to either pretrained or fine-tuned models/tokenizers.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QddFgKLlZlHI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        " # Load tokenizer\n",
        " logger.info(\"Prepare tokenizer, pretrained model and optimizer.\")\n",
        " tokenizer_class = GPT2Tokenizer if \"gpt2\" in model_type else OpenAIGPTTokenizer\n",
        " tokenizer = tokenizer_class.from_pretrained(model_type)\n",
        " # Load model\n",
        " model_class = GPT2LMHeadModel if \"gpt2\" in model_type else OpenAIGPTLMHeadModel\n",
        " model = model_class.from_pretrained(model_type)\n",
        " model.to(device)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xsD0MYlajkFE",
        "colab_type": "text"
      },
      "source": [
        "### Add special tokens\n",
        "Let's see how `<speaker1>` and `<speaker2>` tokens will be tokenized by our current tokenizer:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9s0G97YYjz1D",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "tokenizer.tokenize('<speaker1>')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L7YQdj8ikKKA",
        "colab_type": "text"
      },
      "source": [
        "As you can see the model generates a total of 4 tokens: `['<', 'speaker', '1', '>']`. This doesn't make too much sense for us since the tags should not contain any meaning but should simply indicate who is currently speaking. Luckily there is an easy way to add our speaker tokens to the vocabulary of the tokenizer."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "o11KdvWNjWik",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "ATTR_TO_SPECIAL_TOKEN = {'additional_special_tokens': ('<speaker1>', '<speaker2>')}\n",
        "# Add special tokens if they are not already added\n",
        "def add_special_tokens_(model, tokenizer):\n",
        "   \"\"\" Add special tokens to the tokenizer and the model if they have not already been added. \"\"\"                                                                                   \n",
        "   orig_num_tokens = len(tokenizer.encoder)\n",
        "   num_added_tokens = tokenizer.add_special_tokens(ATTR_TO_SPECIAL_TOKEN) # doesn't add if they are already there                                                                   \n",
        "   if num_added_tokens > 0:\n",
        "       model.resize_token_embeddings(new_num_tokens=orig_num_tokens + num_added_tokens)\n",
        "       \n",
        "add_special_tokens_(model, tokenizer)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oJOfCIAblBH1",
        "colab_type": "text"
      },
      "source": [
        "Now the model should generate a single token:\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zk0ehp_KlDYn",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "tokenizer.tokenize('<speaker1>')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r7hnYlQklwmE",
        "colab_type": "text"
      },
      "source": [
        "# Final setup before training\n",
        "We need to set up a few things before we can start training:\n",
        "* Prepare the data loaders (discussed above)\n",
        "* An optimizer (we will use Adam)\n",
        "* A scheduler to change the learning rate throughout training (we will use a [linear schedule with warmup](https://huggingface.co/transformers/main_classes/optimizer_schedules.html#transformers.get_linear_schedule_with_warmup))"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kdj6wabyjQ0V",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        " # Get data loaders\n",
        " logger.info(\"Prepare datasets\")\n",
        " data_loader = get_data_loader(tokenizer)\n",
        " # Prepare optimizer and schedule (linear warmup and decay)\n",
        " no_decay = [\"bias\", \"LayerNorm.weight\"]\n",
        " optimizer_grouped_parameters = [\n",
        "     {\n",
        "         \"params\": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],\n",
        "         \"weight_decay\": weight_decay,\n",
        "     },\n",
        "     {\"params\": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], \"weight_decay\": 0.0},\n",
        " ]\n",
        " optimizer = AdamW(optimizer_grouped_parameters, lr=lr, eps=adam_epsilon)\n",
        " t_total = len(data_loader) // gradient_accumulation_steps * n_epochs\n",
        " scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps, num_training_steps=t_total)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5TMwH4HnnLGw",
        "colab_type": "text"
      },
      "source": [
        "### Training\n",
        "\n",
        "Finally we can start training!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dOu849kknPnz",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "logger.info(\"***** Running training *****\")\n",
        "global_step = 0\n",
        "epochs_trained = 0\n",
        "steps_trained_in_current_epoch = 0\n",
        "# Check if we are training from a checkpoint or from a pretrained model\n",
        "if os.path.exists(model_type):\n",
        "    # set global_step to gobal_step of last saved checkpoint from model path\n",
        "    global_step = int(model_type.split(\"-\")[-1].split(\"/\")[0])\n",
        "    epochs_trained = global_step // (len(data_loader) // gradient_accumulation_steps)\n",
        "    steps_trained_in_current_epoch = global_step % (len(data_loader) // gradient_accumulation_steps)\n",
        "    logger.info(\"Continuing training from checkpoint, will skip to saved global_step\")\n",
        "    logger.info(f\"Continuing training from epoch {epochs_trained}\")\n",
        "    logger.info(f\"Continuing training from global step {global_step}\")\n",
        "    logger.info(f\"Will skip the first {steps_trained_in_current_epoch} steps in the first epoch\")\n",
        "\n",
        "# Training loop\n",
        "model.zero_grad()\n",
        "epoch_pbar = trange(epochs_trained, int(n_epochs)) # epoch progress bar\n",
        "av_loss = 0\n",
        "for current_epoch in epoch_pbar:\n",
        "    epoch_pbar.set_description(f\"Epoch [{current_epoch+1}/{n_epochs}]\") # description of epoch progress bar\n",
        "    pbar = tqdm(data_loader, position=0) # progress bar\n",
        "    for step, batch in enumerate(pbar):\n",
        "        # Skip past any already trained steps if resuming training\n",
        "        if steps_trained_in_current_epoch > 0:\n",
        "            steps_trained_in_current_epoch -= 1\n",
        "            continue\n",
        "        model.train()\n",
        "        # the language model targets (labels) are the same as the input!\n",
        "        inputs, labels = (batch, batch)\n",
        "        inputs = inputs.to(device)\n",
        "        labels = labels.to(device)\n",
        "        loss, *_ = model(inputs, labels=labels)\n",
        "        loss.backward()\n",
        "        tr_loss = loss.item()\n",
        "        # Compute a running average of the loss\n",
        "        av_loss = (step*av_loss + tr_loss)/(step + 1)\n",
        "        pbar.set_description(f\"Average loss: {av_loss:.4f}\")\n",
        "        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)\n",
        "        if (step + 1) % gradient_accumulation_steps == 0:\n",
        "            optimizer.step()\n",
        "            scheduler.step()  # Update learning rate schedule\n",
        "            model.zero_grad()\n",
        "            global_step += 1\n",
        "            if global_step % save_every == 0 and global_step > 0:\n",
        "                checkpoint_prefix = \"checkpoint\"\n",
        "                output_dir = os.path.join('runs', run_name, \"{}-{}\".format(checkpoint_prefix, global_step))\n",
        "                if not os.path.exists(output_dir):\n",
        "                    os.makedirs(output_dir)\n",
        "                logger.info(f\"Saving model checkpoint to {output_dir}\")\n",
        "                model.save_pretrained(output_dir)\n",
        "                tokenizer.save_pretrained(output_dir)\n",
        "                logger.info(f\"Saving optimizer and scheduler states to {output_dir}\")\n",
        "                torch.save(optimizer.state_dict(), os.path.join(output_dir, \"optimizer.pt\"))\n",
        "                torch.save(scheduler.state_dict(), os.path.join(output_dir, \"scheduler.pt\"))\n",
        "\n",
        "# save model\n",
        "output_dir = os.path.join('runs', run_name)\n",
        "if not os.path.exists(output_dir):\n",
        "    os.makedirs(output_dir)\n",
        "logger.info(f\"Saving model checkpoint to {output_dir}\")\n",
        "model.save_pretrained(output_dir)\n",
        "tokenizer.save_pretrained(output_dir)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8Ew8SGqBB-Te",
        "colab_type": "text"
      },
      "source": [
        "# Tweaking parameters\n",
        "\n",
        "You can change the training parameters to see how they affect the language model: adjust them below, then run the again the cell above."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jRyEvCxPB_Np",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "run_name = 'run1'               # The name of the run (subdirectory in ./runs)\")\n",
        "model_type = 'openai-gpt'       # Initialize model from path to checkpoint or with model name (openai-gpt/openai-gpt2)\"\n",
        "weight_decay = 0                # Weight decay if we apply some.\n",
        "train_batch_size = 4            # Batch size for training\n",
        "gradient_accumulation_steps = 8 # Accumulate gradients on several steps\n",
        "lr = 5e-5                       # Learning rate\n",
        "n_epochs = 1                    # Number of training epochs\n",
        "warmup_steps = 0                # Linear warmup over warmup_steps."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_v4Y2p3RpOBB",
        "colab_type": "text"
      },
      "source": [
        "# Interact with the model\n",
        "The trained model can now be found under `./runs/{run_name}/`.\n",
        "\n",
        "Let's see what happens when we feed in some text to our model:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CiT2mrSAqvWS",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "input_ids = torch.tensor(tokenizer.encode(\"Hello world\", add_special_tokens=True), device=device).unsqueeze(0)\n",
        "print('Input IDs:')\n",
        "print(input_ids)\n",
        "out, = model(input_ids)\n",
        "print('Model output:')\n",
        "print(out)\n",
        "print('Output shape:')\n",
        "print(out.shape) # output shape: (batch size x sequence length x hidden size)\n",
        " "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CuP7ua6Jt_WO",
        "colab_type": "text"
      },
      "source": [
        "As you can see the output of the model outputs a tensor of the size of the number of input tokens. By getting the highest value of the last dimension of the output tensor we can get the most likely next token:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xO7VgD0sqxzu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "tokenizer.decode([torch.argmax(out[:, 1, :])])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IuHx8BTrvVsW",
        "colab_type": "text"
      },
      "source": [
        "## Start chatting\n",
        "As we have seen in task 1 we have several hyperparameters to choose from in order to control how we sample from the output probability distribution. Just choosing always the most likely token will not lead to interesting conversations.\n",
        "\n",
        "Below you see how the different sampling strategies are implemented. By default we will use top-p sampling."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w5UrWZw6u1M7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Constants\n",
        " max_history = 2                  # Number of previous utterances to keep in history\n",
        " no_sample = False                # Set to use greedy decoding instead of sampling\n",
        " max_length = 80                  # Maximum length of the output utterances\n",
        " temperature = 1.0                # Sampling softmax temperature\n",
        " top_k = 0                        # Filter top-k tokens before sampling (<=0: no filtering)\n",
        " top_p = 0.8                      # Nucleus filtering (top-p) before sampling (<=0.0: no filtering)\n",
        " no_info = False                   # Only show conversation output"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QaL732AIxCYl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def top_k_top_p_filtering(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')):\n",
        "    \"\"\" Filter a distribution of logits using top-k and/or nucleus (top-p) filtering\n",
        "        Args:\n",
        "            logits: logits distribution shape (batch size x vocabulary size)\n",
        "            top_k > 0: keep only top k tokens with highest probability (top-k filtering).\n",
        "            top_p > 0.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).\n",
        "                Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)\n",
        "        From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317\n",
        "    \"\"\"\n",
        "    top_k = min(top_k, logits.size(-1))  # Safety check\n",
        "    if top_k > 0:\n",
        "        # Remove all tokens with a probability less than the last token of the top-k\n",
        "        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]\n",
        "        logits[indices_to_remove] = filter_value\n",
        "\n",
        "    if top_p > 0.0:\n",
        "        sorted_logits, sorted_indices = torch.sort(logits, descending=True)\n",
        "        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)\n",
        "\n",
        "        # Remove tokens with cumulative probability above the threshold\n",
        "        sorted_indices_to_remove = cumulative_probs > top_p\n",
        "        # Shift the indices to the right to keep also the first token above the threshold\n",
        "        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()\n",
        "        sorted_indices_to_remove[..., 0] = 0\n",
        "\n",
        "        # scatter sorted tensors to original indexing\n",
        "        indices_to_remove = sorted_indices_to_remove.scatter(dim=1, index=sorted_indices, src=sorted_indices_to_remove)\n",
        "        logits[indices_to_remove] = filter_value\n",
        "    return logits\n",
        "\n",
        "def sample_sequence(conversation, model, num_samples=1):\n",
        "    \"\"\"Generate next tokens from pervious conversation\"\"\"\n",
        "    context = torch.tensor(conversation, dtype=torch.long, device=device)\n",
        "    context = context.unsqueeze(0).repeat(num_samples, 1)\n",
        "    generated = context\n",
        "    with torch.no_grad():\n",
        "        for _ in range(max_length):\n",
        "            inputs = {'input_ids': generated}\n",
        "            outputs = model(**inputs)\n",
        "            # scale by temperature\n",
        "            next_token_logits = outputs[0][:, -1, :] / (temperature if temperature > 0 else 1.) \n",
        "            # filter by top-k/top-p\n",
        "            filtered_logits = top_k_top_p_filtering(next_token_logits, top_k=top_k, top_p=top_p)\n",
        "            if temperature == 0: # greedy sampling:\n",
        "                next_token = torch.argmax(filtered_logits, dim=-1).unsqueeze(-1)\n",
        "            else:\n",
        "                next_token = torch.multinomial(F.softmax(filtered_logits, dim=-1), num_samples=1)\n",
        "            generated = torch.cat((generated, next_token), dim=1)\n",
        "    return generated"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cBpALwQAyWb9",
        "colab_type": "text"
      },
      "source": [
        "We are ready to interact with the model. A few things to note:\n",
        "* We will give it a \"trigger\" to start the conversation. From there the model takes it. Note that the model is \"playing\" speaker1 and you are speaker2. \n",
        "* If the `no_info` flag is set to `False`, the output shows both the input (conversation history) as well as the full output of the model. At the very end the answer which was selected by the model is shown.\n",
        "* You can press `h` (and then ENTER) in order to see the whole history of the chat\n",
        "\n",
        "\n",
        "Enjoy! :)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c4gwwMEoxwAd",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "history = []\n",
        "speaker1_tag = '<speaker1>'\n",
        "speaker2_tag = '<speaker2>'\n",
        "speaker1_tag_id = tokenizer.convert_tokens_to_ids(speaker1_tag)\n",
        "speaker2_tag_id = tokenizer.convert_tokens_to_ids(speaker2_tag)\n",
        "history = f\"\"\"\n",
        "{speaker2_tag} Hi!\n",
        "{speaker1_tag} Hello\n",
        "{speaker2_tag} Are you ready?\n",
        "{speaker1_tag} Yes!\n",
        "{speaker2_tag} Ok let's start chatting\n",
        "{speaker1_tag} Sure, what do you want to talk about?\"\"\"\n",
        "print(history)\n",
        "print('\\n[Chat with the model! Send \"h\" to see the full history]\\n')\n",
        "history = history.split('\\n')\n",
        "while True: \n",
        "    message = None\n",
        "    while not message:\n",
        "        message = input(f'{speaker2_tag} ')\n",
        "        if message == 'h':\n",
        "            print('\\n'.join(history))\n",
        "            message = None\n",
        "    # add new message to history\n",
        "    history.append(f'{speaker2_tag} {message}')\n",
        "    # keep only most recent conversation as input to the model\n",
        "    recent_history = history[-(2*max_history):]\n",
        "    # concatenate history into single string and add trigger word \"bot:\"\n",
        "    history_str = '{}\\n{}'.format('\\n'.join(recent_history), speaker1_tag)\n",
        "    # tokenize text and convert into vocabulary ids (input ids)\n",
        "    history_enc = tokenizer.encode(history_str, add_special_tokens=True)\n",
        "    with torch.no_grad():\n",
        "        out_ids = sample_sequence(history_enc, model)\n",
        "    out_ids = out_ids[:, len(history_enc):].tolist()[0]\n",
        "    if not no_info:\n",
        "        print(20*'-')\n",
        "        print('Output of model:')\n",
        "        full_output = tokenizer.decode(out_ids, clean_up_tokenization_spaces=True)\n",
        "        print(full_output)\n",
        "        print('\\nInput to the model:')\n",
        "        print(history_str)\n",
        "        print(20*'-' + '\\n')\n",
        "    # Select part before speaker tags as answer\n",
        "    for i, out_id in enumerate(out_ids):\n",
        "        if out_id in [speaker1_tag_id, speaker2_tag_id]:\n",
        "            break\n",
        "    answer = '{} {}'.format(speaker1_tag, tokenizer.decode(out_ids[:i]))\n",
        "    print(answer)\n",
        "    # add answer to history\n",
        "    history.append(answer)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FkyH26kEB2Su",
        "colab_type": "text"
      },
      "source": [
        "You can now review this whole conversation:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rNY97GdPB1hn",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "history"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IoqHoD_hW1PW",
        "colab_type": "text"
      },
      "source": [
        "### Save model to Google Drive\n",
        "If you are happy with your model consider saving it to your Google Drive. Note that all data on this notebook will be lost after a certain time of inactivity. Note that the model size is quite big (~500MB) so make sure you have enough space in your Google Drive.\n",
        "\n",
        "This will save only your final model state (from your directory `run_name` directory).\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "g_NMpXM2W17W",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import shutil\n",
        "\n",
        "#@title Save to Google Drive\n",
        "save_to_drive = False #@param {type:\"boolean\"}\n",
        "\n",
        "source_directory = f'./runs/{run_name}'\n",
        "target_directory = f\"/content/drive/My Drive/AMLD/models/task2/{run_name}/\"\n",
        "include_checkpoints = False\n",
        "\n",
        "if save_to_drive:\n",
        "  logger.info(f'Copying from {source_directory} to {target_directory}...')\n",
        "  ignore_pattern = None\n",
        "  if not include_checkpoints:\n",
        "    ignore_pattern = shutil.ignore_patterns('checkpoint-*')\n",
        "  shutil.copytree(source_directory, target_directory, ignore=ignore_pattern)\n",
        "  logger.info('Successfully copied your model!')  "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4B6F1B87W6y8",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}