Lint megatron/data/dataset_utils.py

megatron/data/dataset_utils.py

@@ -132,6 +132,7 @@ def truncate_segments(tokens_a, tokens_b, len_a, len_b, max_num_tokens, np_rng):
             tokens.pop()
     return True
 
+
 def create_tokens_and_tokentypes(tokens_a, tokens_b, cls_id, sep_id):
     """Merge segments A and B, add [CLS] and [SEP] and build tokentypes."""
 
@@ -233,7 +234,7 @@ def create_masked_lm_predictions(tokens,
     for idx in range(len(cand_indexes)):
         ngram_index = []
         for n in ngrams:
-      ngram_index.append(cand_indexes[idx:idx+n])
+            ngram_index.append(cand_indexes[idx:idx + n])
         ngram_indexes.append(ngram_index)
 
     np_rng.shuffle(ngram_indexes)
@@ -367,12 +368,12 @@ def pad_and_convert_to_numpy(tokens, tokentypes, masked_positions,
     assert len(masked_positions) == len(masked_labels)
 
     # Tokens and token types.
-    filler = [pad_id]*padding_length
+    filler = [pad_id] * padding_length
     tokens_np = np.array(tokens + filler, dtype=np.int64)
     tokentypes_np = np.array(tokentypes + filler, dtype=np.int64)
 
     # Padding mask.
-    padding_mask_np = np.array([1]*num_tokens + [0]*padding_length,
+    padding_mask_np = np.array([1] * num_tokens + [0] * padding_length,
                                dtype=np.int64)
 
     # Lables and loss mask.