train_cairo_mistral_7b_instruct.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "pip install -q git+https://github.com/huggingface/trl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The model that you want to train from the Hugging Face hub\n",
    "model_name = \"mistralai/Mistral-7B-Instruct-v0.2\"\n",
    "\n",
    "# The instruction dataset to use\n",
    "dataset_name = \"Conjuror64/C64-inst-format\"\n",
    "\n",
    "# Fine-tuned model name\n",
    "new_model = \"Conjuror64/mistral-C64Wizard-PEFT\"\n",
    "\n",
    "\n",
    "hub_name = \"pechaut/Mistral-C64Wizard-instruct\"\n",
    "\n",
    "import os\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "\n",
    "\n",
    "\n",
    "max_steps = 1000 # to tweak to get the best out of the model \n",
    "nb_epochs = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mpierre-emmanuel-chaut\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.17.0"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>c:\\projets\\Conjuror64\\wandb\\run-20240522_165535-jd5jkfqx</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/pierre-emmanuel-chaut/C64Wizard/runs/jd5jkfqx' target=\"_blank\">misty-spaceship-12</a></strong> to <a href='https://wandb.ai/pierre-emmanuel-chaut/C64Wizard' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/pierre-emmanuel-chaut/C64Wizard' target=\"_blank\">https://wandb.ai/pierre-emmanuel-chaut/C64Wizard</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/pierre-emmanuel-chaut/C64Wizard/runs/jd5jkfqx' target=\"_blank\">https://wandb.ai/pierre-emmanuel-chaut/C64Wizard/runs/jd5jkfqx</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<button onClick=\"this.nextSibling.style.display='block';this.style.display='none';\">Display W&B run</button><iframe src='https://wandb.ai/pierre-emmanuel-chaut/C64Wizard/runs/jd5jkfqx?jupyter=true' style='border:none;width:100%;height:420px;display:none;'></iframe>"
      ],
      "text/plain": [
       "<wandb.sdk.wandb_run.Run at 0x1a23f745d60>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import wandb\n",
    "import random\n",
    "\n",
    "# start a new wandb run to track this script\n",
    "wandb.init(\n",
    "    # set the wandb project where this run will be logged\n",
    "    project=\"C64Wizard\",\n",
    "    \n",
    "    # track hyperparameters and run metadata\n",
    "    config={\n",
    "\n",
    "    \"epochs\":nb_epochs,\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\pe_ch\\anaconda3\\envs\\cairo-llm\\Lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "c:\\Users\\pe_ch\\anaconda3\\envs\\cairo-llm\\Lib\\site-packages\\transformers\\models\\auto\\tokenization_auto.py:769: FutureWarning: The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.\n",
      "  warnings.warn(\n",
      "Loading checkpoint shards: 100%|██████████| 3/3 [00:10<00:00,  3.37s/it]\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from datasets import load_dataset, Dataset\n",
    "from peft import LoraConfig, AutoPeftModelForCausalLM\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig\n",
    "from trl import SFTTrainer\n",
    "import os\n",
    " \n",
    "# Load the tokenizer from the model (llama2)\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True, add_eos_token=True, use_fast=False)\n",
    "tokenizer.padding_side = \"right\"\n",
    "tokenizer.pad_token_id = 18610\n",
    "\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    "    bnb_4bit_compute_dtype=torch.bfloat16,\n",
    "    bnb_4bit_use_double_quant=True,\n",
    ")\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=model_name,\n",
    "                                             trust_remote_code=True,\n",
    "                                             quantization_config=bnb_config,\n",
    "                                             low_cpu_mem_usage=True,\n",
    "                                             device_map=\"auto\"\n",
    "                                             )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset, Dataset\n",
    "from peft import LoraConfig, AutoPeftModelForCausalLM, prepare_model_for_kbit_training, get_peft_model\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments\n",
    "from trl import SFTTrainer\n",
    "import os\n",
    "\n",
    "model.resize_token_embeddings(len(tokenizer))\n",
    "model.config.use_cache=False\n",
    "model.config.pretraining_tp=1\n",
    "model.config.window = 512 \n",
    "model.gradient_checkpointing_enable()\n",
    "model = prepare_model_for_kbit_training(model)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loading Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\pe_ch\\anaconda3\\envs\\cairo-llm\\Lib\\site-packages\\datasets\\load.py:2555: FutureWarning: 'ignore_verifications' was deprecated in favor of 'verification_mode' in version 2.9.1 and will be removed in 3.0.0.\n",
      "You can remove this warning by passing 'verification_mode=no_checks' instead.\n",
      "  warnings.warn(\n",
      "Downloading readme: 100%|██████████| 437/437 [00:00<00:00, 437kB/s]\n",
      "Downloading data: 100%|██████████| 1.10M/1.10M [00:00<00:00, 2.97MB/s]\n",
      "Generating train split: 100%|██████████| 8798/8798 [00:00<00:00, 324985.79 examples/s]\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "\n",
    "# Load the dataset\n",
    "dataset = load_dataset(dataset_name, split=\"train\", download_mode='force_redownload',ignore_verifications=True)\n",
    "split_dataset = dataset.train_test_split(test_size=0.001)\n",
    "test_dataset = split_dataset[\"test\"]\n",
    "train_dataset = split_dataset[\"train\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\pe_ch\\anaconda3\\envs\\cairo-llm\\Lib\\site-packages\\transformers\\training_args.py:1474: FutureWarning: `evaluation_strategy` is deprecated and will be removed in version 4.46 of 🤗 Transformers. Use `eval_strategy` instead\n",
      "  warnings.warn(\n",
      "c:\\Users\\pe_ch\\anaconda3\\envs\\cairo-llm\\Lib\\site-packages\\trl\\trainer\\sft_trainer.py:246: UserWarning: You didn't pass a `max_seq_length` argument to the SFTTrainer, this will default to 1024\n",
      "  warnings.warn(\n",
      "c:\\Users\\pe_ch\\anaconda3\\envs\\cairo-llm\\Lib\\site-packages\\trl\\trainer\\sft_trainer.py:257: UserWarning: You passed a `neftune_noise_alpha` argument to the SFTTrainer, the value you passed will override the one in the `TrainingArguments`.\n",
      "  warnings.warn(\n",
      "Map: 100%|██████████| 8789/8789 [00:02<00:00, 4056.76 examples/s]\n",
      "Map: 100%|██████████| 9/9 [00:00<00:00, 598.58 examples/s]\n"
     ]
    }
   ],
   "source": [
    "\n",
    "\n",
    "\n",
    "peft_config = LoraConfig(\n",
    "    r=64, lora_alpha=16, lora_dropout=0.1, bias=\"none\", task_type=\"CAUSAL_LM\", target_modules=[\n",
    "        \"q_proj\",\n",
    "        \"k_proj\",\n",
    "        \"v_proj\",\n",
    "        \"o_proj\",\n",
    "        \"gate_proj\",\n",
    "\n",
    "    ],inference_mode = False\n",
    ")\n",
    "\n",
    "\n",
    "\n",
    "model = get_peft_model(model, peft_config)\n",
    "training_arguments = TrainingArguments(\n",
    "    output_dir=new_model,\n",
    "    per_device_train_batch_size=2,\n",
    "    gradient_accumulation_steps=4,\n",
    "    gradient_checkpointing = True,\n",
    "\n",
    "    evaluation_strategy=\"steps\",\n",
    "    learning_rate=2e-5,\n",
    "    lr_scheduler_type=\"constant\",\n",
    "    warmup_ratio=0.03,\n",
    "    max_grad_norm=0.3,\n",
    "    save_strategy=\"epoch\",\n",
    "    logging_dir=\"./logs\", \n",
    "    logging_steps=50,\n",
    "    num_train_epochs=nb_epochs,\n",
    "    group_by_length=True,\n",
    "    fp16=False,\n",
    "    report_to=\"wandb\",\n",
    "    push_to_hub=True,\n",
    "    adam_beta2=0.999,\n",
    "         do_train=True,\n",
    ")\n",
    "\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=test_dataset,\n",
    "    peft_config=peft_config,\n",
    "    dataset_text_field=\"text\",\n",
    "    args=training_arguments,\n",
    "    tokenizer=tokenizer,\n",
    "    packing=False,\n",
    "    max_seq_length=None,\n",
    "    neftune_noise_alpha=5\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "trainer.train()\n",
    "trainer.model.push_to_hub(new_model)\n",
    "tokenizer.push_to_hub(new_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "If you want to merge and push now.. done on CPU\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bd9cb112964d4716b453167c881632a3",
       "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": [
    "from huggingface_hub import login, logout\n",
    "login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6b0ca794c04040ef9de8b36bf374dd2f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.model:   0%|          | 0.00/493k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/pechaut/Mistral-C64Wizard-instruct-PEFT/commit/6befc980c8675c43c88361a622c72d7e468ff3fd', commit_message='Upload tokenizer', commit_description='', oid='6befc980c8675c43c88361a622c72d7e468ff3fd', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments\n",
    "import os\n",
    "new_model = \"Conjuror64/mistral-C64Wizard-PEFT\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(new_model, use_fast=True)\n",
    "new_model= \"pechaut/Mistral-C64Wizard-instruct-PEFT\"\n",
    "\n",
    "tokenizer.push_to_hub(new_model)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7b09609448bc47368b120a8889f77dee",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fcb427583bb047f09be85e5f5941a222",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "README.md:   0%|          | 0.00/5.17k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\pe_ch\\anaconda3\\envs\\cairo-llm\\Lib\\site-packages\\huggingface_hub\\file_download.py:157: UserWarning: `huggingface_hub` cache-system uses symlinks by default to efficiently store duplicated files but your machine does not support them in C:\\Users\\pe_ch\\.cache\\huggingface\\hub\\models--pechaut--Mistral-C64Wizard-instruct-PEFT. Caching files will still work but in a degraded version that might require more space on your disk. This warning can be disabled by setting the `HF_HUB_DISABLE_SYMLINKS_WARNING` environment variable. For more details, see https://huggingface.co/docs/huggingface_hub/how-to-cache#limitations.\n",
      "To support symlinks on Windows, you either need to activate Developer Mode or to run Python as an administrator. In order to see activate developer mode, see this article: https://docs.microsoft.com/en-us/windows/apps/get-started/enable-your-device-for-development\n",
      "  warnings.warn(message)\n",
      "c:\\Users\\pe_ch\\anaconda3\\envs\\cairo-llm\\Lib\\site-packages\\transformers\\integrations\\peft.py:399: FutureWarning: The `active_adapter` method is deprecated and will be removed in a future version.\n",
      "  warnings.warn(\n",
      "c:\\Users\\pe_ch\\anaconda3\\envs\\cairo-llm\\Lib\\site-packages\\huggingface_hub\\file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "37accf4c54d8486fa9127008c19a27d2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "adapter_model.safetensors:   0%|          | 0.00/369M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/pechaut/Mistral-C64Wizard-instruct-PEFT/commit/bcb85e18e31601805085364b65dbcd02520fae48', commit_message='Upload MistralForCausalLM', commit_description='', oid='bcb85e18e31601805085364b65dbcd02520fae48', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from datasets import load_dataset, Dataset\n",
    "from peft import LoraConfig, AutoPeftModelForCausalLM\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig\n",
    "\n",
    "new_model = \"Conjuror64/mistral-C64Wizard-PEFT\"\n",
    "model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=new_model,\n",
    "                                             trust_remote_code=True,\n",
    "                                             )\n",
    "new_model= \"pechaut/Mistral-C64Wizard-instruct-PEFT\"\n",
    "model.push_to_hub(new_model,max_shard_size=\"1GB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a0f19da1c2d14352ac8e1e5670009a66",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7078b5559a8c406f872f5c410b6cb354",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00002-of-00015.safetensors:   0%|          | 0.00/990M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bf98e061fd624f58926a8f3ded0d87b2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00001-of-00015.safetensors:   0%|          | 0.00/900M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c91b326a80e141e9b11f1d5f23b2f6b8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00003-of-00015.safetensors:   0%|          | 0.00/990M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ba09a7e5ed08412198190f0f99626836",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00004-of-00015.safetensors:   0%|          | 0.00/956M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1360168592334763b7bf5d32b5a7efd9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00005-of-00015.safetensors:   0%|          | 0.00/990M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4aef5c1ea5014f13ad855929f0d4367d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Upload 15 LFS files:   0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f5b28c686b0b45e0ae82be04c5849a2b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00006-of-00015.safetensors:   0%|          | 0.00/990M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1675f243084249f382f75af411338537",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00007-of-00015.safetensors:   0%|          | 0.00/990M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "37ddc9ddb5ab4ec6816cd10599d55218",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00008-of-00015.safetensors:   0%|          | 0.00/956M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d3419d9d31b94f30b42b27532bfdb3e8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00009-of-00015.safetensors:   0%|          | 0.00/990M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a28bc88be93e46cdbb7a4f7e84522811",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00010-of-00015.safetensors:   0%|          | 0.00/990M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bd8c317b32fb42c7bd53e79b39c71e02",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00011-of-00015.safetensors:   0%|          | 0.00/990M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7c5a7c5a2e1e498b8e8f4a993905330c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00012-of-00015.safetensors:   0%|          | 0.00/956M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2f15aaa0224a47b0855c2a2a8aa869df",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00013-of-00015.safetensors:   0%|          | 0.00/990M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "972b54012a7c4a63b029c5b0064542ed",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00014-of-00015.safetensors:   0%|          | 0.00/990M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "42bdf48680834ff686b0746b7a0afd92",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00015-of-00015.safetensors:   0%|          | 0.00/816M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/pechaut/Mistral-C64Wizard-instruct/commit/d307acf721b4c2d18ffd1057be14bc76135f78df', commit_message='Upload MistralForCausalLM', commit_description='', oid='d307acf721b4c2d18ffd1057be14bc76135f78df', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import AutoModelForCausalLM\n",
    "from peft import PeftModel\n",
    "import torch\n",
    "from transformers import  AutoTokenizer, BitsAndBytesConfig\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    "    bnb_4bit_compute_dtype=\"float16\",\n",
    "    bnb_4bit_use_double_quant=True,\n",
    ")\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=model_name,\n",
    "                                             trust_remote_code=True,\n",
    "                                             low_cpu_mem_usage=True,\n",
    "                                             device_map={\"\": \"cpu\"},\n",
    "                                             torch_dtype=torch.float16\n",
    "                                             )\n",
    "\n",
    "model_to_merge  = PeftModel.from_pretrained(model, new_model,\n",
    "                        torch_dtype=torch.float16, \n",
    "                        device_map={\"\": \"cpu\"}\n",
    "                         )\n",
    "merged_model = model_to_merge.merge_and_unload()\n",
    "#model.save_pretrained(\"cairo-mistral\")\n",
    "merged_model.push_to_hub(hub_name,max_shard_size=\"1GB\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "63b4491c3f3f4376b13b65240e952ae5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)lve/main/config.json:   0%|          | 0.00/622 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "88d6db4a4d77469990482fcd2eed1257",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)model.bin.index.json:   0%|          | 0.00/23.9k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e51f678740d54344ac0f02ec48d0f49f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading shards:   0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "216ab1d9effc44ee8b6391586574999f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)l-00001-of-00015.bin:   0%|          | 0.00/900M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9c4574b91131489ba5c28e8c98614520",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)l-00002-of-00015.bin:   0%|          | 0.00/990M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0145159dd4594623814a5c83ac9fa6a6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)l-00003-of-00015.bin:   0%|          | 0.00/990M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0b1ca65897a74956bc67251fcce9f80d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)l-00004-of-00015.bin:   0%|          | 0.00/956M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "030d5cc318c941e2a05b2664af97e49d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)l-00005-of-00015.bin:   0%|          | 0.00/990M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b4d918bb13794757b972198389193a31",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)l-00006-of-00015.bin:   0%|          | 0.00/990M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1e8f203d7eb6497fad714b0b5b9f6bdc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)l-00007-of-00015.bin:   0%|          | 0.00/990M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d243e57a2d204a07a383e6e8edf5574a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)l-00008-of-00015.bin:   0%|          | 0.00/956M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "62ba3ef714d84fa59e72c24dbc640b84",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)l-00009-of-00015.bin:   0%|          | 0.00/990M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c4eca5e3a2f141f8a4f81affc2b0e3f6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)l-00010-of-00015.bin:   0%|          | 0.00/990M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ac402e9535d344b094f5700ac35e654a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)l-00011-of-00015.bin:   0%|          | 0.00/990M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8811856256b140bba5aee95605c380da",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)l-00012-of-00015.bin:   0%|          | 0.00/956M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5eca8f9d7b214f43954f96271e6e15bb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)l-00013-of-00015.bin:   0%|          | 0.00/990M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "25f3e881125c4a939edf2ff86f364c15",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)l-00014-of-00015.bin:   0%|          | 0.00/990M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "331449bd3d4743b79c5dcaf36a1c7667",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)l-00015-of-00015.bin:   0%|          | 0.00/816M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c266eced910b4980b8de8463f3298e19",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Attention Sinks] Injected Position Shifting into 32 attention classes.\n",
      "[Attention Sinks] Injected Attention Sink KV Cache into 1 model class.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "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): Linear4bit(in_features=4096, out_features=4096, bias=False)\n",
       "          (k_proj): Linear4bit(in_features=4096, out_features=1024, bias=False)\n",
       "          (v_proj): Linear4bit(in_features=4096, out_features=1024, bias=False)\n",
       "          (o_proj): Linear4bit(in_features=4096, out_features=4096, bias=False)\n",
       "          (rotary_emb): MistralRotaryEmbedding()\n",
       "        )\n",
       "        (mlp): MistralMLP(\n",
       "          (gate_proj): Linear4bit(in_features=4096, out_features=14336, bias=False)\n",
       "          (up_proj): Linear4bit(in_features=4096, out_features=14336, bias=False)\n",
       "          (down_proj): Linear4bit(in_features=14336, out_features=4096, bias=False)\n",
       "          (act_fn): SiLUActivation()\n",
       "        )\n",
       "        (input_layernorm): MistralRMSNorm()\n",
       "        (post_attention_layernorm): MistralRMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (norm): MistralRMSNorm()\n",
       "  )\n",
       "  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, TextStreamer, GenerationConfig, BitsAndBytesConfig\n",
    "from attention_sinks import AutoModelForCausalLM\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)\n",
    "tokenizer.pad_token = tokenizer.unk_token\n",
    "\n",
    "\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    "    bnb_4bit_compute_dtype=\"float16\",\n",
    "    bnb_4bit_use_double_quant=True,\n",
    ")\n",
    "model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=hub_name,\n",
    "                                             trust_remote_code=True,\n",
    "                                             device_map={\"\": 0},\n",
    "                                             attention_sink_size=4,\n",
    "                                             quantization_config=bnb_config,\n",
    "                                            attention_sink_window_size=252, # <- Low for the sake of faster generation\n",
    "                                             )\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "quantization_config: \n",
    "load_in_8bit: False \n",
    "load_in_4bit: True \n",
    "llm_int8_threshold: 6.0 \n",
    "llm_int8_skip_modules: None \n",
    "llm_int8_enable_fp32_cpu_offload: False \n",
    "llm_int8_has_fp16_weight: False \n",
    "bnb_4bit_quant_type: \"nf4\" \n",
    "bnb_4bit_use_double_quant: False \n",
    "bnb_4bit_compute_dtype: \"float16\""
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cairo-code",
   "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.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}