Fixing bugs & implement training process

Author: codertimo

.gitignore

@@ -1,3 +1,5 @@
+data/
+
 # Created by .ignore support plugin (hsz.mobi)
 ### Python template
 # Byte-compiled / optimized / DLL files

build_dataset.py

@@ -0,0 +1,23 @@
+from dataset.dataset import BERTDatasetCreator
+from dataset import WordVocab
+import argparse
+import tqdm
+
+parser = argparse.ArgumentParser()
+parser.add_argument("-v", "--vocab_path", required=True, type=str)
+parser.add_argument("-c", "--corpus_path", required=True, type=str)
+parser.add_argument("-e", "--encoding", default="utf-8", type=str)
+parser.add_argument("-o", "--output_path", required=True, type=str)
+args = parser.parse_args()
+
+word_vocab = WordVocab.load_vocab(args.vocab_path)
+builder = BERTDatasetCreator(corpus_path=args.corpus_path, vocab=word_vocab, seq_len=None, encoding=args.encoding)
+
+with open(args.output_path, 'w', encoding=args.encoding) as f:
+    for index in tqdm.tqdm(range(len(builder)), desc="Building Dataset", total=len(builder)):
+        data = builder[index]
+        output_form = "%s\t%s\t%s\t%s\t%d\n"
+        t1_text, t2_text = [" ".join(t) for t in [data["t1_random"], data["t2_random"]]]
+        t1_label, t2_label = [" ".join([str(i) for i in label]) for label in [data["t1_label"], data["t2_label"]]]
+        output = output_form % (t1_text, t2_text, t1_label, t2_label, data["is_next"])
+        f.write(output)

build_vocab.py

@@ -0,0 +1,15 @@
+import argparse
+from dataset import WordVocab
+
+parser = argparse.ArgumentParser()
+parser.add_argument("-c", "--corpus_path", required=True, type=str)
+parser.add_argument("-o", "--output_path", required=True, type=str)
+parser.add_argument("-s", "--vocab_size", type=int, default=None)
+parser.add_argument("-e", "--encoding", type=str, default="utf-8")
+parser.add_argument("-m", "--min_freq", type=int, default=1)
+args = parser.parse_args()
+
+with open(args.corpus_path, "r", encoding=args.encoding) as f:
+    vocab = WordVocab(f, max_size=args.vocab_size, min_freq=args.min_freq)
+
+vocab.save_vocab(args.output_path)

dataset/__init__.py

@@ -0,0 +1,2 @@
+from .dataset import BERTDataset, build_dataset
+from .vocab import WordVocab

dataset/dataset.py

@@ -1,4 +1,5 @@
 from torch.utils.data import Dataset
+from .vocab import WordVocab
 import tqdm
 import random
 import argparse
@@ -14,7 +15,7 @@ def __init__(self, corpus_path, vocab, seq_len, encoding="utf-8"):
         with open(corpus_path, "r", encoding=encoding) as f:
             for line in tqdm.tqdm(f, desc="Loading Dataset"):
                 t1, t2, t1_l, t2_l, is_next = line[:-1].split("\t")
-                t1_l, t2_l = [[int(i) for i in label.split(",")] for label in [t1_l, t2_l]]
+                t1_l, t2_l = [[token for token in label.split(" ")] for label in [t1_l, t2_l]]
                 is_next = int(is_next)
                 self.datas.append({
                     "t1": t1,
@@ -29,12 +30,14 @@ def __len__(self):
 
     def __getitem__(self, item):
         # [CLS] tag = SOS tag, [SEP] tag = EOS tag
-        t1, t1_len = self.vocab.to_seq(self.datas[item]["t1"], seq_len=self.seq_len, with_sos=True, with_eos=True)
-        t2, t2_len = self.vocab.to_seq(self.datas[item]["t2"], seq_len=self.seq_len, with_eos=True)
+        t1 = self.vocab.to_seq(self.datas[item]["t1"], with_sos=True, with_eos=True)
+        t2 = self.vocab.to_seq(self.datas[item]["t2"], with_eos=True)
+
+        t1_label = self.vocab.to_seq(self.datas[item]["t1_label"])
+        t2_label = self.vocab.to_seq(self.datas[item]["t2_label"])
 
         output = {"t1": t1, "t2": t2,
-                  "t1_len": t1_len, "t2_len": t2_len,
-                  "t1_label": self.datas[item]["t1_label"], "t2_label": self.datas[item]["t2_label"],
+                  "t1_label": t1_label, "t2_label": t2_label,
                   "is_next": self.datas[item]["is_next"]}
 
         return {key: torch.tensor(value) for key, value in output.items()}
@@ -79,38 +82,18 @@ def random_word(self, sentence):
     def random_sent(self, index):
         # output_text, label(isNotNext:0, isNext:1)
         if random.random() > 0.5:
-            return self.datas[index][2], 1
+            return self.datas[index][1], 1
         else:
-            return self.datas[random.randrange(len(self.datas))][2], 0
+            return self.datas[random.randrange(len(self.datas))][1], 0
 
     def __getitem__(self, index):
-        t1, (t2, is_next_label) = self.datas[index], self.random_sent(index)
+        t1, (t2, is_next_label) = self.datas[index][0], self.random_sent(index)
         t1_random, t1_label = self.random_word(t1)
         t2_random, t2_label = self.random_word(t2)
 
         return {"t1_random": t1_random, "t2_random": t2_random,
                 "t1_label": t1_label, "t2_label": t2_label,
                 "is_next": is_next_label}
 
-
-if __name__ == "__main__":
-    from .vocab import WordVocab
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument("-v", "--vocab_path", required=True, type=str)
-    parser.add_argument("-c", "--corpus_path", required=True, type=str)
-    parser.add_argument("-e", "--encoding", default="utf-8", type=str)
-    parser.add_argument("-o", "--output_path", required=True, type=str)
-    args = parser.parse_args()
-
-    word_vocab = WordVocab.load_vocab(args.vocab_path)
-    builder = BERTDatasetCreator(corpus_path=args.corpus_path, vocab=word_vocab, seq_len=None, encoding=args.encoding)
-
-    with open(args.output_path, 'w', encoding=args.encoding) as f:
-        for index in tqdm.tqdm(range(len(builder)), desc="Building Dataset", total=len(builder)):
-            data = builder[index]
-            output_form = "%s\t%s\t%s\t%d\n"
-            t1_text, t2_text = [" ".join(t) for t in [data["t1_random"], data["t2_random"]]]
-            t1_label, t2_label = [",".join([str(i) for i in label]) for label in [data["t1_label"], data["t2_label"]]]
-            output = output_form % (t1_text, t2_text, t1_label, t2_label, data["is_next"])
-            f.write(output_form)
+    def __len__(self):
+        return len(self.datas)

dataset/vocab.py

@@ -121,13 +121,20 @@ def __init__(self, texts, max_size=None, min_freq=1):
         print("Building Vocab")
         counter = Counter()
         for line in tqdm.tqdm(texts):
-            words = line.replace("\n", "").replace("\t", "").split()
+            if isinstance(line, list):
+                words = line
+            else:
+                words = line.replace("\n", "").replace("\t", "").split()
+
             for word in words:
                 counter[word] += 1
         super().__init__(counter, max_size=max_size, min_freq=min_freq)
 
     def to_seq(self, sentence, seq_len=None, with_eos=False, with_sos=False, with_len=False):
-        seq = [self.stoi.get(word, self.unk_index) for word in sentence.split()]
+        if isinstance(sentence, str):
+            sentence = sentence.split()
+
+        seq = [self.stoi.get(word, self.unk_index) for word in sentence]
 
         if with_eos:
             seq += [self.eos_index]  # this would be index 1
@@ -158,18 +165,3 @@ def from_seq(self, seq, join=False, with_pad=False):
     def load_vocab(vocab_path: str) -> 'WordVocab':
         with open(vocab_path, "rb") as f:
             return pickle.load(f)
-
-
-if __name__ == "__main__":
-    import argparse
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument("-c", "--corpus_path", required=True, type=str)
-    parser.add_argument("-o", "--output_path", required=True, type=str)
-    parser.add_argument("-s", "--vocab_size", type=int, default=None)
-    parser.add_argument("-e", "--encoding", type=str, default="utf-8")
-    parser.add_argument("-m", "--min_freq", type=int, default=1)
-    args = parser.parse_args()
-
-    with open(args.corpus_path, "r", encoding=args.encoding) as f:
-        vocab = WordVocab(f, max_size=args.vocab_size, min_freq=args.min_freq)

model/__init__.py

@@ -0,0 +1,2 @@
+from .bert import BERT
+from .language_model import BERTLM

model/attention/multi_head.py

@@ -12,7 +12,8 @@ def __init__(self, h, d_model, dropout=0.1):
         self.d_k = d_model // h
         self.h = h
 
-        self.linear_layers = nn.ModuleList([nn.Linear(d_model, d_model) for _ in range(4)])
+        self.linear_layers = nn.ModuleList([nn.Linear(d_model, d_model) for _ in range(3)])
+        self.output_linear = nn.Linear(d_model, d_model)
         self.attention = Attention()
 
         self.attn = None
@@ -37,4 +38,4 @@ def forward(self, query, key, value, mask=None):
         # 3) "Concat" using a view and apply a final linear.
         x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k)
 
-        return self.linears[-1](x)
+        return self.output_linear(x)

model/bert.py

@@ -5,22 +5,31 @@
 
 
 class BERT(nn.Module):
-    def __init__(self, embedding: BERTEmbedding, hidden=768, n_layers=12, attn_heads=12, dropout=0.1):
+    def __init__(self, vocab_size, hidden=768, n_layers=12, attn_heads=12, dropout=0.1):
         super().__init__()
         self.hidden = hidden
         self.n_layers = n_layers
         self.attn_heads = attn_heads
         self.feed_forward_hidden = hidden * 4
 
-        self.embedding: BERTEmbedding = embedding
+        self.embedding = BERTEmbedding(vocab_size=vocab_size, embed_size=hidden)
         self.transformer_blocks = nn.ModuleList(
             [TransformerBlock(hidden=hidden,
                               attn_heads=attn_heads,
                               feed_forward_hidden=hidden * 4,
                               dropout=dropout)
              for _ in range(n_layers)])
 
-    def forward(self, x, mask=None):
+    def forward(self, x):
+        mask = (x > 0).unsqueeze(1).repeat(1, x.size(1), 1)
+
+        # sequence -> embedding : (batch_size, seq_len) -> (batch_size, seq_len, embed_size)
+        x = self.embedding(x)
+
+        # embedding through the transformer self-attention
+        # embedding (batch_size, seq_len, embed_size = hidden) -> transformer_output (batch_size, seq_len, hidden)
+        # loop transformer (batch_size, seq_len, hidden) -> transformer_output (batch_size, seq_len, hidden)
         for transformer in self.transformer_blocks:
             x = transformer.forward(x, mask)
+
         return x

model/embedding/__init__.py

@@ -1 +1 @@
-from .bert_embedding import BERTEmbedding
+from .bert import BERTEmbedding

model/embedding/bert.py

@@ -0,0 +1,15 @@
+import torch.nn as nn
+from .token import TokenEmbedding
+from .position import PositionalEmbedding
+from .segment import SegmentEmbedding
+
+
+class BERTEmbedding(nn.Module):
+    def __init__(self, vocab_size, embed_size, dropout=0.1):
+        super().__init__()
+        self.token = TokenEmbedding(vocab_size=vocab_size, embed_size=embed_size)
+        self.position = PositionalEmbedding(self.token.embedding_dim, dropout=dropout)
+        self.segment = SegmentEmbedding(embed_size=self.token.embedding_dim)
+
+    def forward(self, sequence):
+        return self.position(self.token(sequence))

model/embedding/bert_embedding.py

@@ -11,11 +11,11 @@ def __init__(self, d_model, dropout, max_len=512):
         self.dropout = nn.Dropout(p=dropout)
 
         # Compute the positional encodings once in log space.
-        pe = torch.zeros(max_len, d_model)
+        pe = torch.zeros(max_len, d_model).float()
         pe.require_grad = False
 
-        position = torch.arange(0, max_len).unsqueeze(1)
-        div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
+        position = torch.arange(0, max_len).float().unsqueeze(1)
+        div_term = (torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)).float().exp()
 
         pe[:, 0::2] = torch.sin(position * div_term)
         pe[:, 1::2] = torch.cos(position * div_term)

