Natural_Language_Processing/Text_Classification_with_Transformers/text_classification_with_transformers_KYK.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "20e3da5f",
   "metadata": {},
   "source": [
    "# Text classification with Transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d604ba46",
   "metadata": {},
   "source": [
    "작성: 고유경\n",
    "\n",
    "---\n",
    "\n",
    "[Original Keras Code](https://keras.io/examples/nlp/text_classification_with_transformer/)\n",
    "\n",
    "\n",
    "**Main Task**: IMDB 영화 리뷰 데이터 감성 분석 (긍/부정)\n",
    "\n",
    "reference\n",
    "\n",
    "- https://tutorials.pytorch.kr/beginner/text_sentiment_ngrams_tutorial.html\n",
    "\n",
    "- https://wikidocs.net/60691\n",
    "\n",
    "- https://cpm0722.github.io/pytorch-implementation/transformer\n",
    "\n",
    "---\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "028245bd",
   "metadata": {},
   "source": [
    "## 0. Import"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16964bd2",
   "metadata": {},
   "source": [
    "- `torch == 1.10.0`\n",
    "- `torchtext == 0.11.0`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "44697c6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import os\n",
    "import math\n",
    "from timeit import default_timer as timer\n",
    "from typing import Iterable, List\n",
    "from IPython.display import Image\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn.utils.rnn import pad_sequence\n",
    "from torch import Tensor\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data.dataset import random_split\n",
    "from torchsummary import summary\n",
    "\n",
    "import torchtext\n",
    "from torchtext.data.utils import get_tokenizer\n",
    "from torchtext.data.functional import to_map_style_dataset\n",
    "from torchtext.vocab import build_vocab_from_iterator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9eac9cf1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "NVIDIA GeForce RTX 2080 Ti\n",
      "1\n"
     ]
    }
   ],
   "source": [
    "## GPU check\n",
    "print(torch.cuda.is_available())\n",
    "print(torch.cuda.get_device_name(0))\n",
    "print(torch.cuda.device_count())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "52c24b8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_EPOCHS = 5\n",
    "LEARNING_RATE = 0.01\n",
    "BATCH_SIZE = 64\n",
    "UNK_IDX, PAD_IDX = 0, 1\n",
    "EMB_SIZE = 512\n",
    "NHEAD = 8\n",
    "FFN_HID_DIM = 128\n",
    "\n",
    "#NUM_ENCODER_LAYERS = 3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d8fadc7",
   "metadata": {},
   "source": [
    "## 1. 데이터 로드\n",
    "\n",
    "IMDB 영화 리뷰 데이터\n",
    "\n",
    "- 리뷰 50000개로 구성\n",
    "\n",
    "- torchtext 내장 데이터 (for text classification)\n",
    "\n",
    "- `torchtext.datasets.IMDB`로 data_iterator 간편히 불러오기 가능\n",
    "\n",
    "- data split을 위해 `torchtext.data.functional.to_map_style_dataset`을 통해 iterable-style dataset -> map-style dataset 형태로 변환 ([참고-공식문서](https://pytorch.org/text/stable/data_functional.html))\n",
    "\n",
    "- 7.5:1.5:1로 train, valid, test 분할"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "92f4cca2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Data load\n",
    "train_iter, test_iter = torchtext.datasets.IMDB(root='.data', split=('train', 'test'))\n",
    "\n",
    "# Convert iterable-style dataset to map-style dataset.\n",
    "train_dataset = to_map_style_dataset(train_iter)\n",
    "test_dataset = to_map_style_dataset(test_iter)\n",
    "total_dataset = train_dataset + test_dataset\n",
    "\n",
    "# Data split\n",
    "num_train = int(len(total_dataset) * 0.75)\n",
    "num_val = int(len(total_dataset) * 0.15)\n",
    "num_test = len(total_dataset) - num_train - num_val\n",
    "\n",
    "train_data, val_data, test_data = random_split(total_dataset, [num_train, num_val, num_test])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "91a57902",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37500 training set\n",
      "7500 validation set\n",
      "5000 test set\n"
     ]
    }
   ],
   "source": [
    "print(f\"{len(train_data)} training set\")\n",
    "print(f\"{len(val_data)} validation set\")\n",
    "print(f\"{len(test_data)} test set\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c5082d22",
   "metadata": {},
   "outputs": [],
   "source": [
    "## smaller version(실험용)\n",
    "\n",
    "train_iter = torchtext.datasets.IMDB(root='.data', split=('train'))\n",
    "train_dataset = to_map_style_dataset(train_iter)\n",
    "dataset = train_dataset[12450:12590] #리뷰 140개만 사용\n",
    "train_data, val_data, test_data = random_split(dataset, [100, 20, 20])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "11c6f2bb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100 training set\n",
      "20 validation set\n",
      "20 test set\n"
     ]
    }
   ],
   "source": [
    "print(f\"{len(train_data)} training set\")\n",
    "print(f\"{len(val_data)} validation set\")\n",
    "print(f\"{len(test_data)} test set\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1c78f04d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('pos',\n",
       "  'I just recently watched this 1954 movie starring Vincent Price for the first time on Turner Classic Movies. Price portrays Don Gallico, a magician/inventor who is driven to murder when his boss steals several of his magical inventions (and also his wife, portrayed in a brief role by the lovely Eva Gabor). Even though Price is a murderer, I actually found myself rooting for him, he is a sympathetic character who is driven mad by the greedy people around him who keep taking advantage of him.<br /><br />Although this movie doesn\\'t have the \"horror\" factor of some of his more famous roles (such as my favorite, \"House of Wax\") it nonetheless has enough going for it to keep the viewers interest. <br /><br />This is a must for Vincent Price fans.'),\n",
       " ('neg',\n",
       "  \"Exceedingly complicated and drab. I'm a bright guy, but this was just too much for a tired brain. It would really benefit from a few early clues as to who these people are and what they are doing. Probably better for the US market. GC himself hinted that this alone did not supply his Oscar and you can see why.<br /><br />Still the sand dunes are pretty. The nail pulling is nasty. The attorneys drunk dad is a mystery. The cricket is good to see.<br /><br />Very difficult to write the required ten lines on this, despite it being over 2 hours long. Thank heavens they shortened it. Admittedly we don't get to the pictures much, but the last film we saw, Walk the Line, was 10 times better and I don't really like Johnny Cash. My wife says George still looks good with the beard and a few extra pounds so there's that.....am I nearly there yet ?<br /><br />How about now\")]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## data 예시 출력\n",
    "ex_list = [train_data[1], train_data[3]]\n",
    "ex_list"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81b98ac0",
   "metadata": {},
   "source": [
    "## 2. 텍스트 전처리"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00b17d4e",
   "metadata": {},
   "source": [
    "### 2.1. html 태그, 특수문자, 숫자 삭제"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4ae0d367",
   "metadata": {},
   "outputs": [],
   "source": [
    "def clean_web_text(text):\n",
    "    # text 내 html 태그 삭제\n",
    "    text = text.replace(\"<br />\", \" \")\n",
    "    text = text.replace(\"&quot;\", '\"')\n",
    "    text = text.replace(\"<p>\", \" \")\n",
    "    text = text.replace(\"<a href=\", \" \")\n",
    "    text = text.replace(\"</a>\", \"\")\n",
    "    \n",
    "    # 알파벳 제외 특수문자, 숫자, 공백 삭제\n",
    "    text = text.replace(\"\\\\n\", \" \")\n",
    "    text = re.sub('[^A-Za-z\\s]', '', text)\n",
    "    text = text.replace(\"   \", \"\") ## 3개 이상 공백은 제거\n",
    "    text = text.replace(\"  \", \" \")\n",
    "    text = text.lower()\n",
    "    \n",
    "    return text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "69cf931a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "i just recently watched this movie starring vincent price for the first time on turner classic movies price portrays don gallico a magicianinventor who is driven to murder when his boss steals several of his magical inventions and also his wife portrayed in a brief role by the lovely eva gabor even though price is a murderer i actually found myself rooting for him he is a sympathetic character who is driven mad by the greedy people around him who keep taking advantage of him although this movie doesnt have the horror factor of some of his more famous roles such as my favorite house of wax it nonetheless has enough going for it to keep the viewers interestthis is a must for vincent price fans\n",
      "exceedingly complicated and drab im a bright guy but this was just too much for a tired brain it would really benefit from a few early clues as to who these people are and what they are doing probably better for the us market gc himself hinted that this alone did not supply his oscar and you can see why still the sand dunes are pretty the nail pulling is nasty the attorneys drunk dad is a mystery the cricket is good to see very difficult to write the required ten lines on this despite it being over hours long thank heavens they shortened it admittedly we dont get to the pictures much but the last film we saw walk the line was times better and i dont really like johnny cash my wife says george still looks good with the beard and a few extra pounds so theres thatam i nearly there yethow about now\n"
     ]
    }
   ],
   "source": [
    "for t in ex_list:\n",
    "    print(clean_web_text(t[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e528f54",
   "metadata": {},
   "source": [
    "### 2.2.  토크나이저 정의 및 단어사전 구축\n",
    "\n",
    "- spacy tockenizer 사전 다운로드 필요\n",
    "\n",
    "- generator **`yield_tokens`**를 정의하여 train text 내 문장 정제, 토크나이징 및 단어 집합을 반복적으로 수행할 수 있도록 함\n",
    "\n",
    "    - 제너레이터는 함수 안에서 yield라는 키워드만 사용하여 대용량 반복 수행 시, 메모리 효율적으로 사용 ([참고](https://dojang.io/mod/page/view.php?id=2412))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0077f53f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# spacy 토크나이저 다운\n",
    "#!pip install -U spacy\n",
    "#!python -m spacy download en_core_web_sm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "89e913cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenizer(text):\n",
    "    tokenize = get_tokenizer('spacy', language = \"en_core_web_sm\")\n",
    "    clean_text = clean_web_text(text)\n",
    "    return tokenize(clean_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b87689f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def yield_tokens(data_iter):    \n",
    "    for label, text in data_iter:\n",
    "        yield tokenizer(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "566ee7b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab = build_vocab_from_iterator(yield_tokens(train_data), min_freq=5, specials=[\"<unk>\", \"<pad>\"])\n",
    "vocab.set_default_index(vocab[\"<unk>\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "d62cc7f1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('<unk>', 0),\n",
       " ('<pad>', 1),\n",
       " ('the', 2),\n",
       " ('and', 3),\n",
       " ('a', 4),\n",
       " ('of', 5),\n",
       " ('to', 6),\n",
       " ('is', 7),\n",
       " ('in', 8),\n",
       " ('i', 9),\n",
       " ('it', 10),\n",
       " ('this', 11),\n",
       " ('that', 12),\n",
       " ('was', 13),\n",
       " ('as', 14),\n",
       " ('but', 15),\n",
       " ('with', 16),\n",
       " ('film', 17),\n",
       " ('movie', 18),\n",
       " ('for', 19),\n",
       " ('on', 20),\n",
       " ('he', 21),\n",
       " ('you', 22),\n",
       " ('nt', 23),\n",
       " ('his', 24),\n",
       " ('are', 25),\n",
       " ('have', 26),\n",
       " ('not', 27),\n",
       " ('one', 28),\n",
       " ('who', 29),\n",
       " ('like', 30),\n",
       " ('be', 31),\n",
       " ('they', 32),\n",
       " ('an', 33),\n",
       " ('its', 34),\n",
       " ('about', 35),\n",
       " ('all', 36),\n",
       " ('at', 37),\n",
       " ('or', 38),\n",
       " ('if', 39),\n",
       " ('from', 40),\n",
       " ('good', 41),\n",
       " ('just', 42),\n",
       " ('so', 43),\n",
       " ('when', 44),\n",
       " ('out', 45),\n",
       " ('more', 46),\n",
       " ('by', 47),\n",
       " ('s', 48),\n",
       " ('see', 49),\n",
       " ('what', 50),\n",
       " ('has', 51),\n",
       " ('do', 52),\n",
       " ('my', 53),\n",
       " ('there', 54),\n",
       " ('up', 55),\n",
       " ('other', 56),\n",
       " ('would', 57),\n",
       " ('some', 58),\n",
       " ('very', 59),\n",
       " ('even', 60),\n",
       " ('only', 61),\n",
       " ('time', 62),\n",
       " ('two', 63),\n",
       " ('me', 64),\n",
       " ('could', 65),\n",
       " ('films', 66),\n",
       " ('first', 67),\n",
       " ('had', 68),\n",
       " ('really', 69),\n",
       " ('which', 70),\n",
       " ('her', 71),\n",
       " ('she', 72),\n",
       " ('can', 73),\n",
       " ('will', 74),\n",
       " ('did', 75),\n",
       " ('home', 76),\n",
       " ('no', 77),\n",
       " ('their', 78),\n",
       " ('we', 79),\n",
       " ('how', 80),\n",
       " ('movies', 81),\n",
       " ('them', 82),\n",
       " ('show', 83),\n",
       " ('get', 84),\n",
       " ('into', 85),\n",
       " ('made', 86),\n",
       " ('watch', 87),\n",
       " ('him', 88),\n",
       " ('than', 89),\n",
       " ('story', 90),\n",
       " ('well', 91),\n",
       " ('life', 92),\n",
       " ('many', 93),\n",
       " ('much', 94),\n",
       " ('were', 95),\n",
       " ('through', 96),\n",
       " ('von', 97),\n",
       " ('also', 98),\n",
       " ('because', 99),\n",
       " ('characters', 100),\n",
       " ('been', 101),\n",
       " ('both', 102),\n",
       " ('great', 103),\n",
       " ('kibbutz', 104),\n",
       " ('make', 105),\n",
       " ('seen', 106),\n",
       " ('too', 107),\n",
       " ('way', 108),\n",
       " ('any', 109),\n",
       " ('europa', 110),\n",
       " ('then', 111),\n",
       " ('best', 112),\n",
       " ('does', 113),\n",
       " ('ever', 114),\n",
       " ('people', 115),\n",
       " ('real', 116),\n",
       " ('alone', 117),\n",
       " ('bad', 118),\n",
       " ('being', 119),\n",
       " ('most', 120),\n",
       " ('these', 121),\n",
       " ('think', 122),\n",
       " ('those', 123),\n",
       " ('after', 124),\n",
       " ('now', 125),\n",
       " ('plot', 126),\n",
       " ('acting', 127),\n",
       " ('alicia', 128),\n",
       " ('funny', 129),\n",
       " ('over', 130),\n",
       " ('every', 131),\n",
       " ('man', 132),\n",
       " ('say', 133),\n",
       " ('school', 134),\n",
       " ('should', 135),\n",
       " ('such', 136),\n",
       " ('things', 137),\n",
       " ('world', 138),\n",
       " ('again', 139),\n",
       " ('another', 140),\n",
       " ('deanna', 141),\n",
       " ('here', 142),\n",
       " ('little', 143),\n",
       " ('never', 144),\n",
       " ('where', 145),\n",
       " ('ca', 146),\n",
       " ('off', 147),\n",
       " ('why', 148),\n",
       " ('come', 149),\n",
       " ('girl', 150),\n",
       " ('house', 151),\n",
       " ('m', 152),\n",
       " ('trier', 153),\n",
       " ('better', 154),\n",
       " ('find', 155),\n",
       " ('go', 156),\n",
       " ('know', 157),\n",
       " ('own', 158),\n",
       " ('price', 159),\n",
       " ('saw', 160),\n",
       " ('look', 161),\n",
       " ('love', 162),\n",
       " ('maybe', 163),\n",
       " ('seems', 164),\n",
       " ('ve', 165),\n",
       " ('back', 166),\n",
       " ('before', 167),\n",
       " ('dvd', 168),\n",
       " ('happen', 169),\n",
       " ('old', 170),\n",
       " ('original', 171),\n",
       " ('quite', 172),\n",
       " ('right', 173),\n",
       " ('same', 174),\n",
       " ('scene', 175),\n",
       " ('star', 176),\n",
       " ('still', 177),\n",
       " ('thing', 178),\n",
       " ('though', 179),\n",
       " ('trying', 180),\n",
       " ('while', 181),\n",
       " ('years', 182),\n",
       " ('your', 183),\n",
       " ('around', 184),\n",
       " ('black', 185),\n",
       " ('character', 186),\n",
       " ('enough', 187),\n",
       " ('few', 188),\n",
       " ('job', 189),\n",
       " ('lot', 190),\n",
       " ('room', 191),\n",
       " ('away', 192),\n",
       " ('d', 193),\n",
       " ('excellent', 194),\n",
       " ('must', 195),\n",
       " ('nothing', 196),\n",
       " ('place', 197),\n",
       " ('script', 198),\n",
       " ('shooting', 199),\n",
       " ('something', 200),\n",
       " ('times', 201),\n",
       " ('train', 202),\n",
       " ('watching', 203),\n",
       " ('without', 204),\n",
       " ('work', 205),\n",
       " ('act', 206),\n",
       " ('actors', 207),\n",
       " ('anyone', 208),\n",
       " ('cast', 209),\n",
       " ('comes', 210),\n",
       " ('different', 211),\n",
       " ('director', 212),\n",
       " ('during', 213),\n",
       " ('end', 214),\n",
       " ('give', 215),\n",
       " ('horror', 216),\n",
       " ('human', 217),\n",
       " ('kind', 218),\n",
       " ('point', 219),\n",
       " ('scenes', 220),\n",
       " ('shot', 221),\n",
       " ('since', 222),\n",
       " ('stooges', 223),\n",
       " ('take', 224),\n",
       " ('takes', 225),\n",
       " ('us', 226),\n",
       " ('always', 227),\n",
       " ('american', 228),\n",
       " ('bit', 229),\n",
       " ('cd', 230),\n",
       " ('christmas', 231),\n",
       " ('cinema', 232),\n",
       " ('dark', 233),\n",
       " ('done', 234),\n",
       " ('down', 235),\n",
       " ('each', 236),\n",
       " ('fact', 237),\n",
       " ('fun', 238),\n",
       " ('having', 239),\n",
       " ('head', 240),\n",
       " ('however', 241),\n",
       " ('left', 242),\n",
       " ('makes', 243),\n",
       " ('night', 244),\n",
       " ('once', 245),\n",
       " ('perhaps', 246),\n",
       " ('pretty', 247),\n",
       " ('rather', 248),\n",
       " ('want', 249),\n",
       " ('whole', 250),\n",
       " ('year', 251),\n",
       " ('action', 252),\n",
       " ('although', 253),\n",
       " ('busy', 254),\n",
       " ('camera', 255),\n",
       " ('course', 256),\n",
       " ('feel', 257),\n",
       " ('high', 258),\n",
       " ('instead', 259),\n",
       " ('lars', 260),\n",
       " ('laugh', 261),\n",
       " ('line', 262),\n",
       " ('long', 263),\n",
       " ('lou', 264),\n",
       " ('mad', 265),\n",
       " ('main', 266),\n",
       " ('our', 267),\n",
       " ('performance', 268),\n",
       " ('performances', 269),\n",
       " ('person', 270),\n",
       " ('play', 271),\n",
       " ('probably', 272),\n",
       " ('shows', 273),\n",
       " ('someone', 274),\n",
       " ('terrible', 275),\n",
       " ('thought', 276),\n",
       " ('three', 277),\n",
       " ('tv', 278),\n",
       " ('video', 279),\n",
       " ('viewer', 280),\n",
       " ('wax', 281),\n",
       " ('yet', 282),\n",
       " ('young', 283),\n",
       " ('actually', 284),\n",
       " ('almost', 285),\n",
       " ('amazing', 286),\n",
       " ('between', 287),\n",
       " ('breaking', 288),\n",
       " ('called', 289),\n",
       " ('came', 290),\n",
       " ('christensen', 291),\n",
       " ('effects', 292),\n",
       " ('ending', 293),\n",
       " ('feature', 294),\n",
       " ('girls', 295),\n",
       " ('going', 296),\n",
       " ('got', 297),\n",
       " ('hospital', 298),\n",
       " ('keep', 299),\n",
       " ('last', 300),\n",
       " ('let', 301),\n",
       " ('looking', 302),\n",
       " ('may', 303),\n",
       " ('mind', 304),\n",
       " ('part', 305),\n",
       " ('role', 306),\n",
       " ('roles', 307),\n",
       " ('several', 308),\n",
       " ('style', 309),\n",
       " ('together', 310),\n",
       " ('turn', 311),\n",
       " ('until', 312),\n",
       " ('values', 313),\n",
       " ('vincent', 314),\n",
       " ('war', 315),\n",
       " ('worth', 316),\n",
       " ('writer', 317),\n",
       " ('zentropa', 318),\n",
       " ('actor', 319),\n",
       " ('believe', 320),\n",
       " ('big', 321),\n",
       " ('car', 322),\n",
       " ('case', 323),\n",
       " ('crime', 324),\n",
       " ('day', 325),\n",
       " ('despite', 326),\n",
       " ('drama', 327),\n",
       " ('element', 328),\n",
       " ('erika', 329),\n",
       " ('everything', 330),\n",
       " ('family', 331),\n",
       " ('favorite', 332),\n",
       " ('four', 333),\n",
       " ('hard', 334),\n",
       " ('images', 335),\n",
       " ('inside', 336),\n",
       " ('interesting', 337),\n",
       " ('later', 338),\n",
       " ('leo', 339),\n",
       " ('loved', 340),\n",
       " ('matter', 341),\n",
       " ('might', 342),\n",
       " ('minutes', 343),\n",
       " ('money', 344),\n",
       " ('mother', 345),\n",
       " ('new', 346),\n",
       " ('picture', 347),\n",
       " ('piece', 348),\n",
       " ('played', 349),\n",
       " ('plays', 350),\n",
       " ('poor', 351),\n",
       " ('re', 352),\n",
       " ('remember', 353),\n",
       " ('rest', 354),\n",
       " ('reunion', 355),\n",
       " ('ryan', 356),\n",
       " ('shemp', 357),\n",
       " ('special', 358),\n",
       " ('sure', 359),\n",
       " ('totally', 360),\n",
       " ('uncle', 361),\n",
       " ('waves', 362),\n",
       " ('works', 363),\n",
       " ('alex', 364),\n",
       " ('am', 365),\n",
       " ('anything', 366),\n",
       " ('audience', 367),\n",
       " ('beautiful', 368),\n",
       " ('boring', 369),\n",
       " ('boy', 370),\n",
       " ('boys', 371),\n",
       " ('budget', 372),\n",
       " ('buy', 373),\n",
       " ('chance', 374),\n",
       " ('comedy', 375),\n",
       " ('complete', 376),\n",
       " ('days', 377),\n",
       " ('death', 378),\n",
       " ('doing', 379),\n",
       " ('either', 380),\n",
       " ('else', 381),\n",
       " ('enjoy', 382),\n",
       " ('entirely', 383),\n",
       " ('far', 384),\n",
       " ('fish', 385),\n",
       " ('goodrich', 386),\n",
       " ('himself', 387),\n",
       " ('idea', 388),\n",
       " ('john', 389),\n",
       " ('kid', 390),\n",
       " ('least', 391),\n",
       " ('leopold', 392),\n",
       " ('looked', 393),\n",
       " ('magician', 394),\n",
       " ('mr', 395),\n",
       " ('music', 396),\n",
       " ('next', 397),\n",
       " ('obviously', 398),\n",
       " ('phillips', 399),\n",
       " ('reb', 400),\n",
       " ('recently', 401),\n",
       " ('recommend', 402),\n",
       " ('released', 403),\n",
       " ('richard', 404),\n",
       " ('second', 405),\n",
       " ('seemed', 406),\n",
       " ('sequence', 407),\n",
       " ('series', 408),\n",
       " ('shootings', 409),\n",
       " ('songs', 410),\n",
       " ('state', 411),\n",
       " ('stop', 412),\n",
       " ('students', 413),\n",
       " ('throughout', 414),\n",
       " ('trains', 415),\n",
       " ('understand', 416),\n",
       " ('across', 417),\n",
       " ('against', 418),\n",
       " ('album', 419),\n",
       " ('among', 420),\n",
       " ('appears', 421),\n",
       " ('cat', 422),\n",
       " ('color', 423),\n",
       " ('comic', 424),\n",
       " ('coming', 425),\n",
       " ('culkin', 426),\n",
       " ('deal', 427),\n",
       " ('definitely', 428),\n",
       " ('dialogue', 429),\n",
       " ('direction', 430),\n",
       " ('drugs', 431),\n",
       " ('due', 432),\n",
       " ('emotional', 433),\n",
       " ('entire', 434),\n",
       " ('especially', 435),\n",
       " ('feeling', 436),\n",
       " ('found', 437),\n",
       " ('friends', 438),\n",
       " ('gallico', 439),\n",
       " ('gave', 440),\n",
       " ('germany', 441),\n",
       " ('gets', 442),\n",
       " ('gives', 443),\n",
       " ('ha', 444),\n",
       " ('help', 445),\n",
       " ('horrible', 446),\n",
       " ('huge', 447),\n",
       " ('incredibly', 448),\n",
       " ('involved', 449),\n",
       " ('israel', 450),\n",
       " ('journey', 451),\n",
       " ('kids', 452),\n",
       " ('late', 453),\n",
       " ('lead', 454),\n",
       " ('leads', 455),\n",
       " ('less', 456),\n",
       " ('ll', 457),\n",
       " ('low', 458),\n",
       " ('mark', 459),\n",
       " ('members', 460),\n",
       " ('michael', 461),\n",
       " ('offers', 462),\n",
       " ('pain', 463),\n",
       " ('please', 464),\n",
       " ('police', 465),\n",
       " ('prices', 466),\n",
       " ('put', 467),\n",
       " ('scott', 468),\n",
       " ('screen', 469),\n",
       " ('sense', 470),\n",
       " ('short', 471),\n",
       " ('sort', 472),\n",
       " ('started', 473),\n",
       " ('stupid', 474),\n",
       " ('surprise', 475),\n",
       " ('taking', 476),\n",
       " ('tell', 477),\n",
       " ('ten', 478),\n",
       " ('terrorists', 479),\n",
       " ('tragedy', 480),\n",
       " ('triers', 481),\n",
       " ('turned', 482),\n",
       " ('used', 483),\n",
       " ('version', 484),\n",
       " ('waste', 485),\n",
       " ('white', 486),\n",
       " ('wonderful', 487),\n",
       " ('able', 488),\n",
       " ('age', 489),\n",
       " ('al', 490),\n",
       " ('already', 491),\n",
       " ('apparently', 492),\n",
       " ('art', 493),\n",
       " ('barr', 494),\n",
       " ('bland', 495),\n",
       " ('book', 496),\n",
       " ('boyfriend', 497),\n",
       " ('business', 498),\n",
       " ('change', 499),\n",
       " ('cheesy', 500),\n",
       " ('child', 501),\n",
       " ('cinematography', 502),\n",
       " ('class', 503),\n",
       " ('climax', 504),\n",
       " ('completely', 505),\n",
       " ('control', 506),\n",
       " ('crazy', 507),\n",
       " ('creatures', 508),\n",
       " ('dancer', 509),\n",
       " ('decent', 510),\n",
       " ('denver', 511),\n",
       " ('detective', 512),\n",
       " ('died', 513),\n",
       " ('directing', 514),\n",
       " ('discussion', 515),\n",
       " ('dogma', 516),\n",
       " ('dream', 517),\n",
       " ('editing', 518),\n",
       " ('elements', 519),\n",
       " ('everyone', 520),\n",
       " ('except', 521),\n",
       " ('experience', 522),\n",
       " ('extremely', 523),\n",
       " ('eyes', 524),\n",
       " ('fan', 525),\n",
       " ('final', 526),\n",
       " ('form', 527),\n",
       " ('full', 528),\n",
       " ('genius', 529),\n",
       " ('given', 530),\n",
       " ('gone', 531),\n",
       " ('guy', 532),\n",
       " ('guys', 533),\n",
       " ('happened', 534),\n",
       " ('hidden', 535),\n",
       " ('highly', 536),\n",
       " ('hit', 537),\n",
       " ('hollywood', 538),\n",
       " ('hope', 539),\n",
       " ('imdb', 540),\n",
       " ('including', 541),\n",
       " ('invasion', 542),\n",
       " ('issues', 543),\n",
       " ('kessler', 544),\n",
       " ('larry', 545),\n",
       " ('learn', 546),\n",
       " ('live', 547),\n",
       " ('looks', 548),\n",
       " ('luke', 549),\n",
       " ('lydia', 550),\n",
       " ('major', 551),\n",
       " ('making', 552),\n",
       " ('max', 553),\n",
       " ('mean', 554),\n",
       " ('mental', 555),\n",
       " ('mouth', 556),\n",
       " ('moving', 557),\n",
       " ('mystery', 558),\n",
       " ('near', 559),\n",
       " ('noir', 560),\n",
       " ('often', 561),\n",
       " ('ones', 562),\n",
       " ('oscar', 563),\n",
       " ('others', 564),\n",
       " ('philipps', 565),\n",
       " ('political', 566),\n",
       " ('railway', 567),\n",
       " ('read', 568),\n",
       " ('reason', 569),\n",
       " ('record', 570),\n",
       " ('reed', 571),\n",
       " ('release', 572),\n",
       " ('seeing', 573),\n",
       " ('seem', 574),\n",
       " ('society', 575),\n",
       " ('sound', 576),\n",
       " ('space', 577),\n",
       " ('spies', 578),\n",
       " ('stars', 579),\n",
       " ('storyline', 580),\n",
       " ('stuff', 581),\n",
       " ('taken', 582),\n",
       " ('television', 583),\n",
       " ('third', 584),\n",
       " ('tony', 585),\n",
       " ('took', 586),\n",
       " ('use', 587),\n",
       " ('uses', 588),\n",
       " ('victor', 589),\n",
       " ('visual', 590),\n",
       " ('went', 591),\n",
       " ('woman', 592),\n",
       " ('worst', 593),\n",
       " ('wrestler', 594)]"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocabs = vocab.get_stoi()\n",
    "sorted(vocabs.items(), key = lambda x: x[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "b9abca1b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "단어 집합 크기: 595\n"
     ]
    }
   ],
   "source": [
    "VOCAB_SIZE = len(vocab)\n",
    "print(f\"단어 집합 크기: {VOCAB_SIZE}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c108371",
   "metadata": {},
   "source": [
    "### 2.3. data -> tensor 변환 (Vectorization)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "a1b7c2a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Text -> Tensor w/ index from vacab_dict\n",
    "def text2tensor(text):\n",
    "    tokens = tokenizer(text)\n",
    "    token_ids = vocab(tokens)\n",
    "    return torch.tensor(token_ids, dtype=torch.int64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "849e34ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Encode label -> pos=1, neg=0\n",
    "def label2num(label):\n",
    "    n_label = 1 if label == 'pos' else 0\n",
    "    return n_label"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b61b103",
   "metadata": {},
   "source": [
    "### 2.4. Dataloader 정의"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74a9363d",
   "metadata": {},
   "source": [
    "- `collate_fn`: collate lists of samples into batches ([참고](https://pytorch.org/docs/stable/data.html#dataloader-collate-fn))\n",
    "\n",
    "    - `custom_collate_fn`을 정의하여 text를 input으로 받았을 때 vectorization, padding 등 연속 작업을 수행하여 iterable한 dataloader로 반환할 수 있도록 함\n",
    "    - `torch.utils.data.DataLoader`의 인자로 들어감\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "ed4e972d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda')"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "c10e31d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# DatoLoader에 쓰일 collate_fn 커스터마이징\n",
    "\n",
    "def custom_collate_fn(batch):\n",
    "    label_batch = []\n",
    "    text_batch = []\n",
    "    \n",
    "    for label, text in batch:\n",
    "        label_batch.append(label2num(label))\n",
    "        text_batch.append(text2tensor(text))\n",
    "        \n",
    "    label_batch = torch.tensor(label_batch, dtype=torch.int64)\n",
    "    \n",
    "    # padding\n",
    "    text_batch = pad_sequence(text_batch, padding_value=PAD_IDX)\n",
    "    \n",
    "    # batch_first = True\n",
    "    text_batch = text_batch.transpose(0,1) # (batch_size x max_len)\n",
    "    \n",
    "    return label_batch.to(device), text_batch.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "83dabf4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# train, valid 용 dataloader 생성\n",
    "train_dataloader = DataLoader(train_data, batch_size=BATCH_SIZE, collate_fn=custom_collate_fn)\n",
    "val_dataloader = DataLoader(val_data, batch_size=BATCH_SIZE, collate_fn=custom_collate_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "7380a1e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataloader = DataLoader(train_data, batch_size = 16, collate_fn = custom_collate_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "2d453d4d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(tensor([1, 1, 0, 0, 0, 0, 1, 1, 0, 1, 0, 1, 1, 1, 1, 1], device='cuda:0'), tensor([[ 76, 191, 290,  ..., 215,  11,  18],\n",
      "        [  9,  42, 401,  ...,   1,   1,   1],\n",
      "        [ 58, 137,   0,  ...,   1,   1,   1],\n",
      "        ...,\n",
      "        [  9, 152,  40,  ...,   1,   1,   1],\n",
      "        [  9,  26,   6,  ...,   1,   1,   1],\n",
      "        [  9,  13,  55,  ...,   1,   1,   1]], device='cuda:0'))\n"
     ]
    }
   ],
   "source": [
    "## batch 예시 출력\n",
    "batch_ex = next(iter(train_dataloader))\n",
    "print(batch_ex)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "e830bc16",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([16, 405])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch_ex[1].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cfb10c9",
   "metadata": {},
   "source": [
    "## 3. 모델링(Transformer)\n",
    "\n",
    "![transformer](transformer.png)\n",
    "\n",
    "- classification이기 때문에 Encoder block으로만 구성"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "120caa96",
   "metadata": {},
   "source": [
    "### 1. Positional Embedding\n",
    "\n",
    "- positional encoding + token embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "550a5975",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEmbedding(nn.Module):\n",
    "    def __init__(self, vocab_size=VOCAB_SIZE, emb_size=512, pad_idx=PAD_IDX, device=device, maxlen=1000):\n",
    "        super().__init__()\n",
    "        self.token_embedding = nn.Embedding(vocab_size, emb_size, padding_idx=pad_idx).to(device) #(128, 17, 5)\n",
    "        \n",
    "        pos = torch.arange(0,maxlen).reshape(maxlen,1).to(device)\n",
    "        two_i = torch.arange(0,emb_size,2).to(device)\n",
    "        \n",
    "        self.pos_encoding = torch.zeros((maxlen, emb_size))\n",
    "        self.pos_encoding = self.pos_encoding.to(device)\n",
    "        self.pos_encoding.requires_grad= False\n",
    "        \n",
    "        self.pos_encoding[:,0::2] = torch.sin(pos/ (10000**(two_i/emb_size))) #짝수 index\n",
    "        self.pos_encoding[:,1::2] = torch.cos(pos/ (10000**(two_i/emb_size))) # 홀수 index\n",
    "        \n",
    "    def forward(self, x):\n",
    "        try:\n",
    "            sen_length = x.shape[1] #(batch_size, seq_len, emb_size) 3d tensor 기준\n",
    "        except:\n",
    "            sen_length = x.shape[0] #(batch_size, seq_len, emb_size) 2d tensor 기준\n",
    "\n",
    "        return self.token_embedding(x) + self.pos_encoding[:sen_length,:]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f07b727",
   "metadata": {},
   "source": [
    "### 2. Encoder\n",
    "\n",
    "MSA, FFN 이전(input), 이후 차원 동일\n",
    "\n",
    "\n",
    "- **Multi-head Self-Attention** (토큰 간 관계, 문맥 학습)\n",
    "\n",
    "- Residual Connection\n",
    "- Layer Norm\n",
    "    \n",
    "- **Position-wise FFN** (head 별 mix up)\n",
    "\n",
    "- Residual Connection\n",
    "- Layer Norm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1844f1a2",
   "metadata": {},
   "source": [
    "**Multi-head Self-Attention**\n",
    "\n",
    "- pos_embedding을 거친 임베딩 벡터는 query, key, value로 copy되어 각각 선형 결합을 통해 동일한 차원의 w_q, w_k, w_v로 재표현됨\n",
    "\n",
    "- 이후, `num_head` 개수의 head로 나뉘어 각각 dot product, scaling, masking, softmax를 거쳐 attention score matrix (batch_size, num_head, max_len, d_head)이 도출됨\n",
    "\n",
    "- 최종적으로 head 별 attn_score concat하여 MHA-score 도출"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "e53b3d38",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, num_head=8, emb_size=512):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.num_head = num_head\n",
    "        self.emb_size = emb_size\n",
    "        self.linear_query = nn.Linear(self.emb_size, self.emb_size).to(device)\n",
    "        self.linear_key = nn.Linear(self.emb_size, self.emb_size).to(device)\n",
    "        self.linear_value = nn.Linear(self.emb_size, self.emb_size).to(device)\n",
    "        \n",
    "    \n",
    "    def attention(self, q, k, v, mask=None):\n",
    "        batch_size, num_head, max_len, d_head = q.size()\n",
    "        \n",
    "        # 1. dot product (q*k^T)\n",
    "        k_T = k.transpose(2,3)\n",
    "        attn = torch.matmul(q, k_T) ## -> attn: (batch_size, num_head, max_len, max_len)\n",
    "        \n",
    "        # 2. scaling (divide w/ sqrt(d_k))\n",
    "        attn = attn / math.sqrt(d_head)\n",
    "        \n",
    "        # 3. masking!!!!\n",
    "        if mask is not None:\n",
    "            attn = attn.masked_fill(mask == 0, -1e9)\n",
    "        \n",
    "        # 4. softmax\n",
    "        attn = nn.Softmax(dim=-1)(attn) ## (max_len, max_len) # 행의 합 1이 되도록 dim=-1\n",
    "        \n",
    "        # 5. multipy w/ V\n",
    "        attn_score = torch.matmul(attn, v) ## -> attn_score: (batch_size, num_head, max_len, d_head)\n",
    "        \n",
    "        return attn_score, attn\n",
    "        \n",
    "        \n",
    "    def forward(self, out_pos_emb, mask=None):\n",
    "    \n",
    "        batch_size, max_len, emb_size = out_pos_emb.size()\n",
    "        d_head = emb_size // self.num_head\n",
    "        \n",
    "        # query, key, value #각각의 가중치 벡터와 곱해져 w_q, w_k, w_v 생성\n",
    "        w_q = self.linear_query(out_pos_emb)\n",
    "        w_k = self.linear_key(out_pos_emb)\n",
    "        w_v = self.linear_value(out_pos_emb)\n",
    "        \n",
    "        # split to heads (torch.view, traspose) -> (batch_size, num_head, max_len, d_head)\n",
    "        qs = w_q.view(batch_size, max_len, self.num_head, d_head).transpose(1,2)\n",
    "        ks = w_k.view(batch_size, max_len, self.num_head, d_head).transpose(1,2)\n",
    "        vs = w_v.view(batch_size, max_len, self.num_head, d_head).transpose(1,2)\n",
    "        \n",
    "        # multi-head-attention # attn_score: (batch_size, num_head, max_len, d_head)\n",
    "        attn_score, attention = self.attention(qs, ks, vs, mask)\n",
    "        \n",
    "        # concat -> out_mha: (batch_size, max_len, emb_size)\n",
    "        out_mha = attn_score.transpose(1,2).contiguous().view(batch_size, max_len, emb_size)\n",
    "        \n",
    "        return out_mha, attention\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82e6a794",
   "metadata": {},
   "source": [
    "**Feed Forward Network**\n",
    "\n",
    "- Linear Layer 1 -> ReLU -> Linear Layer2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "bab34e23",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForwardNetwork(nn.Module):\n",
    "    def __init__(self, emb_size=512, d_hid=128):\n",
    "        super().__init__()\n",
    "        self.layer1 = nn.Linear(emb_size, d_hid)\n",
    "        self.layer2 = nn.Linear(d_hid, emb_size)\n",
    "        self.relu = nn.ReLU()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.layer1(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.layer2(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2368c9f1",
   "metadata": {},
   "source": [
    "**Transformer(Encoder) Block for Classification**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "beed38ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer_for_Classification(nn.Module):\n",
    "    def __init__(self, num_class=2, emb_size = 512, num_head = 8,\n",
    "                 layer_norm_eps=1e-05, dropout_rate = 0.2, pad_idx = PAD_IDX, device=device):\n",
    "        super().__init__()\n",
    "        self.emb_size = emb_size\n",
    "        self.num_head = num_head\n",
    "        self.pad_idx = pad_idx\n",
    "        \n",
    "        self.pos_embedding = PositionalEmbedding(vocab_size=VOCAB_SIZE, emb_size=512, pad_idx=self.pad_idx, device=device, maxlen=1000) #argument 채워넣어라\n",
    "        self.mh_attention = MultiHeadAttention(num_head=8, emb_size=512)\n",
    "        self.ffn = FeedForwardNetwork(emb_size=512, d_hid=128)\n",
    "        \n",
    "        self.classifier = nn.Linear(emb_size, num_class)\n",
    "        \n",
    "        self.layer_norm = nn.LayerNorm(self.emb_size, eps=layer_norm_eps)\n",
    "        self.dropout = nn.Dropout(dropout_rate)\n",
    "        \n",
    "        \n",
    "    \n",
    "    def forward(self, data_batch):\n",
    "        \n",
    "        # 0. padding mask 생성\n",
    "        batch_size, max_len = data_batch.size()\n",
    "        batch_mask = (data_batch != self.pad_idx).view(batch_size, 1, 1, max_len) #boolean mask: False -> 0\n",
    "        \n",
    "        # 1. Positional Embedding\n",
    "        out_pos_emb = self.pos_embedding(data_batch)\n",
    "        \n",
    "        \n",
    "        # 2. Multi-Head Attention\n",
    "        x, attn = self.mh_attention(out_pos_emb, batch_mask)\n",
    "        \n",
    "        ### Dropout\n",
    "        x = self.dropout(x)\n",
    "        ### Residual Connection\n",
    "        x = x + out_pos_emb\n",
    "        ### Layer Normalization\n",
    "        out_mha = self.layer_norm(x)\n",
    "        \n",
    "        \n",
    "        # 3. Feed Forward Network\n",
    "        x = self.ffn(out_mha)\n",
    "        \n",
    "        ### Dropout\n",
    "        x = self.dropout(x)\n",
    "        ### Residual Connection\n",
    "        x = x + out_mha\n",
    "        ### Layer Normalization\n",
    "        x = self.layer_norm(x) # -> (batch_size, max_len, emb_size)\n",
    "        \n",
    "        \n",
    "        # 4. Classifier\n",
    "        max_len = x.size(1)\n",
    "        x = nn.AvgPool2d((max_len,1))(x).squeeze() # -> (batch_size, emb_size)\n",
    "        # x.mean(dim=0)\n",
    "        \n",
    "        x = self.classifier(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47682f12",
   "metadata": {},
   "source": [
    "## 4. Training, Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "b20577f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0) #랜덤 시드 고정\n",
    "\n",
    "transformer = Transformer_for_Classification(num_class=2, emb_size = 512, num_head = 8,\n",
    "                 layer_norm_eps=1e-05, dropout_rate = 0.2, pad_idx = PAD_IDX, device=device)\n",
    "transformer = transformer.to(device)\n",
    "\n",
    "\n",
    "# CE Loss\n",
    "criterion = torch.nn.CrossEntropyLoss(ignore_index=PAD_IDX)\n",
    "\n",
    "# Optimizer : SGD\n",
    "optimizer = torch.optim.SGD(transformer.parameters(), lr=LEARNING_RATE)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "15c987e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_epoch(model, optimizer, dataloader):\n",
    "    model.train()\n",
    "    losses = 0\n",
    "    \n",
    "    for label_batch, text_batch in dataloader:\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        logits = model(text_batch)\n",
    "        loss = criterion(logits, label_batch)\n",
    "        loss.backward()\n",
    "        \n",
    "        optimizer.step()\n",
    "        \n",
    "        losses += loss.item()\n",
    "        \n",
    "    return losses / len(dataloader)\n",
    "\n",
    "def evaluate(model, dataloader):\n",
    "    model.eval()\n",
    "    losses = 0\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for label_batch, text_batch in dataloader:\n",
    "            logits = model(text_batch)\n",
    "        \n",
    "            loss = criterion(logits, label_batch)\n",
    "            losses += loss.item()\n",
    "        \n",
    "    return losses / len(dataloader)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "97a11ba0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 1, Train loss: 0.126, Val loss: 0.018, Epoch time = 44.206s\n",
      "Epoch: 2, Train loss: 0.021, Val loss: 0.009, Epoch time = 43.584s\n",
      "Epoch: 3, Train loss: 0.013, Val loss: 0.006, Epoch time = 43.304s\n",
      "Epoch: 4, Train loss: 0.010, Val loss: 0.005, Epoch time = 42.804s\n",
      "Epoch: 5, Train loss: 0.008, Val loss: 0.004, Epoch time = 42.916s\n"
     ]
    }
   ],
   "source": [
    "# Start training! (5 epochs)\n",
    "for epoch in range(1, NUM_EPOCHS+1):\n",
    "    start_time = timer()\n",
    "    train_loss = train_epoch(transformer, optimizer, train_dataloader)\n",
    "    end_time = timer()\n",
    "    val_loss = evaluate(transformer, val_dataloader)\n",
    "    print((f\"Epoch: {epoch}, Train loss: {train_loss:.3f}, Val loss: {val_loss:.3f}, \"f\"Epoch time = {(end_time - start_time):.3f}s\"))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "86499bc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save model.pth\n",
    "\n",
    "path = \"model/\"\n",
    "if not os.path.exists(path):\n",
    "    os.mkdir(path)\n",
    "torch.save(transformer, path+'transformer_cls.pth')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e178f30b",
   "metadata": {},
   "source": [
    "## 5. Test (Predict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "c0e44ad8",
   "metadata": {},
   "outputs": [],
   "source": [
    "### load saved model\n",
    "\n",
    "model = torch.load(path+'transformer_cls.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "3e839558",
   "metadata": {},
   "outputs": [],
   "source": [
    "### define function for prediction(Test)\n",
    "\n",
    "def predict(model, text_num, dataset):\n",
    "    with torch.no_grad():\n",
    "        text = text2tensor(test_data[text_num][1]).unsqueeze(0).to(device)\n",
    "        output = model(text)\n",
    "        pred_label = 'pos' if output.argmax(0).item() == 1 else 'neg'\n",
    "\n",
    "    print( f'[Label] {test_data[text_num][0]} \\n[Pred_Label] {pred_label} \\n [Text] {test_data[text_num][1]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4be4cfcb",
   "metadata": {},
   "source": [
    "**오답 예시**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "ca0a1c53",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Label] pos \n",
      "[Pred_Label] neg \n",
      " [Text] `Europa' (or, as it is also known, `Zentropa') is one of the most visually stunning films I have ever seen. The blend of grayscale and colour photography is near seamless... a true feast for the eyes. The picture was a contender for a 1991's Golden Palm in Canners. The award went to Barton Fink (by Coen brothers); a film stylistically very similar to Zentropa. Here's an exercise in class: rent both films and be a judge for yourself.\n"
     ]
    }
   ],
   "source": [
    "predict(model, 2, test_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8cfbda5",
   "metadata": {},
   "source": [
    "**정답 예시**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "e350b8f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Label] neg \n",
      "[Pred_Label] neg \n",
      " [Text] Sondra Locke stinks in this film, but then she was an awful 'actress' anyway. Unfortunately, she drags everyone else (including then =real life boyfriend Clint Eastwood down the drain with her. But what was Clint Eastwood thinking when he agreed to star in this one? One read of the script should have told him that this one was going to be a real snorer. It's an exceptionally weak story, basically no story or plot at all. Add in bored, poor acting, even from the normally good Eastwood. There's absolutely no action except a couple arguments and as far as I was concerned, this film ranks up at the top of the heap of natural sleep enhancers. Wow! Could a film BE any more boring? I think watching paint dry or the grass grow might be more fun. A real stinker. Don't bother with this one.\n"
     ]
    }
   ],
   "source": [
    "predict(model, 1, test_data)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}