Upload 7 files

bert_dataset.py

@@ -0,0 +1,115 @@
+import torch
+import random
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+from transformers import BertTokenizer
+from data import get_data
+import itertools
+
+tokenizer = BertTokenizer.from_pretrained("bert-it-1/bert-it-vocab.txt")
+
+
+class BERTDataset(Dataset):
+
+    def __init__(self, tokenizer: BertTokenizer=tokenizer, data_pair: list=get_data('datasets/movie_conversations.txt', "datasets/movie_lines.txt"), seq_len: int=128) -> None:
+        super().__init__()
+
+        self.tokenizer = tokenizer
+        self.seq_len = seq_len
+        self.corpus_lines = len(data_pair)
+        self.lines = data_pair
+
+    def __len__(self):
+        return self.corpus_lines
+
+    def __getitem__(self, item):
+
+        # Step 1: get random sentence pair, either negative or positive (saved as is_next_label)
+        t1, t2, is_next_label = self.get_sent(item)
+
+        # Step 2: replace random words in sentence with mask / random words
+        t1_random, t1_label = self.random_word(t1)
+        t2_random, t2_label = self.random_word(t2)
+
+        # Step 3: Adding CLS and SEP tokens to the start and end of sentences
+         # Adding PAD token for labels
+        t1 = [self.tokenizer.vocab['[CLS]']] + t1_random + [self.tokenizer.vocab['[SEP]']]
+        t2 = t2_random + [self.tokenizer.vocab['[SEP]']]
+        t1_label = [self.tokenizer.vocab['[PAD]']] + t1_label + [self.tokenizer.vocab['[PAD]']]
+        t2_label = t2_label + [self.tokenizer.vocab['[PAD]']]
+
+        # Step 4: combine sentence 1 and 2 as one input
+        # adding PAD tokens to make the sentence same length as seq_len
+        segment_label = ([1 for _ in range(len(t1))] + [2 for _ in range(len(t2))])[:self.seq_len]
+        bert_input = (t1 + t2)[:self.seq_len]
+        bert_label = (t1_label + t2_label)[:self.seq_len]
+        padding = [self.tokenizer.vocab['[PAD]'] for _ in range(self.seq_len - len(bert_input))]
+        bert_input.extend(padding), bert_label.extend(padding), segment_label.extend(padding)
+
+        output = {"bert_input": bert_input,
+                  "bert_label": bert_label,
+                  "segment_label": segment_label,
+                  "is_next": is_next_label}
+
+        return {key: torch.tensor(value) for key, value in output.items()}
+
+    def random_word(self, sentence):
+        tokens = sentence.split()
+        output_label = []
+        output = []
+
+        # 15% of the tokens would be replaced
+        for i, token in enumerate(tokens):
+            prob = random.random()
+
+            # remove cls and sep token
+            token_id = self.tokenizer(token)['input_ids'][1:-1]
+
+            if prob < 0.15:
+                prob /= 0.15
+
+                # 80% chance change token to mask token
+                if prob < 0.8:
+                    for i in range(len(token_id)):
+                        output.append(self.tokenizer.vocab['[MASK]'])
+
+                # 10% chance change token to random token
+                elif prob < 0.9:
+                    for i in range(len(token_id)):
+                        output.append(random.randrange(len(self.tokenizer.vocab)))
+
+                # 10% chance change token to current token
+                else:
+                    output.append(token_id)
+
+                output_label.append(token_id)
+
+            else:
+                output.append(token_id)
+                for i in range(len(token_id)):
+                    output_label.append(0)
+
+        # flattening
+        output = list(itertools.chain(*[[x] if not isinstance(x, list) else x for x in output]))
+        output_label = list(itertools.chain(*[[x] if not isinstance(x, list) else x for x in output_label]))
+        assert len(output) == len(output_label)
+        return output, output_label
+
+    def get_sent(self, index):
+        '''return random sentence pair'''
+        t1, t2 = self.get_corpus_line(index)
+
+        # negative or positive pair, for next sentence prediction
+        if random.random() > 0.5:
+            return t1, t2, 1
+        else:
+            return t1, self.get_random_line(), 0
+
+    def get_corpus_line(self, item):
+        '''return sentence pair'''
+        return self.lines[item][0], self.lines[item][1]
+
+    def get_random_line(self):
+        '''return random single sentence'''
+        return self.lines[random.randrange(len(self.lines))][1]

