dialog_train_transformers.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "source": [
    "! pip install datasets transformers\n",
    "! pip install pynvml numba"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "YiKzbtakYRjR",
    "outputId": "ff213bc4-0caa-4f7b-dd00-929c05de189d"
   },
   "id": "YiKzbtakYRjR",
   "execution_count": 1,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
      "Requirement already satisfied: datasets in /usr/local/lib/python3.8/dist-packages (2.8.0)\n",
      "Requirement already satisfied: transformers in /usr/local/lib/python3.8/dist-packages (4.25.1)\n",
      "Requirement already satisfied: fsspec[http]>=2021.11.1 in /usr/local/lib/python3.8/dist-packages (from datasets) (2022.11.0)\n",
      "Requirement already satisfied: xxhash in /usr/local/lib/python3.8/dist-packages (from datasets) (3.2.0)\n",
      "Requirement already satisfied: dill<0.3.7 in /usr/local/lib/python3.8/dist-packages (from datasets) (0.3.6)\n",
      "Requirement already satisfied: pyarrow>=6.0.0 in /usr/local/lib/python3.8/dist-packages (from datasets) (9.0.0)\n",
      "Requirement already satisfied: huggingface-hub<1.0.0,>=0.2.0 in /usr/local/lib/python3.8/dist-packages (from datasets) (0.11.1)\n",
      "Requirement already satisfied: responses<0.19 in /usr/local/lib/python3.8/dist-packages (from datasets) (0.18.0)\n",
      "Requirement already satisfied: pandas in /usr/local/lib/python3.8/dist-packages (from datasets) (1.3.5)\n",
      "Requirement already satisfied: aiohttp in /usr/local/lib/python3.8/dist-packages (from datasets) (3.8.3)\n",
      "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.8/dist-packages (from datasets) (1.21.6)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.8/dist-packages (from datasets) (6.0)\n",
      "Requirement already satisfied: requests>=2.19.0 in /usr/local/lib/python3.8/dist-packages (from datasets) (2.25.1)\n",
      "Requirement already satisfied: tqdm>=4.62.1 in /usr/local/lib/python3.8/dist-packages (from datasets) (4.64.1)\n",
      "Requirement already satisfied: multiprocess in /usr/local/lib/python3.8/dist-packages (from datasets) (0.70.14)\n",
      "Requirement already satisfied: packaging in /usr/local/lib/python3.8/dist-packages (from datasets) (21.3)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.8/dist-packages (from transformers) (2022.6.2)\n",
      "Requirement already satisfied: filelock in /usr/local/lib/python3.8/dist-packages (from transformers) (3.9.0)\n",
      "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /usr/local/lib/python3.8/dist-packages (from transformers) (0.13.2)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets) (1.8.2)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets) (22.2.0)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets) (4.0.2)\n",
      "Requirement already satisfied: charset-normalizer<3.0,>=2.0 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets) (2.1.1)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets) (6.0.4)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets) (1.3.3)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets) (1.3.1)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.8/dist-packages (from huggingface-hub<1.0.0,>=0.2.0->datasets) (4.4.0)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.8/dist-packages (from packaging->datasets) (3.0.9)\n",
      "Requirement already satisfied: chardet<5,>=3.0.2 in /usr/local/lib/python3.8/dist-packages (from requests>=2.19.0->datasets) (4.0.0)\n",
      "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests>=2.19.0->datasets) (2.10)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.8/dist-packages (from requests>=2.19.0->datasets) (2022.12.7)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests>=2.19.0->datasets) (1.26.14)\n",
      "Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.8/dist-packages (from pandas->datasets) (2022.7)\n",
      "Requirement already satisfied: python-dateutil>=2.7.3 in /usr/local/lib/python3.8/dist-packages (from pandas->datasets) (2.8.2)\n",
      "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.8/dist-packages (from python-dateutil>=2.7.3->pandas->datasets) (1.15.0)\n",
      "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
      "Requirement already satisfied: pynvml in /usr/local/lib/python3.8/dist-packages (11.4.1)\n",
      "Requirement already satisfied: numba in /usr/local/lib/python3.8/dist-packages (0.56.4)\n",
      "Requirement already satisfied: llvmlite<0.40,>=0.39.0dev0 in /usr/local/lib/python3.8/dist-packages (from numba) (0.39.1)\n",
      "Requirement already satisfied: numpy<1.24,>=1.18 in /usr/local/lib/python3.8/dist-packages (from numba) (1.21.6)\n",
      "Requirement already satisfied: setuptools in /usr/local/lib/python3.8/dist-packages (from numba) (57.4.0)\n",
      "Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.8/dist-packages (from numba) (6.0.0)\n",
      "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.8/dist-packages (from importlib-metadata->numba) (3.11.0)\n"
     ]
    }
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cb12288b",
   "metadata": {
    "id": "cb12288b"
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "from transformers import DataCollatorForLanguageModeling, TrainingArguments, Trainer\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "\n",
    "import torch\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "from datasets import DatasetDict, Dataset"
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "def on_gpu(f):\n",
    "    def wrapper():\n",
    "        if torch.cuda.is_available():\n",
    "            return f()\n",
    "        else:\n",
    "            print('cuda unavailable')\n",
    "    return wrapper"
   ],
   "metadata": {
    "id": "O73s7d_Lt2FK"
   },
   "id": "O73s7d_Lt2FK",
   "execution_count": 3,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "if torch.cuda.is_available():\n",
    "    from pynvml import *\n",
    "    from numba import cuda\n",
    "\n",
    "@on_gpu\n",
    "def print_gpu_utilization():\n",
    "    try:\n",
    "        nvmlInit()\n",
    "        handle = nvmlDeviceGetHandleByIndex(0)\n",
    "        info = nvmlDeviceGetMemoryInfo(handle)\n",
    "        print(f\"GPU memory occupied: {info.used//1024**2} MB.\")\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "\n",
    "@on_gpu\n",
    "def free_gpu_cache():\n",
    "    print(\"Initial GPU Usage\")\n",
    "    print_gpu_utilization()                          \n",
    "\n",
    "    torch.cuda.empty_cache()\n",
    "\n",
    "    print(\"GPU Usage after emptying the cache\")\n",
    "    print_gpu_utilization()"
   ],
   "metadata": {
    "id": "blBIHGjItula"
   },
   "id": "blBIHGjItula",
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "! mkdir output"
   ],
   "metadata": {
    "id": "Nf0oMkQ4YYhr",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "b51de29a-5d56-4308-97a0-b5942cad85a2"
   },
   "id": "Nf0oMkQ4YYhr",
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "mkdir: cannot create directory ‘output’: File exists\n"
     ]
    }
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "60c073d5",
   "metadata": {
    "id": "60c073d5"
   },
   "outputs": [],
   "source": [
    "BASIC_PATH = './'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d92c1117",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 633
    },
    "id": "d92c1117",
    "outputId": "43983de2-38a2-4971-a7fa-f5cafb989576"
   },
   "outputs": [],
   "source": [
    "data = pd.read_csv('./full.csv', header=0)\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "60b4e18c",
   "metadata": {
    "id": "60b4e18c"
   },
   "outputs": [],
   "source": [
    "X = data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2fb3f341",
   "metadata": {
    "id": "2fb3f341"
   },
   "outputs": [],
   "source": [
    "X_train, X_test = train_test_split(X, test_size=0.15, random_state=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "bd37d65b",
   "metadata": {
    "id": "bd37d65b"
   },
   "outputs": [],
   "source": [
    "MODEL_NAME = 'BlackSamorez/rudialogpt3_medium_based_on_gpt2_2ch'"
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "BOS = '<bos>'\n",
    "EOS = '<eos>'\n",
    "CTX = '<ctx>'\n",
    "RPL = '<rpl>'\n",
    "\n",
    "SPECIAL_TOKENS = {\n",
    "    'bos_token': BOS,\n",
    "    'eos_token': EOS,\n",
    "    'additional_special_tokens': [CTX, RPL]\n",
    "}"
   ],
   "metadata": {
    "id": "fLIbtAckUBmd"
   },
   "id": "fLIbtAckUBmd",
   "execution_count": 11,
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b3511623",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "b3511623",
    "outputId": "cf2fa265-ac5c-43bf-9250-203cbcdb057f"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)"
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "num_new_tokens = tokenizer.add_special_tokens(SPECIAL_TOKENS)\n",
    "model.resize_token_embeddings(len(tokenizer))"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "eNwtqCS8hZsL",
    "outputId": "ee5aa53a-6bac-4ebc-c016-8b8233618df4"
   },
   "id": "eNwtqCS8hZsL",
   "execution_count": 13,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "Embedding(50262, 1024)"
      ]
     },
     "metadata": {},
     "execution_count": 13
    }
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "26accbe4",
   "metadata": {
    "id": "26accbe4"
   },
   "outputs": [],
   "source": [
    "def tokenization(example):\n",
    "    return tokenizer(\n",
    "        BOS + CTX + example['prompt'] + RPL + example['answer'] + EOS, \n",
    "        max_length=1024,\n",
    "        padding='max_length', \n",
    "        truncation=True, \n",
    "        add_special_tokens=False,\n",
    "        # return_tensors='pt'\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "class TelegramDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, data):\n",
    "        self.input_ids = []\n",
    "        self.attn_masks = []\n",
    "        self.labels = []\n",
    "        for idx, row in data.iterrows():\n",
    "            encodings_dict = tokenization(row)\n",
    "            self.input_ids.append(torch.tensor(encodings_dict['input_ids']))\n",
    "            self.attn_masks.append(torch.tensor(encodings_dict['attention_mask']))\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.input_ids)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.input_ids[idx], self.attn_masks[idx]"
   ],
   "metadata": {
    "id": "gHIMBQCkPRAF"
   },
   "id": "gHIMBQCkPRAF",
   "execution_count": 15,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "def collate(data):\n",
    "    return {\n",
    "        'input_ids': torch.stack([f[0] for f in data]),\n",
    "        'attention_mask': torch.stack([f[1] for f in data]),\n",
    "        'labels': torch.stack([f[0] for f in data])\n",
    "    }\n",
    "\n",
    "def load_data_collator(tokenizer, mlm = False):\n",
    "    # data_collator = DataCollatorForLanguageModeling(\n",
    "    #     tokenizer=tokenizer, \n",
    "    #     mlm=mlm,\n",
    "    # )\n",
    "    return collate"
   ],
   "metadata": {
    "id": "7BTvVaccjhR4"
   },
   "id": "7BTvVaccjhR4",
   "execution_count": 16,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "# X_train = X_train[:int(len(X_train) * 0.1)]\n",
    "# X_test = X_test[:int(len(X_test) * 0.1)]"
   ],
   "metadata": {
    "id": "RdWNcu6j0D5W"
   },
   "id": "RdWNcu6j0D5W",
   "execution_count": 17,
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "17c80906",
   "metadata": {
    "id": "17c80906"
   },
   "outputs": [],
   "source": [
    "dataset = DatasetDict({\n",
    "    'train': TelegramDataset(X_train),\n",
    "    'test': TelegramDataset(X_test),\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "# dataset = dataset.map(tokenization, batched=True)\n",
    "# dataset"
   ],
   "metadata": {
    "id": "50madHAuZPuM"
   },
   "id": "50madHAuZPuM",
   "execution_count": 19,
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "257e5e1d",
   "metadata": {
    "id": "257e5e1d"
   },
   "outputs": [],
   "source": [
    "training_args = TrainingArguments(\n",
    "    output_dir=f'{BASIC_PATH}/training',\n",
    "    do_train=True,\n",
    "    do_eval=True,\n",
    "    seed=42,\n",
    "    learning_rate=2e-5,\n",
    "    weight_decay=1e-2,\n",
    "\n",
    "    # 'lr_scheduler_type': 'cosine',\n",
    "    gradient_accumulation_steps=10,\n",
    "    per_device_train_batch_size=1,\n",
    "    num_train_epochs=5,\n",
    "    warmup_steps=0,\n",
    "    max_grad_norm=1000,\n",
    "    adam_epsilon=1e-6,\n",
    "\n",
    "    log_level='debug',\n",
    "    save_strategy='steps',\n",
    "    disable_tqdm=False,\n",
    "    logging_steps=100,\n",
    "    eval_steps=100,\n",
    "    save_steps=500,\n",
    "    resume_from_checkpoint=True,\n",
    "    per_device_eval_batch_size=2,\n",
    "    evaluation_strategy='steps',\n",
    "    fp16=True,\n",
    "    # metric_for_best_model='loss',\n",
    "    load_best_model_at_end=True,\n",
    "    # greater_is_better=False,\n",
    "    \n",
    "    save_total_limit=1,\n",
    "    \n",
    "    report_to='none',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "32a52f5b",
   "metadata": {
    "id": "32a52f5b",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "113e5a2f-cefd-47cf-ee50-52f8b7e9b2d2"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "Using cuda_amp half precision backend\n"
     ]
    }
   ],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    data_collator=load_data_collator(tokenizer),\n",
    "    train_dataset=dataset['train'],\n",
    "    eval_dataset=dataset['test'],\n",
    "    tokenizer=tokenizer,\n",
    "    # compute_metrics=compute_metrics,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19e5418b",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 265
    },
    "id": "19e5418b",
    "outputId": "31359568-df31-464b-8be9-e28094275974"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "/usr/local/lib/python3.8/dist-packages/transformers/optimization.py:306: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n",
      "***** Running training *****\n",
      "  Num examples = 4606\n",
      "  Num Epochs = 5\n",
      "  Instantaneous batch size per device = 1\n",
      "  Total train batch size (w. parallel, distributed & accumulation) = 10\n",
      "  Gradient Accumulation steps = 10\n",
      "  Total optimization steps = 2300\n",
      "  Number of trainable parameters = 355876864\n"
     ]
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='8' max='2300' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [   8/2300 00:26 < 2:50:06, 0.22 it/s, Epoch 0.02/5]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "  </tbody>\n",
       "</table><p>"
      ]
     },
     "metadata": {}
    }
   ],
   "source": [
    "results = trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "assert model.transformer.wte.weight.shape[0] == len(tokenizer)"
   ],
   "metadata": {
    "id": "4xiKQOqoymgh"
   },
   "id": "4xiKQOqoymgh",
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "free_gpu_cache()"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "fr-FGJTGt74Z",
    "outputId": "efd226c2-a88c-4cae-c734-63a4f2cec773"
   },
   "id": "fr-FGJTGt74Z",
   "execution_count": 22,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Initial GPU Usage\n",
      "GPU memory occupied: 15060 MB.\n",
      "GPU Usage after emptying the cache\n",
      "GPU memory occupied: 15060 MB.\n"
     ]
    }
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "results"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "W9XqfB4hZs7w",
    "outputId": "974bd8a2-8e15-48c9-ae86-d0b489c1c54e"
   },
   "id": "W9XqfB4hZs7w",
   "execution_count": null,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "TrainOutput(global_step=462, training_loss=0.6178220856241333, metrics={'train_runtime': 120.6189, 'train_samples_per_second': 15.288, 'train_steps_per_second': 3.83, 'total_flos': 244269883121664.0, 'train_loss': 0.6178220856241333, 'epoch': 2.0})"
      ]
     },
     "metadata": {},
     "execution_count": 21
    }
   ]
  },
  {
   "cell_type": "code",
   "source": [],
   "metadata": {
    "id": "codvgE1umSfW"
   },
   "id": "codvgE1umSfW",
   "execution_count": null,
   "outputs": []
  }
 ],
 "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.8.12"
  },
  "colab": {
   "provenance": []
  },
  "accelerator": "GPU",
  "gpuClass": "standard"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}