Fixing training environment

Author: codertimo

.gitignore

@@ -1,4 +1,5 @@
 data/
+output/
 
 # Created by .ignore support plugin (hsz.mobi)
 ### Python template

README.md

@@ -8,13 +8,123 @@ Pytorch implementation of Google AI's 2018 BERT, with simple annotation
 
 ## Introduction
 
-Currently WIP, with very high speed :)
-But it might be takes some days to validate my code
+Google AI's BERT paper shows the amazing result on various NLP task (new 17 NLP tasks SOTA), 
+including outperform the human F1 score on SQuAD v1.1 QA task. 
+This paper proved that Transformer(self-attention) based encoder can be powerfully used as 
+alternative of previous language model with proper language model training method. 
+And more importantly, they showed us that this pre-trained language model can be transfer 
+into any NLP task without making task specific model architecture.
 
-If you have any comment or question about my code, please leave it to issue.
-I'll reply back as soon as possible.
+This amazing result would be record in NLP history, 
+and I expect many further papers about BERT will be published very soon.
 
-Thank you
+This repo is implementation of BERT. Code is very simple and easy to understand fastly.
+Some of these codes are based on [The Annotated Transformer](http://nlp.seas.harvard.edu/2018/04/03/attention.html)
+
+
+## Language Model Pre-training
+
+In the paper, authors shows the new language model training methods, 
+which are "masked language model" and "predict next sentence".
+
+
+### Masked Language Model 
+
+> Original Paper : 3.3.1 Task #1: Masked LM 
+
+```
+Input Sequence  : The man went to [MASK] store with [MASK] dog
+Target Sequence :                  the                his
+```
+
+#### Rules:
+Randomly 15% of input token will be changed into something, based on under sub-rules
+
+1. Randomly 80% of tokens, gonna be a `[MASK]` token
+2. Randomly 10% of tokens, gonna be a `[RANDOM]` token(another word)
+3. Randomly 10% of tokens, will be remain as same. But need to be predicted.
+
+### Predict Next Sentence
+
+> Original Paper : 3.3.2 Task #2: Next Sentence Prediction
+
+```
+Input : [CLS] the man went to the store [SEP] he bought a gallon of milk [SEP]
+Label : Is Next
+
+Input = [CLS] the man heading to the store [SEP] penguin [MASK] are flight ##less birds [SEP]
+Label = NotNext
+```
+
+"Is this sentence can be continuously connected?"
+
+ understanding the relationship, between two text sentences, which is
+not directly captured by language modeling
+
+#### Rules:
+
+1. Randomly 50% of next sentence, gonna be continuous sentence.
+2. Randomly 50% of next sentence, gonna be unrelated sentence.
+
+
+## Usage
+
+### 1. Building vocab based on your corpus
+```shell
+python build_vocab.py -c data/corpus.small -o data/corpus.small.vocab
+```
+```shell
+usage: build_vocab.py [-h] -c CORPUS_PATH -o OUTPUT_PATH [-s VOCAB_SIZE]
+                      [-e ENCODING] [-m MIN_FREQ]
+
+optional arguments:
+  -h, --help            show this help message and exit
+  -c CORPUS_PATH, --corpus_path CORPUS_PATH
+  -o OUTPUT_PATH, --output_path OUTPUT_PATH
+  -s VOCAB_SIZE, --vocab_size VOCAB_SIZE
+  -e ENCODING, --encoding ENCODING
+  -m MIN_FREQ, --min_freq MIN_FREQ
+
+```
+### 2. Building BERT train dataset with your corpus
+```shell
+python build_dataset.py -d data/corpus.small -v data/corpus.small.vocab -o data/dataset.small
+```
+
+```shell
+usage: build_dataset.py [-h] -v VOCAB_PATH -c CORPUS_PATH [-e ENCODING] -o
+                        OUTPUT_PATH
+
+optional arguments:
+  -h, --help            show this help message and exit
+  -v VOCAB_PATH, --vocab_path VOCAB_PATH
+  -c CORPUS_PATH, --corpus_path CORPUS_PATH
+  -e ENCODING, --encoding ENCODING
+  -o OUTPUT_PATH, --output_path OUTPUT_PATH
+```
+
+### 3. Train your own BERT model
+```shell
+python train.py -d data/dataset.small -v data/corpus.small.vocab -o output/
+```
+```shell
+usage: train.py [-h] -d TRAIN_DATASET [-t TEST_DATASET] -v VOCAB_PATH -o
+                OUTPUT_DIR [-hs HIDDEN] [-n LAYERS] [-a ATTN_HEADS]
+                [-s SEQ_LEN] [-b BATCH_SIZE] [-e EPOCHS]
+
+optional arguments:
+  -h, --help            show this help message and exit
+  -d TRAIN_DATASET, --train_dataset TRAIN_DATASET
+  -t TEST_DATASET, --test_dataset TEST_DATASET
+  -v VOCAB_PATH, --vocab_path VOCAB_PATH
+  -o OUTPUT_DIR, --output_dir OUTPUT_DIR
+  -hs HIDDEN, --hidden HIDDEN
+  -n LAYERS, --layers LAYERS
+  -a ATTN_HEADS, --attn_heads ATTN_HEADS
+  -s SEQ_LEN, --seq_len SEQ_LEN
+  -b BATCH_SIZE, --batch_size BATCH_SIZE
+  -e EPOCHS, --epochs EPOCHS
+```
 
 
 ## Author

dataset/__init__.py

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

dataset/dataset.py

@@ -1,8 +1,6 @@
 from torch.utils.data import Dataset
-from .vocab import WordVocab
 import tqdm
 import random
-import argparse
 import torch
 
 
@@ -17,13 +15,7 @@ def __init__(self, corpus_path, vocab, seq_len, encoding="utf-8"):
                 t1, t2, t1_l, t2_l, is_next = line[:-1].split("\t")
                 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,
-                    "t2": t2,
-                    "t1_label": t1_l,
-                    "t2_label": t2_l,
-                    "is_next": is_next
-                })
+                self.datas.append({"t1": t1, "t2": t2, "t1_label": t1_l, "t2_label": t2_l, "is_next": is_next})
 
     def __len__(self):
         return len(self.datas)
@@ -33,11 +25,19 @@ def __getitem__(self, item):
         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"])
+        t1_label = [0] + self.vocab.to_seq(self.datas[item]["t1_label"]) + [0]
+        t2_label = self.vocab.to_seq(self.datas[item]["t2_label"]) + [0]
 
-        output = {"t1": t1, "t2": t2,
-                  "t1_label": t1_label, "t2_label": t2_label,
+        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.vocab.pad_index for _ in range(self.seq_len - len(t1) - len(t2))]
+        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": self.datas[item]["is_next"]}
 
         return {key: torch.tensor(value) for key, value in output.items()}

model/bert.py

@@ -20,15 +20,11 @@ def __init__(self, vocab_size, hidden=768, n_layers=12, attn_heads=12, dropout=0
                               dropout=dropout)
              for _ in range(n_layers)])
 
-    def forward(self, x):
+    def forward(self, x, segment_info):
         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)
+        x = self.embedding(x, segment_info)
 
-        # 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)
 

model/embedding/bert.py

@@ -11,5 +11,5 @@ def __init__(self, vocab_size, embed_size, dropout=0.1):
         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))
+    def forward(self, sequence, segment_label):
+        return self.position(self.token(sequence)) + self.segment(segment_label)

model/embedding/segment.py

@@ -3,4 +3,4 @@
 
 class SegmentEmbedding(nn.Embedding):
     def __init__(self, embed_size=512):
-        super().__init__(2, embed_size)
+        super().__init__(3, embed_size, padding_idx=0)

model/embedding/token.py

