example.ipynb

{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"name":"bert trial.ipynb","provenance":[{"file_id":"1TpiHs8rTjT-J2iZOov1hbW-UmbUK4stC","timestamp":1654495778360}],"collapsed_sections":[],"machine_shape":"hm","authorship_tag":"ABX9TyM733ktc7wNCgYxrAoPRFlb"},"kernelspec":{"name":"python3","display_name":"Python 3"},"language_info":{"name":"python"},"accelerator":"GPU"},"cells":[{"cell_type":"code","source":["!nvidia-smi"],"metadata":{"id":"ofDHO3Rip5sA"},"execution_count":null,"outputs":[]},{"cell_type":"code","execution_count":null,"metadata":{"id":"DK2DqBiXoeF9"},"outputs":[],"source":["!pip install datasets transformers==4.18.0 sentencepiece\n","!pip install torchtext==0.10.0"]},{"cell_type":"code","source":["from datasets import *\n","import argparse\n","from torch.utils.data import DataLoader\n","from torch.utils.data import Dataset\n","import torch.nn as nn\n","import torch.nn.functional as F\n","from torch.optim import Adam\n","from torch.utils.data import DataLoader\n","import torch\n","import math\n","import numpy as np\n","import tqdm\n","import random\n","import pickle\n","from collections import Counter"],"metadata":{"id":"NfXGS3VBpCXK"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# Connect with Google Drive\n","from google.colab import drive\n","drive.mount('/content/drive', force_remount = True)"],"metadata":{"id":"XW0C5igLCGSk"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["cd drive/My Drive/Colab Notebooks/BERT-pytorch"],"metadata":{"id":"uKrFjOXmDW41"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class Attention(nn.Module):\n","    \"\"\"\n","    Compute 'Scaled Dot Product Attention\n","    \"\"\"\n","\n","    def forward(self, query, key, value, mask=None, dropout=None):\n","        scores = torch.matmul(query, key.transpose(-2, -1)) \\\n","                 / math.sqrt(query.size(-1))\n","\n","        if mask is not None:\n","            scores = scores.masked_fill(mask == 0, -1e9)\n","\n","        p_attn = F.softmax(scores, dim=-1)\n","\n","        if dropout is not None:\n","            p_attn = dropout(p_attn)\n","\n","        return torch.matmul(p_attn, value), p_attn"],"metadata":{"id":"DvGt3Bzd4VI-"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class MultiHeadedAttention(nn.Module):\n","    \"\"\"\n","    Take in model size and number of heads.\n","    \"\"\"\n","\n","    def __init__(self, h, d_model, dropout=0.1):\n","        super().__init__()\n","        assert d_model % h == 0\n","\n","        # We assume d_v always equals d_k\n","        self.d_k = d_model // h\n","        self.h = h\n","\n","        self.linear_layers = nn.ModuleList([nn.Linear(d_model, d_model) for _ in range(3)])\n","        self.output_linear = nn.Linear(d_model, d_model)\n","        self.attention = Attention()\n","\n","        self.dropout = nn.Dropout(p=dropout)\n","\n","    def forward(self, query, key, value, mask=None):\n","        batch_size = query.size(0)\n","\n","        # 1) Do all the linear projections in batch from d_model => h x d_k\n","        query, key, value = [l(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)\n","                             for l, x in zip(self.linear_layers, (query, key, value))]\n","\n","        # 2) Apply attention on all the projected vectors in batch.\n","        x, attn = self.attention(query, key, value, mask=mask, dropout=self.dropout)\n","\n","        # 3) \"Concat\" using a view and apply a final linear.\n","        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k)\n","\n","        return self.output_linear(x)"],"metadata":{"id":"7d1irawD5Nlu"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class LayerNorm(nn.Module):\n","    \"Construct a layernorm module (See citation for details).\"\n","\n","    def __init__(self, features, eps=1e-6):\n","        super(LayerNorm, self).__init__()\n","        self.a_2 = nn.Parameter(torch.ones(features))\n","        self.b_2 = nn.Parameter(torch.zeros(features))\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","        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2"],"metadata":{"id":"ig6kn_X35Red"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class SublayerConnection(nn.Module):\n","    \"\"\"\n","    A residual connection followed by a layer norm.\n","    Note for code simplicity the norm is first as opposed to last.\n","    \"\"\"\n","\n","    def __init__(self, size, dropout):\n","        super(SublayerConnection, self).__init__()\n","        self.norm = LayerNorm(size)\n","        self.dropout = nn.Dropout(dropout)\n","\n","    def forward(self, x, sublayer):\n","        \"Apply residual connection to any sublayer with the same size.\"\n","        return x + self.dropout(sublayer(self.norm(x)))"],"metadata":{"id":"pDFNIA4x7mMQ"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class GELU(nn.Module):\n","    \"\"\"\n","    Paper Section 3.4, last paragraph notice that BERT used the GELU instead of RELU\n","    \"\"\"\n","\n","    def forward(self, x):\n","        return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))"],"metadata":{"id":"50Kdblji7ruV"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class PositionwiseFeedForward(nn.Module):\n","    \"Implements FFN equation.\"\n","\n","    def __init__(self, d_model, d_ff, dropout=0.1):\n","        super(PositionwiseFeedForward, self).__init__()\n","        self.w_1 = nn.Linear(d_model, d_ff)\n","        self.w_2 = nn.Linear(d_ff, d_model)\n","        self.dropout = nn.Dropout(dropout)\n","        self.activation = GELU()\n","\n","    def forward(self, x):\n","        return self.w_2(self.dropout(self.activation(self.w_1(x))))"],"metadata":{"id":"iCwTtkMk7zE-"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class TransformerBlock(nn.Module):\n","    \"\"\"\n","    Bidirectional Encoder = Transformer (self-attention)\n","    Transformer = MultiHead_Attention + Feed_Forward with sublayer connection\n","    \"\"\"\n","\n","    def __init__(self, hidden, attn_heads, feed_forward_hidden, dropout):\n","        \"\"\"\n","        :param hidden: hidden size of transformer\n","        :param attn_heads: head sizes of multi-head attention\n","        :param feed_forward_hidden: feed_forward_hidden, usually 4*hidden_size\n","        :param dropout: dropout rate\n","        \"\"\"\n","\n","        super().__init__()\n","        self.attention = MultiHeadedAttention(h=attn_heads, d_model=hidden)\n","        self.feed_forward = PositionwiseFeedForward(d_model=hidden, d_ff=feed_forward_hidden, dropout=dropout)\n","        self.input_sublayer = SublayerConnection(size=hidden, dropout=dropout)\n","        self.output_sublayer = SublayerConnection(size=hidden, dropout=dropout)\n","        self.dropout = nn.Dropout(p=dropout)\n","\n","    def forward(self, x, mask):\n","        x = self.input_sublayer(x, lambda _x: self.attention.forward(_x, _x, _x, mask=mask))\n","        x = self.output_sublayer(x, self.feed_forward)\n","        return self.dropout(x)"],"metadata":{"id":"YqQgcJBL73vT"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class TokenEmbedding(nn.Embedding):\n","    def __init__(self, vocab_size, embed_size=512):\n","        super().__init__(vocab_size, embed_size, padding_idx=0)"],"metadata":{"id":"R5ZMO5VK8PwL"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class PositionalEmbedding(nn.Module):\n","\n","    def __init__(self, d_model, max_len=512):\n","        super().__init__()\n","\n","        # Compute the positional encodings once in log space.\n","        pe = torch.zeros(max_len, d_model).float()\n","        pe.require_grad = False\n","\n","        position = torch.arange(0, max_len).float().unsqueeze(1)\n","        div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()\n","\n","        pe[:, 0::2] = torch.sin(position * div_term)\n","        pe[:, 1::2] = torch.cos(position * div_term)\n","\n","        pe = pe.unsqueeze(0)\n","        self.register_buffer('pe', pe)\n","\n","    def forward(self, x):\n","        return self.pe[:, :x.size(1)]"],"metadata":{"id":"RmXeScuo8USq"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class SegmentEmbedding(nn.Embedding):\n","    def __init__(self, embed_size=512):\n","        super().__init__(3, embed_size, padding_idx=0)"],"metadata":{"id":"n1fihAPb8VEQ"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class BERTEmbedding(nn.Module):\n","    \"\"\"\n","    BERT Embedding which is consisted with under features\n","        1. TokenEmbedding : normal embedding matrix\n","        2. PositionalEmbedding : adding positional information using sin, cos\n","        2. SegmentEmbedding : adding sentence segment info, (sent_A:1, sent_B:2)\n","        sum of all these features are output of BERTEmbedding\n","    \"\"\"\n","\n","    def __init__(self, vocab_size, embed_size, dropout=0.1):\n","        \"\"\"\n","        :param vocab_size: total vocab size\n","        :param embed_size: embedding size of token embedding\n","        :param dropout: dropout rate\n","        \"\"\"\n","        super().__init__()\n","        self.token = TokenEmbedding(vocab_size=vocab_size, embed_size=embed_size)\n","        self.position = PositionalEmbedding(d_model=self.token.embedding_dim)\n","        self.segment = SegmentEmbedding(embed_size=self.token.embedding_dim)\n","        self.dropout = nn.Dropout(p=dropout)\n","        self.embed_size = embed_size\n","\n","    def forward(self, sequence, segment_label):\n","        x = self.token(sequence) + self.position(sequence) + self.segment(segment_label)\n","        return self.dropout(x)"],"metadata":{"id":"MW2zeaOl8ZKd"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class BERT(nn.Module):\n","    \"\"\"\n","    BERT model : Bidirectional Encoder Representations from Transformers.\n","    \"\"\"\n","\n","    def __init__(self, vocab_size, hidden=768, n_layers=12, attn_heads=12, dropout=0.1):\n","        \"\"\"\n","        :param vocab_size: vocab_size of total words\n","        :param hidden: BERT model hidden size\n","        :param n_layers: numbers of Transformer blocks(layers)\n","        :param attn_heads: number of attention heads\n","        :param dropout: dropout rate\n","        \"\"\"\n","\n","        super().__init__()\n","        self.hidden = hidden\n","        self.n_layers = n_layers\n","        self.attn_heads = attn_heads\n","\n","        # paper noted they used 4*hidden_size for ff_network_hidden_size\n","        self.feed_forward_hidden = hidden * 4\n","\n","        # embedding for BERT, sum of positional, segment, token embeddings\n","        self.embedding = BERTEmbedding(vocab_size=vocab_size, embed_size=hidden)\n","\n","        # multi-layers transformer blocks, deep network\n","        self.transformer_blocks = nn.ModuleList(\n","            [TransformerBlock(hidden, attn_heads, hidden * 4, dropout) for _ in range(n_layers)])\n","\n","    def forward(self, x, segment_info):\n","        # attention masking for padded token\n","        # torch.ByteTensor([batch_size, 1, seq_len, seq_len)\n","        mask = (x > 0).unsqueeze(1).repeat(1, x.size(1), 1).unsqueeze(1)\n","\n","        # embedding the indexed sequence to sequence of vectors\n","        x = self.embedding(x, segment_info)\n","\n","        # running over multiple transformer blocks\n","        for transformer in self.transformer_blocks:\n","            x = transformer.forward(x, mask)\n","\n","        return x"],"metadata":{"id":"9L5JgqP48boJ"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class BERTLM(nn.Module):\n","    \"\"\"\n","    BERT Language Model\n","    Next Sentence Prediction Model + Masked Language Model\n","    \"\"\"\n","\n","    def __init__(self, bert: BERT, vocab_size):\n","        \"\"\"\n","        :param bert: BERT model which should be trained\n","        :param vocab_size: total vocab size for masked_lm\n","        \"\"\"\n","\n","        super().__init__()\n","        self.bert = bert\n","        self.next_sentence = NextSentencePrediction(self.bert.hidden)\n","        self.mask_lm = MaskedLanguageModel(self.bert.hidden, vocab_size)\n","\n","    def forward(self, x, segment_label):\n","        x = self.bert(x, segment_label)\n","        return self.next_sentence(x), self.mask_lm(x)\n","\n","\n","class NextSentencePrediction(nn.Module):\n","    \"\"\"\n","    2-class classification model : is_next, is_not_next\n","    \"\"\"\n","\n","    def __init__(self, hidden):\n","        \"\"\"\n","        :param hidden: BERT model output size\n","        \"\"\"\n","        super().__init__()\n","        self.linear = nn.Linear(hidden, 2)\n","        self.softmax = nn.LogSoftmax(dim=-1)\n","\n","    def forward(self, x):\n","        return self.softmax(self.linear(x[:, 0]))\n","\n","\n","class MaskedLanguageModel(nn.Module):\n","    \"\"\"\n","    predicting origin token from masked input sequence\n","    n-class classification problem, n-class = vocab_size\n","    \"\"\"\n","\n","    def __init__(self, hidden, vocab_size):\n","        \"\"\"\n","        :param hidden: output size of BERT model\n","        :param vocab_size: total vocab size\n","        \"\"\"\n","        super().__init__()\n","        self.linear = nn.Linear(hidden, vocab_size)\n","        self.softmax = nn.LogSoftmax(dim=-1)\n","\n","    def forward(self, x):\n","        return self.softmax(self.linear(x))"],"metadata":{"id":"vpprqjHZ8eUE"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class ScheduledOptim():\n","    '''A simple wrapper class for learning rate scheduling'''\n","\n","    def __init__(self, optimizer, d_model, n_warmup_steps):\n","        self._optimizer = optimizer\n","        self.n_warmup_steps = n_warmup_steps\n","        self.n_current_steps = 0\n","        self.init_lr = np.power(d_model, -0.5)\n","\n","    def step_and_update_lr(self):\n","        \"Step with the inner optimizer\"\n","        self._update_learning_rate()\n","        self._optimizer.step()\n","\n","    def zero_grad(self):\n","        \"Zero out the gradients by the inner optimizer\"\n","        self._optimizer.zero_grad()\n","\n","    def _get_lr_scale(self):\n","        return np.min([\n","            np.power(self.n_current_steps, -0.5),\n","            np.power(self.n_warmup_steps, -1.5) * self.n_current_steps])\n","\n","    def _update_learning_rate(self):\n","        ''' Learning rate scheduling per step '''\n","\n","        self.n_current_steps += 1\n","        lr = self.init_lr * self._get_lr_scale()\n","\n","        for param_group in self._optimizer.param_groups:\n","            param_group['lr'] = lr"],"metadata":{"id":"hwej77Vu9baJ"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class BERTTrainer:\n","    \"\"\"\n","    BERTTrainer make the pretrained BERT model with two LM training method.\n","        1. Masked Language Model : 3.3.1 Task #1: Masked LM\n","        2. Next Sentence prediction : 3.3.2 Task #2: Next Sentence Prediction\n","    please check the details on README.md with simple example.\n","    \"\"\"\n","\n","    def __init__(self, bert: BERT, vocab_size: int,\n","                 train_dataloader: DataLoader, test_dataloader: DataLoader = None,\n","                 lr: float = 1e-4, betas=(0.9, 0.999), weight_decay: float = 0.01, warmup_steps=10000,\n","                 with_cuda: bool = True, cuda_devices=None, log_freq: int = 10):\n","        \"\"\"\n","        :param bert: BERT model which you want to train\n","        :param vocab_size: total word vocab size\n","        :param train_dataloader: train dataset data loader\n","        :param test_dataloader: test dataset data loader [can be None]\n","        :param lr: learning rate of optimizer\n","        :param betas: Adam optimizer betas\n","        :param weight_decay: Adam optimizer weight decay param\n","        :param with_cuda: traning with cuda\n","        :param log_freq: logging frequency of the batch iteration\n","        \"\"\"\n","\n","        # Setup cuda device for BERT training, argument -c, --cuda should be true\n","        cuda_condition = torch.cuda.is_available() and with_cuda\n","        self.device = torch.device(\"cuda:0\" if cuda_condition else \"cpu\")\n","\n","        # This BERT model will be saved every epoch\n","        self.bert = bert\n","        # Initialize the BERT Language Model, with BERT model\n","        self.model = BERTLM(bert, vocab_size).to(self.device)\n","\n","        # Distributed GPU training if CUDA can detect more than 1 GPU\n","        if with_cuda and torch.cuda.device_count() > 1:\n","            print(\"Using %d GPUS for BERT\" % torch.cuda.device_count())\n","            self.model = nn.DataParallel(self.model, device_ids=cuda_devices)\n","\n","        # Setting the train and test data loader\n","        self.train_data = train_dataloader\n","        self.test_data = test_dataloader\n","\n","        # Setting the Adam optimizer with hyper-param\n","        self.optim = Adam(self.model.parameters(), lr=lr, betas=betas, weight_decay=weight_decay)\n","        self.optim_schedule = ScheduledOptim(self.optim, self.bert.hidden, n_warmup_steps=warmup_steps)\n","\n","        # Using Negative Log Likelihood Loss function for predicting the masked_token\n","        self.criterion = nn.NLLLoss(ignore_index=0)\n","\n","        self.log_freq = log_freq\n","\n","        print(\"Total Parameters:\", sum([p.nelement() for p in self.model.parameters()]))\n","\n","    def train(self, epoch):\n","        self.iteration(epoch, self.train_data)\n","\n","    def test(self, epoch):\n","        self.iteration(epoch, self.test_data, train=False)\n","\n","    def iteration(self, epoch, data_loader, train=True):\n","        \"\"\"\n","        loop over the data_loader for training or testing\n","        if on train status, backward operation is activated\n","        and also auto save the model every peoch\n","        :param epoch: current epoch index\n","        :param data_loader: torch.utils.data.DataLoader for iteration\n","        :param train: boolean value of is train or test\n","        :return: None\n","        \"\"\"\n","        str_code = \"train\" if train else \"test\"\n","\n","        # Setting the tqdm progress bar\n","        data_iter = tqdm.tqdm(enumerate(data_loader),\n","                              desc=\"EP_%s:%d\" % (str_code, epoch),\n","                              total=len(data_loader),\n","                              bar_format=\"{l_bar}{r_bar}\")\n","\n","        avg_loss = 0.0\n","        total_correct = 0\n","        total_element = 0\n","\n","        for i, data in data_iter:\n","            # 0. batch_data will be sent into the device(GPU or cpu)\n","            data = {key: value.to(self.device) for key, value in data.items()}\n","\n","            # 1. forward the next_sentence_prediction and masked_lm model\n","            next_sent_output, mask_lm_output = self.model.forward(data[\"bert_input\"], data[\"segment_label\"])\n","\n","            # 2-1. NLL(negative log likelihood) loss of is_next classification result\n","            next_loss = self.criterion(next_sent_output, data[\"is_next\"])\n","\n","            # 2-2. NLLLoss of predicting masked token word\n","            mask_loss = self.criterion(mask_lm_output.transpose(1, 2), data[\"bert_label\"])\n","\n","            # 2-3. Adding next_loss and mask_loss : 3.4 Pre-training Procedure\n","            loss = next_loss + mask_loss\n","\n","            # 3. backward and optimization only in train\n","            if train:\n","                self.optim_schedule.zero_grad()\n","                loss.backward()\n","                self.optim_schedule.step_and_update_lr()\n","\n","            # next sentence prediction accuracy\n","            correct = next_sent_output.argmax(dim=-1).eq(data[\"is_next\"]).sum().item()\n","            avg_loss += loss.item()\n","            total_correct += correct\n","            total_element += data[\"is_next\"].nelement()\n","\n","            post_fix = {\n","                \"epoch\": epoch,\n","                \"iter\": i,\n","                \"avg_loss\": avg_loss / (i + 1),\n","                \"avg_acc\": total_correct / total_element * 100,\n","                \"loss\": loss.item()\n","            }\n","\n","            if i % self.log_freq == 0:\n","                data_iter.write(str(post_fix))\n","\n","        print(\"EP%d_%s, avg_loss=\" % (epoch, str_code), avg_loss / len(data_iter), \"total_acc=\",\n","              total_correct * 100.0 / total_element)\n","\n","    def save(self, epoch, file_path=\"output/bert_trained.model\"):\n","        \"\"\"\n","        Saving the current BERT model on file_path\n","        :param epoch: current epoch number\n","        :param file_path: model output path which gonna be file_path+\"ep%d\" % epoch\n","        :return: final_output_path\n","        \"\"\"\n","        output_path = file_path + \".ep%d\" % epoch\n","        torch.save(self.bert.cpu(), output_path)\n","        self.bert.to(self.device)\n","        print(\"EP:%d Model Saved on:\" % epoch, output_path)\n","        return output_path"],"metadata":{"id":"chkPTUPz-BRc"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class BERTDataset(Dataset):\n","    def __init__(self, corpus_path, vocab, seq_len, encoding=\"utf-8\", corpus_lines=None, on_memory=True):\n","        self.vocab = vocab\n","        self.seq_len = seq_len\n","\n","        self.on_memory = on_memory\n","        self.corpus_lines = corpus_lines\n","        self.corpus_path = corpus_path\n","        self.encoding = encoding\n","\n","        with open(corpus_path, \"r\", encoding=encoding) as f:\n","            if self.corpus_lines is None and not on_memory:\n","                for _ in tqdm.tqdm(f, desc=\"Loading Dataset\", total=corpus_lines):\n","                    self.corpus_lines += 1\n","\n","            if on_memory:\n","                self.lines = [line[:-1].split(\"\\t\")\n","                              for line in tqdm.tqdm(f, desc=\"Loading Dataset\", total=corpus_lines)]\n","                self.corpus_lines = len(self.lines)\n","\n","        if not on_memory:\n","            self.file = open(corpus_path, \"r\", encoding=encoding)\n","            self.random_file = open(corpus_path, \"r\", encoding=encoding)\n","\n","            for _ in range(random.randint(self.corpus_lines if self.corpus_lines < 1000 else 1000)):\n","                self.random_file.__next__()\n","\n","    def __len__(self):\n","        return self.corpus_lines\n","\n","    def __getitem__(self, item):\n","        t1, t2, is_next_label = self.random_sent(item)\n","        t1_random, t1_label = self.random_word(t1)\n","        t2_random, t2_label = self.random_word(t2)\n","\n","        # [CLS] tag = SOS tag, [SEP] tag = EOS tag\n","        t1 = [self.vocab.sos_index] + t1_random + [self.vocab.eos_index]\n","        t2 = t2_random + [self.vocab.eos_index]\n","\n","        t1_label = [self.vocab.pad_index] + t1_label + [self.vocab.pad_index]\n","        t2_label = t2_label + [self.vocab.pad_index]\n","\n","        segment_label = ([1 for _ in range(len(t1))] + [2 for _ in range(len(t2))])[:self.seq_len]\n","        bert_input = (t1 + t2)[:self.seq_len]\n","        bert_label = (t1_label + t2_label)[:self.seq_len]\n","\n","        padding = [self.vocab.pad_index for _ in range(self.seq_len - len(bert_input))]\n","        bert_input.extend(padding), bert_label.extend(padding), segment_label.extend(padding)\n","\n","        output = {\"bert_input\": bert_input,\n","                  \"bert_label\": bert_label,\n","                  \"segment_label\": segment_label,\n","                  \"is_next\": is_next_label}\n","\n","        return {key: torch.tensor(value) for key, value in output.items()}\n","\n","    def random_word(self, sentence):\n","        tokens = sentence.split()\n","        output_label = []\n","\n","        for i, token in enumerate(tokens):\n","            prob = random.random()\n","            if prob < 0.15:\n","                prob /= 0.15\n","\n","                # 80% randomly change token to mask token\n","                if prob < 0.8:\n","                    tokens[i] = self.vocab.mask_index\n","\n","                # 10% randomly change token to random token\n","                elif prob < 0.9:\n","                    tokens[i] = random.randrange(len(self.vocab))\n","\n","                # 10% randomly change token to current token\n","                else:\n","                    tokens[i] = self.vocab.stoi.get(token, self.vocab.unk_index)\n","\n","                output_label.append(self.vocab.stoi.get(token, self.vocab.unk_index))\n","\n","            else:\n","                tokens[i] = self.vocab.stoi.get(token, self.vocab.unk_index)\n","                output_label.append(0)\n","\n","        return tokens, output_label\n","\n","    def random_sent(self, index):\n","        t1, t2 = self.get_corpus_line(index)\n","\n","        # output_text, label(isNotNext:0, isNext:1)\n","        if random.random() > 0.5:\n","            return t1, t2, 1\n","        else:\n","            return t1, self.get_random_line(), 0\n","\n","    def get_corpus_line(self, item):\n","        if self.on_memory:\n","            return self.lines[item][0], self.lines[item][1]\n","        else:\n","            line = self.file.__next__()\n","            if line is None:\n","                self.file.close()\n","                self.file = open(self.corpus_path, \"r\", encoding=self.encoding)\n","                line = self.file.__next__()\n","\n","            t1, t2 = line[:-1].split(\"\\t\")\n","            return t1, t2\n","\n","    def get_random_line(self):\n","        if self.on_memory:\n","            return self.lines[random.randrange(len(self.lines))][1]\n","\n","        line = self.file.__next__()\n","        if line is None:\n","            self.file.close()\n","            self.file = open(self.corpus_path, \"r\", encoding=self.encoding)\n","            for _ in range(random.randint(self.corpus_lines if self.corpus_lines < 1000 else 1000)):\n","                self.random_file.__next__()\n","            line = self.random_file.__next__()\n","        return line[:-1].split(\"\\t\")[1]"],"metadata":{"id":"WNvHxnIe-9Xh"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class TorchVocab(object):\n","    \"\"\"Defines a vocabulary object that will be used to numericalize a field.\n","    Attributes:\n","        freqs: A collections.Counter object holding the frequencies of tokens\n","            in the data used to build the Vocab.\n","        stoi: A collections.defaultdict instance mapping token strings to\n","            numerical identifiers.\n","        itos: A list of token strings indexed by their numerical identifiers.\n","    \"\"\"\n","\n","    def __init__(self, counter, max_size=None, min_freq=1, specials=['<pad>', '<oov>'],\n","                 vectors=None, unk_init=None, vectors_cache=None):\n","        \"\"\"Create a Vocab object from a collections.Counter.\n","        Arguments:\n","            counter: collections.Counter object holding the frequencies of\n","                each value found in the data.\n","            max_size: The maximum size of the vocabulary, or None for no\n","                maximum. Default: None.\n","            min_freq: The minimum frequency needed to include a token in the\n","                vocabulary. Values less than 1 will be set to 1. Default: 1.\n","            specials: The list of special tokens (e.g., padding or eos) that\n","                will be prepended to the vocabulary in addition to an <unk>\n","                token. Default: ['<pad>']\n","            vectors: One of either the available pretrained vectors\n","                or custom pretrained vectors (see Vocab.load_vectors);\n","                or a list of aforementioned vectors\n","            unk_init (callback): by default, initialize out-of-vocabulary word vectors\n","                to zero vectors; can be any function that takes in a Tensor and\n","                returns a Tensor of the same size. Default: torch.Tensor.zero_\n","            vectors_cache: directory for cached vectors. Default: '.vector_cache'\n","        \"\"\"\n","        self.freqs = counter\n","        counter = counter.copy()\n","        min_freq = max(min_freq, 1)\n","\n","        self.itos = list(specials)\n","        # frequencies of special tokens are not counted when building vocabulary\n","        # in frequency order\n","        for tok in specials:\n","            del counter[tok]\n","\n","        max_size = None if max_size is None else max_size + len(self.itos)\n","\n","        # sort by frequency, then alphabetically\n","        words_and_frequencies = sorted(counter.items(), key=lambda tup: tup[0])\n","        words_and_frequencies.sort(key=lambda tup: tup[1], reverse=True)\n","\n","        for word, freq in words_and_frequencies:\n","            if freq < min_freq or len(self.itos) == max_size:\n","                break\n","            self.itos.append(word)\n","\n","        # stoi is simply a reverse dict for itos\n","        self.stoi = {tok: i for i, tok in enumerate(self.itos)}\n","\n","        self.vectors = None\n","        if vectors is not None:\n","            self.load_vectors(vectors, unk_init=unk_init, cache=vectors_cache)\n","        else:\n","            assert unk_init is None and vectors_cache is None\n","\n","    def __eq__(self, other):\n","        if self.freqs != other.freqs:\n","            return False\n","        if self.stoi != other.stoi:\n","            return False\n","        if self.itos != other.itos:\n","            return False\n","        if self.vectors != other.vectors:\n","            return False\n","        return True\n","\n","    def __len__(self):\n","        return len(self.itos)\n","\n","    def vocab_rerank(self):\n","        self.stoi = {word: i for i, word in enumerate(self.itos)}\n","\n","    def extend(self, v, sort=False):\n","        words = sorted(v.itos) if sort else v.itos\n","        for w in words:\n","            if w not in self.stoi:\n","                self.itos.append(w)\n","                self.stoi[w] = len(self.itos) - 1\n","\n","\n","class Vocab(TorchVocab):\n","    def __init__(self, counter, max_size=None, min_freq=1):\n","        self.pad_index = 0\n","        self.unk_index = 1\n","        self.eos_index = 2\n","        self.sos_index = 3\n","        self.mask_index = 4\n","        super().__init__(counter, specials=[\"<pad>\", \"<unk>\", \"<eos>\", \"<sos>\", \"<mask>\"],\n","                         max_size=max_size, min_freq=min_freq)\n","\n","    def to_seq(self, sentece, seq_len, with_eos=False, with_sos=False) -> list:\n","        pass\n","\n","    def from_seq(self, seq, join=False, with_pad=False):\n","        pass\n","\n","    @staticmethod\n","    def load_vocab(vocab_path: str) -> 'Vocab':\n","        with open(vocab_path, \"rb\") as f:\n","            return pickle.load(f)\n","\n","    def save_vocab(self, vocab_path):\n","        with open(vocab_path, \"wb\") as f:\n","            pickle.dump(self, f)\n","\n","\n","# Building Vocab with text files\n","class WordVocab(Vocab):\n","    def __init__(self, texts, max_size=None, min_freq=1):\n","        print(\"Building Vocab\")\n","        counter = Counter()\n","        for line in tqdm.tqdm(texts):\n","            if isinstance(line, list):\n","                words = line\n","            else:\n","                words = line.replace(\"\\n\", \"\").replace(\"\\t\", \"\").split()\n","\n","            for word in words:\n","                counter[word] += 1\n","        super().__init__(counter, max_size=max_size, min_freq=min_freq)\n","\n","    def to_seq(self, sentence, seq_len=None, with_eos=False, with_sos=False, with_len=False):\n","        if isinstance(sentence, str):\n","            sentence = sentence.split()\n","\n","        seq = [self.stoi.get(word, self.unk_index) for word in sentence]\n","\n","        if with_eos:\n","            seq += [self.eos_index]  # this would be index 1\n","        if with_sos:\n","            seq = [self.sos_index] + seq\n","\n","        origin_seq_len = len(seq)\n","\n","        if seq_len is None:\n","            pass\n","        elif len(seq) <= seq_len:\n","            seq += [self.pad_index for _ in range(seq_len - len(seq))]\n","        else:\n","            seq = seq[:seq_len]\n","\n","        return (seq, origin_seq_len) if with_len else seq\n","\n","    def from_seq(self, seq, join=False, with_pad=False):\n","        words = [self.itos[idx]\n","                 if idx < len(self.itos)\n","                 else \"<%d>\" % idx\n","                 for idx in seq\n","                 if not with_pad or idx != self.pad_index]\n","\n","        return \" \".join(words) if join else words\n","\n","    @staticmethod\n","    def load_vocab(vocab_path: str) -> 'WordVocab':\n","        with open(vocab_path, \"rb\") as f:\n","            return pickle.load(f)\n","\n","\n","def build():\n","    import argparse\n","\n","    parser = argparse.ArgumentParser()\n","    parser.add_argument(\"-c\", \"--corpus_path\", required=True, type=str)\n","    parser.add_argument(\"-o\", \"--output_path\", required=True, type=str)\n","    parser.add_argument(\"-s\", \"--vocab_size\", type=int, default=None)\n","    parser.add_argument(\"-e\", \"--encoding\", type=str, default=\"utf-8\")\n","    parser.add_argument(\"-m\", \"--min_freq\", type=int, default=1)\n","    args = parser.parse_args()\n","\n","    with open(args.corpus_path, \"r\", encoding=args.encoding) as f:\n","        vocab = WordVocab(f, max_size=args.vocab_size, min_freq=args.min_freq)\n","\n","    print(\"VOCAB SIZE:\", len(vocab))\n","    vocab.save_vocab(args.output_path)"],"metadata":{"id":"p8F4s7Rg_j63"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class TorchVocab(object):\n","    \"\"\"Defines a vocabulary object that will be used to numericalize a field.\n","    Attributes:\n","        freqs: A collections.Counter object holding the frequencies of tokens\n","            in the data used to build the Vocab.\n","        stoi: A collections.defaultdict instance mapping token strings to\n","            numerical identifiers.\n","        itos: A list of token strings indexed by their numerical identifiers.\n","    \"\"\"\n","\n","    def __init__(self, counter, max_size=None, min_freq=1, specials=['<pad>', '<oov>'],\n","                 vectors=None, unk_init=None, vectors_cache=None):\n","        \"\"\"Create a Vocab object from a collections.Counter.\n","        Arguments:\n","            counter: collections.Counter object holding the frequencies of\n","                each value found in the data.\n","            max_size: The maximum size of the vocabulary, or None for no\n","                maximum. Default: None.\n","            min_freq: The minimum frequency needed to include a token in the\n","                vocabulary. Values less than 1 will be set to 1. Default: 1.\n","            specials: The list of special tokens (e.g., padding or eos) that\n","                will be prepended to the vocabulary in addition to an <unk>\n","                token. Default: ['<pad>']\n","            vectors: One of either the available pretrained vectors\n","                or custom pretrained vectors (see Vocab.load_vectors);\n","                or a list of aforementioned vectors\n","            unk_init (callback): by default, initialize out-of-vocabulary word vectors\n","                to zero vectors; can be any function that takes in a Tensor and\n","                returns a Tensor of the same size. Default: torch.Tensor.zero_\n","            vectors_cache: directory for cached vectors. Default: '.vector_cache'\n","        \"\"\"\n","        self.freqs = counter\n","        counter = counter.copy()\n","        min_freq = max(min_freq, 1)\n","\n","        self.itos = list(specials)\n","        # frequencies of special tokens are not counted when building vocabulary\n","        # in frequency order\n","        for tok in specials:\n","            del counter[tok]\n","\n","        max_size = None if max_size is None else max_size + len(self.itos)\n","\n","        # sort by frequency, then alphabetically\n","        words_and_frequencies = sorted(counter.items(), key=lambda tup: tup[0])\n","        words_and_frequencies.sort(key=lambda tup: tup[1], reverse=True)\n","\n","        for word, freq in words_and_frequencies:\n","            if freq < min_freq or len(self.itos) == max_size:\n","                break\n","            self.itos.append(word)\n","\n","        # stoi is simply a reverse dict for itos\n","        self.stoi = {tok: i for i, tok in enumerate(self.itos)}\n","\n","        self.vectors = None\n","        if vectors is not None:\n","            self.load_vectors(vectors, unk_init=unk_init, cache=vectors_cache)\n","        else:\n","            assert unk_init is None and vectors_cache is None\n","\n","    def __eq__(self, other):\n","        if self.freqs != other.freqs:\n","            return False\n","        if self.stoi != other.stoi:\n","            return False\n","        if self.itos != other.itos:\n","            return False\n","        if self.vectors != other.vectors:\n","            return False\n","        return True\n","\n","    def __len__(self):\n","        return len(self.itos)\n","\n","    def vocab_rerank(self):\n","        self.stoi = {word: i for i, word in enumerate(self.itos)}\n","\n","    def extend(self, v, sort=False):\n","        words = sorted(v.itos) if sort else v.itos\n","        for w in words:\n","            if w not in self.stoi:\n","                self.itos.append(w)\n","                self.stoi[w] = len(self.itos) - 1\n","\n","\n","class Vocab(TorchVocab):\n","    def __init__(self, counter, max_size=None, min_freq=1):\n","        self.pad_index = 0\n","        self.unk_index = 1\n","        self.eos_index = 2\n","        self.sos_index = 3\n","        self.mask_index = 4\n","        super().__init__(counter, specials=[\"<pad>\", \"<unk>\", \"<eos>\", \"<sos>\", \"<mask>\"],\n","                         max_size=max_size, min_freq=min_freq)\n","\n","    def to_seq(self, sentece, seq_len, with_eos=False, with_sos=False) -> list:\n","        pass\n","\n","    def from_seq(self, seq, join=False, with_pad=False):\n","        pass\n","\n","    @staticmethod\n","    def load_vocab(vocab_path: str) -> 'Vocab':\n","        with open(vocab_path, \"rb\") as f:\n","            return pickle.load(f)\n","\n","    def save_vocab(self, vocab_path):\n","        with open(vocab_path, \"wb\") as f:\n","            pickle.dump(self, f)\n","\n","\n","# Building Vocab with text files\n","class WordVocab(Vocab):\n","    def __init__(self, texts, max_size=None, min_freq=1):\n","        print(\"Building Vocab\")\n","        counter = Counter()\n","        for line in tqdm.tqdm(texts):\n","            if isinstance(line, list):\n","                words = line\n","            else:\n","                words = line.replace(\"\\n\", \"\").replace(\"\\t\", \"\").split()\n","\n","            for word in words:\n","                counter[word] += 1\n","        super().__init__(counter, max_size=max_size, min_freq=min_freq)\n","\n","    def to_seq(self, sentence, seq_len=None, with_eos=False, with_sos=False, with_len=False):\n","        if isinstance(sentence, str):\n","            sentence = sentence.split()\n","\n","        seq = [self.stoi.get(word, self.unk_index) for word in sentence]\n","\n","        if with_eos:\n","            seq += [self.eos_index]  # this would be index 1\n","        if with_sos:\n","            seq = [self.sos_index] + seq\n","\n","        origin_seq_len = len(seq)\n","\n","        if seq_len is None:\n","            pass\n","        elif len(seq) <= seq_len:\n","            seq += [self.pad_index for _ in range(seq_len - len(seq))]\n","        else:\n","            seq = seq[:seq_len]\n","\n","        return (seq, origin_seq_len) if with_len else seq\n","\n","    def from_seq(self, seq, join=False, with_pad=False):\n","        words = [self.itos[idx]\n","                 if idx < len(self.itos)\n","                 else \"<%d>\" % idx\n","                 for idx in seq\n","                 if not with_pad or idx != self.pad_index]\n","\n","        return \" \".join(words) if join else words\n","\n","    @staticmethod\n","    def load_vocab(vocab_path: str) -> 'WordVocab':\n","        with open(vocab_path, \"rb\") as f:\n","            return pickle.load(f)"],"metadata":{"id":"yWsr6Jgw_zlG"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["parser = argparse.ArgumentParser()\n","parser.add_argument(\"-c\", \"--corpus_path\", required=True, type=str)\n","parser.add_argument(\"-o\", \"--output_path\", required=True, type=str)\n","parser.add_argument(\"-s\", \"--vocab_size\", type=int, default=None)\n","parser.add_argument(\"-e\", \"--encoding\", type=str, default=\"utf-8\")\n","parser.add_argument(\"-m\", \"--min_freq\", type=int, default=1)\n","args = parser.parse_args(['-c', './Data/corpus.small', '-o', './Data/vocab_train.txt'])\n","\n","with open(args.corpus_path, \"r\", encoding=args.encoding) as f:\n","    vocab = WordVocab(f, max_size=args.vocab_size, min_freq=args.min_freq)\n","\n","print(\"VOCAB SIZE:\", len(vocab))\n","vocab.save_vocab(args.output_path)"],"metadata":{"id":"Lmd0ZariDvnH"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["parser = argparse.ArgumentParser()\n","\n","parser.add_argument(\"-c\", \"--train_dataset\", required=True, type=str, help=\"train dataset for train bert\")\n","parser.add_argument(\"-t\", \"--test_dataset\", type=str, default=None, help=\"test set for evaluate train set\")\n","parser.add_argument(\"-v\", \"--vocab_path\", required=True, type=str, help=\"built vocab model path with bert-vocab\")\n","parser.add_argument(\"-o\", \"--output_path\", required=True, type=str, help=\"ex)output/bert.model\")\n","\n","parser.add_argument(\"-hs\", \"--hidden\", type=int, default=256, help=\"hidden size of transformer model\")\n","parser.add_argument(\"-l\", \"--layers\", type=int, default=8, help=\"number of layers\")\n","parser.add_argument(\"-a\", \"--attn_heads\", type=int, default=8, help=\"number of attention heads\")\n","parser.add_argument(\"-s\", \"--seq_len\", type=int, default=20, help=\"maximum sequence len\")\n","\n","parser.add_argument(\"-b\", \"--batch_size\", type=int, default=64, help=\"number of batch_size\")\n","parser.add_argument(\"-e\", \"--epochs\", type=int, default=10, help=\"number of epochs\")\n","parser.add_argument(\"-w\", \"--num_workers\", type=int, default=10, help=\"dataloader worker size\")\n","\n","parser.add_argument(\"--with_cuda\", type=bool, default=True, help=\"training with CUDA: true, or false\")\n","parser.add_argument(\"--log_freq\", type=int, default=10, help=\"printing loss every n iter: setting n\")\n","parser.add_argument(\"--corpus_lines\", type=int, default=None, help=\"total number of lines in corpus\")\n","parser.add_argument(\"--cuda_devices\", type=int, nargs='+', default=None, help=\"CUDA device ids\")\n","parser.add_argument(\"--on_memory\", type=bool, default=True, help=\"Loading on memory: true or false\")\n","\n","parser.add_argument(\"--lr\", type=float, default=1e-3, help=\"learning rate of adam\")\n","parser.add_argument(\"--adam_weight_decay\", type=float, default=0.01, help=\"weight_decay of adam\")\n","parser.add_argument(\"--adam_beta1\", type=float, default=0.9, help=\"adam first beta value\")\n","parser.add_argument(\"--adam_beta2\", type=float, default=0.999, help=\"adam first beta value\")\n","\n","args = parser.parse_args(['-c', 'Data/corpus.small', '-v', 'Data/vocab_train.txt', '-o', 'output/bert.model'])\n","\n","print(\"Loading Vocab\", args.vocab_path)\n","vocab = WordVocab.load_vocab(args.vocab_path)\n","print(\"Vocab Size: \", len(vocab))\n","\n","print(\"Loading Train Dataset\", args.train_dataset)\n","train_dataset = BERTDataset(args.train_dataset, vocab, seq_len=args.seq_len,\n","                            corpus_lines=args.corpus_lines, on_memory=args.on_memory)\n","\n","print(\"Loading Test Dataset\", args.test_dataset)\n","test_dataset = BERTDataset(args.test_dataset, vocab, seq_len=args.seq_len, on_memory=args.on_memory) \\\n","    if args.test_dataset is not None else None\n","\n","print(\"Creating Dataloader\")\n","train_data_loader = DataLoader(train_dataset, batch_size=args.batch_size, num_workers=args.num_workers)\n","test_data_loader = DataLoader(test_dataset, batch_size=args.batch_size, num_workers=args.num_workers) \\\n","    if test_dataset is not None else None\n","\n","print(\"Building BERT model\")\n","bert = BERT(len(vocab), hidden=args.hidden, n_layers=args.layers, attn_heads=args.attn_heads)\n","\n","print(\"Creating BERT Trainer\")\n","trainer = BERTTrainer(bert, len(vocab), train_dataloader=train_data_loader, test_dataloader=test_data_loader,\n","                      lr=args.lr, betas=(args.adam_beta1, args.adam_beta2), weight_decay=args.adam_weight_decay,\n","                      with_cuda=args.with_cuda, cuda_devices=args.cuda_devices, log_freq=args.log_freq)\n","\n","print(\"Training Start\")\n","for epoch in range(args.epochs):\n","    trainer.train(epoch)\n","    trainer.save(epoch, args.output_path)\n","\n","    if test_data_loader is not None:\n","        trainer.test(epoch)"],"metadata":{"id":"DTvizUOuFUeW"},"execution_count":null,"outputs":[]}]}