bert_model.py

@@ -0,0 +1,250 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+
+class PositionalEmbedding(torch.nn.Module):
+
+    def __init__(self, d_model, max_len=128):
+        super().__init__()
+
+        # Compute the positional encodings once in log space.
+        pe = torch.zeros(max_len, d_model).float()
+        pe.require_grad = False
+
+        for pos in range(max_len):   
+            # for each dimension of the each position
+            for i in range(0, d_model, 2):   
+                pe[pos, i] = math.sin(pos / (10000 ** ((2 * i)/d_model)))
+                pe[pos, i + 1] = math.cos(pos / (10000 ** ((2 * (i + 1))/d_model)))
+
+        # include the batch size
+        self.pe = pe.unsqueeze(0)   
+        # self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        return self.pe
+
+class BERTEmbedding(torch.nn.Module):
+    """
+    BERT Embedding which is consisted with under features
+        1. TokenEmbedding : normal embedding matrix
+        2. PositionalEmbedding : adding positional information using sin, cos
+        2. SegmentEmbedding : adding sentence segment info, (sent_A:1, sent_B:2)
+        sum of all these features are output of BERTEmbedding
+    """
+
+    def __init__(self, vocab_size, embed_size, seq_len=64, dropout=0.1):
+        """
+        :param vocab_size: total vocab size
+        :param embed_size: embedding size of token embedding
+        :param dropout: dropout rate
+        """
+
+        super().__init__()
+        self.embed_size = embed_size
+        # (m, seq_len) --> (m, seq_len, embed_size)
+        # padding_idx is not updated during training, remains as fixed pad (0)
+        self.token = torch.nn.Embedding(vocab_size, embed_size, padding_idx=0)
+        self.segment = torch.nn.Embedding(3, embed_size, padding_idx=0)
+        self.position = PositionalEmbedding(d_model=embed_size, max_len=seq_len)
+        self.dropout = torch.nn.Dropout(p=dropout)
+       
+    def forward(self, sequence, segment_label):
+        x = self.token(sequence) + self.position(sequence) + self.segment(segment_label)
+        return self.dropout(x)
+    
+### attention layers
+class MultiHeadedAttention(torch.nn.Module):
+    
+    def __init__(self, heads, d_model, dropout=0.1):
+        super(MultiHeadedAttention, self).__init__()
+        
+        assert d_model % heads == 0
+        self.d_k = d_model // heads
+        self.heads = heads
+        self.dropout = torch.nn.Dropout(dropout)
+
+        self.query = torch.nn.Linear(d_model, d_model)
+        self.key = torch.nn.Linear(d_model, d_model)
+        self.value = torch.nn.Linear(d_model, d_model)
+        self.output_linear = torch.nn.Linear(d_model, d_model)
+        
+    def forward(self, query, key, value, mask):
+        """
+        query, key, value of shape: (batch_size, max_len, d_model)
+        mask of shape: (batch_size, 1, 1, max_words)
+        """
+        # (batch_size, max_len, d_model)
+        query = self.query(query)
+        key = self.key(key)        
+        value = self.value(value)   
+        
+        # (batch_size, max_len, d_model) --> (batch_size, max_len, h, d_k) --> (batch_size, h, max_len, d_k)
+        query = query.view(query.shape[0], -1, self.heads, self.d_k).permute(0, 2, 1, 3)   
+        key = key.view(key.shape[0], -1, self.heads, self.d_k).permute(0, 2, 1, 3)  
+        value = value.view(value.shape[0], -1, self.heads, self.d_k).permute(0, 2, 1, 3)  
+        
+        # (batch_size, h, max_len, d_k) matmul (batch_size, h, d_k, max_len) --> (batch_size, h, max_len, max_len)
+        scores = torch.matmul(query, key.permute(0, 1, 3, 2)) / math.sqrt(query.size(-1))
+
+        # fill 0 mask with super small number so it wont affect the softmax weight
+        # (batch_size, h, max_len, max_len)
+        scores = scores.masked_fill(mask == 0, -1e9)    
+
+        # (batch_size, h, max_len, max_len)
+        # softmax to put attention weight for all non-pad tokens
+        # max_len X max_len matrix of attention
+        weights = F.softmax(scores, dim=-1)           
+        weights = self.dropout(weights)
+
+        # (batch_size, h, max_len, max_len) matmul (batch_size, h, max_len, d_k) --> (batch_size, h, max_len, d_k)
+        context = torch.matmul(weights, value)
+
+        # (batch_size, h, max_len, d_k) --> (batch_size, max_len, h, d_k) --> (batch_size, max_len, d_model)
+        context = context.permute(0, 2, 1, 3).contiguous().view(context.shape[0], -1, self.heads * self.d_k)
+
+        # (batch_size, max_len, d_model)
+        return self.output_linear(context)
+
+class FeedForward(torch.nn.Module):
+    "Implements FFN equation."
+
+    def __init__(self, d_model, middle_dim=2048, dropout=0.1):
+        super(FeedForward, self).__init__()
+        
+        self.fc1 = torch.nn.Linear(d_model, middle_dim)
+        self.fc2 = torch.nn.Linear(middle_dim, d_model)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.activation = torch.nn.GELU()
+
+    def forward(self, x):
+        out = self.activation(self.fc1(x))
+        out = self.fc2(self.dropout(out))
+        return out
+
+class EncoderLayer(torch.nn.Module):
+    def __init__(
+        self, 
+        d_model=768,
+        heads=12, 
+        feed_forward_hidden=768 * 4, 
+        dropout=0.1
+        ):
+        super(EncoderLayer, self).__init__()
+        self.layernorm = torch.nn.LayerNorm(d_model)
+        self.self_multihead = MultiHeadedAttention(heads, d_model)
+        self.feed_forward = FeedForward(d_model, middle_dim=feed_forward_hidden)
+        self.dropout = torch.nn.Dropout(dropout)
+
+    def forward(self, embeddings, mask):
+        # embeddings: (batch_size, max_len, d_model)
+        # encoder mask: (batch_size, 1, 1, max_len)
+        # result: (batch_size, max_len, d_model)
+        interacted = self.dropout(self.self_multihead(embeddings, embeddings, embeddings, mask))
+        # residual layer
+        interacted = self.layernorm(interacted + embeddings)
+        # bottleneck
+        feed_forward_out = self.dropout(self.feed_forward(interacted))
+        encoded = self.layernorm(feed_forward_out + interacted)
+        return encoded
+    
+
+class BERT(torch.nn.Module):
+    """
+    BERT model : Bidirectional Encoder Representations from Transformers.
+    """
+
+    def __init__(self, vocab_size, d_model=768, n_layers=12, heads=12, dropout=0.1):
+        """
+        :param vocab_size: vocab_size of total words
+        :param hidden: BERT model hidden size
+        :param n_layers: numbers of Transformer blocks(layers)
+        :param attn_heads: number of attention heads
+        :param dropout: dropout rate
+        """
+
+        super().__init__()
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.heads = heads
+
+        # paper noted they used 4 * hidden_size for ff_network_hidden_size
+        self.feed_forward_hidden = d_model * 4
+
+        # embedding for BERT, sum of positional, segment, token embeddings
+        self.embedding = BERTEmbedding(vocab_size=vocab_size, embed_size=d_model)
+
+        # multi-layers transformer blocks, deep network
+        self.encoder_blocks = torch.nn.ModuleList(
+            [EncoderLayer(d_model, heads, d_model * 4, dropout) for _ in range(n_layers)])
+
+    def forward(self, x, segment_info):
+        # attention masking for padded token
+        # (batch_size, 1, seq_len, seq_len)
+        mask = (x > 0).unsqueeze(1).repeat(1, x.size(1), 1).unsqueeze(1)
+
+        # embedding the indexed sequence to sequence of vectors
+        x = self.embedding(x, segment_info)
+
+        # running over multiple transformer blocks
+        for encoder in self.encoder_blocks:
+            x = encoder.forward(x, mask)
+        return x
+
+class NextSentencePrediction(torch.nn.Module):
+    """
+    2-class classification model : is_next, is_not_next
+    """
+
+    def __init__(self, hidden):
+        """
+        :param hidden: BERT model output size
+        """
+        super().__init__()
+        self.linear = torch.nn.Linear(hidden, 2)
+        self.softmax = torch.nn.LogSoftmax(dim=-1)
+
+    def forward(self, x):
+        # use only the first token which is the [CLS]
+        return self.softmax(self.linear(x[:, 0]))
+
+class MaskedLanguageModel(torch.nn.Module):
+    """
+    predicting origin token from masked input sequence
+    n-class classification problem, n-class = vocab_size
+    """
+
+    def __init__(self, hidden, vocab_size):
+        """
+        :param hidden: output size of BERT model
+        :param vocab_size: total vocab size
+        """
+        super().__init__()
+        self.linear = torch.nn.Linear(hidden, vocab_size)
+        self.softmax = torch.nn.LogSoftmax(dim=-1)
+
+    def forward(self, x):
+        return self.softmax(self.linear(x))
+
+class BERTLM(torch.nn.Module):
+    """
+    BERT Language Model
+    Next Sentence Prediction Model + Masked Language Model
+    """
+
+    def __init__(self, bert: BERT, vocab_size):
+        """
+        :param bert: BERT model which should be trained
+        :param vocab_size: total vocab size for masked_lm
+        """
+
+        super().__init__()
+        self.bert = bert
+        self.next_sentence = NextSentencePrediction(self.bert.d_model)
+        self.mask_lm = MaskedLanguageModel(self.bert.d_model, vocab_size)
+
+    def forward(self, x, segment_label):
+        x = self.bert(x, segment_label)
+        return self.next_sentence(x), self.mask_lm(x)

