Upload NomicBertForPreTraining

modeling_hf_nomic_bert.py

@@ -318,6 +318,9 @@ class NomicBertPreTrainedModel(PreTrainedModel):
         remove_cls = cls != NomicBertForPreTraining
         remove_bert_prefix = cls != NomicBertForPreTraining
         ignore_mismatched_shapes = kwargs.pop("ignore_mismatched_sizes", False)
+        num_labels = kwargs.pop("num_labels", None)
+        if num_labels:
+            config.num_labels = num_labels
         model = cls(config, *inputs)
         # TODO: fix this
         # Assuming we know what we're doing when loading from disk
@@ -343,7 +346,7 @@ class NomicBertPreTrainedModel(PreTrainedModel):
                 state_dict,
                 strict=True
             )
-        logger.info(load_return)
+        logger.warning(load_return)
         return model
 
     def _set_gradient_checkpointing(self, module, value=False):