emotion_GPT2_as_classifier.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/aburkov/theLMbook/blob/main/emotion_GPT2_as_classifier.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "<div style=\"display: flex; justify-content: center;\">\n",
        "    <div style=\"background-color: #f4f6f7; padding: 15px; width: 80%;\">\n",
        "        <table style=\"width: 100%\">\n",
        "            <tr>\n",
        "                <td style=\"vertical-align: middle;\">\n",
        "                    <span style=\"font-size: 14px;\">\n",
        "                        A notebook for <a href=\"https://www.thelmbook.com\" target=\"_blank\" rel=\"noopener\">The Hundred-Page Language Models Book</a> by Andriy Burkov<br><br>\n",
        "                        Code repository: <a href=\"https://github.com/aburkov/theLMbook\" target=\"_blank\" rel=\"noopener\">https://github.com/aburkov/theLMbook</a>\n",
        "                    </span>\n",
        "                </td>\n",
        "                <td style=\"vertical-align: middle;\">\n",
        "                    <a href=\"https://www.thelmbook.com\" target=\"_blank\" rel=\"noopener\">\n",
        "                        <img src=\"https://thelmbook.com/img/book.png\" width=\"80px\" alt=\"The Hundred-Page Language Models Book\">\n",
        "                    </a>\n",
        "                </td>\n",
        "            </tr>\n",
        "        </table>\n",
        "    </div>\n",
        "</div>"
      ],
      "metadata": {
        "id": "LpCleINoAMtW"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yy0zjL_2ouOU",
        "outputId": "96ee626b-f141-4438-afa1-ae39e5305184"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Some weights of GPT2ForSequenceClassification were not initialized from the model checkpoint at openai-community/gpt2 and are newly initialized: ['score.weight']\n",
            "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
            "Epoch 1/8: 100%|██████████| 1125/1125 [00:51<00:00, 21.68it/s, Loss=0.488]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Average loss: 0.4882, Test accuracy: 0.9230\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 2/8: 100%|██████████| 1125/1125 [00:51<00:00, 21.73it/s, Loss=0.144]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Average loss: 0.1437, Test accuracy: 0.9310\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 3/8: 100%|██████████| 1125/1125 [00:51<00:00, 21.77it/s, Loss=0.115]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Average loss: 0.1149, Test accuracy: 0.9410\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 4/8: 100%|██████████| 1125/1125 [00:51<00:00, 21.81it/s, Loss=0.104]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Average loss: 0.1038, Test accuracy: 0.9395\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 5/8: 100%|██████████| 1125/1125 [00:51<00:00, 21.78it/s, Loss=0.0963]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Average loss: 0.0963, Test accuracy: 0.9340\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 6/8: 100%|██████████| 1125/1125 [00:51<00:00, 21.71it/s, Loss=0.0851]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Average loss: 0.0851, Test accuracy: 0.9395\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 7/8: 100%|██████████| 1125/1125 [00:51<00:00, 21.74it/s, Loss=0.0806]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Average loss: 0.0806, Test accuracy: 0.9400\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 8/8: 100%|██████████| 1125/1125 [00:51<00:00, 21.74it/s, Loss=0.0767]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Average loss: 0.0767, Test accuracy: 0.9460\n",
            "Input: I'm so happy to be able to finetune an LLM!\n",
            "Predicted emotion: joy\n"
          ]
        }
      ],
      "source": [
        "# Import required libraries\n",
        "import torch           # Main PyTorch library\n",
        "from torch.utils.data import DataLoader  # For dataset handling\n",
        "from torch.optim import AdamW    # Optimizer for training\n",
        "from transformers import AutoTokenizer, AutoModelForSequenceClassification  # Hugging Face components\n",
        "from tqdm import tqdm   # Progress bar utilities\n",
        "import json             # For parsing JSON data\n",
        "import requests         # For downloading dataset from URL\n",
        "import gzip             # For decompressing dataset\n",
        "import random           # For setting seeds and shuffling data\n",
        "\n",
        "def set_seed(seed):\n",
        "    \"\"\"\n",
        "    Sets random seeds for reproducibility across different libraries.\n",
        "\n",
        "    Args:\n",
        "        seed (int): Seed value for random number generation\n",
        "    \"\"\"\n",
        "    # Set Python's built-in random seed\n",
        "    random.seed(seed)\n",
        "    # Set PyTorch's CPU random seed\n",
        "    torch.manual_seed(seed)\n",
        "    # Set seed for all available GPUs\n",
        "    torch.cuda.manual_seed_all(seed)\n",
        "    # Request cuDNN to use deterministic algorithms\n",
        "    torch.backends.cudnn.deterministic = True\n",
        "    # Disable cuDNN's auto-tuner for consistent behavior\n",
        "    torch.backends.cudnn.benchmark = False\n",
        "\n",
        "def load_and_split_dataset(url, test_ratio=0.1):\n",
        "    \"\"\"\n",
        "    Downloads and splits dataset into train and test sets.\n",
        "\n",
        "    Args:\n",
        "        url (str): URL of the dataset\n",
        "        test_ratio (float): Proportion of data for testing\n",
        "\n",
        "    Returns:\n",
        "        tuple: (train_dataset, test_dataset)\n",
        "    \"\"\"\n",
        "    # Download and decompress dataset\n",
        "    response = requests.get(url)\n",
        "    content = gzip.decompress(response.content).decode()\n",
        "\n",
        "    # Parse JSON lines into list of examples\n",
        "    dataset = [json.loads(line) for line in content.splitlines()]\n",
        "\n",
        "    # Shuffle and split dataset\n",
        "    random.shuffle(dataset)\n",
        "    split_index = int(len(dataset) * (1 - test_ratio))\n",
        "\n",
        "    return dataset[:split_index], dataset[split_index:]\n",
        "\n",
        "def load_model_and_tokenizer(model_name, device, label_to_id, id_to_label, unique_labels):\n",
        "    \"\"\"\n",
        "    Loads and configures the model and tokenizer for sequence classification.\n",
        "\n",
        "    Args:\n",
        "        model_name (str): Name of pre-trained model\n",
        "        device: Device to load model on\n",
        "        label_to_id (dict): Mapping from label strings to IDs\n",
        "        id_to_label (dict): Mapping from IDs to label strings\n",
        "        unique_labels (list): List of all possible labels\n",
        "\n",
        "    Returns:\n",
        "        tuple: (model, tokenizer)\n",
        "    \"\"\"\n",
        "    # Initialize model with correct number of output classes\n",
        "    model = AutoModelForSequenceClassification.from_pretrained(\n",
        "        model_name,\n",
        "        num_labels=len(unique_labels)\n",
        "    )\n",
        "\n",
        "    # Configure padding and label mappings\n",
        "    model.config.pad_token_id = model.config.eos_token_id\n",
        "    model.config.id2label = id_to_label\n",
        "    model.config.label2id = label_to_id\n",
        "\n",
        "    # Initialize and configure tokenizer\n",
        "    tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
        "    tokenizer.pad_token = tokenizer.eos_token\n",
        "\n",
        "    return (model.to(device), tokenizer)\n",
        "\n",
        "def encode_text(tokenizer, text, return_tensor=False):\n",
        "    \"\"\"\n",
        "    Encodes text using the provided tokenizer.\n",
        "\n",
        "    Args:\n",
        "        tokenizer: Hugging Face tokenizer\n",
        "        text (str): Text to encode\n",
        "        return_tensor (bool): Whether to return PyTorch tensor\n",
        "\n",
        "    Returns:\n",
        "        List or tensor of token IDs\n",
        "    \"\"\"\n",
        "    # If tensor output is requested, encode with PyTorch tensors\n",
        "    if return_tensor:\n",
        "        return tokenizer.encode(\n",
        "            text, add_special_tokens=False, return_tensors=\"pt\"\n",
        "        )\n",
        "    # Otherwise return list of token IDs\n",
        "    else:\n",
        "        return tokenizer.encode(text, add_special_tokens=False)\n",
        "\n",
        "class TextClassificationDataset(torch.utils.data.Dataset):\n",
        "    \"\"\"\n",
        "    PyTorch Dataset for text classification.\n",
        "    Converts text and labels into model-ready format.\n",
        "\n",
        "    Args:\n",
        "        data (list): List of dictionaries containing text and labels\n",
        "        tokenizer: Hugging Face tokenizer\n",
        "        label_to_id (dict): Mapping from label strings to IDs\n",
        "    \"\"\"\n",
        "    def __init__(self, data, tokenizer, label_to_id):\n",
        "        self.data = data\n",
        "        self.tokenizer = tokenizer\n",
        "        self.label_to_id = label_to_id\n",
        "\n",
        "    def __len__(self):\n",
        "        # Return total number of examples\n",
        "        return len(self.data)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        \"\"\"\n",
        "        Returns a single training example.\n",
        "\n",
        "        Args:\n",
        "            idx (int): Index of the example to fetch\n",
        "\n",
        "        Returns:\n",
        "            dict: Contains input_ids and labels\n",
        "        \"\"\"\n",
        "        # Get example from dataset\n",
        "        item = self.data[idx]\n",
        "        # Convert text to token IDs\n",
        "        input_ids = encode_text(self.tokenizer, item[\"text\"])\n",
        "        # Convert label string to ID\n",
        "        labels = self.label_to_id[item[\"label\"]]\n",
        "\n",
        "        return {\n",
        "            \"input_ids\": input_ids,\n",
        "            \"labels\": labels\n",
        "        }\n",
        "\n",
        "def collate_fn(batch):\n",
        "    \"\"\"\n",
        "    Collates batch of examples into training-ready format.\n",
        "    Handles padding and conversion to tensors.\n",
        "\n",
        "    Args:\n",
        "        batch: List of examples from Dataset\n",
        "\n",
        "    Returns:\n",
        "        dict: Contains input_ids, labels, and attention_mask tensors\n",
        "    \"\"\"\n",
        "    # Find longest sequence for padding\n",
        "    max_length = max(len(item[\"input_ids\"]) for item in batch)\n",
        "\n",
        "    # Pad input sequences with zeros\n",
        "    input_ids = [\n",
        "        item[\"input_ids\"] +\n",
        "        [0] * (max_length - len(item[\"input_ids\"]))\n",
        "        for item in batch\n",
        "    ]\n",
        "\n",
        "    # Create attention masks (1 for tokens, 0 for padding)\n",
        "    attention_mask = [\n",
        "        [1] * len(item[\"input_ids\"]) +\n",
        "        [0] * (max_length - len(item[\"input_ids\"]))\n",
        "        for item in batch\n",
        "    ]\n",
        "\n",
        "    # Collect labels\n",
        "    labels = [item[\"labels\"] for item in batch]\n",
        "\n",
        "    # Convert everything to tensors\n",
        "    return {\n",
        "        \"input_ids\": torch.tensor(input_ids),\n",
        "        \"labels\": torch.tensor(labels),\n",
        "        \"attention_mask\": torch.tensor(attention_mask)\n",
        "    }\n",
        "\n",
        "def generate_label(model, tokenizer, text):\n",
        "    \"\"\"\n",
        "    Generates label prediction for input text.\n",
        "\n",
        "    Args:\n",
        "        model: Fine-tuned model\n",
        "        tokenizer: Associated tokenizer\n",
        "        text (str): Input text to classify\n",
        "\n",
        "    Returns:\n",
        "        str: Predicted label\n",
        "    \"\"\"\n",
        "    # Encode text and move to model's device\n",
        "    input_ids = encode_text(\n",
        "        tokenizer,\n",
        "        text,\n",
        "        return_tensor=True\n",
        "    ).to(model.device)\n",
        "\n",
        "    # Get model predictions\n",
        "    outputs = model(input_ids)\n",
        "    logits = outputs.logits[0]\n",
        "    # Get class with highest probability\n",
        "    predicted_class = logits.argmax().item()\n",
        "    # Convert class ID to label string\n",
        "    return model.config.id2label[predicted_class]\n",
        "\n",
        "def calculate_accuracy(model, dataloader):\n",
        "    \"\"\"\n",
        "    Calculates prediction accuracy on a dataset.\n",
        "\n",
        "    Args:\n",
        "        model: Fine-tuned model\n",
        "        dataloader: DataLoader containing evaluation examples\n",
        "\n",
        "    Returns:\n",
        "        float: Accuracy score\n",
        "    \"\"\"\n",
        "    # Set model to evaluation mode\n",
        "    model.eval()\n",
        "    correct = 0\n",
        "    total = 0\n",
        "\n",
        "    # Disable gradient computation for efficiency\n",
        "    with torch.no_grad():\n",
        "        for batch in dataloader:\n",
        "            # Move batch to device\n",
        "            input_ids = batch[\"input_ids\"].to(model.device)\n",
        "            attention_mask = batch[\"attention_mask\"].to(model.device)\n",
        "            labels = batch[\"labels\"].to(model.device)\n",
        "\n",
        "            # Get model predictions\n",
        "            outputs = model(input_ids=input_ids, attention_mask=attention_mask)\n",
        "            predictions = outputs.logits.argmax(dim=-1)\n",
        "\n",
        "            # Update accuracy counters\n",
        "            correct += (predictions == labels).sum().item()\n",
        "            total += labels.size(0)\n",
        "\n",
        "    # Calculate accuracy\n",
        "    accuracy = correct / total\n",
        "    # Reset model to training mode\n",
        "    model.train()\n",
        "    return accuracy\n",
        "\n",
        "def create_label_mappings(train_dataset):\n",
        "    \"\"\"\n",
        "    Creates mappings between label strings and IDs.\n",
        "\n",
        "    Args:\n",
        "        train_dataset: List of training examples\n",
        "\n",
        "    Returns:\n",
        "        tuple: (label_to_id, id_to_label, unique_labels)\n",
        "    \"\"\"\n",
        "    # Get sorted list of unique labels\n",
        "    unique_labels = sorted(set(item[\"label\"] for item in train_dataset))\n",
        "    # Create mappings between labels and IDs\n",
        "    label_to_id = {label: i for i, label in enumerate(unique_labels)}\n",
        "    id_to_label = {i: label for label, i in label_to_id.items()}\n",
        "    return label_to_id, id_to_label, unique_labels\n",
        "\n",
        "def test_model(model_path, test_input):\n",
        "    \"\"\"\n",
        "    Tests a saved model on a single input.\n",
        "\n",
        "    Args:\n",
        "        model_path (str): Path to saved model\n",
        "        test_input (str): Text to classify\n",
        "    \"\"\"\n",
        "    # Setup device and load model\n",
        "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "    model = AutoModelForSequenceClassification.from_pretrained(model_path).to(device)\n",
        "    tokenizer = AutoTokenizer.from_pretrained(model_path)\n",
        "\n",
        "    # Generate and display prediction\n",
        "    emotion = generate_label(model, tokenizer, test_input)\n",
        "    print(f\"Input: {test_input}\")\n",
        "    print(f\"Predicted emotion: {emotion}\")\n",
        "\n",
        "def download_and_prepare_data(data_url, tokenizer, batch_size):\n",
        "    \"\"\"\n",
        "    Downloads and prepares dataset for training.\n",
        "\n",
        "    Args:\n",
        "        data_url (str): URL of the dataset\n",
        "        tokenizer: Tokenizer for text processing\n",
        "        batch_size (int): Batch size for DataLoader\n",
        "\n",
        "    Returns:\n",
        "        tuple: (train_dataloader, test_dataloader, label_to_id, id_to_label, unique_labels)\n",
        "    \"\"\"\n",
        "    # Load and split dataset\n",
        "    train_dataset, test_dataset = load_and_split_dataset(data_url)\n",
        "\n",
        "    # Create label mappings\n",
        "    label_to_id, id_to_label, unique_labels = create_label_mappings(train_dataset)\n",
        "\n",
        "    # Create datasets\n",
        "    train_data = TextClassificationDataset(\n",
        "        train_dataset,\n",
        "        tokenizer,\n",
        "        label_to_id\n",
        "    )\n",
        "    test_data = TextClassificationDataset(\n",
        "        test_dataset,\n",
        "        tokenizer,\n",
        "        label_to_id\n",
        "    )\n",
        "\n",
        "    # Create dataloaders\n",
        "    train_dataloader = DataLoader(\n",
        "        train_data,\n",
        "        batch_size=batch_size,\n",
        "        shuffle=True,\n",
        "        collate_fn=collate_fn\n",
        "    )\n",
        "    test_dataloader = DataLoader(\n",
        "        test_data,\n",
        "        batch_size=batch_size,\n",
        "        shuffle=True,\n",
        "        collate_fn=collate_fn\n",
        "    )\n",
        "    return train_dataloader, test_dataloader, label_to_id, id_to_label, unique_labels\n",
        "\n",
        "def get_hyperparameters():\n",
        "    \"\"\"\n",
        "    Returns training hyperparameters.\n",
        "\n",
        "    Returns:\n",
        "        tuple: (num_epochs, batch_size, learning_rate)\n",
        "    \"\"\"\n",
        "    # Train for fewer epochs as sequence classification converges faster\n",
        "    num_epochs=8\n",
        "    # Standard batch size that works well with most GPU memory\n",
        "    batch_size=16\n",
        "    # Standard learning rate for fine-tuning transformers\n",
        "    learning_rate=5e-5\n",
        "    return num_epochs, batch_size, learning_rate\n",
        "\n",
        "# Main training script\n",
        "if __name__ == \"__main__\":\n",
        "    # Set random seed for reproducibility\n",
        "    seed = 42\n",
        "    set_seed(seed)\n",
        "\n",
        "    # Configure training parameters\n",
        "    data_url = \"https://www.thelmbook.com/data/emotions\"\n",
        "    model_name = \"openai-community/gpt2\"\n",
        "\n",
        "    # Get hyperparameters\n",
        "    num_epochs, batch_size, learning_rate = get_hyperparameters()\n",
        "\n",
        "    # Setup device\n",
        "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "    # Initialize tokenizer\n",
        "    tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
        "    tokenizer.pad_token = tokenizer.eos_token\n",
        "\n",
        "    # Prepare data and get label mappings\n",
        "    train_loader, test_loader, label_to_id, id_to_label, unique_labels = download_and_prepare_data(\n",
        "        data_url,\n",
        "        tokenizer,\n",
        "        batch_size\n",
        "    )\n",
        "\n",
        "    # Initialize model for sequence classification\n",
        "    model = AutoModelForSequenceClassification.from_pretrained(\n",
        "        model_name,\n",
        "        num_labels=len(unique_labels)\n",
        "    ).to(device)\n",
        "\n",
        "    # Configure model's label handling\n",
        "    model.config.pad_token_id = model.config.eos_token_id\n",
        "    model.config.id2label = id_to_label\n",
        "    model.config.label2id = label_to_id\n",
        "\n",
        "    # Initialize optimizer\n",
        "    optimizer = AdamW(model.parameters(), lr=learning_rate)\n",
        "\n",
        "    # Training loop\n",
        "    for epoch in range(num_epochs):\n",
        "        model.train()\n",
        "        total_loss = 0\n",
        "        num_batches = 0\n",
        "        progress_bar = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{num_epochs}\")\n",
        "\n",
        "        for batch in progress_bar:\n",
        "            # Move batch to device\n",
        "            input_ids = batch[\"input_ids\"].to(device)\n",
        "            attention_mask = batch[\"attention_mask\"].to(device)\n",
        "            labels = batch[\"labels\"].to(device)\n",
        "\n",
        "            # Forward pass\n",
        "            outputs = model(\n",
        "                input_ids=input_ids,\n",
        "                attention_mask=attention_mask,\n",
        "                labels=labels\n",
        "            )\n",
        "\n",
        "            # Backward pass and optimization\n",
        "            loss = outputs.loss\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "            optimizer.zero_grad()\n",
        "\n",
        "            # Update metrics\n",
        "            total_loss += loss.item()\n",
        "            num_batches += 1\n",
        "\n",
        "            progress_bar.set_postfix({\"Loss\": total_loss / num_batches})\n",
        "\n",
        "        # Display epoch metrics\n",
        "        avg_loss = total_loss / num_batches\n",
        "        test_acc = calculate_accuracy(model, test_loader)\n",
        "        print(f\"Average loss: {avg_loss:.4f}, Test accuracy: {test_acc:.4f}\")\n",
        "\n",
        "    # Save the fine-tuned model\n",
        "    model.save_pretrained(\"./finetuned_model\")\n",
        "    tokenizer.save_pretrained(\"./finetuned_model\")\n",
        "\n",
        "    # Test the model\n",
        "    test_input = \"I'm so happy to be able to finetune an LLM!\"\n",
        "    test_model(\"./finetuned_model\", test_input)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "A100",
      "authorship_tag": "ABX9TyOfPNbJZ+ZzhLYc0JRHSozh",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}