solve Infinite loop when masking

Author: graykode

models.py

@@ -1,7 +1,8 @@
-# Copyright 2019 Tae Hwan Jung(@graykode)
-# forked from https://github.com/dhlee347/pytorchic-bert
-# (Strongly inspired by original Google BERT code and Hugging Face's code)
-# Remove All Dropout
+"""
+    Copyright 2019 Tae Hwan Jung
+    ALBERT Implementation with forking
+    Clean Pytorch Code from https://github.com/dhlee347/pytorchic-bert
+"""
 
 """ Transformer Model Classes & Config Class """
 

optim.py

@@ -1,5 +1,8 @@
-# Copyright 2018 The Google AI Language Team Authors and The HugginFace Inc. team,
-# and Dong-Hyun Lee, Kakao Brain.
+"""
+    Copyright 2019 Tae Hwan Jung
+    ALBERT Implementation with forking
+    Clean Pytorch Code from https://github.com/dhlee347/pytorchic-bert
+"""
 
 """ a slightly modified version of Hugging Face's BERTAdam class """
 

pretrain.py

@@ -18,8 +18,7 @@
 import optim
 import train
 
-from utils import set_seeds, get_device, get_random_word, truncate_tokens_pair, \
-    _is_start_piece, _sample_mask
+from utils import set_seeds, get_device, truncate_tokens_pair, _sample_mask
 
 # Input file format :
 # 1. One sentence per line. These should ideally be actual sentences,
@@ -113,11 +112,12 @@ def __call__(self, instance):
 
 class Preprocess4Pretrain(Pipeline):
     """ Pre-processing steps for pretraining transformer """
-    def __init__(self, max_pred, vocab_words, indexer, max_len,
+    def __init__(self, max_pred, mask_prob, vocab_words, indexer, max_len,
                  mask_alpha, mask_beta, max_gram):
         super().__init__()
         self.max_len = max_len
         self.max_pred = max_pred # max tokens of prediction
+        self.mask_prob = mask_prob # masking probability
         self.vocab_words = vocab_words # vocabulary (sub)words
 
         self.indexer = indexer # function from token to token index
@@ -137,9 +137,13 @@ def __call__(self, instance):
         segment_ids = [0]*(len(tokens_a)+2) + [1]*(len(tokens_b)+1)
         input_mask = [1]*len(tokens)
 
+        # the number of prediction is sometimes less than max_pred when sequence is short
+        n_pred = min(self.max_pred, max(1, int(round(len(tokens) * self.mask_prob))))
+
         # For masked Language Models
         masked_tokens, masked_pos, tokens = _sample_mask(tokens, self.mask_alpha,
-                                            self.mask_beta, self.max_gram, self.max_pred)
+                                            self.mask_beta, self.max_gram,
+                                            goal_num_predict=n_pred)
 
         masked_weights = [1]*len(masked_tokens)
 
@@ -213,6 +217,7 @@ def main(args):
     tokenize = lambda x: tokenizer.tokenize(tokenizer.convert_to_unicode(x))
 
     pipeline = [Preprocess4Pretrain(args.max_pred,
+                                    args.mask_prob,
                                     list(tokenizer.vocab.keys()),
                                     tokenizer.convert_tokens_to_ids,
                                     model_cfg.max_len,
@@ -262,11 +267,12 @@ def get_loss(model, batch, global_step): # make sure loss is tensor
 
     # official google-reacher/bert is use 20, but 20/512(=seq_len)*100 make only 3% Mask
     # So, official XLNET zihangdai/xlnet use 85 with name of num_predict(SAME HERE!)
-    parser.add_argument('--max_pred', type=int, default=85)
+    parser.add_argument('--max_pred', type=int, default=76, help='max tokens of prediction')
+    parser.add_argument('--mask_prob', type=float, default=0.15, help='masking probability')
 
     # try to n-gram masking SpanBERT(Joshi et al., 2019)
     parser.add_argument('--mask_alpha', type=int,
-                        default=6, help="How many tokens to form a group.")
+                        default=4, help="How many tokens to form a group.")
     parser.add_argument('--mask_beta', type=int,
                         default=1, help="How many tokens to mask within each group.")
     parser.add_argument('--max_gram', type=int,

tokenization.py

@@ -1,6 +1,8 @@
-# Copyright 2018 The Google AI Language Team Authors and The HugginFace Inc. team.
-
-""" Tokenization classes (It's exactly the same code as Google BERT code """
+"""
+    Copyright 2019 Tae Hwan Jung
+    ALBERT Implementation with forking
+    Clean Pytorch Code from https://github.com/dhlee347/pytorchic-bert
+"""
 
 from __future__ import absolute_import
 from __future__ import division

train.py

@@ -1,4 +1,8 @@
-# Copyright 2018 Dong-Hyun Lee, Kakao Brain.
+"""
+    Copyright 2019 Tae Hwan Jung
+    ALBERT Implementation with forking
+    Clean Pytorch Code from https://github.com/dhlee347/pytorchic-bert
+"""
 
 """ Training Config & Helper Classes  """
 

utils.py

@@ -1,4 +1,8 @@
-# Copyright 2018 Dong-Hyun Lee, Kakao Brain.
+"""
+    Copyright 2019 Tae Hwan Jung
+    ALBERT Implementation with forking
+    Clean Pytorch Code from https://github.com/dhlee347/pytorchic-bert
+"""
 
 """ Utils Functions """
 
@@ -114,12 +118,16 @@ def _sample_mask(seg, mask_alpha, mask_beta,
     pvals /= pvals.sum(keepdims=True) # p(n) = 1/n / sigma(1/k)
 
     cur_len = 0
+
     while cur_len < seg_len:
         if goal_num_predict is not None and num_predict >= goal_num_predict: break
 
         n = np.random.choice(ngrams, p=pvals)
         if goal_num_predict is not None:
             n = min(n, goal_num_predict - num_predict)
+
+        # `mask_alpha` : number of tokens forming group
+        # `mask_beta` : number of tokens to be masked in each groups.
         ctx_size = (n * mask_alpha) // mask_beta
         l_ctx = np.random.choice(ctx_size)
         r_ctx = ctx_size - l_ctx
@@ -164,5 +172,4 @@ def _sample_mask(seg, mask_alpha, mask_beta,
             tokens.append('[MASK]')
         else:
             tokens.append(seg[i])
-
     return masked_tokens, masked_pos, tokens