[선택 과제 -1] Transformer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Transformer 구현 실습\n",
    "\n",
    "- Transformer 모델을 구현해보는 실습입니다.\n",
    "- 논문 : https://arxiv.org/abs/1706.03762\n",
    "- 정답 + 참고 + 원본 : https://github.com/hyunwoongko/transformer\n",
    "![model](https://raw.githubusercontent.com/hyunwoongko/transformer/master/image/model.png)"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "source": [
    "import torch\r\n",
    "import torch.nn as nn\r\n",
    "import torch.functional as F\r\n",
    "import math"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 1. Positional Encoding\n",
    "![](https://raw.githubusercontent.com/hyunwoongko/transformer/master/image/positional_encoding.jpg)"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "source": [
    "class PositionalEncoding(nn.Module):\r\n",
    "    \"\"\"\r\n",
    "    compute sinusoid encoding.\r\n",
    "    \"\"\"\r\n",
    "    def __init__(self, d_model, max_len, device):\r\n",
    "        \"\"\"\r\n",
    "        constructor of sinusoid encoding class\r\n",
    "\r\n",
    "        :param d_model: dimension of model\r\n",
    "        :param max_len: max sequence length\r\n",
    "        :param device: hardware device setting\r\n",
    "        \"\"\"\r\n",
    "        super(PositionalEncoding, self).__init__()\r\n",
    "\r\n",
    "        # same size with input matrix (for adding with input matrix)\r\n",
    "\r\n",
    "        '''\r\n",
    "        TODO를  작성해 주세요.\r\n",
    "        '''\r\n",
    "        \r\n",
    "        self.encoding = torch.zeros(max_len, d_model, device=device)\r\n",
    "        self.encoding.requires_grad = '''TODO 1'''   # we don't need to compute gradient\r\n",
    "\r\n",
    "        pos = torch.arange(0, max_len, device=device)\r\n",
    "        pos = '''TODO 2'''\r\n",
    "        # 1D => 2D unsqueeze to represent word's position\r\n",
    "\r\n",
    "        _2i = torch.arange(0, d_model, step=2, device=device).float()\r\n",
    "        # 'i' means index of d_model (e.g. embedding size = 50, 'i' = [0,50])\r\n",
    "        # \"step=2\" means 'i' multiplied with two (same with 2 * i)\r\n",
    "\r\n",
    "        self.encoding[:, 0::2] = torch.sin(pos / (10000 ** ('''TODO 3''')))\r\n",
    "        self.encoding[:, 1::2] = torch.cos(pos / (10000 ** ('''TODO 3''')))\r\n",
    "        # compute positional encoding to consider positional information of words\r\n",
    "\r\n",
    "    def forward(self, x):\r\n",
    "\r\n",
    "        batch_size, seq_len = x.size()\r\n",
    "\r\n",
    "        return '''TODO 4'''"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "source": [
    "def padding(data):\r\n",
    "    max_len = len(max(data, key=len))\r\n",
    "    pad_id = 0\r\n",
    "\r\n",
    "    for i, seq in enumerate(data):\r\n",
    "        if len(seq) < max_len:\r\n",
    "            data[i] = seq + [pad_id] * (max_len - len(seq))\r\n",
    "\r\n",
    "    return data, max_len\r\n",
    "\r\n",
    "def test():\r\n",
    "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\r\n",
    "    data = [\r\n",
    "        [62, 13, 47, 39, 78, 33, 56, 13],\r\n",
    "        [60, 96, 51, 32, 90],\r\n",
    "        [35, 45, 48, 65, 91, 99, 92, 10, 3, 21],\r\n",
    "        [66, 88, 98, 47],\r\n",
    "        [77, 65, 51, 77, 19, 15, 35, 19, 23]\r\n",
    "    ]\r\n",
    "    vocab_size = 100\r\n",
    "    \r\n",
    "    data, max_len = padding(data)\r\n",
    "    d_model = 512 \r\n",
    "    batch = torch.LongTensor(data)\r\n",
    "    positional_encoder = PositionalEncoding(d_model, max_len, device)\r\n",
    "    \r\n",
    "    encode_data = positional_encoder(batch)\r\n",
    "    \r\n",
    "    assert encode_data.size() == torch.Size([10,512]), 'Worng implementation, Check your shape'\r\n",
    "    print('Correct implementation!')\r\n",
    "    \r\n",
    "test()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 2. Scale Dot Product Attention & Multi-Head Attention\n",
    "![](https://raw.githubusercontent.com/hyunwoongko/transformer/master/image/scale_dot_product_attention.jpg)\n",
    "![](https://raw.githubusercontent.com/hyunwoongko/transformer/master/image/multi_head_attention.jpg)"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "source": [
    "class ScaleDotProductAttention(nn.Module):\r\n",
    "    \"\"\"\r\n",
    "    compute scale dot product attention\r\n",
    "\r\n",
    "    Query : given sentence that we focused on (decoder)\r\n",
    "    Key : every sentence to check relationship with Qeury(encoder)\r\n",
    "    Value : every sentence same with Key (encoder)\r\n",
    "    \"\"\"\r\n",
    "\r\n",
    "    def __init__(self):\r\n",
    "        super(ScaleDotProductAttention, self).__init__()\r\n",
    "        self.softmax = nn.Softmax(dim = 1)\r\n",
    "\r\n",
    "    def forward(self, q, k, v, mask=None, e=1e-12):\r\n",
    "        '''\r\n",
    "        TODO를 작성해 주세요.\r\n",
    "        '''\r\n",
    "        # input is 4 dimension tensor\r\n",
    "        # [batch_size, head, length, d_tensor]\r\n",
    "        batch_size, head, length, d_tensor = '''TODO 5'''\r\n",
    "\r\n",
    "        # 1. dot product Query with Key^T to compute similarity\r\n",
    "        \r\n",
    "        k_t = k.view(batch_size, head, d_tensor, length)  # transpose\r\n",
    "        \r\n",
    "        score = '''TODO 6'''  # scaled dot product\r\n",
    "\r\n",
    "        # 2. apply masking (opt)\r\n",
    "        if mask is not None:\r\n",
    "            score = score.masked_fill(mask == 0, -e)\r\n",
    "\r\n",
    "        # 3. pass them softmax to make [0, 1] range\r\n",
    "        score = '''TODO 7'''\r\n",
    "\r\n",
    "        # 4. multiply with Value\r\n",
    "        v = '''TODO 8'''\r\n",
    "\r\n",
    "        return v, score\r\n",
    "     \r\n",
    "    \r\n",
    "class MultiHeadAttention(nn.Module):\r\n",
    "    '''\r\n",
    "    TODO를 작성해 주세요.\r\n",
    "    ''' \r\n",
    "    def __init__(self, d_model, n_head):\r\n",
    "        super(MultiHeadAttention, self).__init__()\r\n",
    "        self.n_head = n_head\r\n",
    "        self.attention = '''TODO 9'''\r\n",
    "        self.w_q = '''TODO 10'''\r\n",
    "        self.w_k = '''TODO 10'''\r\n",
    "        self.w_v = '''TODO 10'''\r\n",
    "        self.w_concat = '''TODO 10'''\r\n",
    "\r\n",
    "    def forward(self, q, k, v, mask=None):\r\n",
    "        # 1. dot product with weight matrices\r\n",
    "        q, k, v = '''TODO 11'''\r\n",
    "\r\n",
    "        # 2. split tensor by number of heads\r\n",
    "        q, k, v = self.split(q), self.split(k), self.split(v)\r\n",
    "\r\n",
    "        # 3. do scale dot product to compute similarity\r\n",
    "        out, attention = self.attention(q, k, v, mask=mask)\r\n",
    "        \r\n",
    "        # 4. concat and pass to linear layer\r\n",
    "        out = self.concat(out)\r\n",
    "        out = self.w_concat(out)\r\n",
    "\r\n",
    "        return out\r\n",
    "\r\n",
    "    def split(self, tensor):\r\n",
    "        \"\"\"\r\n",
    "        split tensor by number of head\r\n",
    "\r\n",
    "        :param tensor: [batch_size, length, d_model]\r\n",
    "        :return: [batch_size, head, length, d_tensor]\r\n",
    "        \"\"\"\r\n",
    "        batch_size, length, d_model = tensor.size()\r\n",
    "\r\n",
    "        d_tensor = d_model // self.n_head\r\n",
    "        tensor = tensor.view(batch_size, self.n_head, length, d_tensor)\r\n",
    "        # it is similar with group convolution (split by number of heads)\r\n",
    "\r\n",
    "        return tensor\r\n",
    "\r\n",
    "    def concat(self, tensor):\r\n",
    "        \"\"\"\r\n",
    "        inverse function of self.split(tensor : torch.Tensor)\r\n",
    "\r\n",
    "        :param tensor: [batch_size, head, length, d_tensor]\r\n",
    "        :return: [batch_size, length, d_model]\r\n",
    "        \"\"\"\r\n",
    "        batch_size, head, length, d_tensor = tensor.size()\r\n",
    "        d_model = head * d_tensor\r\n",
    "\r\n",
    "        tensor = tensor.view(batch_size, length, d_model)\r\n",
    "        return tensor"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "source": [
    "def test():\r\n",
    "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\r\n",
    "    data = [\r\n",
    "        [62, 13, 47, 39, 78, 33, 56, 13],\r\n",
    "        [60, 96, 51, 32, 90],\r\n",
    "        [35, 45, 48, 65, 91, 99, 92, 10, 3, 21],\r\n",
    "        [66, 88, 98, 47],\r\n",
    "        [77, 65, 51, 77, 19, 15, 35, 19, 23]\r\n",
    "    ]\r\n",
    "    vocab_size = 100\r\n",
    "    data, max_len = padding(data)\r\n",
    "    \r\n",
    "    d_model = 512  # model의 hidden size\r\n",
    "    num_heads = 8 # head\r\n",
    "    \r\n",
    "    embedding = nn.Embedding(vocab_size, d_model)\r\n",
    "    \r\n",
    "    \r\n",
    "    batch = torch.LongTensor(data) # (B, L)\r\n",
    "    batch_emb = embedding(batch).to(device) # ()\r\n",
    "    \r\n",
    "    positional_encoder = PositionalEncoding(d_model, max_len, device)\r\n",
    "    mha = MultiHeadAttention(d_model,num_heads).to(device)\r\n",
    "    \r\n",
    "    encode_emb = positional_encoder(batch)\r\n",
    "    \r\n",
    "    batch_emb = batch_emb + encode_emb\r\n",
    "    \r\n",
    "    out = mha(batch_emb,batch_emb,batch_emb)\r\n",
    "    \r\n",
    "    assert out.shape == torch.Size([5,10,512]), 'Worng implementation, Check your shape'\r\n",
    "    print('Correct implementation!')\r\n",
    "\r\n",
    "test()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 3. Layer Norm\n",
    "![](https://raw.githubusercontent.com/hyunwoongko/transformer/master/image/layer_norm.jpg)"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "source": [
    "class LayerNorm(nn.Module):\n",
    "    '''\n",
    "    TODO를 작성해 주세요.\n",
    "    '''\n",
    "    def __init__(self, d_model, eps=1e-12):\n",
    "        super(LayerNorm, self).__init__()\n",
    "        self.gamma = nn.Parameter(torch.ones(d_model))\n",
    "        self.beta = nn.Parameter(torch.zeros(d_model))\n",
    "        self.eps = eps\n",
    "\n",
    "    def forward(self, x):\n",
    "        mean = x.mean(-1, keepdim=True)\n",
    "        std = x.std(-1, keepdim=True)\n",
    "        # '-1' means last dimension. \n",
    "\n",
    "        out = '''TODO 12'''\n",
    "        out = '''TODO 13'''\n",
    "        return out"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "source": [
    "def test():\n",
    "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "    data = [\n",
    "        [[1.,2.,3.],\n",
    "        [6.,5.,4.]],\n",
    "        \n",
    "        [[10.,11.,12.],\n",
    "        [80.,80.,90.]]\n",
    "    ]\n",
    "    vocab_size = 100\n",
    "    data, max_len = padding(data)\n",
    "    \n",
    "    d_model = 3  # model의 hidden size\n",
    "    \n",
    "    batch = torch.Tensor(data) # (B, L)\n",
    "    batch = batch # ()\n",
    "    \n",
    "    layer_norm = LayerNorm(d_model)\n",
    "    \n",
    "    out = layer_norm(batch)\n",
    "    out = out.long()\n",
    "    \n",
    "    assert out.shape == torch.Size([2,2,3]),'Worng implementation, Check your shape'\n",
    "    assert torch.all(out == torch.tensor([[[-1, 0, 1],[1, 0,-1]],[[-1,0,1],[0 ,0, 1]]])),'Worng implementation, Check your code'\n",
    "    print('Correct implementation!')\n",
    "\n",
    "test()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 4. Positionwise Feed Forward\n",
    "![](https://raw.githubusercontent.com/hyunwoongko/transformer/master/image/positionwise_feed_forward.jpg)"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "source": [
    "class PositionwiseFeedForward(nn.Module):\n",
    "\n",
    "    def __init__(self, d_model, hidden, drop_prob=0.1):\n",
    "        super(PositionwiseFeedForward, self).__init__()\n",
    "        self.linear1 = '''TODO 14'''\n",
    "        self.linear2 = '''TODO 15'''\n",
    "        self.relu = nn.ReLU()\n",
    "        self.dropout = nn.Dropout(p=drop_prob)\n",
    "\n",
    "    def forward(self, x):\n",
    "        '''TODO 16'''\n",
    "        return x"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "source": [
    "def test():\n",
    "    data = [\n",
    "        [62, 13, 47, 39, 78, 33, 56, 13],\n",
    "        [60, 96, 51, 32, 90],\n",
    "        [35, 45, 48, 65, 91, 99, 92, 10, 3, 21],\n",
    "        [66, 88, 98, 47],\n",
    "        [77, 65, 51, 77, 19, 15, 35, 19, 23]\n",
    "    ]\n",
    "    vocab_size = 100\n",
    "    data, max_len = padding(data)\n",
    "    \n",
    "    d_model = 512  # model의 hidden size\n",
    "    embedding = nn.Embedding(vocab_size, d_model)\n",
    "    batch = torch.LongTensor(data)\n",
    "    \n",
    "    batch_emb = embedding(batch)\n",
    "    #batch_emb = torch.LongTensor(emb_data) # (B, L)\n",
    "    \n",
    "    data, max_len = padding(data)\n",
    "    hidden = 1024\n",
    "    \n",
    "    pfn = PositionwiseFeedForward(d_model,hidden,0)\n",
    "    out = pfn(batch_emb)\n",
    "    \n",
    "    assert out.shape == torch.Size([5, 10, 512]),'Worng implementation, Check your shape'\n",
    "    print('Correct implementation!')\n",
    "\n",
    "test()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 5. Encoder & Decoder Structure\n",
    "![](https://raw.githubusercontent.com/hyunwoongko/transformer/master/image/enc_dec.jpg)"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 5.1 Encoder"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "source": [
    "class EncoderLayer(nn.Module):\n",
    "\n",
    "    def __init__(self, d_model, ffn_hidden, n_head, drop_prob):\n",
    "        super(EncoderLayer, self).__init__()\n",
    "        self.attention = '''TODO 17'''\n",
    "        self.norm1 = '''TODO 18'''\n",
    "        self.dropout1 = nn.Dropout(p=drop_prob)\n",
    "\n",
    "        self.ffn = '''TODO 19'''\n",
    "        self.norm2 = '''TODO 18'''\n",
    "        self.dropout2 = nn.Dropout(p=drop_prob)\n",
    "\n",
    "    def forward(self, x, src_mask):\n",
    "        # 1. compute self attention\n",
    "        '''TODO 20'''\n",
    "        \n",
    "        # 2. add and norm\n",
    "        '''TODO 21'''\n",
    "        \n",
    "        # 3. positionwise feed forward network\n",
    "        '''TODO 22'''\n",
    "      \n",
    "        # 4. add and norm\n",
    "        '''TODO 23'''\n",
    "        return x"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "source": [
    "class TokenEmbedding(nn.Embedding):\n",
    "    \"\"\"\n",
    "    Token Embedding using torch.nn\n",
    "    they will dense representation of word using weighted matrix\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, vocab_size, d_model):\n",
    "        \"\"\"\n",
    "        class for token embedding that included positional information\n",
    "        :param vocab_size: size of vocabulary\n",
    "        :param d_model: dimensions of model\n",
    "        \"\"\"\n",
    "        super(TokenEmbedding, self).__init__(vocab_size, d_model, padding_idx=1)\n",
    "        \n",
    "class TransformerEmbedding(nn.Module):\n",
    "    \"\"\"\n",
    "    token embedding + positional encoding (sinusoid)\n",
    "    positional encoding can give positional information to network\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, vocab_size, d_model, max_len, drop_prob, device):\n",
    "        \"\"\"\n",
    "        class for word embedding that included positional information\n",
    "        :param vocab_size: size of vocabulary\n",
    "        :param d_model: dimensions of model\n",
    "        \"\"\"\n",
    "        super(TransformerEmbedding, self).__init__()\n",
    "        self.tok_emb = '''TODO 24'''\n",
    "        self.pos_emb = '''TODO 25'''\n",
    "        self.drop_out = nn.Dropout(p=drop_prob)\n",
    "\n",
    "    def forward(self, x):\n",
    "        tok_emb = '''TODO 26'''\n",
    "        pos_emb = '''TODO 27'''\n",
    "        return self.drop_out(tok_emb + pos_emb)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "source": [
    "class Encoder(nn.Module):\n",
    "\n",
    "    def __init__(self, enc_voc_size, max_len, d_model, ffn_hidden, n_head, n_layers, drop_prob, device):\n",
    "        super().__init__()\n",
    "        self.emb = '''TODO 28'''\n",
    "\n",
    "        self.layers = nn.ModuleList(['''TODO 29'''])\n",
    "\n",
    "    def forward(self, x, src_mask):\n",
    "        x = self.emb(x)\n",
    "\n",
    "        '''TODO 30'''\n",
    "\n",
    "        return x"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "source": [
    "def test():\n",
    "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "    data = [\n",
    "        [62, 13, 47, 39, 78, 33, 56, 13],\n",
    "        [60, 96, 51, 32, 90],\n",
    "        [35, 45, 48, 65, 91, 99, 92, 10, 3, 21],\n",
    "        [66, 88, 98, 47],\n",
    "        [77, 65, 51, 77, 19, 15, 35, 19, 23]\n",
    "    ]\n",
    "    data, max_len = padding(data)\n",
    "    vocab_size = 100\n",
    "    data, max_len = padding(data)\n",
    "    n_layers = 6\n",
    "    \n",
    "    d_model = 100  # model의 hidden size\n",
    "    batch = torch.LongTensor(data).to(device)\n",
    "    \n",
    "    encoder = Encoder(enc_voc_size=vocab_size, max_len = max_len,ffn_hidden = 20,d_model=d_model,device=device,drop_prob=0,n_head=5,n_layers=n_layers)\n",
    "    encoder.to(device)\n",
    "    test = encoder(batch,None)\n",
    "    \n",
    "    assert len(encoder.layers) == n_layers, 'Worng implementation, Check your Encoder''s layers'\n",
    "    print('Correct Implementation!')\n",
    "test()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 5.2 Decoder"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "source": [
    "class DecoderLayer(nn.Module):\n",
    "\n",
    "    def __init__(self, d_model, ffn_hidden, n_head, drop_prob):\n",
    "        super(DecoderLayer, self).__init__()\n",
    "        self.self_attention = '''TODO 31'''\n",
    "        self.norm1 = '''TODO 32'''\n",
    "        self.dropout1 = nn.Dropout(p=drop_prob)\n",
    "\n",
    "        self.enc_dec_attention = '''TODO 31'''\n",
    "        self.norm2 = '''TODO 32'''\n",
    "        self.dropout2 = nn.Dropout(p=drop_prob)\n",
    "\n",
    "        self.ffn = '''TODO 33'''\n",
    "        self.norm3 = '''TODO 32'''\n",
    "        self.dropout3 = nn.Dropout(p=drop_prob)\n",
    "\n",
    "    def forward(self, dec, enc, trg_mask, src_mask):    \n",
    "        # 1. compute self attention\n",
    "        '''TODO 34'''\n",
    "        \n",
    "        # 2. add and norm\n",
    "        '''TODO 35'''\n",
    "\n",
    "        if enc is not None:\n",
    "            # 3. compute encoder - decoder attention\n",
    "            '''TODO 36'''\n",
    "            \n",
    "            # 4. add and norm\n",
    "            '''TODO 37'''\n",
    "\n",
    "        # 5. positionwise feed forward network\n",
    "        '''TODO 38'''\n",
    "        \n",
    "        # 6. add and norm\n",
    "        '''TODO 39'''\n",
    "        return x"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, dec_voc_size, max_len, d_model, ffn_hidden, n_head, n_layers, drop_prob, device):\n",
    "        super().__init__()\n",
    "        self.emb = TransformerEmbedding(d_model=d_model,\n",
    "                                        drop_prob=drop_prob,\n",
    "                                        max_len=max_len,\n",
    "                                        vocab_size=dec_voc_size,\n",
    "                                        device=device)\n",
    "\n",
    "        self.layers = nn.ModuleList(['''TODO 40'''])\n",
    "\n",
    "        self.linear = nn.Linear(d_model, dec_voc_size)\n",
    "\n",
    "    def forward(self, trg, src, trg_mask, src_mask):\n",
    "        trg = self.emb(trg)\n",
    "\n",
    "        '''TODO 41'''\n",
    "\n",
    "        # pass to LM head\n",
    "        output = self.linear(trg)\n",
    "        return output"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "source": [
    "def test():\n",
    "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "    data = [\n",
    "        [62, 13, 47, 39, 78, 33, 56, 13],\n",
    "        [60, 96, 51, 32, 90],\n",
    "        [35, 45, 48, 65, 91, 99, 92, 10, 3, 21],\n",
    "        [66, 88, 98, 47],\n",
    "        [77, 65, 51, 77, 19, 15, 35, 19, 23]\n",
    "    ]\n",
    "    data, max_len = padding(data)\n",
    "    vocab_size = 100\n",
    "    data, max_len = padding(data)\n",
    "    n_layers = 6\n",
    "    \n",
    "    d_model = 100  # model의 hidden size\n",
    "    batch = torch.LongTensor(data).to(device)\n",
    "    \n",
    "    encoder = Encoder(enc_voc_size=vocab_size, max_len = max_len,ffn_hidden = 20,d_model=d_model,device=device,drop_prob=0,n_head=5,n_layers=n_layers)\n",
    "    encoder.to(device)\n",
    "    src = encoder(batch,None)\n",
    "    \n",
    "    decoder = Decoder(dec_voc_size =vocab_size, max_len=max_len, d_model=d_model, ffn_hidden=20, n_head=5, n_layers=n_layers, drop_prob=0, device=device)\n",
    "    decoder.to(device)\n",
    "    output = decoder(batch, src,None,None)\n",
    "    \n",
    "    print('Correct Implementation!')\n",
    "test()"
   ],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}