data.py

@@ -0,0 +1,30 @@
+def get_data(conversations: str, movie_lines: str, max_len: int=64) -> list: 
+
+    with open(conversations, 'r', encoding='iso-8859-1') as c:
+        conv = c.readlines()
+    with open(movie_lines, 'r', encoding='iso-8859-1') as l:
+        lines = l.readlines()
+
+    ### splitting text using special lines
+    lines_dic = {}
+    for line in lines:
+        objects = line.split(" +++$+++ ")
+        lines_dic[objects[0]] = objects[-1]
+
+    ### generate question answer pairs
+    pairs = []
+    for con in conv:
+        ids = eval(con.split(" +++$+++ ")[-1])
+        for i in range(len(ids)):
+            qa_pairs = []
+            
+            if i == len(ids) - 1:
+                break
+
+            first = lines_dic[ids[i]].strip()  
+            second = lines_dic[ids[i+1]].strip() 
+
+            qa_pairs.append(' '.join(first.split()[:max_len]))
+            qa_pairs.append(' '.join(second.split()[:max_len]))
+            pairs.append(qa_pairs)
+    return pairs

optimizer_schedule.py

@@ -0,0 +1,33 @@
+import numpy as np
+
+class ScheduledOptim():
+    '''A simple wrapper class for learning rate scheduling'''
+
+    def __init__(self, optimizer, d_model, n_warmup_steps):
+        self._optimizer = optimizer
+        self.n_warmup_steps = n_warmup_steps
+        self.n_current_steps = 0
+        self.init_lr = np.power(d_model, -0.5)
+
+    def step_and_update_lr(self):
+        "Step with the inner optimizer"
+        self._update_learning_rate()
+        self._optimizer.step()
+
+    def zero_grad(self):
+        "Zero out the gradients by the inner optimizer"
+        self._optimizer.zero_grad()
+
+    def _get_lr_scale(self):
+        return np.min([
+            np.power(self.n_current_steps, -0.5),
+            np.power(self.n_warmup_steps, -1.5) * self.n_current_steps])
+
+    def _update_learning_rate(self):
+        ''' Learning rate scheduling per step '''
+
+        self.n_current_steps += 1
+        lr = self.init_lr * self._get_lr_scale()
+
+        for param_group in self._optimizer.param_groups:
+            param_group['lr'] = lr

