Train_a_language_model.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/yotamnahum/Mamram-Language-Modelling-Workshop/blob/main/Train_a_language_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": "X4cRE8IbIrIV"
      },
      "source": [
        "If you're opening this Notebook on colab, you will probably need to install 🤗 Transformers and 🤗 Datasets. Uncomment the following cell and run it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "MOsHUjgdIrIW"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "! pip install datasets transformers accelerate -U"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FjBU65E4V833"
      },
      "source": [
        "Make sure your version of Transformers is at least 4.11.0 since the functionality was introduced in that version:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "F93KlnMXV834",
        "outputId": "e365753b-240f-46a3-f876-bbb4cc738b73",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "4.33.1\n",
            "cuda\n"
          ]
        }
      ],
      "source": [
        "import transformers\n",
        "import torch\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "print(transformers.__version__)\n",
        "print(device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a3KD3WXU3l-O"
      },
      "source": [
        "# Fine-tuning a language model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1r_n9OWV3l-Q"
      },
      "source": [
        "## Preparing the dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kswRMhPc3l-Q"
      },
      "source": [
        "For each of those tasks, we will use the `english_quotes` dataset as an example. You can load it very easily with the 🤗 Datasets library.\n",
        "\n",
        "You can choose other dataset [here](https://huggingface.co/models?pipeline_tag=text-generation&sort=trending), or upload your own data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "n2ZRs1cL3l-R",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "028f49d5-02ab-46b3-bf89-5bdd316db0b2"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DatasetDict({\n",
              "    train: Dataset({\n",
              "        features: ['quote', 'author', 'tags'],\n",
              "        num_rows: 2257\n",
              "    })\n",
              "    test: Dataset({\n",
              "        features: ['quote', 'author', 'tags'],\n",
              "        num_rows: 251\n",
              "    })\n",
              "})"
            ]
          },
          "metadata": {},
          "execution_count": 2
        }
      ],
      "source": [
        "from datasets import load_dataset\n",
        "datasets = load_dataset(\"Abirate/english_quotes\", split='train')\n",
        "datasets = datasets.train_test_split(test_size=0.1, shuffle=True, seed=253)\n",
        "datasets"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WHUmphG3IrI3"
      },
      "source": [
        "To get a sense of what the data looks like, the following function will show some examples picked randomly in the dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "ur5sNUcZ3l-g",
        "outputId": "73bb764d-5f28-45af-88b8-be565246f786",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 449
        }
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>quote</th>\n",
              "      <th>author</th>\n",
              "      <th>tags</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>“I'm quite illiterate, but I read a lot. ”</td>\n",
              "      <td>J.D. Salinger,</td>\n",
              "      <td>[holden]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>“Kiss me, and you will see how important I am.”</td>\n",
              "      <td>Sylvia Plath,</td>\n",
              "      <td>[importance, kiss, kissing]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>“Make your own Bible. Select and collect all the words and sentences that in all your readings have been to you like the blast of a trumpet.”</td>\n",
              "      <td>Ralph Waldo Emerson</td>\n",
              "      <td>[spirituality]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>“There is more than one way to burn a book. And the world is full of people running about with lit matches.”</td>\n",
              "      <td>Ray Bradbury</td>\n",
              "      <td>[censorship]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>“I am too intelligent, too demanding, and too resourceful for anyone to be able to take charge of me entirely. No one knows me or loves me completely. I have only myself”</td>\n",
              "      <td>Simone de Beauvoir</td>\n",
              "      <td>[feminism]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>5</th>\n",
              "      <td>“Happiness in intelligent people is the rarest thing I know.”</td>\n",
              "      <td>Ernest Hemingway,</td>\n",
              "      <td>[happiness]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>6</th>\n",
              "      <td>“People, generally, suck.”</td>\n",
              "      <td>Christopher Moore,</td>\n",
              "      <td>[humor]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>7</th>\n",
              "      <td>“And now Iâ€™m looking at you,â€� he said, â€œand youâ€™re asking me if I still want you, as if I could stop loving you. As if I would want to give up the thing that makes me stronger than anything else ever has. I never dared give much of myself to anyone before â€“ bits of myself to the Lightwoods, to Isabelle and Alec, but it took years to do it â€“ but, Clary, since the first time I saw you, I have belonged to you completely. I still do. If you want me.”</td>\n",
              "      <td>Cassandra Clare,</td>\n",
              "      <td>[cassandra-clare, city-of-glass, clary, jace, love, mortal-instruments]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>8</th>\n",
              "      <td>“The scar had not pained Harry for nineteen years. All was well.”</td>\n",
              "      <td>J.K. Rowling,</td>\n",
              "      <td>[harry-potter]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>9</th>\n",
              "      <td>“Ah! There is nothing like staying at home, for real comfort.”</td>\n",
              "      <td>Jane Austen</td>\n",
              "      <td>[relaxation]</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>"
            ]
          },
          "metadata": {}
        }
      ],
      "source": [
        "from datasets import ClassLabel\n",
        "import random\n",
        "import pandas as pd\n",
        "from IPython.display import display, HTML\n",
        "\n",
        "def show_random_elements(dataset, num_examples=10):\n",
        "    assert num_examples <= len(dataset), \"Can't pick more elements than there are in the dataset.\"\n",
        "    picks = []\n",
        "    for _ in range(num_examples):\n",
        "        pick = random.randint(0, len(dataset)-1)\n",
        "        while pick in picks:\n",
        "            pick = random.randint(0, len(dataset)-1)\n",
        "        picks.append(pick)\n",
        "\n",
        "    df = pd.DataFrame(dataset[picks])\n",
        "    for column, typ in dataset.features.items():\n",
        "        if isinstance(typ, ClassLabel):\n",
        "            df[column] = df[column].transform(lambda i: typ.names[i])\n",
        "    display(HTML(df.to_html()))\n",
        "\n",
        "show_random_elements(datasets[\"train\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CKerdF353l-o"
      },
      "source": [
        "As we can see, some of the texts are a full paragraph of a Wikipedia article while others are just titles or empty lines."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JEA1ju653l-p"
      },
      "source": [
        "## Causal Language modeling"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Choose [here](https://huggingface.co/models?pipeline_tag=text-generation&sort=trending) a model to start from"
      ],
      "metadata": {
        "id": "dvU7yhJ_iA24"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "-WGBCO343l-q"
      },
      "outputs": [],
      "source": [
        "model_checkpoint = \"distilgpt2\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5io6fY_d3l-u"
      },
      "source": [
        "To tokenize all our texts with the same vocabulary that was used when training the model, we have to download a pretrained tokenizer. This is all done by the `AutoTokenizer` class:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "iAYlS40Z3l-v"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoTokenizer\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rpOiBrJ13l-y"
      },
      "source": [
        "We can now call the tokenizer on all our texts. This is very simple, using the [`map`](https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map) method from the Datasets library. First we define a function that call the tokenizer on our texts:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "lS2m25YM3l-z"
      },
      "outputs": [],
      "source": [
        "def tokenize_function(examples):\n",
        "    return tokenizer(examples[\"quote\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M9xVAa3s3l-2"
      },
      "source": [
        "Then we apply it to all the splits in our `datasets` object, using `batched=True` and 4 processes to speed up the preprocessing. We won't need the `text` column afterward, so we discard it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "NVAO0H8u3l-3"
      },
      "outputs": [],
      "source": [
        "tokenized_datasets = datasets.map(tokenize_function, batched=True, num_proc=4, remove_columns=[\"quote\", \"author\", \"tags\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8qik3J_C3l-7"
      },
      "source": [
        "If we now look at an element of our datasets, we will see the text have been replaced by the `input_ids` the model will need:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "nYv_mcKk3l-7",
        "outputId": "a63be071-8ae3-4cbb-9251-d0eef06ead2e",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'input_ids': [447, 250, 1858, 743, 307, 1661, 618, 356, 389, 34209, 284, 2948, 21942, 11, 475, 612, 1276, 1239, 307, 257, 640, 618, 356, 2038, 284, 5402, 13, 447, 251], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}\n"
          ]
        }
      ],
      "source": [
        "print(tokenized_datasets[\"train\"][1])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "obvgcXda3l--"
      },
      "source": [
        "Now for the harder part: we need to concatenate all our texts together then split the result in small chunks of a certain `block_size`. To do this, we will use the `map` method again, with the option `batched=True`. This option actually lets us change the number of examples in the datasets by returning a different number of examples than we got. This way, we can create our new samples from a batch of examples.\n",
        "\n",
        "First, we grab the maximum length our model was pretrained with. This might be a big too big to fit in your GPU RAM, so here we take a bit less at just 128."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "DVHs5aCA3l-_"
      },
      "outputs": [],
      "source": [
        "# block_size = tokenizer.model_max_length\n",
        "block_size = 128"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RpNfGiMw3l_A"
      },
      "source": [
        "Then we write the preprocessing function that will group our texts:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "iaAJy5Hu3l_B"
      },
      "outputs": [],
      "source": [
        "def group_texts(examples):\n",
        "    # Concatenate all texts.\n",
        "    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}\n",
        "    total_length = len(concatenated_examples[list(examples.keys())[0]])\n",
        "    # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can\n",
        "        # customize this part to your needs.\n",
        "    total_length = (total_length // block_size) * block_size\n",
        "    # Split by chunks of max_len.\n",
        "    result = {\n",
        "        k: [t[i : i + block_size] for i in range(0, total_length, block_size)]\n",
        "        for k, t in concatenated_examples.items()\n",
        "    }\n",
        "    result[\"labels\"] = result[\"input_ids\"].copy()\n",
        "    return result"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LGJWXtNv3l_C"
      },
      "source": [
        "First note that we duplicate the inputs for our labels. This is because the model of the 🤗 Transformers library apply the shifting to the right, so we don't need to do it manually.\n",
        "\n",
        "Also note that by default, the `map` method will send a batch of 1,000 examples to be treated by the preprocessing function. So here, we will drop the remainder to make the concatenated tokenized texts a multiple of `block_size` every 1,000 examples. You can adjust this behavior by passing a higher batch size (which will also be processed slower). You can also speed-up the preprocessing by using multiprocessing:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "gXUSfBrq3l_C"
      },
      "outputs": [],
      "source": [
        "lm_datasets = tokenized_datasets.map(\n",
        "    group_texts,\n",
        "    batched=True,\n",
        "    batch_size=1000,\n",
        "    num_proc=4,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6n84V8Gc3l_G"
      },
      "source": [
        "And we can check our datasets have changed: now the samples contain chunks of `block_size` contiguous tokens, potentially spanning over several of our original texts."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "hTeGCLl_3l_G",
        "outputId": "42959730-a8f2-4739-ace8-7c5686acc150",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 53
        }
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "' groups, parties, nations and epochs, it is the rule.”“Is it so bad, then, to be misunderstood? Pythagoras was misunderstood, and Socrates, and Jesus, and Luther, and Copernicus, and Galileo, and Newton, and every pure and wise spirit that ever took flesh. To be great is to be misunderstood.”“Try not to become a man of success. Rather become a man of value.”“When he died, all things soft and beautiful and bright would be buried with him.”“Things need not have happened to be true.'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 13
        }
      ],
      "source": [
        "tokenizer.decode(lm_datasets[\"train\"][1][\"input_ids\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iEmeQ7Xm3l_H"
      },
      "source": [
        "Now that the data has been cleaned, we're ready to instantiate our `Trainer`. We will a model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "sPqQA3TT3l_I"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoModelForCausalLM\n",
        "model = AutoModelForCausalLM.from_pretrained(model_checkpoint).to(device)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def generate_text(text, **kwargs):\n",
        "    input_ids = tokenizer(text, return_tensors='pt').input_ids.to(device)\n",
        "    output = model.generate(input_ids, pad_token_id=tokenizer.eos_token_id, **kwargs)\n",
        "    return tokenizer.decode(output[0], skip_special_tokens=True)\n",
        "\n",
        "# Usage example:\n",
        "# Assuming `tokenizer`, `model`, and `device` are already defined\n",
        "text = \"AI is\".strip()\n",
        "\n",
        "generated_text = generate_text(text, do_sample=True, max_length=64, top_p=0.95, top_k=0)\n",
        "print(generated_text)"
      ],
      "metadata": {
        "id": "x1NmCAmhdn4x",
        "outputId": "3ff2eb81-26ff-4fac-de9d-f1164085422a",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "AI is solid, smart and challenging new technology that really makes it easy to manufacture something, does a little bit different, and are definitely worth buying.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "generated_text = generate_text(text, do_sample=True, max_length=65, penalty_alpha=0.6, top_k=30)\n",
        "print(generated_text)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ctkrF2tahMlr",
        "outputId": "6a392423-015f-4286-a393-fb46c2ffa581"
      },
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "AI is the one that is being tested, and we've got an important problem. That's why we have built this software for the mobile version of the Android platform. We've got our product in development.\n",
            "\n",
            "\n",
            "\n",
            "\n",
            "That's a shame for us, though. Even though we do the right thing when it\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VyPQTOF_3l_J"
      },
      "source": [
        "And some `TrainingArguments`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 40,
      "metadata": {
        "id": "YbSwEhQ63l_L"
      },
      "outputs": [],
      "source": [
        "from transformers import Trainer, TrainingArguments\n",
        "model_name = model_checkpoint.split(\"/\")[-1]\n",
        "training_args = TrainingArguments(\n",
        "    f\"{model_name}-finetuned\",\n",
        "    learning_rate=4e-5,\n",
        "    # weight_decay=0.01,\n",
        "    push_to_hub=False,\n",
        "    #max_steps=20, # for testing\n",
        "\tnum_train_epochs=1, # for testing\n",
        "    # logging & evaluation strategies\n",
        "    evaluation_strategy=\"steps\",\n",
        "    eval_steps=10,\n",
        "    logging_strategy=\"steps\",\n",
        "    logging_steps=5,\n",
        "    report_to=\"tensorboard\",\n",
        "    save_total_limit=2,\n",
        "    load_best_model_at_end=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sZRbT9ui3l_N"
      },
      "source": [
        "We pass along all of those to the `Trainer` class:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 41,
      "metadata": {
        "id": "OEuqwIra3l_N"
      },
      "outputs": [],
      "source": [
        "trainer = Trainer(\n",
        "    model=model,\n",
        "    args=training_args,\n",
        "    train_dataset=lm_datasets[\"train\"],\n",
        "    eval_dataset=lm_datasets[\"test\"],\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Vvz34Td3l_O"
      },
      "source": [
        "And we can train our model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 42,
      "metadata": {
        "id": "NyZvu_MF3l_P",
        "outputId": "e8370a2b-364a-42e3-9bdd-0f9d5c64a476",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 971
        }
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "    <div>\n",
              "      \n",
              "      <progress value='282' max='282' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
              "      [282/282 00:19, Epoch 3/3]\n",
              "    </div>\n",
              "    <table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              " <tr style=\"text-align: left;\">\n",
              "      <th>Step</th>\n",
              "      <th>Training Loss</th>\n",
              "      <th>Validation Loss</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <td>10</td>\n",
              "      <td>2.797400</td>\n",
              "      <td>3.751385</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>20</td>\n",
              "      <td>2.847100</td>\n",
              "      <td>3.722462</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>30</td>\n",
              "      <td>2.875800</td>\n",
              "      <td>3.727870</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>40</td>\n",
              "      <td>2.948600</td>\n",
              "      <td>3.735564</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>50</td>\n",
              "      <td>3.081300</td>\n",
              "      <td>3.657136</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>60</td>\n",
              "      <td>3.139600</td>\n",
              "      <td>3.635551</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>70</td>\n",
              "      <td>3.405500</td>\n",
              "      <td>3.554279</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>80</td>\n",
              "      <td>3.376900</td>\n",
              "      <td>3.569230</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>90</td>\n",
              "      <td>3.575600</td>\n",
              "      <td>3.562459</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>100</td>\n",
              "      <td>2.962800</td>\n",
              "      <td>3.627850</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>110</td>\n",
              "      <td>2.662000</td>\n",
              "      <td>3.705742</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>120</td>\n",
              "      <td>2.781200</td>\n",
              "      <td>3.641534</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>130</td>\n",
              "      <td>2.753200</td>\n",
              "      <td>3.670249</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>140</td>\n",
              "      <td>2.847900</td>\n",
              "      <td>3.651517</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>150</td>\n",
              "      <td>2.784200</td>\n",
              "      <td>3.623347</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>160</td>\n",
              "      <td>2.898900</td>\n",
              "      <td>3.598698</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>170</td>\n",
              "      <td>2.779900</td>\n",
              "      <td>3.612564</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>180</td>\n",
              "      <td>2.773500</td>\n",
              "      <td>3.610681</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>190</td>\n",
              "      <td>2.630100</td>\n",
              "      <td>3.616312</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>200</td>\n",
              "      <td>2.340600</td>\n",
              "      <td>3.771531</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>210</td>\n",
              "      <td>2.303100</td>\n",
              "      <td>3.766255</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>220</td>\n",
              "      <td>2.474300</td>\n",
              "      <td>3.748519</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>230</td>\n",
              "      <td>2.437500</td>\n",
              "      <td>3.758205</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>240</td>\n",
              "      <td>2.388400</td>\n",
              "      <td>3.738277</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>250</td>\n",
              "      <td>2.395800</td>\n",
              "      <td>3.734878</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>260</td>\n",
              "      <td>2.397900</td>\n",
              "      <td>3.744394</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>270</td>\n",
              "      <td>2.277100</td>\n",
              "      <td>3.749841</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>280</td>\n",
              "      <td>2.397300</td>\n",
              "      <td>3.748843</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table><p>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "TrainOutput(global_step=282, training_loss=2.760936219641503, metrics={'train_runtime': 19.864, 'train_samples_per_second': 112.968, 'train_steps_per_second': 14.197, 'total_flos': 73293738541056.0, 'train_loss': 2.760936219641503, 'epoch': 3.0})"
            ]
          },
          "metadata": {},
          "execution_count": 42
        }
      ],
      "source": [
        "trainer.train()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3APq-vUc3l_R"
      },
      "source": [
        "Once the training is completed, we can evaluate our model and get its perplexity on the validation set like this:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 43,
      "metadata": {
        "id": "diKZnB1I3l_R",
        "outputId": "80bb7a1d-8a72-41b9-b5c4-1ef605d147aa",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 55
        }
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "    <div>\n",
              "      \n",
              "      <progress value='12' max='12' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
              "      [12/12 00:00]\n",
              "    </div>\n",
              "    "
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Perplexity: 42.47\n"
          ]
        }
      ],
      "source": [
        "import math\n",
        "eval_results = trainer.evaluate()\n",
        "print(f\"Perplexity: {math.exp(eval_results['eval_loss']):.2f}\")"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "text = \"AI is\".strip()\n",
        "\n",
        "generated_text = generate_text(text, do_sample=True, num_beams=3, max_length=64, top_p=0.95, top_k=0)\n",
        "print(generated_text)"
      ],
      "metadata": {
        "id": "2zDDEmFpYvYk",
        "outputId": "2450a175-a0bd-4ba7-efe5-89c164ba4781",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 44,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "AI is a writer, not a writer.”“It is better to have a friend who is not your friend, than to have a friend who is not your friend, than to have a friend who is not your friend, than to have a friend who is not your friend, than to have a friend\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "generated_text = generate_text(text, do_sample=True, max_length=65, penalty_alpha=0.6, top_k=0)\n",
        "print(generated_text)"
      ],
      "metadata": {
        "id": "pN_G6dd-Zk5m",
        "outputId": "9c991e6b-6224-47ab-d344-1c13c89cf663",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 45,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "AI is a woman. I am a woman. I am a rich woman. I am a poor woman. I am a beautiful, delicate, resourceful woman.I pour out my income and put out my retirement plans and everything. So I don't have to waste much time on my recent fortune. What am I going\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "AupawhWzb8g6"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "A100",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.8"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}