notebooks/GPT2_Fine_Tuning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fine Tuning GPT-2 on text from children's stories.\n",
    "\n",
    "Notebook for fine-tuning GPT-2. This notebook is inspired by this blog post\n",
    "https://towardsdatascience.com/film-script-generation-with-gpt-2-58601b00d371 which I found very helpful. I also\n",
    "of course did a lot of reading of the PyTorch docs.\n",
    "\n",
    "The main library used is the HuggingFace transformers library. https://huggingface.co/transformers/ They strive to make using transformer models dead simple and have a lot of helper functions and simple apis.  That's not what my goal was though, I wanted to understand the details of what I was doing. So I made this harder than it probabably needed to be by working through the fine-tuning process in \"raw\" PyTorch. In the end, while I used the tokenizer and model from the HuggingFace library, I wrote the preprocessing, Dataset class and training loop myself. This notebook is very step by step with lots of comments about why I'm doing what I'm doing.\n",
    "\n",
    "Note that the utilities for text pre-processing (tokenization) in are defined in */src/text_tokenization_utils.py*. \n",
    "\n",
    "The pre-processed text is from project Gutenberg (Brothers Grimm and Beatrix Potter) It's in the */textdata/* directory\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import packages\n",
    "\n",
    "import os, time\n",
    "import pickle\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from transformers import GPT2Tokenizer, PreTrainedTokenizer, GPT2LMHeadModel\n",
    "from transformers.optimization import AdamW, get_linear_schedule_with_warmup\n",
    "\n",
    "# The code for tokenizing the text and storing it in files is imported here\n",
    "import sys\n",
    "sys.path.append(\"../src\") \n",
    "from text_tokenization_utils import make_tokenized_examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tokenize the raw text.\n",
    "\n",
    "I got raw text from the Project Gutenberg site https://www.gutenberg.org/ It has a lot of books that are old enough to be in the public domain. I didnt' make any webscraping code because this was a toy project and it was quicker to literally cut and paste some text and manually clean it a little bit ... very old school.\n",
    "\n",
    "This tokenization only needs to be done once. The tokenization wrapper function *make_tokenized_examples* will read in the text and then save the tokenized files. It take care of making the subdirectory tree for you.  \n",
    "\n",
    "See the docstring for *make_tokenized_examples* for the correct directory structure needed. You can also simply look in the */textdata* directory of the repo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# relative path to the root directory of the text data. change this if you put the text data somewhere else.\n",
    "PWD_TextData_Root = \"../textdata/\"\n",
    "\n",
    "# these are the names of authors for which I provide data. Each author has their own subdirectory\n",
    "authors = [\"Grimm\", # Brother's Grimm fairy tales\n",
    "           \"Beatrix_Potter\", # Peter Rabbit and other stories\n",
    "           \"Lewis_Carrol\"] # Alice and Wonderland and Through the Looking Glass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import the tokenizer from the Transformers library\n",
    "gpt2_tokenizer = GPT2Tokenizer.from_pretrained(\"gpt2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# and loop over the authors to make the tokenized examples\n",
    "# note these will be saved locally in a *tokenized_data* directory as .pkl files. The tokenized files are not \n",
    "# in the github repository because there is a *.pkl in the .gitignore file\n",
    "\n",
    "for author in authors: \n",
    "    file_path = os.path.join(PWD_TextData_Root, author)\n",
    "    #print(file_path)\n",
    "    make_tokenized_examples(gpt2_tokenizer, 256,file_path, examples_file=None);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Make the Dataset and Dataloader\n",
    "\n",
    "This is what PyTorch will use to supply data to the training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class StoryData(torch.utils.data.Dataset):\n",
    "    '''This is a class for loading in a list of tokenized gpt2 examples from a list of file paths'''\n",
    "\n",
    "    def __init__(\n",
    "            self,\n",
    "            file_paths: list):\n",
    "\n",
    "        for fpath in file_paths:\n",
    "            assert os.path.isfile(fpath), \"{} does not exist\".format(fpath)\n",
    "\n",
    "        self.examples = []\n",
    "\n",
    "        for fpath in file_paths:\n",
    "            with open(fpath, 'rb') as f:\n",
    "                examps = pickle.load(f)\n",
    "            self.examples.extend(examps)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.examples)\n",
    "\n",
    "    def __getitem__(self, item):\n",
    "        return torch.tensor(self.examples[item], dtype=torch.long)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#  Now choose which Authors we want to use, comment out lines of the author_list below.\n",
    "# Switching authors will change the nature of the generated text after training.\n",
    "\n",
    "author_list = [\n",
    "    #\"Grimm\",\n",
    "    #\"Lewis_Carrol\",\n",
    "    \"Beatrix_Potter\"\n",
    "]\n",
    "\n",
    "# you will need to modify the paths a bit if you choose another block size\n",
    "file_paths = []\n",
    "for author in author_list:\n",
    "  file_paths.append(os.path.join(PWD_TextData_Root, \n",
    "                                 author, \n",
    "                                 \"tokenized_examples/examples_gpt2_blocksize_256_\" + author + \".pkl\"))\n",
    "print(file_paths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make the actual instance of the Dataset class using the chosen authors \n",
    "story_dataset = StoryData(file_paths)\n",
    "\n",
    "# make the data loader\n",
    "#NOTE: the batch_size is 1 because we will be doing gradient accumulation. This is to get around the fact\n",
    "# that I am using a 8GB RTX 2070 Super GPU which is small\n",
    "story_dataloader = DataLoader(story_dataset, batch_size =1, shuffle = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get ready to fit the model\n",
    "\n",
    "We will choose a pre-trained model from the HuggingFace model hub to fine tune, pick some hyperparameters, make an optimizer and also a learning rate scheduler. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pick a model to train\n",
    "\n",
    "# distilgpt2 will on my 8GB GPU\n",
    "model = GPT2LMHeadModel.from_pretrained('distilgpt2')\n",
    "\n",
    "# gpt2-medium will not train on an 8GB GPU ... but you can generate text with the pre-trained model if you like.\n",
    "# if you have a larger GPU, say a 24 GB RTX 3090 you may wish to try training though\n",
    "# model = GPT2LMHeadModel.from_pretrained('gpt2-medium')\n",
    "\n",
    "# NOTE: look into gradient checkpointing and see if that will allow for training within 8GB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set epochs and batch size\n",
    "N_EPOCHS = 5\n",
    "BATCH_SIZE = 8 # note we are actually going to use gradient accumulation because these models are so big\n",
    "\n",
    "# for the scheduler\n",
    "LEARNING_RATE = 0.0001 #0.00002\n",
    "WARMUP_STEPS = 100 # 10000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# put the model on the gpu. note if this doesn't say you're using the gpu this will not train!\n",
    "\n",
    "device = 'cpu'\n",
    "if torch.cuda.is_available():\n",
    "    print('using gpu')\n",
    "    device = 'cuda'\n",
    "print(device)\n",
    "\n",
    "model.to(device)\n",
    "\n",
    "# create optimizer\n",
    "optimizer = AdamW(model.parameters(), lr=LEARNING_RATE)\n",
    "\n",
    "# make scheduler (for varying learning rate over time)\n",
    "scheduler = get_linear_schedule_with_warmup(optimizer=optimizer, num_warmup_steps=WARMUP_STEPS, num_training_steps=-1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Before fine tuning the model, let's generate some text from it to see what it produces.\n",
    "\n",
    "You can run this and generate your own text from the pre-trained model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make a function to generate text from the model\n",
    "\n",
    "def generate_test_text(model, max_length=256, input_text=None):\n",
    "    model.eval() # put the model in eval mode\n",
    "    if input_text is None:\n",
    "        input_text = \"Once upon a time there was a little mouse.\"\n",
    "    input_ids = gpt2_tokenizer.encode(input_text, return_tensors='pt')\n",
    "    input_ids = input_ids.to('cuda')\n",
    "    output_ids = model.generate(input_ids, \n",
    "                                pad_token_id=gpt2_tokenizer.eos_token_id,\n",
    "                                max_length=max_length, \n",
    "                                do_sample=True, \n",
    "                                top_p=0.95, \n",
    "                                top_k=60,\n",
    "                                num_return_sequences=1)\n",
    "\n",
    "    output_text = gpt2_tokenizer.decode(output_ids[0])\n",
    "    return output_text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# and do some generation. The output should be different every time you run this cell \n",
    "# and of all sorts of topics and styles\n",
    "prompt =\"once upon a time there was a little mouse\" \n",
    "print(generate_test_text(model,input_text=prompt, max_length=256))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A couple examples\n",
    "\n",
    "You can see that the style is of general text you'd read in a book and the content is highly variable across independent generations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Example 1**\n",
    "\n",
    "once upon a time there was a little mouse in the sky and I looked at you from here, and you saw the thing which was really not a bird, I had always seen a bird and all those things but something with the mouse there was something. This was something that could be fixed and you just looked at what you were saying. It had a lot of little fish in it, and I knew this to be very precise. The bird had an eye on its mouth and its mouth were almost like the size of a tooth and not a long blade, and it looked very good and very well-rounded.\n",
    "\n",
    "So what has your new favourite piece of bird writing in The Sims?\n",
    "The author's favourite piece of bird writing in The Sims is the paper bird - and I would like to start by saying that the only way to do that with the new one is by putting this one on. So you're talking about one of your favourite words, which is a line from the Sim guide. How has it been in your life since The Sims first started?\n",
    "I didn't really think it would be that easy, but there were the other things to try and do. I was really curious for how it would be if you could get yourself into this and then"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Example2**\n",
    "\n",
    "once upon a time there was a little mouse on the road and an idea for an alternate timeline of the future we knew we were all going to be fighting for.\n",
    "The idea was made for an alternate timeline of a future where humans, and of course mankind, had different beliefs regarding the future, and the world was not based on the current order, but on the present and future world. The future was a great place and it was clear that there was something that was going to be found on Earth...\n",
    "The future changed completely but it was far from new. It wasn't until a long time ago that everyone in the future had different beliefs about the future, and this made it even more important. However, I do realize some thought it might not seem like the entire universe had to be built on Earth.\n",
    "This was just one of the things I was so excited to work with.\n",
    "I felt that I was getting the best of everything from the world to the planet... but, to this day, I'm having my second go at it.\n",
    "The world in this scenario was made up of a world that is mostly earth like ours, inhabited by people who could relate to the history of the galaxy and the past. One such thing happened when I met an"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fine tune the model\n",
    "\n",
    "This is the training loop. It shouldn't take too long to run ... a few minutes or so depending on your GPU\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is the link to the gradient accumulation documentation\n",
    "# https://pytorch.org/docs/stable/notes/amp_examples.html#gradient-accumulation# \n",
    "\n",
    "#TODO: look into gradient checkpointing as well\n",
    "\n",
    "epoch_loss = 0.0 # used to track loss for each epoch\n",
    "\n",
    "internal_batch_count = 0 # used to track number of examples within each batch. This is necessary \n",
    "                         # because of the gradient accumulation hack (to deal with my 8GB GPU memory)\n",
    "    \n",
    "# make FP16 scaler for faster training\n",
    "scaler = torch.cuda.amp.GradScaler()\n",
    "    \n",
    "# put the model into training mode\n",
    "model.train()\n",
    "\n",
    "for epoch in range(N_EPOCHS): # iterate over epochs\n",
    "    \n",
    "    print(\"started epoch {}\".format(epoch))\n",
    "    \n",
    "    for idx, text in enumerate(story_dataloader):  # this data loader is set up to shuffle automatically\n",
    "        \n",
    "        # Do the forward propagation. \n",
    "        # Don't forget to put the text onto the gpu.\n",
    "        # you have to put in labels to get the loss as an output.\n",
    "        # because GPT is an autogregessive model the input is the output for training purposes\n",
    "        \n",
    "        with torch.cuda.amp.autocast():\n",
    "            outputs = model(text.to(device), labels=(text.to(device)))\n",
    "        \n",
    "            # get the loss out so we can do backwards propagation\n",
    "            loss, logits = outputs[:2]\n",
    "            loss = loss / BATCH_SIZE \n",
    "        \n",
    "        # do backpropagation. yay autodifferentiation!\n",
    "        # note the use of the scaler for the FP16 \n",
    "        scaler.scale(loss).backward()\n",
    "        \n",
    "        # keep track of the loss\n",
    "        epoch_loss = epoch_loss + loss.detach().cpu().numpy()  # need to detach the gradients \n",
    "                                                               # because we only care about the numerical value\n",
    "                                                               # also store the epoch loss on the cpu as numpy\n",
    "            \n",
    "        # increment the internal_batch_count\n",
    "        internal_batch_count = internal_batch_count + 1\n",
    "        \n",
    "        # Now, if we have run through a full batch, take some optimizer and gradient steps\n",
    "        if internal_batch_count == BATCH_SIZE:\n",
    "            internal_batch_count = 0 # reset this\n",
    "            \n",
    "            # take an optimizer step. note the use of the scaler for FP16\n",
    "            scaler.step(optimizer) \n",
    "            scaler.update()\n",
    "            optimizer.zero_grad() # zero out the gradients in the optimizer\n",
    "            \n",
    "            model.zero_grad() # zero out the gradients we've been accumulating in the model\n",
    "            \n",
    "            scheduler.step() # take a scheduler step\n",
    "            \n",
    "            \n",
    "    # Now that we've gone through an epoch, let's see what the loss is and what some generated text looks like\n",
    "    \n",
    "    # put the model into evaluation mode\n",
    "    model.eval()\n",
    "    \n",
    "    # print the loss\n",
    "    print(\"Epoch {} has loss {}\".format(epoch, epoch_loss))\n",
    "    # reset the loss\n",
    "    epoch_loss = 0.0\n",
    "    \n",
    "    # uncomment this if you want to print some test text after each epoch\n",
    "    #print(generate_test_text(model,input_text=prompt))\n",
    "    \n",
    "    \n",
    "    # put the model back in training mode\n",
    "    model.train()\n",
    "        \n",
    "        \n",
    "        \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## OK now generate some text!  \n",
    "\n",
    "That's pretty much it for fine tuning, you can now play around and generate text from the trained model.The text is generally gramatically correct but doesn't hang together as a story and is pretty rambling. distilgpt2 is a rather small language model. Text generated from larger language models tends to be very fluent however.\n",
    "\n",
    "What is however clear, is that the fine-tuning is doing matching the authors upon who's text we are tuning. Beatrix Potter and Brother's Grimm are very different afterall!  Farther below I give some example generated text that demonstrates the differences resulting from fine tuning on different authors, but you can of course generate your own text using the next cell. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pick a prompt\n",
    "prompt =\"Once upon a time there was a little mouse\" \n",
    "\n",
    "# And generate text!\n",
    "print(generate_test_text(model,input_text=prompt, max_length=256))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Here are two examples of text generated after fine tuning on Beatrix Potter stories:\n",
    "\n",
    "Note that while the text isn't very sensible, it matches the general structure and linguistic style of Beatrix Potter stories. Lots of short sentences separated into different lines. It's also pretty light hearted. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Example 1**\n",
    "Once upon a time there was a little mouse on the edge of a wood-row, and a young\n",
    "mouse!\n",
    "\n",
    "\"I shall see; I wish ye had\n",
    "never seen another mouse!\" said\n",
    "Pumpkin.\n",
    "\n",
    "Then there was a very fine table-cloth covered with\n",
    "dry dishes and very little chairs!\n",
    "\n",
    "\"A beautiful table-cloth,\" said a big lump-handling gentleman\n",
    "who has been at the table, with a nice brown\n",
    "and a very fine wooden plate.\n",
    "\n",
    "\"I am sorry,\" said Mr. Pickles. \"\n",
    "Pumpkin\n",
    "said that he made the meal which I thought\n",
    "was a fat ham!\n",
    "\n",
    "\"What a nice little ham!\" said\n",
    "Mr. Pickles.\n",
    "\n",
    "Pumpkin looked out upon the table, and came upon the\n",
    "dark wood-row, with a little knife and\n",
    "mouse.\n",
    "\n",
    "\"It is quite a hard mouse; it should be too cold\n",
    "if you will. It is quite hard; it is really a\n",
    "little mouse, and is very strong; and\n",
    "hard\"\n",
    "should we have eaten with it?\" said Mr. Pickles, to\n",
    "pumpkin; \"I will have eaten it,\" said Pick\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Example 2**\n",
    "\n",
    "Once upon a time there was a little mouse who could not quite get out. When they met that little person, they bumped into a hole in the wall, where it was very large\n",
    "in a little pottery basket, and began pulling out wood from beneath.\n",
    "\n",
    "One day something\n",
    "dishfuls and a nut had been broken into.\n",
    "\n",
    "\"Who will run in my day?\" said a mouse, looking for\n",
    "another mouse. They had been at a party and got in there; he\n",
    "could only wait until his day came round.\n",
    "\n",
    "\"Come into the garden in the morning!\"\n",
    "\n",
    "The mice took notice, but they\n",
    "felt much too early to hide it.\n",
    "\n",
    "\"This is a little fish, called\n",
    "Tobit!\"\n",
    "\n",
    "The mouse could see a hole in the\n",
    "door, and saw something big\n",
    "in a hole, and they came out of nowhere.\n",
    "\n",
    "\"No one wants to go in here,\" said\n",
    "\"Old Mr. Kitten!\"\n",
    "\n",
    "\n",
    "There came a little mouse whose life was almost complete without him; the mouse\n",
    "dishful was quite empty and\n",
    "very little\n",
    "on account of the rats. It looked rather beautiful, but\n",
    "there was no one who could see\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### And some examples from fine-tuning on Brother's Grimm\n",
    "\n",
    "These have a very different, paragraph type structure ... and they are pretty dark, as we'd expect from Brothers Grimm."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Grimm Example 1**\n",
    "\n",
    "Once upon a time there was a little mouse, and there\n",
    "was some dwarf in his room that had been holding it. ‘If it could stay, then it would just\n",
    "be enough to eat and eat with,’ said she, and the creature said to the dwarf,\n",
    " ‘if you will, I would just eat in a second, or I would\n",
    "eat your meat or something in my kitchen.’\n",
    "\n",
    "So the dwarf went, and the dwarf fell on top of him, and the dwarf\n",
    "was knocked off his head, and the dwarf’s head flew on the ground\n",
    "while the dwarf lay in a heap of ashes.\n",
    "The dwarf went about this, and the dwarf’s head flew in\n",
    "heathers, and at last he fell backwards and he had fallen down upon\n",
    "his face, and now the dwarf went round the castle, and the dwarf\n",
    "fell a little, and then the dwarf fell with him to the fire. The dwarf said,\n",
    "‘Why not? How come back down on the stairs again? When you\n",
    "re on the stair and it seems as if you can have a beautiful old tower.’ So the dwarf threw himself\n",
    "down and fell back, and was forced"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Grimm Example 2**\n",
    "\n",
    "Once upon a time there was a little mouse at the bottom and he would cut a hole\n",
    "into one and lay down in it, and when they were tired, they lay\n",
    "still on the ground under the tree and with their feet beneath it,\n",
    "and as soon as they were tired, they said: ‘Don’t let us kill you.’ They\n",
    "said: ‘Take the mouse that we are, or the mouse that’s cut in half. I’ll cut my mouse,’\n",
    "and he will cut my foot in half to a length, and with the same head cut in\n",
    "half, then the mouse will cut it into the hole, and\n",
    "he will cut the other one in half.’\n",
    "\n",
    "When the mice came in the hole, they found a young man lying still at the top of a tree\n",
    "in which was standing, with a little boy lying lying dead. The\n",
    "young man was lying down in his sack with his little son lying\n",
    "down on the ground, and the old man lying lying on the floor\n",
    "sitting in a dead wood bench. The old man’s face\n",
    "disturbed, but there the old man was crying and\n",
    "snorting and crying"
   ]
  }
 ],
 "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}