A local-ized version of Max Woolf's GPT-2 training Colab notebook

LOCAL-Train_a_GPT_2_Text_Generating_Model_w_GPU.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "H7LoMj4GA4n_"
   },
   "source": [
    "#  Train a GPT-2 Text-Generating Model w/GPU For About €1600 In GPU Acquisition Costs \n",
    "\n",
    "by [Max Woolf](http://minimaxir.com), updated for local use by [Theodric](https://github.com/theodric)\n",
    "\n",
    "*Last update: 2025-01-30*\n",
    "\n",
    "Retrain an advanced text generating neural network on any text dataset **on your local machine** using `gpt-2-simple`!\n",
    "\n",
    "For more about `gpt-2-simple`, you can visit [this GitHub repository](https://github.com/minimaxir/gpt-2-simple).\n",
    "\n",
    "\n",
    "To get started:\n",
    "\n",
    "1. Copy this notebook to a working directory on your machine, which obviously runs Linux because you're not a Windows luser or an Apple paypig.\n",
    "2. Place a text file containing your finetune training data in the same directory\n",
    "3. Create a virtualenv (`virtualenv venv`)\n",
    "4. Spin up the virtualenv (`source venv/bin/activate`)\n",
    "5. Install CUDA prerequisites per [the TensorFlow instructions](https://www.tensorflow.org/install/pip#linux_setup) (`pip install tensorflow[and-cuda]`)\n",
    "6. Install notebook-specific prerequisites (`pip3 install jupyter datetime gpt_2_simple`)\n",
    "7. Run the jupyter notebook server (`jupyter notebook`)\n",
    "8. Open the link farted out by the jupyter notebook server in your web browser of choice\n",
    "9. Run the cells below:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "KBkpRgBCBS2_",
    "outputId": "a03d4834-ee68-441e-a1f9-c48e1949594c"
   },
   "outputs": [],
   "source": [
    "!pip install -q gpt_2_simple\n",
    "import gpt_2_simple as gpt2\n",
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Bj2IJLHP3KwE"
   },
   "source": [
    "## Verify GPU\n",
    "\n",
    "You can verify which GPU is active and available in CUDA by running the cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 306
    },
    "colab_type": "code",
    "id": "sUmTooTW3osf",
    "outputId": "850a6a7c-2d33-4b4a-b1b4-a5b8b096d0da"
   },
   "outputs": [],
   "source": [
    "!nvidia-smi\n",
    "!python3 -c \"import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "0wXB05bPDYxS"
   },
   "source": [
    "## Downloading GPT-2\n",
    "\n",
    "If you're retraining a model on new text, you need to download the GPT-2 model first. \n",
    "\n",
    "There are two sizes of GPT-2:\n",
    "\n",
    "* `117M` (default): the \"small\" model, 500MB on disk.\n",
    "* `345M`: the \"medium\" model, 1.5GB on disk.\n",
    "\n",
    "Larger models have more knowledge, but take longer to finetune and longer to generate text. You can specify which base model to use by changing `model_name` in the cells below.\n",
    "\n",
    "The next cell downloads it from Google Cloud Storage and saves it to your hard drive at `/models/<model_name>`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 136
    },
    "colab_type": "code",
    "id": "P8wSlgXoDPCR",
    "outputId": "03cdcf24-2beb-4829-f7a5-7795e7bb1749"
   },
   "outputs": [],
   "source": [
    "gpt2.download_gpt2(model_name=\"345M\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "HeeSKtNWUedE"
   },
   "source": [
    "Next, specify the text file containing the data you wish to use for finetuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_name = \"hey_bozo_populate_your_source_data_filename_here.txt\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "LdpZQXknFNY3"
   },
   "source": [
    "## Finetune GPT-2\n",
    "\n",
    "The next cell will start the actual finetuning of GPT-2. It creates a persistent TensorFlow session which stores the training config, then runs the training for the specified number of `steps`. (to have the finetuning run indefinitely, set `steps = -1`)\n",
    "\n",
    "The model checkpoints will be saved in `/checkpoint/run1` by default. The checkpoints are saved every 500 steps (can be changed) and when the cell is stopped.\n",
    "\n",
    "**IMPORTANT NOTE:** If you want to rerun this cell, **restart the kernel first**. You will need to rerun imports but not recopy files.\n",
    "\n",
    "Other optional-but-helpful parameters for `gpt2.finetune`:\n",
    "\n",
    "\n",
    "*  **`restore_from`**: Set to `fresh` to start training from the base GPT-2, or set to `latest` to restart training from an existing checkpoint.\n",
    "* **`sample_every`**: Number of steps to print example output\n",
    "* **`print_every`**: Number of steps to print training progress.\n",
    "* **`learning_rate`**:  Learning rate for the training. (default `1e-4`, can lower to `1e-5` if you have <1MB input data)\n",
    "*  **`run_name`**: subfolder within `checkpoint` to save the model. This is useful if you want to work with multiple models (will also need to specify  `run_name` when loading the model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 9098
    },
    "colab_type": "code",
    "id": "aeXshJM-Cuaf",
    "outputId": "dfd723ab-086c-4cbd-af52-3a6ca0dce510"
   },
   "outputs": [],
   "source": [
    "sess = gpt2.start_tf_sess()\n",
    "\n",
    "gpt2.finetune(sess,\n",
    "              learning_rate=1e-5,\n",
    "              dataset=file_name,\n",
    "              model_name='345M',\n",
    "              steps=250,\n",
    "              restore_from='fresh',\n",
    "              print_every=10,\n",
    "              sample_every=50,\n",
    "            save_every=50\n",
    "              )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "IXSuTNERaw6K"
   },
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "IXSuTNERaw6K"
   },
   "source": [
    "After the model is trained, you will have a trained model. Incredible!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "qQJgV_b4bmzd"
   },
   "source": [
    "You're done! Feel free to go to the **Generate Text From The Trained Model** section to generate text based on your retrained model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "pel-uBULXO2L"
   },
   "source": [
    "## Load a Trained Model Checkpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "RTa6zf3e_9gV"
   },
   "source": [
    "The next cell will allow you to load the retrained model checkpoint + metadata necessary to generate text.\n",
    "\n",
    "**IMPORTANT NOTE:** If you want to rerun this cell, **restart the kernel first**. You will need to rerun imports but not recopy files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-fxL77nvAMAX"
   },
   "outputs": [],
   "source": [
    "sess = gpt2.start_tf_sess()\n",
    "gpt2.load_gpt2(sess)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ClJwpF_ACONp"
   },
   "source": [
    "## Generate Text From The Trained Model\n",
    "\n",
    "After you've trained the model or loaded a retrained model from checkpoint, you can now generate text. `generate` generates a single text from the loaded model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 2740
    },
    "colab_type": "code",
    "id": "4RNY6RBI9LmL",
    "outputId": "593aa543-9d85-42ff-c9ce-d654bf623889"
   },
   "outputs": [],
   "source": [
    "gpt2.generate(sess)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "oF4-PqF0Fl7R"
   },
   "source": [
    "If you're creating an API based on your model and need to pass the generated text elsewhere, you can do `text = gpt2.generate(sess, return_as_list=True)[0]`\n",
    "\n",
    "You can also pass in a `prefix` to the generate function to force the text to start with a given character sequence and generate text from there (good if you add an indicator when the text starts).\n",
    "\n",
    "You can also generate multiple texts at a time by specifing `nsamples`. Unique to GPT-2, you can pass a `batch_size` to generate multiple samples in parallel, giving a massive speedup (in Colaboratory, set a maximum of 20 for `batch_size`).\n",
    "\n",
    "Other optional-but-helpful parameters for `gpt2.generate` and friends:\n",
    "\n",
    "*  **`length`**: Number of tokens to generate (default 1023, the maximum)\n",
    "* **`temperature`**: The higher the temperature, the crazier the text (default 0.7, recommended to keep between 0.7 and 1.0)\n",
    "* **`truncate`**: Truncates the input text until a given sequence, excluding that sequence (e.g. if `truncate='<|endoftext|>'`, the returned text will include everything before the first `<|endoftext|>`). It may be useful to combine this with a smaller `length` if the input texts are short.\n",
    "*  **`include_prefix`**: If using `truncate` and `include_prefix=False`, the specified `prefix` will not be included in the returned text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 238667
    },
    "colab_type": "code",
    "id": "8DKMc0fiej4N",
    "outputId": "3bd90377-68d5-4777-fc65-145e6e597d74"
   },
   "outputs": [],
   "source": [
    "gpt2.generate(sess,\n",
    "              length=250,\n",
    "              temperature=1.3,\n",
    "              prefix=\"My UNIX WORKSTATION is \",\n",
    "              nsamples=7000,\n",
    "              batch_size=5\n",
    "              )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "zjjEN2Tafhl2"
   },
   "source": [
    "For bulk generation, you can generate a large amount of text to a file and sort out the samples locally on your computer. The next cell will generate a generated text file with a unique timestamp and then download it.\n",
    "\n",
    "You can rerun the cell as many times as you want for even more generated texts!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Fa6p6arifSL0"
   },
   "outputs": [],
   "source": [
    "gen_file = 'gpt2_gentext_{:%Y%m%d_%H%M%S}.txt'.format(datetime.utcnow())\n",
    "\n",
    "gpt2.generate_to_file(sess,\n",
    "                      destination_path=gen_file,\n",
    "                      length=500,\n",
    "                      temperature=1,\n",
    "                      nsamples=100,\n",
    "                      batch_size=20\n",
    "                      )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "wmTXWNUygS5E"
   },
   "source": [
    "# LICENSE\n",
    "\n",
    "MIT License\n",
    "\n",
    "Copyright (c) 2019 Max Woolf\n",
    "\n",
    "Permission is hereby granted, free of charge, to any person obtaining a copy\n",
    "of this software and associated documentation files (the \"Software\"), to deal\n",
    "in the Software without restriction, including without limitation the rights\n",
    "to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n",
    "copies of the Software, and to permit persons to whom the Software is\n",
    "furnished to do so, subject to the following conditions:\n",
    "\n",
    "The above copyright notice and this permission notice shall be included in all\n",
    "copies or substantial portions of the Software.\n",
    "\n",
    "THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n",
    "IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n",
    "FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n",
    "AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n",
    "LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n",
    "OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n",
    "SOFTWARE.\n",
    "\n",
    "Ass and balls :) Thank you for reading!!"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Train a GPT-2 Text-Generating Model w/ GPU",
   "provenance": [],
   "toc_visible": true,
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}