add bert prefix

automodel.py

@@ -52,8 +52,9 @@ class MosaicBertForEmbeddingGeneration(BertPreTrainedModel):
         # Load the state_dict
         state_dict = load_file(archive_file)
 
-        # remove `model` prefix to avoid error
-        consume_prefix_in_state_dict_if_present(state_dict, prefix="model.")
+        # add missing bert prefix
+        state_dict = {f'bert.{key}': value for key, value in state_dict.items()}
+        
         missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
 
         if len(missing_keys) > 0:
@@ -131,8 +132,9 @@ class ClinicalMosaicForSequenceClassification(BertPreTrainedModel):
         # Load the state_dict
         state_dict = load_file(archive_file)
 
-        # remove `model` prefix to avoid error
-        consume_prefix_in_state_dict_if_present(state_dict, prefix="model.")
+        # add missing bert prefix
+        state_dict = {f'bert.{key}': value for key, value in state_dict.items()}
+        
         missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
 
         if len(missing_keys) > 0: