Update dnabert_layer.py

dnabert_layer.py

@@ -3,9 +3,21 @@ from typing import List, Optional, Tuple, Union
 import torch
 import torch.nn as nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-from transformers.models.bert.modeling_bert import BertPreTrainedModel, BertModel
+from transformers.models.bert.modeling_bert import BertModel as TransformersBertModel
+from transformers.models.bert.modeling_bert import BertForMaskedLM as TransformersBertForMaskedLM
+from transformers.models.bert.modeling_bert import BertPreTrainedModel
 from transformers.modeling_outputs import SequenceClassifierOutput
 
+
+class BertModel(TransformersBertModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+class BertForMaskedLM(TransformersBertForMaskedLM):
+    def __init__(self, config):
+        super().__init__(config)
+
+
 class DNABertForSequenceClassification(BertPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)