Merge pull request #12 from codertimo/alpha0.0.1a3

alpha0.0.1a3 version update

README.md

@@ -1,9 +1,9 @@
 # BERT-pytorch
 
-[![LICENSE](https://img.shields.io/github/license/codertimo/BERT-pytorch.svg)](https://github.com/kor2vec/kor2vec/blob/master/LICENSE)
+[![LICENSE](https://img.shields.io/github/license/codertimo/BERT-pytorch.svg)](https://github.com/codertimo/BERT-pytorch/blob/master/LICENSE)
 ![GitHub issues](https://img.shields.io/github/issues/codertimo/BERT-pytorch.svg)
-[![GitHub stars](https://img.shields.io/github/stars/codertimo/BERT-pytorch.svg)](https://github.com/kor2vec/kor2vec/stargazers)
-[![CircleCI](https://circleci.com/gh/codertimo/BERT-pytorch.svg?style=shield)](https://circleci.com/gh/kor2vec/kor2vec)
+[![GitHub stars](https://img.shields.io/github/stars/codertimo/BERT-pytorch.svg)](https://github.com/codertimo/BERT-pytorch/stargazers)
+[![CircleCI](https://circleci.com/gh/codertimo/BERT-pytorch.svg?style=shield)](https://circleci.com/gh/codertimo/BERT-pytorch)
 [![PyPI](https://img.shields.io/pypi/v/bert-pytorch.svg)](https://pypi.org/project/bert_pytorch/)
 [![PyPI - Status](https://img.shields.io/pypi/status/bert-pytorch.svg)](https://pypi.org/project/bert_pytorch/)
 [![Documentation Status](https://readthedocs.org/projects/bert-pytorch/badge/?version=latest)](https://bert-pytorch.readthedocs.io/en/latest/?badge=latest)
@@ -39,24 +39,28 @@ pip install bert-pytorch
 ## Quickstart
 
 **NOTICE : Your corpus should be prepared with two sentences in one line with tab(\t) separator**
+
+### 0. Prepare your corpus
 ```
-Welcome to the \t the jungle \n
-I can stay \t here all night \n
+Welcome to the \t the jungle\n
+I can stay \t here all night\n
 ```
 
-### 1. Building vocab based on your corpus
-```shell
-bert-vocab -c data/corpus.small -o data/corpus.small.vocab
+or tokenized corpus (tokenization is not in package)
 ```
+Wel_ _come _to _the \t _the _jungle\n
+_I _can _stay \t _here _all _night\n
+```
+
 
-### 2. Building BERT train dataset with your corpus
+### 1. Building vocab based on your corpus
 ```shell
-bert-dataset -d data/corpus.small -v data/corpus.small.vocab -o data/dataset.small
+bert-vocab -c data/corpus.small -o data/vocab.small
 ```
 
-### 3. Train your own BERT model
+### 2. Train your own BERT model
 ```shell
-bert -d data/dataset.small -v data/corpus.small.vocab -o output/bert.model
+bert -c data/dataset.small -v data/vocab.small -o output/bert.model
 ```
 
 ## Language Model Pre-training

bert_pytorch/__main__.py

@@ -0,0 +1,67 @@
+import argparse
+
+from torch.utils.data import DataLoader
+
+from .model import BERT
+from .trainer import BERTTrainer
+from .dataset import BERTDataset, WordVocab
+
+
+def train():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("-c", "--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_path", required=True, type=str)
+
+    parser.add_argument("-hs", "--hidden", type=int, default=256)
+    parser.add_argument("-l", "--layers", type=int, default=8)
+    parser.add_argument("-a", "--attn_heads", type=int, default=8)
+    parser.add_argument("-s", "--seq_len", type=int, default=20)
+
+    parser.add_argument("-b", "--batch_size", type=int, default=64)
+    parser.add_argument("-e", "--epochs", type=int, default=10)
+    parser.add_argument("-w", "--num_workers", type=int, default=5)
+    parser.add_argument("--with_cuda", type=bool, default=True)
+    parser.add_argument("--log_freq", type=int, default=10)
+    parser.add_argument("--corpus_lines", type=int, default=None)
+
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--adam_weight_decay", type=float, default=0.01)
+    parser.add_argument("--adam_beta1", type=float, default=0.9)
+    parser.add_argument("--adam_beta2", type=float, default=0.999)
+
+    args = parser.parse_args()
+
+    print("Loading Vocab", args.vocab_path)
+    vocab = WordVocab.load_vocab(args.vocab_path)
+    print("Vocab Size: ", len(vocab))
+
+    print("Loading Train Dataset", args.train_dataset)
+    train_dataset = BERTDataset(args.train_dataset, vocab, seq_len=args.seq_len, corpus_lines=args.corpus_lines)
+
+    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
+
+    print("Creating Dataloader")
+    train_data_loader = DataLoader(train_dataset, batch_size=args.batch_size, num_workers=args.num_workers)
+    test_data_loader = DataLoader(test_dataset, batch_size=args.batch_size, num_workers=args.num_workers) \
+        if test_dataset is not None else None
+
+    print("Building BERT model")
+    bert = BERT(len(vocab), hidden=args.hidden, n_layers=args.layers, attn_heads=args.attn_heads)
+
+    print("Creating BERT Trainer")
+    trainer = BERTTrainer(bert, len(vocab), train_dataloader=train_data_loader, test_dataloader=test_data_loader,
+                          lr=args.lr, betas=(args.adam_beta1, args.adam_beta2), weight_decay=args.adam_weight_decay,
+                          with_cuda=args.with_cuda, log_freq=args.log_freq)
+
+    print("Training Start")
+    for epoch in range(args.epochs):
+        trainer.train(epoch)
+        trainer.save(epoch, args.output_path)
+
+        if test_data_loader is not None:
+            trainer.test(epoch)

bert_pytorch/dataset/__init__.py

@@ -1,3 +1,2 @@
 from .dataset import BERTDataset
-from .creator import BERTDatasetCreator
 from .vocab import WordVocab

bert_pytorch/dataset/dataset.py

@@ -1,32 +1,32 @@
 from torch.utils.data import Dataset
 import tqdm
 import torch
+import random
 
 
 class BERTDataset(Dataset):
     def __init__(self, corpus_path, vocab, seq_len, encoding="utf-8", corpus_lines=None):
         self.vocab = vocab
         self.seq_len = seq_len
 
-        self.datas = []
         with open(corpus_path, "r", encoding=encoding) as f:
-            for line in tqdm.tqdm(f, desc="Loading Dataset", total=corpus_lines):
-                t1, t2, t1_l, t2_l, is_next = line[:-1].split("\t")
-                t1, t2 = [[int(token) for token in t.split(",")] for t in [t1, t2]]
-                t1_l, t2_l = [[int(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 = [line[:-1].split("\t")
+                          for line in tqdm.tqdm(f, desc="Loading Dataset", total=corpus_lines)]
 
     def __len__(self):
         return len(self.datas)
 
     def __getitem__(self, item):
+        t1, (t2, is_next_label) = self.datas[item][0], self.random_sent(item)
+        t1_random, t1_label = self.random_word(t1)
+        t2_random, t2_label = self.random_word(t2)
+
         # [CLS] tag = SOS tag, [SEP] tag = EOS tag
-        t1 = [self.vocab.sos_index] + self.datas[item]["t1"] + [self.vocab.eos_index]
-        t2 = self.datas[item]["t2"] + [self.vocab.eos_index]
+        t1 = [self.vocab.sos_index] + t1_random + [self.vocab.eos_index]
+        t2 = t2_random + [self.vocab.eos_index]
 
-        t1_label = [0] + self.datas[item]["t1_label"] + [0]
-        t2_label = self.datas[item]["t2_label"] + [0]
+        t1_label = [self.vocab.pad_index] + t1_label + [self.vocab.pad_index]
+        t2_label = t2_label + [self.vocab.pad_index]
 
         segment_label = ([1 for _ in range(len(t1))] + [2 for _ in range(len(t2))])[:self.seq_len]
         bert_input = (t1 + t2)[:self.seq_len]
@@ -38,6 +38,40 @@ def __getitem__(self, item):
         output = {"bert_input": bert_input,
                   "bert_label": bert_label,
                   "segment_label": segment_label,
-                  "is_next": self.datas[item]["is_next"]}
+                  "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 = []
+
+        for i, token in enumerate(tokens):
+            prob = random.random()
+            if prob < 0.15:
+                # 80% randomly change token to make token
+                if prob < prob * 0.8:
+                    tokens[i] = self.vocab.mask_index
+
+                # 10% randomly change token to random token
+                elif prob * 0.8 <= prob < prob * 0.9:
+                    tokens[i] = random.randrange(len(self.vocab))
+
+                # 10% randomly change token to current token
+                elif prob >= prob * 0.9:
+                    tokens[i] = self.vocab.stoi.get(token, self.vocab.unk_index)
+
+                output_label.append(self.vocab.stoi.get(token, self.vocab.unk_index))
+
+            else:
+                tokens[i] = self.vocab.stoi.get(token, self.vocab.unk_index)
+                output_label.append(0)
+
+        return tokens, output_label
+
+    def random_sent(self, index):
+        # output_text, label(isNotNext:0, isNext:1)
+        if random.random() > 0.5:
+            return self.datas[index][1], 1
+        else:
+            return self.datas[random.randrange(len(self.datas))][1], 0

bert_pytorch/dataset/vocab.py

@@ -165,3 +165,21 @@ 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)
+
+
+def build():
+    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)
+
+    print("VOCAB SIZE:", len(vocab))
+    vocab.save_vocab(args.output_path)