Update modeling_hf_nomic_bert.py

modeling_hf_nomic_bert.py

@@ -316,7 +316,9 @@ class NomicBertPreTrainedModel(PreTrainedModel):
         num_labels = kwargs.pop("num_labels", None)
         rotary_scaling_factor = kwargs.pop("rotary_scaling_factor", None)
         strict = kwargs.pop("strict", True)
-        config.rotary_scaling_factor = rotary_scaling_factor
+        if rotary_scaling_factor:
+            config.rotary_scaling_factor = rotary_scaling_factor
+
         if config.n_positions <= 0 and config.rotary_emb_fraction > 0:
             config.n_positions = 2048
         if num_labels:
@@ -325,7 +327,10 @@ class NomicBertPreTrainedModel(PreTrainedModel):
         if "add_pooling_layer" in kwargs:
             model = cls(config, *inputs, add_pooling_layer=kwargs.pop("add_pooling_layer"))
         else:
-            model = cls(config, *inputs)
+            if cls == NomicBertModel:
+                model = cls(config, *inputs, add_pooling_layer=False)
+            else:
+                model = cls(config, *inputs)
         # TODO: fix this
         # Assuming we know what we're doing when loading from disk
         # Prob a bad assumption but i'm tired and want to train this asap
@@ -344,7 +349,7 @@ class NomicBertPreTrainedModel(PreTrainedModel):
             load_return = model.load_state_dict(state_dict, strict=False)
         else:
             # TODO: can probably check config class and see if we need to remap from a bert model
-            state_dict = state_dict_from_pretrained(model_name)
+            state_dict = state_dict_from_pretrained(model_name, safe_serialization=kwargs.get("safe_serialization", False))
             state_dict = remap_bert_state_dict(
                 state_dict,
                 config,
@@ -1057,9 +1062,9 @@ class NomicBertModel(NomicBertPreTrainedModel):
     def forward(
         self,
         input_ids,
+        attention_mask=None,
         position_ids=None,
         token_type_ids=None,
-        attention_mask=None,
         return_dict=None,
         matryoshka_dim=None,
     ):