Update modeling_hf_nomic_bert.py

modeling_hf_nomic_bert.py

@@ -105,13 +105,7 @@ def filter_shapes(state_dict, model):
     return filtered_state_dict
 
 
-def remap_bert_state_dict(
-    state_dict,
-    config,
-    remove_bert=False,
-    remove_cls_weights=False,
-    add_pooling_layer=False,
-):
+def remap_bert_state_dict(state_dict, config, remove_bert=False, remove_cls_weights=False, add_pooling_layer=False):
     """
     Map the state_dict of a Huggingface BERT model to be flash_attn compatible.
     """
@@ -311,12 +305,13 @@ class NomicBertPreTrainedModel(PreTrainedModel):
         if config is None:
             config = cls.config_class.from_pretrained(model_name)
         remove_cls = cls != NomicBertForPreTraining
-        remove_bert_prefix = cls != NomicBertForPreTraining and cls != NomicBertForSequenceClassification
+        remove_bert_prefix = cls != NomicBertForPreTraining
         ignore_mismatched_shapes = kwargs.pop("ignore_mismatched_sizes", False)
         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 +320,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 +342,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,
@@ -355,7 +353,7 @@ class NomicBertPreTrainedModel(PreTrainedModel):
             if ignore_mismatched_shapes:
                 state_dict = filter_shapes(state_dict, model)
 
-            load_return = model.load_state_dict(state_dict, strict=strict)
+            load_return = model.load_state_dict(state_dict, strict=True)
         logger.warning(load_return)
         return model
 
@@ -726,7 +724,7 @@ class NomicBertAttention(nn.Module):
 
         self.rotary_emb_dim = self.head_dim * config.rotary_emb_fraction
         if self.rotary_emb_dim > 0:
-            if getattr(config, "rotary_scaling_factor", None):
+            if config.rotary_scaling_factor:
                 self.rotary_emb = NomicBertDynamicNTKRotaryEmbedding(
                     dim=self.rotary_emb_dim,
                     base=config.rotary_emb_base,
@@ -1057,11 +1055,10 @@ class NomicBertModel(NomicBertPreTrainedModel):
     def forward(
         self,
         input_ids,
-        position_ids=None,
-        token_type_ids=None,
         attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
         return_dict=None,
-        matryoshka_dim=None,
     ):
         if token_type_ids is None:
             token_type_ids = torch.zeros_like(input_ids)
@@ -1074,9 +1071,6 @@ class NomicBertModel(NomicBertPreTrainedModel):
 
         pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
 
-        if matryoshka_dim:
-            sequence_output = sequence_output[:, :matryoshka_dim]
-
         return BaseModelOutputWithPoolingAndCrossAttentions(
             last_hidden_state=sequence_output,
             pooler_output=pooled_output,
@@ -1224,4 +1218,4 @@ class NomicBertForSequenceClassification(NomicBertPreTrainedModel):
             logits=logits,
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
-        )
+        )