feat: added top-level docstring, made it compatible with AutoModel

modeling_bert.py

@@ -1,3 +1,10 @@
+""" Implementation of BERT, using ALiBi and Flash Attention
+
+The implementation was adopted from
+https://github.com/Dao-AILab/flash-attention/blob/43950dda456e095969d842fca7a73c5bfe3cecd0/flash_attn/models/bert.py
+and made modifications to use ALiBi.
+"""
+
 # Copyright (c) 2022, Tri Dao.
 # This BERT implementation is based on our MLPerf 2.0 and MLPerf 2.1 BERT implementation.
 # https://github.com/mlcommons/training_results_v2.0/blob/main/HazyResearch/benchmarks/bert/implementations/pytorch/modeling.py
@@ -297,12 +304,10 @@ class BertPreTrainedModel(nn.Module):
 
     def __init__(self, config, *inputs, **kwargs):
         super().__init__()
-        if not isinstance(config, JinaBertConfig):
+        if not config.__class__.__name__ == 'JinaBertConfig':
             raise ValueError(
-                "Parameter config in `{}(config)` should be an instance of class `BertConfig`. "
-                "To create a model from a Google pretrained model use "
-                "`model = {}.from_pretrained(PRETRAINED_MODEL_NAME)`".format(
-                    self.__class__.__name__, self.__class__.__name__
+                "Parameter config in `{}(config)` should be an instance of class `JinaBertConfig`.".format(
+                    self.__class__.__name__,
                 )
             )
         self.config = config