train_BertForPreTraining.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "eef28ac0-4918-4e7e-b5fb-ced2fc8f86d9",
   "metadata": {},
   "source": [
    "```\n",
    "Copyright (C) 2023 Donggeon Lee\n",
    " \n",
    "This program is free software: you can redistribute it and/or modify\n",
    "it under the terms of the GNU Affero General Public License as\n",
    "published by the Free Software Foundation, either version 3 of the\n",
    "License, or (at your option) any later version.\n",
    " \n",
    "This program is distributed in the hope that it will be useful,\n",
    "but WITHOUT ANY WARRANTY; without even the implied warranty of\n",
    "MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the\n",
    "GNU Affero General Public License for more details.\n",
    " \n",
    "You should have received a copy of the GNU Affero General Public License\n",
    "along with this program. If not, see <http://www.gnu.org/licenses/>.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "21737c46-7955-4254-890e-67d78b1449e9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-11-21 11:19:17.990208: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-11-21 11:19:18.186614: E tensorflow/stream_executor/cuda/cuda_blas.cc:2981] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2023-11-21 11:21:28.557506: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/nvidia/lib:/usr/local/nvidia/lib64\n",
      "2023-11-21 11:21:28.557791: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/nvidia/lib:/usr/local/nvidia/lib64\n",
      "2023-11-21 11:21:28.557817: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = str(0)\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "import random\n",
    "import pandas as pd\n",
    "from sklearn.model_selection import train_test_split\n",
    "from transformers import BertTokenizer, BertForPreTraining, Trainer, TrainingArguments, DataCollatorForLanguageModeling\n",
    "from tqdm import tqdm\n",
    "\n",
    "def set_seed(seed=42):\n",
    "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "\n",
    "SEED = 42\n",
    "MAX_SEQ_LEN = 512\n",
    "set_seed(SEED)\n",
    "\n",
    "\n",
    "class BertDataset(Dataset):\n",
    "    def __init__(self, corpus, tokenizer, max_length):\n",
    "        self.corpus = corpus\n",
    "        self.tokenizer = tokenizer\n",
    "        self.max_length = max_length\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.corpus)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        report = self.corpus[idx]\n",
    "        sentences = report.split('\\n')\n",
    "        tokenized_sentences = [self.tokenizer.encode(sentence, add_special_tokens=False) for sentence in sentences]\n",
    "\n",
    "        # Choose two sentences for NSP\n",
    "        if len(tokenized_sentences) >= 2:\n",
    "            # 50% of the time, the second sentence is a random sentence\n",
    "            if random.random() > 0.5:\n",
    "                second_sentence = random.choice(tokenized_sentences)\n",
    "                tokenized_sentences[1] = second_sentence\n",
    "                next_sentence_label = 0\n",
    "            else:\n",
    "                next_sentence_label = 1\n",
    "        else:\n",
    "            next_sentence_label = 0\n",
    "\n",
    "        # Combine sentences and add [CLS] and [SEP] tokens\n",
    "        tokens = [self.tokenizer.cls_token_id]\n",
    "        for sentence in tokenized_sentences:\n",
    "            tokens += sentence + [self.tokenizer.sep_token_id]\n",
    "        tokens = tokens[:self.max_length]\n",
    "\n",
    "        # Create segment ids\n",
    "        if self.tokenizer.sep_token_id in tokens:\n",
    "            token_type_ids = [0 if i <= tokens.index(self.tokenizer.sep_token_id) else 1 for i in range(len(tokens))]\n",
    "        else:\n",
    "            token_type_ids = [0] * len(tokens)\n",
    "\n",
    "\n",
    "        # MLM\n",
    "        probability_matrix = torch.full((len(tokens),), 0.15)\n",
    "        special_tokens_mask = [self.tokenizer.get_special_tokens_mask([val], already_has_special_tokens=True)[0] for val in tokens]\n",
    "        probability_matrix.masked_fill_(torch.tensor(special_tokens_mask, dtype=torch.bool), value=0.0)\n",
    "        masked_indices = torch.bernoulli(probability_matrix).bool()\n",
    "        labels = torch.where(masked_indices, torch.tensor(tokens), torch.tensor(-100))\n",
    "\n",
    "\n",
    "\n",
    "        # Padding\n",
    "        padding_length = self.max_length - len(tokens)\n",
    "        tokens = tokens + ([self.tokenizer.pad_token_id] * padding_length)\n",
    "        labels = labels.tolist() + ([-100] * padding_length)\n",
    "        token_type_ids = token_type_ids + ([0] * padding_length)\n",
    "\n",
    "        attention_mask = [1 if token != self.tokenizer.pad_token_id else 0 for token in tokens]\n",
    "\n",
    "        return {\"input_ids\": tokens, \"attention_mask\": attention_mask, \"token_type_ids\": token_type_ids, \"next_sentence_label\": next_sentence_label, \"labels\": labels}\n",
    "\n",
    "\n",
    "class CustomTrainer(Trainer):\n",
    "    def compute_loss(self, model, inputs, return_outputs=False):\n",
    "        labels = inputs.pop(\"labels\")\n",
    "        next_sentence_label = inputs.pop(\"next_sentence_label\")\n",
    "        outputs = model(**inputs)\n",
    "        prediction_logits = outputs.prediction_logits\n",
    "        next_sentence_logits = outputs.seq_relationship_logits\n",
    "\n",
    "        nsp_loss = torch.nn.CrossEntropyLoss()(next_sentence_logits.view(-1, 2), next_sentence_label.view(-1))\n",
    "        mlm_loss = torch.nn.CrossEntropyLoss()(prediction_logits.view(-1, self.model.config.vocab_size), labels.view(-1))\n",
    "\n",
    "        loss = nsp_loss + mlm_loss\n",
    "\n",
    "        return (loss, outputs) if return_outputs else loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6168cf2b-3aa5-4f61-b776-f5bc26c16128",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Token indices sequence length is longer than the specified maximum sequence length for this model (622 > 512). Running this sequence through the model will result in indexing errors\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='6624' max='6624' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [6624/6624 39:51, Epoch 3/3]\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",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>500</td>\n",
       "      <td>0.772200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1000</td>\n",
       "      <td>0.413600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1500</td>\n",
       "      <td>0.387500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2000</td>\n",
       "      <td>0.368100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2500</td>\n",
       "      <td>0.337900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3000</td>\n",
       "      <td>0.352900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3500</td>\n",
       "      <td>0.338000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4000</td>\n",
       "      <td>0.326700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4500</td>\n",
       "      <td>0.333800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5000</td>\n",
       "      <td>0.333400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5500</td>\n",
       "      <td>0.299400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6000</td>\n",
       "      <td>0.310800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6500</td>\n",
       "      <td>0.310300</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=6624, training_loss=0.37455422809158545, metrics={'train_runtime': 2517.2658, 'train_samples_per_second': 21.044, 'train_steps_per_second': 2.631, 'total_flos': 1.4039613220737024e+16, 'train_loss': 0.37455422809158545, 'epoch': 3.0})"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# BERT 모델 및 토크나이저 불러오기\n",
    "model = BertForPreTraining.from_pretrained('klue/bert-base')\n",
    "tokenizer = BertTokenizer.from_pretrained('klue/bert-base')\n",
    "\n",
    "# csv 파일 로드\n",
    "df = pd.read_csv('data/train/BERT_trainset.csv')\n",
    "\n",
    "# 결측치 제거\n",
    "df = df.dropna(axis=0)\n",
    "\n",
    "# 'text' 열의 문장들을 리스트로 변환\n",
    "corpus = df['text'].tolist()\n",
    "\n",
    "# Split the data\n",
    "train_corpus, val_corpus = train_test_split(corpus, test_size=0.1, random_state=SEED)\n",
    "\n",
    "# Create datasets\n",
    "train_dataset = BertDataset(train_corpus, tokenizer, max_length=MAX_SEQ_LEN)\n",
    "val_dataset = BertDataset(val_corpus, tokenizer, max_length=MAX_SEQ_LEN)\n",
    "\n",
    "# Training arguments\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./results\",\n",
    "    num_train_epochs=3,\n",
    "    per_device_train_batch_size=8,\n",
    "    per_device_eval_batch_size=8,\n",
    "    warmup_steps=500,\n",
    "    weight_decay=0.01,\n",
    "    logging_dir='./logs',\n",
    "    save_total_limit=3\n",
    ")\n",
    "\n",
    "# Create trainer\n",
    "trainer = CustomTrainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=val_dataset\n",
    ")\n",
    "\n",
    "# Train the model\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1afb67df-1d98-405c-a3e8-3b7ae038af4c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BertForPreTraining(\n",
       "  (bert): BertModel(\n",
       "    (embeddings): BertEmbeddings(\n",
       "      (word_embeddings): Embedding(32000, 768, padding_idx=0)\n",
       "      (position_embeddings): Embedding(512, 768)\n",
       "      (token_type_embeddings): Embedding(2, 768)\n",
       "      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (encoder): BertEncoder(\n",
       "      (layer): ModuleList(\n",
       "        (0-11): 12 x BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (intermediate_act_fn): GELUActivation()\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (pooler): BertPooler(\n",
       "      (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "      (activation): Tanh()\n",
       "    )\n",
       "  )\n",
       "  (cls): BertPreTrainingHeads(\n",
       "    (predictions): BertLMPredictionHead(\n",
       "      (transform): BertPredictionHeadTransform(\n",
       "        (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (transform_act_fn): GELUActivation()\n",
       "        (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "      )\n",
       "      (decoder): Linear(in_features=768, out_features=32000, bias=True)\n",
       "    )\n",
       "    (seq_relationship): Linear(in_features=768, out_features=2, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "34cf94b2-137a-4fc4-8939-3524b6a1aa9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained(\"model\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "148c7192-a384-402e-a364-0890f3002073",
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import notebook_login"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2965079d-e57f-4146-928b-256fd8ee0141",
   "metadata": {},
   "outputs": [],
   "source": [
    "# notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "63d835d3-ee87-4be7-b3b5-d2e2012e09e9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c025b30699734f2ca95452463fa94aae",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pytorch_model.bin:   0%|          | 0.00/445M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/oneonlee/KoAirBERT/commit/811ee837de8be44405a7907b58b0613d00f84fc1', commit_message='Upload BertForPreTraining', commit_description='', oid='811ee837de8be44405a7907b58b0613d00f84fc1', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.push_to_hub(\"KoAirBERT\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9b4115c6-ea28-40bc-be73-a9a30a3a1087",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters:  111243010\n"
     ]
    }
   ],
   "source": [
    "num_parameters = model.num_parameters()\n",
    "print(\"Number of parameters: \", num_parameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "269f3fec-990a-459b-a378-71f8a105fc0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "exit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "035451dd-9e27-4990-b65d-ef7a1c5018f8",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}