day09/ft-llm-example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "be94e6d6-4096-4d1a-aa58-5afd89f33bff",
   "metadata": {},
   "source": [
    "# Fine-tuning Sandbox\n",
    "\n",
    "Code authored by: Shawhin Talebi <br>\n",
    "Video link: https://www.youtube.com/watch?v=eC6Hd1hFvos <br>\n",
    "Blog link: https://medium.com/towards-data-science/fine-tuning-large-language-models-llms-23473d763b91"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a265c214",
   "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"
     ]
    }
   ],
   "source": [
    "!pip install --quiet transformers peft accelerate datasets scikit-learn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4ef8ea85-d04d-4217-99a3-21c446bf2ffa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "from transformers import (\n",
    "    AutoTokenizer,\n",
    "    AutoModelForSequenceClassification,\n",
    "    DataCollatorWithPadding,\n",
    "    TrainingArguments,\n",
    "    Trainer)\n",
    "\n",
    "from peft import PeftModel, PeftConfig, get_peft_model, LoraConfig\n",
    "import evaluate\n",
    "import torch\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa6a4484-07d8-49dd-81ef-672105f53ebe",
   "metadata": {},
   "source": [
    "### dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fa9722d3-0609-4aea-9585-9aa2cfc1fc9a",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# # how dataset was generated\n",
    "\n",
    "# # load imdb data\n",
    "# imdb_dataset = load_dataset(\"imdb\")\n",
    "\n",
    "# # define subsample size\n",
    "# N = 1000 \n",
    "# # generate indexes for random subsample\n",
    "# rand_idx = np.random.randint(24999, size=N) \n",
    "\n",
    "# # extract train and test data\n",
    "# x_train = imdb_dataset['train'][rand_idx]['text']\n",
    "# y_train = imdb_dataset['train'][rand_idx]['label']\n",
    "\n",
    "# x_test = imdb_dataset['test'][rand_idx]['text']\n",
    "# y_test = imdb_dataset['test'][rand_idx]['label']\n",
    "\n",
    "# # create new dataset\n",
    "# dataset = DatasetDict({'train':Dataset.from_dict({'label':y_train,'text':x_train}),\n",
    "#                              'validation':Dataset.from_dict({'label':y_test,'text':x_test})})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "de226234-c521-4577-802c-0e7079ef4364",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['label', 'text'],\n",
       "        num_rows: 1000\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['label', 'text'],\n",
       "        num_rows: 1000\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load dataset\n",
    "dataset = load_dataset('shawhin/imdb-truncated')\n",
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d5625faa-5fea-4334-bd38-b77de983d8a8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.5"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# display % of training data with label=1\n",
    "np.array(dataset['train']['label']).sum()/len(dataset['train']['label'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3644c68d-9adf-48a4-90a2-8fd89555a302",
   "metadata": {},
   "source": [
    "### model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a60dd1fe-8144-4678-b018-20891e49237a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of DistilBertForSequenceClassification were not initialized from the model checkpoint at distilbert-base-uncased and are newly initialized: ['classifier.bias', 'classifier.weight', 'pre_classifier.bias', 'pre_classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "model_checkpoint = 'distilbert-base-uncased'\n",
    "# model_checkpoint = 'roberta-base' # you can alternatively use roberta-base but this model is bigger thus training will take longer\n",
    "\n",
    "# define label maps\n",
    "id2label = {0: \"Negative\", 1: \"Positive\"}\n",
    "label2id = {\"Negative\":0, \"Positive\":1}\n",
    "\n",
    "# generate classification model from model_checkpoint\n",
    "model = AutoModelForSequenceClassification.from_pretrained(\n",
    "    model_checkpoint, num_labels=2, id2label=id2label, label2id=label2id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "853002f8-d39c-4bc4-8d07-e44a47de3b47",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DistilBertForSequenceClassification(\n",
       "  (distilbert): DistilBertModel(\n",
       "    (embeddings): Embeddings(\n",
       "      (word_embeddings): Embedding(30522, 768, padding_idx=0)\n",
       "      (position_embeddings): Embedding(512, 768)\n",
       "      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (transformer): Transformer(\n",
       "      (layer): ModuleList(\n",
       "        (0-5): 6 x TransformerBlock(\n",
       "          (attention): MultiHeadSelfAttention(\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "            (q_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (k_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (v_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (out_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (sa_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "          (ffn): FFN(\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "            (lin1): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (lin2): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (activation): GELUActivation()\n",
       "          )\n",
       "          (output_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (pre_classifier): Linear(in_features=768, out_features=768, bias=True)\n",
       "  (classifier): Linear(in_features=768, out_features=2, bias=True)\n",
       "  (dropout): Dropout(p=0.2, inplace=False)\n",
       ")"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# display architecture\n",
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bc98609-873d-455c-bac4-155632cda484",
   "metadata": {},
   "source": [
    "### preprocess data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7fe08707-657f-4e66-aa72-84899c54bf8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, add_prefix_space=True)\n",
    "\n",
    "# add pad token if none exists\n",
    "if tokenizer.pad_token is None:\n",
    "    tokenizer.add_special_tokens({'pad_token': '[PAD]'})\n",
    "    model.resize_token_embeddings(len(tokenizer))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "20f4adb9-ce8f-4f54-9b94-300c9daae1b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create tokenize function\n",
    "def tokenize_function(examples):\n",
    "    # extract text\n",
    "    text = examples[\"text\"]\n",
    "\n",
    "    #tokenize and truncate text\n",
    "    tokenizer.truncation_side = \"left\"\n",
    "    tokenized_inputs = tokenizer(\n",
    "        text,\n",
    "        return_tensors=\"np\",\n",
    "        truncation=True,\n",
    "        max_length=512\n",
    "    )\n",
    "\n",
    "    return tokenized_inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "b7600bcd-7e93-4fb4-bd8d-ffc76bed1ac2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Map:   0%|          | 0/1000 [00:00<?, ? examples/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Map: 100%|██████████| 1000/1000 [00:00<00:00, 4898.04 examples/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['label', 'text', 'input_ids', 'attention_mask'],\n",
       "        num_rows: 1000\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['label', 'text', 'input_ids', 'attention_mask'],\n",
       "        num_rows: 1000\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# tokenize training and validation datasets\n",
    "tokenized_dataset = dataset.map(tokenize_function, batched=True)\n",
    "tokenized_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "3f8e85f9-1804-4f49-a783-4da59580ea1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create data collator\n",
    "data_collator = DataCollatorWithPadding(tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cd9a120-580d-470c-a981-7c7e22604865",
   "metadata": {},
   "source": [
    "### evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "2a894819-2e9c-4a53-9790-32130c182bca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import accuracy evaluation metric\n",
    "accuracy = evaluate.load(\"accuracy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c07b9be2-a3f6-4b38-b9e8-6a2bc8aa945a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define an evaluation function to pass into trainer later\n",
    "def compute_metrics(p):\n",
    "    predictions, labels = p\n",
    "    predictions = np.argmax(predictions, axis=1)\n",
    "\n",
    "    return {\"accuracy\": accuracy.compute(predictions=predictions, references=labels)}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47500035-a555-46e0-83dc-440586d96b7e",
   "metadata": {},
   "source": [
    "### Apply untrained model to text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "8f3761c1-a297-45c8-882e-d74856259810",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Untrained model predictions:\n",
      "----------------------------\n",
      "It was good. - Negative\n",
      "Not a fan, don't recommed. - Negative\n",
      "Better than the first one. - Negative\n",
      "This is not worth watching even once. - Negative\n",
      "This one is a pass. - Negative\n"
     ]
    }
   ],
   "source": [
    "# define list of examples\n",
    "text_list = [\"It was good.\", \"Not a fan, don't recommed.\", \"Better than the first one.\", \"This is not worth watching even once.\", \"This one is a pass.\"]\n",
    "\n",
    "print(\"Untrained model predictions:\")\n",
    "print(\"----------------------------\")\n",
    "for text in text_list:\n",
    "    # tokenize text\n",
    "    inputs = tokenizer.encode(text, return_tensors=\"pt\")\n",
    "    # compute logits\n",
    "    logits = model(inputs).logits\n",
    "    # convert logits to label\n",
    "    predictions = torch.argmax(logits)\n",
    "\n",
    "    print(text + \" - \" + id2label[predictions.tolist()])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff356f78-c9fd-4f2b-8f5b-097cf29c1c08",
   "metadata": {},
   "source": [
    "### Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "e4dde538-cd7f-4ab5-a96d-c30f3003822e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# see target modules from \"model\" printout\n",
    "peft_config = LoraConfig(task_type=\"SEQ_CLS\",\n",
    "                        r=4,\n",
    "                        lora_alpha=32,\n",
    "                        lora_dropout=0.01,\n",
    "                        # target_modules = ['q_lin'],\n",
    "                        target_modules = ['q_lin', 'k_lin', 'v_lin']   # better performance\n",
    "                        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "f1391303-1e16-4d5c-b2b4-799997eff9f8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LoraConfig(peft_type=<PeftType.LORA: 'LORA'>, auto_mapping=None, base_model_name_or_path=None, revision=None, task_type='SEQ_CLS', inference_mode=False, r=4, target_modules={'v_lin', 'q_lin', 'k_lin'}, lora_alpha=32, lora_dropout=0.01, fan_in_fan_out=False, bias='none', use_rslora=False, modules_to_save=None, init_lora_weights=True, layers_to_transform=None, layers_pattern=None, rank_pattern={}, alpha_pattern={}, megatron_config=None, megatron_core='megatron.core', loftq_config={}, use_dora=False, layer_replication=None, runtime_config=LoraRuntimeConfig(ephemeral_gpu_offload=False))"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "peft_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "3e0d9408-9fc4-4bd3-8d35-4d8217fe01e2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 702,722 || all params: 67,657,732 || trainable%: 1.0386\n"
     ]
    }
   ],
   "source": [
    "model = get_peft_model(model, peft_config)\n",
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "5db78059-e5ae-4807-89db-b58ef6abedd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# hyperparameters\n",
    "lr = 1e-3\n",
    "batch_size = 4\n",
    "num_epochs = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "9244ed55-65a4-4c66-8388-55efd87bceb8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/transformers/training_args.py:1525: FutureWarning: `evaluation_strategy` is deprecated and will be removed in version 4.46 of 🤗 Transformers. Use `eval_strategy` instead\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "# define training arguments\n",
    "training_args = TrainingArguments(\n",
    "    output_dir= model_checkpoint + \"-lora-text-classification\",\n",
    "    learning_rate=lr,\n",
    "    per_device_train_batch_size=batch_size,\n",
    "    per_device_eval_batch_size=batch_size,\n",
    "    num_train_epochs=num_epochs,\n",
    "    weight_decay=0.01,\n",
    "    evaluation_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "    load_best_model_at_end=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "fc8bc705-5dd7-4305-a797-399b2b0fa2c7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/_functions.py:68: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='630' max='630' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [630/630 03:09, Epoch 10/10]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Accuracy</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>No log</td>\n",
       "      <td>0.262720</td>\n",
       "      <td>{'accuracy': 0.902}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>No log</td>\n",
       "      <td>0.260622</td>\n",
       "      <td>{'accuracy': 0.897}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>No log</td>\n",
       "      <td>0.333514</td>\n",
       "      <td>{'accuracy': 0.892}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>No log</td>\n",
       "      <td>0.440060</td>\n",
       "      <td>{'accuracy': 0.899}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5</td>\n",
       "      <td>No log</td>\n",
       "      <td>0.574066</td>\n",
       "      <td>{'accuracy': 0.9}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6</td>\n",
       "      <td>No log</td>\n",
       "      <td>0.577979</td>\n",
       "      <td>{'accuracy': 0.904}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>7</td>\n",
       "      <td>No log</td>\n",
       "      <td>0.654253</td>\n",
       "      <td>{'accuracy': 0.901}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>8</td>\n",
       "      <td>0.133900</td>\n",
       "      <td>0.713301</td>\n",
       "      <td>{'accuracy': 0.908}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>9</td>\n",
       "      <td>0.133900</td>\n",
       "      <td>0.694802</td>\n",
       "      <td>{'accuracy': 0.905}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>0.133900</td>\n",
       "      <td>0.691969</td>\n",
       "      <td>{'accuracy': 0.907}</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/_functions.py:68: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n",
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/_functions.py:68: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n",
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/_functions.py:68: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n",
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/_functions.py:68: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n",
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/_functions.py:68: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n",
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/_functions.py:68: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n",
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/_functions.py:68: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n",
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/_functions.py:68: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n",
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/_functions.py:68: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n",
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
      "/root/anaconda3/envs/dlbootcamp/lib/python3.11/site-packages/torch/nn/parallel/_functions.py:68: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=630, training_loss=0.10742026841829694, metrics={'train_runtime': 189.8528, 'train_samples_per_second': 52.672, 'train_steps_per_second': 3.318, 'total_flos': 1328711824240320.0, 'train_loss': 0.10742026841829694, 'epoch': 10.0})"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# creater trainer object\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=tokenized_dataset[\"train\"],\n",
    "    eval_dataset=tokenized_dataset[\"validation\"],\n",
    "    tokenizer=tokenizer,\n",
    "    data_collator=data_collator, # this will dynamically pad examples in each batch to be equal length\n",
    "    compute_metrics=compute_metrics,\n",
    ")\n",
    "\n",
    "# train model\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f5664d1-9bd2-4ce1-bc24-cab5adf80f49",
   "metadata": {},
   "source": [
    "### Generate prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "e5dc029e-1c16-491d-a3f1-715f9e0adf52",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Trained model predictions:\n",
      "--------------------------\n",
      "It was good. - Positive\n",
      "Not a fan, don't recommed. - Negative\n",
      "Better than the first one. - Positive\n",
      "This is not worth watching even once. - Negative\n",
      "This one is a pass. - Positive\n"
     ]
    }
   ],
   "source": [
    "device = \"cuda\"   # mps for Mac, or \"cpu\" / \"cuda\" for other machines\n",
    "\n",
    "model.to(device) \n",
    "\n",
    "print(\"Trained model predictions:\")\n",
    "print(\"--------------------------\")\n",
    "for text in text_list:\n",
    "    inputs = tokenizer.encode(text, return_tensors=\"pt\").to(device) # moving to mps for Mac (can alternatively do 'cpu')\n",
    "\n",
    "    logits = model(inputs).logits\n",
    "    predictions = torch.max(logits,1).indices\n",
    "\n",
    "    print(text + \" - \" + id2label[predictions.tolist()[0]])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c084bd9e-f7b1-4979-b753-73335ee0cede",
   "metadata": {},
   "source": [
    "### Optional: push model to hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "159eb49a-dd0d-4c9e-b9ab-27e06585fd84",
   "metadata": {},
   "outputs": [],
   "source": [
    "# option 1: notebook login\n",
    "from huggingface_hub import notebook_login\n",
    "notebook_login() # ensure token gives write access\n",
    "\n",
    "# # option 2: key login\n",
    "# from huggingface_hub import login\n",
    "# write_key = 'hf_' # paste token here\n",
    "# login(write_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09496307-e253-47e3-a46f-3f28a84c89a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "hf_name = 'shawhin' # your hf username or org name\n",
    "model_id = hf_name + \"/\" + model_checkpoint + \"-lora-text-classification\" # you can name the model whatever you want"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c56ea581-0ea3-45f3-af21-362e9093ee37",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.push_to_hub(model_id) # save model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f487331a-8552-4fb2-867f-985b8fe1d1ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.push_to_hub(model_id) # save trainer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00e7feaa-b70e-4b1d-a118-23c616d14639",
   "metadata": {},
   "source": [
    "### Optional: load peft model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19cffa01-25a4-4c86-a7fa-a84353b8caae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# how to load peft model from hub for inference\n",
    "config = PeftConfig.from_pretrained(model_id)\n",
    "inference_model = AutoModelForSequenceClassification.from_pretrained(\n",
    "    config.base_model_name_or_path, num_labels=2, id2label=id2label, label2id=label2id\n",
    ")\n",
    "tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)\n",
    "model = PeftModel.from_pretrained(inference_model, model_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77c6ed42-8ec3-4343-9e42-405feac052ba",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}