tokenizer.py

@@ -0,0 +1,59 @@
+import os 
+from pathlib import Path
+from tokenizers import BertWordPieceTokenizer
+from transformers import BertTokenizer
+import tqdm
+
+from data import get_data
+
+
+import re
+import transformers, datasets
+import numpy as np
+from torch.optim import Adam
+import math
+
+ 
+pairs = get_data('datasets/movie_conversations.txt', "datasets/movie_lines.txt")
+
+# WordPiece tokenizer
+
+### save data as txt file
+os.mkdir('data')
+text_data = []
+file_count = 0
+
+
+for sample in tqdm.tqdm([x[0] for x in pairs]):
+    text_data.append(sample)
+
+    # once we hit the 10K mark, save to file
+    if len(text_data) == 10000:
+        with open(f'data/text_{file_count}.txt', 'w', encoding='utf-8') as fp:
+            fp.write('\n'.join(text_data))
+        text_data = []
+        file_count += 1
+
+paths = [str(x) for x in Path('data').glob('**/*.txt')]
+
+
+### Training own tokenizer
+tokenizer = BertWordPieceTokenizer(
+    clean_text=True,
+    handle_chinese_chars=False,
+    strip_accents=False,
+    lowercase=True
+)
+
+
+tokenizer.train(
+    files=paths,
+    min_frequency=5,
+    limit_alphabet=1000,
+    wordpieces_prefix="##",
+    special_tokens=["[PAD]", "[CLS]", "[SEP]", "[MASK]", "[UNK]"]
+)
+
+
+os.mkdir("bert-it-1")
+tokenizer.save_model("bert-it-1", "bert-it")

