mistral-custom-qLoRA.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "317bfd0d-32fa-4d17-bf47-9c72dc90226e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!pip install -q -U bitsandbytes\n",
    "!pip install -q -U git+https://github.com/huggingface/transformers.git\n",
    "!pip install -q -U git+https://github.com/huggingface/peft.git\n",
    "!pip install -q -U git+https://github.com/huggingface/accelerate.git\n",
    "!pip install -q -U datasets scipy ipywidgets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "eec36d2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q -U datasets "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6340ac5d-3dda-4ee6-a4e6-e313ab338338",
   "metadata": {},
   "outputs": [],
   "source": [
    "from accelerate import FullyShardedDataParallelPlugin, Accelerator\n",
    "from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig\n",
    "fsdp_plugin = FullyShardedDataParallelPlugin(\n",
    "    state_dict_config=FullStateDictConfig(offload_to_cpu=True, rank0_only=False),\n",
    "    optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=False),\n",
    ")\n",
    "accelerator = Accelerator(fsdp_plugin=fsdp_plugin)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "34a1f1fd-d5a8-4d6a-a05e-1119f98fc18c",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from datasets import load_dataset\n",
    "\n",
    "\n",
    "train_dataset = load_dataset('squad_v2', split='train')\n",
    "eval_dataset = load_dataset('squad_v2', split='validation')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d09f4968",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Source</th>\n",
       "      <th>Output</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>11Lecture Self-Check Question \\nAssign a varia...</td>\n",
       "      <td>{\\n    \"Output\": {\\n        \"question\": \"Where...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>•Give the program \\nadditional conditions. \\n•...</td>\n",
       "      <td>{\\n \"Output\": {\\n  \"question\": \"What type of r...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>•Java Packages•The String Class•Using System.o...</td>\n",
       "      <td>{\\n    \"Output\": {\\n        \"question\": \"Which...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>TYPES OF NETWORK\\nPeer-to-peer\\nClient-Server\\...</td>\n",
       "      <td>{\\n    \"Output\": {\\n        \"question\": \"Which...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Physical Topologies\\n6Bus\\nRing\\nStarExtended ...</td>\n",
       "      <td>{\\n    \"Output\": {\\n        \"question\": \"Which...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                              Source  \\\n",
       "0  11Lecture Self-Check Question \\nAssign a varia...   \n",
       "1  •Give the program \\nadditional conditions. \\n•...   \n",
       "2  •Java Packages•The String Class•Using System.o...   \n",
       "3  TYPES OF NETWORK\\nPeer-to-peer\\nClient-Server\\...   \n",
       "4  Physical Topologies\\n6Bus\\nRing\\nStarExtended ...   \n",
       "\n",
       "                                              Output  \n",
       "0  {\\n    \"Output\": {\\n        \"question\": \"Where...  \n",
       "1  {\\n \"Output\": {\\n  \"question\": \"What type of r...  \n",
       "2  {\\n    \"Output\": {\\n        \"question\": \"Which...  \n",
       "3  {\\n    \"Output\": {\\n        \"question\": \"Which...  \n",
       "4  {\\n    \"Output\": {\\n        \"question\": \"Which...  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "df = pd.read_csv('./modified_dataset.csv')\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "477af1cc-d217-4e37-b0e1-f46bdc6add6a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset({\n",
      "    features: ['id', 'title', 'context', 'question', 'answers'],\n",
      "    num_rows: 130319\n",
      "})\n",
      "Dataset({\n",
      "    features: ['id', 'title', 'context', 'question', 'answers'],\n",
      "    num_rows: 11873\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "print(train_dataset)\n",
    "print(eval_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f5535de1-05cd-4c01-8b8d-239002ea485e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "53d0ef405d424fef86a3840a282c92d8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
    "\n",
    "\n",
    "base_model_id = \"mistralai/Mistral-7B-v0.1\"\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_use_double_quant=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    "    bnb_4bit_compute_dtype=torch.bfloat16\n",
    ")\n",
    "model = AutoModelForCausalLM.from_pretrained(base_model_id, quantization_config=bnb_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b0bad471-ba9d-441e-9658-5a1cf7f1dcfe",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    base_model_id,\n",
    "    model_max_length=512,\n",
    "    padding_side=\"left\",\n",
    "    add_eos_token=True)\n",
    "tokenizer.pad_token = tokenizer.eos_token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9f98f462-9edd-40c6-b4db-3f385a1e53c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize(prompt):\n",
    "    result = tokenizer(\n",
    "        prompt,\n",
    "        truncation=True,\n",
    "        max_length=512,\n",
    "        padding=\"max_length\",\n",
    "    )\n",
    "    result[\"labels\"] = result[\"input_ids\"].copy()\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6a0b46e-e92d-4215-957c-e21f53e31500",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f96fcef4-7756-4248-b416-08523090aa30",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_and_tokenize_prompt(row):\n",
    "    full_prompt = f\"\"\"\n",
    "    Instruct: Generate multiple MCQ questions and correct answer and incorrect answers from the given context. The question should be general purpose and should be answerable without the context. Output should be an array of JSON in the below format.\n",
    "                      [\n",
    "  {{\n",
    "    \"question\": \"string\",\n",
    "    \"correct_answer\": \"string\",\n",
    "    \"incorrect_answers\": [\"string\", \"string\", \"string\"]\n",
    "  }}, \n",
    "]\n",
    "    Context: {row[\"Source\"]}\n",
    "Output: {row[\"Output\"]}\n",
    "    \"\"\"\n",
    "    return tokenize(full_prompt)\n",
    "\n",
    "tokenized_data = df.apply(generate_and_tokenize_prompt, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "473022f4-0818-453c-afdf-d09428d5b520",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e0a75749e7f446a99b4be8f2f0bf384c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/130319 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "tokenized_train_dataset = train_dataset.map(generate_and_tokenize_prompt)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "4cc6b906-342c-4b52-b384-8ae81e176cd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_val_dataset = eval_dataset.map(generate_and_tokenize_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d87a131b-389b-4c26-a912-4ab21a71e8bf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'id': '56bf6b0f3aeaaa14008c9602', 'title': 'Beyoncé', 'context': 'Beyoncé Giselle Knowles-Carter (/biːˈjɒnseɪ/ bee-YON-say) (born September 4, 1981) is an American singer, songwriter, record producer and actress. Born and raised in Houston, Texas, she performed in various singing and dancing competitions as a child, and rose to fame in the late 1990s as lead singer of R&B girl-group Destiny\\'s Child. Managed by her father, Mathew Knowles, the group became one of the world\\'s best-selling girl groups of all time. Their hiatus saw the release of Beyoncé\\'s debut album, Dangerously in Love (2003), which established her as a solo artist worldwide, earned five Grammy Awards and featured the Billboard Hot 100 number-one singles \"Crazy in Love\" and \"Baby Boy\".', 'question': 'In which decade did Beyonce become famous?', 'answers': {'text': ['late 1990s'], 'answer_start': [276]}, 'input_ids': [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 345, 27554, 1374, 272, 3857, 2758, 304, 8270, 264, 8598, 2996, 304, 4372, 28723, 28705, 13, 2287, 3433, 8261, 272, 2758, 778, 6421, 8076, 298, 2380, 272, 24685, 3030, 28723, 13, 2287, 4849, 6212, 264, 2996, 369, 12728, 597, 3900, 477, 272, 2758, 28723, 13, 2287, 8587, 264, 1001, 21891, 4372, 486, 19135, 16441, 17807, 2373, 272, 2758, 28723, 13, 13, 2287, 774, 14268, 28747, 13, 2287, 16015, 266, 22034, 420, 278, 3222, 11308, 867, 28733, 28743, 4136, 325, 28748, 6309, 29240, 29097, 28768, 31099, 28711, 331, 29603, 28748, 347, 28706, 28733, 28802, 832, 28733, 21205, 28731, 325, 6363, 4074, 28705, 28781, 28725, 28705, 28740, 28774, 28783, 28740, 28731, 349, 396, 2556, 15307, 28725, 4034, 11660, 28725, 2395, 13876, 304, 18334, 28723, 20109, 304, 6333, 297, 16904, 28725, 7826, 28725, 630, 7885, 297, 4118, 14493, 304, 17319, 3440, 2065, 390, 264, 1502, 28725, 304, 8536, 298, 24665, 297, 272, 3909, 28705, 28740, 28774, 28774, 28734, 28713, 390, 1736, 15307, 302, 399, 28800, 28760, 2746, 28733, 3063, 14466, 4279, 28742, 28713, 8183, 28723, 2213, 2569, 486, 559, 3140, 28725, 5702, 10345, 11308, 867, 28725, 272, 2071, 3246, 624, 302, 272, 1526, 28742, 28713, 1489, 28733, 24943, 2746, 4938, 302, 544, 727, 28723, 6723, 12014, 1589, 2672, 272, 5614, 302, 16015, 266, 22034, 28742, 28713, 9928, 4708, 28725, 384, 3689, 5019, 297, 7481, 325, 28750, 28734, 28734, 28770, 557, 690, 6740, 559, 390, 264, 8375, 7325, 15245, 28725, 12839, 3359, 21629, 1916, 10882, 304, 12526, 272, 6502, 3516, 6895, 28705, 28740, 28734, 28734, 1474, 28733, 538, 16975, 345, 28743, 7853, 297, 7481, 28739, 304, 345, 28760, 4541, 11036, 2586, 13, 260, 13, 260, 13, 2287, 774, 22478, 28747, 13, 2287, 560, 690, 12250, 863, 16015, 11069, 2727, 8376, 28804, 13, 2287, 774, 26307, 13, 2287, 3909, 28705, 28740, 28774, 28774, 28734, 28713, 13, 260, 2], 'attention_mask': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], 'labels': [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 345, 27554, 1374, 272, 3857, 2758, 304, 8270, 264, 8598, 2996, 304, 4372, 28723, 28705, 13, 2287, 3433, 8261, 272, 2758, 778, 6421, 8076, 298, 2380, 272, 24685, 3030, 28723, 13, 2287, 4849, 6212, 264, 2996, 369, 12728, 597, 3900, 477, 272, 2758, 28723, 13, 2287, 8587, 264, 1001, 21891, 4372, 486, 19135, 16441, 17807, 2373, 272, 2758, 28723, 13, 13, 2287, 774, 14268, 28747, 13, 2287, 16015, 266, 22034, 420, 278, 3222, 11308, 867, 28733, 28743, 4136, 325, 28748, 6309, 29240, 29097, 28768, 31099, 28711, 331, 29603, 28748, 347, 28706, 28733, 28802, 832, 28733, 21205, 28731, 325, 6363, 4074, 28705, 28781, 28725, 28705, 28740, 28774, 28783, 28740, 28731, 349, 396, 2556, 15307, 28725, 4034, 11660, 28725, 2395, 13876, 304, 18334, 28723, 20109, 304, 6333, 297, 16904, 28725, 7826, 28725, 630, 7885, 297, 4118, 14493, 304, 17319, 3440, 2065, 390, 264, 1502, 28725, 304, 8536, 298, 24665, 297, 272, 3909, 28705, 28740, 28774, 28774, 28734, 28713, 390, 1736, 15307, 302, 399, 28800, 28760, 2746, 28733, 3063, 14466, 4279, 28742, 28713, 8183, 28723, 2213, 2569, 486, 559, 3140, 28725, 5702, 10345, 11308, 867, 28725, 272, 2071, 3246, 624, 302, 272, 1526, 28742, 28713, 1489, 28733, 24943, 2746, 4938, 302, 544, 727, 28723, 6723, 12014, 1589, 2672, 272, 5614, 302, 16015, 266, 22034, 28742, 28713, 9928, 4708, 28725, 384, 3689, 5019, 297, 7481, 325, 28750, 28734, 28734, 28770, 557, 690, 6740, 559, 390, 264, 8375, 7325, 15245, 28725, 12839, 3359, 21629, 1916, 10882, 304, 12526, 272, 6502, 3516, 6895, 28705, 28740, 28734, 28734, 1474, 28733, 538, 16975, 345, 28743, 7853, 297, 7481, 28739, 304, 345, 28760, 4541, 11036, 2586, 13, 260, 13, 260, 13, 2287, 774, 22478, 28747, 13, 2287, 560, 690, 12250, 863, 16015, 11069, 2727, 8376, 28804, 13, 2287, 774, 26307, 13, 2287, 3909, 28705, 28740, 28774, 28774, 28734, 28713, 13, 260, 2]}\n"
     ]
    }
   ],
   "source": [
    "print(tokenized_train_dataset[4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "1cd74cd4-c139-44b4-9a2e-2596d9493964",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "512\n"
     ]
    }
   ],
   "source": [
    "print(len(tokenized_train_dataset[4]['input_ids']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "8192d670-0d41-4d5b-97c7-1500083568aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import prepare_model_for_kbit_training\n",
    "model.gradient_checkpointing_enable()\n",
    "model = prepare_model_for_kbit_training(model)\n",
    "\n",
    "def print_trainable_parameters(model):\n",
    "    \"\"\"\n",
    "    Prints the number of trainable parameters in the model.\n",
    "    \"\"\"\n",
    "    trainable_params = 0\n",
    "    all_param = 0\n",
    "    for _, param in model.named_parameters():\n",
    "        all_param += param.numel()\n",
    "        if param.requires_grad:\n",
    "            trainable_params += param.numel()\n",
    "    print(\n",
    "        f\"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "030164fe-4f61-41b4-bdf0-356c4c1fd82d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 21260288 || all params: 3773331456 || trainable%: 0.5634354746703705\n",
      "PeftModelForCausalLM(\n",
      "  (base_model): LoraModel(\n",
      "    (model): MistralForCausalLM(\n",
      "      (model): MistralModel(\n",
      "        (embed_tokens): Embedding(32000, 4096)\n",
      "        (layers): ModuleList(\n",
      "          (0-31): 32 x MistralDecoderLayer(\n",
      "            (self_attn): MistralAttention(\n",
      "              (q_proj): lora.Linear4bit(\n",
      "                (base_layer): Linear4bit(in_features=4096, out_features=4096, bias=False)\n",
      "                (lora_dropout): ModuleDict(\n",
      "                  (default): Dropout(p=0.05, inplace=False)\n",
      "                )\n",
      "                (lora_A): ModuleDict(\n",
      "                  (default): Linear(in_features=4096, out_features=8, bias=False)\n",
      "                )\n",
      "                (lora_B): ModuleDict(\n",
      "                  (default): Linear(in_features=8, out_features=4096, bias=False)\n",
      "                )\n",
      "                (lora_embedding_A): ParameterDict()\n",
      "                (lora_embedding_B): ParameterDict()\n",
      "              )\n",
      "              (k_proj): lora.Linear4bit(\n",
      "                (base_layer): Linear4bit(in_features=4096, out_features=1024, bias=False)\n",
      "                (lora_dropout): ModuleDict(\n",
      "                  (default): Dropout(p=0.05, inplace=False)\n",
      "                )\n",
      "                (lora_A): ModuleDict(\n",
      "                  (default): Linear(in_features=4096, out_features=8, bias=False)\n",
      "                )\n",
      "                (lora_B): ModuleDict(\n",
      "                  (default): Linear(in_features=8, out_features=1024, bias=False)\n",
      "                )\n",
      "                (lora_embedding_A): ParameterDict()\n",
      "                (lora_embedding_B): ParameterDict()\n",
      "              )\n",
      "              (v_proj): lora.Linear4bit(\n",
      "                (base_layer): Linear4bit(in_features=4096, out_features=1024, bias=False)\n",
      "                (lora_dropout): ModuleDict(\n",
      "                  (default): Dropout(p=0.05, inplace=False)\n",
      "                )\n",
      "                (lora_A): ModuleDict(\n",
      "                  (default): Linear(in_features=4096, out_features=8, bias=False)\n",
      "                )\n",
      "                (lora_B): ModuleDict(\n",
      "                  (default): Linear(in_features=8, out_features=1024, bias=False)\n",
      "                )\n",
      "                (lora_embedding_A): ParameterDict()\n",
      "                (lora_embedding_B): ParameterDict()\n",
      "              )\n",
      "              (o_proj): lora.Linear4bit(\n",
      "                (base_layer): Linear4bit(in_features=4096, out_features=4096, bias=False)\n",
      "                (lora_dropout): ModuleDict(\n",
      "                  (default): Dropout(p=0.05, inplace=False)\n",
      "                )\n",
      "                (lora_A): ModuleDict(\n",
      "                  (default): Linear(in_features=4096, out_features=8, bias=False)\n",
      "                )\n",
      "                (lora_B): ModuleDict(\n",
      "                  (default): Linear(in_features=8, out_features=4096, bias=False)\n",
      "                )\n",
      "                (lora_embedding_A): ParameterDict()\n",
      "                (lora_embedding_B): ParameterDict()\n",
      "              )\n",
      "              (rotary_emb): MistralRotaryEmbedding()\n",
      "            )\n",
      "            (mlp): MistralMLP(\n",
      "              (gate_proj): lora.Linear4bit(\n",
      "                (base_layer): Linear4bit(in_features=4096, out_features=14336, bias=False)\n",
      "                (lora_dropout): ModuleDict(\n",
      "                  (default): Dropout(p=0.05, inplace=False)\n",
      "                )\n",
      "                (lora_A): ModuleDict(\n",
      "                  (default): Linear(in_features=4096, out_features=8, bias=False)\n",
      "                )\n",
      "                (lora_B): ModuleDict(\n",
      "                  (default): Linear(in_features=8, out_features=14336, bias=False)\n",
      "                )\n",
      "                (lora_embedding_A): ParameterDict()\n",
      "                (lora_embedding_B): ParameterDict()\n",
      "              )\n",
      "              (up_proj): lora.Linear4bit(\n",
      "                (base_layer): Linear4bit(in_features=4096, out_features=14336, bias=False)\n",
      "                (lora_dropout): ModuleDict(\n",
      "                  (default): Dropout(p=0.05, inplace=False)\n",
      "                )\n",
      "                (lora_A): ModuleDict(\n",
      "                  (default): Linear(in_features=4096, out_features=8, bias=False)\n",
      "                )\n",
      "                (lora_B): ModuleDict(\n",
      "                  (default): Linear(in_features=8, out_features=14336, bias=False)\n",
      "                )\n",
      "                (lora_embedding_A): ParameterDict()\n",
      "                (lora_embedding_B): ParameterDict()\n",
      "              )\n",
      "              (down_proj): lora.Linear4bit(\n",
      "                (base_layer): Linear4bit(in_features=14336, out_features=4096, bias=False)\n",
      "                (lora_dropout): ModuleDict(\n",
      "                  (default): Dropout(p=0.05, inplace=False)\n",
      "                )\n",
      "                (lora_A): ModuleDict(\n",
      "                  (default): Linear(in_features=14336, out_features=8, bias=False)\n",
      "                )\n",
      "                (lora_B): ModuleDict(\n",
      "                  (default): Linear(in_features=8, out_features=4096, bias=False)\n",
      "                )\n",
      "                (lora_embedding_A): ParameterDict()\n",
      "                (lora_embedding_B): ParameterDict()\n",
      "              )\n",
      "              (act_fn): SiLU()\n",
      "            )\n",
      "            (input_layernorm): MistralRMSNorm()\n",
      "            (post_attention_layernorm): MistralRMSNorm()\n",
      "          )\n",
      "        )\n",
      "        (norm): MistralRMSNorm()\n",
      "      )\n",
      "      (lm_head): lora.Linear(\n",
      "        (base_layer): Linear(in_features=4096, out_features=32000, bias=False)\n",
      "        (lora_dropout): ModuleDict(\n",
      "          (default): Dropout(p=0.05, inplace=False)\n",
      "        )\n",
      "        (lora_A): ModuleDict(\n",
      "          (default): Linear(in_features=4096, out_features=8, bias=False)\n",
      "        )\n",
      "        (lora_B): ModuleDict(\n",
      "          (default): Linear(in_features=8, out_features=32000, bias=False)\n",
      "        )\n",
      "        (lora_embedding_A): ParameterDict()\n",
      "        (lora_embedding_B): ParameterDict()\n",
      "      )\n",
      "    )\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "from peft import LoraConfig, get_peft_model\n",
    "config = LoraConfig(\n",
    "    r=8,\n",
    "    lora_alpha=16,\n",
    "    target_modules=[\n",
    "        \"q_proj\",\n",
    "        \"k_proj\",\n",
    "        \"v_proj\",\n",
    "        \"o_proj\",\n",
    "        \"gate_proj\",\n",
    "        \"up_proj\",\n",
    "        \"down_proj\",\n",
    "        \"lm_head\",\n",
    "    ],\n",
    "    bias=\"none\",\n",
    "    lora_dropout=0.05,  # Conventional\n",
    "    task_type=\"CAUSAL_LM\",\n",
    ")\n",
    "model = get_peft_model(model, config)\n",
    "print_trainable_parameters(model)\n",
    "# Apply the accelerator. You can comment this out to remove the accelerator.\n",
    "model = accelerator.prepare_model(model)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d175d06f-bb22-4730-b4fc-b03a3999747f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mvpmb\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    }
   ],
   "source": [
    "!pip install -q wandb -U\n",
    "\n",
    "import wandb, os\n",
    "wandb.login()\n",
    "\n",
    "wandb_project = \"mistral-finetune-v1.7\"\n",
    "if len(wandb_project) > 0:\n",
    "    os.environ[\"WANDB_PROJECT\"] = wandb_project\n",
    "    os.environ[\"WANDB_RESUME\"] = \"must\"\n",
    "    os.environ[\"WANDB_RUN_ID\"] = \"f5rpfdmw\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "9add748c-13fe-4974-8f25-65b083d9e3ff",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "Changes to your `wandb` environment variables will be ignored because your `wandb` session has already started. For more information on how to modify your settings with `wandb.init()` arguments, please refer to <a href='https://wandb.me/wandb-init' target=\"_blank\">the W&B docs</a>."
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "221d46e7d6a94ca7a08e221343ef3e50",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='Waiting for wandb.init()...\\r'), FloatProgress(value=0.011113907299780597, max=1.0…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Problem at: /tmp/ipykernel_220/2155807707.py 1 <module>\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[28], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mwandb\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minit\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mid\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mf5rpfdmw\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mresume\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mmust\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/usr/local/lib/python3.10/dist-packages/wandb/sdk/wandb_init.py:1193\u001b[0m, in \u001b[0;36minit\u001b[0;34m(job_type, dir, config, project, entity, reinit, tags, group, name, notes, magic, config_exclude_keys, config_include_keys, anonymous, mode, allow_val_change, resume, force, tensorboard, sync_tensorboard, monitor_gym, save_code, id, settings)\u001b[0m\n\u001b[1;32m   1191\u001b[0m     \u001b[38;5;28;01massert\u001b[39;00m logger\n\u001b[1;32m   1192\u001b[0m     logger\u001b[38;5;241m.\u001b[39mwarning(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124minterrupted\u001b[39m\u001b[38;5;124m\"\u001b[39m, exc_info\u001b[38;5;241m=\u001b[39me)\n\u001b[0;32m-> 1193\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m e\n\u001b[1;32m   1194\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m   1195\u001b[0m     error_seen \u001b[38;5;241m=\u001b[39m e\n",
      "File \u001b[0;32m/usr/local/lib/python3.10/dist-packages/wandb/sdk/wandb_init.py:1170\u001b[0m, in \u001b[0;36minit\u001b[0;34m(job_type, dir, config, project, entity, reinit, tags, group, name, notes, magic, config_exclude_keys, config_include_keys, anonymous, mode, allow_val_change, resume, force, tensorboard, sync_tensorboard, monitor_gym, save_code, id, settings)\u001b[0m\n\u001b[1;32m   1168\u001b[0m except_exit \u001b[38;5;241m=\u001b[39m wi\u001b[38;5;241m.\u001b[39msettings\u001b[38;5;241m.\u001b[39m_except_exit\n\u001b[1;32m   1169\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 1170\u001b[0m     run \u001b[38;5;241m=\u001b[39m \u001b[43mwi\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minit\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1171\u001b[0m     except_exit \u001b[38;5;241m=\u001b[39m wi\u001b[38;5;241m.\u001b[39msettings\u001b[38;5;241m.\u001b[39m_except_exit\n\u001b[1;32m   1172\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m (\u001b[38;5;167;01mKeyboardInterrupt\u001b[39;00m, \u001b[38;5;167;01mException\u001b[39;00m) \u001b[38;5;28;01mas\u001b[39;00m e:\n",
      "File \u001b[0;32m/usr/local/lib/python3.10/dist-packages/wandb/sdk/wandb_init.py:752\u001b[0m, in \u001b[0;36m_WandbInit.init\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    749\u001b[0m logger\u001b[38;5;241m.\u001b[39minfo(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcommunicating run to backend with \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtimeout\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m second timeout\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    751\u001b[0m run_init_handle \u001b[38;5;241m=\u001b[39m backend\u001b[38;5;241m.\u001b[39minterface\u001b[38;5;241m.\u001b[39mdeliver_run(run)\n\u001b[0;32m--> 752\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[43mrun_init_handle\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwait\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    753\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtimeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtimeout\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    754\u001b[0m \u001b[43m    \u001b[49m\u001b[43mon_progress\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_on_progress_init\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    755\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcancel\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    756\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    757\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m result:\n\u001b[1;32m    758\u001b[0m     run_result \u001b[38;5;241m=\u001b[39m result\u001b[38;5;241m.\u001b[39mrun_result\n",
      "File \u001b[0;32m/usr/local/lib/python3.10/dist-packages/wandb/sdk/lib/mailbox.py:283\u001b[0m, in \u001b[0;36mMailboxHandle.wait\u001b[0;34m(self, timeout, on_probe, on_progress, release, cancel)\u001b[0m\n\u001b[1;32m    280\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_interface\u001b[38;5;241m.\u001b[39m_transport_keepalive_failed():\n\u001b[1;32m    281\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m MailboxError(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtransport failed\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m--> 283\u001b[0m found, abandoned \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_slot\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_get_and_clear\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtimeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwait_timeout\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    284\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m found:\n\u001b[1;32m    285\u001b[0m     \u001b[38;5;66;03m# Always update progress to 100% when done\u001b[39;00m\n\u001b[1;32m    286\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m on_progress \u001b[38;5;129;01mand\u001b[39;00m progress_handle \u001b[38;5;129;01mand\u001b[39;00m progress_sent:\n",
      "File \u001b[0;32m/usr/local/lib/python3.10/dist-packages/wandb/sdk/lib/mailbox.py:130\u001b[0m, in \u001b[0;36m_MailboxSlot._get_and_clear\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    128\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_get_and_clear\u001b[39m(\u001b[38;5;28mself\u001b[39m, timeout: \u001b[38;5;28mfloat\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tuple[Optional[pb\u001b[38;5;241m.\u001b[39mResult], \u001b[38;5;28mbool\u001b[39m]:\n\u001b[1;32m    129\u001b[0m     found \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 130\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_wait\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtimeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m:\n\u001b[1;32m    131\u001b[0m         \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_lock:\n\u001b[1;32m    132\u001b[0m             found \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_result\n",
      "File \u001b[0;32m/usr/local/lib/python3.10/dist-packages/wandb/sdk/lib/mailbox.py:126\u001b[0m, in \u001b[0;36m_MailboxSlot._wait\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    125\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_wait\u001b[39m(\u001b[38;5;28mself\u001b[39m, timeout: \u001b[38;5;28mfloat\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28mbool\u001b[39m:\n\u001b[0;32m--> 126\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_event\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwait\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtimeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/usr/lib/python3.10/threading.py:607\u001b[0m, in \u001b[0;36mEvent.wait\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    605\u001b[0m signaled \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_flag\n\u001b[1;32m    606\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m signaled:\n\u001b[0;32m--> 607\u001b[0m     signaled \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_cond\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwait\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    608\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m signaled\n",
      "File \u001b[0;32m/usr/lib/python3.10/threading.py:324\u001b[0m, in \u001b[0;36mCondition.wait\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    322\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    323\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m timeout \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[0;32m--> 324\u001b[0m         gotit \u001b[38;5;241m=\u001b[39m \u001b[43mwaiter\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43macquire\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    325\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    326\u001b[0m         gotit \u001b[38;5;241m=\u001b[39m waiter\u001b[38;5;241m.\u001b[39macquire(\u001b[38;5;28;01mFalse\u001b[39;00m)\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "wandb.init(id=\"f5rpfdmw\", resume=\"must\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "258da017-0a41-4958-b705-01a8d7a06e0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "33ce26ac-99a9-430c-ac0c-e3d31d40fa39",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    }
   ],
   "source": [
    "!wandb init --reset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6ec77b4-b7cb-4d4a-b243-b1613e1a5926",
   "metadata": {},
   "outputs": [],
   "source": [
    "import transformers\n",
    "import numpy as np\n",
    "from torch.utils.data import Subset\n",
    "from datetime import datetime\n",
    "\n",
    "\n",
    "project = \"qa-finetune\"\n",
    "base_model_name = \"mistral-7b\"\n",
    "run_name = base_model_name + \"-\" + project\n",
    "output_dir = \"./\" + run_name\n",
    "\n",
    "\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "# eval_batch_size = 10\n",
    "# desired_eval_steps = 25\n",
    "# desired_eval_size = eval_batch_size * desired_eval_steps\n",
    "\n",
    "# subset_indices = np.random.choice(len(tokenized_val_dataset), desired_eval_size, replace=False)\n",
    "# eval_subset = Subset(tokenized_val_dataset, subset_indices)\n",
    "\n",
    "# resume_from_checkpoint = True\n",
    "trainer = transformers.Trainer(\n",
    "    model=model,\n",
    "    train_dataset=tokenized_train_dataset,\n",
    "    args=transformers.TrainingArguments(\n",
    "        output_dir=output_dir,\n",
    "        warmup_steps=5,\n",
    "        per_device_train_batch_size=2,\n",
    "        gradient_accumulation_steps=1,\n",
    "        gradient_checkpointing=True,\n",
    "        max_steps=500,\n",
    "        learning_rate=2.5e-5, # Want about 10x smaller than the Mistral learning rate\n",
    "        logging_steps=50,\n",
    "        bf16=True,\n",
    "        resume_from_checkpoint = True,\n",
    "        save_total_limit = 3,\n",
    "        load_best_model_at_end=True,\n",
    "        optim=\"paged_adamw_8bit\",\n",
    "        logging_dir=\"./logs\",        # Directory for storing logs\n",
    "        save_strategy=\"steps\",       # Save the model checkpoint every logging step\n",
    "        save_steps=200,                # Save checkpoints every 50 steps\n",
    "        # evaluation_strategy=\"steps\", # Evaluate the model every logging step\n",
    "        # eval_steps=50,               # Evaluate and save checkpoints every 50 steps\n",
    "        # do_eval=True,                # Perform evaluation at the end of training\n",
    "        report_to=\"wandb\",           # Comment this out if you don't want to use weights & baises\n",
    "        run_name=f\"{run_name}-{datetime.now().strftime('%Y-%m-%d-%H-%M')}\"          # Name of the W&B run (optional)\n",
    "    ),\n",
    "    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),\n",
    ")\n",
    "model.config.use_cache = False  # silence the warnings. Please re-enable for inference!\n",
    "trainer.train(resume_from_checkpoint=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "77b94ec0-375d-4535-9704-eea700db1b1f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/transformers/models/auto/auto_factory.py:472: FutureWarning: The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c39f37d7bb224b48ac7b17e48af10153",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from peft import PeftModel\n",
    "\n",
    "base_model = AutoModelForCausalLM.from_pretrained(\n",
    "    base_model_id,  # Mistral, same as before\n",
    "    quantization_config=bnb_config,  # Same quantization config as before\n",
    "    device_map=\"auto\",\n",
    "    trust_remote_code=True,\n",
    "    use_auth_token=True\n",
    ")\n",
    "tokenizer = AutoTokenizer.from_pretrained(base_model_id, trust_remote_code=True)\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "98ed6b84-0877-45d1-ab14-f29f03393bc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import PeftModel\n",
    "ft_model = PeftModel.from_pretrained(base_model, \"mistral-qa-finetune/checkpoint-9200\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ba59e64-7138-4c1f-8bcb-8976865cbddd",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "eval_prompt = \"\"\" \n",
    "\n",
    "    \n",
    "### Context : The deal includes more than $1 billion in new compensation and benefit plan funding, along with outsized gains to the traditional residuals formulas. It offers a new compensation model for performers working in streaming, with a substantial bonus on top of existing residuals structures, plus compensation escalation for principal and background actors. Additionally, the deal establishes detailed informed consent and compensation guardrails for the use of AI, hair and makeup equity, meaningful protections for the casting process, sexual harassment prevention protections and more.\n",
    "### Instructions : Generate diverse sets of questions and answers based on the context\n",
    "### Question :\n",
    "\"\"\"\n",
    "\n",
    "model_input = tokenizer(eval_prompt, return_tensors=\"pt\").to(\"cuda\")\n",
    "\n",
    "ft_model.eval()\n",
    "with torch.no_grad():\n",
    "    print(tokenizer.decode(ft_model.generate(**model_input, max_new_tokens=256, pad_token_id=2,)[0],repetition_penalty=2,skip_special_tokens=False))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}