Gemma_Fine_tuning.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/jonbaer/googlecolab/blob/master/Gemma_Fine_tuning.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qbZ_zjyeu0C9"
      },
      "source": [
        "### LoRA Fine-tuning Gemma-2B\n",
        "\n",
        "This notebook is made for LoRA fine-tuning Gemma-2B. LoRA is a parameter efficient fine-tuning technique that only adjusts few parameters instead of full fine-tuning of the model, thus, it's faster. We will be using [VMWare/open-instruct](https://huggingface.co/datasets/VMware/open-instruct) dataset that has instructions. To apply LoRA, we'll use [PEFT](https://huggingface.co/docs/peft/index) library and for supervised instruction tuning, we will use `SFTTrainer` from [TRL](https://huggingface.co/docs/trl/en/index)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sP8XDTKxrzEm"
      },
      "outputs": [],
      "source": [
        "!pip install -q -U transformers peft accelerate datasets trl bitsandbytes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aVZ6CNzksJmI"
      },
      "outputs": [],
      "source": [
        "from peft import LoraConfig\n",
        "\n",
        "lora_config = LoraConfig(\n",
        "    r=8,\n",
        "    target_modules=[\"q_proj\", \"o_proj\", \"k_proj\", \"v_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
        "    task_type=\"CAUSAL_LM\",\n",
        ")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Login to Hugging Face Hub, since Gemma-2B has gated access and login confirms that you have access to the model. If you don't have an access, get it from the model repository [here](https://huggingface.co/google/gemma-2b) your request will shortly be accepted."
      ],
      "metadata": {
        "id": "p4gt3em5gGEU"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uBwftrQXuafj"
      },
      "outputs": [],
      "source": [
        "from huggingface_hub import notebook_login\n",
        "notebook_login()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "We'll shrink the model even further by loading it in 4bit using `bitsandbytes`. Then initialize the model with the CausalLM head and initialize the tokenizer."
      ],
      "metadata": {
        "id": "0gbZfYBdjfLm"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xmSm_3AisFvJ"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
        "import os\n",
        "\n",
        "model_id = \"google/gemma-2b\"\n",
        "bnb_config = BitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    bnb_4bit_quant_type=\"nf4\",\n",
        "    bnb_4bit_compute_dtype=torch.bfloat16\n",
        ")\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
        "model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map={\"\":0})\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Load the dataset."
      ],
      "metadata": {
        "id": "3XJK1RHYe9ex"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eZrKzsKesNeA"
      },
      "outputs": [],
      "source": [
        "from datasets import load_dataset\n",
        "\n",
        "data = load_dataset(\"VMware/open-instruct\", split=\"train\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Concat Alpaca prompt with responses."
      ],
      "metadata": {
        "id": "H96WEyOzfBrV"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4vi45Qepydwd"
      },
      "outputs": [],
      "source": [
        "texts = []\n",
        "for prompt, response in zip(data[\"alpaca_prompt\"], data[\"response\"]):\n",
        "  text = prompt + response\n",
        "  texts.append(text)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Remove unnecessary columns."
      ],
      "metadata": {
        "id": "aAijXn3sfFXN"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "77j7X9SUyEB8"
      },
      "outputs": [],
      "source": [
        "data = data.remove_columns([\"source\", \"alpaca_prompt\", \"response\", \"task_name\", \"template_type\", \"instruction\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Add the concatenated column back."
      ],
      "metadata": {
        "id": "jMtZU4ULjhrP"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DwhOALMjzOtq"
      },
      "outputs": [],
      "source": [
        "data = data.add_column(\"text_column\", texts)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Depending on your dataset prompts, you might want to truncate and handle overflowing tokens like below. If you keep it like this, your prompts will be truncated though and you will have bad results. 😔 So adjust the below cell depending on what you need."
      ],
      "metadata": {
        "id": "T0-Mi0qUfILc"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Rgqw0kukvzit"
      },
      "outputs": [],
      "source": [
        "def tokenize_dataset(ds):\n",
        "  result = tokenizer(ds[\"text_column\"],truncation=True,\n",
        "                       max_length=512)\n",
        "  #sample_map = result.pop(\"overflow_to_sample_mapping\")\n",
        "  #for key, values in ds.items():\n",
        "  #  result[key] = [values[i] for i in sample_map]\n",
        "  #  print(result[key])\n",
        "  return result"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iZfVr5xoweiP"
      },
      "outputs": [],
      "source": [
        "ds = data.map(tokenize_dataset)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "ds"
      ],
      "metadata": {
        "id": "mqTd6rLI05u5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Initializing `SFTTrainer` from TRL is all you need!\n",
        "\n",
        "Small note: if your dataset needs formatting, you can write a formatting function and pass it. You need to either pass `formatting_func` or `dataset_text_field` if your dataset text field doesn't need any formatting and you did your preprocessing beforehand.\n",
        "\n",
        "Then simply call ` train`. Note that this notebook is built for educational purposes so you might need to adjust the hyperparameters to your own use case."
      ],
      "metadata": {
        "id": "oxrNGaUKfaUW"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EMXGufOmsR07"
      },
      "outputs": [],
      "source": [
        "import transformers\n",
        "from trl import SFTTrainer\n",
        "\n",
        "\n",
        "trainer = SFTTrainer(\n",
        "    model=model,\n",
        "    train_dataset=ds,\n",
        "    dataset_text_field=\"text_column\",\n",
        "    args=transformers.TrainingArguments(\n",
        "        per_device_train_batch_size=1,\n",
        "        gradient_accumulation_steps=4,\n",
        "        warmup_steps=2,\n",
        "        max_steps=30,\n",
        "        learning_rate=2e-4,\n",
        "        fp16=True,\n",
        "        logging_steps=1,\n",
        "        output_dir=\"outputs\",\n",
        "        optim=\"paged_adamw_8bit\"\n",
        "    ),\n",
        "    peft_config=lora_config,\n",
        "    #formatting_func=formatting_func,\n",
        ")\n",
        "trainer.train()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "58ck-m-CsT7S"
      },
      "outputs": [],
      "source": [
        "text = \"Write a news style post about a fake event, like aliens from Mars landing on Earth. It is meant to be funny but also be written in the authoritative style of a news report, kind of like The Onion. ### Response:\"\n",
        "device = \"cuda:0\"\n",
        "inputs = tokenizer(text, return_tensors=\"pt\").to(device)\n",
        "\n",
        "outputs = model.generate(**inputs, max_new_tokens=20)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(tokenizer.decode(outputs[0], skip_special_tokens=True))"
      ],
      "metadata": {
        "id": "wU4nOPGcfqWL"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "machine_shape": "hm",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}