Fine_Tune_LLM_with_Hugging_Face_falcon_7b_finetune.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "T4",
      "collapsed_sections": [
        "kf_Vy70HsKGQ"
      ],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/erdult/Fine-Tune-LLM/blob/main/Fine_Tune_LLM_with_Hugging_Face_falcon_7b_finetune.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CkLQo3pmo-Jk"
      },
      "outputs": [],
      "source": [
        "!pip install -Uqqq pip\n",
        "!pip install -qqq bitsandbytes==0.39.0\n",
        "!pip install -qqq torch==2.0.1\n",
        "!pip install -qqq -U git+https://github.com/huggingface/transformers.git@e03a9cc\n",
        "!pip install -qqq -U git+https://github.com/huggingface/peft.git@42a184f\n",
        "!pip install -qqq -U git+https://github.com/huggingface/accelerate.git@c9fbb71\n",
        "!pip install -qqq datasets==2.12.0\n",
        "!pip install -qqq loralib==0.1.1\n",
        "!pip install -qqq einops==0.6.1"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "import os\n",
        "from pprint import pprint\n",
        "import bitsandbytes as bnb\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import transformers\n",
        "from datasets import load_dataset\n",
        "from huggingface_hub import notebook_login\n",
        "from peft import (\n",
        "    LoraConfig,\n",
        "    PeftConfig,\n",
        "    PeftModel,\n",
        "    get_peft_model,\n",
        "    prepare_model_for_kbit_training\n",
        ")\n",
        "from transformers import (\n",
        "    AutoConfig,\n",
        "    AutoModelForCausalLM,\n",
        "    AutoTokenizer,\n",
        "    BitsAndBytesConfig\n",
        ")\n",
        "\n",
        "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
      ],
      "metadata": {
        "id": "YNvT2Uynpit-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "notebook_login()"
      ],
      "metadata": {
        "id": "5bdd0Xi0ppev"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "id": "WiZCJjbya3mI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "above lines mounts google drive folder\n"
      ],
      "metadata": {
        "id": "Ebv4BaqtcB3a"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# LOAD FALCON MODEL & TOKENIZER"
      ],
      "metadata": {
        "id": "2bh3J5Ndp1lO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "MODEL_NAME = \"vilsonrodrigues/falcon-7b-instruct-sharded\"\n",
        "\n",
        "bnb_config = BitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    bnb_4bit_use_double_quant=True,\n",
        "    bnb_4bit_quant_type=\"nf4\",\n",
        "    bnb_4bit_compute_dtype=torch.bfloat16\n",
        ")\n",
        "\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    MODEL_NAME,\n",
        "    device_map=\"auto\",\n",
        "    trust_remote_code=True,\n",
        "    quantization_config=bnb_config\n",
        ")\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n",
        "tokenizer.pad_token = tokenizer.eos_token"
      ],
      "metadata": {
        "id": "1GUD7mBRp2qH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def print_trainable_parameters(model):\n",
        "  \"\"\"\n",
        "  Prints the number of trainable parameters in the model.\n",
        "  \"\"\"\n",
        "  trainable_params = 0\n",
        "  all_param = 0\n",
        "  for _, param in model.named_parameters():\n",
        "    all_param += param.numel()\n",
        "    if param.requires_grad:\n",
        "      trainable_params += param.numel()\n",
        "  print(\n",
        "      f\"trainable params: {trainable_params} || all params: {all_param} || trainables%: {100 * trainable_params / all_param}\"\n",
        "  )"
      ],
      "metadata": {
        "id": "RsxPa1WmmktC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model.gradient_checkpointing_enable()\n",
        "model = prepare_model_for_kbit_training(model)"
      ],
      "metadata": {
        "id": "VsvSHI6lml1S"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "config = LoraConfig(\n",
        "    r=16,\n",
        "    lora_alpha=32,\n",
        "    target_modules=[\"query_key_value\"],\n",
        "    lora_dropout=0.05,\n",
        "    bias=\"none\",\n",
        "    task_type=\"CAUSAL_LM\"\n",
        ")\n",
        "\n",
        "model = get_peft_model(model, config)\n",
        "print_trainable_parameters(model)"
      ],
      "metadata": {
        "id": "6IgfANNJmnGx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Test original model"
      ],
      "metadata": {
        "id": "NvUWzHR9p8nz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = \"\"\"\n",
        "<human>: midjourney prompt for a girl sit on the mountain\n",
        "<assistant>:\n",
        "\"\"\".strip()"
      ],
      "metadata": {
        "id": "VB1lR2Yep-Im"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "generation_config = model.generation_config\n",
        "generation_config.max_new_tokens = 200\n",
        "generation_config.temperature = 0.7\n",
        "generation_config.top_p = 0.7\n",
        "generation_config.num_return_sequences = 1\n",
        "generation_config.pad_token_id = tokenizer.eos_token_id\n",
        "generation_config.eos_token_id = tokenizer.eos_token_id"
      ],
      "metadata": {
        "id": "Yu4xood_p_eg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "device = \"cuda:0\"\n",
        "\n",
        "encoding = tokenizer(prompt, return_tensors=\"pt\").to(device)\n",
        "with torch.inference_mode():\n",
        "  outputs = model.generate(\n",
        "      input_ids = encoding.input_ids,\n",
        "      attention_mask = encoding.attention_mask,\n",
        "      generation_config = generation_config\n",
        "  )\n",
        "\n",
        "print(tokenizer.decode(outputs[0], skip_special_tokens=True))"
      ],
      "metadata": {
        "id": "3l8_DfXuqBLu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Prep dataset"
      ],
      "metadata": {
        "id": "lnnzcwuKqGn0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "data = load_dataset(\"csv\", data_files=\"midjourney_prompt_dataset.csv\")"
      ],
      "metadata": {
        "id": "XVpOQJnSqJYp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "data"
      ],
      "metadata": {
        "id": "JPBLGZDcqX6g"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "data[\"train\"][0]"
      ],
      "metadata": {
        "id": "GbjbejfPqbYz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def generate_prompt(data_point):\n",
        "  return f\"\"\"\n",
        "<human>: {data_point[\"User\"]}\n",
        "<assistant>: {data_point[\"Prompt\"]}\n",
        "\"\"\".strip()\n",
        "\n",
        "def generate_and_tokenize_prompt(data_point):\n",
        "  full_prompt = generate_prompt(data_point)\n",
        "  tokenized_full_prompt = tokenizer(full_prompt, padding=True, truncation=True)\n",
        "  return tokenized_full_prompt"
      ],
      "metadata": {
        "id": "UcNlEuwnqh4m"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "data = data[\"train\"].shuffle().map(generate_and_tokenize_prompt)"
      ],
      "metadata": {
        "id": "i4jK1V20qiac"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "data"
      ],
      "metadata": {
        "id": "v4w31zK0qk7A"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Finetune the model"
      ],
      "metadata": {
        "id": "kf_Vy70HsKGQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "training_args = transformers.TrainingArguments(\n",
        "      per_device_train_batch_size=1,\n",
        "      gradient_accumulation_steps=4,\n",
        "      num_train_epochs=1,\n",
        "      learning_rate=2e-4,\n",
        "      fp16=True,\n",
        "      save_total_limit=3,\n",
        "      logging_steps=1,\n",
        "      output_dir=\"experiments\",\n",
        "      optim=\"paged_adamw_8bit\",\n",
        "      lr_scheduler_type=\"cosine\",\n",
        "      warmup_ratio=0.05,\n",
        ")\n",
        "\n",
        "trainer = transformers.Trainer(\n",
        "    model=model,\n",
        "    train_dataset=data,\n",
        "    args=training_args,\n",
        "    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False)\n",
        ")\n",
        "model.config.use_cache = False\n",
        "trainer.train()"
      ],
      "metadata": {
        "id": "UTjSAkYUsQqg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Save trained model model.save saves it on colab"
      ],
      "metadata": {
        "id": "S0JaaueJs3Tt"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model.save_pretrained(\"trained-model\")"
      ],
      "metadata": {
        "id": "bFCETjids62_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# save it on hugging face"
      ],
      "metadata": {
        "id": "1V3UB9y_b3QM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "PEFT_MODEL = \"midjourney-falcon-7b\"\n",
        "\n",
        "model.push_to_hub(\n",
        "    PEFT_MODEL, use_auth_token=True\n",
        ")"
      ],
      "metadata": {
        "id": "sl_IzzKdtFtG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "config = PeftConfig.from_pretrained(PEFT_MODEL)\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    config.base_model_name_or_path,\n",
        "    return_dict=True,\n",
        "    quantization_config=bnb_config,\n",
        "    device_map=\"auto\",\n",
        "    trust_remote_code=True\n",
        ")\n",
        "\n",
        "tokenizer=AutoTokenizer.from_pretrained(config.base_model_name_or_path)\n",
        "tokenizer.pad_token = tokenizer.eos_token\n",
        "\n",
        "model = PeftModel.from_pretrained(model, PEFT_MODEL)"
      ],
      "metadata": {
        "id": "1mUR2fpbtI7-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Run the finetuned model"
      ],
      "metadata": {
        "id": "jIs2DDuetQzF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "generation_config = model.generation_config\n",
        "generation_config.max_new_tokens = 200\n",
        "generation_config.temperature = 0.7\n",
        "generation_config.top_p = 0.7\n",
        "generation_config.num_return_sequences = 1\n",
        "generation_config.pad_token_id = tokenizer.eos_token_id\n",
        "generation_config.eos_token_id = tokenizer.eos_token_id"
      ],
      "metadata": {
        "id": "ArdOBIwgtRQx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "device = \"cuda:0\"\n",
        "\n",
        "prompt = \"\"\"\n",
        "<human>: midjourney prompt for a boy running in the snow\n",
        "<assistant>:\n",
        "\"\"\".strip()\n",
        "\n",
        "encoding = tokenizer(prompt, return_tensors=\"pt\").to(device)\n",
        "with torch.inference_mode():\n",
        "  outputs = model.generate(\n",
        "      input_ids = encoding.input_ids,\n",
        "      attention_mask = encoding.attention_mask,\n",
        "      generation_config = generation_config\n",
        "  )\n",
        "\n",
        "print(tokenizer.decode(outputs[0], skip_special_tokens=True))"
      ],
      "metadata": {
        "id": "WvzZUmTttSgR"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}