{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import PeftModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "import torch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the base model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained(\"openlm-research/open_llama_3b_v2\", torch_dtype=torch.float16, device_map=\"auto\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"ashrielbrian/openllama_3b_v2-teknium-GPT4-LLM-Cleaned\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "inp = tokenizer(\"write a function that takes in two integers, and returns its modulo.\", return_tensors=\"pt\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': tensor([[    1,  2786,   260,  1155,   347,  2976,   293,   846,  1146,  6014,\n",
       "         29522,   295,  5729,   737,   966, 19795, 29520]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': tensor([[    1,  2786,   260,  1155,   347,  2976,   293,   846,  1146,  6014,\n",
       "          29522,   295,  5729,   737,   966, 19795, 29520]], device='cuda:0'),\n",
       " 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], device='cuda:0')}"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inputs = {k: v.to(\"cuda\") for k, v in inp.items()}\n",
    "inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    generate_ids = model.generate(**inputs, max_length=1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'write a function that takes in two integers, and returns its modulo.\\nFor example, if you have the integers 10 and 2, then the modulo of 10 and 2 is 2.\\nThe modulo of 10 and 3 is 3.\\nThe modulo of 10 and 4 is 4.\\nThe modulo of 10 and 5 is 5.\\nThe modulo of 10 and 6 is 6.\\nThe modulo of 10 and 7 is 7.\\nThe modulo of 10 and 8 is 8.\\nThe modulo of 10 and 9 is 9.\\nThe modulo of 10 and 10 is 10.\\nThe modulo of 10 and 11 is 11.\\nThe modulo of 10 and 12 is 12.\\nThe modulo of 10 and 13 is 13.\\nThe modulo of 10 and 14 is 14.\\nThe modulo of 10 and 15 is 15.\\nThe modulo of 10 and 16 is 16.\\nThe modulo of 10 and 17 is 17.\\nThe modulo of 10 and 18 is 18.\\nThe modulo of 10 and 19 is 19.\\nThe modulo of 10 and 20 is 20.\\nThe modulo of 10 and 21 is 21.\\nThe modulo of 10 and 22 is 22.\\nThe modulo of 10 and 23 is 23.\\nThe modulo of 10 and 24 is 24.\\nThe modulo of 10 and 25 is 25.\\nThe modulo of 10 and 26 is 26.\\nThe modulo of 10 and 27 is 27.\\nThe modulo of 10 and 28 is 28.\\nThe modulo of 10 and 29 is 29.\\nThe modulo of 10 and 30 is 30.\\nThe modulo of 10 and 31 is 31.\\nThe modulo of 10 and 32 is 32.\\nThe modulo of 10 and 33 is 33.\\nThe modulo of 10 and 34 is 34.\\nThe modulo of 10 and 35 is 35.\\nThe modulo of 10 and 36 is 36.\\nThe modulo of 10 and 37 is 37.\\nThe modulo of 10 and 38 is 38.\\nThe modulo of 10 and 39 is 39.\\nThe modulo of 10 and 40 is 40.\\nThe modulo of 10 and 41 is 41.\\nThe modulo of 10 and 42 is 42.\\nThe modulo of 10 and 43 is 43.\\nThe modulo of 10 and 44 is 44.\\nThe modulo of 10 and 45 is 45.\\nThe modulo of 10 and 46 is 46.\\nThe modulo of 10 and 47 is 47.\\nThe modulo of 10 and 48 is 48.\\nThe modulo of 10 and 49 is 49.\\nThe modulo of 10 and 50 is 50.\\nThe modulo of 10 and 51 is 51.\\nThe modulo of 10 and 52 is 52.\\nThe modulo of 10 and 53 is 53.\\nThe modulo of 10 and 54 is 54.\\nThe modulo of 10 and 55 is 55.\\nThe modulo of 10 and 56 is 56.\\nThe modulo of 10 and 57 is 57.\\nThe modulo of 10 and 58 is 58.\\nThe modulo of 10 and 59 is'"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# the output is practically gibberish because it was trained as a completion model, and is NOT instruction-tuned.\n",
    "outputs = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]\n",
    "outputs\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load adapters without merging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "peft_model_id = \"ashrielbrian/openllama_3b_v2-teknium-GPT4-LLM-Cleaned\"\n",
    "model = PeftModel.from_pretrained(model, peft_model_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'write a function that takes in two integers, and returns its modulo.\\n\\nHere is one way to write the function in Python:\\n\\n```python\\ndef modulo(a, b):\\n    return a % b\\n```\\n\\nThis function takes in two arguments, `a` and `b`, and returns the remainder of `a` when divided by `b`. The remainder is the value that remains after the division.'"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    generate_ids = model.generate(**inputs, max_length=1000)\n",
    "\n",
    "outputs = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]\n",
    "outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Merge the adapter into the base model\n",
    "Helpful resource on [Huggingface](https://huggingface.co/docs/peft/main/en/developer_guides/lora)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PeftModelForCausalLM(\n",
       "  (base_model): LoraModel(\n",
       "    (model): LlamaForCausalLM(\n",
       "      (model): LlamaModel(\n",
       "        (embed_tokens): Embedding(32000, 3200, padding_idx=0)\n",
       "        (layers): ModuleList(\n",
       "          (0-25): 26 x LlamaDecoderLayer(\n",
       "            (self_attn): LlamaSdpaAttention(\n",
       "              (q_proj): Linear(in_features=3200, out_features=3200, bias=False)\n",
       "              (k_proj): Linear(in_features=3200, out_features=3200, bias=False)\n",
       "              (v_proj): Linear(in_features=3200, out_features=3200, bias=False)\n",
       "              (o_proj): Linear(in_features=3200, out_features=3200, bias=False)\n",
       "              (rotary_emb): LlamaRotaryEmbedding()\n",
       "            )\n",
       "            (mlp): LlamaMLP(\n",
       "              (gate_proj): Linear(in_features=3200, out_features=8640, bias=False)\n",
       "              (up_proj): Linear(in_features=3200, out_features=8640, bias=False)\n",
       "              (down_proj): Linear(in_features=8640, out_features=3200, bias=False)\n",
       "              (act_fn): SiLU()\n",
       "            )\n",
       "            (input_layernorm): LlamaRMSNorm()\n",
       "            (post_attention_layernorm): LlamaRMSNorm()\n",
       "          )\n",
       "        )\n",
       "        (norm): LlamaRMSNorm()\n",
       "      )\n",
       "      (lm_head): Linear(in_features=3200, out_features=32000, bias=False)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.merge_and_unload()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'write a function that takes in two integers, and returns its modulo.\\n\\nHere is one way to write the function in Python:\\n\\n```python\\ndef modulo(a, b):\\n    return a % b\\n```\\n\\nThis function takes in two arguments, `a` and `b`, and returns the remainder of `a` when divided by `b`. The remainder is the value that remains after the division.'"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    generate_ids = model.generate(**inputs, max_length=1000)\n",
    "\n",
    "# inference latency here is lower than if we kept the adapter separate as in the previous step\n",
    "# comparing walltime between the unmerged adapter model, with the merged LORA weights, are 1.3s and 0.9s respectively.\n",
    "outputs = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]\n",
    "outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "axolotl",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}