Try to subclass PretrainedModel

modeling_bert.py

@@ -22,7 +22,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from einops import rearrange
-from transformers import PretrainedConfig
+from transformers import PretrainedModel
 from .configuration_bert import JinaBertConfig
 from transformers.models.bert.modeling_bert import (
     BaseModelOutputWithPoolingAndCrossAttentions,
@@ -295,7 +295,7 @@ class BertPreTrainingHeads(nn.Module):
         return prediction_scores, seq_relationship_score
 
 
-class BertPreTrainedModel(nn.Module):
+class BertPreTrainedModel(PretrainedModel):
     """An abstract class to handle weights initialization and
     a simple interface for dowloading and loading pretrained models.
     """
@@ -310,33 +310,6 @@ class BertPreTrainedModel(nn.Module):
             )
         self.config = config
 
-    @classmethod
-    def from_pretrained(cls, model_name, config, *inputs, **kwargs):
-        """
-        Instantiate a BertPreTrainedModel from a pre-trained model file or a pytorch state dict.
-        Download and cache the pre-trained model file if needed.
-
-        Params:
-            pretrained_model_name_or_path: either:
-                - a path or url to a pretrained model archive containing:
-                    . `bert_config.json` a configuration file for the model
-                    . `pytorch_model.bin` a PyTorch dump of a BertForPretraining instance
-                - a path or url to a pretrained model archive containing:
-                    . `bert_config.json` a configuration file for the model
-                    . `model.chkpt` a TensorFlow checkpoint
-            *inputs, **kwargs: additional input for the specific Bert class
-                (ex: num_labels for BertForSequenceClassification)
-        """
-        # Instantiate model.
-        model = cls(config, *inputs, **kwargs)
-        load_return = model.load_state_dict(state_dict_from_pretrained(model_name), strict=True)
-        logger.info(load_return)
-        return model
-
-    @classmethod
-    def _from_config(cls, config, **kwargs):
-        return cls(config, **kwargs)
-
 
 class BertModel(BertPreTrainedModel):
     def __init__(self, config: JinaBertConfig, add_pooling_layer=True):