train.ipynb

@@ -0,0 +1,71 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from bert_dataset import BERTDataset\n",
+    "from torch.utils.data import DataLoader\n",
+    "from bert_model import BERT, BERTLM\n",
+    "from trainer import BERTTrainer\n",
+    "from transformers import BertTokenizer\n",
+    "from data import get_data\n",
+    "\n",
+    "MAX_LEN = 128\n",
+    "\n",
+    "pairs = get_data('datasets/movie_conversations.txt', \"datasets/movie_lines.txt\")\n",
+    "tokenizer = BertTokenizer.from_pretrained(\"bert-it-1/bert-it-vocab.txt\")\n",
+    "\n",
+    "train_data = BERTDataset()\n",
+    "\n",
+    "train_loader = DataLoader(\n",
+    "   train_data, batch_size=32, shuffle=True, pin_memory=True)\n",
+    "\n",
+    "bert_model = BERT(\n",
+    "  vocab_size=len(tokenizer.vocab),\n",
+    "  d_model=768,\n",
+    "  n_layers=2,\n",
+    "  heads=12,\n",
+    "  dropout=0.1\n",
+    ")\n",
+    "\n",
+    "bert_lm = BERTLM(bert=bert_model, vocab_size=len(tokenizer.vocab))\n",
+    "bert_trainer = BERTTrainer(bert_lm, train_loader, device='cpu')\n",
+    "epochs = 20\n",
+    "\n",
+    "for epoch in range(epochs):\n",
+    "  bert_trainer.train(epoch)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "base",
+   "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.11.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

