emotion_GPT2_as_text_generator_LoRA.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_text_generator_LoRA.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yy0zjL_2ouOU",
        "outputId": "e12e16c1-7815-411f-db02-e40a3bd05659"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Using device: cuda\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 1/18: 100%|██████████| 1125/1125 [00:53<00:00, 20.98it/s, Loss=0.613]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 1 - Average loss: 0.6127, Test accuracy: 0.7605\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 2/18: 100%|██████████| 1125/1125 [00:54<00:00, 20.62it/s, Loss=0.353]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 2 - Average loss: 0.3532, Test accuracy: 0.7970\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 3/18: 100%|██████████| 1125/1125 [00:54<00:00, 20.61it/s, Loss=0.237]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 3 - Average loss: 0.2375, Test accuracy: 0.8530\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 4/18: 100%|██████████| 1125/1125 [00:55<00:00, 20.45it/s, Loss=0.184]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 4 - Average loss: 0.1843, Test accuracy: 0.8985\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 5/18: 100%|██████████| 1125/1125 [00:54<00:00, 20.68it/s, Loss=0.146]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 5 - Average loss: 0.1457, Test accuracy: 0.9175\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 6/18: 100%|██████████| 1125/1125 [00:54<00:00, 20.63it/s, Loss=0.121]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 6 - Average loss: 0.1208, Test accuracy: 0.9215\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 7/18: 100%|██████████| 1125/1125 [00:54<00:00, 20.50it/s, Loss=0.103]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 7 - Average loss: 0.1028, Test accuracy: 0.9260\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 8/18: 100%|██████████| 1125/1125 [00:54<00:00, 20.62it/s, Loss=0.0927]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 8 - Average loss: 0.0927, Test accuracy: 0.9260\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 9/18: 100%|██████████| 1125/1125 [00:54<00:00, 20.59it/s, Loss=0.0887]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 9 - Average loss: 0.0887, Test accuracy: 0.9330\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 10/18: 100%|██████████| 1125/1125 [00:54<00:00, 20.58it/s, Loss=0.079]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 10 - Average loss: 0.0790, Test accuracy: 0.9315\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 11/18: 100%|██████████| 1125/1125 [00:54<00:00, 20.69it/s, Loss=0.0771]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 11 - Average loss: 0.0771, Test accuracy: 0.9325\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 12/18: 100%|██████████| 1125/1125 [00:54<00:00, 20.60it/s, Loss=0.0699]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 12 - Average loss: 0.0699, Test accuracy: 0.9345\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 13/18: 100%|██████████| 1125/1125 [00:54<00:00, 20.60it/s, Loss=0.0663]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 13 - Average loss: 0.0663, Test accuracy: 0.9265\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 14/18: 100%|██████████| 1125/1125 [00:54<00:00, 20.61it/s, Loss=0.064]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 14 - Average loss: 0.0640, Test accuracy: 0.9375\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 15/18: 100%|██████████| 1125/1125 [00:54<00:00, 20.60it/s, Loss=0.0633]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 15 - Average loss: 0.0633, Test accuracy: 0.9380\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 16/18: 100%|██████████| 1125/1125 [00:54<00:00, 20.75it/s, Loss=0.0605]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 16 - Average loss: 0.0605, Test accuracy: 0.9390\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 17/18: 100%|██████████| 1125/1125 [00:53<00:00, 20.85it/s, Loss=0.0571]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 17 - Average loss: 0.0571, Test accuracy: 0.9370\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Epoch 18/18: 100%|██████████| 1125/1125 [00:55<00:00, 20.26it/s, Loss=0.0574]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 18 - Average loss: 0.0574, Test accuracy: 0.9420\n",
            "Training accuracy: 0.9424\n",
            "Test accuracy: 0.9420\n",
            "Using device: cuda\n",
            "Input: I'm so happy to be able to finetune an LLM!\n",
            "Generated emotion: joy\n"
          ]
        }
      ],
      "source": [
        "# Import required libraries\n",
        "import json            # For parsing JSON data\n",
        "import random          # For setting seeds and shuffling data\n",
        "import gzip            # For decompressing dataset\n",
        "import requests        # For downloading dataset from URL\n",
        "import torch           # Main PyTorch library\n",
        "from peft import get_peft_model, LoraConfig, TaskType  # For efficient finetuning using LoRA\n",
        "from torch.utils.data import Dataset, DataLoader  # For dataset handling\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM  # Hugging Face model components\n",
        "from torch.optim import AdamW    # Optimizer for training\n",
        "from tqdm import tqdm   # Progress bar utilities\n",
        "import re               # For text normalization\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 build_prompt(text):\n",
        "    \"\"\"\n",
        "    Creates a standardized prompt for emotion classification.\n",
        "\n",
        "    Args:\n",
        "        text (str): Input text to classify\n",
        "\n",
        "    Returns:\n",
        "        str: Formatted prompt for the model\n",
        "    \"\"\"\n",
        "    # Format the input text into a consistent prompt structure\n",
        "    return f\"Predict the emotion for the following text: {text}\\nEmotion:\"\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",
        "def decode_text(tokenizer, token_ids):\n",
        "    \"\"\"\n",
        "    Decodes token IDs back to text.\n",
        "\n",
        "    Args:\n",
        "        tokenizer: Hugging Face tokenizer\n",
        "        token_ids: List or tensor of token IDs\n",
        "\n",
        "    Returns:\n",
        "        str: Decoded text\n",
        "    \"\"\"\n",
        "    # Convert token IDs back to text, skipping special tokens\n",
        "    return tokenizer.decode(token_ids, skip_special_tokens=True)\n",
        "\n",
        "class PromptCompletionDataset(Dataset):\n",
        "    \"\"\"\n",
        "    PyTorch Dataset for prompt-completion pairs.\n",
        "    Handles the conversion of text data into model-ready format.\n",
        "\n",
        "    Args:\n",
        "        data (list): List of dictionaries containing prompts and completions\n",
        "        tokenizer: Hugging Face tokenizer\n",
        "    \"\"\"\n",
        "    def __init__(self, data, tokenizer):\n",
        "        # Store the raw data and tokenizer for later use\n",
        "        self.data = data\n",
        "        self.tokenizer = tokenizer\n",
        "\n",
        "    def __len__(self):\n",
        "        # Return the total number of examples in the dataset\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, labels, prompt, and expected completion\n",
        "        \"\"\"\n",
        "        # Get the specific example from our dataset\n",
        "        item = self.data[idx]\n",
        "        prompt = item[\"prompt\"]\n",
        "        completion = item[\"completion\"]\n",
        "\n",
        "        # Convert text to token IDs for both prompt and completion\n",
        "        encoded_prompt = encode_text(self.tokenizer, prompt)\n",
        "        encoded_completion = encode_text(self.tokenizer, completion)\n",
        "        # Get the end-of-sequence token ID\n",
        "        eos_token = self.tokenizer.eos_token_id\n",
        "\n",
        "        # Combine prompt and completion tokens with EOS token\n",
        "        input_ids = encoded_prompt + encoded_completion + [eos_token]\n",
        "        # Create labels: -100 for prompt (ignored in loss), completion tokens for learning\n",
        "        labels = [-100] * len(encoded_prompt) + encoded_completion + [eos_token]\n",
        "\n",
        "        return {\n",
        "            \"input_ids\": input_ids,\n",
        "            \"labels\": labels,\n",
        "            \"prompt\": prompt,\n",
        "            \"expected_completion\": completion\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",
        "        tuple: (input_ids, attention_mask, labels, prompts, expected_completions)\n",
        "    \"\"\"\n",
        "    # Find the longest sequence in the batch for padding\n",
        "    max_length = max(len(item[\"input_ids\"]) for item in batch)\n",
        "\n",
        "    # Pad input sequences to max_length with pad token\n",
        "    input_ids = [\n",
        "        item[\"input_ids\"] +\n",
        "        [tokenizer.pad_token_id] * (max_length - len(item[\"input_ids\"]))\n",
        "        for item in batch\n",
        "    ]\n",
        "\n",
        "    # Pad label sequences with -100 (ignored in loss calculation)\n",
        "    labels = [\n",
        "        item[\"labels\"] +\n",
        "        [-100] * (max_length - len(item[\"labels\"]))\n",
        "        for item in batch\n",
        "    ]\n",
        "\n",
        "    # Create attention masks: 1 for real 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",
        "    # Keep original prompts and completions for evaluation\n",
        "    prompts = [item[\"prompt\"] for item in batch]\n",
        "    expected_completions = [item[\"expected_completion\"] for item in batch]\n",
        "\n",
        "    # Convert everything to PyTorch tensors except text\n",
        "    return (\n",
        "        torch.tensor(input_ids),\n",
        "        torch.tensor(attention_mask),\n",
        "        torch.tensor(labels),\n",
        "        prompts,\n",
        "        expected_completions\n",
        "    )\n",
        "\n",
        "def normalize_text(text):\n",
        "    \"\"\"\n",
        "    Normalizes text for consistent comparison.\n",
        "\n",
        "    Args:\n",
        "        text (str): Input text\n",
        "\n",
        "    Returns:\n",
        "        str: Normalized text\n",
        "    \"\"\"\n",
        "    # Remove leading/trailing whitespace and convert to lowercase\n",
        "    text = text.strip().lower()\n",
        "    # Replace multiple whitespace characters with single space\n",
        "    text = re.sub(r\"\\s+\", ' ', text)\n",
        "    return text\n",
        "\n",
        "def calculate_accuracy(model, tokenizer, loader):\n",
        "    \"\"\"\n",
        "    Calculates prediction accuracy on a dataset.\n",
        "\n",
        "    Args:\n",
        "        model: Finetuned model\n",
        "        tokenizer: Associated tokenizer\n",
        "        loader: DataLoader containing evaluation examples\n",
        "\n",
        "    Returns:\n",
        "        float: Accuracy score\n",
        "    \"\"\"\n",
        "    # Set model to evaluation mode\n",
        "    model.eval()\n",
        "    # Initialize counters for accuracy calculation\n",
        "    correct = 0\n",
        "    total = 0\n",
        "\n",
        "    # Disable gradient computation for efficiency\n",
        "    with torch.no_grad():\n",
        "        for input_ids, attention_mask, labels, prompts, expected_completions in loader:\n",
        "            for prompt, expected_completion in zip(prompts, expected_completions):\n",
        "                # Generate model's prediction\n",
        "                generated_text = generate_text(model, tokenizer, prompt)\n",
        "                # Compare normalized versions of prediction and target\n",
        "                if normalize_text(generated_text) == normalize_text(expected_completion):\n",
        "                    correct += 1\n",
        "                total += 1\n",
        "\n",
        "    # Calculate accuracy, handling empty dataset case\n",
        "    accuracy = correct / total if total > 0 else 0\n",
        "    # Reset model to training mode\n",
        "    model.train()\n",
        "    return accuracy\n",
        "\n",
        "def generate_text(model, tokenizer, prompt, max_new_tokens=50):\n",
        "    \"\"\"\n",
        "    Generates text completion for a given prompt.\n",
        "\n",
        "    Args:\n",
        "        model: Finetuned model\n",
        "        tokenizer: Associated tokenizer\n",
        "        prompt (str): Input prompt\n",
        "        max_new_tokens (int): Maximum number of tokens to generate\n",
        "\n",
        "    Returns:\n",
        "        str: Generated completion\n",
        "    \"\"\"\n",
        "    # Encode prompt and move to model's device\n",
        "    input_ids = tokenizer(prompt, return_tensors=\"pt\").to(model.device)\n",
        "\n",
        "    # Generate completion using model's generate method\n",
        "    output_ids = model.generate(\n",
        "        input_ids=input_ids[\"input_ids\"],\n",
        "        attention_mask=input_ids[\"attention_mask\"],\n",
        "        max_new_tokens=max_new_tokens,\n",
        "        pad_token_id=tokenizer.pad_token_id,\n",
        "        eos_token_id=tokenizer.eos_token_id\n",
        "    )[0]\n",
        "\n",
        "    # Extract and decode only the generated part (excluding prompt)\n",
        "    generated_text = decode_text(tokenizer, output_ids[input_ids[\"input_ids\"].shape[1]:])\n",
        "    return generated_text.strip()\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",
        "    # Determine device (GPU if available, else CPU)\n",
        "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "    print(f\"Using device: {device}\")\n",
        "\n",
        "    # Load saved model and tokenizer\n",
        "    model = AutoModelForCausalLM.from_pretrained(model_path).to(device)\n",
        "    tokenizer = AutoTokenizer.from_pretrained(model_path)\n",
        "\n",
        "    # Configure padding token\n",
        "    if tokenizer.pad_token is None:\n",
        "        tokenizer.pad_token = tokenizer.eos_token\n",
        "    model.config.pad_token_id = tokenizer.pad_token_id\n",
        "\n",
        "    # Generate and display prediction\n",
        "    prompt = build_prompt(test_input)\n",
        "    generated_text = generate_text(model, tokenizer, prompt)\n",
        "\n",
        "    print(f\"Input: {test_input}\")\n",
        "    print(f\"Generated emotion: {generated_text}\")\n",
        "\n",
        "def download_and_prepare_data(data_url, tokenizer, batch_size, test_ratio=0.1):\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",
        "        test_ratio (float): Proportion of data for testing\n",
        "\n",
        "    Returns:\n",
        "        tuple: (train_loader, test_loader)\n",
        "    \"\"\"\n",
        "    # Download and decompress dataset\n",
        "    response = requests.get(data_url)\n",
        "    content = gzip.decompress(response.content).decode()\n",
        "\n",
        "    # Process each example into prompt-completion pairs\n",
        "    dataset = []\n",
        "    for entry in map(json.loads, content.splitlines()):\n",
        "        dataset.append({\n",
        "            \"prompt\": build_prompt(entry['text']),\n",
        "            \"completion\": entry[\"label\"].strip()\n",
        "        })\n",
        "\n",
        "    # Split into train and test sets\n",
        "    random.shuffle(dataset)\n",
        "    split_index = int(len(dataset) * (1 - test_ratio))\n",
        "    train_data = dataset[:split_index]\n",
        "    test_data = dataset[split_index:]\n",
        "\n",
        "    # Create datasets\n",
        "    train_dataset = PromptCompletionDataset(train_data, tokenizer)\n",
        "    test_dataset = PromptCompletionDataset(test_data, tokenizer)\n",
        "\n",
        "    # Create data loaders\n",
        "    train_loader = DataLoader(\n",
        "        train_dataset,\n",
        "        batch_size=batch_size,\n",
        "        shuffle=True,\n",
        "        collate_fn=collate_fn\n",
        "    )\n",
        "    test_loader = DataLoader(\n",
        "        test_dataset,\n",
        "        batch_size=batch_size,\n",
        "        shuffle=False,\n",
        "        collate_fn=collate_fn\n",
        "    )\n",
        "\n",
        "    return train_loader, test_loader\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 more epochs with LoRA as it's more efficient\n",
        "    num_epochs = 18\n",
        "    # Batch size\n",
        "    batch_size = 16\n",
        "    # Standard learning rate for finetuning transformers\n",
        "    learning_rate = 5e-5\n",
        "\n",
        "    return num_epochs, batch_size, learning_rate\n",
        "\n",
        "# Main training script\n",
        "if __name__ == \"__main__\":\n",
        "    # Set random seeds for reproducibility\n",
        "    set_seed(42)\n",
        "\n",
        "    # Configure basic training parameters\n",
        "    data_url = \"https://www.thelmbook.com/data/emotions\"\n",
        "    model_name = \"openai-community/gpt2\"\n",
        "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "    print(f\"Using device: {device}\")\n",
        "\n",
        "    # Initialize tokenizer\n",
        "    tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
        "    tokenizer.pad_token = tokenizer.eos_token\n",
        "\n",
        "    # Configure LoRA parameters\n",
        "    peft_config = LoraConfig(\n",
        "        task_type = TaskType.CAUSAL_LM,  # Set task type for causal language modeling\n",
        "        inference_mode = False,          # Enable training mode\n",
        "        r = 16,                          # Rank of LoRA update matrices\n",
        "        lora_alpha = 32                  # LoRA scaling factor\n",
        "    )\n",
        "\n",
        "    # Load model and apply LoRA configuration\n",
        "    model = AutoModelForCausalLM.from_pretrained(model_name).to(device)\n",
        "    model = get_peft_model(model, peft_config)\n",
        "\n",
        "    # Get hyperparameters and prepare data\n",
        "    num_epochs, batch_size, learning_rate = get_hyperparameters()\n",
        "    train_loader, test_loader = download_and_prepare_data(data_url, tokenizer, batch_size)\n",
        "\n",
        "    # Initialize optimizer\n",
        "    optimizer = AdamW(model.parameters(), lr=learning_rate)\n",
        "\n",
        "    # Training loop\n",
        "    for epoch in range(num_epochs):\n",
        "        total_loss = 0\n",
        "        num_batches = 0\n",
        "        progress_bar = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{num_epochs}\")\n",
        "\n",
        "        for input_ids, attention_mask, labels, _, _ in progress_bar:\n",
        "            # Move batch to device\n",
        "            input_ids = input_ids.to(device)\n",
        "            attention_mask = attention_mask.to(device)\n",
        "            labels = 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",
        "            loss = outputs.loss\n",
        "\n",
        "            # Backward pass and optimization\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "            optimizer.zero_grad()\n",
        "\n",
        "            # Update metrics\n",
        "            total_loss += loss.item()\n",
        "            num_batches += 1\n",
        "            progress_bar.set_postfix({\"Loss\": total_loss / num_batches})\n",
        "\n",
        "        # Calculate and display epoch metrics\n",
        "        avg_loss = total_loss / num_batches\n",
        "        test_acc = calculate_accuracy(model, tokenizer, test_loader)\n",
        "        print(f\"Epoch {epoch+1} - Average loss: {avg_loss:.4f}, Test accuracy: {test_acc:.4f}\")\n",
        "\n",
        "    # Calculate final model performance\n",
        "    train_acc = calculate_accuracy(model, tokenizer, train_loader)\n",
        "    print(f\"Training accuracy: {train_acc:.4f}\")\n",
        "    print(f\"Test accuracy: {test_acc:.4f}\")\n",
        "\n",
        "    # Save the LoRA-tuned model and tokenizer\n",
        "    model.save_pretrained(\"./finetuned_model\")\n",
        "    tokenizer.save_pretrained(\"./finetuned_model\")\n",
        "\n",
        "    # Test the finetuned model with a sample input\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": "ABX9TyOJ8D8foCRGxJa2Ehsu6evL",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}