best_model(45).ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "18dde1a7-cfba-4b88-8431-664f6014a909",
   "metadata": {},
   "source": [
    "## Load_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19a9eff1-6362-40a1-a38a-be376dbcf510",
   "metadata": {},
   "outputs": [],
   "source": [
    "!df -h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "19fcefbd-bcbd-40d2-9011-5d44e2bac431",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle as pickle\n",
    "import os\n",
    "import pandas as pd\n",
    "import torch\n",
    "from sklearn.model_selection import train_test_split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "dd9d9a40-f84e-429b-84ab-84a81f80c315",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda:0\n"
     ]
    }
   ],
   "source": [
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8875affb-c6fa-4097-ad86-5bd00746a327",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RE_Dataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, tokenized_dataset, labels):\n",
    "        self.tokenized_dataset = tokenized_dataset\n",
    "        self.labels = labels\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        item = {key: torch.tensor(val[idx]) for key, val in self.tokenized_dataset.items()}\n",
    "        item['labels'] = torch.tensor(self.labels[idx])\n",
    "        return item\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4ca06956-77a3-4ea9-8947-121291aef38e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocessing_dataset(dataset, label_type):\n",
    "    label = []\n",
    "    for i in dataset[8]:\n",
    "        if i == 'blind':\n",
    "            label.append(100)\n",
    "        else:\n",
    "            label.append(label_type[i])\n",
    "    out_dataset = pd.DataFrame({'sentence':dataset[1],'entity_01':dataset[2],'entity_02':dataset[5],'label':label,})\n",
    "    return out_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a8c73b40-068c-4e64-a789-40e4f7d08c3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_data(dataset_dir):\n",
    "    # load label_type, classes\n",
    "    with open('/opt/ml/input/data/label_type.pkl', 'rb') as f:\n",
    "        label_type = pickle.load(f)\n",
    "    # load dataset\n",
    "    dataset = pd.read_csv(dataset_dir, delimiter='\\t', header=None)\n",
    "    # preprecessing dataset\n",
    "    dataset = preprocessing_dataset(dataset, label_type)\n",
    "  \n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6d63608d-d16a-458b-b793-3af270fbd651",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenized_dataset(dataset, tokenizer):\n",
    "    concat_entity = []\n",
    "    for e01, e02 in zip(dataset['entity_01'], dataset['entity_02']):\n",
    "        temp = ''\n",
    "        temp = e01 + ' </s> ' + e02\n",
    "        concat_entity.append(temp)\n",
    "    tokenized_sentences = tokenizer(\n",
    "        concat_entity,\n",
    "        list(dataset['sentence']),\n",
    "        return_tensors=\"pt\",\n",
    "        padding=True,\n",
    "        truncation='only_second',\n",
    "        max_length=128,\n",
    "        add_special_tokens=True,\n",
    "    )\n",
    "    return tokenized_sentences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "29ad6558-8414-46a1-838a-f28067b1e627",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train = load_data('./input/data/train/train.tsv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc20e3bb-6003-4abb-9b5a-eaab04216875",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(len(df_train)):\n",
    "    df_train['sentence'][i] = df_train['sentence'][i] + f\" [SEP] 이 문장에서 {df_train['entity_01'][i]}과 {df_train['entity_02'][i]}는 어떤 관계야?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c324d56-24b3-4f0b-9909-830045d56342",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 중복 제거\n",
    "df_train.drop_duplicates(subset=['sentence'], inplace=True)\n",
    "len(df_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f8df0aa9-601a-4af9-a7f7-09801f06624b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "8985"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_train = df_train[~df_train['label'].isin((41, 37, 40, 29))]\n",
    "len(df_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96a77d1b-c68e-4a61-ac7f-177b605934f0",
   "metadata": {},
   "source": [
    "## Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3b790638-b8af-4c71-b31b-bc2aa30fdadb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle as pickle\n",
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import re\n",
    "import random\n",
    "import torch\n",
    "from sklearn.metrics import accuracy_score\n",
    "from transformers import AutoTokenizer, BertForSequenceClassification, Trainer, TrainingArguments, BertConfig, ElectraTokenizer\n",
    "# from tokenization_kobert import KoBertTokenizer\n",
    "# from load_data import *\n",
    "\n",
    "import argparse\n",
    "from importlib import import_module\n",
    "from pathlib import Path\n",
    "import glob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8a9cb335-dd19-4047-9919-419f1e103b09",
   "metadata": {},
   "outputs": [],
   "source": [
    "def seed_everything(seed):\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed)  # if use multi-GPU\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    torch.backends.cudnn.benchmark = False\n",
    "    np.random.seed(seed)\n",
    "    random.seed(seed)\n",
    "    \n",
    "seed_everything(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4110ec26-7aee-4de9-b87d-2d69aaed72bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_metrics(pred):\n",
    "    labels = pred.label_ids\n",
    "    preds = pred.predictions.argmax(-1)\n",
    "    # calculate accuracy using sklearn's function\n",
    "    acc = accuracy_score(labels, preds)\n",
    "    return {\n",
    "        'accuracy': acc,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65464003-bc15-4010-944a-cc296d1b6440",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "32c8e509-d52a-4db3-bc30-d0ffe371ea12",
   "metadata": {},
   "outputs": [],
   "source": [
    "def increment_output_dir(output_path, exist_ok=False):\n",
    "    path = Path(output_path)\n",
    "    if (path.exists() and exist_ok) or (not path.exists()):\n",
    "        return str(path)\n",
    "    else:\n",
    "        dirs = glob.glob(f\"{path}*\")\n",
    "        matches = [re.search(rf\"%s(\\d+)\" %path.stem, d) for d in dirs]\n",
    "        i = [int(m.groups()[0]) for m in matches if m]\n",
    "        n = max(i) + 1 if i else 2\n",
    "        return f\"{path}{n}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "1f0a2fec-3bfb-490a-ae72-d4261091fabf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, BertForSequenceClassification, Trainer, TrainingArguments, BertConfig, ElectraTokenizer, ElectraConfig, ElectraForSequenceClassification, XLMRobertaTokenizer, XLMRobertaForSequenceClassification, XLMRobertaForTokenClassification, XLMRobertaConfig\n",
    "\n",
    "MODEL_NAME = \"xlm-roberta-large\"\n",
    "tokenizer = XLMRobertaTokenizer.from_pretrained(MODEL_NAME)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1ccfb8a3-cfa6-4828-ba13-8ddb4985ff26",
   "metadata": {},
   "outputs": [],
   "source": [
    "train, val = train_test_split(df_train, test_size=0.15, random_state=0)\n",
    "\n",
    "tokenized_train = tokenized_dataset(train, tokenizer)\n",
    "tokenized_val = tokenized_dataset(val, tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "63620b81-35d0-4f16-a74b-e8be6a61ab55",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_label = train['label'].values\n",
    "val_label = val['label'].values\n",
    "\n",
    "train_dataset = RE_Dataset(tokenized_train, train_label)\n",
    "val_dataset = RE_Dataset(tokenized_val, val_label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "a8dc1781-f2a8-4d41-be75-2be8c805fb64",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at xlm-roberta-large were not used when initializing XLMRobertaForSequenceClassification: ['lm_head.bias', 'lm_head.dense.weight', 'lm_head.dense.bias', 'lm_head.layer_norm.weight', 'lm_head.layer_norm.bias', 'lm_head.decoder.weight', 'roberta.pooler.dense.weight', 'roberta.pooler.dense.bias']\n",
      "- This IS expected if you are initializing XLMRobertaForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing XLMRobertaForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of XLMRobertaForSequenceClassification were not initialized from the model checkpoint at xlm-roberta-large and are newly initialized: ['classifier.dense.weight', 'classifier.dense.bias', 'classifier.out_proj.weight', 'classifier.out_proj.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "XLMRobertaForSequenceClassification(\n",
       "  (roberta): RobertaModel(\n",
       "    (embeddings): RobertaEmbeddings(\n",
       "      (word_embeddings): Embedding(250002, 1024, padding_idx=1)\n",
       "      (position_embeddings): Embedding(514, 1024, padding_idx=1)\n",
       "      (token_type_embeddings): Embedding(1, 1024)\n",
       "      (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (encoder): RobertaEncoder(\n",
       "      (layer): ModuleList(\n",
       "        (0): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (1): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (2): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (3): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (4): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (5): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (6): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (7): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (8): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (9): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (10): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (11): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (12): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (13): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (14): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (15): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (16): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (17): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (18): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (19): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (20): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (21): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (22): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (23): RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (key): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (value): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (classifier): RobertaClassificationHead(\n",
       "    (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "    (dropout): Dropout(p=0.1, inplace=False)\n",
       "    (out_proj): Linear(in_features=1024, out_features=42, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_config = XLMRobertaConfig.from_pretrained(MODEL_NAME)\n",
    "model_config.num_labels = 42\n",
    "model = XLMRobertaForSequenceClassification.from_pretrained(MODEL_NAME, config=model_config)\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "3f963b6f-0bd0-4e10-a99a-b78d2149b234",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_dir = increment_output_dir('./results')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90bea097-bcc9-4135-932f-f8b23b3b2729",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_args = TrainingArguments(\n",
    "    output_dir=output_dir,          # output directory\n",
    "    save_total_limit=3,              # number of total save model.\n",
    "    save_steps=500,                 # model saving step.\n",
    "    num_train_epochs=20,              # total number of training epochs\n",
    "    learning_rate=5e-5,               # learning_rate\n",
    "    per_device_train_batch_size=32,  # batch size per device during training\n",
    "    per_device_eval_batch_size=32,   # batch size for evaluation\n",
    "    warmup_steps=500,                # number of warmup steps for learning rate scheduler\n",
    "    weight_decay=0.01,               # strength of weight decay\n",
    "    logging_dir='./logs',            # directory for storing logs\n",
    "    logging_steps=100,              # log saving step.\n",
    "    metric_for_best_model='accuracy',\n",
    "    evaluation_strategy='steps', # evaluation strategy to adopt during training\n",
    "                                # `no`: No evaluation during training.\n",
    "                                # `steps`: Evaluate every `eval_steps`.\n",
    "                                # `epoch`: Evaluate every end of epoch.\n",
    "    eval_steps = 500,            # evaluation step.\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ebd6aca-4560-49d2-a2a3-b167422f92c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,                         # the instantiated 🤗 Transformers model to be trained\n",
    "    args=training_args,                  # training arguments, defined above\n",
    "    train_dataset=train_dataset,         # training dataset\n",
    "    eval_dataset=val_dataset,             # evaluation dataset\n",
    "    compute_metrics=compute_metrics,         # define metrics function\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd825e6a-6f78-47b5-9bfe-66471cbe001c",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1375d4eb-a4b4-4757-9881-cd333f84885c",
   "metadata": {},
   "source": [
    "## Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ef15ed0-70a6-4286-a7a0-6ac45ced554c",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.evaluate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a35d3b56-5eeb-4aea-93e3-9909e2455e9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# trainer.save_model('./results/checkpoint-4500/')\n",
    "# trainer.save_state()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7376aa1d-5b43-48c0-ba7c-00421a7dfcc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, BertForSequenceClassification, Trainer, TrainingArguments, BertConfig, BertTokenizer, ElectraTokenizer, ElectraForSequenceClassification\n",
    "# from tokenization_kobert import KoBertTokenizer\n",
    "from torch.utils.data import DataLoader\n",
    "# from load_data import *\n",
    "import pandas as pd\n",
    "import torch\n",
    "import pickle as pickle\n",
    "import numpy as np\n",
    "import argparse\n",
    "\n",
    "# import argparse\n",
    "from importlib import import_module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1827e83-3e4e-4e66-99a9-0bb9b7549527",
   "metadata": {},
   "outputs": [],
   "source": [
    "def inference(model, tokenized_sent, device):\n",
    "    dataloader = DataLoader(tokenized_sent, batch_size=40, shuffle=False)\n",
    "    model.eval()\n",
    "    output_pred = []\n",
    "  \n",
    "    for i, data in enumerate(dataloader):\n",
    "        with torch.no_grad():\n",
    "            outputs = model(\n",
    "                input_ids=data['input_ids'].to(device),\n",
    "                attention_mask=data['attention_mask'].to(device),\n",
    "#                 token_type_ids=data['token_type_ids'].to(device)\n",
    "            )\n",
    "        logits = outputs[0]\n",
    "        logits = logits.detach().cpu().numpy()\n",
    "        result = np.argmax(logits, axis=-1)\n",
    "\n",
    "        output_pred.append(result)\n",
    "  \n",
    "    return np.array(output_pred).flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cf908e8-abc0-4c0d-b14a-ef76197c0e8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_test_dataset(dataset_dir, tokenizer):\n",
    "    test_dataset = load_data(dataset_dir)\n",
    "    test_label = test_dataset['label'].values\n",
    "    # tokenizing dataset\n",
    "    tokenized_test = tokenized_dataset(test_dataset, tokenizer)\n",
    "    return tokenized_test, test_label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cde29e4-4b75-4d0f-b9e6-c588facf1d7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "TOK_NAME = \"xlm-roberta-large\"\n",
    "tokenizer = XLMRobertaTokenizer.from_pretrained(TOK_NAME)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bab6bb1-d423-4f0e-88ce-a979131f5998",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = XLMRobertaForSequenceClassification.from_pretrained('./results/checkpoint-4500/')\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91cc9970-faef-4ff3-b2e4-8d13cbe663db",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset_dir = \"/opt/ml/input/data/test/test.tsv\"\n",
    "test_data = load_data(test_dataset_dir)\n",
    "test_label = test_data['label']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9ccf44c-a815-4632-8e23-70efde64683d",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a6fdeef-da58-4afe-b41f-5ebf463e82c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(len(test_data)):\n",
    "    test_data['sentence'][i] = test_data['sentence'][i] + f\" [SEP] 이 문장에서 {test_data['entity_01'][i]}과 {test_data['entity_02'][i]}는 어떤 관계야?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05c1be7e-9bdd-4e03-9de0-6174a16bfb80",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_test = tokenized_dataset(test_data, tokenizer)\n",
    "test_dataset = RE_Dataset(tokenized_test ,test_label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "522a6f12-c25c-491f-bd58-dce669dfb445",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_answer = inference(model, test_dataset, device)\n",
    "output = pd.DataFrame(pred_answer, columns=['pred'])\n",
    "output.to_csv(\"./prediction/submission.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "272dbb07-d3b2-4f3b-bd12-44cb4e9dd4d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Audio\n",
    "Audio(\"https://upload.wikimedia.org/wikipedia/commons/0/05/Beep-09.ogg\", autoplay=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d3ee78b-f292-4148-b717-17aaefa01403",
   "metadata": {},
   "source": [
    "### check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "006aa72c-e20b-46e6-ace8-7928d5ee8732",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train = load_data('./input/data/train/train.tsv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fc56037-b18b-42c6-b371-4e08494450f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred = pd.read_csv('./prediction/submission.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fb67fcb-44b2-4cb3-9cd0-c92de079f629",
   "metadata": {},
   "outputs": [],
   "source": [
    "cnt = pred['pred'] == df_train['label']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf5be60d-781d-40c6-93e0-457314a21392",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(cnt.sum()/len(cnt))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc9a1eb7-216f-4257-a8fc-8293280eb090",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}