fine-tune.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Experimenting with Fine-tuning a Base Language Model for Instruction Following\n",
    "I.e. a chatbot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Requirements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['HF_HOME'] = '/home/stefanwebb/models/hf'\n",
    "import torch\n",
    "torch.random.manual_seed(0)\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, BitsAndBytesConfig\n",
    "from peft import LoraConfig, prepare_model_for_kbit_training, PeftModel\n",
    "from datasets import load_dataset\n",
    "from trl import SFTTrainer\n",
    "\n",
    "if torch.cuda.is_bf16_supported():\n",
    "  compute_dtype = torch.bfloat16\n",
    "else:\n",
    "  compute_dtype = torch.float16"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Pretrained Model\n",
    "Using Gemma 2B from Google"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "use_4bit = True\n",
    "bnb_4bit_quant_type = \"nf4\"\n",
    "use_double_quant = True\n",
    "\n",
    "compute_dtype = torch.bfloat16\n",
    "attn_implementation = 'flash_attention_2'\n",
    "\n",
    "# 'target_modules' is a list of the modules that should be targeted by LoRA.\n",
    "target_modules= ['k_proj', 'q_proj', 'v_proj', 'o_proj', \"gate_proj\", \"down_proj\", \"up_proj\"]\n",
    "\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "        load_in_4bit=use_4bit,\n",
    "        bnb_4bit_quant_type=bnb_4bit_quant_type,\n",
    "        bnb_4bit_compute_dtype=compute_dtype,\n",
    "        bnb_4bit_use_double_quant=use_double_quant,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`config.hidden_act` is ignored, you should use `config.hidden_activation` instead.\n",
      "Gemma's activation function will be set to `gelu_pytorch_tanh`. Please, use\n",
      "`config.hidden_activation` if you want to override this behaviour.\n",
      "See https://github.com/huggingface/transformers/pull/29402 for more details.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3c94df6933294641bb352aae3730b535",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "MODEL_ID = \"google/gemma-2b\"\n",
    "NEW_MODEL_NAME = \"stefans-gemma-2b-instruct\"\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    MODEL_ID,\n",
    "    device_map='cuda',\n",
    "    torch_dtype=\"auto\",\n",
    "    quantization_config=bnb_config,\n",
    "    attn_implementation=attn_implementation\n",
    ")\n",
    "\n",
    "model = prepare_model_for_kbit_training(model)\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)\n",
    "tokenizer.padding_side = 'right' # not sure if this is necessary...\n",
    "tokenizer.pad_token_id = tokenizer.eos_token_id"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Examining model and tokenizer object to understand more deeply"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GemmaForCausalLM(\n",
      "  (model): GemmaModel(\n",
      "    (embed_tokens): Embedding(256000, 2048, padding_idx=0)\n",
      "    (layers): ModuleList(\n",
      "      (0-17): 18 x GemmaDecoderLayer(\n",
      "        (self_attn): GemmaSdpaAttention(\n",
      "          (q_proj): Linear(in_features=2048, out_features=2048, bias=False)\n",
      "          (k_proj): Linear(in_features=2048, out_features=256, bias=False)\n",
      "          (v_proj): Linear(in_features=2048, out_features=256, bias=False)\n",
      "          (o_proj): Linear(in_features=2048, out_features=2048, bias=False)\n",
      "          (rotary_emb): GemmaRotaryEmbedding()\n",
      "        )\n",
      "        (mlp): GemmaMLP(\n",
      "          (gate_proj): Linear(in_features=2048, out_features=16384, bias=False)\n",
      "          (up_proj): Linear(in_features=2048, out_features=16384, bias=False)\n",
      "          (down_proj): Linear(in_features=16384, out_features=2048, bias=False)\n",
      "          (act_fn): PytorchGELUTanh()\n",
      "        )\n",
      "        (input_layernorm): GemmaRMSNorm()\n",
      "        (post_attention_layernorm): GemmaRMSNorm()\n",
      "      )\n",
      "    )\n",
      "    (norm): GemmaRMSNorm()\n",
      "  )\n",
      "  (lm_head): Linear(in_features=2048, out_features=256000, bias=False)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "# Structure of model\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Long output...\n",
    "# for x in dir(model):\n",
    "#     print(type(getattr(model, x)), x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'transformers.models.gemma.modeling_gemma.GemmaModel'>\n",
      "<class 'transformers.models.gemma.modeling_gemma.GemmaModel'>\n",
      "<class 'transformers.models.gemma.modeling_gemma.GemmaPreTrainedModel'>\n",
      "<class 'transformers.modeling_utils.PreTrainedModel'>\n",
      "<class 'torch.nn.modules.module.Module'>\n",
      "<class 'transformers.modeling_utils.ModuleUtilsMixin'>\n",
      "<class 'transformers.generation.utils.GenerationMixin'>\n",
      "<class 'transformers.utils.hub.PushToHubMixin'>\n",
      "<class 'transformers.integrations.peft.PeftAdapterMixin'>\n",
      "<class 'object'>\n"
     ]
    }
   ],
   "source": [
    "# HuggingFace models inherit from torch.nn.Module\n",
    "# TODO: Go over code to implement GemmaModel class\n",
    "from inspect import getmro\n",
    "\n",
    "for cls in getmro(type(model.model)):\n",
    "    print(cls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'transformers.models.gemma.configuration_gemma.GemmaConfig'>\n",
      "<class 'transformers.configuration_utils.PretrainedConfig'>\n",
      "<class 'transformers.utils.hub.PushToHubMixin'>\n",
      "<class 'object'>\n"
     ]
    }
   ],
   "source": [
    "# However, PreTrainedConfig is specific to Transformers library\n",
    "for cls in getmro(type(model.config)):\n",
    "    print(cls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'transformers.models.gemma.tokenization_gemma_fast.GemmaTokenizerFast'> GemmaTokenizerFast(name_or_path='google/gemma-2b', vocab_size=256000, model_max_length=1000000000000000019884624838656, is_fast=True, padding_side='left', truncation_side='right', special_tokens={'bos_token': '<bos>', 'eos_token': '<eos>', 'unk_token': '<unk>', 'pad_token': '<pad>', 'additional_special_tokens': ['<start_of_turn>', '<end_of_turn>']}, clean_up_tokenization_spaces=False),  added_tokens_decoder={\n",
      "\t0: AddedToken(\"<pad>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
      "\t1: AddedToken(\"<eos>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
      "\t2: AddedToken(\"<bos>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
      "\t3: AddedToken(\"<unk>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
      "\t4: AddedToken(\"<mask>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t5: AddedToken(\"<2mass>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t6: AddedToken(\"[@BOS@]\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t7: AddedToken(\"<unused0>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t8: AddedToken(\"<unused1>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t9: AddedToken(\"<unused2>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t10: AddedToken(\"<unused3>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t11: AddedToken(\"<unused4>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t12: AddedToken(\"<unused5>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t13: AddedToken(\"<unused6>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t14: AddedToken(\"<unused7>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t15: AddedToken(\"<unused8>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t16: AddedToken(\"<unused9>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t17: AddedToken(\"<unused10>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t18: AddedToken(\"<unused11>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t19: AddedToken(\"<unused12>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t20: AddedToken(\"<unused13>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t21: AddedToken(\"<unused14>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t22: AddedToken(\"<unused15>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t23: AddedToken(\"<unused16>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t24: AddedToken(\"<unused17>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t25: AddedToken(\"<unused18>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t26: AddedToken(\"<unused19>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t27: AddedToken(\"<unused20>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t28: AddedToken(\"<unused21>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t29: AddedToken(\"<unused22>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t30: AddedToken(\"<unused23>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t31: AddedToken(\"<unused24>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t32: AddedToken(\"<unused25>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t33: AddedToken(\"<unused26>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t34: AddedToken(\"<unused27>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t35: AddedToken(\"<unused28>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t36: AddedToken(\"<unused29>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t37: AddedToken(\"<unused30>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t38: AddedToken(\"<unused31>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t39: AddedToken(\"<unused32>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t40: AddedToken(\"<unused33>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t41: AddedToken(\"<unused34>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t42: AddedToken(\"<unused35>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t43: AddedToken(\"<unused36>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t44: AddedToken(\"<unused37>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t45: AddedToken(\"<unused38>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t46: AddedToken(\"<unused39>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t47: AddedToken(\"<unused40>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t48: AddedToken(\"<unused41>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t49: AddedToken(\"<unused42>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t50: AddedToken(\"<unused43>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t51: AddedToken(\"<unused44>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t52: AddedToken(\"<unused45>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t53: AddedToken(\"<unused46>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t54: AddedToken(\"<unused47>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t55: AddedToken(\"<unused48>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t56: AddedToken(\"<unused49>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t57: AddedToken(\"<unused50>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t58: AddedToken(\"<unused51>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t59: AddedToken(\"<unused52>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t60: AddedToken(\"<unused53>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t61: AddedToken(\"<unused54>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t62: AddedToken(\"<unused55>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t63: AddedToken(\"<unused56>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t64: AddedToken(\"<unused57>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t65: AddedToken(\"<unused58>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t66: AddedToken(\"<unused59>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t67: AddedToken(\"<unused60>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t68: AddedToken(\"<unused61>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t69: AddedToken(\"<unused62>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t70: AddedToken(\"<unused63>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t71: AddedToken(\"<unused64>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t72: AddedToken(\"<unused65>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t73: AddedToken(\"<unused66>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t74: AddedToken(\"<unused67>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t75: AddedToken(\"<unused68>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t76: AddedToken(\"<unused69>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t77: AddedToken(\"<unused70>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t78: AddedToken(\"<unused71>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t79: AddedToken(\"<unused72>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t80: AddedToken(\"<unused73>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t81: AddedToken(\"<unused74>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t82: AddedToken(\"<unused75>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t83: AddedToken(\"<unused76>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t84: AddedToken(\"<unused77>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t85: AddedToken(\"<unused78>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t86: AddedToken(\"<unused79>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t87: AddedToken(\"<unused80>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t88: AddedToken(\"<unused81>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t89: AddedToken(\"<unused82>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t90: AddedToken(\"<unused83>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t91: AddedToken(\"<unused84>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t92: AddedToken(\"<unused85>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t93: AddedToken(\"<unused86>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t94: AddedToken(\"<unused87>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t95: AddedToken(\"<unused88>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t96: AddedToken(\"<unused89>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t97: AddedToken(\"<unused90>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t98: AddedToken(\"<unused91>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t99: AddedToken(\"<unused92>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t100: AddedToken(\"<unused93>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t101: AddedToken(\"<unused94>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t102: AddedToken(\"<unused95>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t103: AddedToken(\"<unused96>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t104: AddedToken(\"<unused97>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t105: AddedToken(\"<unused98>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t106: AddedToken(\"<start_of_turn>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
      "\t107: AddedToken(\"<end_of_turn>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
      "\t108: AddedToken(\"\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t109: AddedToken(\"\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t110: AddedToken(\"\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t111: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t112: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t113: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t114: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t115: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t116: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t117: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t118: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t119: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t120: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t121: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t122: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t123: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t124: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t125: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t126: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t127: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t128: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t129: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t130: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t131: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t132: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t133: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t134: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t135: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t136: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t137: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t138: AddedToken(\"\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t139: AddedToken(\"▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t140: AddedToken(\"▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t141: AddedToken(\"▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t142: AddedToken(\"▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t143: AddedToken(\"▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t144: AddedToken(\"▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t145: AddedToken(\"▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t146: AddedToken(\"▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t147: AddedToken(\"▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t148: AddedToken(\"▁▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t149: AddedToken(\"▁▁▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t150: AddedToken(\"▁▁▁▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t151: AddedToken(\"▁▁▁▁▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t152: AddedToken(\"▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t153: AddedToken(\"▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t154: AddedToken(\"▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t155: AddedToken(\"▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t156: AddedToken(\"▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t157: AddedToken(\"▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t158: AddedToken(\"▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t159: AddedToken(\"▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t160: AddedToken(\"▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t161: AddedToken(\"▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t162: AddedToken(\"▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t163: AddedToken(\"▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t164: AddedToken(\"▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t165: AddedToken(\"▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t166: AddedToken(\"▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t167: AddedToken(\"▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t168: AddedToken(\"▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t169: AddedToken(\"<table>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t170: AddedToken(\"<caption>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t171: AddedToken(\"<thead>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t172: AddedToken(\"<tbody>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t173: AddedToken(\"<tfoot>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t174: AddedToken(\"<tr>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t175: AddedToken(\"<th>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t176: AddedToken(\"<td>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t177: AddedToken(\"</table>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t178: AddedToken(\"</caption>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t179: AddedToken(\"</thead>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t180: AddedToken(\"</tbody>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t181: AddedToken(\"</tfoot>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t182: AddedToken(\"</tr>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t183: AddedToken(\"</th>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t184: AddedToken(\"</td>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t185: AddedToken(\"<h1>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t186: AddedToken(\"<h2>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t187: AddedToken(\"<h3>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t188: AddedToken(\"<h4>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t189: AddedToken(\"<h5>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t190: AddedToken(\"<h6>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t191: AddedToken(\"<blockquote>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t192: AddedToken(\"</h1>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t193: AddedToken(\"</h2>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t194: AddedToken(\"</h3>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t195: AddedToken(\"</h4>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t196: AddedToken(\"</h5>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t197: AddedToken(\"</h6>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t198: AddedToken(\"</blockquote>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t199: AddedToken(\"<strong>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t200: AddedToken(\"<em>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t201: AddedToken(\"<b>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t202: AddedToken(\"<i>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t203: AddedToken(\"<u>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t204: AddedToken(\"<s>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t205: AddedToken(\"<sub>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t206: AddedToken(\"<sup>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t207: AddedToken(\"<code>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t208: AddedToken(\"</strong>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t209: AddedToken(\"</em>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t210: AddedToken(\"</b>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t211: AddedToken(\"</i>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t212: AddedToken(\"</u>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t213: AddedToken(\"</s>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t214: AddedToken(\"</sub>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t215: AddedToken(\"</sup>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "\t216: AddedToken(\"</code>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=False),\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "print(type(tokenizer), tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'transformers.models.gemma.tokenization_gemma_fast.GemmaTokenizerFast'>\n",
      "<class 'transformers.tokenization_utils_fast.PreTrainedTokenizerFast'>\n",
      "<class 'transformers.tokenization_utils_base.PreTrainedTokenizerBase'>\n",
      "<class 'transformers.tokenization_utils_base.SpecialTokensMixin'>\n",
      "<class 'transformers.utils.hub.PushToHubMixin'>\n",
      "<class 'object'>\n"
     ]
    }
   ],
   "source": [
    "# TODO: Check out how a tokenizer is implemented\n",
    "for cls in getmro(type(tokenizer)):\n",
    "    print(cls)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Quantize model with bitsandbytes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "lora_r = 16\n",
    "\n",
    "# 'lora_alpha' is the alpha parameter for LoRA scaling.\n",
    "lora_alpha = 16\n",
    "\n",
    "# 'lora_dropout' is the dropout probability for LoRA layers.\n",
    "lora_dropout = 0.05\n",
    "\n",
    "# 'target_modules' is a list of the modules that should be targeted by LoRA.\n",
    "target_modules= ['k_proj', 'q_proj', 'v_proj', 'o_proj', \"gate_proj\", \"down_proj\", \"up_proj\"]\n",
    "\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "        load_in_4bit=use_4bit,\n",
    "        bnb_4bit_quant_type=bnb_4bit_quant_type,\n",
    "        bnb_4bit_compute_dtype=compute_dtype,\n",
    "        bnb_4bit_use_double_quant=use_double_quant,\n",
    ")\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Data\n",
    "https://huggingface.co/datasets/tatsu-lab/alpaca"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "54c80ec3e6a94933bb50e6f6ef06f4ec",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading readme:   0%|          | 0.00/7.47k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "193e1439a054424280edd6fdbf47fff0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/24.2M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "deddda6f1487423e82c1629e1a150c31",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split:   0%|          | 0/52002 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset({\n",
      "    features: ['instruction', 'input', 'output', 'text'],\n",
      "    num_rows: 52002\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "DATASET_NAME = \"tatsu-lab/alpaca\"\n",
    "SPLIT = \"train\"\n",
    "MAX_SEQ_LENGTH = 8192 # 2048\n",
    "EOS_TOKEN = tokenizer.eos_token_id\n",
    "\n",
    "# train is the only data split for this dataset\n",
    "dataset = load_dataset(DATASET_NAME, split=\"train\")\n",
    "print(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### An aside: playing around with dataset object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'datasets.arrow_dataset.Dataset'>\n",
      "<class 'datasets.arrow_dataset.DatasetInfoMixin'>\n",
      "<class 'datasets.search.IndexableMixin'>\n",
      "<class 'datasets.arrow_dataset.TensorflowDatasetMixin'>\n",
      "<class 'object'>\n"
     ]
    }
   ],
   "source": [
    "# Playing around with dataset\n",
    "for cls in getmro(type(dataset)):\n",
    "    print(cls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "instruction\n",
      "\n",
      "Give three tips for staying healthy.\n",
      "\n",
      "input\n",
      "\n",
      "\n",
      "\n",
      "output\n",
      "\n",
      "1.Eat a balanced diet and make sure to include plenty of fruits and vegetables. \n",
      "2. Exercise regularly to keep your body active and strong. \n",
      "3. Get enough sleep and maintain a consistent sleep schedule.\n",
      "\n",
      "text\n",
      "\n",
      "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n",
      "\n",
      "### Instruction:\n",
      "Give three tips for staying healthy.\n",
      "\n",
      "### Response:\n",
      "1.Eat a balanced diet and make sure to include plenty of fruits and vegetables. \n",
      "2. Exercise regularly to keep your body active and strong. \n",
      "3. Get enough sleep and maintain a consistent sleep schedule.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "Fields of dataset\n",
    "\n",
    "    instruction: describes the task the model should perform. Each of the 52K instructions is unique.\n",
    "    input: optional context or input for the task. For example, when the instruction is \"Summarize the following article\", the input is the article. Around 40% of the examples have an input.\n",
    "    output: the answer to the instruction as generated by text-davinci-003.\n",
    "    text: the instruction, input and output formatted with the prompt template used by the authors for fine-tuning their models.\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "# Examine a sample\n",
    "for k,v in dataset[0].items():\n",
    "    print(k)\n",
    "    print()\n",
    "    print(v)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Back to pre-processing our dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n",
      "\n",
      "### Instruction:\n",
      "Render a 3D model of a house\n",
      "\n",
      "### Response:\n",
      "<nooutput> This type of instruction cannot be fulfilled by a GPT model.\n"
     ]
    }
   ],
   "source": [
    "# Select a subset of the data for faster processing\n",
    "dataset = dataset.select(range(100))\n",
    "\n",
    "# Print the 9th example from the 'text' field of the dataset to check the result.\n",
    "print(dataset['text'][8])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def formatting_prompts_func(example):\n",
    "    output_texts = []\n",
    "    for i in range(len(example['instruction'])):\n",
    "        text = f\"{example['input'][i]} ### Question: {example['instruction'][i]}\\n ### Answer: {example['output'][i]}\"\n",
    "        output_texts.append(text)\n",
    "    return output_texts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "# Define a function to format the prompts in the dataset.\n",
    "# This function takes a batch of examples and returns a dictionary with the key 'text' and the value being a list of formatted texts.\n",
    "def formatting_prompts_func(examples):\n",
    "    # Extract the conversations from the examples.\n",
    "    convos = examples[\"conversations\"]\n",
    "    # Initialize an empty list to store the formatted texts.\n",
    "    texts = []\n",
    "    # Define a dictionary to map the 'from' field in the conversation to a prefix.\n",
    "    mapper = {\"system\": \"system\\n\", \"human\": \"\\nuser\\n\", \"gpt\": \"\\nassistant\\n\"}\n",
    "    # Define a dictionary to map the 'from' field in the conversation to a suffix.\n",
    "    end_mapper = {\"system\": \"\", \"human\": \"\", \"gpt\": \"\"}\n",
    "    # Iterate over each conversation.\n",
    "    for convo in convos:\n",
    "        # Format the conversation by joining each turn with its corresponding prefix and suffix.\n",
    "        # Append the EOS token to the end of the conversation.\n",
    "        text = \"\".join(f\"{mapper[(turn := x['from'])]} {x['value']}\\n{end_mapper[turn]}\" for x in convo)\n",
    "        texts.append(f\"{text}{EOS_TOKEN}\")\n",
    "    # Return the formatted texts.\n",
    "    return {\"text\": texts}\n",
    "\n",
    "# Apply the formatting function to the dataset using the map method.\n",
    "# The 'batched=True' argument means that the function is applied to batches of examples.\n",
    "dataset = dataset.map(formatting_prompts_func, batched=True)\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import LoraConfig\n",
    " \n",
    "# LoRA config based on QLoRA paper & Sebastian Raschka experiment\n",
    "\"\"\"peft_config = LoraConfig(\n",
    "        lora_alpha=8,\n",
    "        lora_dropout=0.05,\n",
    "        r=6,\n",
    "        bias=\"none\",\n",
    "        target_modules=\"all-linear\",\n",
    "        task_type=\"CAUSAL_LM\",\n",
    ")\"\"\"\n",
    "\n",
    "lora_config = LoraConfig(\n",
    "    r=6,\n",
    "    target_modules=[\"q_proj\", \"o_proj\", \"k_proj\", \"v_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
    "    bias=\"none\",\n",
    "    task_type=\"CAUSAL_LM\",\n",
    "    lora_alpha=8,\n",
    "    lora_dropout=0.05\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/stefanwebb/anaconda3/envs/huggingface/lib/python3.11/site-packages/transformers/training_args.py:1474: FutureWarning: `evaluation_strategy` is deprecated and will be removed in version 4.46 of 🤗 Transformers. Use `eval_strategy` instead\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "# Create a TrainingArguments object, which is used to define the parameters for model training.\n",
    "args = TrainingArguments(\n",
    "    # 'evaluation_strategy' is set to \"steps\", which means evaluation is done at each logging step.\n",
    "    evaluation_strategy=\"steps\",\n",
    "\n",
    "    # 'per_device_train_batch_size' is set to 7, which means each training batch will contain 7 samples per device.\n",
    "    per_device_train_batch_size=7,\n",
    "\n",
    "    # 'gradient_accumulation_steps' is set to 4, which means gradients are accumulated for 4 steps before performing a backward/update pass.\n",
    "    gradient_accumulation_steps=4,\n",
    "\n",
    "    # 'gradient_checkpointing' is set to True, which means model gradients are stored in memory during training to reduce memory usage.\n",
    "    gradient_checkpointing=True,\n",
    "\n",
    "    # 'learning_rate' is set to 1e-4, which is the learning rate for the optimizer.\n",
    "    learning_rate=1e-4,\n",
    "\n",
    "    # 'fp16' is set to True if bfloat16 is not supported, which means the model will use 16-bit floating point precision for training if possible.\n",
    "    fp16 = not torch.cuda.is_bf16_supported(),\n",
    "\n",
    "    # 'bf16' is set to True if bfloat16 is supported, which means the model will use bfloat16 precision for training if possible.\n",
    "    bf16 = torch.cuda.is_bf16_supported(),\n",
    "\n",
    "    # 'max_steps' is set to -1, which means there is no maximum number of training steps.\n",
    "    max_steps=-1,\n",
    "\n",
    "    # 'num_train_epochs' is set to 3, which means the training process will go through the entire dataset 3 times.\n",
    "    num_train_epochs=3,\n",
    "\n",
    "    # 'save_strategy' is set to \"epoch\", which means the model is saved at the end of each epoch.\n",
    "    save_strategy=\"epoch\",\n",
    "\n",
    "    # 'logging_steps' is set to 10, which means logging is done every 10 steps.\n",
    "    logging_steps=10,\n",
    "\n",
    "    # 'output_dir' is set to NEW_MODEL_NAME, which is the directory where the model and its configuration will be saved.\n",
    "    output_dir=NEW_MODEL_NAME,\n",
    "\n",
    "    # 'optim' is set to \"paged_adamw_32bit\", which is the optimizer to be used for training.\n",
    "    optim=\"paged_adamw_32bit\",\n",
    "\n",
    "    # 'lr_scheduler_type' is set to \"linear\", which means the learning rate scheduler type is linear.\n",
    "    lr_scheduler_type=\"linear\"\n",
    ")\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = TrainingArguments(\n",
    "    output_dir=NEW_MODEL_NAME, # directory to save and repository id\n",
    "    num_train_epochs=3,                     # number of training epochs\n",
    "    per_device_train_batch_size=2,          # batch size per device during training\n",
    "    gradient_accumulation_steps=2,          # number of steps before performing a backward/update pass\n",
    "    gradient_checkpointing=True,            # use gradient checkpointing to save memory\n",
    "    optim=\"paged_adamw_32bit\",              # use fused adamw optimizer\n",
    "    logging_steps=10,                       # log every 10 steps\n",
    "    save_strategy=\"epoch\",                  # save checkpoint every epoch\n",
    "    bf16=True,                              # use bfloat16 precision\n",
    "    tf32=True,                              # use tf32 precision\n",
    "    learning_rate=2e-4,                     # learning rate, based on QLoRA paper\n",
    "    max_grad_norm=0.3,                      # max gradient norm based on QLoRA paper\n",
    "    warmup_ratio=0.03,                      # warmup ratio based on QLoRA paper\n",
    "    lr_scheduler_type=\"constant\",           # use constant learning rate scheduler\n",
    "    push_to_hub=False,                       # push model to hub\n",
    "    # report_to=\"tensorboard\",                # report metrics to tensorboard\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an instance of the SFTTrainer class, which is used to fine-tune the model.\n",
    "\n",
    "trainer = SFTTrainer(\n",
    "    # 'model' is the pre-trained model that will be fine-tuned.\n",
    "    model=model,\n",
    "\n",
    "    # 'args' are the training arguments that specify the training parameters.\n",
    "    args=args,\n",
    "\n",
    "    # 'train_dataset' is the dataset that will be used for training.\n",
    "    train_dataset=dataset,\n",
    "\n",
    "    # 'dataset_text_field' is the key in the dataset that contains the text data.\n",
    "    dataset_text_field=\"text\",\n",
    "\n",
    "    # 'max_seq_length' is the maximum length of the sequences that the model will handle.\n",
    "    max_seq_length=128,\n",
    "\n",
    "    # 'formatting_func' is the function that will be used to format the prompts in the dataset.\n",
    "    formatting_func=formatting_prompts_func\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 'device' is set to 'cuda', which means the CUDA device will be used for computations if available.\n",
    "device = 'cuda'\n",
    "\n",
    "# Import the 'gc' module, which provides an interface to the garbage collector.\n",
    "import gc\n",
    "\n",
    "# Import the 'os' module, which provides a way of using operating system dependent functionality.\n",
    "import os\n",
    "\n",
    "# Call the 'collect' method of the 'gc' module to start a garbage collection, which can help free up memory.\n",
    "gc.collect()\n",
    "\n",
    "# Call the 'empty_cache' method of 'torch.cuda' to release all unused cached memory from PyTorch so that it can be used by other GPU applications.\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Call the 'train' method of the 'trainer' object to start the training process.\n",
    "# This method will fine-tune the model on the training dataset according to the parameters specified in the 'args' object.\n",
    "trainer.train()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "huggingface",
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}