Train_the_GPT_2_model.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Train the GPT-2 model.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [
        "cmMZSSMseqCB",
        "1kXoufzIevAh"
      ],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ShaneZhong/train-gpt-2-model/blob/master/Train_the_GPT_2_model.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "67rQWoqaDxE3",
        "colab_type": "text"
      },
      "source": [
        "# Training the GPT-2 model from an input text"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pzxl1vYX-1kk",
        "colab_type": "text"
      },
      "source": [
        "Setup:\n",
        "\n",
        "1) Make sure GPU is enabled, go to edit->notebook settings->Hardware Accelerator GPU\n",
        "\n",
        "2) make a copy to your google drive, click on copy to drive in panel"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iW0abT07ZkhZ",
        "colab_type": "text"
      },
      "source": [
        "Note: colab will reset after 12 hours make sure to save your model checkpoints to google drive around 10-11 hours mark or before, then go to runtime->reset all runtimes. Now copy your train model back into colab and start training again from the previous checkpoint."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tlvy3k-geezj",
        "colab_type": "text"
      },
      "source": [
        "## Environment Setup\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "__YFNxW2s5Zm",
        "colab_type": "text"
      },
      "source": [
        "### Git clone and install dependencies"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-iieAryNbziM",
        "colab_type": "code",
        "cellView": "both",
        "colab": {}
      },
      "source": [
        "#@title Clone or pull train-gpt-2-model from Github\n",
        "\n",
        "%cd /content\n",
        "Mode = \"clone\" #@param [\"clone\", \"pull\"]\n",
        "\n",
        "if (Mode == \"clone\"):\n",
        "  !git clone https://github.com/ShaneZhong/train-gpt-2-model.git\n",
        "else:\n",
        "  %cd /content/train-gpt-2-model\n",
        "  !git pull\n",
        "\n",
        "!pip3 install -r /content/train-gpt-2-model/requirements.txt"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6eEIs3ApZUVO",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "cd /content/train-gpt-2-model"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y2sxRIarf6zc",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "ls"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WvUQhgK3PQ4L",
        "colab_type": "text"
      },
      "source": [
        "### GDrive Mount"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FNpf6R4ahYSN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Run this cell to mount your Google Drive.\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "A498TySgHYyF",
        "colab": {}
      },
      "source": [
        "# download the models\n",
        "!python3 download_model.py 117M\n",
        "!python3 download_model.py 345M"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7oJPQtdLbbeK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!export PYTHONIOENCODING=UTF-8"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0KzSbAvePgsI",
        "colab_type": "text"
      },
      "source": [
        "### Fetch previous model checkpoints in google drive (optional)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mywabcPPP8sW",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# check if you have any content in your GDrive directory\n",
        "ls /content/drive/My\\ Drive/Models/GPT-2/checkpoint/"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cA2Wk7yIPmS6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# If the above is not empty, you can copy the previously\n",
        "# saved model to your project directory.\n",
        "!cp -r /content/drive/My\\ Drive/Models/GPT-2/checkpoint/ /content/train-gpt-2-model/ "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cmMZSSMseqCB",
        "colab_type": "text"
      },
      "source": [
        "## Get the training dataset\n",
        "\n",
        "lets get our text to train on, in this case from project gutenberg, A Tale of Two Cities, by Charles Dickens"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QOCvrs-DHvxa",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "  !wget https://www.gutenberg.org/files/98/98-0.txt"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2k9vxViSQi6E",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "ls /content/train-gpt-2-model/"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1kXoufzIevAh",
        "colab_type": "text"
      },
      "source": [
        "### Alternative: Trump tweets\n",
        "The Trump tweets is already saved in the /data directory"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qA-pk5N8d83r",
        "colab_type": "code",
        "outputId": "606e173c-03ff-40ac-aaab-9a219ef03241",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "ls /content/train-gpt-2-model/data"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "trump_tweets.txt\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vNBkFcWnetp9",
        "colab_type": "text"
      },
      "source": [
        "## Training the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yPfJ5b3CQXqr",
        "colab_type": "text"
      },
      "source": [
        "Select either the 117M or 345M model to train. \n",
        "\n",
        "**IMPORTANT**: After running the cell below, it does not stop automatically. To stop training, you need to click the stop button. The saved model will be generated in the checkpoint directory (`e.g. Saving checkpoint/run1/model-289`)\n",
        "\n",
        "\n",
        "Many parameters can be tunned in this model. You can find the reference here: [link](https://github.com/ShaneZhong/train-gpt-2-model/blob/master/train.py) <br><br>\n",
        "\n",
        "\n",
        "\n",
        "Regarding the input_data, here are the two files you can choose:\n",
        "* '/content/train-gpt-2-model/data/trump_tweets.txt'\n",
        "*'/content/train-gpt-2-model/98-0.txt'"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YQLK2LdSbYSC",
        "colab_type": "code",
        "cellView": "both",
        "colab": {}
      },
      "source": [
        "#@title Train the model with the following parameters\n",
        "input_data = '/content/train-gpt-2-model/data/trump_tweets.txt' #@param\n",
        "model = \"345M\" #@param [\"117M\",\"345M\"]\n",
        "Samples_per_N_steps = 100 #@param\n",
        "Folder_Name = 'Trump-tweets' #@param\n",
        "\n",
        "!PYTHONPATH=src /content/train-gpt-2-model/train.py --dataset $input_data --model_name $model --sample_every $Samples_per_N_steps --run_name $Folder_Name"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vS1RJJDFOPnb",
        "colab_type": "text"
      },
      "source": [
        "### Save the trained model to GDrive\n",
        "By default, the trained model is saved in the `checkpoint` folder under your your GDrive root directory."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JretqG1zOXdi",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!cp -r /content/train-gpt-2-model/checkpoint/ /content/drive/My\\ Drive/Models/GPT-2/"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vdAESkrKfEsS",
        "colab_type": "text"
      },
      "source": [
        "## Apply the trained model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6D-i7vERWbNS",
        "colab_type": "text"
      },
      "source": [
        "### Fetch the trained model\n",
        "The trained model (117M or 345M) is pasted to the model directory."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VeETvWvrbKga",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#!cp -r /content/train-gpt-2-model/checkpoint/run1/* /content/train-gpt-2-model/models/117M/\n",
        "!cp -r /content/train-gpt-2-model/checkpoint/Trump-tweets/* /content/train-gpt-2-model/models/345M/"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5D1hdernJrDH",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# load the instruction\n",
        "!python3 src/interactive_conditional_samples.py -- --help"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GmnSrXqtfRbq",
        "colab_type": "text"
      },
      "source": [
        "### Conditional samples"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "utJj-iY4gHwE",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#!python3 src/interactive_conditional_samples.py --model_name='117M' --nsamples=2 --top_k=40 --temperature=0.7\n",
        "!python3 src/interactive_conditional_samples.py --model_name='345M' --nsamples=2 --top_k=40 --temperature=0.7"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K8rSqkGxg5OK",
        "colab_type": "text"
      },
      "source": [
        "### Unconditional samples"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LaQUEnRxWc3c",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!python3 src/generate_unconditional_samples.py --model_name='117M' --nsamples=2 --top_k=40 --temperature=0.7\n",
        "!python3 src/generate_unconditional_samples.py --model_name='345M' --nsamples=2 --top_k=40 --temperature=0.7"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}