{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Gzg8SopX8EWH"
   },
   "source": [
    "# Alpaca OpenLLaMa 3B LoRA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "YSW4FUQPwIYu",
    "outputId": "421da5b9-be5e-4661-bf98-6c57fbf1f6fd"
   },
   "outputs": [],
   "source": [
    "!git clone https://github.com/tloen/alpaca-lora.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "s1xm2uERx_st",
    "outputId": "c77b42e1-202b-45ac-aa2c-3fdf2bcde155"
   },
   "outputs": [],
   "source": [
    "%cd alpaca-lora/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JCB9UzMVwsSM",
    "outputId": "7f0688e4-f360-4da6-a4d3-0d59d4135649"
   },
   "outputs": [],
   "source": [
    "!pip install -q datasets loralib sentencepiece\n",
    "\n",
    "!pip install -q git+https://github.com/huggingface/transformers.git\n",
    "!pip install -q git+https://github.com/huggingface/peft.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "qCnXTszZxE2T",
    "outputId": "e619573c-3ac2-4ac3-bed5-caab57d00f9a"
   },
   "outputs": [],
   "source": [
    "!pip install bitsandbytes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "NyxvpDjcrUvv",
    "outputId": "abe581a2-70ce-4956-b3dd-593fa4f3d8ef"
   },
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9w0aSCzhxxQf"
   },
   "source": [
    "## Data Check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 150,
     "referenced_widgets": [
      "76eba77b0fc9499e9b4015393156153e",
      "c42123767d4849ffbcdd48f09580ea13",
      "e89f07c1f5744328b071f96a3bdc9532",
      "da36c71ba5bc469eb32bde730c5038d3",
      "6cc6b7c564fe48f086152c78f8a9a915",
      "c6a3adfa76dd4bbd8cc5e75ee28ef3e8",
      "0c65f0e913b74de88ab712d65b9027f7",
      "b07c7f3b6bc24f10821b88e254e88e43",
      "42ea2834ee074029abe8cf8f9ff79a16",
      "ec4205b0491a4f53ad1cba42689d59c4",
      "f8b9df09c8c54d53a22310e64e30772f",
      "e69b8912d3384180842ff0ead6096803",
      "550a2ef0612e4d8396462d1b1c43608d",
      "f8d82b4d016f46ca8bab1a2fe591861e",
      "975d908a1874425198e0d62ae884a3ef",
      "152aa8231b3f4b7db4e6f27bb61065b9",
      "27b48982ae244c66bccd855e7e33a00d",
      "a896cb94d1fc4641ad224c7a9ca619b1",
      "be5d02d7001a4a25a848d6b0c582b42b",
      "609437910022418e83543fe6c3d16501",
      "c3b5ef42dae64f95a2ef055db7e50fa8",
      "ee0f0aff0c254b82854d1a6f4b367161",
      "fd90085d416440579f409cdd9b29e053",
      "fdd55d56775a4641925fc88e43f946b8",
      "a54ebe30a12b4c21b67b2f8d43aea2d3",
      "6f041f053ff240cbbb52f6d117ebfcf1",
      "b7d5f8aaee0047178661e58aa59433ee",
      "563ca6cac1594ba0a7c168ee31ccc982",
      "e10e7830d1404baea72e81e0df295239",
      "acc07aa077fd49d38d74e0e10fa3991a",
      "441dfdd43ef0442894c19cad81a478db",
      "9efe7ac9ee0946388eeb85f9a4603ffe",
      "ae9996dcefcf4cfc8ea2f7488629c5ec",
      "db463b34f98c488e8282d7a59c421347",
      "2ba524e48f35421c9afc9cc81e9ce3fb",
      "1829c96471c64a03a4872e1b7a16551c",
      "f639a121fab0407595ebdf04015d9267",
      "fa5ff44ae5f64ce4bbf9b59e300b7ba3",
      "e0d701b8c09744858db4fda6011b69df",
      "7ef61952e0f2463bad3dd1a814bfb95a",
      "75263b1d19ce4cb0b95cacb6138b6fcb",
      "c3f35ce818e34b05b247fae76200c02e",
      "be8dddb19dde441da1193d5f8738413d",
      "4be8d304ddce49319488f20ab36e7f53"
     ]
    },
    "id": "OdgRTo5YxyRL",
    "outputId": "e6e64014-796b-4fa9-9635-3cb29f2dab7d"
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from transformers import LlamaTokenizer\n",
    "\n",
    "\n",
    "tokenizer = LlamaTokenizer.from_pretrained(\"openlm-research/open_llama_3b_600bt_preview\", add_eos_token=True)\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "tokenizer.pad_token_id = tokenizer.eos_token_id\n",
    "\n",
    "data = load_dataset(\"json\", data_files=\"/content/drive/MyDrive/alpaca-data.json\")\n",
    "\n",
    "\n",
    "def generate_prompt(data_point):\n",
    "    # sorry about the formatting disaster gotta move fast\n",
    "    if data_point[\"input\"]:\n",
    "        return f\"\"\"Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n",
    "\n",
    "### Instruction:\n",
    "{data_point[\"instruction\"]}\n",
    "\n",
    "### Input:\n",
    "{data_point[\"input\"]}\n",
    "\n",
    "### Response:\n",
    "{data_point[\"output\"]}\"\"\"\n",
    "    else:\n",
    "        return f\"\"\"Below is an instruction that describes a task. Write a response that appropriately completes the request.\n",
    "\n",
    "### Instruction:\n",
    "{data_point[\"instruction\"]}\n",
    "\n",
    "### Response:\n",
    "{data_point[\"output\"]}\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "j7qZqe0YxG2c"
   },
   "source": [
    "## Fine Tuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "EQyg3MygxKCW",
    "outputId": "82c54601-2cf6-48f1-e3ca-994ae2b8b405"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import bitsandbytes as bnb\n",
    "import transformers\n",
    "from transformers import AutoTokenizer, AutoConfig, LlamaForCausalLM\n",
    "from peft import prepare_model_for_int8_training, LoraConfig, get_peft_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "76iZTtGJy26x"
   },
   "outputs": [],
   "source": [
    "# Setting for A100 - For 3090 \n",
    "MICRO_BATCH_SIZE = 3  # change to 4 for 3090\n",
    "BATCH_SIZE = 128\n",
    "GRADIENT_ACCUMULATION_STEPS = BATCH_SIZE // MICRO_BATCH_SIZE\n",
    "EPOCHS = 1  # paper uses 3\n",
    "LEARNING_RATE = 2e-5  # from the original paper\n",
    "CUTOFF_LEN = 256  # 256 accounts for about 96% of the data\n",
    "LORA_R = 4\n",
    "LORA_ALPHA = 16\n",
    "LORA_DROPOUT = 0.05"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 150,
     "referenced_widgets": [
      "984896bacc39446091e2b6573facfa3e",
      "f15179aad6ad44dfa2dfbd703dc7c319",
      "f92c68902b984e648337da7e63fcc775",
      "5ce279d7dca74e55acae1c70efae6cd5",
      "b187d889a6654fd7ad948d66a8bce197",
      "e131a64ebaa84a809754b3c762525c37",
      "72819a3c606c41f19d1493539d64aff5",
      "6761867d6e2b415b932bebd33a3144f3",
      "72c9aa3bb9614dfc8f316bc45a3455a0",
      "58e5e895215d4910b55d8cbe1417a111",
      "b6bad7f690c54cb5814d09081e306393",
      "ae3921fcb8e44862a42031f284f608c3",
      "aaac19cd3f524821a656ceffa57fb337",
      "a77af2c0771b4a20b42dade07c40021c",
      "7e5ced10225544399add5e68dda80595",
      "ccf73891ca1941ca8b77c87ba84058f8",
      "bfc1e1979f0741a09e6c9f4181ae88e7",
      "8ce25329fdf04a6e8c4ad8097975111e",
      "479d18fc0f9448a49b07b93ac608ae1d",
      "85ffc3306cb4474aa5bfac22f6356352",
      "0d5765bde1fa470eb087df5fe1e45e56",
      "6873f0a521cd4e4390563e27046c1aa6",
      "0832a0aa65ac4f70ba2f303f9b867d25",
      "436e9deafae04955a13f2ff0fc0e6348",
      "83427a912e3740008c0a7ffb9d20a1bd",
      "084f196c105c4c428253e0ce6dce8a21",
      "403990397a4041abb8aff7bdcc5d8e8a",
      "f957482962954f03b383d7c4f93c7b04",
      "19a731496867404b8888986f3d9c83ef",
      "e0b09d3aa9ae4719bcdfad8e75ad86c5",
      "9d41c197843644858919a396dc927936",
      "67f9f877dee04d538ffbe2de6c4dd181",
      "001bb332cdb94a0c96dd2a322ff8718f"
     ]
    },
    "id": "tI5Ta-gQy56c",
    "outputId": "a07d8bea-809e-44e8-f3ea-28345af97157"
   },
   "outputs": [],
   "source": [
    "model = LlamaForCausalLM.from_pretrained(\n",
    "    \"openlm-research/open_llama_3b_600bt_preview\",\n",
    "    load_in_8bit=True,\n",
    "    device_map=\"auto\",\n",
    ")\n",
    "\n",
    "model = prepare_model_for_int8_training(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Oh-jJFKNzBpw"
   },
   "outputs": [],
   "source": [
    "config = LoraConfig(\n",
    "    r=LORA_R,\n",
    "    lora_alpha=LORA_ALPHA,\n",
    "    target_modules=[\"q_proj\", \"v_proj\"],\n",
    "    lora_dropout=LORA_DROPOUT,\n",
    "    bias=\"none\",\n",
    "    task_type=\"CAUSAL_LM\",\n",
    ")\n",
    "model = get_peft_model(model, config)\n",
    "tokenizer.pad_token_id = 0  # unk. we want this to be different from the eos token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000,
     "referenced_widgets": [
      "310970764a4a4c9382e2334a4abbbd77",
      "2cd59f5acc4948988814eb9b3d909aed",
      "8939309032ca4787a9d3c7b54b889926",
      "fcd31f8d1e5a482d9267ab63fd723e86",
      "2d126272a4284ba9af39de65e9680fc8",
      "d8a0e85fa9fc4a9291f9303861fc9efc",
      "d21f1cfb275b4bb38baf7022665e0208",
      "81e25933b9af43a492a6da4a5e3560f6",
      "5fe27494333b4e07a632769696744210",
      "35ca13104a3e4263a192889362dbc05e",
      "85ae96c237754c4db014e5c09898e162"
     ]
    },
    "id": "kWP89TPIwRkK",
    "outputId": "fb4d3012-795b-46c2-90a0-f72fc8185a11"
   },
   "outputs": [],
   "source": [
    "data = data.shuffle().map(\n",
    "    lambda data_point: tokenizer(\n",
    "        generate_prompt(data_point),\n",
    "        truncation=True,\n",
    "        max_length=CUTOFF_LEN,\n",
    "        padding=\"max_length\",\n",
    "    )\n",
    ")\n",
    "\n",
    "trainer = transformers.Trainer(\n",
    "    model=model,\n",
    "    train_dataset=data[\"train\"],\n",
    "    args=transformers.TrainingArguments(\n",
    "        per_device_train_batch_size=MICRO_BATCH_SIZE,\n",
    "        gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,\n",
    "        warmup_steps=100,\n",
    "        num_train_epochs=EPOCHS,\n",
    "        learning_rate=LEARNING_RATE,\n",
    "        fp16=True,\n",
    "        logging_steps=1,\n",
    "        output_dir=\"lora-alpaca\",\n",
    "        save_total_limit=3,\n",
    "    ),\n",
    "    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),\n",
    ")\n",
    "model.config.use_cache = False\n",
    "trainer.train(resume_from_checkpoint=False)\n",
    "\n",
    "model.save_pretrained(\"lora-alpaca\")\n",
    "\n",
    "!cp -rv lora-alpaca /content/drive/MyDrive/"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "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.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}