03 - nmt_pretrain_ParaBART.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install torch==1.6.0\n",
    "# !pip install transformers==3.0.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'transformers.modeling_bart'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "\u001b[1;32mc:\\Users\\Guntsv\\Documents\\GitHub\\Thai-Paraphase\\03 - nmt_pretrain_ParaBART.ipynb Cell 2\u001b[0m in \u001b[0;36m<cell line: 5>\u001b[1;34m()\u001b[0m\n\u001b[0;32m      <a href='vscode-notebook-cell:/c%3A/Users/Guntsv/Documents/GitHub/Thai-Paraphase/03%20-%20nmt_pretrain_ParaBART.ipynb#W1sZmlsZQ%3D%3D?line=1'>2</a>\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mtorch\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mnn\u001b[39;00m \u001b[39mas\u001b[39;00m \u001b[39mnn\u001b[39;00m\n\u001b[0;32m      <a href='vscode-notebook-cell:/c%3A/Users/Guntsv/Documents/GitHub/Thai-Paraphase/03%20-%20nmt_pretrain_ParaBART.ipynb#W1sZmlsZQ%3D%3D?line=2'>3</a>\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mtorch\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mnn\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mfunctional\u001b[39;00m \u001b[39mas\u001b[39;00m \u001b[39mF\u001b[39;00m\n\u001b[1;32m----> <a href='vscode-notebook-cell:/c%3A/Users/Guntsv/Documents/GitHub/Thai-Paraphase/03%20-%20nmt_pretrain_ParaBART.ipynb#W1sZmlsZQ%3D%3D?line=4'>5</a>\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mtransformers\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mmodeling_bart\u001b[39;00m \u001b[39mimport\u001b[39;00m (\n\u001b[0;32m      <a href='vscode-notebook-cell:/c%3A/Users/Guntsv/Documents/GitHub/Thai-Paraphase/03%20-%20nmt_pretrain_ParaBART.ipynb#W1sZmlsZQ%3D%3D?line=5'>6</a>\u001b[0m     PretrainedBartModel,  \n\u001b[0;32m      <a href='vscode-notebook-cell:/c%3A/Users/Guntsv/Documents/GitHub/Thai-Paraphase/03%20-%20nmt_pretrain_ParaBART.ipynb#W1sZmlsZQ%3D%3D?line=6'>7</a>\u001b[0m     LayerNorm, \n\u001b[0;32m      <a href='vscode-notebook-cell:/c%3A/Users/Guntsv/Documents/GitHub/Thai-Paraphase/03%20-%20nmt_pretrain_ParaBART.ipynb#W1sZmlsZQ%3D%3D?line=7'>8</a>\u001b[0m     EncoderLayer, \n\u001b[0;32m      <a href='vscode-notebook-cell:/c%3A/Users/Guntsv/Documents/GitHub/Thai-Paraphase/03%20-%20nmt_pretrain_ParaBART.ipynb#W1sZmlsZQ%3D%3D?line=8'>9</a>\u001b[0m     DecoderLayer, \n\u001b[0;32m     <a href='vscode-notebook-cell:/c%3A/Users/Guntsv/Documents/GitHub/Thai-Paraphase/03%20-%20nmt_pretrain_ParaBART.ipynb#W1sZmlsZQ%3D%3D?line=9'>10</a>\u001b[0m     LearnedPositionalEmbedding,\n\u001b[0;32m     <a href='vscode-notebook-cell:/c%3A/Users/Guntsv/Documents/GitHub/Thai-Paraphase/03%20-%20nmt_pretrain_ParaBART.ipynb#W1sZmlsZQ%3D%3D?line=10'>11</a>\u001b[0m     _prepare_bart_decoder_inputs,\n\u001b[0;32m     <a href='vscode-notebook-cell:/c%3A/Users/Guntsv/Documents/GitHub/Thai-Paraphase/03%20-%20nmt_pretrain_ParaBART.ipynb#W1sZmlsZQ%3D%3D?line=11'>12</a>\u001b[0m     _make_linear_from_emb\n\u001b[0;32m     <a href='vscode-notebook-cell:/c%3A/Users/Guntsv/Documents/GitHub/Thai-Paraphase/03%20-%20nmt_pretrain_ParaBART.ipynb#W1sZmlsZQ%3D%3D?line=12'>13</a>\u001b[0m )\n\u001b[0;32m     <a href='vscode-notebook-cell:/c%3A/Users/Guntsv/Documents/GitHub/Thai-Paraphase/03%20-%20nmt_pretrain_ParaBART.ipynb#W1sZmlsZQ%3D%3D?line=14'>15</a>\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mos\u001b[39;00m\u001b[39m,\u001b[39m \u001b[39margparse\u001b[39;00m\u001b[39m,\u001b[39m \u001b[39mpickle\u001b[39;00m\u001b[39m,\u001b[39m \u001b[39mh5py\u001b[39;00m\n\u001b[0;32m     <a href='vscode-notebook-cell:/c%3A/Users/Guntsv/Documents/GitHub/Thai-Paraphase/03%20-%20nmt_pretrain_ParaBART.ipynb#W1sZmlsZQ%3D%3D?line=15'>16</a>\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mpandas\u001b[39;00m \u001b[39mas\u001b[39;00m \u001b[39mpd\u001b[39;00m\n",
      "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'transformers.modeling_bart'"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from transformers.modeling_bart import (\n",
    "    PretrainedBartModel,  \n",
    "    LayerNorm, \n",
    "    EncoderLayer, \n",
    "    DecoderLayer, \n",
    "    LearnedPositionalEmbedding,\n",
    "    _prepare_bart_decoder_inputs,\n",
    "    _make_linear_from_emb\n",
    ")\n",
    "\n",
    "import os, argparse, pickle, h5py\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader, random_split\n",
    "\n",
    "from utils import Timer, make_path, deleaf\n",
    "from pprint import pprint\n",
    "from tqdm import tqdm\n",
    "from transformers import BartTokenizer, BartConfig, BartModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from tqdm import tqdm\n",
    "from utils import *\n",
    "\n",
    "def prepare_dataset(para_data, tokenizer, num, synt_vocab):\n",
    "    max_sent_len = 40\n",
    "    max_synt_len = 160\n",
    "    sents1 = list(para_data['sents1'][:num])\n",
    "    synts1 = list(para_data['synts1'][:num])\n",
    "    sents2 = list(para_data['sents2'][:num])\n",
    "    synts2 = list(para_data['synts2'][:num])\n",
    "\n",
    "    sent1_token_ids = torch.ones((num, max_sent_len+2), dtype=torch.long) \n",
    "    sent2_token_ids = torch.ones((num, max_sent_len+2), dtype=torch.long)    \t\t\n",
    "    synt1_token_ids = torch.ones((num, max_synt_len+2), dtype=torch.long) \n",
    "    synt2_token_ids = torch.ones((num, max_synt_len+2), dtype=torch.long)\n",
    "    synt1_bow = torch.ones((num, 74))\n",
    "    synt2_bow = torch.ones((num, 74))\n",
    "        \n",
    "    bsz = 64\n",
    "    \n",
    "    for i in tqdm(range(0, num, bsz)):\n",
    "        sent1_inputs = tokenizer([s.decode() for s in sents1[i:i+bsz]], padding='max_length', truncation=True, max_length=max_sent_len+2, return_tensors=\"pt\")\n",
    "        sent2_inputs = tokenizer([s.decode() for s in sents2[i:i+bsz]], padding='max_length', truncation=True, max_length=max_sent_len+2, return_tensors=\"pt\")\n",
    "        sent1_token_ids[i:i+bsz] = sent1_inputs['input_ids']\n",
    "        sent2_token_ids[i:i+bsz] = sent2_inputs['input_ids']\n",
    "\n",
    "    for i in tqdm(range(num)):\n",
    "        synt1 = ['<s>'] + deleaf(synts1[i].decode()) + ['</s>']\n",
    "        synt1_token_ids[i, :len(synt1)] = torch.tensor([synt_vocab[tag] for tag in synt1])[:max_synt_len+2]\n",
    "        synt2 = ['<s>'] + deleaf(synts2[i].decode()) + ['</s>']\n",
    "        synt2_token_ids[i, :len(synt2)] = torch.tensor([synt_vocab[tag] for tag in synt2])[:max_synt_len+2]\n",
    "        \n",
    "        for tag in synt1:\n",
    "            if tag != '<s>' and tag != '</s>':\n",
    "                synt1_bow[i][synt_vocab[tag]-3] += 1\n",
    "        for tag in synt2:\n",
    "            if tag != '<s>' and tag != '</s>':\n",
    "                synt2_bow[i][synt_vocab[tag]-3] += 1\n",
    "\n",
    "    synt1_bow /= synt1_bow.sum(1, keepdim=True)\n",
    "    synt2_bow /= synt2_bow.sum(1, keepdim=True)\n",
    "    \n",
    "    sum = 0\n",
    "    for i in range(num):\n",
    "        if torch.equal(synt1_bow[i], synt2_bow[i]):\n",
    "            sum += 1\n",
    "\n",
    "    return {'sent1':sent1_token_ids, 'sent2':sent2_token_ids, 'synt1': synt1_token_ids, 'synt2': synt2_token_ids,\n",
    "            'synt1bow': synt1_bow, 'synt2bow': synt2_bow}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==== loading data ====\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<KeysViewHDF5 ['sents1', 'sents2', 'synts1', 'synts2']>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import h5py, os\n",
    "print(\"==== loading data ====\")\n",
    "mrpc_set = h5py.File(os.path.join('./test_data/test_data_mrpc.h5'), 'r')\n",
    "mrpc_set.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==== preparing data ====\n"
     ]
    }
   ],
   "source": [
    "from transformers import BartTokenizer, BartConfig, BartModel\n",
    "from utils import Timer, make_path, deleaf\n",
    "import pickle\n",
    "\n",
    "print(\"==== preparing data ====\")\n",
    "make_path('./bart-base/')\n",
    "tokenizer = BartTokenizer.from_pretrained('facebook/bart-base', cache_dir='./bart-base/')\n",
    "\n",
    "with open('synt_vocab.pkl', 'rb') as f:\n",
    "    synt_vocab = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the vocabulary from the tokenizer\n",
    "vocab = tokenizer.get_vocab()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 30/30 [00:00<00:00, 40.75it/s]\n",
      "100%|██████████| 1920/1920 [00:03<00:00, 507.78it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "dict_keys(['sent1', 'sent2', 'synt1', 'synt2', 'synt1bow', 'synt2bow'])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "num = 1920\n",
    "dataset = prepare_dataset(mrpc_set, tokenizer, num, synt_vocab)\n",
    "dataset.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "test_idxs = random.sample(range(0, 1920), 1920)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loader = DataLoader(test_idxs, batch_size=16, shuffle=False)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from transformers.modeling_bart import (\n",
    "    PretrainedBartModel,  \n",
    "    LayerNorm, \n",
    "    EncoderLayer, \n",
    "    DecoderLayer, \n",
    "    LearnedPositionalEmbedding,\n",
    "    _prepare_bart_decoder_inputs,\n",
    "    _make_linear_from_emb\n",
    ")\n",
    "\n",
    "class ParaBart(PretrainedBartModel):\n",
    "    def __init__(self, config):\n",
    "        super().__init__(config)\n",
    "        \n",
    "        self.shared = nn.Embedding(config.vocab_size, config.d_model, config.pad_token_id)\n",
    "\n",
    "        self.encoder = ParaBartEncoder(config, self.shared)\n",
    "        self.decoder = ParaBartDecoder(config, self.shared)\n",
    "                \n",
    "        self.linear = nn.Linear(config.d_model, config.vocab_size)\n",
    "        \n",
    "        self.adversary = Discriminator(config)\n",
    "        \n",
    "        self.init_weights()\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        input_ids,      \n",
    "        decoder_input_ids,\n",
    "        attention_mask=None,\n",
    "        decoder_padding_mask=None,\n",
    "        encoder_outputs=None,\n",
    "        return_encoder_outputs=False,\n",
    "    ):\n",
    "        if attention_mask is None:\n",
    "            attention_mask = input_ids == self.config.pad_token_id\n",
    "        \n",
    "        if encoder_outputs is None:\n",
    "            encoder_outputs = self.encoder(input_ids, attention_mask=attention_mask)\n",
    "            \n",
    "        if return_encoder_outputs:\n",
    "            return encoder_outputs\n",
    "        \n",
    "        assert encoder_outputs is not None\n",
    "        assert decoder_input_ids is not None\n",
    "\n",
    "        decoder_input_ids = decoder_input_ids[:, :-1]\n",
    "                \n",
    "        _, decoder_padding_mask, decoder_causal_mask = _prepare_bart_decoder_inputs(\n",
    "            self.config,\n",
    "            input_ids=None,\n",
    "            decoder_input_ids=decoder_input_ids,\n",
    "            decoder_padding_mask=decoder_padding_mask,\n",
    "            causal_mask_dtype=self.shared.weight.dtype,\n",
    "        )    \n",
    "\n",
    "        attention_mask2 = torch.cat((torch.zeros(input_ids.shape[0], 1).bool().cuda(), attention_mask[:, self.config.max_sent_len+2:]), dim=1)\n",
    "           \n",
    "        # decoder\n",
    "        decoder_outputs = self.decoder(\n",
    "            decoder_input_ids,\n",
    "            torch.cat((encoder_outputs[1], encoder_outputs[0][:, self.config.max_sent_len+2:]), dim=1),           \n",
    "            decoder_padding_mask=decoder_padding_mask,\n",
    "            decoder_causal_mask=decoder_causal_mask,\n",
    "            encoder_attention_mask=attention_mask2,\n",
    "        )[0]\n",
    "        \n",
    "       \n",
    "        batch_size = decoder_outputs.shape[0]\n",
    "        outputs = self.linear(decoder_outputs.contiguous().view(-1, self.config.d_model))\n",
    "        outputs = outputs.view(batch_size, -1, self.config.vocab_size)\n",
    "        \n",
    "        # discriminator\n",
    "        for p in self.adversary.parameters():\n",
    "            p.required_grad=False\n",
    "        adv_outputs = self.adversary(encoder_outputs[1])        \n",
    "        \n",
    "        return outputs, adv_outputs\n",
    "    \n",
    "    def prepare_inputs_for_generation(self, decoder_input_ids, past, attention_mask, use_cache, **kwargs):\n",
    "        assert past is not None, \"past has to be defined for encoder_outputs\"\n",
    "\n",
    "        encoder_outputs = past[0]\n",
    "        return {\n",
    "            \"input_ids\": None,  # encoder_outputs is defined. input_ids not needed\n",
    "            \"encoder_outputs\": encoder_outputs,\n",
    "            \"decoder_input_ids\": torch.cat((decoder_input_ids, torch.zeros((decoder_input_ids.shape[0], 1), dtype=torch.long).cuda()), 1),\n",
    "            \"attention_mask\": attention_mask,\n",
    "        }\n",
    "\n",
    "    def get_encoder(self):\n",
    "        return self.encoder\n",
    "\n",
    "    def get_output_embeddings(self):\n",
    "        return _make_linear_from_emb(self.shared)\n",
    "    \n",
    "    def get_input_embeddings(self):\n",
    "        return self.shared\n",
    "    \n",
    "    @staticmethod\n",
    "    def _reorder_cache(past, beam_idx):\n",
    "        enc_out = past[0][0]\n",
    "\n",
    "        new_enc_out = enc_out.index_select(0, beam_idx)\n",
    "\n",
    "        past = ((new_enc_out, ), )\n",
    "        return past\n",
    "\n",
    "    def forward_adv(\n",
    "        self,\n",
    "        input_token_ids,      \n",
    "        attention_mask=None,\n",
    "        decoder_padding_mask=None\n",
    "    ):\n",
    "        for p in self.adversary.parameters():\n",
    "            p.required_grad=True\n",
    "        sent_embeds = self.encoder.embed(input_token_ids, attention_mask=attention_mask).detach()\n",
    "        adv_outputs = self.adversary(sent_embeds)\n",
    "\n",
    "        return adv_outputs\n",
    "\n",
    "    def generate(self, input_ids, decoder_input_ids, attention_mask=None,decoder_padding_mask=None,\n",
    "                 encoder_outputs=None,return_encoder_outputs=False, \n",
    "                 max_len = 40, sample=True, temp=0.5):\n",
    "        \n",
    "        max_targ_len = decoder_input_ids.size(1) - 2\n",
    "        batch_size   = decoder_input_ids.size(0)\n",
    "        # output index starts with <sos>\n",
    "        idxs = torch.zeros((batch_size, max_targ_len +2), dtype=torch.long).cuda()\n",
    "        idxs[:, 0] = 0\n",
    "\n",
    "        if attention_mask is None:\n",
    "            attention_mask = input_ids == self.config.pad_token_id\n",
    "        \n",
    "        if encoder_outputs is None:\n",
    "            encoder_outputs = self.encoder(input_ids, attention_mask=attention_mask)\n",
    "            \n",
    "        if return_encoder_outputs:\n",
    "            return encoder_outputs\n",
    "        \n",
    "        assert encoder_outputs is not None\n",
    "        assert decoder_input_ids is not None\n",
    "\n",
    "        # decoder_input_ids = decoder_input_ids[:, :1]\n",
    "                \n",
    "        # _, decoder_padding_mask, decoder_causal_mask = _prepare_bart_decoder_inputs(\n",
    "        #     self.config,\n",
    "        #     input_ids=None,\n",
    "        #     decoder_input_ids=decoder_input_ids,\n",
    "        #     decoder_padding_mask=decoder_padding_mask,\n",
    "        #     causal_mask_dtype=self.shared.weight.dtype,\n",
    "        # )    \n",
    "\n",
    "        attention_mask2 = torch.cat((torch.zeros(input_ids.shape[0], 1).bool().cuda(), attention_mask[:, self.config.max_sent_len+2:]), dim=1)\n",
    "           \n",
    "        # decoder\n",
    "        decoder_outputs = self.decoder(\n",
    "            idxs[:, :1],\n",
    "            torch.cat((encoder_outputs[1], encoder_outputs[0][:, self.config.max_sent_len+2:]), dim=1),           \n",
    "            decoder_padding_mask = decoder_padding_mask,\n",
    "            decoder_causal_mask = None,\n",
    "            encoder_attention_mask = attention_mask2,\n",
    "        )[0].transpose(0,1)\n",
    "        # print('decoder_outputs',decoder_outputs.shape)\n",
    "        \n",
    "        # output index starts with <sos>\n",
    "        idxs = torch.zeros((batch_size, max_targ_len+2), dtype=torch.long).to(self.device)\n",
    "        idxs[:, 0] = 1\n",
    "\n",
    "        # auto-regressively generate output\n",
    "        for i in range(1, max_targ_len+2):\n",
    "            batch_size = decoder_outputs.shape[0]\n",
    "            outputs = self.linear(decoder_outputs[-1].contiguous().view(-1, self.config.d_model))\n",
    "            # outputs = outputs.view(batch_size, -1, self.config.vocab_size)\n",
    "            # print('outputs', outputs.shape)\n",
    "            # get argmax index or sample index\n",
    "            if not sample:\n",
    "                values, idx = torch.max(outputs, 1)\n",
    "            else:\n",
    "                probs = F.softmax(outputs/temp, dim=1)\n",
    "                idx = torch.multinomial(probs, 1).squeeze(1)\n",
    "                # print('idx',idx.shape)\n",
    "            # save to output index\n",
    "            idxs[:, i] = idx   \n",
    "      \n",
    "            attention_mask2 = torch.cat((torch.zeros(input_ids.shape[0], 1).bool().cuda(), attention_mask[:, self.config.max_sent_len+2:]), dim=1)\n",
    "           \n",
    "            # decoder\n",
    "            decoder_outputs = self.decoder(\n",
    "                idxs[:, :i+1],\n",
    "                torch.cat((encoder_outputs[1], encoder_outputs[0][:, self.config.max_sent_len+2:]), dim=1),           \n",
    "                decoder_padding_mask = decoder_padding_mask,\n",
    "                decoder_causal_mask =  None,\n",
    "                encoder_attention_mask = attention_mask2,\n",
    "            )[0].transpose(0,1)\n",
    "            # print('decoder_outputs',decoder_outputs.shape)\n",
    "        \n",
    "        return idxs[:, 1:]\n",
    "\n",
    "class ParaBartEncoder(nn.Module):\n",
    "    def __init__(self, config, embed_tokens):\n",
    "        super().__init__()\n",
    "        self.config = config\n",
    "\n",
    "        self.dropout = config.dropout\n",
    "        self.embed_tokens = embed_tokens\n",
    "                \n",
    "        self.embed_synt = nn.Embedding(77, config.d_model, config.pad_token_id)       \n",
    "        self.embed_synt.weight.data.normal_(mean=0.0, std=config.init_std)\n",
    "        self.embed_synt.weight.data[config.pad_token_id].zero_()\n",
    "\n",
    "        self.embed_positions = LearnedPositionalEmbedding(\n",
    "            config.max_position_embeddings, config.d_model, config.pad_token_id, config.extra_pos_embeddings\n",
    "        )\n",
    "        \n",
    "        self.layers = nn.ModuleList([EncoderLayer(config) for _ in range(config.encoder_layers)])\n",
    "        self.synt_layers = nn.ModuleList([EncoderLayer(config) for _ in range(1)])\n",
    "\n",
    "        self.layernorm_embedding = LayerNorm(config.d_model) \n",
    "\n",
    "        self.synt_layernorm_embedding = LayerNorm(config.d_model)\n",
    "        \n",
    "        self.pooling = MeanPooling(config)\n",
    "        \n",
    "\n",
    "    def forward(self, input_ids, attention_mask): \n",
    "        \n",
    "        input_token_ids, input_synt_ids = torch.split(input_ids, [self.config.max_sent_len+2, self.config.max_synt_len+2], dim=1)\n",
    "        input_token_mask, input_synt_mask = torch.split(attention_mask, [self.config.max_sent_len+2, self.config.max_synt_len+2], dim=1)\n",
    "        \n",
    "        x = self.forward_token(input_token_ids, input_token_mask)\n",
    "        y = self.forward_synt(input_synt_ids, input_synt_mask)\n",
    "                \n",
    "        encoder_outputs = torch.cat((x,y), dim=1)\n",
    "\n",
    "        sent_embeds = self.pooling(x, input_token_ids)\n",
    "\n",
    "        return encoder_outputs, sent_embeds\n",
    "    \n",
    "    def forward_token(self, input_token_ids, attention_mask):\n",
    "        if self.training:\n",
    "            drop_mask = torch.bernoulli(self.config.word_dropout*torch.ones(input_token_ids.shape)).bool().cuda()\n",
    "            input_token_ids = input_token_ids.masked_fill(drop_mask, 50264)\n",
    "               \n",
    "        input_token_embeds = self.embed_tokens(input_token_ids) + self.embed_positions(input_token_ids)\n",
    "        x = self.layernorm_embedding(input_token_embeds)\n",
    "        x = F.dropout(x, p=self.dropout, training=self.training)\n",
    "        \n",
    "        x = x.transpose(0, 1)\n",
    "        \n",
    "        for encoder_layer in self.layers:\n",
    "            x, _ = encoder_layer(x, encoder_padding_mask=attention_mask)\n",
    "            \n",
    "        x = x.transpose(0, 1)\n",
    "        return x\n",
    "        \n",
    "    def forward_synt(self, input_synt_ids, attention_mask):\n",
    "        input_synt_embeds = self.embed_synt(input_synt_ids) + self.embed_positions(input_synt_ids)        \n",
    "        y = self.synt_layernorm_embedding(input_synt_embeds)        \n",
    "        y = F.dropout(y, p=self.dropout, training=self.training)\n",
    "\n",
    "        # B x T x C -> T x B x C\n",
    "        y = y.transpose(0, 1)\n",
    "            \n",
    "        for encoder_synt_layer in self.synt_layers:\n",
    "            y, _ = encoder_synt_layer(y, encoder_padding_mask=attention_mask)\n",
    "\n",
    "        # T x B x C -> B x T x C\n",
    "        y = y.transpose(0, 1)\n",
    "        return y\n",
    "        \n",
    "\n",
    "    def embed(self, input_token_ids, attention_mask=None, pool='mean'):\n",
    "        if attention_mask is None:\n",
    "            attention_mask = input_token_ids == self.config.pad_token_id\n",
    "            \n",
    "        x = self.forward_token(input_token_ids, attention_mask)\n",
    "        \n",
    "        sent_embeds = self.pooling(x, input_token_ids)\n",
    "        return sent_embeds\n",
    "            \n",
    "class MeanPooling(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.config = config\n",
    "        \n",
    "    def forward(self, x, input_token_ids):\n",
    "        mask = input_token_ids != self.config.pad_token_id\n",
    "        mean_mask = mask.float()/mask.float().sum(1, keepdim=True)\n",
    "        x = (x*mean_mask.unsqueeze(2)).sum(1, keepdim=True)\n",
    "        return x\n",
    "\n",
    "\n",
    "class ParaBartDecoder(nn.Module):\n",
    "    def __init__(self, config, embed_tokens):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.dropout = config.dropout\n",
    "        \n",
    "        self.embed_tokens = embed_tokens\n",
    "        \n",
    "        self.embed_positions = LearnedPositionalEmbedding(\n",
    "            config.max_position_embeddings, config.d_model, config.pad_token_id, config.extra_pos_embeddings\n",
    "        )\n",
    "        \n",
    "        self.layers = nn.ModuleList([DecoderLayer(config) for _ in range(1)]) \n",
    "        self.layernorm_embedding = LayerNorm(config.d_model)\n",
    "\n",
    "    def forward(\n",
    "        self, \n",
    "        decoder_input_ids, \n",
    "        encoder_hidden_states,  \n",
    "        decoder_padding_mask, \n",
    "        decoder_causal_mask,  \n",
    "        encoder_attention_mask\n",
    "    ):        \n",
    "\t\t\n",
    "        x = self.embed_tokens(decoder_input_ids) + self.embed_positions(decoder_input_ids)\n",
    "        x = self.layernorm_embedding(x)\n",
    "        x = F.dropout(x, p=self.dropout, training=self.training)\n",
    "\n",
    "        x = x.transpose(0, 1)\n",
    "        encoder_hidden_states = encoder_hidden_states.transpose(0, 1)\n",
    "\n",
    "        for idx, decoder_layer in enumerate(self.layers):\n",
    "            x, _, _ = decoder_layer(\n",
    "                x, \n",
    "                encoder_hidden_states,\n",
    "                encoder_attn_mask=encoder_attention_mask,\n",
    "                decoder_padding_mask=decoder_padding_mask,\n",
    "                causal_mask=decoder_causal_mask)\n",
    "\n",
    "        x = x.transpose(0, 1)\n",
    "        return x,\n",
    "    \n",
    "class Discriminator(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.sent_layernorm_embedding = LayerNorm(config.d_model, elementwise_affine=False)\n",
    "        self.adv = nn.Linear(config.d_model, 74)\n",
    "        \n",
    "    def forward(self, sent_embeds):\n",
    "        x = self.sent_layernorm_embedding(sent_embeds).squeeze(1)\n",
    "        x = self.adv(x)\n",
    "        return x\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==== loading model ====\n"
     ]
    }
   ],
   "source": [
    "print(\"==== loading model ====\")\n",
    "config = BartConfig.from_pretrained('facebook/bart-base', cache_dir='./bart-base/')\n",
    "config.word_dropout = 0.2\n",
    "config.max_sent_len = 40\n",
    "config.max_synt_len = 160\n",
    "\n",
    "bart = BartModel.from_pretrained('facebook/bart-base', cache_dir='./bart-base/')\n",
    "model = ParaBart(config)\n",
    "# model.load_state_dict(bart.state_dict(), strict=False)\n",
    "model.zero_grad()\n",
    "del bart"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "idx2word = {v:k for k,v in vocab.items()}\n",
    "span_vocab = vocab.copy()\n",
    "span_vocab.pop('<s>')\n",
    "span_vocab.pop('<pad>')\n",
    "span_vocab.pop('</s>')\n",
    "span_vocab.pop('<unk>')\n",
    "span_vocab.pop('<mask>')\n",
    "\n",
    "def reverse_bpe(sent):\n",
    "    x = []\n",
    "    cache = ''\n",
    "\n",
    "    for w in sent:\n",
    "        if w.startswith('Ġ'):\n",
    "            cache += w.replace('Ġ', '')\n",
    "            # cache = cache.strip()\n",
    "        elif cache != '':\n",
    "            x.append(cache + w)\n",
    "            cache = ''\n",
    "        else:\n",
    "            x.append(w)\n",
    "\n",
    "    return ' '.join(x)\n",
    "\n",
    "def sent2str(sent, vocab):\n",
    "    return \" \".join([idx2word[i] for i in sent if i != vocab[\"<pad>\"]])\n",
    "\n",
    "def synt2str(synt, vocab):\n",
    "    eos_pos = np.where(synt==vocab[\"</s>\"])[0]\n",
    "    eos_pos = eos_pos[0] if len(eos_pos) > 0 else len(synt)\n",
    "    return \" \".join([idx2word[i][1:-1] if i in span_vocab.values() else idx2word[i] for i in synt[:eos_pos]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==== loading model ====\n"
     ]
    }
   ],
   "source": [
    "print(\"==== loading model ====\")\n",
    "config = BartConfig.from_pretrained('facebook/bart-base', cache_dir=\"./bart-base/\")\n",
    "embed_model = ParaBart(config)\n",
    "tokenizer = BartTokenizer.from_pretrained('facebook/bart-base', cache_dir=\"./bart-base/\")\n",
    "model.load_state_dict(torch.load(os.path.join(\"./models/model.pt\"), map_location='cpu'))\n",
    "model = model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sent2str(sent,something=None):\n",
    "    return \" \".join([idx2word[i].replace('Ġ','') for i in sent if i != vocab[\"<pad>\"]]).split('</s>')[0]\n",
    "\n",
    "\n",
    "def generate(model, loader, vocab_transform):\n",
    "    #turn off dropout (and batch norm if used)\n",
    "    model.eval()\n",
    "    epoch_loss = 0\n",
    "    with open(\"./eval_via2/target_sents_parabart.txt\", \"w\") as target_sent, \\\n",
    "         open(\"./eval_via2/target_synts_parabart.txt\", \"w\") as syntax_keep, \\\n",
    "         open(\"./eval_via2/outputs_parabart.txt\", \"w\") as output_sentence,\\\n",
    "         open(\"./eval_via2/inputs_parabart.txt\", \"w\") as input_sentences:\n",
    "        with torch.no_grad():\n",
    "            for idxs in tqdm(loader):\n",
    "                \n",
    "                sent1_token_ids = dataset['sent1'][idxs].cuda()\n",
    "                synt1_token_ids = dataset['synt1'][idxs].cuda()\n",
    "                sent2_token_ids = dataset['sent2'][idxs].cuda()\n",
    "                synt2_token_ids = dataset['synt2'][idxs].cuda()\n",
    "                synt1_bow = dataset['synt1bow'][idxs].cuda()\n",
    "                synt2_bow = dataset['synt2bow'][idxs].cuda()\n",
    "\n",
    "                # generate\n",
    "                idxs = model.generate(torch.cat((sent1_token_ids, synt2_token_ids),1), sent1_token_ids, temp=0.5)\n",
    "                \n",
    "                for sent, idx,sent2, synt,synt2 in zip(sent1_token_ids.cpu().numpy(), idxs.cpu().numpy(),sent2_token_ids.cpu().numpy(), synt1_token_ids.cpu().numpy(), synt2_token_ids.cpu().numpy()):\n",
    "                    \n",
    "                    convert_idx_out =sent2str(idx, None)\n",
    "                    targetSent = sent2str(sent2[1:-1],None) \n",
    "                    inputSente = sent2str(sent[1:-1],None) \n",
    "                    input_sentences.write(inputSente+'\\n')\n",
    "                    target_sent.write(targetSent+'\\n') \n",
    "                    output_sentence.write(convert_idx_out+'\\n')\n",
    "\n",
    "                    # fp1.write(convert_sent)\n",
    "                    # fp2.write(convert_synt)\n",
    "                    # fp3.write(convert_idx)\n",
    "                    \n",
    "                    # fp1.write(sent2str(sent, vocab_transform) +'\\n')\n",
    "                    # fp2.write(synt2str(synt[1:], vocab_transform)+'\\n')\n",
    "                    # fp3.write(reverse_bpe(synt2str(idx, vocab_transform).replace(\"<pad>\", \"\")) +'\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 120/120 [00:21<00:00,  5.63it/s]\n"
     ]
    }
   ],
   "source": [
    "generate(model, test_loader, vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from nltk.translate.bleu_score import sentence_bleu\n",
    "\n",
    "def cal_bleu(hypothesis, reference, n):\n",
    "    hypothesis = hypothesis.strip().split(' ')\n",
    "    reference = reference.strip().split(' ')\n",
    "\n",
    "    if n == 0:\n",
    "        return sentence_bleu([reference], hypothesis)\n",
    "    elif n == 1:\n",
    "        weights = (1, 0, 0, 0)\n",
    "    elif n == 2:\n",
    "        weights = (0, 1, 0, 0)\n",
    "    elif n == 3:\n",
    "        weights = (0, 0, 1, 0)\n",
    "    elif n == 4:\n",
    "        weights = (0, 0, 0, 1)\n",
    "\n",
    "    return sentence_bleu([reference], hypothesis, weights=weights)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "number of examples: 1920 , 1920\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "with open('./eval_via2/target_sents_parabart.txt') as fp:\n",
    "    targs = fp.readlines()\n",
    "with open('./eval_via2/outputs_parabart.txt') as fp:\n",
    "    preds = fp.readlines()\n",
    "with open('./eval_via2/inputs_parabart.txt') as fp:\n",
    "    inps = fp.readlines()\n",
    "\n",
    "print(f\"number of examples: {len(preds)} , {len(targs)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('the terrorist activity will be awareness awareness awareness . \\n',\n",
       " 'the new name will be terrorism information awareness . \\n')"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "preds[4], targs[4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1920it [00:00, 7774.23it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BLEU: 45.794625299054395\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "scores0 = [cal_bleu(pred, targ, 1) for pred, targ in tqdm(zip(preds, targs))]\n",
    "print(f\"BLEU: {np.mean(scores0)*100.0}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1920it [00:02, 696.31it/s] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "METHEO: 40.84170222329394\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from nltk.translate import meteor\n",
    "\n",
    "def cal_meteor(hypothesis, reference):\n",
    "    hypothesis = hypothesis.strip().split(' ')\n",
    "    reference = reference.strip().split(' ')\n",
    "\n",
    "    return meteor([reference], hypothesis)   \n",
    "\n",
    "scoresm = [cal_meteor(pred, targ) for pred, targ in tqdm(zip(preds, targs))]\n",
    "print(f\"METHEO: {np.mean(scoresm)*100.0}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1920it [00:00, 2442.85it/s]\n",
      "1920it [00:00, 2505.09it/s]\n",
      "1920it [00:00, 2494.01it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Rouge-r: 45.352399615722454\n",
      "Rouge-p: 49.067626041670145\n",
      "Rouge-f: 46.99901940638212\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from rouge import Rouge\n",
    "scorer = Rouge()\n",
    "scoresR = [scorer.get_scores(pred,refs= targ)[0]['rouge-1']['r'] for pred, targ in tqdm(zip(preds, targs))]\n",
    "scoresP = [scorer.get_scores(pred,refs= targ)[0]['rouge-1']['p'] for pred, targ in tqdm(zip(preds, targs))]\n",
    "scoresF = [scorer.get_scores(pred,refs= targ)[0]['rouge-1']['f'] for pred, targ in tqdm(zip(preds, targs))]\n",
    "\n",
    "# # for ind,k in enumerate(scoresโร้ก):\n",
    "print(f\"Rouge-r: {np.mean(scoresR)*100.0}\") \n",
    "print(f\"Rouge-p: {np.mean(scoresP)*100.0}\") \n",
    "print(f\"Rouge-f: {np.mean(scoresF)*100.0}\") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}