Style fixes

Signed-off-by: MaximumEntropy <[email protected]>
NVIDIA · Jan 26, 2022 · 3b28c8d · 3b28c8d
1 parent f6a8ea0
commit 3b28c8d
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 4 deletions.
diff --git a/nemo/collections/common/tokenizers/sentencepiece_tokenizer.py b/nemo/collections/common/tokenizers/sentencepiece_tokenizer.py
@@ -240,7 +240,9 @@ def unk_id(self):
     @property
     def additional_special_tokens_ids(self):
         """Returns a list of the additional special tokens (excluding bos, eos, pad, unk). Used to return sentinel tokens for e.g. T5."""
-        special_tokens = set([self.bos_token, self.eos_token, self.pad_token, self.mask_token, self.cls_token, self.sep_token])
+        special_tokens = set(
+            [self.bos_token, self.eos_token, self.pad_token, self.mask_token, self.cls_token, self.sep_token]
+        )
         return [v for k, v in self.special_token_to_id.items() if k not in special_tokens]
 
     @property

diff --git a/nemo/collections/nlp/models/language_modeling/megatron_t5_model.py b/nemo/collections/nlp/models/language_modeling/megatron_t5_model.py
@@ -504,9 +504,11 @@ def _add_special_tokens_to_tokenizer(self):
                 self.tokenizer.add_special_tokens({'eos_token': '</s>'})
 
             # Special check to see if <extra_id_{}> is already present in the tokenizer. If it is, only modify the additional_special_tokens function.
-            for i in range (self.num_sentinel_tokens):
+            for i in range(self.num_sentinel_tokens):
                 if f'▁<extra_id_{i}>' in self.tokenizer.vocab:
-                    self.tokenizer.special_token_to_id[f'<extra_id_{i}>'] = self.tokenizer.text_to_ids(f'<extra_id_{i}>')[0]
+                    self.tokenizer.special_token_to_id[f'<extra_id_{i}>'] = self.tokenizer.text_to_ids(
+                        f'<extra_id_{i}>'
+                    )[0]
                 else:
                     self.tokenizer.add_special_tokens([f'<extra_id_{i}>'])
 

diff --git a/nemo/collections/nlp/modules/common/lm_utils.py b/nemo/collections/nlp/modules/common/lm_utils.py
@@ -87,9 +87,10 @@ def get_lm_model(
         )
 
     if nemo_file is not None:
-        from nemo.collections.nlp.models.language_modeling.megatron_bert_model import MegatronBertModel
         import torch
 
+        from nemo.collections.nlp.models.language_modeling.megatron_bert_model import MegatronBertModel
+
         class Identity(torch.nn.Module):
             def __init__(self):
                 super(Identity, self).__init__()