@@ -3,4 +3,4 @@
 
 class TokenEmbedding(nn.Embedding):
     def __init__(self, vocab_size, embed_size=512):
-        super().__init__(vocab_size, embed_size)
+        super().__init__(vocab_size, embed_size, padding_idx=0)

model/language_model.py

@@ -5,29 +5,29 @@
 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)
+        self.bert = bert
+        self.next_sentence = BERTNextSentence(self.bert.hidden)
+        self.mask_lm = BERTMaskLM(self.bert.hidden, vocab_size)
 
-    def forward(self, x):
+    def forward(self, x, segment_label):
+        x = self.bert(x, segment_label)
         return self.next_sentence(x), self.mask_lm(x)
 
 
 class BERTNextSentence(nn.Module):
-    def __init__(self, bert: BERT):
+    def __init__(self, hidden):
         super().__init__()
-        self.bert = bert
-        self.linear = nn.Linear(self.bert.hidden, 2)
+        self.linear = nn.Linear(hidden, 2)
         self.softmax = nn.LogSoftmax(dim=-1)
 
     def forward(self, x):
-        return self.softmax(self.linear(x))
+        return self.softmax(self.linear(x[:, 0]))
 
 
 class BERTMaskLM(nn.Module):
-    def __init__(self, bert: BERT, vocab_size):
+    def __init__(self, hidden, vocab_size):
         super().__init__()
-        self.bert = bert
-        self.linear = nn.Linear(self.bert.hidden, vocab_size)
+        self.linear = nn.Linear(hidden, vocab_size)
         self.softmax = nn.LogSoftmax(dim=-1)
 
     def forward(self, x):

train.py

@@ -1,20 +1,44 @@
 import argparse
-from dataset.dataset import BERTDataset, WordVocab
+
 from torch.utils.data import DataLoader
-from model import BERT, BERTLM
+
+from model import BERT
+from trainer import BERTTrainer
+from dataset import BERTDataset, WordVocab
 
 parser = argparse.ArgumentParser()
-parser.add_argument("-d", "--dataset_path", required=True, type=str)
+
+parser.add_argument("-d", "--train_dataset", required=True, type=str)
+parser.add_argument("-t", "--test_dataset", type=str, default=None)
 parser.add_argument("-v", "--vocab_path", required=True, type=str)
+parser.add_argument("-o", "--output_dir", required=True, type=str)
+
+parser.add_argument("-hs", "--hidden", type=int, default=128)
+parser.add_argument("-n", "--layers", type=int, default=2)
+parser.add_argument("-a", "--attn_heads", type=int, default=4)
+parser.add_argument("-s", "--seq_len", type=int, default=10)
+
+parser.add_argument("-b", "--batch_size", type=int, default=64)
+parser.add_argument("-e", "--epochs", type=int, default=10)
+
 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)
+print("Loading Train Dataset", args.train_dataset)
+train_dataset = BERTDataset(args.train_dataset, vocab, seq_len=args.seq_len)
+print("Loading Test Dataset", args.test_dataset)
+test_dataset = BERTDataset(args.test_dataset, vocab, seq_len=args.seq_len) if args.test_dataset is not None else None
+
+train_data_loader = DataLoader(train_dataset, batch_size=16)
+test_data_loader = DataLoader(test_dataset) if test_dataset is not None else None
+
+bert = BERT(len(vocab), hidden=args.hidden, n_layers=args.layers, attn_heads=args.attn_heads)
+trainer = BERTTrainer(bert, len(vocab), train_dataloader=train_data_loader, test_dataloader=test_data_loader)
 
+for epoch in range(args.epochs):
+    trainer.train(epoch)
+    trainer.save(args.output_dir, epoch)
 
-for data in data_loader:
-    x = model.forward(data["t1"])
-    print(x.size())
+    if test_data_loader is not None:
+        trainer.test(epoch)

trainer/__init__.py

@@ -0,0 +1 @@
+from .pretrain import BERTTrainer