Natural_Language_Processing/Table_pretraining_with_TAPAS/Table_Pretraining_with_TapasForMaskedLM_KYK.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4b321e96-19ff-4470-9b22-f156e2af2210",
   "metadata": {},
   "source": [
    "# Table Pre-training with `TapasForMaskedLM`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5fd9857-1b30-414b-96ee-f4d267336a2a",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "**Author**: [Yookyung Kho](https://github.com/yookyungkho)\n",
    "\n",
    "**Date presented**: 2022/07/25, DSBA keras2torch Study\n",
    "\n",
    "**Task description**: Table-aware Masked Language Model\n",
    "\n",
    "   - Pre-Training Bert-based Tapas Model with KorWikiTabular dataset for Korean Table MRC(ex. QA)\n",
    "\n",
    "\n",
    "**References**:\n",
    "\n",
    "- https://huggingface.co/docs/transformers/model_doc/tapas\n",
    "\n",
    "- https://github.com/huggingface/transformers/blob/v4.11.3/src/transformers/models/tapas/tokenization_tapas.py\n",
    "\n",
    "- https://github.com/huggingface/transformers/blob/v4.11.3/src/transformers/data/data_collator.py\n",
    "\n",
    "- https://github.com/hwk0702/keras2torch/blob/main/Natural_Language_Processing/Question_Answering_Huggingface/QA_huggingface_KYK.ipynb\n",
    "\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "719b8693-db19-4284-a902-0802fc860387",
   "metadata": {},
   "source": [
    "## 📜Prerequisite: About Table QA"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d216b32-e543-4889-8527-7bb0fe566e56",
   "metadata": {},
   "source": [
    "<img src=\"imgs/table_qa_in_tapas.PNG\" width=\"900\" height=\"600\">\n",
    "\n",
    "Table Question Answering이란, 주어진 query(질문)에 대한 정답을 질문과 매핑되어 있는 table에서 찾아내는 task를 일컫습니다.\n",
    "\n",
    "Table QA를 잘 풀기 위해서는 text와 table의 정보를 잘 함축한 joint representation을 학습해야 합니다.\n",
    "\n",
    "이를 위해서 아래와 같이 Downstream task인 Table QA를 풀기 이전에 대량의 table-text 데이터를 가지고 사전학습을 진행합니다.\n",
    "\n",
    "<img src=\"imgs/tapas_pretraining.PNG\" width=\"900\" height=\"500\">\n",
    "\n",
    "\n",
    "<br/>\n",
    "\n",
    "기존 table 연구 흐름을 살펴보면, 대표적으로 google research에서 table을 text와 함께 인코딩하는 방식에 대해 다양한 연구를 진행했습니다.\n",
    "\n",
    "이 노트북에서 다루게 될 모델인 Tapas 역시 google에서 발표한 구조입니다. ([paper](https://arxiv.org/pdf/2004.02349), [github](https://github.com/google-research/tapas))\n",
    "\n",
    "google은 자사의 github 페이지에 tensorflow로 구현된 코드를 공개하여 pytorch 유저의 입장에서는 활용도가 높지 않았는데, 얼마 전 huggingface 플랫폼을 통해 Tapas의 Pytorch 구현체가 공개되었습니다.(야호!😁)\n",
    "\n",
    "오늘은 Table을 Text와 함께 인코딩하는 가장 기본적인 방법론을 다루고 있는 TAPAS에 대해 살펴보겠습니다!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68bf953d-0d85-44f8-b7ac-d557246795f3",
   "metadata": {},
   "source": [
    "### 📦Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39ed8526-49ab-4e3c-82c6-bf3644141049",
   "metadata": {},
   "source": [
    "2022년 드디어! LG AI Research에서 한국어 버전의 Table MRC 데이터셋을 공개하였습니다!\n",
    "\n",
    "데이터셋과 함께 공개한 paper는 LREC 2022에 accept되었습니다.([Full paper 보러 가기](https://arxiv.org/abs/2201.06223))\n",
    "\n",
    "<img src=\"imgs/table_mrc.PNG\" width=\"900\" height=\"450\">\n",
    "\n",
    "공개된 데이터셋은 Pre-training(사전학습)용 데이터인 `KorWikiTabular`과 Fine-tuning(QA)용 데이터 `KorWikiTQ`, 이렇게 2가지로 나뉩니다.\n",
    "\n",
    "\n",
    "- 참고로, 데이터는 [LG AI Research Github](https://github.com/LG-NLP)의 [KorWikiTableQuestions Repository](https://github.com/LG-NLP/KorWikiTableQuestions)에서 다운받으실 수 있습니다."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54413f48-8e8a-4657-b2db-de92968ef627",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bc7a46f-af20-4af3-9ad9-78a460f752bd",
   "metadata": {},
   "source": [
    "## 0. 준비 과정"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "785ce089-870c-4f62-9620-40b24581497e",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 0.1. 데이터 샘플링"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3e42c779-2bfc-4c57-8cb0-e8af3cc8e1f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "file_path = \"data/KorWikiTabular.json\"\n",
    "\n",
    "with open(file_path, 'r', encoding=\"UTF-8\") as f:\n",
    "    tables = json.loads(f.read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e0af543c-b05c-4547-8f66-ad36acaf20d9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1196306"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(tables['data'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68a76d7a-9f15-40af-9288-86e9a0417bf2",
   "metadata": {},
   "source": [
    "앞서 소개한 한국어 table 사전학습 데이터(`KorWikiTabular`)는 무려 약 **120만 개**의 샘플을 포함하고 있습니다.😱\n",
    "\n",
    "따라서, 모든 데이터를 가지고 모델링을 진행하기에는 무리가 있어 보입니다.(_GPU 메모리도 부족할 뿐더러, 시간이 너무 오래 걸리겠죠?_)\n",
    "\n",
    "원활한 실험을 위해 **200개의 샘플만 랜덤으로 추출**하여 `sample_200_KorWikiTabular.json`의 형태로 `data` 경로에 저장하겠습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ae966c8-3860-4adf-a5f4-22f2fafc33f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "np.random.seed(602)\n",
    "\n",
    "randn_idxs = np.random.choice(list(range(len(tables['data']))), size=200, replace=False) #비복원추출\n",
    "\n",
    "sample_tables = {'data': []}\n",
    "\n",
    "for new_id, org_id in enumerate(randn_idxs):\n",
    "    sample_tables['data'].append(tables['data'][org_id])\n",
    "    sample_tables['data'][new_id]['org_idx'] = int(org_id)\n",
    "    \n",
    "file_path = \"data/sample_200_KorWikiTabular.json\"\n",
    "json.dump(sample_tables, open(file_path,'w'), indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ffecaa7-d121-401f-839d-8b90a3540fa2",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(sample_tables['data'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "157ab71d-9b3c-4bd1-995a-aa778fd64bb4",
   "metadata": {},
   "source": [
    "저장 완료되었습니다!\n",
    "\n",
    "앞으로는 위 4개 cell 실행 없이, sample data만 바로 불러들여와서 실험 진행하도록 하겠습니다."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb48b5d3-6928-4b04-b333-bbe8715cb8be",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4fe3839b-d860-43e3-858d-d8e82beff4fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "file_path = \"data/sample_200_KorWikiTabular.json\"\n",
    "\n",
    "with open(file_path, \"r\") as json_file:\n",
    "    sample_data = json.load(json_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "97e40c00-0b46-4d55-95a6-dce022773037",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "200"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(sample_data['data'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7906d161-595e-489c-9670-46c21d0256e6",
   "metadata": {},
   "source": [
    "### 0.2. Module Import, GPU 세팅"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "976a24e7-499d-49a6-b5a7-64c730f96b1f",
   "metadata": {},
   "source": [
    "지금처럼 소량(200개)의 샘플 데이터만 뽑아 쓰는 경우가 아니라 Full dataset을 모두 돌릴 때에는, local gpu만으로 MLM 학습이 불가능합니다.(이유: 메모리 부족)\n",
    "\n",
    "따라서, multi-gpu를 장착한 서버 환경에서 돌려야 하는데, 이 경우에 필요한 코드도 함께 적어두었으니 참고해주세요!\n",
    "\n",
    "\n",
    "\n",
    "> 💙Requirements\n",
    "> \n",
    "> - `torch == 1.12.0+cu113`\n",
    "> - `torch-scatter == 2.0.9`(for `TapasModel`)\n",
    "> - `transformers == 4.11.3`\n",
    "> - `pandas == 1.3.5`\n",
    "\n",
    "- 이 중 `torch-scatter` 설치 관련해서는 [이 link](https://github.com/rusty1s/pytorch_scatter)를 참고해주세요."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bb95a402-8324-4362-9b6f-09c15b01e7fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import wget\n",
    "import os\n",
    "import random\n",
    "import json\n",
    "import argparse\n",
    "import wandb\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "from transformers import TapasConfig, TapasTokenizer, TapasForMaskedLM, AdamW, get_scheduler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0a7c4aed-a5ac-4491-8caa-1d8acd59ee63",
   "metadata": {},
   "outputs": [],
   "source": [
    "# multi gpu\n",
    "\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"  # Arrange GPU devices starting from 0\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]= \"0,1,2,3\"  # Set the GPUs to use"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "aba2a7cb-0460-425f-ba62-903f461bfae1",
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser()\n",
    "\n",
    "parser.add_argument(\"--vocab_model_name\", default=\"klue/bert-base\", type=str)\n",
    "parser.add_argument(\"--tok_path\", default=\"table_tokenizer\", type=str)\n",
    "\n",
    "parser.add_argument(\"--max_seq_len\", default=512, type=int)\n",
    "parser.add_argument(\"--max_query_len\", default=470, type=int)\n",
    "parser.add_argument(\"--row_del_ratio\", default=0.9, type=float)\n",
    "parser.add_argument(\"--col_del_ratio\", default=0.9, type=float)\n",
    "parser.add_argument(\"--mlm_prob\", default=0.15, type=float)\n",
    "\n",
    "parser.add_argument(\"--random_seed\", default=602, type=int)\n",
    "\n",
    "parser.add_argument(\"--epoch\", default=5, type=int)\n",
    "parser.add_argument(\"--batch_size\", default=16, type=int)\n",
    "parser.add_argument(\"--learning_rate\", default=4e-4, type=float)\n",
    "parser.add_argument(\"--weight_decay\", default=0.0, type=float)\n",
    "parser.add_argument(\"--lr_scheduler_type\", default=\"linear\", type=str)\n",
    "parser.add_argument(\"--num_warmup_steps\", default=0, type=int)\n",
    "parser.add_argument(\"--eval_step\", default=10, type=int)\n",
    "\n",
    "parser.add_argument(\"--wandb_project\", default=\"Table Pretraining\", type=str)\n",
    "parser.add_argument(\"--wandb_name\", default=\"tapas-base-mlm-clean\", type=str)\n",
    "parser.add_argument(\"--wandb_entity\", default=\"yookyungkho\", type=str)\n",
    "\n",
    "parser.add_argument(\"--output_dir\", default=\"models/\", type=str)\n",
    "\n",
    "args = parser.parse_args([])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "eb3a2e46-87ec-4e1c-908f-7cd4ee62ac28",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Device: cuda\n",
      "Current cuda device: 0\n",
      "Count of using GPUs: 4\n",
      "number of gpu:  4\n"
     ]
    }
   ],
   "source": [
    "args.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\") #multi_gpu\n",
    "\n",
    "print('Device:', args.device)\n",
    "print('Current cuda device:', torch.cuda.current_device())\n",
    "print('Count of using GPUs:', torch.cuda.device_count())\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    n_gpu = torch.cuda.device_count()\n",
    "    print(\"number of gpu: \", n_gpu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ddfee8cb-c4ab-4af2-854d-7a6bb9d176cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# seed 고정\n",
    "def seed_everything(seed):\n",
    "    # random.seed(seed) #masking dynamics를 위해 이 부분은 주석 처리\n",
    "    np.random.seed(seed)\n",
    "    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)  # type: ignore\n",
    "    torch.backends.cudnn.deterministic = True  # type: ignore\n",
    "    torch.backends.cudnn.benchmark = True  # type: ignore\n",
    "\n",
    "seed_everything(args.random_seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d539b39-60f4-4624-88fc-8e66a3599294",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "이제 모델링을 위한 준비를 마쳤습니다.\n",
    "\n",
    "지금부터는 본격적으로 사전학습을 위한 데이터셋을 전처리하고 모델을 구축한 뒤, 실제 학습을 진행하도록 하겠습니다.\n",
    "\n",
    "렛츠공~💛\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e7e7d14-d379-4e9b-ad94-974ef27032a8",
   "metadata": {},
   "source": [
    "## 1. 데이터 전처리"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1a4e518-74ef-4ade-b1f6-e20a68a372d2",
   "metadata": {},
   "source": [
    "Tapas의 사전학습 task는 Table-aware Masked Language Model입니다.\n",
    "\n",
    "text와 table을 input으로 받아 하나의 입력 시퀀스를 생성하고, 시퀀스 내 특정 비율(ex. 15%)의 토큰에 총 2가지 방식의 masking을 수행합니다.\n",
    "\n",
    "<img src=\"imgs/tapas_masking.png\" width=\"1000\" height=\"600\">\n",
    "\n",
    "\n",
    "1. **Whole Word Masking**: 랜덤으로 선택된 단어의 모든 토큰을 masking\n",
    "\n",
    "2. **Whole Cell Masking**: 랜덤으로 선택된 table cell의 모든 토큰을 masking\n",
    "\n",
    "<br/>\n",
    "\n",
    "이렇게 2가지 방식의 masking을 통해 위 그림과 같은 Table MLM의 input을 생성하게 됩니다.\n",
    "\n",
    "지금부터 한단계씩 input 형태를 갖춰가는 과정을 살펴보겠습니다."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "284505cf-1611-4f4d-bcab-1bba5ee5498d",
   "metadata": {},
   "source": [
    "### 1.1. `TableTokenizer` 정의"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95006c93-06ae-4542-a4d0-3ef40841bbe2",
   "metadata": {},
   "source": [
    "우선, table 전용 토크나이저인 `TapasTokenizer`부터 불러오겠습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f767fd11-b0c8-48c8-a6a8-b2a63861bb13",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TapasTokenizer\n",
    "\n",
    "tapas_tokenizer = TapasTokenizer.from_pretrained(args.vocab_model_name)\n",
    "tapas_tokenizer.save_pretrained(args.tok_path)\n",
    "\n",
    "tokenizer = TapasTokenizer.from_pretrained(args.tok_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0ae13a98-9d79-43c0-aaf8-037e406f07a0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.vocab['[PAD]']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7d56003d-74d3-4fed-9e57-36d817079f06",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "32001"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(tokenizer) #vocab 파일에 명시되어있지는 않지만 special token으로는 포함되어 있다!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd926fa2-9a1c-4249-9083-749761cfb84c",
   "metadata": {},
   "source": [
    "\n",
    "우리는 이 `TapasTokenizer`를 그대로 쓰지 않고, Table-aware MLM 학습에 맞게 변형해서 활용할 것입니다.\n",
    "\n",
    "이를 위해, `TapasTokenizer` class를 상속받아 필요한 method를 직접 정의하는 작업이 필요합니다.\n",
    "\n",
    "아래처럼 말이죠!\n",
    "\n",
    "> **❣ 주의 ❣**\n",
    "> \n",
    "> _transformers 라이브러리의 버전을 꼭 확인하세요!! 버전 마다 구현체가 조금씩 다르기 때문에, 설치된 버전과 다른 버전에서 정의된 변수나 메소드를 불러올 시 에러가 발생하게 됩니다! 현재 설치되어 있는 버전에 맞는 소스코드를 참고하세요! 참고로 제가 사용한 버전은 `transformers==4.11.3` 입니다._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6e12b882-f7f3-4d6b-a46b-57a56846b62e",
   "metadata": {},
   "outputs": [],
   "source": [
    "## tokenizer_utils.py에 위치할 것\n",
    "\n",
    "def create_token_word_idx_list(word_list):\n",
    "    # token 별로 몇번째 word에 속해있는지 파악하기 위한 index 리스트 생성\n",
    "    n = 0\n",
    "    word_idx = []\n",
    "    for r in range(len(word_list)):\n",
    "        toks = tokenizer.tokenize(word_list[r])\n",
    "        idxs = []\n",
    "        for _ in range(len(toks)):\n",
    "            n += 1\n",
    "            idxs.append(n)\n",
    "        word_idx.append(idxs)\n",
    "\n",
    "    assert len(word_idx) == len(word_list)\n",
    "    \n",
    "    return word_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "759cd87f-7aee-46ac-a498-38841fa0c3be",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TableTokenizer(TapasTokenizer):\n",
    "    def __init__(\n",
    "        self,\n",
    "        **kwargs\n",
    "    ):\n",
    "        super().__init__(**kwargs)\n",
    "        \n",
    "        \n",
    "    def query_truncate(self, query, max_query_len):\n",
    "        \"\"\"\n",
    "        TapasTokenizer에는 max_query_len에 맞게 query를 truncate하는 기능이 없습니다.\n",
    "        query token의 길이가 max_query_len을 넘겼을 때나, min_query_len보다 짧을 때,\n",
    "        query를 인코딩하지 않고 전부 버립니다.(아래 코드처럼)\n",
    "            \n",
    "            def _get_question_tokens(self, query):\n",
    "                # Tokenizes the query, taking into account the max and min question length.\n",
    "                query_tokens = self.tokenize(query)\n",
    "                if self.max_question_length is not None and len(query_tokens) > self.max_question_length:\n",
    "                    logger.warning(\"Skipping query as its tokens are longer than the max question length\")\n",
    "                    return \"\", []\n",
    "                if self.min_question_length is not None and len(query_tokens) < self.min_question_length:\n",
    "                    logger.warning(\"Skipping query as its tokens are shorter than the min question length\")\n",
    "                    return \"\", []\n",
    "\n",
    "                return query, query_tokens\n",
    "        \n",
    "        왜 전부 버려야 하죠? 조금은 남길 수 있는 거잖아요?\n",
    "        그래서 query_truncate 메소드를 새로 정의해보았습니다.\n",
    "        구현 방식에서 더 좋은 아이디어가 있다면 톡~투~미~\n",
    "        \"\"\"\n",
    "        query_tokens = self.tokenize(query)\n",
    "        \n",
    "        if len(query_tokens) > max_query_len:\n",
    "            query_words = query.split()\n",
    "            query_word_idx = create_token_word_idx_list(query_words)\n",
    "            for i, query_tokens in enumerate(query_word_idx):\n",
    "                if max_query_len in query_tokens:\n",
    "                    end_word_idx = i\n",
    "            \n",
    "            query_txt = \" \".join(query_words[:end_word_idx])\n",
    "        \n",
    "            return query_txt\n",
    "            \n",
    "            ### table 길이도 고려해서 max query len 자동 지정하는 방안도 괜찮을거같아요~(나중에)\n",
    "        else:\n",
    "            return query\n",
    "        \n",
    "        \n",
    "    def get_idx_features_for_masking(self, table, query, max_length = 512):\n",
    "        # 1) table 각 cell의 좌표 정보를 비롯하여 query, table을 각각 인코딩한 결과를 받아옵니다.\n",
    "        table_data, query_ids, table_ids = self.get_coordinates(table, query, max_length)\n",
    "        \n",
    "        # 2) query와 table이 몇개의 토큰으로 구성되어 있는지 시퀀스 길이 정보를 받아옵니다.\n",
    "        len_query, len_table = self.get_length_before_pad(query_ids, table_ids)\n",
    "        len_info = [len_query, len_table]\n",
    "        \n",
    "        # 3) whole word masking(text)을 위해 단어 별 토큰의 위치 정보를 받아옵니다.\n",
    "        whole_word_idxs = self.get_whole_word_info(query_ids)\n",
    "        \n",
    "        # 4) whole cell masking(table)을 위해 cell 별 토큰의 위치 정보를 받아옵니다.\n",
    "        whole_cell_idxs = self.get_whole_cell_info(len_query, table_data)\n",
    "        \n",
    "        whole_word_idxs.extend(whole_cell_idxs)\n",
    "        # print(whole_word_idxs, len_info)\n",
    "        \n",
    "        return whole_word_idxs, len_info, query_ids, table_ids\n",
    "            \n",
    "    # 1) table 각 cell의 좌표 정보를 비롯하여 query, table을 각각 인코딩한 결과를 받아옵니다.\n",
    "    def get_coordinates(self, table, query, max_length = 512, truncation=\"drop_rows_to_fit\"):\n",
    "        ## truncation : DROP_ROWS_TO_FIT = \"drop_rows_to_fit\",  DO_NOT_TRUNCATE = \"do_not_truncate\"\n",
    "        \n",
    "        ## https://github.com/huggingface/transformers/blob/v4.11.3/src/transformers/models/tapas/tokenization_tapas.py#L1039\n",
    "        table_tokens = self._tokenize_table(table)\n",
    "        query_tokens = self.tokenize(query) # query, query_tokens = self._get_question_tokens(query) #(4.xx.x 버전)\n",
    "        \n",
    "        ## https://github.com/huggingface/transformers/blob/v4.11.3/src/transformers/models/tapas/tokenization_tapas.py#L1138\n",
    "        num_rows = self._get_num_rows(table, truncation != \"do_not_truncate\") #row idx는 0부터 시작\n",
    "        num_columns = self._get_num_columns(table) #col idx는 1부터 시작\n",
    "        _, _, num_tokens = self._get_table_boundaries(table_tokens)\n",
    "\n",
    "        if truncation != \"do_not_truncate\":\n",
    "            num_rows, num_tokens = self._get_truncated_table_rows(\n",
    "                query_tokens, table_tokens, num_rows, num_columns, max_length, truncation_strategy=truncation\n",
    "            )\n",
    "        \n",
    "        table_data = list(self._get_table_values(table_tokens, num_columns, num_rows, num_tokens))\n",
    "        # [TokenValue(token='층', column_id=1, row_id=0),\n",
    "        #  TokenValue(token='##수', column_id=1, row_id=0),\n",
    "        #  TokenValue(token='시설', column_id=2, row_id=0),\n",
    "        #  TokenValue(token='비', column_id=3, row_id=0),\n",
    "        #  TokenValue(token='##고', column_id=3, row_id=0),\n",
    "        # ...\n",
    "        \n",
    "        query_ids = self.convert_tokens_to_ids(query_tokens)\n",
    "        \n",
    "        table_ids = list(zip(*table_data))[0] if len(table_data) > 0 else list(zip(*table_data))\n",
    "        table_ids = self.convert_tokens_to_ids(list(table_ids))\n",
    "    \n",
    "        return table_data, query_ids, table_ids\n",
    "    \n",
    "    \n",
    "    # 2) query와 table이 몇개의 토큰으로 구성되어 있는지 시퀀스 길이 정보를 받아옵니다.\n",
    "    def get_length_before_pad(self, query_data, table_data):\n",
    "        return len(query_data), len(table_data)\n",
    "        \n",
    "        \n",
    "    # 3) whole word masking(text)을 위해 단어 별 토큰의 위치 정보를 받아옵니다.\n",
    "    def get_whole_word_info(self, query_ids):\n",
    "        '''whole word index information for text masking'''\n",
    "        ref_tokens = []\n",
    "        for n, token_idx in enumerate(query_ids):\n",
    "            token = tokenizer._convert_id_to_token(token_idx)\n",
    "            ref_tokens.append(token)\n",
    "            \n",
    "        cand_indexes = []\n",
    "        for (i, token) in enumerate(ref_tokens):\n",
    "            if len(cand_indexes) >= 1 and token.startswith(\"##\"):\n",
    "                cand_indexes[-1].append(i+1)\n",
    "            else:\n",
    "                cand_indexes.append([i+1])\n",
    "        \n",
    "        return cand_indexes\n",
    "    \n",
    "    \n",
    "    # 4) whole cell masking(table)을 위해 cell 별 토큰의 위치 정보를 받아옵니다.\n",
    "    def get_whole_cell_info(self, len_query, table_cell_info):\n",
    "        '''whole cell index information fir table masking'''\n",
    "        idx_order = len_query + 2\n",
    "        total_cell_cands, one_cell = [], []\n",
    "        start_row = 0\n",
    "        start_col = 1\n",
    "        \n",
    "        for tok_idx, tok_info in enumerate(table_cell_info):\n",
    "            if tok_idx == len(table_cell_info) - 1:\n",
    "                one_cell.append(idx_order+tok_idx)\n",
    "                total_cell_cands.append(one_cell)\n",
    "            else:\n",
    "                if tok_info.row_id == start_row:\n",
    "                    if tok_info.column_id == start_col:\n",
    "                        one_cell.append(idx_order+tok_idx)\n",
    "                    else:\n",
    "                        total_cell_cands.append(one_cell)\n",
    "                        one_cell = []\n",
    "                        start_col += 1\n",
    "                        one_cell.append(idx_order+tok_idx)\n",
    "                else:\n",
    "                    total_cell_cands.append(one_cell)\n",
    "                    one_cell = []\n",
    "                    start_row += 1\n",
    "                    start_col = 1\n",
    "                    one_cell.append(idx_order+tok_idx)\n",
    "        \n",
    "        return total_cell_cands\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9af7661-dff3-443d-8668-5d724223ac5f",
   "metadata": {},
   "source": [
    "새롭게 정의된 TableTokenizer가 어떤 식으로 작동하는지 아래 예시를 통해 살펴보겠습니다!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b3e212e1-54b9-4a46-80ab-d46041058dd0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. \n",
      "The tokenizer class you load from this checkpoint is 'TapasTokenizer'. \n",
      "The class this function is called from is 'TableTokenizer'.\n"
     ]
    }
   ],
   "source": [
    "tokenizer = TableTokenizer.from_pretrained(args.tok_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "21c9487b-534e-4f34-88dc-3d2126866896",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'N서울타워의 층수는 P0, P1, P2, EZ, T1, T2, T3, T5로 총 8개 층으로 되어있다. P는 플라자의 약자이며 출입구와 약간의 상가로 구성되어있다. EZ는 익스프레스 존의 약자이며 흰색 기둥부분을 가리킨다. T는 타워의 약자이며 전망대와 스낵코너, 그리고 식당으로 구성되어 있다. 지하 1층에서 지상 5층까지는 서울타워 플라자의 시설이 있고 5층부터 타워 1층에서 타워 5층까지는 N서울타워의 시설이 있다. 남산 케이블카도 유명하다.'"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ex_text = \"N서울타워의 층수는 P0, P1, P2, EZ, T1, T2, T3, T5로 총 8개 층으로 되어있다. P는 플라자의 약자이며 출입구와 약간의 상가로 구성되어있다. EZ는 익스프레스 존의 약자이며 흰색 기둥부분을 가리킨다. T는 타워의 약자이며 전망대와 스낵코너, 그리고 식당으로 구성되어 있다. 지하 1층에서 지상 5층까지는 서울타워 플라자의 시설이 있고 5층부터 타워 1층에서 타워 5층까지는 N서울타워의 시설이 있다. 남산 케이블카도 유명하다.\"\n",
    "ex_text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "0367412a-1e3a-45c4-914f-959e38f50149",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>층수</th>\n",
       "      <th>시설</th>\n",
       "      <th>비고</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>nan</td>\n",
       "      <td>엔그릴/기계실</td>\n",
       "      <td>양식당 '엔그릴'이며 이곳에서는 개성과 인천의 관측도 가능하다...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>타워 6층</td>\n",
       "      <td>N칼국수,전망대</td>\n",
       "      <td>... 휴전선까지 관측 가능하다</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>타워 5층</td>\n",
       "      <td>nan</td>\n",
       "      <td>디지털 전망대와 상행 엘리베이터가 있다.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>타워 4층</td>\n",
       "      <td>전망대, N포토 스튜디오, 하늘 화장실, 투썸커피</td>\n",
       "      <td>아날로그 전망대와 하행 엘리베이터가 있다.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>타워 3층</td>\n",
       "      <td>nan</td>\n",
       "      <td>이곳에서는 서울 시내까지만 보인다.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>타워 2층</td>\n",
       "      <td>루프테라스, 더 플레이스 다이닝</td>\n",
       "      <td>루프테라스, 더 플레이스 다이닝이 있다.</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "      층수                           시설                                     비고\n",
       "0    nan                      엔그릴/기계실  양식당 '엔그릴'이며 이곳에서는 개성과 인천의 관측도 가능하다...\n",
       "1  타워 6층                     N칼국수,전망대                      ... 휴전선까지 관측 가능하다\n",
       "2  타워 5층                          nan                 디지털 전망대와 상행 엘리베이터가 있다.\n",
       "3  타워 4층  전망대, N포토 스튜디오, 하늘 화장실, 투썸커피                아날로그 전망대와 하행 엘리베이터가 있다.\n",
       "4  타워 3층                          nan                    이곳에서는 서울 시내까지만 보인다.\n",
       "5  타워 2층            루프테라스, 더 플레이스 다이닝                 루프테라스, 더 플레이스 다이닝이 있다."
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ex_table = [[\"층수\", \"시설\", \"비고\"],\n",
    "[\"\", \"엔그릴/기계실\", \"양식당 '엔그릴'이며 이곳에서는 개성과 인천의 관측도 가능하다...\"],\n",
    "[\"타워 6층\", \"N칼국수,전망대\", \"... 휴전선까지 관측 가능하다\"],\n",
    "[\"타워 5층\", \"\", \"디지털 전망대와 상행 엘리베이터가 있다.\"],\n",
    "[\"타워 4층\", \"전망대, N포토 스튜디오, 하늘 화장실, 투썸커피\", \"아날로그 전망대와 하행 엘리베이터가 있다.\"],\n",
    "[\"타워 3층\", \"\", \"이곳에서는 서울 시내까지만 보인다.\"],\n",
    "[\"타워 2층\", \"루프테라스, 더 플레이스 다이닝\", \"루프테라스, 더 플레이스 다이닝이 있다.\"]]\n",
    "\n",
    "ex_tbl_df = pd.DataFrame(ex_table)\n",
    "ex_tbl_df = ex_tbl_df.rename(columns=ex_tbl_df.iloc[0])\n",
    "ex_tbl_df = ex_tbl_df.drop(ex_tbl_df.index[0])\n",
    "ex_tbl_df.reset_index(drop=True, inplace=True)\n",
    "ex_tbl_df = ex_tbl_df.astype('str')\n",
    "\n",
    "for row_index, row in ex_tbl_df.iterrows():\n",
    "    for col_index, cell in enumerate(row):\n",
    "        #print(row_index, col_index, ex_tbl_df.iloc[row_index, col_index])\n",
    "        if ex_tbl_df.iloc[row_index, col_index] == \"\":\n",
    "            ex_tbl_df.iloc[row_index, col_index] = \"nan\"\n",
    "        else:\n",
    "            continue\n",
    "\n",
    "ex_tbl_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "04d3f254-8fa1-4735-83ed-54fa0d226130",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': [2, 50, 28750, 2256, 2667, 2079, 1688, 2113, 2259, 52, 2082, 16, 52, 2083, 16, 52, 2302, 16, 41, 2611, 16, 56, 2083, 16, 56, 2302, 16, 56, 2195, 16, 56, 2049, 2200, 1668, 28, 2019, 1688, 6233, 859, 2051, 2689, 2062, 18, 52, 2259, 16597, 2079, 9383, 2052, 2307, 18455, 2522, 4943, 2079, 6682, 2200, 3896, 2496, 2051, 2689, 2062, 18, 41, 2611, 2259, 24284, 1554, 2079, 9383, 2052, 2307, 12003, 9856, 12547, 2069, 16519, 18, 56, 2259, 8203, 2079, 9383, 2052, 2307, 14822, 2522, 1, 16, 3673, 5499, 6233, 3896, 2496, 2051, 1513, 2062, 18, 4670, 21, 2624, 27135, 5377, 25, 2624, 2299, 2118, 2259, 3671, 2256, 2667, 16597, 2079, 3953, 2052, 1513, 2088, 25, 2624, 3797, 8203, 21, 2624, 27135, 8203, 25, 2624, 2299, 2118, 2259, 50, 28750, 2256, 2667, 2079, 3953, 2052, 1513, 2062, 18, 12103, 15879, 2119, 4455, 2205, 2062, 18, 3, 1688, 2113, 3953, 1187, 2088, 32000, 1423, 2029, 2388, 19, 5276, 2477, 6277, 2481, 11, 1423, 2029, 2388, 11, 1504, 2307, 4441, 27135, 2259, 5879, 2145, 4068, 2079, 6541, 2119, 3662, 2205, 2062, 18, 18, 18, 8203, 26, 2624, 50, 2600, 9473, 16, 14822, 18, 18, 18, 31099, 2299, 2118, 6541, 3662, 2205, 2062, 8203, 25, 2624, 32000, 5476, 14822, 2522, 1242, 2375, 10874, 2116, 1513, 2062, 18, 8203, 24, 2624, 14822, 16, 50, 2208, 2386, 9238, 16, 4573, 7047, 16, 1801, 3428, 20468, 16834, 14822, 2522, 1889, 2375, 10874, 2116, 1513, 2062, 18, 8203, 23, 2624, 32000, 4441, 27135, 2259, 3671, 6011, 2299, 3683, 4090, 18, 8203, 22, 2624, 19283, 2201, 5822, 16, 831, 18312, 5970, 2944, 19283, 2201, 5822, 16, 831, 18312, 5970, 2944, 2052, 1513, 2062, 18, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'token_type_ids': [[0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [1, 1, 0, 0, 0, 0, 0], [1, 1, 0, 0, 0, 0, 0], [1, 2, 0, 0, 0, 0, 0], [1, 3, 0, 0, 0, 0, 0], [1, 3, 0, 0, 0, 0, 0], [1, 1, 1, 0, 0, 0, 0], [1, 2, 1, 0, 0, 0, 0], [1, 2, 1, 0, 0, 0, 0], [1, 2, 1, 0, 0, 0, 0], [1, 2, 1, 0, 0, 0, 0], [1, 2, 1, 0, 0, 0, 0], [1, 2, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 3, 1, 0, 0, 0, 0], [1, 1, 2, 0, 5, 1, 6], [1, 1, 2, 0, 5, 1, 6], [1, 1, 2, 0, 5, 1, 6], [1, 2, 2, 0, 0, 0, 0], [1, 2, 2, 0, 0, 0, 0], [1, 2, 2, 0, 0, 0, 0], [1, 2, 2, 0, 0, 0, 0], [1, 2, 2, 0, 0, 0, 0], [1, 3, 2, 0, 0, 0, 0], [1, 3, 2, 0, 0, 0, 0], [1, 3, 2, 0, 0, 0, 0], [1, 3, 2, 0, 0, 0, 0], [1, 3, 2, 0, 0, 0, 0], [1, 3, 2, 0, 0, 0, 0], [1, 3, 2, 0, 0, 0, 0], [1, 3, 2, 0, 0, 0, 0], [1, 3, 2, 0, 0, 0, 0], [1, 3, 2, 0, 0, 0, 0], [1, 1, 3, 0, 4, 2, 7], [1, 1, 3, 0, 4, 2, 7], [1, 1, 3, 0, 4, 2, 7], [1, 2, 3, 0, 0, 0, 0], [1, 3, 3, 0, 0, 0, 0], [1, 3, 3, 0, 0, 0, 0], [1, 3, 3, 0, 0, 0, 0], [1, 3, 3, 0, 0, 0, 0], [1, 3, 3, 0, 0, 0, 0], [1, 3, 3, 0, 0, 0, 0], [1, 3, 3, 0, 0, 0, 0], [1, 3, 3, 0, 0, 0, 0], [1, 3, 3, 0, 0, 0, 0], [1, 3, 3, 0, 0, 0, 0], [1, 1, 4, 0, 3, 3, 6], [1, 1, 4, 0, 3, 3, 6], [1, 1, 4, 0, 3, 3, 6], [1, 2, 4, 0, 0, 0, 0], [1, 2, 4, 0, 0, 0, 0], [1, 2, 4, 0, 0, 0, 0], [1, 2, 4, 0, 0, 0, 0], [1, 2, 4, 0, 0, 0, 0], [1, 2, 4, 0, 0, 0, 0], [1, 2, 4, 0, 0, 0, 0], [1, 2, 4, 0, 0, 0, 0], [1, 2, 4, 0, 0, 0, 0], [1, 2, 4, 0, 0, 0, 0], [1, 2, 4, 0, 0, 0, 0], [1, 2, 4, 0, 0, 0, 0], [1, 2, 4, 0, 0, 0, 0], [1, 3, 4, 0, 0, 0, 0], [1, 3, 4, 0, 0, 0, 0], [1, 3, 4, 0, 0, 0, 0], [1, 3, 4, 0, 0, 0, 0], [1, 3, 4, 0, 0, 0, 0], [1, 3, 4, 0, 0, 0, 0], [1, 3, 4, 0, 0, 0, 0], [1, 3, 4, 0, 0, 0, 0], [1, 3, 4, 0, 0, 0, 0], [1, 3, 4, 0, 0, 0, 0], [1, 1, 5, 0, 2, 4, 7], [1, 1, 5, 0, 2, 4, 7], [1, 1, 5, 0, 2, 4, 7], [1, 2, 5, 0, 0, 0, 0], [1, 3, 5, 0, 0, 0, 0], [1, 3, 5, 0, 0, 0, 0], [1, 3, 5, 0, 0, 0, 0], [1, 3, 5, 0, 0, 0, 0], [1, 3, 5, 0, 0, 0, 0], [1, 3, 5, 0, 0, 0, 0], [1, 3, 5, 0, 0, 0, 0], [1, 3, 5, 0, 0, 0, 0], [1, 3, 5, 0, 0, 0, 0], [1, 1, 6, 0, 1, 5, 7], [1, 1, 6, 0, 1, 5, 7], [1, 1, 6, 0, 1, 5, 7], [1, 2, 6, 0, 0, 0, 0], [1, 2, 6, 0, 0, 0, 0], [1, 2, 6, 0, 0, 0, 0], [1, 2, 6, 0, 0, 0, 0], [1, 2, 6, 0, 0, 0, 0], [1, 2, 6, 0, 0, 0, 0], [1, 2, 6, 0, 0, 0, 0], [1, 2, 6, 0, 0, 0, 0], [1, 3, 6, 0, 0, 0, 0], [1, 3, 6, 0, 0, 0, 0], [1, 3, 6, 0, 0, 0, 0], [1, 3, 6, 0, 0, 0, 0], [1, 3, 6, 0, 0, 0, 0], [1, 3, 6, 0, 0, 0, 0], [1, 3, 6, 0, 0, 0, 0], [1, 3, 6, 0, 0, 0, 0], [1, 3, 6, 0, 0, 0, 0], [1, 3, 6, 0, 0, 0, 0], [1, 3, 6, 0, 0, 0, 0], [1, 3, 6, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0]]}"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ex_toks = tokenizer(table=ex_tbl_df, queries=ex_text, max_length=300, padding=True)\n",
    "ex_toks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd5916c6-3b0f-43ed-a3be-79a768eb8570",
   "metadata": {},
   "source": [
    "- `attention_mask`: table&text 토큰=1, padding=0\n",
    "\n",
    "- `token_type_ids`: TapasTokenizer에는 테이블 구조를 반영하기 위한 7개의 token type id가 추가되어 있습니다. 각 id에 대한 설명은 아래 Tapas 공식 문서를 발췌한 부분을 참고해주세요!\n",
    "\n",
    "    1. `segment_ids`: indicate **whether a token belongs to the question (0) or the table (1)**. 0 for special tokens and padding.\n",
    "    2. `column_ids`: indicate to **which column of the table** a token belongs (starting from 1). Is 0 for all question tokens, special tokens and padding.\n",
    "    3. `row_ids`: indicate to **which row of the table** a token belongs (starting from 1). Is 0 for all question tokens, special tokens and padding. Tokens of column headers are also 0.\n",
    "    4. `prev_labels`: indicate **whether a token was (part of) an answer to the previous question** (1) or not (0). Useful in a conversational setup (such as SQA dataset).\n",
    "    5. `column_ranks`: indicate the **rank of a table token relative to a column**, **if applicable**. For example, if you have a column \"number of movies\" with values 87, 53 and 69, then the column ranks of these tokens are 3, 1 and 2 respectively. 0 for all question tokens, special tokens and padding.\n",
    "    6. `inv_column_ranks`: indicate the **inverse rank** of a table token relative to a column, **if applicable**. For example, if you have a column \"number of movies\" with values 87, 53 and 69, then the inverse column ranks of these tokens are 1, 3 and 2 respectively. 0 for all question tokens, special tokens and padding.\n",
    "    7. `numeric_relations`: indicate numeric relations between the question and the tokens of the table. 0 for all question tokens, special tokens and padding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "8f10ba2c-ab9a-4c13-947b-f46b919a9b8e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"[CLS] N서울타워의 층수는 P0, P1, P2, EZ, T1, T2, T3, T5로 총 8개 층으로 되어있다. P는 플라자의 약자이며 출입구와 약간의 상가로 구성되어있다. EZ는 익스프레스 존의 약자이며 흰색 기둥부분을 가리킨다. T는 타워의 약자이며 전망대와 [UNK], 그리고 식당으로 구성되어 있다. 지하 1층에서 지상 5층까지는 서울타워 플라자의 시설이 있고 5층부터 타워 1층에서 타워 5층까지는 N서울타워의 시설이 있다. 남산 케이블카도 유명하다. [SEP] 층수 시설 비고 [EMPTY] 엔그릴 / 기계실 양식당'엔그릴'이며 이곳에서는 개성과 인천의 관측도 가능하다... 타워 6층 N칼국수, 전망대... 휴전선까지 관측 가능하다 타워 5층 [EMPTY] 디지털 전망대와 상행 엘리베이터가 있다. 타워 4층 전망대, N포토 스튜디오, 하늘 화장실, 투썸커피 아날로그 전망대와 하행 엘리베이터가 있다. 타워 3층 [EMPTY] 이곳에서는 서울 시내까지만 보인다. 타워 2층 루프테라스, 더 플레이스 다이닝 루프테라스, 더 플레이스 다이닝이 있다. [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD]\""
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.decode(ex_toks['input_ids'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "6827fe8d-489e-4334-96ac-031d133e206f",
   "metadata": {},
   "outputs": [],
   "source": [
    "idx_info, len_info, query_ids, table_ids = tokenizer.get_idx_features_for_masking(table=ex_tbl_df, query=ex_text, max_length=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "2daf7b3e-80da-4682-82d5-abca71fe4389",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[1, 2, 3, 4, 5], [6, 7, 8], [9, 10], [11], [12, 13], [14], [15, 16], [17], [18, 19], [20], [21, 22], [23], [24, 25], [26], [27, 28], [29], [30, 31, 32], [33], [34, 35], [36, 37], [38, 39, 40, 41], [42], [43, 44], [45, 46], [47, 48, 49], [50, 51], [52, 53], [54, 55], [56, 57, 58, 59, 60], [61], [62, 63, 64], [65], [66, 67], [68, 69, 70], [71], [72, 73, 74], [75], [76], [77, 78], [79, 80], [81, 82, 83], [84, 85], [86], [87], [88], [89, 90], [91, 92, 93], [94, 95], [96], [97], [98, 99, 100], [101], [102, 103, 104, 105, 106], [107, 108, 109], [110, 111], [112, 113], [114, 115], [116, 117, 118], [119], [120, 121, 122], [123], [124, 125, 126, 127, 128], [129, 130, 131, 132, 133], [134, 135], [136, 137], [138], [139], [140, 141], [142, 143, 144], [145], [147, 148], [149], [150, 151], [152], [153, 154, 155, 156, 157, 158], [159, 160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181, 182], [183, 184, 185], [186, 187, 188, 189, 190], [191, 192, 193, 194, 195, 196, 197, 198, 199, 200], [201, 202, 203], [204], [205, 206, 207, 208, 209, 210, 211, 212, 213, 214], [215, 216, 217], [218, 219, 220, 221, 222, 223, 224, 225, 226, 227, 228, 229, 230], [231, 232, 233, 234, 235, 236, 237, 238, 239, 240], [241, 242, 243], [244], [245, 246, 247, 248, 249, 250, 251, 252, 253], [254, 255, 256], [257, 258, 259, 260, 261, 262, 263, 264], [265, 266, 267, 268, 269, 270, 271, 272, 273, 274, 275, 276]]\n"
     ]
    }
   ],
   "source": [
    "print(idx_info)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85798b1c-c705-466d-9fb9-a6b8295910f2",
   "metadata": {},
   "source": [
    "각 하위 리스트들은 단어, cell에 속한 토큰들의 위치 정보를 포함합니다.\n",
    "\n",
    "예를 들어, query(text)의 첫번째 단어(`[1, 2, 3, 4, 5]`)는 첫번째 토큰부터 다섯번째 토큰으로 구성되어 있습니다.\n",
    "\n",
    "참고로 `[CLS]` 토큰에 해당하는 0번 인덱스는 이 리스트에 포함되어 있지 않으며, 마찬가지로 text와 table을 분리하는 `[SEP]` 토큰 역시 제외된 상태입니다.\n",
    "\n",
    "중간에 토큰의 위치 인덱스가 비어있는 지점(위 예시에서는 146번 인덱스)가 text에서 table로 넘어가는 `[SEP]` 토큰의 위치라고 보시면 됩니다.\n",
    "\n",
    "<br/>\n",
    "\n",
    "또한, 아래와 같이 query(text)와 table의 길이 정보, 인코딩 결과가 함께 제공됩니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "8264ef4c-1951-4ac4-85f7-dab38b86cf5c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[145, 130]\n"
     ]
    }
   ],
   "source": [
    "print(len_info)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "89a63daa-8f53-430c-a4bf-c401b7990a1f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[50, 28750, 2256, 2667, 2079, 1688, 2113, 2259, 52, 2082, 16, 52, 2083, 16, 52, 2302, 16, 41, 2611, 16, 56, 2083, 16, 56, 2302, 16, 56, 2195, 16, 56, 2049, 2200, 1668, 28, 2019, 1688, 6233, 859, 2051, 2689, 2062, 18, 52, 2259, 16597, 2079, 9383, 2052, 2307, 18455, 2522, 4943, 2079, 6682, 2200, 3896, 2496, 2051, 2689, 2062, 18, 41, 2611, 2259, 24284, 1554, 2079, 9383, 2052, 2307, 12003, 9856, 12547, 2069, 16519, 18, 56, 2259, 8203, 2079, 9383, 2052, 2307, 14822, 2522, 1, 16, 3673, 5499, 6233, 3896, 2496, 2051, 1513, 2062, 18, 4670, 21, 2624, 27135, 5377, 25, 2624, 2299, 2118, 2259, 3671, 2256, 2667, 16597, 2079, 3953, 2052, 1513, 2088, 25, 2624, 3797, 8203, 21, 2624, 27135, 8203, 25, 2624, 2299, 2118, 2259, 50, 28750, 2256, 2667, 2079, 3953, 2052, 1513, 2062, 18, 12103, 15879, 2119, 4455, 2205, 2062, 18]\n"
     ]
    }
   ],
   "source": [
    "print(query_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "0068b293-01b7-49a1-b583-fecc72ab27e4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1688, 2113, 3953, 1187, 2088, 32000, 1423, 2029, 2388, 19, 5276, 2477, 6277, 2481, 11, 1423, 2029, 2388, 11, 1504, 2307, 4441, 27135, 2259, 5879, 2145, 4068, 2079, 6541, 2119, 3662, 2205, 2062, 18, 18, 18, 8203, 26, 2624, 50, 2600, 9473, 16, 14822, 18, 18, 18, 31099, 2299, 2118, 6541, 3662, 2205, 2062, 8203, 25, 2624, 32000, 5476, 14822, 2522, 1242, 2375, 10874, 2116, 1513, 2062, 18, 8203, 24, 2624, 14822, 16, 50, 2208, 2386, 9238, 16, 4573, 7047, 16, 1801, 3428, 20468, 16834, 14822, 2522, 1889, 2375, 10874, 2116, 1513, 2062, 18, 8203, 23, 2624, 32000, 4441, 27135, 2259, 3671, 6011, 2299, 3683, 4090, 18, 8203, 22, 2624, 19283, 2201, 5822, 16, 831, 18312, 5970, 2944, 19283, 2201, 5822, 16, 831, 18312, 5970, 2944, 2052, 1513, 2062, 18]\n"
     ]
    }
   ],
   "source": [
    "print(table_ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c192164d-d74a-4f54-9eee-3196b93e7cdb",
   "metadata": {},
   "source": [
    "### 1.2. `TableDataset` 정의"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a708718-215e-4947-9aa9-40a9db40be7c",
   "metadata": {},
   "source": [
    "위에서 커스터마이징한 Table Tokenizer를 활용하여 학습에 필요한 TableDataset을 정의하겠습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "099e822e-f74a-47ba-83a1-462326379bf5",
   "metadata": {},
   "outputs": [],
   "source": [
    "## dataset_utils.py에 위치할 것\n",
    "\n",
    "def none_to_nan(name):\n",
    "    return \"nan\" if name == None else name\n",
    "\n",
    "def del_row_or_col(df, row_del_ratio, col_del_ratio):\n",
    "    \"\"\"nan 개수가 일정 비율 이상이면 해당 column/row 전부 삭제하기 위한 함수입니다.\"\"\"\n",
    "    \n",
    "    n_r = df.shape[0]\n",
    "    n_c = df.shape[1]\n",
    "    \n",
    "    c_cand_del = []\n",
    "    for i_c in range(n_c):\n",
    "        n_col_nan = df.iloc[:,i_c].isnull().sum()\n",
    "        if n_col_nan > 0:\n",
    "            col_nan_ratio = n_col_nan/n_r\n",
    "            if col_nan_ratio > col_del_ratio:\n",
    "                # print(f\"col {i_c}({df.columns[i_c]}) insnull ratio: {col_nan_ratio}\")\n",
    "                c_cand_del.append(df.columns[i_c])\n",
    "    \n",
    "    r_cand_del = []\n",
    "    for i_r in range(n_r):\n",
    "        n_row_nan = df.iloc[i_r,:].isnull().sum()\n",
    "        if n_row_nan > 0:\n",
    "            row_nan_ratio = n_row_nan/n_c\n",
    "            if row_nan_ratio > row_del_ratio:\n",
    "                # print(f\"row {i_r}({df.index[i_r]}) insnull ratio: {row_nan_ratio}\")\n",
    "                r_cand_del.append(df.index[i_r])    \n",
    "    \n",
    "    # del nan col\n",
    "    df = df.drop(c_cand_del, axis=1)\n",
    "    # del nan row\n",
    "    df = df.drop(r_cand_del)\n",
    "    \n",
    "    # 다 지운 담에는 np.nan -> \"nan\"으로 변경해서 [EMPTY] 토큰으로 인식되도록!\n",
    "    df = df.fillna('nan')\n",
    "    \n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "f9352580-dbe9-42f0-a4bd-44ed27060ecb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "import pandas as pd\n",
    "import random\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "from transformers import TapasTokenizer\n",
    "# from table_tokenizer import TableTokenizer\n",
    "\n",
    "\n",
    "class TableDataset(Dataset):\n",
    "    def __init__(self, data, args=None, tokenizer=None):\n",
    "        self.data = data\n",
    "        self.args = args\n",
    "        self.tokenizer = tokenizer\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "    \n",
    "    def __getitem__(self, data_idx):\n",
    "        data_dict = self.data[data_idx]\n",
    "        text = data_dict['Description']\n",
    "        table = self.convert_list_to_df(data_dict['TBL'], self.args.row_del_ratio, self.args.col_del_ratio)\n",
    "        \n",
    "        real_text = self.tokenizer.query_truncate(text, self.args.max_query_len) ###max_query_len\n",
    "        \n",
    "        idx_info, len_info, _, _ = self.tokenizer.get_idx_features_for_masking(table=table, query=real_text, max_length=self.args.max_seq_len)\n",
    "        \n",
    "        inputs = self.tokenizer(table=table, queries=real_text, padding=True, max_length = self.args.max_seq_len, truncation=True)\n",
    "        # idx_info, len_info, _, _ = self.tokenizer.get_idx_features_for_masking(table=table, query=text)\n",
    "        # ## [[1, 2, 3, 4, 5], [6, 7, 8], [9, 10], [11], [12, 13], [14], [15, 16], ...]\n",
    "        \n",
    "        idx_info_copy = copy.deepcopy(idx_info)\n",
    "        \n",
    "        input_ids, labels = self.create_features_for_mlm(inputs['input_ids'], idx_info, len_info)\n",
    "        \n",
    "        return [input_ids, inputs['attention_mask'], inputs['token_type_ids'], labels, idx_info_copy]\n",
    "    \n",
    "    \n",
    "    def convert_list_to_df(self, table_list, row_del_ratio, col_del_ratio):\n",
    "        \"\"\"\n",
    "        이 메소드는 list로 주어진 table input을 받아 pandas dataframe 형식으로 변환 및 간단한 전처리를 수행합니다.\n",
    "        구체적인 프로세스는 주석을 참고하세요!\n",
    "        \"\"\"\n",
    "        # 1. list -> dataframe\n",
    "        tbl_df = pd.DataFrame(table_list)\n",
    "        # 2. 첫번째 행은 column 이름!\n",
    "        tbl_df = tbl_df.rename(columns=tbl_df.iloc[0])\n",
    "        # 3. column 이름으로 올려놨으니 첫번째 행은 지우기!\n",
    "        tbl_df = tbl_df.drop(tbl_df.index[0])\n",
    "        # 4. row index 초기화!\n",
    "        tbl_df.reset_index(drop=True, inplace=True)\n",
    "        # 5. 데이터 타입 문자열로 변환해야 None -> \"None\"으로 바뀌고, tokenizing 가능!\n",
    "        tbl_df = tbl_df.astype('str')\n",
    "\n",
    "        # 6. row, column 별로 돌면서 비어있는 값은 모두 결측치(np.nan)로 변환 -> 결측치 많은 행, 열은 제거할 것이기 때문!\n",
    "        for row_index, row in tbl_df.iterrows():\n",
    "            for col_index, cell in enumerate(row):\n",
    "                if tbl_df.iloc[row_index, col_index].strip() == \"\" or tbl_df.iloc[row_index, col_index] == \"None\":\n",
    "                    tbl_df.iloc[row_index, col_index] = np.nan ###\n",
    "                else:\n",
    "                    continue\n",
    "\n",
    "        # 7. 결측치가 너무 많은 row와 column은 제거!\n",
    "        tbl_df = del_row_or_col(tbl_df, row_del_ratio, col_del_ratio)\n",
    "\n",
    "        # 8. column 이름이 None인 경우 Nonetype error 방지를 위해 nan으로 rename!\n",
    "        tbl_df.columns = [none_to_nan(col) for col in list(tbl_df.columns)]\n",
    "\n",
    "        return tbl_df\n",
    "    \n",
    "    \n",
    "    def create_features_for_mlm(self, org_input_ids, idx_info, len_info):\n",
    "        \n",
    "        # masking ref: https://github.com/huggingface/transformers/blob/v4.11.3/src/transformers/data/data_collator.py#L924\n",
    "        \n",
    "        len_query_table = sum(len_info)\n",
    "    \n",
    "        num_to_predict = min(self.args.max_seq_len, max(1, int(round(len_query_table * self.args.mlm_prob))))\n",
    "\n",
    "        random.shuffle(idx_info) #랜덤으로 섞고 앞에서부터 차례로 masking할 토큰수 만큼 채움\n",
    "        ## [[1, 2, 3, 4, 5], [6, 7, 8], [9, 10], [11], [12, 13], [14], [15, 16], [17], ...] -> [[94, 95], [11], [62, 63, 64], ...]\n",
    "\n",
    "        masked_lms = []\n",
    "        covered_indexes = set()\n",
    "        for index_set in idx_info:\n",
    "            if len(masked_lms) >= num_to_predict:\n",
    "                break\n",
    "            # If adding a whole-word mask would exceed the maximum number of predictions, then just skip this candidate.\n",
    "            if len(masked_lms) + len(index_set) > num_to_predict:\n",
    "                continue\n",
    "            is_any_index_covered = False\n",
    "            for index in index_set:\n",
    "                if index in covered_indexes:\n",
    "                    is_any_index_covered = True\n",
    "                    break\n",
    "            if is_any_index_covered:\n",
    "                continue\n",
    "            for index in index_set:\n",
    "                covered_indexes.add(index)\n",
    "                masked_lms.append(index)\n",
    "\n",
    "        assert len(covered_indexes) == len(masked_lms)\n",
    "        mask_labels = [1 if i in covered_indexes else 0 for i in range(len_query_table + 2)] #+2: [CLS]와 [SEP] 개수도 포함!!\n",
    "\n",
    "        for _ in range(self.args.max_seq_len - len(mask_labels)):\n",
    "            mask_labels.append(self.tokenizer.pad_token_id) # 0\n",
    "\n",
    "\n",
    "        # 1. labels\n",
    "        inputs = torch.Tensor(org_input_ids)\n",
    "        labels = inputs.clone()\n",
    "\n",
    "        masked_indices = torch.Tensor(mask_labels).bool()\n",
    "\n",
    "        labels[~masked_indices] = -100 #masking된 토큰을 제외하고서는 loss 연산에서 제외\n",
    "        \n",
    "        real_labels = labels.tolist()\n",
    "        \n",
    "        real_labels = [int(lab) for lab in real_labels]\n",
    "\n",
    "        # 2. input_ids\n",
    "        input_ids = inputs.tolist()\n",
    "        for mask_pos in masked_lms:\n",
    "            input_ids[mask_pos] = self.tokenizer.mask_token_id # 4\n",
    "            \n",
    "        input_ids = [int(inp) for inp in input_ids]\n",
    "\n",
    "        return input_ids, real_labels\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fa1d6ae-398c-4265-8dc8-40e052b341b8",
   "metadata": {},
   "source": [
    "### 1.3. `table collate 함수` 및 `DatoLoader` 정의"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "43726a5b-c6ad-4a9c-b9fb-c0cb9a5b276d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "def table_collate_fn(batch):\n",
    "    \n",
    "    features = {\n",
    "        'input_ids': torch.LongTensor([sample[0] for sample in batch]),\n",
    "        'attention_mask': torch.LongTensor([sample[1] for sample in batch]),\n",
    "        'token_type_ids': torch.LongTensor([sample[2] for sample in batch]),\n",
    "        'labels': torch.LongTensor([sample[3] for sample in batch]),\n",
    "        'offsets': [sample[4] for sample in batch]\n",
    "    }\n",
    "    \n",
    "    return features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "c541fe71-e88b-40f1-a33b-bb31b5676f5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "\n",
    "train_data = sample_data['data'][:170]\n",
    "valid_data = sample_data['data'][170:]\n",
    "\n",
    "train_dataset = TableDataset(train_data, args, tokenizer)\n",
    "valid_dataset = TableDataset(valid_data, args, tokenizer)\n",
    "\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, collate_fn=table_collate_fn)\n",
    "valid_dataloader = DataLoader(valid_dataset, batch_size=args.batch_size, collate_fn=table_collate_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a43ded29-4506-4e29-a19d-188d26eb8a92",
   "metadata": {},
   "source": [
    "아래 cell은 batch 예시를 출력해보기 위해 선언한 가상의 dataloader입니다.\n",
    "\n",
    "실제 학습 때는 실행하지 않으셔도 됩니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "8857997e-f070-4a8c-8148-f85d74bfa3a9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'input_ids': tensor([[   2,  170, 1376,  ...,    0,    0,    0],\n",
      "        [   2,   14,  809,  ...,    0,    0,    0]]), 'attention_mask': tensor([[1, 1, 1,  ..., 0, 0, 0],\n",
      "        [1, 1, 1,  ..., 0, 0, 0]]), 'token_type_ids': tensor([[[0, 0, 0,  ..., 0, 0, 0],\n",
      "         [0, 0, 0,  ..., 0, 0, 0],\n",
      "         [0, 0, 0,  ..., 0, 0, 0],\n",
      "         ...,\n",
      "         [0, 0, 0,  ..., 0, 0, 0],\n",
      "         [0, 0, 0,  ..., 0, 0, 0],\n",
      "         [0, 0, 0,  ..., 0, 0, 0]],\n",
      "\n",
      "        [[0, 0, 0,  ..., 0, 0, 0],\n",
      "         [0, 0, 0,  ..., 0, 0, 0],\n",
      "         [0, 0, 0,  ..., 0, 0, 0],\n",
      "         ...,\n",
      "         [0, 0, 0,  ..., 0, 0, 0],\n",
      "         [0, 0, 0,  ..., 0, 0, 0],\n",
      "         [0, 0, 0,  ..., 0, 0, 0]]]), 'labels': tensor([[-100, -100, -100,  ..., -100, -100, -100],\n",
      "        [-100, -100, -100,  ..., -100, -100, -100]]), 'offsets': [[[1], [2, 3], [4], [5], [6], [7], [8], [9, 10, 11], [12, 13, 14], [15], [16], [17, 18], [19, 20], [21], [22, 23], [24, 25], [26], [27, 28], [29], [31, 32, 33, 34, 35, 36, 37, 38, 39, 40], [41, 42, 43, 44, 45, 46, 47, 48, 49, 50], [51], [52, 53, 54, 55], [56], [57, 58, 59, 60, 61, 62, 63, 64], [65], [66, 67, 68, 69, 70, 71, 72, 73], [74], [75, 76, 77, 78, 79, 80, 81, 82], [83], [84, 85, 86, 87, 88, 89, 90, 91, 92], [93, 94], [95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112], [113], [114, 115], [116], [117, 118], [119, 120]], [[1], [2, 3, 4], [5, 6], [7, 8], [9], [10, 11, 12, 13], [14, 15], [16, 17], [18, 19], [20], [21], [23], [24], [25], [26], [27, 28], [29], [30, 31, 32, 33, 34, 35], [36], [37], [38, 39, 40]]]}\n"
     ]
    }
   ],
   "source": [
    "ex_dataloader = DataLoader(train_dataset, batch_size=2, collate_fn=table_collate_fn)\n",
    "\n",
    "batch_ex = next(iter(ex_dataloader))\n",
    "print(batch_ex)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "469ae66e-6339-4d3d-8e78-f5e9dc7cd5ea",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1. input ids: \n",
      "tensor([    2,   170,  1376,  2320,  8347,   171,    12,     4,    30, 26384,\n",
      "         2054,  4704, 28495, 15869,  2041,    13,   793, 17591,  2615,  9296,\n",
      "         2170,     4,     4,     4,  6837,  2897,  9296,  3771, 28674,    18,\n",
      "            3,  1376,  2320,  8347,  2106,  2008,  7088,  4704, 28495, 15869,\n",
      "         2041,  1376,  2320,  8347,  2106,  2008,  7088,  4704, 28495, 15869,\n",
      "         2041,  3871,     4,     4,     4,     4,  4734, 10073,  1050,  2506,\n",
      "        16555,  2548,  4358, 22330, 17194,  4152, 14113, 13429,  2029,  2234,\n",
      "         3193,  2059,  2269,  2255, 12300,     4,     4,     4,     4,     4,\n",
      "            4,     4,     4, 28157, 27357,  1883,  2294,  2255, 11398,  1883,\n",
      "         2294,  2255,  5551,  6837,  2210,  4840,  2440,  3718,  2429,  3718,\n",
      "         2210,  4840,  2440,  3718,  2429,  4136,  2210,  4840,  2440,  3718,\n",
      "         2429,  3912,  2210,  3641,  7815,  2377,  3728,     4,     4,  4686,\n",
      "         4612,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "            0,     0])\n",
      "-----------------------------------------\n",
      "2. labels: \n",
      "tensor([ -100,  -100,  -100,  -100,  -100,  -100,  -100,  4612,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  4249,  4840,  2440,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100, 12740,  1739, 22661, 14570,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100, 17591,  2615,  9296, 24569,  1883,\n",
      "         2294,  2255,  9296,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  4353,  3666,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100])\n",
      "-----------------------------------------\n"
     ]
    }
   ],
   "source": [
    "print(f\"1. input ids: \\n{batch_ex['input_ids'][0]}\")\n",
    "print(\"-----------------------------------------\")\n",
    "print(f\"2. labels: \\n{batch_ex['labels'][0]}\")\n",
    "print(\"-----------------------------------------\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d88975e-7b7a-4180-84d6-06e3c54cef7d",
   "metadata": {},
   "source": [
    "`input ids`에서 masking된 token은 masking 인덱스인 4를 부여받고, `labels`에서는 maksing된 token 외 모든 토큰들이 전부 -100의 인덱스를 부여받은 것을 확인할 수 있습니다.\n",
    "\n",
    "이제 MLM 학습을 위한 데이터 구성은 모두 마쳤습니다!\n",
    "\n",
    "모델링 파트로 넘어가보아요~!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3eacec47-7b45-40bc-81be-0cba324a6cb8",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b73538c2-39a1-4d91-a371-8a31ed7d20aa",
   "metadata": {},
   "source": [
    "## 2. 모델링"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc003a38-3f9a-4811-86eb-01b2f502effa",
   "metadata": {},
   "source": [
    "Tapas는 embedding을 제외하고서 BERT와 동일한 구조로 구성되어 있으며, 코드 역시 BERT 구현체와 상당히 유사합니다.\n",
    "\n",
    "따라서, huggingface의 `BertModel` 계열 소스코드를을 접해보신 분들이라면 익숙하실 겁니다.\n",
    "\n",
    "Tapas에 MLM을 추가한 구조는 `TapasForMaskedLM`을 통해 손쉽게 불러올 수 있습니다.\n",
    "\n",
    "그 전에 모델의 정보를 담은 `TapasConfig`를 호출하겠습니다."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6e63dda-4c4c-4420-8de1-c8e26c153d99",
   "metadata": {},
   "source": [
    "### 2.1. `TapasConfig` 수정"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6dc14135-797d-4982-8c28-d750144ef355",
   "metadata": {},
   "source": [
    "> _\"그냥 모델 가져다 쓰면 된다면서요. 왜 config를 수정하는거죠?\"_\n",
    "\n",
    "TapasConfig는 기준 언어가 영어로 맞춰져 있어서 모델 초반 Embedding Layer를 관할하는 `vocab_size`가 `bert-base-uncased`의 vocab 개수인 30522로 설정되어 있습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "7fd8d900-4c42-464f-ac93-ae28d204fb52",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TapasConfig\n",
    "config = TapasConfig()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "824b5e49-0e8f-4d43-bd76-996664a78df3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "30522"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "config.vocab_size"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "349f08f3-50d0-451a-87c3-18a137ec4795",
   "metadata": {},
   "source": [
    "우리는 한국어 table mrc를 위한 tapas 모델을 구축해야 하기 때문에, 일전에 `TableTokenizer`를 로드할 때 사용했던 `klue/bert-base`에 맞는 vocab_size를 새로 정의해주어야 합니다.\n",
    "\n",
    "`TapasTokenizer`에는 special token인 `[EMPTY]`가 추가로 부여되어 있는 관계로, 기존 `klue/bert-base`의 vocab size(32000)에서 1을 추가한 32001로 수정하겠습니다!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "d5a5cae4-a9e4-4206-885f-da2c1265535c",
   "metadata": {},
   "outputs": [],
   "source": [
    "config.vocab_size = 32001"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "363bc474-1fae-452a-8304-183d3fbc6523",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TapasConfig {\n",
       "  \"aggregation_labels\": null,\n",
       "  \"aggregation_loss_weight\": 1.0,\n",
       "  \"aggregation_temperature\": 1.0,\n",
       "  \"allow_empty_column_selection\": false,\n",
       "  \"answer_loss_cutoff\": null,\n",
       "  \"answer_loss_importance\": 1.0,\n",
       "  \"attention_probs_dropout_prob\": 0.1,\n",
       "  \"average_approximation_function\": \"ratio\",\n",
       "  \"average_logits_per_cell\": false,\n",
       "  \"cell_selection_preference\": null,\n",
       "  \"disable_per_token_loss\": false,\n",
       "  \"hidden_act\": \"gelu\",\n",
       "  \"hidden_dropout_prob\": 0.1,\n",
       "  \"hidden_size\": 768,\n",
       "  \"huber_loss_delta\": null,\n",
       "  \"init_cell_selection_weights_to_zero\": false,\n",
       "  \"initializer_range\": 0.02,\n",
       "  \"intermediate_size\": 3072,\n",
       "  \"layer_norm_eps\": 1e-12,\n",
       "  \"max_num_columns\": 32,\n",
       "  \"max_num_rows\": 64,\n",
       "  \"max_position_embeddings\": 1024,\n",
       "  \"model_type\": \"tapas\",\n",
       "  \"no_aggregation_label_index\": null,\n",
       "  \"num_aggregation_labels\": 0,\n",
       "  \"num_attention_heads\": 12,\n",
       "  \"num_hidden_layers\": 12,\n",
       "  \"pad_token_id\": 0,\n",
       "  \"positive_label_weight\": 10.0,\n",
       "  \"reset_position_index_per_cell\": true,\n",
       "  \"select_one_column\": true,\n",
       "  \"transformers_version\": \"4.11.3\",\n",
       "  \"type_vocab_sizes\": [\n",
       "    3,\n",
       "    256,\n",
       "    256,\n",
       "    2,\n",
       "    256,\n",
       "    256,\n",
       "    10\n",
       "  ],\n",
       "  \"use_answer_as_supervision\": null,\n",
       "  \"use_gumbel_for_aggregation\": false,\n",
       "  \"use_gumbel_for_cells\": false,\n",
       "  \"use_normalized_answer_loss\": false,\n",
       "  \"vocab_size\": 32001\n",
       "}"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "config"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "272c698c-b599-4749-855a-077d8432bef8",
   "metadata": {},
   "source": [
    "맨 아래 `vocab_size`를 보니 잘 수정이 되었네요.\n",
    "\n",
    "- 바로 아래 `2.2. TapasForMaskedLM 소환`에서 model 구조를 출력한 셀을 보면 Word Embedding layer의 input 차원 역시 32001로 바뀐 것을 확인할 수 있습니다!\n",
    "\n",
    "이제 모델을 불러옵시다."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7616880f-87ae-4567-bbd0-bb5dc8969b29",
   "metadata": {},
   "source": [
    "### 2.2. `TapasForMaskedLM` 소환"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05532f39-c03b-4ff4-9a77-ff9c4bd27be3",
   "metadata": {},
   "source": [
    "위에서 수정한 Tapas config를 입력 변수로 넣어 미리 잘 구축된 Tapas MLM 모델을 불러오겠습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "b4bd9992-7b1a-4088-af62-3c2e6f5abfb7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TapasForMaskedLM\n",
    "\n",
    "model = TapasForMaskedLM(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "d94455eb-402f-49b8-907e-23f9c0bcbe34",
   "metadata": {},
   "outputs": [],
   "source": [
    "if n_gpu > 1:\n",
    "    model = torch.nn.DataParallel(model).to(args.device)\n",
    "    \n",
    "else:\n",
    "    model = model.to(args.device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39827412-a682-43ea-b129-2243da37a437",
   "metadata": {},
   "source": [
    "앞서 말씀드렸 듯이 MLM은 맨 마지막 classifier에서 vocab size 만큼의 output label 수를 가지기 때문에 발생되는 파라미터 수가 굉장히 큽니다.\n",
    "\n",
    "따라서, 메모리 용량이 큰 single GPU나 여러 개의 GPU가 할당된 서버가 아닌 이상, 학습을 수행하기 어렵습니다.\n",
    "\n",
    "제 경우에는 multi GPU인 `Titan-RTX-4way`(24GB*4)를 사용하였습니다.\n",
    "\n",
    "multi-gpu를 활용하는 경우, 분산 학습을 위해 모델은 DataParallel이라는 객체에 묶이게 됩니다.\n",
    "\n",
    "아래 model 구조를 참고해주세요!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "80ef4702-c77e-43da-b9b8-19505c046440",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DataParallel(\n",
       "  (module): TapasForMaskedLM(\n",
       "    (tapas): TapasModel(\n",
       "      (embeddings): TapasEmbeddings(\n",
       "        (word_embeddings): Embedding(32001, 768, padding_idx=0)\n",
       "        (position_embeddings): Embedding(1024, 768)\n",
       "        (token_type_embeddings_0): Embedding(3, 768)\n",
       "        (token_type_embeddings_1): Embedding(256, 768)\n",
       "        (token_type_embeddings_2): Embedding(256, 768)\n",
       "        (token_type_embeddings_3): Embedding(2, 768)\n",
       "        (token_type_embeddings_4): Embedding(256, 768)\n",
       "        (token_type_embeddings_5): Embedding(256, 768)\n",
       "        (token_type_embeddings_6): Embedding(10, 768)\n",
       "        (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (encoder): TapasEncoder(\n",
       "        (layer): ModuleList(\n",
       "          (0): TapasLayer(\n",
       "            (attention): TapasAttention(\n",
       "              (self): TapasSelfAttention(\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): TapasSelfOutput(\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): TapasIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): TapasOutput(\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",
       "          (1): TapasLayer(\n",
       "            (attention): TapasAttention(\n",
       "              (self): TapasSelfAttention(\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): TapasSelfOutput(\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): TapasIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): TapasOutput(\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",
       "          (2): TapasLayer(\n",
       "            (attention): TapasAttention(\n",
       "              (self): TapasSelfAttention(\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): TapasSelfOutput(\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): TapasIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): TapasOutput(\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",
       "          (3): TapasLayer(\n",
       "            (attention): TapasAttention(\n",
       "              (self): TapasSelfAttention(\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): TapasSelfOutput(\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): TapasIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): TapasOutput(\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",
       "          (4): TapasLayer(\n",
       "            (attention): TapasAttention(\n",
       "              (self): TapasSelfAttention(\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): TapasSelfOutput(\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): TapasIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): TapasOutput(\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",
       "          (5): TapasLayer(\n",
       "            (attention): TapasAttention(\n",
       "              (self): TapasSelfAttention(\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): TapasSelfOutput(\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): TapasIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): TapasOutput(\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",
       "          (6): TapasLayer(\n",
       "            (attention): TapasAttention(\n",
       "              (self): TapasSelfAttention(\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): TapasSelfOutput(\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): TapasIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): TapasOutput(\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",
       "          (7): TapasLayer(\n",
       "            (attention): TapasAttention(\n",
       "              (self): TapasSelfAttention(\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): TapasSelfOutput(\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): TapasIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): TapasOutput(\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",
       "          (8): TapasLayer(\n",
       "            (attention): TapasAttention(\n",
       "              (self): TapasSelfAttention(\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): TapasSelfOutput(\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): TapasIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): TapasOutput(\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",
       "          (9): TapasLayer(\n",
       "            (attention): TapasAttention(\n",
       "              (self): TapasSelfAttention(\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): TapasSelfOutput(\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): TapasIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): TapasOutput(\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",
       "          (10): TapasLayer(\n",
       "            (attention): TapasAttention(\n",
       "              (self): TapasSelfAttention(\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): TapasSelfOutput(\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): TapasIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): TapasOutput(\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",
       "          (11): TapasLayer(\n",
       "            (attention): TapasAttention(\n",
       "              (self): TapasSelfAttention(\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): TapasSelfOutput(\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): TapasIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): TapasOutput(\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",
       "    )\n",
       "    (cls): TapasOnlyMLMHead(\n",
       "      (predictions): TapasLMPredictionHead(\n",
       "        (transform): TapasPredictionHeadTransform(\n",
       "          (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "        )\n",
       "        (decoder): Linear(in_features=768, out_features=32001, bias=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68fd82d1-2868-4c08-baa0-2e7e2c078bcc",
   "metadata": {},
   "source": [
    "우리가 사전학습을 잘 수행하고서 풀어야 할 downstream task는 결국 Table Question-Answering입니다.\n",
    "\n",
    "Fine-tuning(Table QA 학습)을 할 때, Pre-training에 쓰였던 모든 파라미터가 다 필요하지는 않습니다.\n",
    "\n",
    "- query와 table 각 토큰들의 Embedding을 수행하고 representation을 생성하는 `TapasModel`만이 fine-tuning 모델 구축에 활용됩니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2913dc54-c9c6-437d-aff4-439ff276c31c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 이 cell을 실행하시면 TapasModel의 구조만 따로 확인할 수 있씁니다.\n",
    "model.module.tapas"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f9882f7-f2f2-4af7-b639-3704fa23b131",
   "metadata": {},
   "source": [
    "따라서, 모델 학습 과정 중에 저장할 파라미터는 `torch.save(model.module.tapas.state_dict(), output_dir)`와 같은 형식으로 정의됩니다.\n",
    "\n",
    "- 이에 대한 코드는 아래 `3. 학습` 파트의 `train` 함수에서 (*) 부분을 보시면 됩니다!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92bca5f2-ac4b-4441-a0da-bc5e83cb3ddf",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8d563dc-5264-46a7-aeb8-67bb4f7345a2",
   "metadata": {},
   "source": [
    "## 3. 학습"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6171232c-8b0f-446c-880c-63c8d4f4df90",
   "metadata": {},
   "source": [
    "모델을 정의했으니, 이제 학습과 평가에 필요한 train, evaluate 함수를 정의하고, 실제 학습을 수행할 단계입니다.\n",
    "\n",
    "학습 과정 중 기록되는 log는 `wandb` 플랫폼에서 시각화 되도록 설계했습니다."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59688b40-28b0-4a88-9624-b89fdd61bdc6",
   "metadata": {},
   "source": [
    "### 3.1. 학습 준비"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f4742e2-c351-40f6-bd4c-f822c04a124d",
   "metadata": {},
   "source": [
    "#### 3.1.1. train, eval 함수 정의"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "d230fe5c-990d-4720-9714-8535f5f1b630",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(args, model, train_dataloader, valid_dataloader, optimizer, lr_scheduler):\n",
    "    \n",
    "    global_step = 0\n",
    "    best_loss = 100\n",
    "\n",
    "    wandb.init(\n",
    "        project=args.wandb_project,\n",
    "        name=args.wandb_name,\n",
    "        entity=args.wandb_entity\n",
    "    )\n",
    "    \n",
    "    wandb.config.update(\n",
    "        {\n",
    "            \"epochs\": args.epoch,\n",
    "            \"batch_size\": args.batch_size,\n",
    "            \"learning_rate\": args.learning_rate,\n",
    "        }\n",
    "    )\n",
    "    wandb.watch(model, log=\"all\")\n",
    "\n",
    "    for epoch in tqdm(range(1, args.epoch+1)):\n",
    "        \n",
    "        model.train()\n",
    "        \n",
    "        losses = 0\n",
    "        for batch_idx, batch in enumerate(train_dataloader):\n",
    "            \n",
    "            inputs = {\n",
    "                'input_ids': batch['input_ids'].to(args.device),\n",
    "                'attention_mask': batch['attention_mask'].to(args.device),\n",
    "                'token_type_ids': batch['token_type_ids'].to(args.device),\n",
    "                'labels': batch['labels'].to(args.device),\n",
    "            }\n",
    "            \n",
    "            output = model(**inputs)\n",
    "\n",
    "            loss = output.loss.mean()\n",
    "            losses += loss.item()\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            lr_scheduler.step()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            global_step += 1\n",
    "\n",
    "            \n",
    "            if global_step % args.eval_step == 0:\n",
    "                \n",
    "                eval_loss = evaluate(\n",
    "                    args, model,\n",
    "                    valid_dataloader\n",
    "                )\n",
    "                \n",
    "                # train logging\n",
    "                wandb.log({\n",
    "                    'train_mlm_loss': loss.item(),\n",
    "                    'eval_mlm_loss': eval_loss\n",
    "                })\n",
    "                \n",
    "\n",
    "                if eval_loss < best_loss:\n",
    "                    best_loss = eval_loss\n",
    "                    output_dir = os.path.join(args.output_dir, \"best_tapas_model.pt\")\n",
    "                    \n",
    "                    # (*) 모델 저장\n",
    "                    # torch.save(model.tapas.state_dict(), output_dir) #single GPU\n",
    "                    torch.save(model.module.tapas.state_dict(), output_dir) # multi GPU\n",
    "\n",
    "        \n",
    "        epoch_eval_loss = evaluate(args, model,valid_dataloader)\n",
    "        \n",
    "        print(f\"[Epoch{epoch}] Train mlm loss: {losses/len(train_dataloader)}, Eval mlm loss: {epoch_eval_loss}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "de9df6e4-3cd0-423c-a5dc-99d9913ada7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(args, model, eval_dataloader):\n",
    "    \n",
    "    model.eval()    # close drop out, batch normalization\n",
    "\n",
    "    eval_loss = 0\n",
    "\n",
    "    for batch_idx, batch in enumerate(eval_dataloader):\n",
    "        with torch.no_grad():\n",
    "            inputs = {\n",
    "                'input_ids': batch['input_ids'].to(args.device),\n",
    "                'attention_mask': batch['attention_mask'].to(args.device),\n",
    "                'token_type_ids': batch['token_type_ids'].to(args.device),\n",
    "                'labels': batch['labels'].to(args.device),\n",
    "            }\n",
    "            \n",
    "            output = model(**inputs)\n",
    "\n",
    "            loss = output.loss.mean()\n",
    "\n",
    "            eval_loss += loss.item()\n",
    "            \n",
    "    return eval_loss/len(eval_dataloader)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2954ab6d-b9c7-4846-9548-08c3312480cf",
   "metadata": {},
   "source": [
    "#### 3.1.2. optimizer, lr_scheduler 정의"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "837f5645-cdd6-4292-a72f-f75514e60479",
   "metadata": {},
   "outputs": [],
   "source": [
    "no_decay = [\"bias\", \"LayerNorm.weight\"]\n",
    "optimizer_grouped_parameters = [\n",
    "    {\n",
    "        \"params\": [\n",
    "            p\n",
    "            for n, p in model.named_parameters()\n",
    "            if not any(nd in n for nd in no_decay)\n",
    "        ],\n",
    "        \"weight_decay\": args.weight_decay,\n",
    "    },\n",
    "    {\n",
    "        \"params\": [\n",
    "            p\n",
    "            for n, p in model.named_parameters()\n",
    "            if any(nd in n for nd in no_decay)\n",
    "        ],\n",
    "        \"weight_decay\": args.weight_decay,\n",
    "    },\n",
    "]\n",
    "optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate)\n",
    "\n",
    "train_step = args.epoch * len(train_dataloader)\n",
    "\n",
    "lr_scheduler = get_scheduler(\n",
    "        name=args.lr_scheduler_type,\n",
    "        optimizer=optimizer,\n",
    "        num_warmup_steps=args.num_warmup_steps,\n",
    "        num_training_steps=train_step,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "845dabd1-0621-4ba6-8a80-014ef99ec1b2",
   "metadata": {},
   "source": [
    "### 3.2. Let's train!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "10ff5145-5314-4d81-b443-b031ce11d76c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "\n",
    "# 경고 메시지를 무시하고 숨기거나\n",
    "warnings.filterwarnings(action='ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "0b673acb-2e9e-4c15-a135-cf5cab49a830",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Logging into wandb.ai. (Learn how to deploy a W&B server locally: https://wandb.me/wandb-server)\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: You can find your API key in your browser here: https://wandb.ai/authorize\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Paste an API key from your profile and hit enter, or press ctrl+c to quit:"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "  ········\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[32m\u001b[41mERROR\u001b[0m API key must be 40 characters long, yours was 7\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Logging into wandb.ai. (Learn how to deploy a W&B server locally: https://wandb.me/wandb-server)\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: You can find your API key in your browser here: https://wandb.ai/authorize\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Paste an API key from your profile and hit enter, or press ctrl+c to quit:"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "  ········\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Appending key for api.wandb.ai to your netrc file: /home/go60/.netrc\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.12.21"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/workspace/codes/lab602/tablemrc/wandb/run-20220728_080052-32o4lkqn</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href=\"https://wandb.ai/yookyungkho/Table%20Pretraining/runs/32o4lkqn\" target=\"_blank\">tapas-base-mlm-clean</a></strong> to <a href=\"https://wandb.ai/yookyungkho/Table%20Pretraining\" target=\"_blank\">Weights & Biases</a> (<a href=\"https://wandb.me/run\" target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|                                                                                                                                                                                                                  | 0/5 [00:00<?, ?it/s]/opt/conda/lib/python3.7/site-packages/torch/nn/parallel/_functions.py:68: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n",
      " 20%|████████████████████████████████████████▍                                                                                                                                                                 | 1/5 [00:27<01:49, 27.41s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Epoch1] Train mlm loss: 9.664233554493297, Eval mlm loss: 8.657608985900879\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|████████████████████████████████████████████████████████████████████████████████▊                                                                                                                         | 2/5 [00:51<01:15, 25.25s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Epoch2] Train mlm loss: 8.181182167746805, Eval mlm loss: 7.921424865722656\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 60%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▏                                                                                | 3/5 [01:13<00:47, 23.88s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Epoch3] Train mlm loss: 7.727213035930287, Eval mlm loss: 7.8767523765563965\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 80%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▌                                        | 4/5 [01:37<00:23, 23.92s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Epoch4] Train mlm loss: 7.543686736713756, Eval mlm loss: 7.736621379852295\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [01:59<00:00, 23.95s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Epoch5] Train mlm loss: 7.48064695705067, Eval mlm loss: 7.823406219482422\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "train(args, model, train_dataloader, valid_dataloader, optimizer, lr_scheduler)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18a88a3b-2c4c-4d52-af9e-2a3e3ed2250c",
   "metadata": {},
   "source": [
    "학습 중간에 출력된 경고 메세지는 DDP가 아닌 DP를 사용해서 그렇습니다.([관련 link](https://github.com/huggingface/transformers/issues/14128))\n",
    "\n",
    "향후 DDP 버전으로 업데이트할 예정이니 참고해주세요!"
   ]
  }
 ],
 "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.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}