update kwargs for pretrainer

geneformer/pretrainer.py

@@ -106,9 +106,8 @@ class TensorType(ExplicitEnum):
 
 class GeneformerPreCollator(SpecialTokensMixin):
     def __init__(self, *args, **kwargs) -> None:
-        
-        super().__init__(mask_token = "<mask>", pad_token = "<pad>")
-        
+        super().__init__(mask_token="<mask>", pad_token="<pad>")
+
         self.token_dictionary = kwargs.get("token_dictionary")
         # self.mask_token = "<mask>"
         # self.mask_token_id = self.token_dictionary.get("<mask>")
@@ -120,8 +119,8 @@ class GeneformerPreCollator(SpecialTokensMixin):
         #     self.token_dictionary.get("<pad>"),
         # ]
         self.model_input_names = ["input_ids"]
-    
-    def convert_ids_to_tokens(self,value):
+
+    def convert_ids_to_tokens(self, value):
         return self.token_dictionary.get(value)
 
     def _get_padding_truncation_strategies(
@@ -391,7 +390,6 @@ class GeneformerPreCollator(SpecialTokensMixin):
 
             for key, value in encoded_inputs.items():
                 encoded_inputs[key] = to_py_obj(value)
-                
 
         # Convert padding_strategy in PaddingStrategy
         padding_strategy, _, max_length, _ = self._get_padding_truncation_strategies(
@@ -596,15 +594,17 @@ class GeneformerPreCollator(SpecialTokensMixin):
 
 class GeneformerPretrainer(Trainer):
     def __init__(self, *args, **kwargs):
-        data_collator = kwargs.get("data_collator",None)
+        data_collator = kwargs.get("data_collator", None)
         token_dictionary = kwargs.pop("token_dictionary")
+        mlm = kwargs.pop("mlm", True)
+        mlm_probability = kwargs.pop("mlm_probability", 0.15)
 
         if data_collator is None:
             precollator = GeneformerPreCollator(token_dictionary=token_dictionary)
 
             # # Data Collator Functions
             data_collator = DataCollatorForLanguageModeling(
-                tokenizer=precollator, mlm=True, mlm_probability=0.15
+                tokenizer=precollator, mlm=mlm, mlm_probability=mlm_probability
             )
             kwargs["data_collator"] = data_collator
 
@@ -694,6 +694,7 @@ class CustomDistributedLengthGroupedSampler(DistributedLengthGroupedSampler):
     Distributed Sampler that samples indices in a way that groups together features of the dataset of roughly the same
     length while keeping a bit of randomness.
     """
+
     # Copied and adapted from PyTorch DistributedSampler.
     def __init__(
         self,
@@ -757,7 +758,7 @@ class CustomDistributedLengthGroupedSampler(DistributedLengthGroupedSampler):
         # Deterministically shuffle based on epoch and seed
         g = torch.Generator()
         g.manual_seed(self.seed + self.epoch)
-        
+
         indices = get_length_grouped_indices(self.lengths, self.batch_size, generator=g)
 
         if not self.drop_last: