notebooks/llama2_7b_akkadian_peft_finetuning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4b250baf-43f9-4e33-a79d-74231a111aa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "21a44402",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ae78544001074d7ba532e7bb63e3b320",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(HTML(value='<center> <img\\nsrc=https://huggingface.co/front/assets/huggingface_logo-noborder.sv…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# HuggingFace Hub login required for Llama-2 models\n",
    "from huggingface_hub import notebook_login\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d234f956-eac9-47c9-a45e-4b6f5b0e2645",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET = 'data/cleaned_akkadian_en.txt'\n",
    "TOTAL_PROPORTION = 0.10\n",
    "TRAIN_PROPORTION = 0.90\n",
    "\n",
    "CONTEXT_SIZE = 128\n",
    "\n",
    "NUM_TRAIN_STEPS = 200\n",
    "EVAL_STEPS = NUM_TRAIN_STEPS // 20\n",
    "VAL_DATASET_SIZE = 100\n",
    "\n",
    "LOGGING_STEPS = NUM_TRAIN_STEPS // 20\n",
    "\n",
    "CHECKPOINT_FOLDER = 'llama2_akkadian_peft'\n",
    "\n",
    "BASE_MODEL_ID = \"meta-llama/Llama-2-7b-hf\"\n",
    "PEFT_MODEL_ID = \"veezbo/LLama-2-7b-hf-akkadian\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e83f44e5-f7d9-49ec-a510-8a1a335aeb8e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using pad_token, but it is not set yet.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f186cb3fc6d14fb780d33845d89dd7b7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Loading model and tokenizer for use\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "import bitsandbytes as bnb\n",
    "import torch.nn as nn\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)\n",
    "\n",
    "# TODO:vkumar May need to do something like this: tokenizer.pad_token_id = tokenizer.eos_token_id\n",
    "# Above is in order to actually pad the inputs during inference\n",
    "if tokenizer.pad_token is None:\n",
    "    tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    BASE_MODEL_ID, \n",
    "    load_in_8bit=True, ## Using 8-bit precision to load LLAMA-7b on 24GB GPU to fine-tune with PEFT\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7936e17a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Post-processing on the model\n",
    "\n",
    "for param in model.parameters():\n",
    "  param.requires_grad = False  # freeze the model - train adapters later\n",
    "  if param.ndim == 1:\n",
    "    # cast the small parameters (e.g. layernorm) to fp32 for stability\n",
    "    param.data = param.data.to(torch.float32)\n",
    "\n",
    "model.gradient_checkpointing_enable()  # reduce number of stored activations\n",
    "model.enable_input_require_grads()\n",
    "\n",
    "# Ensure that the outputs are still fp32\n",
    "class CastOutputToFloat(nn.Sequential):\n",
    "  def forward(self, x): return super().forward(x).to(torch.float32)\n",
    "model.lm_head = CastOutputToFloat(model.lm_head)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b7469927",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_trainable_parameters(model: nn.Module):\n",
    "    \"\"\"\n",
    "    Prints the number of trainable parameters in the model.\n",
    "    \"\"\"\n",
    "    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "    all_params = sum(p.numel() for p in model.parameters())\n",
    "    print(\n",
    "        f\"trainable params: {trainable_params} || all params: {all_params} || trainable%: {100 * trainable_params / all_params}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3bc99734",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 8388608 || all params: 6746804224 || trainable%: 0.12433454005023165\n"
     ]
    }
   ],
   "source": [
    "from peft import LoraConfig, get_peft_model, TaskType\n",
    "\n",
    "# TODO: Try this with a prompt-type finetuning instead of LORA directly\n",
    "config = LoraConfig(\n",
    "    r=16,\n",
    "    lora_alpha=32,\n",
    "    target_modules=[\"q_proj\", \"v_proj\"],\n",
    "    lora_dropout=0.05,\n",
    "    bias=\"none\",\n",
    "    task_type=TaskType.CAUSAL_LM,  # TODO: Do we want to instead use a Question/Answers task? Data must be really different, though.\n",
    ")\n",
    "\n",
    "model = get_peft_model(model, config)\n",
    "print_trainable_parameters(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2e5be555-6d9a-434c-bc33-fe31578b2b7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import evaluate\n",
    "\n",
    "metric = evaluate.load(\"accuracy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "72a02445-1c8a-494a-8ae9-8d130ff96829",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_metrics(eval_pred):\n",
    "    logits, labels = eval_pred\n",
    "    predictions = np.argmax(logits, axis=-1)\n",
    "    return metric.compute(predictions=predictions, references=labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8924a2fd-cf0a-45ba-8af0-6909c933d492",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9721 874 98\n"
     ]
    }
   ],
   "source": [
    "# Load training and evaluation datasets\n",
    "import random\n",
    "random.seed(42)\n",
    "\n",
    "data = open(DATASET, 'r').read().split('\\n')\n",
    "random.shuffle(data)\n",
    "\n",
    "overall_max_index = int(len(data) * TOTAL_PROPORTION)\n",
    "train_max_index = int(overall_max_index * TRAIN_PROPORTION)\n",
    "train_data = data[:train_max_index]\n",
    "val_data = data[train_max_index:overall_max_index]\n",
    "\n",
    "print(len(data), len(train_data), len(val_data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "535bb097-246a-4e6c-9637-f5e9c5b78b85",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['I took the city Elenzaš as a royal city and a fortress for that district, then I changed its former name and called it Kār-Sennacherib. I settled therein the people of the lands that I had conquered. I placed it under the authority of a eunuch of mine, the governor of the city Ḫarḫar, and thus enlarged my land.', '1 fruit and vegetable garden in the city of Harran: 300 fruit plants therein; 150 poplars and willows: a total of 450.', 'And he said: \"I will destroy Elam; its army shall be levelled to the ground. In this manner I will finish Elam.\"', 'For my lordly pleasure, I had a portico, a replica of a Hittite palace, which is called bīt-ḫilāni in the language of the land Amurru, constructed inside them.', 'Whoever in the future, at any time, lodges a complaint and breaks the contract whether Salmanu-imme or his sons or his  grandsons or his brothers or his nephews or']\n"
     ]
    }
   ],
   "source": [
    "print(train_data[:5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c8629df2-eb9c-44dc-9e8c-9dd7d1339cfe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 60088]) torch.Size([1, 6012])\n"
     ]
    }
   ],
   "source": [
    "train_ids = tokenizer.encode(train_data, return_tensors='pt', is_split_into_words=True)\n",
    "val_ids = tokenizer.encode(val_data, return_tensors='pt', is_split_into_words=True)\n",
    "\n",
    "print(train_ids.shape, val_ids.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "6083a89f-403b-4164-bada-9fc224186797",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.bos_token_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "368b8c44-c14f-4b9c-a559-96d4cbad0c87",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.eos_token_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8abcd993-d7a3-40d5-bb46-acf34fe621f0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[    1,  2180,   393,  ...,  1407,  4549, 29889]])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "val_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f4c2b0e3-7a5a-4fd2-bbee-eb7a37845a9a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 6012])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "val_ids.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "abefb8b9-470f-4f15-a27e-5df138328324",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([6012])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "val_ids.squeeze().shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "a59b19d1-fd05-42e1-bace-44d00d75a2c5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 2180,   393,   931,  ...,  1407,  4549, 29889])"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "val_ids.squeeze()[1:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "df7ab173-a1f1-4ba9-b1f9-771a21aba4ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch import Tensor\n",
    "\n",
    "class AkkadianDatasetforLLM(Dataset):\n",
    "    def __init__(self, input_ids: Tensor, context_size: int):\n",
    "        self.input_ids = input_ids.squeeze()[1:]  # Make tensor 1D and remove the bos_token\n",
    "        self.context_size = context_size\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.input_ids) - self.context_size\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        return {'input_ids': torch.cat((torch.tensor([tokenizer.bos_token_id]), self.input_ids[idx:idx + self.context_size - 1])),\n",
    "                'labels': torch.cat((self.input_ids[idx:idx + self.context_size - 1], torch.tensor([tokenizer.eos_token_id])))}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "691609df-7fdd-4712-b8fc-65f5c4d908ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = AkkadianDatasetforLLM(train_ids, CONTEXT_SIZE)\n",
    "val_dataset = AkkadianDatasetforLLM(val_ids, CONTEXT_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "c8d7277a-28ae-494c-9fa8-c5a325030006",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5884"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(val_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "3f477c30-b0b3-4e9c-960c-2886326437b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Subset\n",
    "\n",
    "indices = random.sample(range(len(val_dataset)), VAL_DATASET_SIZE)\n",
    "\n",
    "val_dataset = Subset(val_dataset, indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "ec121c02-56a0-4f06-8748-3666de00c119",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "100"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(val_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "ba2862fb-863b-4026-bab9-903cf58fd01d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import Trainer, TrainingArguments\n",
    "from transformers.trainer_pt_utils import get_parameter_names\n",
    "from torch import nn\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=CHECKPOINT_FOLDER, \n",
    "    max_steps=NUM_TRAIN_STEPS,\n",
    "    fp16=True,\n",
    "    warmup_steps=100,\n",
    "    per_device_train_batch_size=16,  # TODO: Try a larger size for this batch\n",
    "    gradient_accumulation_steps=4,\n",
    "    learning_rate=5e-5,\n",
    "    \n",
    "    # LOGGING PARAMS\n",
    "    logging_strategy=\"steps\",\n",
    "    logging_steps=LOGGING_STEPS,\n",
    "    \n",
    "    # EVAL Params (TODO: These do not work with error ValueError: Predictions and/or references don't match the expected format. )\n",
    "    # per_device_eval_batch_size=4,\n",
    "    # eval_accumulation_steps=5,\n",
    "    # evaluation_strategy=\"steps\",\n",
    "    # eval_steps=EVAL_STEPS,\n",
    "    \n",
    "    # UNUSED PARAMS\n",
    "    # gradient_checkpointing=True,\n",
    "    # optim=\"adamw_bnb_8bit\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "c0107b8e-6717-4c6a-8893-e9fb36fc5da2",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset,\n",
    "    # eval_dataset=val_dataset,\n",
    "    compute_metrics=compute_metrics\n",
    ")\n",
    "\n",
    "# TODO: Get validation to work without the ValueError above to ensure we're not overfitting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "119dc0e0-151d-44b7-bb24-49b2781797e8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='200' max='200' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [200/200 52:37, Epoch 0/1]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>5.026700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>20</td>\n",
       "      <td>4.990800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>30</td>\n",
       "      <td>4.979700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>40</td>\n",
       "      <td>4.942300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>50</td>\n",
       "      <td>4.912000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>60</td>\n",
       "      <td>4.891600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>70</td>\n",
       "      <td>4.843700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>80</td>\n",
       "      <td>4.751700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>90</td>\n",
       "      <td>4.700700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>100</td>\n",
       "      <td>4.646600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>110</td>\n",
       "      <td>4.559800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>120</td>\n",
       "      <td>4.488600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>130</td>\n",
       "      <td>4.498600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>140</td>\n",
       "      <td>4.381600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>150</td>\n",
       "      <td>4.389100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>160</td>\n",
       "      <td>4.374600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>170</td>\n",
       "      <td>4.354400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>180</td>\n",
       "      <td>4.303600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>190</td>\n",
       "      <td>4.298100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>200</td>\n",
       "      <td>4.328100</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=200, training_loss=4.633113670349121, metrics={'train_runtime': 3172.6487, 'train_samples_per_second': 4.034, 'train_steps_per_second': 0.063, 'total_flos': 6.50352940548096e+16, 'train_loss': 4.633113670349121, 'epoch': 0.21})"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "12dd7707",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "13fda1820e0c45d189292d276918f07e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "adapter_model.bin:   0%|          | 0.00/33.6M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/veezbo/LLama-2-7b-hf-akkadian/commit/950e683f2a46905bd366fe3e327b7701169eda5a', commit_message='Upload model', commit_description='', oid='950e683f2a46905bd366fe3e327b7701169eda5a', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.push_to_hub(PEFT_MODEL_ID, use_auth_token=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "32f6459a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dcda54bb569143838957b5ca8c3da541",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)/adapter_config.json:   0%|          | 0.00/447 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "437f50ae676946fea4fc01660188e4b0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using pad_token, but it is not set yet.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b79da31dfc7f4797820214c0d3ed79cb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading adapter_model.bin:   0%|          | 0.00/33.6M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from peft import PeftModel, PeftConfig\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "import bitsandbytes as bnb\n",
    "import torch.nn as nn\n",
    "\n",
    "config = PeftConfig.from_pretrained(PEFT_MODEL_ID)\n",
    "base_model_name = config.base_model_name_or_path\n",
    "if not base_model_name:\n",
    "    base_model_name = BASE_MODEL_ID\n",
    "    \n",
    "model = AutoModelForCausalLM.from_pretrained(base_model_name, return_dict=True, load_in_8bit=True, device_map='auto')\n",
    "tokenizer = AutoTokenizer.from_pretrained(base_model_name)\n",
    "if tokenizer.pad_token is None:\n",
    "    tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "# Load the Lora model\n",
    "model = PeftModel.from_pretrained(model, PEFT_MODEL_ID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "f6def076-8446-4220-bb3f-ff13b9d45b47",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LoraConfig(peft_type='LORA', auto_mapping=None, base_model_name_or_path='meta-llama/Llama-2-7b-hf', revision=None, task_type='CAUSAL_LM', inference_mode=True, r=16, target_modules=['q_proj', 'v_proj'], lora_alpha=32, lora_dropout=0.05, fan_in_fan_out=False, bias='none', modules_to_save=None, init_lora_weights=True, layers_to_transform=None, layers_pattern=None)\n"
     ]
    }
   ],
   "source": [
    "print(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "a46a8d31-c5b4-4523-b503-925a349e9ae0",
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def generate_from_model(text: str) -> str:\n",
    "    model.eval()\n",
    "    with torch.cuda.amp.autocast():\n",
    "        # tokenizer_output = tokenizer(text, padding='max_length', truncation=True, max_length=CONTEXT_SIZE, return_tensors='pt')\n",
    "        # input_ids = tokenizer_output['input_ids'].cuda()\n",
    "        # pad_mask = tokenizer_output['attention_mask'].cuda()\n",
    "        input_ids = tokenizer.encode(text, return_tensors='pt').cuda()\n",
    "        \n",
    "    # TODO:vkumar Make sure this has the BOS, EOS, and padding tokens as needed.\n",
    "    # Also ensure that the model has the appropriate attention mask (with the padding)\n",
    "    gen_output = model.generate(input_ids=input_ids, max_new_tokens=150)\n",
    "    \n",
    "    str_output = tokenizer.decode(gen_output[0])\n",
    "    \n",
    "    eos_token_index = str_output.find(tokenizer.eos_token)\n",
    "    if eos_token_index != -1:\n",
    "        str_output = str_output[:eos_token_index]\n",
    "    \n",
    "    return str_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "16192f57-3049-492c-8441-55bb288b4635",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'</s>'"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.pad_token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "770128c4-621a-4ad5-9e13-18f984a99681",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ---------------\n",
    "# Below is from 200 training steps with batch size 4, followed by 200 training steps with batch size 16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "7233eff6-9117-455e-8e19-a0d2ccef89fe",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<s> it rains because of king lord the of the Aad, the ofard, the ofil the of-, of- and of- of cityamuu the of ofar of cityama the of theu; the of the 10 ofings the of ofḫḫ, ofšḫ,šḫ,,šu,šl,šur  the of theing ofš,š,šr  the of theing ofšš,šlš,šr  the the ofššlš,šr  the the ofšššlšš The of\n"
     ]
    }
   ],
   "source": [
    "print(generate_from_model(\"it rains because\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "5891279b-795c-4883-9d53-68c4d5215105",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<s> one can live a virtuous life by, the of, and of Šš, king the, of cityḫa, of cityḫ, of cityī the, of city-, of city- of city, the of city the of city the the of-ions the the of-u-- of- of- of--- of-- of- of- of-- of- of- of of the , of the, the of the of, the of of the the of  of-- of- of--- of-- of- of- of of  the of  of-- of--- of- of of of -- of - of - of - of - of - of \n"
     ]
    }
   ],
   "source": [
    "print(generate_from_model(\"one can live a virtuous life by\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "6a6dd0e1-645f-4ad4-bb55-e835a71efda5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<s> a society can thrive by. and the of Ašur the, ofard andull of gods the, of who not a of the of gods the is to of. the of Ašur Nû----, of who not a of the of gods is to of. the ofššu-- of Šš, of whom not a ofšu of gods is to of. theū ofšu the ofš of, of who not ašu ofš of is to of. the ofš--ūš, of whom notš ofš is to of.ū-ūš the\n"
     ]
    }
   ],
   "source": [
    "print(generate_from_model(\"a society can thrive by\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "d0d6a735-3592-47e5-87be-a3298e5cd17f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<s> May any future prince, during whose reign this work falls into disrepair (and) sustains damage the the, should him! the of city the ofḫ-ḫš, is in. the ofšu. the of, king the the the, is in hands. the of---- of cityḪšš is in hands The of-- of, king the of Aad is in hands The of- of kingur of landam is in. the of- of king the '-- ofḫš is in hands The of- of king the -- ofš\n"
     ]
    }
   ],
   "source": [
    "print(generate_from_model(\"May any future prince, during whose reign this work falls into disrepair (and) sustains damage\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b0787bc-b016-44fa-9716-0c77ba3f5af0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ---------------\n",
    "# Below is from only 200 training steps with batch size 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "996bbc7f-625e-4ad2-a3c5-c1613a6ddb99",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<s> it rains because it is raining.\n",
      "it rains because it is raining.\n",
      "it rains because it is raining.\n",
      "it rains because it is raining.\n",
      "it rains because it is raining.\n",
      "it rains because it is raining.\n",
      "it rains because it is raining.\n",
      "it rains because it is raining.\n",
      "it rains because it is raining.\n",
      "it rains because it is raining.\n",
      "it rains because it is raining.\n",
      "it rains because it is raining.\n",
      "it rains because it is raining.\n",
      "it rains because it is raining.\n",
      "it rains because it is raining.\n",
      "it rains because\n"
     ]
    }
   ],
   "source": [
    "print(generate_from_model(\"it rains because\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "f2f83058-56dd-44bd-a850-dfc811e83e58",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<s> one can live a virtuous life by following the precepts of Buddhism.\n",
      "This is a simple and clear statement of a Buddhist doctrine. But what does it mean? What are the precepts of Buddhism?\n",
      "The Buddha gave us eight precepts that we can follow.\n",
      "The first precept is to abstain from taking life.\n",
      "The second precept is to abstain from taking what is not given.\n",
      "The third precept is to abstain from sexual misconduct.\n",
      "The fourth precept is to abstain from false speech.\n",
      "The fifth precept is to abstain from intoxicants.\n",
      "The sixth precept is to abstain from eating at the wrong time.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(generate_from_model(\"one can live a virtuous life by\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "c57f334c-ce8c-4eba-bf6e-cd36e9b0ac0d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<s> a society can thrive by being open to all kinds of ideas, regardless of their source.\n",
      "The following are some of the ideas that have influenced my life.\n",
      "The idea that the world is a beautiful place, that it is worth protecting, and that we have a responsibility to do so.\n",
      "I am a strong believer in the power of nature, and I believe that we should all be doing our best to protect it.\n",
      "The idea that we can change the world for the better, and that we should all be working together to do so.\n",
      "I believe that we should all be working together to create a better world, and that we should all be doing our best to protect the environment.\n",
      "I believe that we should all be doing our best\n"
     ]
    }
   ],
   "source": [
    "print(generate_from_model(\"a society can thrive by\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "e32bb8db-f6f0-4fe5-a8f0-aa4d5a8f3882",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<s> An Akkadian would say that a society can thrive by its ability to harness the power of the sun and the water. Akkadians would say that the power of the sun and the water is the source of all life.\n",
      "Akkadians would say that the sun and the water are the source of all life.\n",
      "Akkadians would say that the sun and the water are the source of all life. Akkadians would say that the sun and the water are the source of all life. Akkadians would say that the sun and the water are the source of all life.\n",
      "Akkadians would say that the sun and the water are the source of all life. Akkadians would say that the sun and the water are the source of\n"
     ]
    }
   ],
   "source": [
    "print(generate_from_model(\"An Akkadian would say that a society can thrive by\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "9f059caf-1201-4b7b-9fa8-a61b29ead01b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<s> An Akkadian would say that one can live a virtuous life by following the principles of the \"Law of the Land.\"\n",
      "The Law of the Land was a concept that was developed by the Sumerian King of Ur-Nammu (2047–2030 BCE). It was based on the belief that all men were created equal and should be treated as such. This meant that everyone was to be given the same rights and privileges, regardless of their social status or wealth.\n",
      "The Law of the Land was one of the first laws to be written down, and it was used as a guide for many other laws that were created throughout history. It is still used today as a basis for many modern laws, such as those in the United States and Canada.\n",
      "The Law\n"
     ]
    }
   ],
   "source": [
    "print(generate_from_model(\"An Akkadian would say that one can live a virtuous life by\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "36fbabca-ca1b-45cf-a791-dc508065b9da",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<s> I named it Kār-Aššur, set up the weapon of the god Aššur, my lord, therein, and settled the people of foreign lands conquered by me therein. I imposed upon them the yoke of the god Aššur, my lord, and I took from them the tribute of the land and the grain of the soil. I built a temple to the god Aššur, my lord, and I dedicated it to him. I built a temple to the god Aššur, my lord, and I dedicated it to him. I built a temple to the god Aššur, my lord, and I dedicated it to him. I built a temple to the god Aššur, my lord, and I dedicated it to him. I built a temple to the god Aššur, my lord, and I dedicated it to him. I built a temple to the god\n"
     ]
    }
   ],
   "source": [
    "print(generate_from_model(\"I named it Kār-Aššur, set up the weapon of the god Aššur, my lord, therein, and settled the people of foreign lands conquered by me therein. I imposed upon them\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "a84b0e9b-ecd2-4a7e-92b5-ba1ca7052713",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(generate_from_model(\"how does one live a virtuous life?\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "86975624-b0c3-4328-b446-f962973bad2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(generate_from_model(\"where can I live?\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59063e46-107e-41c6-a325-8d57f0d5c430",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}