prefix_tuning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/xzrderek/.local/lib/python3.11/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"
     ]
    }
   ],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, default_data_collator, get_linear_schedule_with_warmup, BitsAndBytesConfig, AutoModelForCausalLM\n",
    "from peft import get_peft_config, get_peft_model, get_peft_model_state_dict, PrefixTuningConfig, TaskType\n",
    "from datasets import load_dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from tqdm import tqdm\n",
    "import torch\n",
    "import os\n",
    "\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0, 1\"\n",
    "\n",
    "device = \"cuda\"\n",
    "# model_name_or_path =\"tiiuae/falcon-7b-instruct\"\n",
    "# tokenizer_name_or_path = \"tiiuae/falcon-7b-instruct\"\n",
    "\n",
    "text_column = \"sentence\"\n",
    "label_column = \"text_label\"\n",
    "max_length = 128\n",
    "lr = 1e-2\n",
    "num_epochs = 5\n",
    "batch_size = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading checkpoint shards: 100%|██████████| 2/2 [00:14<00:00,  7.35s/it]\n"
     ]
    }
   ],
   "source": [
    "from peft import prepare_model_for_kbit_training\n",
    "\n",
    "model_id = \"tiiuae/falcon-7b-instruct\"\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",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map=\"auto\")\n",
    "\n",
    "model.gradient_checkpointing_enable()\n",
    "model = prepare_model_for_kbit_training(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "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": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 5,816,320 || all params: 6,927,537,024 || trainable%: 0.0839594213621629\n"
     ]
    }
   ],
   "source": [
    "peft_config = PrefixTuningConfig(task_type=TaskType.CAUSAL_LM, inference_mode=False, num_virtual_tokens=20)\n",
    "\n",
    "model = get_peft_model(model, peft_config)\n",
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DatasetDict({\n",
      "    train: Dataset({\n",
      "        features: ['SpongeBob Lines', 'input_ids', 'attention_mask'],\n",
      "        num_rows: 9663\n",
      "    })\n",
      "    test: Dataset({\n",
      "        features: ['SpongeBob Lines', 'input_ids', 'attention_mask'],\n",
      "        num_rows: 2416\n",
      "    })\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "dataset = load_dataset('csv', data_files='data/spongebob_lines2.csv')\n",
    "dataset = dataset.map(lambda samples: tokenizer(samples[\"SpongeBob Lines\"]), batched=True)\n",
    "dataset = dataset['train'].train_test_split(test_size=0.2)\n",
    "\n",
    "print(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_function(examples):\n",
    "    # print(examples)\n",
    "    inputs = examples[\"SpongeBob Lines\"]\n",
    "    # targets = examples[label_column]\n",
    "    model_inputs = tokenizer(inputs, max_length=max_length, padding=\"max_length\", truncation=True, return_tensors=\"pt\")\n",
    "    # labels = tokenizer(targets, max_length=2, padding=\"max_length\", truncation=True, return_tensors=\"pt\")\n",
    "    # labels = labels[\"input_ids\"]\n",
    "    # labels[labels == tokenizer.pad_token_id] = -100\n",
    "    # model_inputs[\"labels\"] = labels\n",
    "    return model_inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Running tokenizer on dataset:   0%|          | 0/9663 [00:00<?, ? examples/s]\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "Asking to pad but the tokenizer does not have a padding token. Please select a token to use as `pad_token` `(tokenizer.pad_token = tokenizer.eos_token e.g.)` or add a new pad token via `tokenizer.add_special_tokens({'pad_token': '[PAD]'})`.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[1;32m/home/xzrderek/cs182-project/prefix_tuning.ipynb Cell 7\u001b[0m line \u001b[0;36m1\n\u001b[0;32m----> <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W1sdnNjb2RlLXJlbW90ZQ%3D%3D?line=0'>1</a>\u001b[0m processed_datasets \u001b[39m=\u001b[39m dataset\u001b[39m.\u001b[39;49mmap(\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W1sdnNjb2RlLXJlbW90ZQ%3D%3D?line=1'>2</a>\u001b[0m     preprocess_function,\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W1sdnNjb2RlLXJlbW90ZQ%3D%3D?line=2'>3</a>\u001b[0m     batched\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m,\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W1sdnNjb2RlLXJlbW90ZQ%3D%3D?line=3'>4</a>\u001b[0m     num_proc\u001b[39m=\u001b[39;49m\u001b[39m1\u001b[39;49m,\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W1sdnNjb2RlLXJlbW90ZQ%3D%3D?line=4'>5</a>\u001b[0m     remove_columns\u001b[39m=\u001b[39;49mdataset[\u001b[39m\"\u001b[39;49m\u001b[39mtrain\u001b[39;49m\u001b[39m\"\u001b[39;49m]\u001b[39m.\u001b[39;49mcolumn_names,\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W1sdnNjb2RlLXJlbW90ZQ%3D%3D?line=5'>6</a>\u001b[0m     load_from_cache_file\u001b[39m=\u001b[39;49m\u001b[39mFalse\u001b[39;49;00m,\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W1sdnNjb2RlLXJlbW90ZQ%3D%3D?line=6'>7</a>\u001b[0m     desc\u001b[39m=\u001b[39;49m\u001b[39m\"\u001b[39;49m\u001b[39mRunning tokenizer on dataset\u001b[39;49m\u001b[39m\"\u001b[39;49m,\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W1sdnNjb2RlLXJlbW90ZQ%3D%3D?line=7'>8</a>\u001b[0m )\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W1sdnNjb2RlLXJlbW90ZQ%3D%3D?line=9'>10</a>\u001b[0m train_dataset \u001b[39m=\u001b[39m processed_datasets[\u001b[39m\"\u001b[39m\u001b[39mtrain\u001b[39m\u001b[39m\"\u001b[39m]\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W1sdnNjb2RlLXJlbW90ZQ%3D%3D?line=10'>11</a>\u001b[0m eval_dataset \u001b[39m=\u001b[39m processed_datasets[\u001b[39m\"\u001b[39m\u001b[39mtest\u001b[39m\u001b[39m\"\u001b[39m]\n",
      "File \u001b[0;32m~/miniconda3/envs/kibsi/lib/python3.11/site-packages/datasets/dataset_dict.py:855\u001b[0m, in \u001b[0;36mDatasetDict.map\u001b[0;34m(self, function, with_indices, with_rank, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_names, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, desc)\u001b[0m\n\u001b[1;32m    852\u001b[0m \u001b[39mif\u001b[39;00m cache_file_names \u001b[39mis\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m    853\u001b[0m     cache_file_names \u001b[39m=\u001b[39m {k: \u001b[39mNone\u001b[39;00m \u001b[39mfor\u001b[39;00m k \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m}\n\u001b[1;32m    854\u001b[0m \u001b[39mreturn\u001b[39;00m DatasetDict(\n\u001b[0;32m--> 855\u001b[0m     {\n\u001b[1;32m    856\u001b[0m         k: dataset\u001b[39m.\u001b[39;49mmap(\n\u001b[1;32m    857\u001b[0m             function\u001b[39m=\u001b[39;49mfunction,\n\u001b[1;32m    858\u001b[0m             with_indices\u001b[39m=\u001b[39;49mwith_indices,\n\u001b[1;32m    859\u001b[0m             with_rank\u001b[39m=\u001b[39;49mwith_rank,\n\u001b[1;32m    860\u001b[0m             input_columns\u001b[39m=\u001b[39;49minput_columns,\n\u001b[1;32m    861\u001b[0m             batched\u001b[39m=\u001b[39;49mbatched,\n\u001b[1;32m    862\u001b[0m             batch_size\u001b[39m=\u001b[39;49mbatch_size,\n\u001b[1;32m    863\u001b[0m             drop_last_batch\u001b[39m=\u001b[39;49mdrop_last_batch,\n\u001b[1;32m    864\u001b[0m             remove_columns\u001b[39m=\u001b[39;49mremove_columns,\n\u001b[1;32m    865\u001b[0m             keep_in_memory\u001b[39m=\u001b[39;49mkeep_in_memory,\n\u001b[1;32m    866\u001b[0m             load_from_cache_file\u001b[39m=\u001b[39;49mload_from_cache_file,\n\u001b[1;32m    867\u001b[0m             cache_file_name\u001b[39m=\u001b[39;49mcache_file_names[k],\n\u001b[1;32m    868\u001b[0m             writer_batch_size\u001b[39m=\u001b[39;49mwriter_batch_size,\n\u001b[1;32m    869\u001b[0m             features\u001b[39m=\u001b[39;49mfeatures,\n\u001b[1;32m    870\u001b[0m             disable_nullable\u001b[39m=\u001b[39;49mdisable_nullable,\n\u001b[1;32m    871\u001b[0m             fn_kwargs\u001b[39m=\u001b[39;49mfn_kwargs,\n\u001b[1;32m    872\u001b[0m             num_proc\u001b[39m=\u001b[39;49mnum_proc,\n\u001b[1;32m    873\u001b[0m             desc\u001b[39m=\u001b[39;49mdesc,\n\u001b[1;32m    874\u001b[0m         )\n\u001b[1;32m    875\u001b[0m         \u001b[39mfor\u001b[39;49;00m k, dataset \u001b[39min\u001b[39;49;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mitems()\n\u001b[1;32m    876\u001b[0m     }\n\u001b[1;32m    877\u001b[0m )\n",
      "File \u001b[0;32m~/miniconda3/envs/kibsi/lib/python3.11/site-packages/datasets/dataset_dict.py:856\u001b[0m, in \u001b[0;36m<dictcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m    852\u001b[0m \u001b[39mif\u001b[39;00m cache_file_names \u001b[39mis\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m    853\u001b[0m     cache_file_names \u001b[39m=\u001b[39m {k: \u001b[39mNone\u001b[39;00m \u001b[39mfor\u001b[39;00m k \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m}\n\u001b[1;32m    854\u001b[0m \u001b[39mreturn\u001b[39;00m DatasetDict(\n\u001b[1;32m    855\u001b[0m     {\n\u001b[0;32m--> 856\u001b[0m         k: dataset\u001b[39m.\u001b[39;49mmap(\n\u001b[1;32m    857\u001b[0m             function\u001b[39m=\u001b[39;49mfunction,\n\u001b[1;32m    858\u001b[0m             with_indices\u001b[39m=\u001b[39;49mwith_indices,\n\u001b[1;32m    859\u001b[0m             with_rank\u001b[39m=\u001b[39;49mwith_rank,\n\u001b[1;32m    860\u001b[0m             input_columns\u001b[39m=\u001b[39;49minput_columns,\n\u001b[1;32m    861\u001b[0m             batched\u001b[39m=\u001b[39;49mbatched,\n\u001b[1;32m    862\u001b[0m             batch_size\u001b[39m=\u001b[39;49mbatch_size,\n\u001b[1;32m    863\u001b[0m             drop_last_batch\u001b[39m=\u001b[39;49mdrop_last_batch,\n\u001b[1;32m    864\u001b[0m             remove_columns\u001b[39m=\u001b[39;49mremove_columns,\n\u001b[1;32m    865\u001b[0m             keep_in_memory\u001b[39m=\u001b[39;49mkeep_in_memory,\n\u001b[1;32m    866\u001b[0m             load_from_cache_file\u001b[39m=\u001b[39;49mload_from_cache_file,\n\u001b[1;32m    867\u001b[0m             cache_file_name\u001b[39m=\u001b[39;49mcache_file_names[k],\n\u001b[1;32m    868\u001b[0m             writer_batch_size\u001b[39m=\u001b[39;49mwriter_batch_size,\n\u001b[1;32m    869\u001b[0m             features\u001b[39m=\u001b[39;49mfeatures,\n\u001b[1;32m    870\u001b[0m             disable_nullable\u001b[39m=\u001b[39;49mdisable_nullable,\n\u001b[1;32m    871\u001b[0m             fn_kwargs\u001b[39m=\u001b[39;49mfn_kwargs,\n\u001b[1;32m    872\u001b[0m             num_proc\u001b[39m=\u001b[39;49mnum_proc,\n\u001b[1;32m    873\u001b[0m             desc\u001b[39m=\u001b[39;49mdesc,\n\u001b[1;32m    874\u001b[0m         )\n\u001b[1;32m    875\u001b[0m         \u001b[39mfor\u001b[39;00m k, dataset \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mitems()\n\u001b[1;32m    876\u001b[0m     }\n\u001b[1;32m    877\u001b[0m )\n",
      "File \u001b[0;32m~/miniconda3/envs/kibsi/lib/python3.11/site-packages/datasets/arrow_dataset.py:591\u001b[0m, in \u001b[0;36mtransmit_tasks.<locals>.wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    589\u001b[0m     \u001b[39mself\u001b[39m: \u001b[39m\"\u001b[39m\u001b[39mDataset\u001b[39m\u001b[39m\"\u001b[39m \u001b[39m=\u001b[39m kwargs\u001b[39m.\u001b[39mpop(\u001b[39m\"\u001b[39m\u001b[39mself\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[1;32m    590\u001b[0m \u001b[39m# apply actual function\u001b[39;00m\n\u001b[0;32m--> 591\u001b[0m out: Union[\u001b[39m\"\u001b[39m\u001b[39mDataset\u001b[39m\u001b[39m\"\u001b[39m, \u001b[39m\"\u001b[39m\u001b[39mDatasetDict\u001b[39m\u001b[39m\"\u001b[39m] \u001b[39m=\u001b[39m func(\u001b[39mself\u001b[39;49m, \u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m    592\u001b[0m datasets: List[\u001b[39m\"\u001b[39m\u001b[39mDataset\u001b[39m\u001b[39m\"\u001b[39m] \u001b[39m=\u001b[39m \u001b[39mlist\u001b[39m(out\u001b[39m.\u001b[39mvalues()) \u001b[39mif\u001b[39;00m \u001b[39misinstance\u001b[39m(out, \u001b[39mdict\u001b[39m) \u001b[39melse\u001b[39;00m [out]\n\u001b[1;32m    593\u001b[0m \u001b[39mfor\u001b[39;00m dataset \u001b[39min\u001b[39;00m datasets:\n\u001b[1;32m    594\u001b[0m     \u001b[39m# Remove task templates if a column mapping of the template is no longer valid\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/kibsi/lib/python3.11/site-packages/datasets/arrow_dataset.py:556\u001b[0m, in \u001b[0;36mtransmit_format.<locals>.wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    549\u001b[0m self_format \u001b[39m=\u001b[39m {\n\u001b[1;32m    550\u001b[0m     \u001b[39m\"\u001b[39m\u001b[39mtype\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_format_type,\n\u001b[1;32m    551\u001b[0m     \u001b[39m\"\u001b[39m\u001b[39mformat_kwargs\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_format_kwargs,\n\u001b[1;32m    552\u001b[0m     \u001b[39m\"\u001b[39m\u001b[39mcolumns\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_format_columns,\n\u001b[1;32m    553\u001b[0m     \u001b[39m\"\u001b[39m\u001b[39moutput_all_columns\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_output_all_columns,\n\u001b[1;32m    554\u001b[0m }\n\u001b[1;32m    555\u001b[0m \u001b[39m# apply actual function\u001b[39;00m\n\u001b[0;32m--> 556\u001b[0m out: Union[\u001b[39m\"\u001b[39m\u001b[39mDataset\u001b[39m\u001b[39m\"\u001b[39m, \u001b[39m\"\u001b[39m\u001b[39mDatasetDict\u001b[39m\u001b[39m\"\u001b[39m] \u001b[39m=\u001b[39m func(\u001b[39mself\u001b[39;49m, \u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m    557\u001b[0m datasets: List[\u001b[39m\"\u001b[39m\u001b[39mDataset\u001b[39m\u001b[39m\"\u001b[39m] \u001b[39m=\u001b[39m \u001b[39mlist\u001b[39m(out\u001b[39m.\u001b[39mvalues()) \u001b[39mif\u001b[39;00m \u001b[39misinstance\u001b[39m(out, \u001b[39mdict\u001b[39m) \u001b[39melse\u001b[39;00m [out]\n\u001b[1;32m    558\u001b[0m \u001b[39m# re-apply format to the output\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/kibsi/lib/python3.11/site-packages/datasets/arrow_dataset.py:3089\u001b[0m, in \u001b[0;36mDataset.map\u001b[0;34m(self, function, with_indices, with_rank, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, suffix_template, new_fingerprint, desc)\u001b[0m\n\u001b[1;32m   3082\u001b[0m \u001b[39mif\u001b[39;00m transformed_dataset \u001b[39mis\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m   3083\u001b[0m     \u001b[39mwith\u001b[39;00m logging\u001b[39m.\u001b[39mtqdm(\n\u001b[1;32m   3084\u001b[0m         disable\u001b[39m=\u001b[39m\u001b[39mnot\u001b[39;00m logging\u001b[39m.\u001b[39mis_progress_bar_enabled(),\n\u001b[1;32m   3085\u001b[0m         unit\u001b[39m=\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m examples\u001b[39m\u001b[39m\"\u001b[39m,\n\u001b[1;32m   3086\u001b[0m         total\u001b[39m=\u001b[39mpbar_total,\n\u001b[1;32m   3087\u001b[0m         desc\u001b[39m=\u001b[39mdesc \u001b[39mor\u001b[39;00m \u001b[39m\"\u001b[39m\u001b[39mMap\u001b[39m\u001b[39m\"\u001b[39m,\n\u001b[1;32m   3088\u001b[0m     ) \u001b[39mas\u001b[39;00m pbar:\n\u001b[0;32m-> 3089\u001b[0m         \u001b[39mfor\u001b[39;49;00m rank, done, content \u001b[39min\u001b[39;49;00m Dataset\u001b[39m.\u001b[39;49m_map_single(\u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mdataset_kwargs):\n\u001b[1;32m   3090\u001b[0m             \u001b[39mif\u001b[39;49;00m done:\n\u001b[1;32m   3091\u001b[0m                 shards_done \u001b[39m+\u001b[39;49m\u001b[39m=\u001b[39;49m \u001b[39m1\u001b[39;49m\n",
      "File \u001b[0;32m~/miniconda3/envs/kibsi/lib/python3.11/site-packages/datasets/arrow_dataset.py:3466\u001b[0m, in \u001b[0;36mDataset._map_single\u001b[0;34m(shard, function, with_indices, with_rank, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, new_fingerprint, rank, offset)\u001b[0m\n\u001b[1;32m   3462\u001b[0m indices \u001b[39m=\u001b[39m \u001b[39mlist\u001b[39m(\n\u001b[1;32m   3463\u001b[0m     \u001b[39mrange\u001b[39m(\u001b[39m*\u001b[39m(\u001b[39mslice\u001b[39m(i, i \u001b[39m+\u001b[39m batch_size)\u001b[39m.\u001b[39mindices(shard\u001b[39m.\u001b[39mnum_rows)))\n\u001b[1;32m   3464\u001b[0m )  \u001b[39m# Something simpler?\u001b[39;00m\n\u001b[1;32m   3465\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m-> 3466\u001b[0m     batch \u001b[39m=\u001b[39m apply_function_on_filtered_inputs(\n\u001b[1;32m   3467\u001b[0m         batch,\n\u001b[1;32m   3468\u001b[0m         indices,\n\u001b[1;32m   3469\u001b[0m         check_same_num_examples\u001b[39m=\u001b[39;49m\u001b[39mlen\u001b[39;49m(shard\u001b[39m.\u001b[39;49mlist_indexes()) \u001b[39m>\u001b[39;49m \u001b[39m0\u001b[39;49m,\n\u001b[1;32m   3470\u001b[0m         offset\u001b[39m=\u001b[39;49moffset,\n\u001b[1;32m   3471\u001b[0m     )\n\u001b[1;32m   3472\u001b[0m \u001b[39mexcept\u001b[39;00m NumExamplesMismatchError:\n\u001b[1;32m   3473\u001b[0m     \u001b[39mraise\u001b[39;00m DatasetTransformationNotAllowedError(\n\u001b[1;32m   3474\u001b[0m         \u001b[39m\"\u001b[39m\u001b[39mUsing `.map` in batched mode on a dataset with attached indexes is allowed only if it doesn\u001b[39m\u001b[39m'\u001b[39m\u001b[39mt create or remove existing examples. You can first run `.drop_index() to remove your index and then re-add it.\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m   3475\u001b[0m     ) \u001b[39mfrom\u001b[39;00m \u001b[39mNone\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/kibsi/lib/python3.11/site-packages/datasets/arrow_dataset.py:3345\u001b[0m, in \u001b[0;36mDataset._map_single.<locals>.apply_function_on_filtered_inputs\u001b[0;34m(pa_inputs, indices, check_same_num_examples, offset)\u001b[0m\n\u001b[1;32m   3343\u001b[0m \u001b[39mif\u001b[39;00m with_rank:\n\u001b[1;32m   3344\u001b[0m     additional_args \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m (rank,)\n\u001b[0;32m-> 3345\u001b[0m processed_inputs \u001b[39m=\u001b[39m function(\u001b[39m*\u001b[39;49mfn_args, \u001b[39m*\u001b[39;49madditional_args, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mfn_kwargs)\n\u001b[1;32m   3346\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39misinstance\u001b[39m(processed_inputs, LazyDict):\n\u001b[1;32m   3347\u001b[0m     processed_inputs \u001b[39m=\u001b[39m {\n\u001b[1;32m   3348\u001b[0m         k: v \u001b[39mfor\u001b[39;00m k, v \u001b[39min\u001b[39;00m processed_inputs\u001b[39m.\u001b[39mdata\u001b[39m.\u001b[39mitems() \u001b[39mif\u001b[39;00m k \u001b[39mnot\u001b[39;00m \u001b[39min\u001b[39;00m processed_inputs\u001b[39m.\u001b[39mkeys_to_format\n\u001b[1;32m   3349\u001b[0m     }\n",
      "\u001b[1;32m/home/xzrderek/cs182-project/prefix_tuning.ipynb Cell 7\u001b[0m line \u001b[0;36m5\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W1sdnNjb2RlLXJlbW90ZQ%3D%3D?line=2'>3</a>\u001b[0m inputs \u001b[39m=\u001b[39m examples[\u001b[39m\"\u001b[39m\u001b[39mSpongeBob Lines\u001b[39m\u001b[39m\"\u001b[39m]\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W1sdnNjb2RlLXJlbW90ZQ%3D%3D?line=3'>4</a>\u001b[0m \u001b[39m# targets = examples[label_column]\u001b[39;00m\n\u001b[0;32m----> <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W1sdnNjb2RlLXJlbW90ZQ%3D%3D?line=4'>5</a>\u001b[0m model_inputs \u001b[39m=\u001b[39m tokenizer(inputs, max_length\u001b[39m=\u001b[39;49mmax_length, padding\u001b[39m=\u001b[39;49m\u001b[39m\"\u001b[39;49m\u001b[39mmax_length\u001b[39;49m\u001b[39m\"\u001b[39;49m, truncation\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m, return_tensors\u001b[39m=\u001b[39;49m\u001b[39m\"\u001b[39;49m\u001b[39mpt\u001b[39;49m\u001b[39m\"\u001b[39;49m)\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W1sdnNjb2RlLXJlbW90ZQ%3D%3D?line=5'>6</a>\u001b[0m \u001b[39m# labels = tokenizer(targets, max_length=2, padding=\"max_length\", truncation=True, return_tensors=\"pt\")\u001b[39;00m\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W1sdnNjb2RlLXJlbW90ZQ%3D%3D?line=6'>7</a>\u001b[0m \u001b[39m# labels = labels[\"input_ids\"]\u001b[39;00m\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W1sdnNjb2RlLXJlbW90ZQ%3D%3D?line=7'>8</a>\u001b[0m \u001b[39m# labels[labels == tokenizer.pad_token_id] = -100\u001b[39;00m\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W1sdnNjb2RlLXJlbW90ZQ%3D%3D?line=8'>9</a>\u001b[0m \u001b[39m# model_inputs[\"labels\"] = labels\u001b[39;00m\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W1sdnNjb2RlLXJlbW90ZQ%3D%3D?line=9'>10</a>\u001b[0m \u001b[39mreturn\u001b[39;00m model_inputs\n",
      "File \u001b[0;32m~/.local/lib/python3.11/site-packages/transformers/tokenization_utils_base.py:2798\u001b[0m, in \u001b[0;36mPreTrainedTokenizerBase.__call__\u001b[0;34m(self, text, text_pair, text_target, text_pair_target, add_special_tokens, padding, truncation, max_length, stride, is_split_into_words, pad_to_multiple_of, return_tensors, return_token_type_ids, return_attention_mask, return_overflowing_tokens, return_special_tokens_mask, return_offsets_mapping, return_length, verbose, **kwargs)\u001b[0m\n\u001b[1;32m   2796\u001b[0m     \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_in_target_context_manager:\n\u001b[1;32m   2797\u001b[0m         \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_switch_to_input_mode()\n\u001b[0;32m-> 2798\u001b[0m     encodings \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_call_one(text\u001b[39m=\u001b[39;49mtext, text_pair\u001b[39m=\u001b[39;49mtext_pair, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mall_kwargs)\n\u001b[1;32m   2799\u001b[0m \u001b[39mif\u001b[39;00m text_target \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m   2800\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_switch_to_target_mode()\n",
      "File \u001b[0;32m~/.local/lib/python3.11/site-packages/transformers/tokenization_utils_base.py:2884\u001b[0m, in \u001b[0;36mPreTrainedTokenizerBase._call_one\u001b[0;34m(self, text, text_pair, add_special_tokens, padding, truncation, max_length, stride, is_split_into_words, pad_to_multiple_of, return_tensors, return_token_type_ids, return_attention_mask, return_overflowing_tokens, return_special_tokens_mask, return_offsets_mapping, return_length, verbose, **kwargs)\u001b[0m\n\u001b[1;32m   2879\u001b[0m         \u001b[39mraise\u001b[39;00m \u001b[39mValueError\u001b[39;00m(\n\u001b[1;32m   2880\u001b[0m             \u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mbatch length of `text`: \u001b[39m\u001b[39m{\u001b[39;00m\u001b[39mlen\u001b[39m(text)\u001b[39m}\u001b[39;00m\u001b[39m does not match batch length of `text_pair`:\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m   2881\u001b[0m             \u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m \u001b[39m\u001b[39m{\u001b[39;00m\u001b[39mlen\u001b[39m(text_pair)\u001b[39m}\u001b[39;00m\u001b[39m.\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m   2882\u001b[0m         )\n\u001b[1;32m   2883\u001b[0m     batch_text_or_text_pairs \u001b[39m=\u001b[39m \u001b[39mlist\u001b[39m(\u001b[39mzip\u001b[39m(text, text_pair)) \u001b[39mif\u001b[39;00m text_pair \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m \u001b[39melse\u001b[39;00m text\n\u001b[0;32m-> 2884\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mbatch_encode_plus(\n\u001b[1;32m   2885\u001b[0m         batch_text_or_text_pairs\u001b[39m=\u001b[39;49mbatch_text_or_text_pairs,\n\u001b[1;32m   2886\u001b[0m         add_special_tokens\u001b[39m=\u001b[39;49madd_special_tokens,\n\u001b[1;32m   2887\u001b[0m         padding\u001b[39m=\u001b[39;49mpadding,\n\u001b[1;32m   2888\u001b[0m         truncation\u001b[39m=\u001b[39;49mtruncation,\n\u001b[1;32m   2889\u001b[0m         max_length\u001b[39m=\u001b[39;49mmax_length,\n\u001b[1;32m   2890\u001b[0m         stride\u001b[39m=\u001b[39;49mstride,\n\u001b[1;32m   2891\u001b[0m         is_split_into_words\u001b[39m=\u001b[39;49mis_split_into_words,\n\u001b[1;32m   2892\u001b[0m         pad_to_multiple_of\u001b[39m=\u001b[39;49mpad_to_multiple_of,\n\u001b[1;32m   2893\u001b[0m         return_tensors\u001b[39m=\u001b[39;49mreturn_tensors,\n\u001b[1;32m   2894\u001b[0m         return_token_type_ids\u001b[39m=\u001b[39;49mreturn_token_type_ids,\n\u001b[1;32m   2895\u001b[0m         return_attention_mask\u001b[39m=\u001b[39;49mreturn_attention_mask,\n\u001b[1;32m   2896\u001b[0m         return_overflowing_tokens\u001b[39m=\u001b[39;49mreturn_overflowing_tokens,\n\u001b[1;32m   2897\u001b[0m         return_special_tokens_mask\u001b[39m=\u001b[39;49mreturn_special_tokens_mask,\n\u001b[1;32m   2898\u001b[0m         return_offsets_mapping\u001b[39m=\u001b[39;49mreturn_offsets_mapping,\n\u001b[1;32m   2899\u001b[0m         return_length\u001b[39m=\u001b[39;49mreturn_length,\n\u001b[1;32m   2900\u001b[0m         verbose\u001b[39m=\u001b[39;49mverbose,\n\u001b[1;32m   2901\u001b[0m         \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs,\n\u001b[1;32m   2902\u001b[0m     )\n\u001b[1;32m   2903\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m   2904\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mencode_plus(\n\u001b[1;32m   2905\u001b[0m         text\u001b[39m=\u001b[39mtext,\n\u001b[1;32m   2906\u001b[0m         text_pair\u001b[39m=\u001b[39mtext_pair,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   2922\u001b[0m         \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs,\n\u001b[1;32m   2923\u001b[0m     )\n",
      "File \u001b[0;32m~/.local/lib/python3.11/site-packages/transformers/tokenization_utils_base.py:3066\u001b[0m, in \u001b[0;36mPreTrainedTokenizerBase.batch_encode_plus\u001b[0;34m(self, batch_text_or_text_pairs, add_special_tokens, padding, truncation, max_length, stride, is_split_into_words, pad_to_multiple_of, return_tensors, return_token_type_ids, return_attention_mask, return_overflowing_tokens, return_special_tokens_mask, return_offsets_mapping, return_length, verbose, **kwargs)\u001b[0m\n\u001b[1;32m   3049\u001b[0m \u001b[39m\u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m   3050\u001b[0m \u001b[39mTokenize and prepare for the model a list of sequences or a list of pairs of sequences.\u001b[39;00m\n\u001b[1;32m   3051\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   3062\u001b[0m \u001b[39m        details in `encode_plus`).\u001b[39;00m\n\u001b[1;32m   3063\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m   3065\u001b[0m \u001b[39m# Backward compatibility for 'truncation_strategy', 'pad_to_max_length'\u001b[39;00m\n\u001b[0;32m-> 3066\u001b[0m padding_strategy, truncation_strategy, max_length, kwargs \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_get_padding_truncation_strategies(\n\u001b[1;32m   3067\u001b[0m     padding\u001b[39m=\u001b[39;49mpadding,\n\u001b[1;32m   3068\u001b[0m     truncation\u001b[39m=\u001b[39;49mtruncation,\n\u001b[1;32m   3069\u001b[0m     max_length\u001b[39m=\u001b[39;49mmax_length,\n\u001b[1;32m   3070\u001b[0m     pad_to_multiple_of\u001b[39m=\u001b[39;49mpad_to_multiple_of,\n\u001b[1;32m   3071\u001b[0m     verbose\u001b[39m=\u001b[39;49mverbose,\n\u001b[1;32m   3072\u001b[0m     \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs,\n\u001b[1;32m   3073\u001b[0m )\n\u001b[1;32m   3075\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_batch_encode_plus(\n\u001b[1;32m   3076\u001b[0m     batch_text_or_text_pairs\u001b[39m=\u001b[39mbatch_text_or_text_pairs,\n\u001b[1;32m   3077\u001b[0m     add_special_tokens\u001b[39m=\u001b[39madd_special_tokens,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   3092\u001b[0m     \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs,\n\u001b[1;32m   3093\u001b[0m )\n",
      "File \u001b[0;32m~/.local/lib/python3.11/site-packages/transformers/tokenization_utils_base.py:2703\u001b[0m, in \u001b[0;36mPreTrainedTokenizerBase._get_padding_truncation_strategies\u001b[0;34m(self, padding, truncation, max_length, pad_to_multiple_of, verbose, **kwargs)\u001b[0m\n\u001b[1;32m   2701\u001b[0m \u001b[39m# Test if we have a padding token\u001b[39;00m\n\u001b[1;32m   2702\u001b[0m \u001b[39mif\u001b[39;00m padding_strategy \u001b[39m!=\u001b[39m PaddingStrategy\u001b[39m.\u001b[39mDO_NOT_PAD \u001b[39mand\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mpad_token \u001b[39mis\u001b[39;00m \u001b[39mNone\u001b[39;00m \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mpad_token_id \u001b[39m<\u001b[39m \u001b[39m0\u001b[39m):\n\u001b[0;32m-> 2703\u001b[0m     \u001b[39mraise\u001b[39;00m \u001b[39mValueError\u001b[39;00m(\n\u001b[1;32m   2704\u001b[0m         \u001b[39m\"\u001b[39m\u001b[39mAsking to pad but the tokenizer does not have a padding token. \u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m   2705\u001b[0m         \u001b[39m\"\u001b[39m\u001b[39mPlease select a token to use as `pad_token` `(tokenizer.pad_token = tokenizer.eos_token e.g.)` \u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m   2706\u001b[0m         \u001b[39m\"\u001b[39m\u001b[39mor add a new pad token via `tokenizer.add_special_tokens(\u001b[39m\u001b[39m{\u001b[39m\u001b[39m'\u001b[39m\u001b[39mpad_token\u001b[39m\u001b[39m'\u001b[39m\u001b[39m: \u001b[39m\u001b[39m'\u001b[39m\u001b[39m[PAD]\u001b[39m\u001b[39m'\u001b[39m\u001b[39m})`.\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m   2707\u001b[0m     )\n\u001b[1;32m   2709\u001b[0m \u001b[39m# Check that we will truncate to a multiple of pad_to_multiple_of if both are provided\u001b[39;00m\n\u001b[1;32m   2710\u001b[0m \u001b[39mif\u001b[39;00m (\n\u001b[1;32m   2711\u001b[0m     truncation_strategy \u001b[39m!=\u001b[39m TruncationStrategy\u001b[39m.\u001b[39mDO_NOT_TRUNCATE\n\u001b[1;32m   2712\u001b[0m     \u001b[39mand\u001b[39;00m padding_strategy \u001b[39m!=\u001b[39m PaddingStrategy\u001b[39m.\u001b[39mDO_NOT_PAD\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   2715\u001b[0m     \u001b[39mand\u001b[39;00m (max_length \u001b[39m%\u001b[39m pad_to_multiple_of \u001b[39m!=\u001b[39m \u001b[39m0\u001b[39m)\n\u001b[1;32m   2716\u001b[0m ):\n",
      "\u001b[0;31mValueError\u001b[0m: Asking to pad but the tokenizer does not have a padding token. Please select a token to use as `pad_token` `(tokenizer.pad_token = tokenizer.eos_token e.g.)` or add a new pad token via `tokenizer.add_special_tokens({'pad_token': '[PAD]'})`."
     ]
    }
   ],
   "source": [
    "processed_datasets = dataset.map(\n",
    "    preprocess_function,\n",
    "    batched=True,\n",
    "    num_proc=1,\n",
    "    remove_columns=dataset[\"train\"].column_names,\n",
    "    load_from_cache_file=False,\n",
    "    desc=\"Running tokenizer on dataset\",\n",
    ")\n",
    "\n",
    "train_dataset = processed_datasets[\"train\"]\n",
    "eval_dataset = processed_datasets[\"test\"]\n",
    "\n",
    "train_dataloader = DataLoader(\n",
    "    dataset['train'], shuffle=True, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True\n",
    ")\n",
    "eval_dataloader = DataLoader(dataset['test'], collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Detected kernel version 5.4.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "You can't train a model that has been loaded in 8-bit precision on a different device than the one you're training on. Make sure you loaded the model on the correct device using for example `device_map={'':torch.cuda.current_device()}you're training on. Make sure you loaded the model on the correct device using for example `device_map={'':torch.cuda.current_device() or device_map={'':torch.xpu.current_device()}",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[1;32m/home/xzrderek/cs182-project/prefix_tuning.ipynb Cell 8\u001b[0m line \u001b[0;36m2\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#X15sdnNjb2RlLXJlbW90ZQ%3D%3D?line=4'>5</a>\u001b[0m trainer \u001b[39m=\u001b[39m transformers\u001b[39m.\u001b[39mTrainer(\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#X15sdnNjb2RlLXJlbW90ZQ%3D%3D?line=5'>6</a>\u001b[0m     model\u001b[39m=\u001b[39mmodel,\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#X15sdnNjb2RlLXJlbW90ZQ%3D%3D?line=6'>7</a>\u001b[0m     train_dataset\u001b[39m=\u001b[39mdataset[\u001b[39m\"\u001b[39m\u001b[39mtrain\u001b[39m\u001b[39m\"\u001b[39m],\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#X15sdnNjb2RlLXJlbW90ZQ%3D%3D?line=18'>19</a>\u001b[0m     data_collator\u001b[39m=\u001b[39mtransformers\u001b[39m.\u001b[39mDataCollatorForLanguageModeling(tokenizer, mlm\u001b[39m=\u001b[39m\u001b[39mFalse\u001b[39;00m),\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#X15sdnNjb2RlLXJlbW90ZQ%3D%3D?line=19'>20</a>\u001b[0m )\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#X15sdnNjb2RlLXJlbW90ZQ%3D%3D?line=20'>21</a>\u001b[0m model\u001b[39m.\u001b[39mconfig\u001b[39m.\u001b[39muse_cache \u001b[39m=\u001b[39m \u001b[39mFalse\u001b[39;00m  \u001b[39m# silence the warnings. Please re-enable for inference!\u001b[39;00m\n\u001b[0;32m---> <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#X15sdnNjb2RlLXJlbW90ZQ%3D%3D?line=21'>22</a>\u001b[0m trainer\u001b[39m.\u001b[39;49mtrain()\n",
      "File \u001b[0;32m~/.local/lib/python3.11/site-packages/transformers/trainer.py:1555\u001b[0m, in \u001b[0;36mTrainer.train\u001b[0;34m(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)\u001b[0m\n\u001b[1;32m   1553\u001b[0m         hf_hub_utils\u001b[39m.\u001b[39menable_progress_bars()\n\u001b[1;32m   1554\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m-> 1555\u001b[0m     \u001b[39mreturn\u001b[39;00m inner_training_loop(\n\u001b[1;32m   1556\u001b[0m         args\u001b[39m=\u001b[39;49margs,\n\u001b[1;32m   1557\u001b[0m         resume_from_checkpoint\u001b[39m=\u001b[39;49mresume_from_checkpoint,\n\u001b[1;32m   1558\u001b[0m         trial\u001b[39m=\u001b[39;49mtrial,\n\u001b[1;32m   1559\u001b[0m         ignore_keys_for_eval\u001b[39m=\u001b[39;49mignore_keys_for_eval,\n\u001b[1;32m   1560\u001b[0m     )\n",
      "File \u001b[0;32m~/.local/lib/python3.11/site-packages/transformers/trainer.py:1689\u001b[0m, in \u001b[0;36mTrainer._inner_training_loop\u001b[0;34m(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)\u001b[0m\n\u001b[1;32m   1687\u001b[0m         model \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39maccelerator\u001b[39m.\u001b[39mprepare(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmodel)\n\u001b[1;32m   1688\u001b[0m     \u001b[39melse\u001b[39;00m:\n\u001b[0;32m-> 1689\u001b[0m         model, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39moptimizer \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49maccelerator\u001b[39m.\u001b[39;49mprepare(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mmodel, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49moptimizer)\n\u001b[1;32m   1690\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m   1691\u001b[0m     \u001b[39m# to handle cases wherein we pass \"DummyScheduler\" such as when it is specified in DeepSpeed config.\u001b[39;00m\n\u001b[1;32m   1692\u001b[0m     model, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39moptimizer, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mlr_scheduler \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39maccelerator\u001b[39m.\u001b[39mprepare(\n\u001b[1;32m   1693\u001b[0m         \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmodel, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39moptimizer, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mlr_scheduler\n\u001b[1;32m   1694\u001b[0m     )\n",
      "File \u001b[0;32m~/.local/lib/python3.11/site-packages/accelerate/accelerator.py:1288\u001b[0m, in \u001b[0;36mAccelerator.prepare\u001b[0;34m(self, device_placement, *args)\u001b[0m\n\u001b[1;32m   1286\u001b[0m     result \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_prepare_megatron_lm(\u001b[39m*\u001b[39margs)\n\u001b[1;32m   1287\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m-> 1288\u001b[0m     result \u001b[39m=\u001b[39m \u001b[39mtuple\u001b[39m(\n\u001b[1;32m   1289\u001b[0m         \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_prepare_one(obj, first_pass\u001b[39m=\u001b[39m\u001b[39mTrue\u001b[39;00m, device_placement\u001b[39m=\u001b[39md) \u001b[39mfor\u001b[39;00m obj, d \u001b[39min\u001b[39;00m \u001b[39mzip\u001b[39m(args, device_placement)\n\u001b[1;32m   1290\u001b[0m     )\n\u001b[1;32m   1291\u001b[0m     result \u001b[39m=\u001b[39m \u001b[39mtuple\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_prepare_one(obj, device_placement\u001b[39m=\u001b[39md) \u001b[39mfor\u001b[39;00m obj, d \u001b[39min\u001b[39;00m \u001b[39mzip\u001b[39m(result, device_placement))\n\u001b[1;32m   1293\u001b[0m \u001b[39mif\u001b[39;00m tpu_should_fix_optimizer \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmixed_precision \u001b[39m==\u001b[39m \u001b[39m\"\u001b[39m\u001b[39mfp8\u001b[39m\u001b[39m\"\u001b[39m:\n\u001b[1;32m   1294\u001b[0m     \u001b[39m# 2. grabbing new model parameters\u001b[39;00m\n",
      "File \u001b[0;32m~/.local/lib/python3.11/site-packages/accelerate/accelerator.py:1289\u001b[0m, in \u001b[0;36m<genexpr>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m   1286\u001b[0m     result \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_prepare_megatron_lm(\u001b[39m*\u001b[39margs)\n\u001b[1;32m   1287\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m   1288\u001b[0m     result \u001b[39m=\u001b[39m \u001b[39mtuple\u001b[39m(\n\u001b[0;32m-> 1289\u001b[0m         \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_prepare_one(obj, first_pass\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m, device_placement\u001b[39m=\u001b[39;49md) \u001b[39mfor\u001b[39;00m obj, d \u001b[39min\u001b[39;00m \u001b[39mzip\u001b[39m(args, device_placement)\n\u001b[1;32m   1290\u001b[0m     )\n\u001b[1;32m   1291\u001b[0m     result \u001b[39m=\u001b[39m \u001b[39mtuple\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_prepare_one(obj, device_placement\u001b[39m=\u001b[39md) \u001b[39mfor\u001b[39;00m obj, d \u001b[39min\u001b[39;00m \u001b[39mzip\u001b[39m(result, device_placement))\n\u001b[1;32m   1293\u001b[0m \u001b[39mif\u001b[39;00m tpu_should_fix_optimizer \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmixed_precision \u001b[39m==\u001b[39m \u001b[39m\"\u001b[39m\u001b[39mfp8\u001b[39m\u001b[39m\"\u001b[39m:\n\u001b[1;32m   1294\u001b[0m     \u001b[39m# 2. grabbing new model parameters\u001b[39;00m\n",
      "File \u001b[0;32m~/.local/lib/python3.11/site-packages/accelerate/accelerator.py:1094\u001b[0m, in \u001b[0;36mAccelerator._prepare_one\u001b[0;34m(self, obj, first_pass, device_placement)\u001b[0m\n\u001b[1;32m   1092\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mprepare_data_loader(obj, device_placement\u001b[39m=\u001b[39mdevice_placement)\n\u001b[1;32m   1093\u001b[0m \u001b[39melif\u001b[39;00m \u001b[39misinstance\u001b[39m(obj, torch\u001b[39m.\u001b[39mnn\u001b[39m.\u001b[39mModule):\n\u001b[0;32m-> 1094\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mprepare_model(obj, device_placement\u001b[39m=\u001b[39;49mdevice_placement)\n\u001b[1;32m   1095\u001b[0m \u001b[39melif\u001b[39;00m \u001b[39misinstance\u001b[39m(obj, torch\u001b[39m.\u001b[39moptim\u001b[39m.\u001b[39mOptimizer):\n\u001b[1;32m   1096\u001b[0m     optimizer \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mprepare_optimizer(obj, device_placement\u001b[39m=\u001b[39mdevice_placement)\n",
      "File \u001b[0;32m~/.local/lib/python3.11/site-packages/accelerate/accelerator.py:1376\u001b[0m, in \u001b[0;36mAccelerator.prepare_model\u001b[0;34m(self, model, device_placement, evaluation_mode)\u001b[0m\n\u001b[1;32m   1373\u001b[0m \u001b[39mif\u001b[39;00m torch\u001b[39m.\u001b[39mdevice(current_device_index) \u001b[39m!=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdevice:\n\u001b[1;32m   1374\u001b[0m     \u001b[39m# if on the first device (GPU 0) we don't care\u001b[39;00m\n\u001b[1;32m   1375\u001b[0m     \u001b[39mif\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdevice\u001b[39m.\u001b[39mindex \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m) \u001b[39mor\u001b[39;00m (current_device_index \u001b[39m!=\u001b[39m \u001b[39m0\u001b[39m):\n\u001b[0;32m-> 1376\u001b[0m         \u001b[39mraise\u001b[39;00m \u001b[39mValueError\u001b[39;00m(\n\u001b[1;32m   1377\u001b[0m             \u001b[39m\"\u001b[39m\u001b[39mYou can\u001b[39m\u001b[39m'\u001b[39m\u001b[39mt train a model that has been loaded in 8-bit precision on a different device than the one \u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m   1378\u001b[0m             \u001b[39m\"\u001b[39m\u001b[39myou\u001b[39m\u001b[39m'\u001b[39m\u001b[39mre training on. Make sure you loaded the model on the correct device using for example `device_map=\u001b[39m\u001b[39m{\u001b[39m\u001b[39m'\u001b[39m\u001b[39m'\u001b[39m\u001b[39m:torch.cuda.current_device()}\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m   1379\u001b[0m             \u001b[39m\"\u001b[39m\u001b[39myou\u001b[39m\u001b[39m'\u001b[39m\u001b[39mre training on. Make sure you loaded the model on the correct device using for example `device_map=\u001b[39m\u001b[39m{\u001b[39m\u001b[39m'\u001b[39m\u001b[39m'\u001b[39m\u001b[39m:torch.cuda.current_device() or device_map=\u001b[39m\u001b[39m{\u001b[39m\u001b[39m'\u001b[39m\u001b[39m'\u001b[39m\u001b[39m:torch.xpu.current_device()}\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m   1380\u001b[0m         )\n\u001b[1;32m   1382\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39m\"\u001b[39m\u001b[39mcpu\u001b[39m\u001b[39m\"\u001b[39m \u001b[39min\u001b[39;00m model_devices \u001b[39mor\u001b[39;00m \u001b[39m\"\u001b[39m\u001b[39mdisk\u001b[39m\u001b[39m\"\u001b[39m \u001b[39min\u001b[39;00m model_devices:\n\u001b[1;32m   1383\u001b[0m     \u001b[39mraise\u001b[39;00m \u001b[39mValueError\u001b[39;00m(\n\u001b[1;32m   1384\u001b[0m         \u001b[39m\"\u001b[39m\u001b[39mYou can\u001b[39m\u001b[39m'\u001b[39m\u001b[39mt train a model that has been loaded in 8-bit precision with CPU or disk offload.\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m   1385\u001b[0m     )\n",
      "\u001b[0;31mValueError\u001b[0m: You can't train a model that has been loaded in 8-bit precision on a different device than the one you're training on. Make sure you loaded the model on the correct device using for example `device_map={'':torch.cuda.current_device()}you're training on. Make sure you loaded the model on the correct device using for example `device_map={'':torch.cuda.current_device() or device_map={'':torch.xpu.current_device()}"
     ]
    }
   ],
   "source": [
    "import transformers\n",
    "\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "trainer = transformers.Trainer(\n",
    "    model=model,\n",
    "    train_dataset=dataset[\"train\"],\n",
    "    args=transformers.TrainingArguments(\n",
    "        per_device_train_batch_size=1,\n",
    "        gradient_accumulation_steps=4,\n",
    "        warmup_steps=2,\n",
    "        max_steps=10,\n",
    "        learning_rate=2e-4,\n",
    "        fp16=True,\n",
    "        logging_steps=1,\n",
    "        output_dir=\"outputs\",\n",
    "        optim=\"paged_adamw_8bit\"\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()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/1208 [00:00<?, ?it/s]\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "expected sequence of length 30 at dim 1 (got 28)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[1;32m/home/xzrderek/cs182-project/prefix_tuning.ipynb Cell 9\u001b[0m line \u001b[0;36m1\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W6sdnNjb2RlLXJlbW90ZQ%3D%3D?line=10'>11</a>\u001b[0m model\u001b[39m.\u001b[39mtrain()\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W6sdnNjb2RlLXJlbW90ZQ%3D%3D?line=11'>12</a>\u001b[0m total_loss \u001b[39m=\u001b[39m \u001b[39m0\u001b[39m\n\u001b[0;32m---> <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W6sdnNjb2RlLXJlbW90ZQ%3D%3D?line=12'>13</a>\u001b[0m \u001b[39mfor\u001b[39;49;00m step, batch \u001b[39min\u001b[39;49;00m \u001b[39menumerate\u001b[39;49m(tqdm(train_dataloader)):\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W6sdnNjb2RlLXJlbW90ZQ%3D%3D?line=13'>14</a>\u001b[0m     \u001b[39m# batch = {k: v.to(device) for k, v in batch.items()}\u001b[39;49;00m\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W6sdnNjb2RlLXJlbW90ZQ%3D%3D?line=14'>15</a>\u001b[0m     outputs \u001b[39m=\u001b[39;49m model(\u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mbatch)\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2Bhoneydew/home/xzrderek/cs182-project/prefix_tuning.ipynb#W6sdnNjb2RlLXJlbW90ZQ%3D%3D?line=15'>16</a>\u001b[0m     loss \u001b[39m=\u001b[39;49m outputs\u001b[39m.\u001b[39;49mloss\n",
      "File \u001b[0;32m~/.local/lib/python3.11/site-packages/tqdm/std.py:1182\u001b[0m, in \u001b[0;36mtqdm.__iter__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1179\u001b[0m time \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_time\n\u001b[1;32m   1181\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m-> 1182\u001b[0m     \u001b[39mfor\u001b[39;49;00m obj \u001b[39min\u001b[39;49;00m iterable:\n\u001b[1;32m   1183\u001b[0m         \u001b[39myield\u001b[39;49;00m obj\n\u001b[1;32m   1184\u001b[0m         \u001b[39m# Update and possibly print the progressbar.\u001b[39;49;00m\n\u001b[1;32m   1185\u001b[0m         \u001b[39m# Note: does not call self.update(1) for speed optimisation.\u001b[39;49;00m\n",
      "File \u001b[0;32m~/.local/lib/python3.11/site-packages/torch/utils/data/dataloader.py:630\u001b[0m, in \u001b[0;36m_BaseDataLoaderIter.__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    627\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_sampler_iter \u001b[39mis\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m    628\u001b[0m     \u001b[39m# TODO(https://github.com/pytorch/pytorch/issues/76750)\u001b[39;00m\n\u001b[1;32m    629\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_reset()  \u001b[39m# type: ignore[call-arg]\u001b[39;00m\n\u001b[0;32m--> 630\u001b[0m data \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_next_data()\n\u001b[1;32m    631\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_num_yielded \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m \u001b[39m1\u001b[39m\n\u001b[1;32m    632\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_dataset_kind \u001b[39m==\u001b[39m _DatasetKind\u001b[39m.\u001b[39mIterable \u001b[39mand\u001b[39;00m \\\n\u001b[1;32m    633\u001b[0m         \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_IterableDataset_len_called \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m \u001b[39mand\u001b[39;00m \\\n\u001b[1;32m    634\u001b[0m         \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_num_yielded \u001b[39m>\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_IterableDataset_len_called:\n",
      "File \u001b[0;32m~/.local/lib/python3.11/site-packages/torch/utils/data/dataloader.py:674\u001b[0m, in \u001b[0;36m_SingleProcessDataLoaderIter._next_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    672\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m_next_data\u001b[39m(\u001b[39mself\u001b[39m):\n\u001b[1;32m    673\u001b[0m     index \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_next_index()  \u001b[39m# may raise StopIteration\u001b[39;00m\n\u001b[0;32m--> 674\u001b[0m     data \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_dataset_fetcher\u001b[39m.\u001b[39;49mfetch(index)  \u001b[39m# may raise StopIteration\u001b[39;00m\n\u001b[1;32m    675\u001b[0m     \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_pin_memory:\n\u001b[1;32m    676\u001b[0m         data \u001b[39m=\u001b[39m _utils\u001b[39m.\u001b[39mpin_memory\u001b[39m.\u001b[39mpin_memory(data, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_pin_memory_device)\n",
      "File \u001b[0;32m~/.local/lib/python3.11/site-packages/torch/utils/data/_utils/fetch.py:54\u001b[0m, in \u001b[0;36m_MapDatasetFetcher.fetch\u001b[0;34m(self, possibly_batched_index)\u001b[0m\n\u001b[1;32m     52\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m     53\u001b[0m     data \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdataset[possibly_batched_index]\n\u001b[0;32m---> 54\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mcollate_fn(data)\n",
      "File \u001b[0;32m~/.local/lib/python3.11/site-packages/transformers/data/data_collator.py:70\u001b[0m, in \u001b[0;36mdefault_data_collator\u001b[0;34m(features, return_tensors)\u001b[0m\n\u001b[1;32m     64\u001b[0m \u001b[39m# In this function we'll make the assumption that all `features` in the batch\u001b[39;00m\n\u001b[1;32m     65\u001b[0m \u001b[39m# have the same attributes.\u001b[39;00m\n\u001b[1;32m     66\u001b[0m \u001b[39m# So we will look at the first element as a proxy for what attributes exist\u001b[39;00m\n\u001b[1;32m     67\u001b[0m \u001b[39m# on the whole batch.\u001b[39;00m\n\u001b[1;32m     69\u001b[0m \u001b[39mif\u001b[39;00m return_tensors \u001b[39m==\u001b[39m \u001b[39m\"\u001b[39m\u001b[39mpt\u001b[39m\u001b[39m\"\u001b[39m:\n\u001b[0;32m---> 70\u001b[0m     \u001b[39mreturn\u001b[39;00m torch_default_data_collator(features)\n\u001b[1;32m     71\u001b[0m \u001b[39melif\u001b[39;00m return_tensors \u001b[39m==\u001b[39m \u001b[39m\"\u001b[39m\u001b[39mtf\u001b[39m\u001b[39m\"\u001b[39m:\n\u001b[1;32m     72\u001b[0m     \u001b[39mreturn\u001b[39;00m tf_default_data_collator(features)\n",
      "File \u001b[0;32m~/.local/lib/python3.11/site-packages/transformers/data/data_collator.py:136\u001b[0m, in \u001b[0;36mtorch_default_data_collator\u001b[0;34m(features)\u001b[0m\n\u001b[1;32m    134\u001b[0m             batch[k] \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mtensor(np\u001b[39m.\u001b[39mstack([f[k] \u001b[39mfor\u001b[39;00m f \u001b[39min\u001b[39;00m features]))\n\u001b[1;32m    135\u001b[0m         \u001b[39melse\u001b[39;00m:\n\u001b[0;32m--> 136\u001b[0m             batch[k] \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39;49mtensor([f[k] \u001b[39mfor\u001b[39;49;00m f \u001b[39min\u001b[39;49;00m features])\n\u001b[1;32m    138\u001b[0m \u001b[39mreturn\u001b[39;00m batch\n",
      "\u001b[0;31mValueError\u001b[0m: expected sequence of length 30 at dim 1 (got 28)"
     ]
    }
   ],
   "source": [
    "optimizer = torch.optim.AdamW(model.parameters(), lr=lr)\n",
    "lr_scheduler = get_linear_schedule_with_warmup(\n",
    "    optimizer=optimizer,\n",
    "    num_warmup_steps=0,\n",
    "    num_training_steps=(len(train_dataloader) * num_epochs),\n",
    ")\n",
    "\n",
    "model = model.to(device)\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    total_loss = 0\n",
    "    for step, batch in enumerate(tqdm(train_dataloader)):\n",
    "        batch = {k: v.to(device) for k, v in batch.items()}\n",
    "        outputs = model(**batch)\n",
    "        loss = outputs.loss\n",
    "        total_loss += loss.detach().float()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        lr_scheduler.step()\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "    model.eval()\n",
    "    eval_loss = 0\n",
    "    eval_preds = []\n",
    "    for step, batch in enumerate(tqdm(eval_dataloader)):\n",
    "        batch = {k: v.to(device) for k, v in batch.items()}\n",
    "        with torch.no_grad():\n",
    "            outputs = model(**batch)\n",
    "        loss = outputs.loss\n",
    "        eval_loss += loss.detach().float()\n",
    "        eval_preds.extend(\n",
    "            tokenizer.batch_decode(torch.argmax(outputs.logits, -1).detach().cpu().numpy(), skip_special_tokens=True)\n",
    "        )\n",
    "\n",
    "    eval_epoch_loss = eval_loss / len(eval_dataloader)\n",
    "    eval_ppl = torch.exp(eval_epoch_loss)\n",
    "    train_epoch_loss = total_loss / len(train_dataloader)\n",
    "    train_ppl = torch.exp(train_epoch_loss)\n",
    "    print(f\"{epoch=}: {train_ppl=} {train_epoch_loss=} {eval_ppl=} {eval_epoch_loss=}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "kibsi",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}