model/embedding/position.py

@@ -0,0 +1,34 @@
+from .bert import BERT
+import torch.nn as nn
+
+
+class BERTLM(nn.Module):
+    def __init__(self, bert: BERT, vocab_size):
+        super().__init__()
+        self.next_sentence = BERTNextSentence(bert)
+        self.mask_lm = BERTMaskLM(bert, vocab_size)
+
+    def forward(self, x):
+        return self.next_sentence(x), self.mask_lm(x)
+
+
+class BERTNextSentence(nn.Module):
+    def __init__(self, bert: BERT):
+        super().__init__()
+        self.bert = bert
+        self.linear = nn.Linear(self.bert.hidden, 2)
+        self.softmax = nn.LogSoftmax(dim=-1)
+
+    def forward(self, x):
+        return self.softmax(self.linear(x))
+
+
+class BERTMaskLM(nn.Module):
+    def __init__(self, bert: BERT, vocab_size):
+        super().__init__()
+        self.bert = bert
+        self.linear = nn.Linear(self.bert.hidden, vocab_size)
+        self.softmax = nn.LogSoftmax(dim=-1)
+
+    def forward(self, x):
+        return self.softmax(self.linear(x))

model/language_model.py

@@ -13,8 +13,6 @@ def __init__(self, hidden, attn_heads, feed_forward_hidden, dropout):
         self.output_sublayer = SublayerConnection(size=hidden, dropout=dropout)
 
     def forward(self, x, mask):
