Fine-Tuning GPT-2 on a Custom Dataset

fine-tuning-gpt-2-on-a-custom-dataset.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Fine-Tuning GPT-2 on a Custom Dataset",
      "provenance": [],
      "collapsed_sections": [],
      "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/gist/MattPitlyk/45541145ad48b93da395f0a72ec2e7dc/fine-tuning-gpt-2-on-a-custom-dataset.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nXQ9wn5dfiBz",
        "colab_type": "text"
      },
      "source": [
        "# Fine-tuning GPT-2 on a Custom Dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sL4ht6hRHj1W",
        "colab_type": "text"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iSLvhw96EXK5",
        "colab_type": "text"
      },
      "source": [
        "Make sure you enable a GPU or TPU in the runtime  \n",
        "Runtime -> Change Runtime -> Hardware Accelerator "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y-2RF_rof4o0",
        "colab_type": "text"
      },
      "source": [
        "Install helper library for fine-tuning."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KBkpRgBCBS2_",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install -q gpt-2-simple"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FYfTaj9lfgDa",
        "colab_type": "text"
      },
      "source": [
        "Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dJOSDm7EF00R",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Colab pre-installs many common libraries including TensorFlow.\n",
        "# Use Colab magic command to tell Colab to not use TensorFlow 2.0.\n",
        "%tensorflow_version 1.x\n",
        "import gpt_2_simple as gpt2\n",
        "from datetime import datetime\n",
        "from google.colab import files"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gvlc3CXwhDal",
        "colab_type": "text"
      },
      "source": [
        "Download GPT-2 models.  \n",
        "When fine-tuning on a single GPU, only the 124M and 335M size models can be used."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P8wSlgXoDPCR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#gpt2.download_gpt2(model_name=\"124M\")\n",
        "gpt2.download_gpt2(model_name=\"355M\")\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-mLQqkm5G-A1",
        "colab_type": "text"
      },
      "source": [
        "Set variables"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6OFnPCLADfll",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "file_name = \"my_custom_dataset.txt\"\n",
        "run_name = 'fine_tuning_run_1'\n",
        "model_size = '355M'"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d5nqoJ_IhcAj",
        "colab_type": "text"
      },
      "source": [
        "If your custom data is stored in your G-Drive, mount your drive and you can copy the data to Colab with the code below.  \n",
        "Alternatively, you can upload your dataset directly to Colab using the Colab \"Files\" menu on the left (not the \"File\" menu above).  \n",
        "Training examples in the dataset file should be separated with a blank line."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "puq4iC6vUAHc",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "gpt2.mount_gdrive()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-Z6okFD8VKtS",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "gpt2.copy_file_from_gdrive(file_name)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L0qodq58HZmO",
        "colab_type": "text"
      },
      "source": [
        "## Perform fine-tuning"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aeXshJM-Cuaf",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "sess = gpt2.start_tf_sess()\n",
        "\n",
        "gpt2.finetune(sess,\n",
        "              dataset=file_name,\n",
        "              model_name=model_size,\n",
        "              steps=200,\n",
        "              restore_from='fresh',\n",
        "              run_name = run_name,\n",
        "              print_every=10,\n",
        "              sample_every=50,\n",
        "              save_every=50)\n",
        "             # , learning_rate=.00003)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VHdTL8NDbAh3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# If G-Drive is mounted, save the fine-tuned model to the drive.\n",
        "gpt2.copy_checkpoint_to_gdrive(run_name)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rvL8l2rCcp-V",
        "colab_type": "text"
      },
      "source": [
        "End of training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7P_W0Ir-itMu",
        "colab_type": "text"
      },
      "source": [
        "## Explore results"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lWGNVgMekNQr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print(run_name)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DCcx5u7sbPTD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        " # Copy the model from G-Drive if it wasn't trained in this Colab session.\n",
        "gpt2.copy_checkpoint_from_gdrive(run_name=run_name)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KkSL_kfii591",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Must reset the graph if training was done in this Colab session.\n",
        "import tensorflow as tf\n",
        "tf.reset_default_graph()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-fxL77nvAMAX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Load the model.\n",
        "sess = gpt2.start_tf_sess()\n",
        "gpt2.load_gpt2(sess, run_name=run_name)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Iofztc2f-58U",
        "colab_type": "text"
      },
      "source": [
        "Unconditional generation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4RNY6RBI9LmL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "gpt2.generate(sess, run_name=run_name, temperature=.7, length=100, prefix=None, top_k=40, nsamples=10)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "15Es5HiU_GGm",
        "colab_type": "text"
      },
      "source": [
        "Conditional generation (give the model an input prompt)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "82hy6qlX_FtR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "input_prompt = \"Today's weather is\""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_WY9OUZmMqEr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "gpt2.generate(sess, run_name=run_name, temperature=.7, length=100, prefix=input_prompt, top_k=40, nsamples=10)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uTBDSjm_Onwm",
        "colab_type": "text"
      },
      "source": [
        "## Copy model to Google Cloud Storage  \n",
        "We have already copied the model to your G-Drive.  \n",
        "This code will copy the model to your Google Cloud Platform Cloud Storage account (not free), where it can be used in data science deployment pipelines.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NLuEs3LtMp0z",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "project_id = 'your_gcp_project_id'\n",
        "bucket_name = 'your_gcp_bucket_name'"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JXaejaPOOncO",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from google.colab import auth\n",
        "auth.authenticate_user()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z0EmRyEYOnT6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!gcloud config set project {project_id}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "D0VyAFzVXkFe",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Need to zip just the inner folder.\n",
        "# Currently, the name of the folder inside the zip needs to be `run1`.\n",
        "!cd checkpoint; zip -r ../your_model_name.zip ."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aQH4FMhxYCO7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!gsutil -m cp your_model_name.zip gs://{bucket_name}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VlNmSlvSMNXr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bnZrnHLjF0of",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}

What would the format be for the dataset? Could it just be sentences? JSON Format or Text format? This is the stuff that one needs to know, at least for beginners in machine learning. If you're still active on GitHub, would love a reply. Thanks!

Just a text file with each example sentence on its own line. No json needed.

If I want to create a text summary ml model by fine tuning GPT2, then how the text file should be formatted?

the text will get tokenized in its own?
we just have to pass the text file, what kind of formatting should be done,

Question: 'the ques'

Answer: 'the answer'

will this format work?
And secondly my colab session is crashing when we train the model, what can be the solution to this?

It requires >20gb memory so you can subscribe to colab plus or use a free trial virtual machine.

Is that gpt2 fine tuning approach effective btw?