trainer.py

@@ -0,0 +1,106 @@
+import torch
+import torch.nn as nn
+import torch.nn. functional as F
+from optimizer_schedule import ScheduledOptim
+import tqdm
+from torch.optim import Adam
+
+
+class BERTTrainer:
+    def __init__(
+        self, 
+        model, 
+        train_dataloader, 
+        test_dataloader=None, 
+        lr= 1e-4,
+        weight_decay=0.01,
+        betas=(0.9, 0.999),
+        warmup_steps=10000,
+        log_freq=10,
+        device='cuda'
+        ):
+
+        self.device = device
+        self.model = model
+        self.train_data = train_dataloader
+        self.test_data = test_dataloader
+
+        # Setting the Adam optimizer with hyper-param
+        self.optim = Adam(self.model.parameters(), lr=lr, betas=betas, weight_decay=weight_decay)
+        self.optim_schedule = ScheduledOptim(
+            self.optim, self.model.bert.d_model, n_warmup_steps=warmup_steps
+            )
+
+        # Using Negative Log Likelihood Loss function for predicting the masked_token
+        self.criterion = torch.nn.NLLLoss(ignore_index=0)
+        self.log_freq = log_freq
+        print("Total Parameters:", sum([p.nelement() for p in self.model.parameters()]))
+    
+    def train(self, epoch):
+        self.iteration(epoch, self.train_data)
+
+    def test(self, epoch):
+        self.iteration(epoch, self.test_data, train=False)
+
+    def iteration(self, epoch, data_loader, train=True):
+        
+        avg_loss = 0.0
+        total_correct = 0
+        total_element = 0
+        
+        mode = "train" if train else "test"
+
+        # progress bar
+        data_iter = tqdm.tqdm(
+            enumerate(data_loader),
+            desc="EP_%s:%d" % (mode, epoch),
+            total=len(data_loader),
+            bar_format="{l_bar}{r_bar}"
+        )
+
+        for i, data in data_iter:
+
+            # 0. batch_data will be sent into the device(GPU or cpu)
+            data = {key: value.to(self.device) for key, value in data.items()}
+
+            # 1. forward the next_sentence_prediction and masked_lm model
+            next_sent_output, mask_lm_output = self.model.forward(data["bert_input"], data["segment_label"])
+
+            # 2-1. NLL(negative log likelihood) loss of is_next classification result
+            next_loss = self.criterion(next_sent_output, data["is_next"])
+
+            # 2-2. NLLLoss of predicting masked token word
+            # transpose to (m, vocab_size, seq_len) vs (m, seq_len)
+            # criterion(mask_lm_output.view(-1, mask_lm_output.size(-1)), data["bert_label"].view(-1))
+            mask_loss = self.criterion(mask_lm_output.transpose(1, 2), data["bert_label"])
+
+            # 2-3. Adding next_loss and mask_loss : 3.4 Pre-training Procedure
+            loss = next_loss + mask_loss
+
+            # 3. backward and optimization only in train
+            if train:
+                self.optim_schedule.zero_grad()
+                loss.backward()
+                self.optim_schedule.step_and_update_lr()
+
+            # next sentence prediction accuracy
+            correct = next_sent_output.argmax(dim=-1).eq(data["is_next"]).sum().item()
+            avg_loss += loss.item()
+            total_correct += correct
+            total_element += data["is_next"].nelement()
+
+            post_fix = {
+                "epoch": epoch,
+                "iter": i,
+                "avg_loss": avg_loss / (i + 1),
+                "avg_acc": total_correct / total_element * 100,
+                "loss": loss.item()
+            }
+
+            if i % self.log_freq == 0:
+                data_iter.write(str(post_fix))
+        print(
+            f"EP{epoch}, {mode}: \
+            avg_loss={avg_loss / len(data_iter)}, \
+            total_acc={total_correct * 100.0 / total_element}"
+        )