-        _x = self.attention.forward(x, x, x, mask=mask)
-        x = self.input_sublayer(x, _x)
-        _x = self.feed_forward(x)
-        x = self.output_sublayer(x, _x)
+        x = self.input_sublayer(x, lambda _x: self.attention.forward(_x, _x, _x, mask=mask))
+        x = self.output_sublayer(x, self.feed_forward)
         return x

model/transformer.py

@@ -0,0 +1,20 @@
+import argparse
+from dataset.dataset import BERTDataset, WordVocab
+from torch.utils.data import DataLoader
+from model import BERT, BERTLM
+
+parser = argparse.ArgumentParser()
+parser.add_argument("-d", "--dataset_path", required=True, type=str)
+parser.add_argument("-v", "--vocab_path", required=True, type=str)
+args = parser.parse_args()
+
+vocab = WordVocab.load_vocab(args.vocab_path)
+dataset = BERTDataset(args.dataset_path, vocab, seq_len=10)
+data_loader = DataLoader(dataset, batch_size=16)
+
+bert = BERT(len(vocab), hidden=128, n_layers=2, attn_heads=4)
+
+
+for data in data_loader:
+    x = model.forward(data["t1"])
+    print(x.size())

train.py

@@ -0,0 +1,24 @@
+import torch.nn as nn
+from torch.optim import Adam
+from model import BERTLM, BERT
+
+
+class BERTTrainer:
+    def __init__(self, bert: BERT, vocab_size, train_dataloader, test_dataloader=None):
+        self.bert = bert
+        self.lm = BERTLM(bert, vocab_size)
+
+        self.train_data = train_dataloader
+        self.test_data = test_dataloader
+
+        self.optim = Adam(self.lm.parameters())
+        self.criterion = nn.NLLLoss()
+
+    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):
+        pass